diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 95489e4ee..ac67e00ca 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -10,6 +10,7 @@ areas: github_team: "@ai-dynamo/access-aisimulate-maintain" path_globs: - AGENTS.md + - CLAUDE.md - CODE_OF_CONDUCT.md - DEVELOPMENT.md - LICENSE @@ -34,6 +35,7 @@ areas: - python/aisimulate/src/aisimulate/capacity.py - tests/test_aic.py - tests/test_fpm_profile_workflow.py + - tests/test_fpm_grouped_workflow.py - tests/test_fpm_query_evidence.py - label: sweeper github_team: "@ai-dynamo/aisimulate-sweeper-codeowners" @@ -162,6 +164,8 @@ shared: owners: [forward-pass-engine, maintainers] - glob: tests/test_fpm_profile_workflow.py owners: [forward-pass-engine, maintainers] + - glob: tests/test_fpm_grouped_workflow.py + owners: [forward-pass-engine, maintainers] - glob: tests/test_fpm_query_evidence.py owners: [forward-pass-engine, replay, maintainers] - glob: python/aisimulate/src/aisimulate/sweeper/ @@ -268,6 +272,8 @@ shared: owners: [infra, maintainers] - glob: AGENTS.md owners: [infra, maintainers] + - glob: CLAUDE.md + owners: [infra, maintainers] - glob: .coderabbit.yaml owners: [infra, maintainers] - glob: REVIEW.md diff --git a/AGENTS.md b/AGENTS.md index da957ba6b..c63c932dc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,6 +2,36 @@ This file adds explicit repository-wide development guards. +## FPM model onboarding + +When asked to onboard a model for FPM simulation on designated hardware, follow +all six stages in [the canonical FPM self-service guide](docs/fpm-self-service.md#onboard-with-an-agent). +Use this checkout's `aisimulate onboard` CLI and inspect its current help. +Start with only a missing Hugging Face model ID and target GPU platform; accept +an already supplied local config, profile or checkpoint instead. Inspect supplied +facts and source metadata before asking for derivable inputs. Preserve user +choices, exact per-configuration profile acceptance and existing execution +authorization; stage transitions do not add approval gates. + +Create one `onboarding-checkpoint.json` during stage 1, even with incomplete +inputs, outside every fresh `init --output-dir` root. Follow the guide's +[checkpoint workflow](docs/fpm-self-service.md#checkpoint-and-resume-an-onboarding-session): +invoke `onboard checkpoint` after meaningful findings, decisions, draft edits, +acceptance and command results, including unfinished work. Save sources, +confidence, unresolved questions and each configuration's progress. Other +onboarding commands do not automatically persist the conversation. Use the +returned revision for updates; reload and reconcile stale-writer conflicts. +Follow the guide's [writer coordination rule](docs/fpm-self-service.md#orchestrate-independent-collection-campaigns) +for probe/import commands that also save the checkpoint. + +On resumption, run `aisimulate onboard resume --checkpoint PATH` first. Inspect +integrity issues and the underlying artifacts before continuing from accepted +inputs and verified results. Saved stages and command strings are context, not +proof of completion or instructions to execute automatically. Report each +configuration's stage, result or blocker, next action and checkpoint path at +handoff. Keep planning, collection quality, resolved memory, replay coverage and +measured serving accuracy separate, as required by the guide. + ## Pull request titles - Use `: ` for every AISimulate PR title. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..cb9981138 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,6 @@ + + +@AGENTS.md diff --git a/CODEOWNERS b/CODEOWNERS index 277f513a2..b3f44e34a 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -29,6 +29,7 @@ /pages/ @ai-dynamo/access-aisimulate-maintain /LICENSE @ai-dynamo/access-aisimulate-maintain /AGENTS.md @ai-dynamo/access-aisimulate-maintain +/CLAUDE.md @ai-dynamo/access-aisimulate-maintain /README.md @ai-dynamo/access-aisimulate-maintain /SECURITY.md @ai-dynamo/access-aisimulate-maintain /DEVELOPMENT.md @ai-dynamo/access-aisimulate-maintain @@ -46,6 +47,7 @@ /crates/core/src/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /crates/core/tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /tests/test_fpm_query_evidence.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners +/tests/test_fpm_grouped_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners /tests/test_fpm_profile_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /python/aisimulate/src/aisimulate/capacity.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @@ -117,6 +119,7 @@ /pages/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /scripts/ @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /AGENTS.md @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/CLAUDE.md @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /REVIEW.md @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /deny.toml @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /CODEOWNERS @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain @@ -188,6 +191,7 @@ /python/aisimulate/.coderabbit.yaml @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/eslint.config.js @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_nightly.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_fpm_grouped_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_fpm_profile_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_resource_supervision.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/access-aisimulate-maintain /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain diff --git a/README.md b/README.md index 2bd24ab03..e10be7363 100644 --- a/README.md +++ b/README.md @@ -266,7 +266,7 @@ APIs: - [Estimator/FPE Python and Rust SDK](docs/core-api.md) - [FPM self-service: onboard a model on target hardware](docs/fpm-self-service.md) - [AIC-compatible modeled-power contract (semantics only)](docs/power-model.md) -- [Self-benchmarking and FPM onboarding guide](python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md) +- [FPM self-benchmarking and onboarding guide](python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md) - [Replay SDK and artifact contract](crates/core/src/replay/README.md) - [Sweeper SDK](docs/sweeper/overview.md) - [Legacy CLI reference](docs/cli/legacy-aic-user-guide.md) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index fa727a448..07aa9ab91 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -157,6 +157,29 @@ https://gitlab-master.nvidia.com/dl/sglang/sglang/-/blob/02c5a855aceb968c310e6fb ## vLLM +The grouped-cache config derivation and metadata regression tests in +`src/aisimulate/support/config_profile.py`, +`tests/unit/test_support_config_profile.py`, and +`tests/unit/test_support_topology.py` are adapted and modified from vLLM's +Inkling cache layout, decoder state construction, and configuration defaults at +commit `98dff2a81d747d1dba01a47f939f48c3526d4206`: + +- [`vllm/models/inkling/nvidia/sconv_swa_attn.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py) +- [`vllm/models/inkling/nvidia/model.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/model.py) +- [`vllm/models/inkling/configs.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/configs.py) + +Copyright contributors to the vLLM project. + +This material is licensed under Apache-2.0; the upstream license at the +identified revision is available at: +https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/LICENSE + +AISimulate modifies these contracts for metadata-only rank-local aggregate +cache sizing and validation, without constructing runtime tensors or a model. +The derived tests use original synthetic geometries and do not copy a +checkpoint fixture. The derived files retain the upstream copyright and +identify the modified source material. + The repository-root files `crates/core/src/engine/scheduler/vllm/core.rs`, `crates/core/src/engine/kv_manager/{state_cache_manager,vllm_backend}.rs`, `crates/core/src/engine/cache/vllm_block_pool.rs`, diff --git a/crates/core/src/engine/common/protocols.rs b/crates/core/src/engine/common/protocols.rs index 1f82d17fd..1a9e6d3a2 100644 --- a/crates/core/src/engine/common/protocols.rs +++ b/crates/core/src/engine/common/protocols.rs @@ -209,6 +209,10 @@ pub(crate) struct MockEngineArgs { pub kv_transfer_bytes_per_token: Option, #[builder(default = "None")] pub kv_cache_bytes_per_token: Option, + #[builder(default)] + pub kv_cache_groups: Vec, + #[builder(default = "None")] + pub kv_cache_capacity_bytes: Option, #[builder(default = "None")] pub native_host_offload: Option, #[builder(default = "None")] diff --git a/crates/core/src/engine/config.rs b/crates/core/src/engine/config.rs index 00a78e645..a7ce6b663 100644 --- a/crates/core/src/engine/config.rs +++ b/crates/core/src/engine/config.rs @@ -11,6 +11,7 @@ use serde::{Deserialize, Deserializer, Serialize}; use crate::engine::common::speculative::normalize_conditional_accept_rates; use crate::engine::handoff::TransferTimingMode; use crate::engine::timing::{TimingModel, TimingModelConfig, built_in_timing_model}; +use crate::perfmodel::FpmCacheGroup; const DEFAULT_MAX_PREFILL_TOKENS: usize = 16_384; const DEFAULT_CHUNKED_PREFILL_SIZE: usize = 8_192; @@ -371,7 +372,8 @@ pub struct EngineConfig { /// Use [`Self::for_backend`] instead of changing this field on /// [`Self::default`] when backend-dependent defaults are desired. pub backend: Backend, - /// Physical G1 capacity in blocks. + /// Physical G1 capacity in blocks for a linear cache. Retained as a + /// positive compatibility field and ignored when `kv_cache_groups` is set. #[serde(default = "default_num_gpu_blocks")] pub num_gpu_blocks: usize, /// KV block size in tokens. @@ -429,6 +431,12 @@ pub struct EngineConfig { /// Physical KV-cache bytes occupied by one token for host offload. #[serde(skip_serializing_if = "Option::is_none")] pub kv_cache_bytes_per_token: Option, + /// Rank-local cache group geometry. An empty list retains the linear pool. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub kv_cache_groups: Vec, + /// Shared physical budget for `kv_cache_groups`, including page padding. + #[serde(skip_serializing_if = "Option::is_none")] + pub kv_cache_capacity_bytes: Option, /// Optional framework-native host-offload simulation. #[serde(skip_serializing_if = "Option::is_none")] pub native_host_offload: Option, @@ -502,6 +510,10 @@ struct EngineConfigWire { kv_transfer_bytes_per_token: Option, #[serde(default, skip_serializing_if = "Option::is_none")] kv_cache_bytes_per_token: Option, + #[serde(default)] + kv_cache_groups: Vec, + #[serde(default)] + kv_cache_capacity_bytes: Option, #[serde(default, skip_serializing_if = "Option::is_none")] native_host_offload: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -563,6 +575,8 @@ impl<'de> Deserialize<'de> for EngineConfig { emit_kv_token_ids: wire.emit_kv_token_ids, kv_transfer_bytes_per_token: wire.kv_transfer_bytes_per_token, kv_cache_bytes_per_token: wire.kv_cache_bytes_per_token, + kv_cache_groups: wire.kv_cache_groups, + kv_cache_capacity_bytes: wire.kv_cache_capacity_bytes, native_host_offload: wire.native_host_offload, g3_offload: wire.g3_offload, state_cache: wire.state_cache, @@ -604,6 +618,8 @@ impl Default for EngineConfig { emit_kv_token_ids: false, kv_transfer_bytes_per_token: None, kv_cache_bytes_per_token: None, + kv_cache_groups: Vec::new(), + kv_cache_capacity_bytes: None, native_host_offload: None, g3_offload: None, state_cache: None, @@ -761,6 +777,49 @@ impl EngineConfig { self.kv_cache_bytes_per_token.is_none_or(|bytes| bytes > 0), "kv_cache_bytes_per_token must be positive" ); + if self.kv_cache_groups.is_empty() { + ensure!( + self.kv_cache_capacity_bytes.is_none(), + "kv_cache_capacity_bytes requires kv_cache_groups" + ); + } else { + ensure!( + self.backend == Backend::Vllm && self.worker_type == WorkerType::Aggregated, + "kv_cache_groups currently require backend=vllm and worker_type=aggregated" + ); + ensure!( + !self.enable_prefix_caching, + "kv_cache_groups currently require enable_prefix_caching=false" + ); + ensure!( + self.native_host_offload.is_none() && self.g3_offload.is_none(), + "kv_cache_groups currently support only HBM without host or G3 offload" + ); + ensure!( + self.aic_nextn.is_none(), + "kv_cache_groups currently do not support speculative decoding" + ); + ensure!( + !self.emit_kv_events && !self.emit_kv_token_ids, + "kv_cache_groups currently do not support KV event or token-ID emission" + ); + ensure!( + self.kv_cache_bytes_per_token.is_none() + && self.kv_transfer_bytes_per_token.is_none() + && self.kv_transfer_bandwidth.is_none(), + "kv_cache_groups cannot use linear cache or KV transfer byte geometry" + ); + ensure!( + self.kv_cache_capacity_bytes + .is_some_and(|bytes| bytes > 0 && bytes <= (1u64 << 53)), + "kv_cache_groups require positive kv_cache_capacity_bytes at most 2^53" + ); + let mut names = std::collections::HashSet::new(); + for group in &self.kv_cache_groups { + group.validate()?; + ensure!(names.insert(&group.name), "duplicate KV cache group name"); + } + } if let Some(g3) = &self.g3_offload { g3.validate()?; anyhow::ensure!( @@ -1184,6 +1243,63 @@ mod tests { assert!(error.to_string().contains("unknown field")); } + #[test] + fn grouped_cache_round_trips_and_rejects_unsupported_runtime_modes() { + let wire = serde_json::json!({ + "backend": "vllm", + "enable_prefix_caching": false, + "kv_cache_capacity_bytes": 1024, + "kv_cache_groups": [{ + "name": "window", "kind": "attention", "num_layers": 4, + "block_size_tokens": 4, "page_size_bytes": 64, "sliding_window": 8 + }] + }); + let config: EngineConfig = serde_json::from_value(wire.clone()).unwrap(); + config.validate().unwrap(); + let reloaded: EngineConfig = + serde_json::from_value(serde_json::to_value(&config).unwrap()).unwrap(); + assert_eq!(config, reloaded); + for (key, value, expected) in [ + ("backend", serde_json::json!("sglang"), "backend=vllm"), + ( + "worker_type", + serde_json::json!("decode"), + "worker_type=aggregated", + ), + ( + "enable_prefix_caching", + serde_json::json!(true), + "enable_prefix_caching=false", + ), + ("aic_nextn", serde_json::json!(1), "speculative"), + ("emit_kv_events", serde_json::json!(true), "KV event"), + ( + "kv_cache_bytes_per_token", + serde_json::json!(64), + "linear cache", + ), + ( + "kv_cache_capacity_bytes", + serde_json::json!(0), + "positive kv_cache_capacity_bytes", + ), + ( + "native_host_offload", + serde_json::json!({"num_host_blocks": 1}), + "only HBM", + ), + ] { + let mut unsupported = wire.clone(); + unsupported[key] = value; + let config: EngineConfig = serde_json::from_value(unsupported).unwrap(); + let error = config.validate().unwrap_err(); + assert!(error.to_string().contains(expected), "{key}: {error}"); + } + let linear = serde_json::to_value(EngineConfig::default()).unwrap(); + assert!(linear.get("kv_cache_groups").is_none()); + assert!(linear.get("kv_cache_capacity_bytes").is_none()); + } + #[test] fn native_host_offload_deserializes_with_default_bandwidths() { assert_eq!(DEFAULT_HOST_OFFLOAD_BANDWIDTH_GBPS, 32.0); diff --git a/crates/core/src/engine/kv_manager/g1_manager.rs b/crates/core/src/engine/kv_manager/g1_manager.rs index 80d794461..243eec9f0 100644 --- a/crates/core/src/engine/kv_manager/g1_manager.rs +++ b/crates/core/src/engine/kv_manager/g1_manager.rs @@ -18,7 +18,7 @@ use super::vllm_backend::{ StoreSourceSnapshot as VllmStoreSourceSnapshot, VllmAcquire, VllmKvManager, }; pub(crate) use super::vllm_backend::{NativeAllocation, SourceReuseDependency}; -use super::{AllocationRequirement, DestinationReservationMode, G1Acquire}; +use super::{AllocationRequirement, DestinationReservationMode, G1Acquire, GroupedKvPool}; fn into_g1_acquire(outcome: VllmAcquire) -> G1Acquire { match outcome { @@ -70,6 +70,14 @@ pub(crate) struct G1Manager { } impl G1Manager { + pub(crate) fn set_grouped_cache(&mut self, groups: GroupedKvPool) { + self.inner.set_grouped_cache(groups); + } + + pub(crate) fn grouped(&self) -> Option<&GroupedKvPool> { + self.inner.grouped() + } + pub(crate) fn set_belady_oracle(&mut self, oracle: BeladyOracle) { self.inner.set_belady_oracle(oracle); } diff --git a/crates/core/src/engine/kv_manager/grouped.rs b/crates/core/src/engine/kv_manager/grouped.rs new file mode 100644 index 000000000..e048e8c02 --- /dev/null +++ b/crates/core/src/engine/kv_manager/grouped.rs @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Cold, request-owned cache groups sharing one rank-local physical byte budget. + +use crate::perfmodel::FpmCacheGroup; + +#[derive(Debug, Default, PartialEq, Eq)] +pub(crate) struct GroupedLease { + /// Absolute logical block ranges, one per configured group. Full-history + /// groups start at zero; expired window blocks are absent. + ranges: Vec>, + bytes: u64, +} + +impl GroupedLease { + pub(crate) fn blocks(&self) -> usize { + self.ranges.iter().map(std::ops::Range::len).sum() + } +} + +pub(crate) struct GroupedKvPool { + groups: Vec, + capacity_bytes: u64, + used_bytes: u64, + active_blocks: usize, +} + +impl GroupedKvPool { + pub(crate) fn new(groups: Vec, capacity_bytes: u64) -> Self { + Self { + groups, + capacity_bytes, + used_bytes: 0, + active_blocks: 0, + } + } + + fn project(&self, computed_before: usize, computed_after: usize) -> Option { + if computed_after < computed_before { + return None; + } + let mut projected = GroupedLease::default(); + for group in &self.groups { + // A forward must retain the history needed by its first query, + // together with every token in the scheduled chunk. Freeing to + // the final query's window here would undercount prefill storage. + let range = group + .block_range(computed_before as u64, computed_after as u64) + .ok()?; + let range = range.start as usize..range.end as usize; + let bytes = (range.len() as u64).checked_mul(group.page_size_bytes)?; + projected.bytes = projected.bytes.checked_add(bytes)?; + projected.ranges.push(range); + } + Some(projected) + } + + pub(crate) fn required_bytes( + &self, + computed_before: usize, + computed_after: usize, + ) -> Option { + self.project(computed_before, computed_after) + .map(|lease| lease.bytes) + } + + /// Reserve every group atomically. Expired blocks become reusable in this + /// transition, but failed attempts leave both the lease and pool unchanged. + pub(crate) fn allocate( + &mut self, + lease: &mut Option, + computed_before: usize, + computed_after: usize, + ) -> bool { + let Some(projected) = self.project(computed_before, computed_after) else { + return false; + }; + let previous_bytes = lease.as_ref().map_or(0, |lease| lease.bytes); + let other_bytes = self.used_bytes - previous_bytes; + if projected.bytes > self.capacity_bytes - other_bytes { + return false; + } + self.used_bytes = other_bytes + projected.bytes; + self.active_blocks = self.active_blocks - lease.as_ref().map_or(0, GroupedLease::blocks) + + projected.blocks(); + *lease = Some(projected); + true + } + + pub(crate) fn release(&mut self, lease: &mut Option) { + if let Some(lease) = lease.take() { + self.used_bytes -= lease.bytes; + self.active_blocks -= lease.blocks(); + } + } + + pub(crate) fn capacity_bytes(&self) -> u64 { + self.capacity_bytes + } + + pub(crate) fn used_bytes(&self) -> u64 { + self.used_bytes + } + + pub(crate) fn active_blocks(&self) -> usize { + self.active_blocks + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::perfmodel::FpmCacheKind; + + #[test] + fn mixed_groups_reserve_atomically_and_reuse_expired_pages() { + let mut pool = GroupedKvPool::new( + vec![ + FpmCacheGroup { + name: "full".into(), + kind: FpmCacheKind::Attention, + num_layers: 1, + block_size_tokens: 4, + page_size_bytes: 8, + sliding_window: None, + }, + FpmCacheGroup { + name: "window".into(), + kind: FpmCacheKind::Attention, + num_layers: 1, + block_size_tokens: 2, + page_size_bytes: 6, + sliding_window: Some(3), + }, + ], + 42, + ); + let mut first = None; + let mut second = None; + assert!(pool.allocate(&mut first, 0, 4)); // 8 + 2 * 6 = 20 + assert!(pool.allocate(&mut second, 0, 2)); // 8 + 6 = 14 + assert!(!pool.allocate(&mut first, 4, 8)); // 16 + 3 * 6 = 34, plus 14 + assert_eq!(first.as_ref().unwrap().ranges, vec![0..1, 0..2]); + assert_eq!(pool.used_bytes(), 34); + pool.release(&mut second); + assert!(pool.allocate(&mut first, 4, 8)); + assert_eq!(first.as_ref().unwrap().ranges, vec![0..2, 1..4]); + assert_eq!(pool.used_bytes(), 34); + assert!(pool.allocate(&mut first, 8, 9)); + assert_eq!(first.as_ref().unwrap().ranges, vec![0..3, 3..5]); + assert_eq!(pool.used_bytes(), 36); + pool.release(&mut first); + pool.release(&mut first); + assert_eq!(pool.used_bytes(), 0); + assert_eq!(pool.active_blocks(), 0); + } +} diff --git a/crates/core/src/engine/kv_manager/mod.rs b/crates/core/src/engine/kv_manager/mod.rs index 200a66ac3..a111b8370 100644 --- a/crates/core/src/engine/kv_manager/mod.rs +++ b/crates/core/src/engine/kv_manager/mod.rs @@ -35,6 +35,7 @@ impl G1Acquire { } mod g1_manager; +mod grouped; pub(crate) mod sglang_backend; mod state_cache_manager; mod vllm_backend; @@ -44,6 +45,7 @@ mod vllm_firewall_tests; pub(crate) use g1_manager::{ DestinationReservation, G1Manager, NativeAllocation, SourceReuseDependency, }; +pub(crate) use grouped::GroupedKvPool; pub(crate) use sglang_backend::SglangKvManager; pub(crate) use vllm_backend::BlockRequestLease; diff --git a/crates/core/src/engine/kv_manager/vllm_backend.rs b/crates/core/src/engine/kv_manager/vllm_backend.rs index 1542d8989..33d3f9197 100644 --- a/crates/core/src/engine/kv_manager/vllm_backend.rs +++ b/crates/core/src/engine/kv_manager/vllm_backend.rs @@ -10,6 +10,7 @@ use uuid::Uuid; use super::AllocationRequirement; +use super::grouped::{GroupedKvPool, GroupedLease}; use super::state_cache_manager::{StateCacheManager, StateRequest}; use crate::engine::belady::BeladyOracle; pub(crate) use crate::engine::cache::vllm_block_pool::SourceReuseDependency; @@ -132,6 +133,7 @@ pub(crate) struct BlockRequestLease { /// Present only while newly acquired capacity awaits a scheduler-installed /// write fence. The default path does not allocate this state. pending_capacity_writes: Option>, + grouped: Option, state: Option>, } @@ -150,6 +152,7 @@ impl BlockRequestLease { allocated_tokens: 0, prefix_hashes: None, pending_capacity_writes: None, + grouped: None, state: None, } } @@ -205,6 +208,9 @@ impl BlockRequestLease { } pub(crate) fn resident_block_count(&self) -> usize { + if let Some(grouped) = &self.grouped { + return grouped.blocks(); + } self.entries .iter() .filter(|entry| entry.copy.is_some()) @@ -357,11 +363,21 @@ pub(crate) struct VllmKvManager { kv_event_publishers: KvEventPublishers, dp_rank: u32, next_event_id: u64, + grouped: Option, state_cache: Option, prefix_match_unit: Option, } impl VllmKvManager { + pub(crate) fn set_grouped_cache(&mut self, grouped: GroupedKvPool) { + assert!(!self.enable_prefix_caching); + self.grouped = Some(grouped); + } + + pub(crate) fn grouped(&self) -> Option<&GroupedKvPool> { + self.grouped.as_ref() + } + pub(crate) fn set_belady_oracle(&mut self, oracle: BeladyOracle) { self.pool.set_belady_oracle(oracle); } @@ -384,6 +400,7 @@ impl VllmKvManager { kv_event_publishers, dp_rank, next_event_id: 0, + grouped: None, state_cache: None, prefix_match_unit: None, } @@ -592,6 +609,22 @@ impl VllmKvManager { reusable_prefix_blocks: usize, ) -> NativeAllocation { lease.debug_assert_owner(owner); + if let Some(grouped) = &mut self.grouped { + assert_eq!(reusable_prefix_blocks, 0); + let previous_blocks = lease.resident_block_count(); + if !grouped.allocate( + &mut lease.grouped, + lease.allocated_tokens, + cumulative_tokens, + ) { + return NativeAllocation::CapacityExhausted; + } + lease.allocated_tokens = cumulative_tokens; + return NativeAllocation::Ready { + value: lease.resident_block_count().saturating_sub(previous_blocks), + dependencies: Vec::new(), + }; + } if self.state_cache.is_some() { return self.allocate_state_lease(lease, cumulative_tokens, reusable_prefix_blocks); } @@ -1271,6 +1304,9 @@ impl VllmKvManager { } fn release_lease_entries(&mut self, lease: &mut BlockRequestLease) { + if let Some(grouped) = &mut self.grouped { + grouped.release(&mut lease.grouped); + } if let (Some(manager), Some(state)) = (&mut self.state_cache, &mut lease.state) { manager.release(&mut self.pool, state); } @@ -1662,7 +1698,9 @@ impl VllmKvManager { } pub(crate) fn num_active_blocks(&self) -> usize { - self.pool.num_active() + self.grouped + .as_ref() + .map_or_else(|| self.pool.num_active(), GroupedKvPool::active_blocks) } pub(crate) fn num_inactive_blocks(&self) -> usize { diff --git a/crates/core/src/engine/protocol.rs b/crates/core/src/engine/protocol.rs index 12b3e6aeb..ffc5e940a 100644 --- a/crates/core/src/engine/protocol.rs +++ b/crates/core/src/engine/protocol.rs @@ -232,8 +232,14 @@ pub struct Metrics { /// currently populated by vLLM; SGLang reports zero because its legacy /// occupied count already includes radix-resident pages. pub inactive_blocks: u64, + /// Zero for grouped caches: heterogeneous page sizes have no single + /// block capacity. Use `kv_cache_capacity_bytes` in that case. pub total_blocks: u64, - /// `active_blocks / total_blocks`, with backend-native semantics above. + /// Physical grouped-cache occupancy, including scheduled chunk storage. + pub kv_cache_used_bytes: Option, + /// Physical grouped-cache capacity; absent for the legacy linear pool. + pub kv_cache_capacity_bytes: Option, + /// Byte occupancy for grouped caches; otherwise `active_blocks / total_blocks`. pub cache_usage: f64, /// Physical resident fraction. This includes inactive reusable vLLM blocks /// and equals `cache_usage` for SGLang's legacy occupied-page metric. diff --git a/crates/core/src/engine/scheduler/mod.rs b/crates/core/src/engine/scheduler/mod.rs index 23ac6f222..2c0721618 100644 --- a/crates/core/src/engine/scheduler/mod.rs +++ b/crates/core/src/engine/scheduler/mod.rs @@ -134,7 +134,11 @@ pub struct MockerMetrics { /// Resident reusable blocks excluded from `active_decode_blocks`. vLLM /// populates this; SGLang reports zero by construction. pub inactive_decode_blocks: u64, + /// Zero for a grouped cache, whose heterogeneous pages have no fixed + /// total block capacity. The byte fields are authoritative in that case. pub total_blocks: u64, + pub kv_cache_used_bytes: Option, + pub kv_cache_capacity_bytes: Option, pub gpu_cache_usage_perc: f64, pub physical_gpu_cache_usage_perc: f64, pub running_requests: u64, @@ -199,6 +203,8 @@ impl MockerMetrics { active_decode_blocks, inactive_decode_blocks, total_blocks, + kv_cache_used_bytes: None, + kv_cache_capacity_bytes: None, gpu_cache_usage_perc, physical_gpu_cache_usage_perc, running_requests, diff --git a/crates/core/src/engine/scheduler/rank.rs b/crates/core/src/engine/scheduler/rank.rs index 53e6c32a4..736580594 100644 --- a/crates/core/src/engine/scheduler/rank.rs +++ b/crates/core/src/engine/scheduler/rank.rs @@ -429,6 +429,8 @@ fn core_args(config: &EngineConfig, timing: Arc) -> MockEngineA aic_mtp_seed: config.aic_mtp_seed, kv_transfer_bytes_per_token: config.kv_transfer_bytes_per_token, kv_cache_bytes_per_token: config.kv_cache_bytes_per_token, + kv_cache_groups: config.kv_cache_groups.clone(), + kv_cache_capacity_bytes: config.kv_cache_capacity_bytes, native_host_offload: config.native_host_offload, state_cache: config.state_cache, prefix_match_unit: config.prefix_match_unit, @@ -576,6 +578,8 @@ fn map_metrics(metrics: MockerMetrics) -> Metrics { active_blocks: metrics.active_decode_blocks, inactive_blocks: metrics.inactive_decode_blocks, total_blocks: metrics.total_blocks, + kv_cache_used_bytes: metrics.kv_cache_used_bytes, + kv_cache_capacity_bytes: metrics.kv_cache_capacity_bytes, cache_usage: metrics.gpu_cache_usage_perc, physical_cache_usage: metrics.physical_gpu_cache_usage_perc, running_requests: metrics.running_requests, diff --git a/crates/core/src/engine/scheduler/vllm/core.rs b/crates/core/src/engine/scheduler/vllm/core.rs index 3973ef0bb..1306fb635 100644 --- a/crates/core/src/engine/scheduler/vllm/core.rs +++ b/crates/core/src/engine/scheduler/vllm/core.rs @@ -19,8 +19,9 @@ use crate::engine::common::speculative::{ use crate::engine::common::utils::{ compute_prefill_handoff_delay_ms, prefill_handoff_transfer_timing, }; -use crate::engine::kv_manager::{AllocationRequirement, G1Manager}; +use crate::engine::kv_manager::AllocationRequirement; use crate::engine::kv_manager::{DestinationReservation, G1Acquire, NativeAllocation}; +use crate::engine::kv_manager::{G1Manager, GroupedKvPool}; use crate::engine::scheduler::queue_metrics::{QueueStats, QueuedLength}; use crate::engine::scheduler::vllm::host_offload::{ CompletedLoad, HostLookup, StartLoad, VllmHostOffloadAdapter, VllmHostRequestState, @@ -805,20 +806,28 @@ impl VllmCore { ) .expect("validated native host-offload configuration must construct") }); + let mut kv_manager = G1Manager::new_with_caching( + args.num_gpu_blocks, + args.block_size, + kv_event_publishers, + dp_rank, + args.enable_prefix_caching, + ) + .with_state_cache( + args.state_cache, + args.kv_cache_bytes_per_token, + args.aic_nextn.is_some(), + ) + .with_prefix_match_unit(args.prefix_match_unit); + if !args.kv_cache_groups.is_empty() { + kv_manager.set_grouped_cache(GroupedKvPool::new( + args.kv_cache_groups.clone(), + args.kv_cache_capacity_bytes + .expect("validated grouped cache capacity"), + )); + } Self { - kv_manager: G1Manager::new_with_caching( - args.num_gpu_blocks, - args.block_size, - kv_event_publishers, - dp_rank, - args.enable_prefix_caching, - ) - .with_state_cache( - args.state_cache, - args.kv_cache_bytes_per_token, - args.aic_nextn.is_some(), - ) - .with_prefix_match_unit(args.prefix_match_unit), + kv_manager, belady_oracle: None, args, dp_rank, @@ -920,6 +929,14 @@ impl VllmCore { // fallible command is validated, violating RankEngine error atomicity. // While a model pass is in flight, also keep its physical transfer // completions hidden until complete_engine_boundary(). + anyhow::ensure!( + self.kv_manager.grouped().is_none() + || matches!( + &command, + SchedulerCommand::Submit(_) | SchedulerCommand::CancelRequest { .. } + ), + "grouped KV caches support only cold aggregated requests, not KV handoff commands" + ); let mutation_now_ms = command_now_ms.filter(|_| allow_destination_admission); match command { SchedulerCommand::Submit(mut request) => { @@ -1236,11 +1253,14 @@ impl VllmCore { /// logical `max_output_tokens`; scheduling separately enforces the model /// and physical KV limits. fn output_capacity_hint(&self, prompt_len: usize, max_output_tokens: usize) -> usize { - let kv_remaining = self - .args - .num_gpu_blocks - .saturating_mul(self.args.block_size) - .saturating_sub(prompt_len); + let kv_remaining = if self.kv_manager.grouped().is_some() { + max_output_tokens + } else { + self.args + .num_gpu_blocks + .saturating_mul(self.args.block_size) + .saturating_sub(prompt_len) + }; let model_remaining = self .args .max_model_len @@ -1654,7 +1674,7 @@ impl VllmCore { pub(crate) fn mocker_metrics(&self) -> MockerMetrics { let preactivation_destinations = self.pending_destinations.len() + self.destination_holds.len(); - MockerMetrics::from_parts_with_inactive( + let mut metrics = MockerMetrics::from_parts_with_inactive( self.dp_rank, self.kv_manager.num_active_blocks() as u64, self.kv_manager.num_inactive_blocks() as u64, @@ -1666,7 +1686,15 @@ impl VllmCore { self.state.preemptions_total, 0, 0, - ) + ); + if let Some(pool) = self.kv_manager.grouped() { + metrics.total_blocks = 0; + metrics.kv_cache_used_bytes = Some(pool.used_bytes()); + metrics.kv_cache_capacity_bytes = Some(pool.capacity_bytes()); + metrics.gpu_cache_usage_perc = pool.used_bytes() as f64 / pool.capacity_bytes() as f64; + metrics.physical_gpu_cache_usage_perc = metrics.gpu_cache_usage_perc; + } + metrics } #[cfg(test)] @@ -1780,7 +1808,7 @@ impl VllmCore { &mut batch_total_prefix, &mut preempted_any, now_ms, - ) { + )? { ScheduleOutcome::Scheduled { admission, .. } => { if let Some(admission) = admission { if let Some(collector) = collector.as_deref_mut() { @@ -1937,35 +1965,47 @@ impl VllmCore { unreachable!("pending destination stopped before host lookup") } AdmissionStage::FreshKv => { - let is_fresh = request.status == RequestStatus::Waiting; - let reserved_request_blocks = request - .host_offload - .as_deref() - .map(VllmHostRequestState::reserved_blocks) - .unwrap_or_default(); - let inflight_prefill_reserved_blocks = if host_hit.is_some() { - pass_prefill_reserved_blocks - } else { - 0 - }; - policy::decide_waiting_admission_with_cost( - policy::WaitingAdmissionConfig { - policy: scheduling_policy, - num_gpu_blocks: self.args.num_gpu_blocks, - block_size: self.args.block_size, - mtp_enabled: self.args.aic_nextn.is_some(), - }, + if let Some(decision) = policy::decide_grouped_waiting_admission( &request.sequence, - is_fresh, - running_seqs, &self.kv_manager, - reserved_request_blocks, - inflight_prefill_reserved_blocks, - held_completion_blocks, - activated_waiting_completion_blocks, + token_budget, + self.args.enable_chunked_prefill, raw_prefill_cost + .as_ref() .expect("fresh admission must retain its G1 lookup"), - ) + ) { + decision + } else { + let is_fresh = request.status == RequestStatus::Waiting; + let reserved_request_blocks = request + .host_offload + .as_deref() + .map(VllmHostRequestState::reserved_blocks) + .unwrap_or_default(); + let inflight_prefill_reserved_blocks = if host_hit.is_some() { + pass_prefill_reserved_blocks + } else { + 0 + }; + policy::decide_waiting_admission_with_cost( + policy::WaitingAdmissionConfig { + policy: scheduling_policy, + num_gpu_blocks: self.args.num_gpu_blocks, + block_size: self.args.block_size, + mtp_enabled: self.args.aic_nextn.is_some(), + }, + &request.sequence, + is_fresh, + running_seqs, + &self.kv_manager, + reserved_request_blocks, + inflight_prefill_reserved_blocks, + held_completion_blocks, + activated_waiting_completion_blocks, + raw_prefill_cost + .expect("fresh admission must retain its G1 lookup"), + ) + } } } } @@ -2064,7 +2104,7 @@ impl VllmCore { &mut batch_total_prefix, &mut preempted_any, now_ms, - ); + )?; self.update_pass_prefill_reservation( &mut pass_prefill_reserved_blocks, uuid, @@ -2304,12 +2344,7 @@ impl VllmCore { .state .requests .get(&selected) - .map(|request| { - request - .sequence - .num_allocated_tokens() - .div_ceil(self.args.block_size) - }) + .map(|request| request.sequence.lease.resident_block_count()) .unwrap_or_default(); if let (Some(adapter), Some(host)) = ( self.native_host_offload.as_mut(), @@ -2485,7 +2520,7 @@ impl VllmCore { batch_total_prefix: &mut usize, preempted_any: &mut bool, pressure_at_ms: f64, - ) -> ScheduleOutcome { + ) -> anyhow::Result { let request = self .state .requests @@ -2525,7 +2560,7 @@ impl VllmCore { && !self.args.enable_chunked_prefill && prompt_remaining > *token_budget { - return ScheduleOutcome::Blocked; + return Ok(ScheduleOutcome::Blocked); } let desired_tokens = self.align_state_prefill( @@ -2535,14 +2570,24 @@ impl VllmCore { remaining_known_tokens.min(*token_budget), ); if desired_tokens == 0 && remaining_known_tokens > 0 { - return if *token_budget > 0 { + return Ok(if *token_budget > 0 { ScheduleOutcome::SkippedForAlignment } else { ScheduleOutcome::Blocked - }; + }); } let desired_computed_after = effective_computed_before + desired_tokens; + if let Some(pool) = self.kv_manager.grouped() { + let required = pool + .required_bytes(effective_computed_before, desired_computed_after) + .ok_or_else(|| anyhow::anyhow!("grouped KV forward byte calculation overflowed"))?; + anyhow::ensure!( + required <= pool.capacity_bytes(), + "request {uuid} grouped KV forward needs {required} bytes for {desired_tokens} scheduled tokens after {effective_computed_before} computed tokens, exceeding the entire {}-byte cache pool; reduce max_num_batched_tokens or increase cache memory", + pool.capacity_bytes() + ); + } let mut actual_computed_after = desired_computed_after; let host_offload_enabled = self.native_host_offload.is_some(); @@ -2621,7 +2666,7 @@ impl VllmCore { } } if preempted.uuid == uuid { - return ScheduleOutcome::CurrentPreempted; + return Ok(ScheduleOutcome::CurrentPreempted); } } @@ -2630,7 +2675,7 @@ impl VllmCore { } let tokens_used = actual_computed_after.saturating_sub(effective_computed_before); if tokens_used == 0 && actual_computed_after < self.state.request_sequence_len(uuid) { - return ScheduleOutcome::Blocked; + return Ok(ScheduleOutcome::Blocked); } // vLLM's allocate_slots() caches full blocks through this request's @@ -2710,9 +2755,20 @@ impl VllmCore { } else { None }; - ScheduleOutcome::Scheduled { + Ok(ScheduleOutcome::Scheduled { tokens_used, admission, + }) + } + + /// The legacy timing argument is a logical-token clamp, not grouped + /// physical capacity. Never clip whole-forward queries to a retained + /// window or the unused linear pool's block count. + fn decode_kv_context_bound(&self, active_kv_tokens: usize) -> usize { + if self.kv_manager.grouped().is_some() { + active_kv_tokens + } else { + self.args.num_gpu_blocks * self.args.block_size } } @@ -2822,8 +2878,8 @@ impl VllmCore { if self.args.worker_type == WorkerType::Prefill || decode_count == 0 { (Duration::ZERO, decode_start_ms) } else { - let total_kv_tokens = self.args.num_gpu_blocks * self.args.block_size; let active_kv_tokens = total_length; + let total_kv_tokens = self.decode_kv_context_bound(active_kv_tokens); let context_length = total_length / decode_count; let decode_ms = self.args.perf_model.predict_decode_time( decode_count, @@ -3002,8 +3058,8 @@ impl VllmCore { let (decode_time, decode_end_ms) = if self.args.worker_type == WorkerType::Prefill { (Duration::ZERO, decode_start_ms) } else { - let total_kv_tokens = self.args.num_gpu_blocks * self.args.block_size; let active_kv_tokens = total_length; + let total_kv_tokens = self.decode_kv_context_bound(active_kv_tokens); let context_length = total_length / ready.len(); let decode_ms = self.args.perf_model.predict_decode_time( ready.len(), @@ -3922,7 +3978,10 @@ mod state_cache_tests { &mut preempted, 0.0, ); - assert!(matches!(outcome, ScheduleOutcome::Scheduled { .. })); + assert!(matches!( + outcome.unwrap(), + ScheduleOutcome::Scheduled { .. } + )); assert_eq!(core.state.requests[&id].num_computed_tokens, 9); assert_eq!(prefix, 8); assert!(!preempted); diff --git a/crates/core/src/engine/scheduler/vllm/grouped_tests.rs b/crates/core/src/engine/scheduler/vllm/grouped_tests.rs new file mode 100644 index 000000000..6b7bc3b5f --- /dev/null +++ b/crates/core/src/engine/scheduler/vllm/grouped_tests.rs @@ -0,0 +1,261 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; +use uuid::Uuid; + +use crate::engine::common::perf_model::PerfModel; +use crate::engine::common::protocols::{DirectRequest, MockEngineArgs}; +use crate::engine::scheduler::SchedulerCommand; +use crate::engine::trace::TraceCollector; +use crate::perfmodel::{FpmCacheGroup, FpmCacheKind}; + +use super::core::VllmCore; + +#[derive(Default)] +struct CapturedDecodeQueries(Mutex>); + +impl crate::engine::TimingModel for CapturedDecodeQueries { + fn prefill_batch_validation_can_fail(&self) -> bool { + false + } + + fn predict_prefill_ms(&self, _: usize, _: usize, _: usize) -> anyhow::Result { + Ok(0.0) + } + + fn predict_decode_ms( + &self, + _: usize, + active_tokens: usize, + mean_context: usize, + token_bound: usize, + ) -> anyhow::Result { + self.0 + .lock() + .unwrap() + .push((active_tokens, mean_context, token_bound)); + Ok(0.0) + } +} + +fn group(name: &str, block: u32, bytes: u64, window: Option) -> FpmCacheGroup { + FpmCacheGroup { + name: name.into(), + kind: FpmCacheKind::Attention, + num_layers: 1, + block_size_tokens: block, + page_size_bytes: bytes, + sliding_window: window, + } +} + +fn core(groups: Vec, capacity: u64, chunk: usize) -> VllmCore { + VllmCore::new( + MockEngineArgs::builder() + .block_size(4) + // Deliberately smaller than the prompts: this linear capacity + // must not decide admission or accounting for grouped caches. + .num_gpu_blocks(1) + .max_num_batched_tokens(Some(chunk)) + .max_num_seqs(Some(3)) + .enable_prefix_caching(false) + .kv_cache_groups(groups) + .kv_cache_capacity_bytes(Some(capacity)) + .speedup_ratio(0.0) + .build() + .unwrap(), + ) +} + +fn submit(core: &mut VllmCore, id: u128, prompt: usize, output: usize) -> Uuid { + core.receive(DirectRequest { + uuid: Some(Uuid::from_u128(id)), + tokens: (0..prompt as u32).collect(), + max_output_tokens: output, + ..Default::default() + }) +} + +#[test] +fn grouped_window_replays_long_context_with_bounded_storage_and_logical_fpm_lengths() { + let mut core = core(vec![group("window", 4, 8, Some(4))], 24, 8); + let captured = Arc::new(CapturedDecodeQueries::default()); + core.args.perf_model = Arc::new(PerfModel::External { + timing: captured.clone(), + }); + let id = submit(&mut core, 101, 32, 5); + let mut trace = TraceCollector::default(); + let mut outputs = 0; + let mut prefixes = Vec::new(); + let mut decode_contexts = Vec::new(); + let mut bytes = Vec::new(); + for step in 0..20 { + let pass = core.execute_pass(&mut trace, step as f64); + let fpm = pass.fpm.unwrap(); + if fpm.num_prefill_requests > 0 { + prefixes.push(fpm.sum_prefill_kv_tokens); + assert_eq!(fpm.sum_prefill_tokens, 8); + } + if fpm.num_decode_requests > 0 { + decode_contexts.push(fpm.sum_decode_kv_tokens); + } + bytes.push(pass.mocker_metrics.kv_cache_used_bytes.unwrap()); + assert_eq!(pass.mocker_metrics.total_blocks, 0); + assert_eq!(pass.mocker_metrics.kv_cache_capacity_bytes, Some(24)); + assert!(pass.mocker_metrics.physical_gpu_cache_usage_perc <= 1.0); + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0); + for output in pass.output_signals { + assert_eq!(output.uuid, id); + assert!(!output.rejected); + outputs += usize::from(output.token_id.is_some()); + } + if core.is_drained() { + break; + } + } + assert!(core.is_drained()); + assert_eq!(outputs, 5); + assert_eq!(prefixes, vec![0, 8, 16, 24]); + assert_eq!(decode_contexts, vec![33, 34, 35, 36]); + assert_eq!( + *captured.0.lock().unwrap(), + // Prefill produces output token one. The four decode forwards include + // their current input token and retain the full logical context. + vec![(33, 33, 33), (34, 34, 34), (35, 35, 35), (36, 36, 36)] + ); + assert_eq!(&bytes[..4], &[16, 24, 24, 24]); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); + assert_eq!(core.kv_manager.num_active_blocks(), 0); +} + +#[test] +fn grouped_full_window_and_convolution_charge_distinct_padded_pages() { + let mut convolution = group("convolution", 2, 2, Some(2)); + convolution.kind = FpmCacheKind::Convolution; + let mut core = core( + vec![ + group("full", 4, 8, None), + group("window", 4, 4, Some(4)), + convolution, + ], + 100, + 4, + ); + let id = submit(&mut core, 102, 16, 5); + let mut trace = TraceCollector::default(); + let mut bytes = Vec::new(); + for step in 0..5 { + core.execute_pass(&mut trace, step as f64); + bytes.push(core.mocker_metrics().kv_cache_used_bytes.unwrap()); + } + assert_eq!(bytes, vec![16, 30, 38, 46, 52]); + assert_eq!(core.state.requests[&id].num_computed_tokens, 17); + assert_eq!(core.state.requests[&id].sequence.num_allocated_tokens(), 17); + assert_eq!(core.state.requests[&id].sequence.len(), 18); + core.apply_command(SchedulerCommand::CancelRequest { request_id: id }) + .unwrap(); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); + assert_eq!(core.kv_manager.num_active_blocks(), 0); +} + +#[test] +fn grouped_contention_preempts_and_retries_without_leaking_capacity() { + let mut core = core( + vec![group("full", 4, 4, None), group("window", 4, 4, Some(4))], + 24, + 8, + ); + let first = submit(&mut core, 103, 4, 6); + let second = submit(&mut core, 104, 4, 6); + let mut trace = TraceCollector::default(); + let initial = core.execute_pass(&mut trace, 0.0); + assert_eq!(initial.admissions.len(), 2); + assert_eq!(initial.mocker_metrics.kv_cache_used_bytes, Some(16)); + let pressure = core.execute_pass(&mut trace, 1.0); + assert_eq!(pressure.mocker_metrics.vllm_preemptions_total, 1); + assert_eq!(pressure.mocker_metrics.kv_cache_used_bytes, Some(16)); + assert_eq!(core.state.requests[&second].num_computed_tokens, 0); + assert_eq!( + core.state.requests[&second].sequence.num_allocated_tokens(), + 0 + ); + assert_eq!(core.state.requests[&first].num_computed_tokens, 5); + let mut completed = Vec::new(); + for step in 2..40 { + let pass = core.execute_pass(&mut trace, step as f64); + assert!(pass.mocker_metrics.kv_cache_used_bytes.unwrap() <= 24); + for output in pass.output_signals { + assert!(!output.rejected); + if output.completed { + completed.push(output.uuid); + } + } + if core.is_drained() { + break; + } + } + assert_eq!(completed, vec![first, second]); + assert!(core.is_drained()); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); +} + +#[test] +fn grouped_waiting_request_cannot_preempt_running_work() { + let mut core = core(vec![group("window", 4, 4, Some(4))], 12, 8); + let first = submit(&mut core, 105, 8, 2); + let second = submit(&mut core, 106, 8, 2); + let mut trace = TraceCollector::default(); + let initial = core.execute_pass(&mut trace, 0.0); + assert_eq!(initial.admissions[0].uuid, first); + assert_eq!(core.state.requests[&second].num_computed_tokens, 0); + let next = core.execute_pass(&mut trace, 1.0); + assert!(next.admissions.is_empty()); + assert_eq!(next.mocker_metrics.vllm_preemptions_total, 0); + assert_eq!(next.mocker_metrics.kv_cache_used_bytes, Some(0)); + assert_eq!(core.state.requests[&second].num_computed_tokens, 0); + let admitted = core.execute_pass(&mut trace, 2.0); + assert_eq!(admitted.admissions[0].uuid, second); +} + +#[test] +fn grouped_oversized_chunk_fails_without_partial_allocation_or_endless_preemption() { + let mut core = core(vec![group("window", 4, 4, Some(4))], 8, 8); + let id = submit(&mut core, 107, 16, 2); + let mut trace = TraceCollector::default(); + core.execute_pass(&mut trace, 0.0); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(8)); + let error = core.try_execute_pass(&mut trace, 1.0).unwrap_err(); + assert!( + error + .to_string() + .contains("exceeding the entire 8-byte cache pool") + ); + assert!(error.to_string().contains("reduce max_num_batched_tokens")); + assert_eq!(core.state.requests[&id].num_computed_tokens, 8); + assert_eq!(core.state.requests[&id].sequence.num_allocated_tokens(), 8); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(8)); + assert_eq!(core.mocker_metrics().vllm_preemptions_total, 0); + core.apply_command(SchedulerCommand::CancelRequest { request_id: id }) + .unwrap(); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); +} + +#[test] +fn grouped_handoff_commands_fail_before_acquiring_state() { + let mut core = core(vec![group("window", 4, 4, Some(4))], 24, 8); + let handoff_id = crate::engine::HandoffId::new(Uuid::from_u128(108)); + let error = core + .apply_command(SchedulerCommand::SubmitHandoffPrefill { + handoff_id, + request: DirectRequest { + tokens: vec![0, 1, 2, 3], + max_output_tokens: 2, + ..Default::default() + }, + }) + .unwrap_err(); + assert!(error.to_string().contains("not KV handoff commands")); + assert!(core.is_drained()); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); +} diff --git a/crates/core/src/engine/scheduler/vllm/mod.rs b/crates/core/src/engine/scheduler/vllm/mod.rs index dff3bbef0..8e825ee7c 100644 --- a/crates/core/src/engine/scheduler/vllm/mod.rs +++ b/crates/core/src/engine/scheduler/vllm/mod.rs @@ -19,3 +19,6 @@ pub(crate) use core::RequestStatus; #[cfg(test)] mod tests; + +#[cfg(test)] +mod grouped_tests; diff --git a/crates/core/src/engine/scheduler/vllm/policy.rs b/crates/core/src/engine/scheduler/vllm/policy.rs index 991429aad..678acc137 100644 --- a/crates/core/src/engine/scheduler/vllm/policy.rs +++ b/crates/core/src/engine/scheduler/vllm/policy.rs @@ -81,6 +81,42 @@ pub(super) enum AdmissionDecision { Reject, } +/// Grouped cold-cache admission uses the pending forward's physical bytes. +/// Logical prompt length remains independent of bounded cache residency. +pub(super) fn decide_grouped_waiting_admission( + sequence: &S, + kv_manager: &G1Manager, + token_budget: usize, + enable_chunked_prefill: bool, + prefill_cost: &PrefillCost, +) -> Option { + let pool = kv_manager.grouped()?; + let required_tokens = sequence.len(); + // A known context whose smallest possible final forward cannot fit will + // never be admitted, regardless of chunking or competing requests. + let minimum = pool.required_bytes(required_tokens.saturating_sub(1), required_tokens); + if minimum.is_none_or(|bytes| bytes > pool.capacity_bytes()) { + return Some(AdmissionDecision::Reject); + } + if token_budget == 0 || (!enable_chunked_prefill && required_tokens > token_budget) { + return Some(AdmissionDecision::Wait); + } + let target = required_tokens.min(token_budget); + let needed = pool.required_bytes(0, target); + // A request whose configured chunk alone exceeds the entire pool must + // enter schedule_request, which reports an explicit configuration error. + // Contention alone preserves FIFO and must never preempt running work. + if needed.is_some_and(|bytes| { + bytes <= pool.capacity_bytes() && bytes > pool.capacity_bytes() - pool.used_bytes() + }) { + Some(AdmissionDecision::Wait) + } else { + Some(AdmissionDecision::Admit { + prefill_cost: prefill_cost.clone(), + }) + } +} + #[derive(Debug, Clone, Copy)] pub(super) struct WaitingAdmissionConfig { pub(super) policy: SchedulingPolicy, diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index d918d4afa..917951103 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -49,7 +49,8 @@ pub use perfmodel::{ }; pub use perfmodel::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, - OpLevelConfig, RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, + FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, }; #[cfg(feature = "python")] diff --git a/crates/core/src/perfmodel/engine/readiness.rs b/crates/core/src/perfmodel/engine/readiness.rs index b2aa3e8c2..f9e80d3fd 100644 --- a/crates/core/src/perfmodel/engine/readiness.rs +++ b/crates/core/src/perfmodel/engine/readiness.rs @@ -394,7 +394,7 @@ mod tests { (vec![op], vec![]) }; Engine::build(EngineSpec::new(config, context, generation), Arc::new(db))? - .validate_forward_pass_readiness() + .validate_forward_pass_readiness(crate::ForwardPassWorkerType::Aggregated) } #[test] diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index a3dd55eb2..9b2c7300f 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -22,8 +22,10 @@ use std::sync::Arc; +use crate::FpmQueryCoverage; use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; +use crate::operators::RuntimeContext; use crate::operators::base::PerformanceResult; use crate::operators::{FpmForwardOp, FpmPhase, Op}; use crate::perf_database::PerfDatabase; @@ -630,7 +632,10 @@ impl Engine { Ok(()) } - pub(crate) fn validate_forward_pass_readiness(&self) -> Result<(), AicError> { + pub(crate) fn validate_forward_pass_readiness( + &self, + worker_type: crate::ForwardPassWorkerType, + ) -> Result<(), AicError> { if self.prefill_graph_profile { self.db.prefill_graph.validate()?; self.db.prefill_graph.validate_sources(&self.db)?; @@ -638,7 +643,14 @@ impl Engine { } super::readiness::validate( &self.db, - self.context_ops.iter().chain(&self.generation_ops), + self.context_ops + .iter() + .filter(|_| worker_type != crate::ForwardPassWorkerType::Decode) + .chain( + self.generation_ops + .iter() + .filter(|_| worker_type != crate::ForwardPassWorkerType::Prefill), + ), ) } @@ -735,13 +747,14 @@ impl Engine { /// Python `_run_context_phase` (`base_backend.py:144`): `effective_isl = /// isl - prefix`, validate `> 0`, then one full pass over `context_ops`. fn run_context_phase(&self, runtime: &RuntimeConfig) -> Result { - self.run_context_phase_with(runtime, |_, _| {}) + self.run_context_phase_with(runtime, None, |_, _| {}) } fn run_context_phase_with( &self, runtime: &RuntimeConfig, - on_op: impl FnMut(&Op, PerformanceResult), + coverage: Option<&mut FpmQueryCoverage>, + mut on_op: impl FnMut(&Op, PerformanceResult), ) -> Result { // Python raises `ValueError` when `effective_isl <= 0`; mirror that. if runtime.prefix >= runtime.isl { @@ -751,6 +764,34 @@ impl Engine { ))); } let effective_isl = runtime.isl - runtime.prefix; + if let Some(coverage) = coverage { + let Some((prefill, _, ctx_tail, gen_tail)) = self.fpm_split() else { + return Err(AicError::InvalidEngineConfig( + "query coverage requires a direct FPM engine".into(), + )); + }; + if !ctx_tail.is_empty() || !gen_tail.is_empty() || self.nextn > 0 { + return Err(AicError::InvalidEngineConfig( + "query coverage does not support speculative FPM".into(), + )); + } + let result = prefill.query_with_coverage( + &self.db, + &RuntimeContext { + batch_size: runtime.batch_size, + beam_width: 1, + s: effective_isl, + prefix: runtime.prefix, + num_tokens: runtime.batch_size * effective_isl, + seq_imbalance_correction_scale: runtime.seq_imbalance_correction_scale, + ..Default::default() + }, + Some(coverage), + )?; + let latency_ms = result.latency_ms; + on_op(&self.context_ops[0], result); + return Ok(latency_ms); + } run_context_ops_with( &self.context_ops, &self.db, @@ -847,6 +888,16 @@ impl Engine { /// Thin shim over [`Self::run_static`] with `mode=Context` (osl is /// irrelevant for the context phase, so it is fixed at 1). pub fn predict_prefill_latency(&self, bs: u32, isl: u32, prefix: u32) -> Result { + self.predict_prefill_latency_with_coverage(bs, isl, prefix, None) + } + + pub(crate) fn predict_prefill_latency_with_coverage( + &self, + bs: u32, + isl: u32, + prefix: u32, + coverage: Option<&mut FpmQueryCoverage>, + ) -> Result { let rt = RuntimeConfig { batch_size: bs, isl, @@ -856,12 +907,8 @@ impl Engine { }; if self.prefill_graph_profile { crate::perf_database::prefill_graph::public_shape(bs, isl, prefix)?; - self.run_context_phase(&rt) - } else { - Ok(self - .run_static(&rt, StaticMode::Context, DEFAULT_STATIC_STRIDE)? - .total_ms) } + self.run_context_phase_with(&rt, coverage, |_, _| {}) } /// Preserve static prefill geometry, including explicitly equal per-request @@ -872,6 +919,7 @@ impl Engine { bs: u32, isl: u32, prefix: u32, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { let runtime = RuntimeConfig { batch_size: bs, @@ -881,7 +929,7 @@ impl Engine { ..Default::default() }; let mut queries = Vec::new(); - let prefill_ms = self.run_context_phase_with(&runtime, |_, result| { + let prefill_ms = self.run_context_phase_with(&runtime, coverage, |_, result| { queries.extend(result.fpm_queries.into_iter().flat_map(|queries| *queries)); })?; let latency_ms = prefill_ms + 0.0; // same context + generation sum as run_static @@ -925,16 +973,28 @@ impl Engine { &self, batch_size: u32, total_past_kv_tokens: u32, + ) -> Result { + self.predict_decode_latency_total_with_coverage(batch_size, total_past_kv_tokens, None) + } + + pub(crate) fn predict_decode_latency_total_with_coverage( + &self, + batch_size: u32, + total_past_kv_tokens: u32, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { self.require_general_engine()?; - self.forward_pass_time_ms(&[ForwardPassMetrics { - scheduled_requests: crate::ScheduledRequestMetrics { - num_decode_requests: batch_size, - sum_decode_kv_tokens: total_past_kv_tokens, + self.forward_pass_time_ms_with_coverage( + &[ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_decode_requests: batch_size, + sum_decode_kv_tokens: total_past_kv_tokens, + ..Default::default() + }, ..Default::default() - }, - ..Default::default() - }]) + }], + coverage, + ) } /// Highest decode KV-read total covered by a compiled FPM engine. @@ -2201,7 +2261,15 @@ impl Engine { &self, metrics_by_rank: &[ForwardPassMetrics], ) -> Result { - self.forward_pass_estimate(metrics_by_rank, false) + self.forward_pass_time_ms_with_coverage(metrics_by_rank, None) + } + + pub(crate) fn forward_pass_time_ms_with_coverage( + &self, + metrics_by_rank: &[ForwardPassMetrics], + coverage: Option<&mut FpmQueryCoverage>, + ) -> Result { + self.forward_pass_estimate(metrics_by_rank, false, coverage) .map(|estimate| estimate.native_latency_ms.unwrap()) } @@ -2209,6 +2277,7 @@ impl Engine { &self, metrics_by_rank: &[ForwardPassMetrics], capture: bool, + mut coverage: Option<&mut FpmQueryCoverage>, ) -> Result { self.require_general_engine()?; if metrics_by_rank.is_empty() { @@ -2227,7 +2296,11 @@ impl Engine { rank, ..Default::default() }; - let rank_latency = self.rank_latency_ms(metrics, capture.then_some(&mut evidence))?; + let rank_latency = self.rank_latency_ms( + metrics, + capture.then_some(&mut evidence), + coverage.as_deref_mut(), + )?; if rank_latency > max_latency { max_latency = rank_latency; max_rank = Some(rank); @@ -2256,6 +2329,7 @@ impl Engine { &self, metrics: &ForwardPassMetrics, mut evidence: Option<&mut crate::FpmRankEstimate>, + mut coverage: Option<&mut FpmQueryCoverage>, ) -> Result { let sched = &metrics.scheduled_requests; // Token-based dispatch, aligned with `IterationFeatures` (fpm/model.rs): @@ -2321,13 +2395,14 @@ impl Engine { // division on each axis. let mut total = 0.0_f64; if has_prefill { - let result = prefill_op.query_totals( + let result = prefill_op.query_totals_with_coverage( &self.db, &[ sched.num_prefill_requests as f64, sched.sum_prefill_tokens as f64, sched.sum_prefill_kv_tokens as f64, ], + coverage.as_deref_mut(), )?; total += result.latency_ms; if let Some(evidence) = evidence.as_deref_mut() { @@ -2338,12 +2413,13 @@ impl Engine { } } if has_decode { - let result = decode_op.query_totals( + let result = decode_op.query_totals_with_coverage( &self.db, &[ sched.num_decode_requests as f64, sched.sum_decode_kv_tokens as f64, ], + coverage.as_deref_mut(), )?; let decode_ms = result.latency_ms; if let Some(evidence) = evidence.as_deref_mut() { @@ -2356,10 +2432,11 @@ impl Engine { // Mixed rank: marginal-decode composition, mirroring // `_get_fpm_mix_step_latency` (counts already packed, no // `(nextn + 1)` — speculative FPM was rejected above). - let baseline = decode_op.query_pass_baseline( + let baseline = decode_op.query_pass_baseline_with_coverage( &self.db, sched.num_decode_requests, sched.sum_decode_kv_tokens as f64, + coverage.as_deref_mut(), )?; let baseline_ms = baseline.latency_ms; total += (decode_ms - baseline_ms).max(0.0); @@ -3992,7 +4069,7 @@ mod tests { 43.0 ); let detailed = engine - .predict_prefill_detailed(input.batch_size, input.isl, input.prefix) + .predict_prefill_detailed(input.batch_size, input.isl, input.prefix, None) .unwrap(); assert_eq!(detailed.latency_ms, Some(43.0)); assert_eq!( @@ -4005,12 +4082,16 @@ mod tests { // telemetry guard for multi-request bounded decoder replay. assert!( engine - .forward_pass_estimate(&[fpm_replay_guard_metrics(0.0)], true) + .forward_pass_estimate(&[fpm_replay_guard_metrics(0.0)], true, None) .is_err() ); for (isl, prefix) in [(0, 0), (512, 512)] { assert!(engine.predict_prefill_latency(2, isl, prefix).is_err()); - assert!(engine.predict_prefill_detailed(2, isl, prefix).is_err()); + assert!( + engine + .predict_prefill_detailed(2, isl, prefix, None) + .is_err() + ); } } @@ -4861,7 +4942,8 @@ mod tests { "0.25.1", )); let engine = Engine::build(spec, Arc::new(db)).unwrap(); - let result = engine.validate_forward_pass_readiness(); + let result = + engine.validate_forward_pass_readiness(crate::ForwardPassWorkerType::Aggregated); if missing_gemm { assert!( result diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index eacde9f89..c8641e9ca 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -271,6 +271,7 @@ impl ForwardPassPerfModelConfig { self.kvcache_quant_mode = Some(facts.kvcache_quant_mode); self.comm_quant_mode = Some(facts.comm_quant_mode); self.attention_backend = Some(facts.attention_backend); + self.fpm_resources()?; Some(facts.registered) } #[cfg(not(feature = "python"))] @@ -297,6 +298,16 @@ impl ForwardPassPerfModelConfig { } pub(crate) fn validate(&self) -> Result<(), AicError> { + if self.estimator_config.fpm_interpolation.collect_coverage + && (self.estimation_mode != EstimationMode::FpmInterpolation + || self.fallback_policy != ForwardPassFallbackPolicy::Deny + || self.estimator_config.fpm_interpolation.method + != super::FpmInterpolationMethod::Direct) + { + return Err(invalid_config( + "collect_coverage requires estimation_mode='fpm_interpolation', method='direct', and fallback_policy='deny'", + )); + } self.validate_prefill_graph_profile()?; if let Some(distribution) = &self.estimator_config.op_level.decode_workload_distribution { let selected = @@ -572,6 +583,66 @@ mod tests { ); } + #[test] + fn coverage_controls_preserve_default_serialization_and_require_strict_direct() { + let default = serde_json::to_value(EstimatorConfig::default()).unwrap(); + assert_eq!( + default["fpm_interpolation"], + serde_json::json!({"method": "auto"}) + ); + let cfg = config(serde_json::json!({ + "estimation_mode": "fpm_interpolation", + "estimator_config": {"fpm_interpolation": {"method": "direct", "collect_coverage": true}} + })); + cfg.validate().unwrap(); + let reloaded: ForwardPassPerfModelConfig = + serde_json::from_value(serde_json::to_value(&cfg).unwrap()).unwrap(); + assert_eq!(cfg, reloaded); + for mode in [ + EstimationMode::Auto, + EstimationMode::OpLevel, + EstimationMode::FpmRegression, + ] { + let mut invalid = cfg.clone(); + invalid.estimation_mode = mode; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + for method in [ + super::super::FpmInterpolationMethod::Auto, + super::super::FpmInterpolationMethod::Sol, + ] { + let mut invalid = cfg.clone(); + invalid.estimator_config.fpm_interpolation.method = method; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + for fallback in [ + ForwardPassFallbackPolicy::Allow, + ForwardPassFallbackPolicy::LegacyRegression, + ] { + let mut invalid = cfg.clone(); + invalid.fallback_policy = fallback; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + } + #[test] fn fpm_selector_is_validated_and_preserved_in_canonical_identity() { let cfg = config(serde_json::json!({ diff --git a/crates/core/src/perfmodel/fpm/coverage.rs b/crates/core/src/perfmodel/fpm/coverage.rs new file mode 100644 index 000000000..2eecba262 --- /dev/null +++ b/crates/core/src/perfmodel/fpm/coverage.rs @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Bounded evidence from actual direct-FPM lookups. This is not a replay +//! completion or prediction-accuracy assessment. + +use std::collections::BTreeMap; +use std::sync::{Mutex, MutexGuard}; + +use serde::{Deserialize, Serialize}; + +use crate::AicError; +use crate::operators::fpm_forward::{FpmForwardOp, FpmPhase}; +use crate::perf_database::fpm_forward::{ + FPM_CELL_MATCH_COLUMNS, FPM_DECODE_AXES, FPM_PREFILL_AXES, +}; + +const GAP_LIMIT: usize = 128; + +/// Counts of native table lookup resolutions, not requests or iterations. +/// Reusing a cached timing does not perform another native lookup. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct FpmQueryCoverageCounts { + pub measured: u64, + pub interpolated: u64, + pub unsupported: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmQueryPurpose { + ForwardPass, + /// A mixed pass subtracts the measured decode curve floor, or an + /// interpolation of floors, using the same covered rows as decode. + MixedDecodeBaseline, +} + +/// One unsupported query, retained after the estimation call returns its error. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmQueryGap { + pub phase: String, + pub purpose: FpmQueryPurpose, + pub model_path: String, + pub cell_identity: BTreeMap, + pub cell_ids: Vec, + pub coordinates: BTreeMap, + pub reason: String, + pub occurrences: u64, +} + +/// Cumulative evidence for one model instance. A snapshot cannot certify that +/// replay completed: the replay consumer must report completion separately. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmQueryCoverage { + /// Explicit unit prevents callers from interpreting these as replay counts. + pub counting_unit: String, + pub queries: FpmQueryCoverageCounts, + pub prefill: FpmQueryCoverageCounts, + pub decode: FpmQueryCoverageCounts, + pub mixed_decode_baseline: FpmQueryCoverageCounts, + pub gaps: Vec, + pub gap_limit: usize, + /// Failed lookup occurrences omitted after the distinct-gap limit. + pub omitted_gap_queries: u64, +} + +impl Default for FpmQueryCoverage { + fn default() -> Self { + Self { + counting_unit: "native_lookup_resolutions".into(), + queries: Default::default(), + prefill: Default::default(), + decode: Default::default(), + mixed_decode_baseline: Default::default(), + gaps: Vec::new(), + gap_limit: GAP_LIMIT, + omitted_gap_queries: 0, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum DirectFpmResolution { + Measured, + Interpolated, +} + +impl FpmQueryCoverageCounts { + fn record(&mut self, result: Result) { + match result { + Ok(DirectFpmResolution::Measured) => self.measured += 1, + Ok(DirectFpmResolution::Interpolated) => self.interpolated += 1, + Err(_) => self.unsupported += 1, + } + } +} + +impl FpmQueryCoverage { + pub(crate) fn record( + &mut self, + op: &FpmForwardOp, + cell_ids: &[String], + coords: &[f64], + purpose: FpmQueryPurpose, + result: Result, + ) { + self.queries.record(result); + match purpose { + FpmQueryPurpose::MixedDecodeBaseline => &mut self.mixed_decode_baseline, + FpmQueryPurpose::ForwardPass => match op.phase { + FpmPhase::Prefill => &mut self.prefill, + FpmPhase::Decode => &mut self.decode, + }, + } + .record(result); + let Err(error) = result else { return }; + let axes: &[&str] = match op.phase { + FpmPhase::Prefill => &FPM_PREFILL_AXES, + FpmPhase::Decode => &FPM_DECODE_AXES, + }; + let mut gap = FpmQueryGap { + phase: op.phase.as_str().into(), + purpose, + model_path: op.model_path.clone(), + cell_identity: FPM_CELL_MATCH_COLUMNS + .iter() + .zip(&op.match_identity) + .map(|(name, value)| ((*name).into(), value.clone())) + .collect(), + cell_ids: cell_ids.to_vec(), + coordinates: axes + .iter() + .zip(coords) + .map(|(k, v)| ((*k).into(), *v)) + .collect(), + reason: error.to_string(), + occurrences: 0, + }; + if let Some(existing) = self.gaps.iter_mut().find(|entry| { + entry.phase == gap.phase + && entry.purpose == gap.purpose + && entry.model_path == gap.model_path + && entry.cell_identity == gap.cell_identity + && entry.cell_ids == gap.cell_ids + && entry.coordinates == gap.coordinates + && entry.reason == gap.reason + }) { + existing.occurrences += 1; + } else if self.gaps.len() < GAP_LIMIT { + gap.occurrences = 1; + self.gaps.push(gap); + } else { + self.omitted_gap_queries += 1; + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct FpmCoverageState(Mutex); + +impl Clone for FpmCoverageState { + fn clone(&self) -> Self { + // Model clones may be used for another replay/candidate. The loaded + // engine is shared, but diagnostics must start empty for each clone. + Self::default() + } +} + +impl FpmCoverageState { + pub(crate) fn lock(&self) -> Result, AicError> { + self.0.lock().map_err(|_| { + AicError::InvalidEngineConfig("FPM query coverage accumulator was poisoned".into()) + }) + } +} diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index 485622bc2..b2c85bba5 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -44,6 +44,10 @@ pub struct FpmInterpolationConfig { /// External parquet and its same-stem metadata sidecar. #[serde(skip_serializing_if = "Option::is_none")] pub fpm_parquet_path: Option, + /// Record bounded query coverage on the returned canonical model. + /// Requires explicit direct FPM with fallback denied. + #[serde(skip_serializing_if = "is_false")] + pub collect_coverage: bool, /// Match null profile identities only for these unspecified quant modes. #[serde(skip_serializing_if = "Vec::is_empty")] pub unrecorded_quant_modes: Vec, @@ -489,7 +493,9 @@ fn explicit_fields( fields .iter() .map(|(key, value)| { - // Compact serde output may omit an explicit default; keep it supplied. + // Defaults omitted by serialization remain explicit input, + // such as collect_coverage=false. The typed parse above has + // already validated every supplied value. ( key.clone(), explicit_fields(value, normalized.get(key).unwrap_or(value)), @@ -591,6 +597,19 @@ mod migration_tests { assert_eq!(reloaded, config); } + #[test] + fn migration_preserves_explicit_coverage_defaults() { + for collect_coverage in [false, true] { + let saved = format!( + r#"{{"fpm_interpolation": {{"method": "direct", "collect_coverage": {collect_coverage}, "fpm_parquet_path": null}}}}"# + ); + let config = + EstimatorConfig::migrate_legacy_inputs(Some(&saved), None, None, None).unwrap(); + assert_eq!(config.fpm_interpolation.collect_coverage, collect_coverage); + assert_eq!(config.fpm_interpolation.fpm_parquet_path, None); + } + } + #[test] fn migration_accepts_equal_numeric_array_and_null_controls() { let saved = r#"{ diff --git a/crates/core/src/perfmodel/fpm/mod.rs b/crates/core/src/perfmodel/fpm/mod.rs index 945862f9c..992e50bad 100644 --- a/crates/core/src/perfmodel/fpm/mod.rs +++ b/crates/core/src/perfmodel/fpm/mod.rs @@ -24,12 +24,14 @@ mod config; mod correction; +pub(crate) mod coverage; mod estimator; mod evidence; mod metrics; mod model; mod options; mod regression; +mod resources; mod samples; #[cfg(test)] @@ -39,6 +41,7 @@ pub use config::{ EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig, ForwardPassSpeculationConfig, }; +pub use coverage::{FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose}; pub use estimator::*; pub(crate) use evidence::accumulate_fpm_estimates; pub use evidence::*; @@ -50,3 +53,7 @@ pub use model::{ ForwardPassRegressionWorkloadKind, ForwardPassWorkerType, }; pub use options::ForwardPassPerfOptions; +pub use resources::{ + FpmCacheBudget, FpmCacheBudgetAdjusted, FpmCacheBudgetRequest, FpmCacheGroup, FpmCacheKind, + FpmCacheLayout, FpmResourceConfig, FpmRuntimeMemoryConfig, +}; diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 4e14e1b67..f7de86b9c 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -18,6 +18,7 @@ use crate::{AicError, ForwardPassMetrics}; use super::config::{EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig}; use super::correction::CorrectionBuckets; +use super::coverage::{FpmCoverageState, FpmQueryCoverage}; use super::metrics::validate_forward_pass_metrics; use super::options::{ForwardPassPerfOptions, validate_regression_options}; use super::regression::BucketedRegression; @@ -248,6 +249,7 @@ pub struct ForwardPassPerfModel { last_warning: Option, provenance: Option, is_correction_enabled: bool, + pub(super) query_coverage: Option, } #[derive(Clone, Debug)] @@ -282,6 +284,7 @@ impl ForwardPassPerfModel { last_warning: None, provenance: None, is_correction_enabled: true, + query_coverage: None, } } @@ -308,6 +311,7 @@ impl ForwardPassPerfModel { last_warning: None, provenance: None, is_correction_enabled: true, + query_coverage: None, }) } @@ -392,6 +396,11 @@ impl ForwardPassPerfModel { candidate.estimator_config.correction_options(), ); model.is_correction_enabled = candidate.estimator_config.correction.enabled; + model.query_coverage = candidate + .estimator_config + .fpm_interpolation + .collect_coverage + .then(FpmCoverageState::default); model.last_warning = (!failures.is_empty()).then(|| failures.join("; ")); model.provenance = Some(ForwardPassPerfProvenance { config: candidate, @@ -447,7 +456,13 @@ impl ForwardPassPerfModel { let Some(feature) = IterationFeatures::from_metrics(metrics_by_rank)? else { return Ok(Some(0.0)); }; - let native = engine.forward_pass_time_ms(metrics_by_rank)?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + let native = engine + .forward_pass_time_ms_with_coverage(metrics_by_rank, coverage.as_deref_mut())?; let corrected = native * self.correction_factor(corrections, &feature); Ok(Some(corrected)) } @@ -470,7 +485,8 @@ impl ForwardPassPerfModel { /// Estimate the same iteration with executed direct-FPM support. The /// native rank maximum and online correction are explicit; measurement - /// weights describe raw lookup values only. No query state is retained. + /// weights describe raw lookup values only. Detailed query state is returned; + /// bounded coverage is retained only when explicitly enabled. pub fn estimate_forward_pass_detailed( &self, metrics_by_rank: &[ForwardPassMetrics], @@ -490,7 +506,13 @@ impl ForwardPassPerfModel { ranks: Vec::new(), }); }; - let mut estimate = engine.forward_pass_estimate(metrics_by_rank, true)?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + let mut estimate = + engine.forward_pass_estimate(metrics_by_rank, true, coverage.as_deref_mut())?; let factor = self.correction_factor(corrections, &feature); estimate.correction_factor = Some(factor); estimate.latency_ms = estimate.native_latency_ms.map(|native| native * factor); @@ -541,6 +563,7 @@ impl ForwardPassPerfModel { mode, options, is_correction_enabled, + query_coverage, .. } = self; for metrics_by_rank in iterations { @@ -557,7 +580,14 @@ impl ForwardPassPerfModel { else { continue; }; - let native = engine.forward_pass_time_ms(metrics_by_rank)?; + let mut coverage = query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + let native = engine.forward_pass_time_ms_with_coverage( + metrics_by_rank, + coverage.as_deref_mut(), + )?; corrections .store_mut(observation.feature.workload_kind) .add_observation(observation.feature.x, observation.wall_time_ms, native); @@ -691,6 +721,75 @@ impl ForwardPassPerfModel { &self.options } + /// Uncorrected static prefill latency, preserving the compiled engine's + /// existing `(batch, full input length, cached prefix)` contract. Unlike + /// telemetry estimation this does not apply learned online correction. + /// Enabled coverage records the actual native lookup on this model. + pub fn predict_prefill_latency( + &self, + batch_size: u32, + isl: u32, + prefix: u32, + ) -> Result { + let engine = self.require_native_engine()?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + engine.predict_prefill_latency_with_coverage( + batch_size, + isl, + prefix, + coverage.as_deref_mut(), + ) + } + + /// Uncorrected decode latency using exact past-KV totals, excluding the + /// current input token per request. Shares the existing engine contract. + pub fn predict_decode_latency_total( + &self, + batch_size: u32, + total_past_kv_tokens: u32, + ) -> Result { + let engine = self.require_native_engine()?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + engine.predict_decode_latency_total_with_coverage( + batch_size, + total_past_kv_tokens, + coverage.as_deref_mut(), + ) + } + + /// Largest collected decode KV total for the selected FPM cell. Reading + /// this bound is not a timing lookup and does not add coverage evidence. + pub fn fpm_decode_kv_ceiling(&self) -> Result, AicError> { + self.require_native_engine()?.fpm_decode_kv_ceiling() + } + + /// Snapshot of actual native direct-FPM lookup resolutions, including + /// failed lookups. `None` means collection was disabled. Counts do not + /// include timings reused by an external cache, or calls on a separate + /// low-level engine handle. Replay completion must be checked separately. + pub fn fpm_query_coverage(&self) -> Result, AicError> { + self.query_coverage + .as_ref() + .map(|state| state.lock().map(|value| value.clone())) + .transpose() + } + + fn require_native_engine(&self) -> Result, AicError> { + self.native_engine().ok_or_else(|| { + AicError::InvalidEngineConfig( + "static timing requires a native forward-pass estimator".into(), + ) + }) + } + /// Static phase latency before online correction, using the native engine's /// existing integration. Decode returns the total for all generated tokens. pub fn static_phase_latency( @@ -718,11 +817,15 @@ impl ForwardPassPerfModel { prefix: u32, ) -> Result { self.require_general_forward_api()?; - self.native_engine() - .ok_or_else(|| { - AicError::InvalidEngineConfig("static prefill requires a native estimator".into()) - })? - .predict_prefill_detailed(batch_size, input_tokens, prefix) + let engine = self.native_engine().ok_or_else(|| { + AicError::InvalidEngineConfig("static prefill requires a native estimator".into()) + })?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + engine.predict_prefill_detailed(batch_size, input_tokens, prefix, coverage.as_deref_mut()) } /// Native operation evidence for one static prefill or decode step. Values @@ -755,27 +858,6 @@ impl ForwardPassPerfModel { .static_phase_diagnostics(batch_size, context_length, prefix, prefill) } - /// Latency in milliseconds for an exact homogeneous measured prefill shape. - /// `isl` includes the cached prefix; the engine subtracts it exactly once. - /// This latency-only surface is currently qualified only for the selected - /// SGLang GLM-5.2 NVFP4 VR200 graph profile. - pub fn predict_prefill_latency( - &self, - batch_size: u32, - isl: u32, - prefix: u32, - ) -> Result { - let engine = self - .native_engine() - .filter(|engine| engine.has_prefill_graph_profile()) - .ok_or_else(|| { - crate::perf_database::prefill_graph::error( - "direct prefill latency requires the qualified graph profile", - ) - })?; - engine.predict_prefill_latency(batch_size, isl, prefix) - } - pub(crate) fn has_prefill_graph_profile(&self) -> bool { match &self.mode { ForwardPassPerfMode::Native { engine, .. } => engine.has_prefill_graph_profile(), @@ -851,7 +933,7 @@ fn build_native_candidate( let mut last_error = None; for root in resolve_systems_roots(config)? { match build_engine_via_python(config, &root).and_then(|engine| { - engine.validate_forward_pass_readiness()?; + engine.validate_forward_pass_readiness(config.worker_type)?; Ok(engine) }) { Ok(engine) => return Ok((engine, root)), diff --git a/crates/core/src/perfmodel/fpm/resources.rs b/crates/core/src/perfmodel/fpm/resources.rs new file mode 100644 index 000000000..e8c732041 --- /dev/null +++ b/crates/core/src/perfmodel/fpm/resources.rs @@ -0,0 +1,977 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Rank-local cache descriptors and timing-independent profile memory sizing. + +use serde::{Deserialize, Serialize}; + +use super::{ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassWorkerType}; +use crate::{AicError, BackendKind, MemoryBreakdown}; + +const MAX_EXACT_BYTES: u64 = 1 << 53; + +fn invalid(message: impl Into) -> AicError { + AicError::InvalidEngineConfig(message.into()) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmCacheKind { + Attention, + Convolution, +} + +/// One cache group on one TP rank. `page_size_bytes` includes every layer in +/// this group and any runtime padding; it is not a per-layer or logical-token +/// rate. Window lengths describe retained history, independently of FPM timing +/// query context lengths. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmCacheGroup { + pub name: String, + pub kind: FpmCacheKind, + pub num_layers: u32, + pub block_size_tokens: u32, + pub page_size_bytes: u64, + pub sliding_window: Option, +} + +impl FpmCacheGroup { + pub fn validate(&self) -> Result<(), AicError> { + if self.name.trim().is_empty() + || self.num_layers == 0 + || self.block_size_tokens == 0 + || self.page_size_bytes == 0 + || self.page_size_bytes > MAX_EXACT_BYTES + || self.sliding_window == Some(0) + { + return Err(invalid( + "cache group requires a name, positive layers/block/page/window, and page_size_bytes <= 2**53", + )); + } + if self.kind == FpmCacheKind::Convolution && self.sliding_window.is_none() { + return Err(invalid("convolution cache groups require a sliding_window")); + } + Ok(()) + } + + /// Blocks live during one forward pass. All newly scheduled tokens remain + /// allocated until the next pass, including temporary prefill storage. + pub fn block_range( + &self, + computed_tokens: u64, + target_tokens: u64, + ) -> Result, AicError> { + self.validate()?; + if target_tokens < computed_tokens { + return Err(invalid("cache target_tokens must be >= computed_tokens")); + } + let block = u64::from(self.block_size_tokens); + let first = self.sliding_window.map_or(0, |window| { + computed_tokens.saturating_sub(u64::from(window) - 1) / block + }); + Ok(first..target_tokens.div_ceil(block)) + } + + pub fn bytes_for_forward( + &self, + computed_tokens: u64, + target_tokens: u64, + ) -> Result { + let range = self.block_range(computed_tokens, target_tokens)?; + (range.end - range.start) + .checked_mul(self.page_size_bytes) + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| invalid("cache allocation exceeds 2**53 bytes")) + } + + /// Conservative maximum for a request up to `context_length`, with forward + /// chunks no larger than `max_num_tokens`. The extra block alignment can + /// coexist with the entire new chunk; clipping only to the retention window + /// would underestimate prefill memory. + pub fn peak_request_bytes( + &self, + context_length: u64, + max_num_tokens: u32, + ) -> Result { + self.validate()?; + if context_length == 0 || max_num_tokens == 0 { + return Err(invalid("cache context and max_num_tokens must be positive")); + } + let block = u64::from(self.block_size_tokens); + let full = context_length.div_ceil(block); + let pages = if let Some(window) = self.sliding_window { + let span = u64::from(window) - 1 + u64::from(max_num_tokens) + block - 1; + full.min(span.div_ceil(block)) + } else { + full + }; + pages + .checked_mul(self.page_size_bytes) + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| invalid("cache request allocation exceeds 2**53 bytes")) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmCacheLayout { + Linear, + Grouped, +} + +/// Cache capacity observed for the exact deployment and scheduler settings. +/// This includes the runtime's graph/workspace reservations; it is not a +/// non-KV memory limit or a capacity that can be scaled to another setting. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmRuntimeMemoryConfig { + pub kv_cache_bytes: u64, + pub gpu_memory_utilization: f64, + pub max_model_len: u64, + pub provenance: String, +} + +impl FpmRuntimeMemoryConfig { + pub fn validate(&self) -> Result<(), AicError> { + if self.kv_cache_bytes == 0 + || self.kv_cache_bytes > MAX_EXACT_BYTES + || !self.gpu_memory_utilization.is_finite() + || self.gpu_memory_utilization <= 0.0 + || self.gpu_memory_utilization > 1.0 + || self.max_model_len == 0 + || self.provenance.trim().is_empty() + { + return Err(invalid( + "runtime memory requires positive kv_cache_bytes <= 2**53, finite gpu_memory_utilization in (0, 1], positive max_model_len and provenance", + )); + } + Ok(()) + } +} + +/// The typed resource section of the existing profile wire schema. +/// Missing non-KV components leave memory pending until runtime observation. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmResourceConfig { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub weights_bytes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activations_bytes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_overhead_bytes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comm_overhead_bytes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_memory: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_bytes_per_token: Option, + pub cache_layout: FpmCacheLayout, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub cache_groups: Vec, + pub max_num_tokens: u32, + pub max_batch_size: u32, + pub provenance: String, +} + +impl FpmResourceConfig { + pub fn validate(&self) -> Result<(), AicError> { + if self.max_num_tokens == 0 || self.max_batch_size == 0 || self.provenance.trim().is_empty() + { + return Err(invalid( + "FPM resources require positive scheduler bounds and provenance", + )); + } + let declared = self.declared_bytes(); + declared + .iter() + .flatten() + .try_fold(0u64, |total, value| total.checked_add(*value)) + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| invalid("total non-KV resource bytes must not exceed 2**53"))?; + if let Some(runtime) = &self.runtime_memory { + runtime.validate()?; + if declared.iter().any(Option::is_some) { + return Err(invalid( + "runtime_memory cannot be combined with declared non-KV resource bytes", + )); + } + } + match self.cache_layout { + FpmCacheLayout::Linear + if self.cache_groups.is_empty() + && self + .kv_bytes_per_token + .is_some_and(|bytes| bytes > 0 && bytes <= MAX_EXACT_BYTES) => {} + FpmCacheLayout::Grouped + if self.kv_bytes_per_token.is_none() && !self.cache_groups.is_empty() => + { + let mut names = std::collections::HashSet::new(); + for group in &self.cache_groups { + group.validate()?; + if !names.insert(group.name.as_str()) { + return Err(invalid("cache group names must be unique")); + } + } + } + _ => { + return Err(invalid( + "linear cache requires kv_bytes_per_token only; grouped cache requires cache_groups without a scalar token rate", + )); + } + } + Ok(()) + } + + fn declared_bytes(&self) -> [Option; 4] { + [ + self.weights_bytes, + self.activations_bytes, + self.runtime_overhead_bytes, + self.comm_overhead_bytes, + ] + } + + fn non_kv_bytes(&self) -> Result { + if self.runtime_memory.is_some() { + return Err(invalid( + "runtime memory records cache capacity, not a non-KV byte breakdown", + )); + } + self.require_memory()?; + self.declared_bytes() + .into_iter() + .flatten() + .try_fold(0u64, |total, value| total.checked_add(value)) + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| invalid("total non-KV resource bytes must not exceed 2**53")) + } + + pub fn require_memory(&self) -> Result<(), AicError> { + if self.runtime_memory.is_none() && self.declared_bytes().iter().any(Option::is_none) { + return Err(invalid( + "FPM memory is pending runtime profiling; collect and finalize runtime memory or provide all four declared non-KV resource values before simulation", + )); + } + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmCacheBudgetRequest { + pub total_gpu_capacity_bytes: u64, + pub memory_fraction_kind: String, + pub memory_fraction_value: f64, + pub max_num_tokens: u32, + pub max_batch_size: u32, + #[serde(default)] + pub context_length: Option, + /// Optional logical request length for steady, one-token decode occupancy. + /// Scheduler limits and context_length still control resource admission. + #[serde(default)] + pub request_occupancy_tokens: Option, + #[serde(default)] + pub cuda_graph_reserved_bytes: u64, + #[serde(default)] + pub tolerance_fraction: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmCacheBudgetAdjusted { + pub tolerance_fraction: f64, + pub total_kv_size_bytes: u64, + pub total_kv_size_tokens: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmCacheBudget { + pub total_gpu_capacity_bytes: u64, + pub total_kv_size_bytes: u64, + pub kv_size_per_token_bytes: Option, + pub total_kv_size_tokens: Option, + pub source: String, + pub memory_breakdown: Option, + pub tolerance_adjusted: Option, + pub resource_provenance: String, + pub cache_layout: FpmCacheLayout, + pub cache_groups: Vec, + /// Per-request upper bound, including window block alignment and the + /// configured prefill chunk. This is not an aggregate token capacity. + pub request_peak_cache_bytes: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub request_occupancy_cache_bytes: Option, +} + +impl ForwardPassPerfModelConfig { + /// Select exact profile resources without constructing an analytical model + /// or reading timings. Identity/schema normalization is the same as for + /// `best_available`. + pub fn estimate_cache_budget( + &self, + request: &FpmCacheBudgetRequest, + ) -> Result { + let resolved = self.clone().resolve()?; + let resources = resolved + .fpm_resources()? + .ok_or_else(|| invalid("cache budget requires fpm_profile resources"))?; + let context = resolved + .fpm_profile + .as_ref() + .and_then(|p| p.get("context_length")) + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| invalid("fpm_profile requires positive context_length"))?; + let requested_context = request.context_length.unwrap_or(context); + if requested_context == 0 || requested_context > context { + return Err(invalid( + "context_length must be positive and within the FPM profile context bound", + )); + } + resources.estimate_budget(request, requested_context) + } + + pub fn fpm_resources(&self) -> Result, AicError> { + let Some(profile) = &self.fpm_profile else { + return Ok(None); + }; + if profile.get("model").and_then(serde_json::Value::as_str) != Some(self.model.as_str()) { + return Err(invalid("FPM profile model identity mismatch")); + } + if self.backend != BackendKind::Vllm + || self.pp != 1 + || self.nextn != 0 + || self.speculation.is_some() + { + return Err(invalid( + "FPM resources require vLLM PP1 plain autoregressive execution", + )); + } + let deployments = profile + .get("deployments") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| invalid("fpm_profile requires deployments"))?; + let matching_identity: Vec<_> = deployments + .iter() + .filter(|d| { + d.get("system").and_then(serde_json::Value::as_str) == Some(self.system.as_str()) + && d.get("backend").and_then(serde_json::Value::as_str) + == Some(self.backend.as_str()) + && d.get("backend_version").and_then(serde_json::Value::as_str) + == self.backend_version.as_deref() + && [ + ("tp", self.tp), + ("pp", self.pp), + ("dp", self.attention_dp), + ("moe_tp", self.moe_tp_size.unwrap_or(1)), + ("moe_ep", self.moe_ep_size.unwrap_or(1)), + ("cp", 1), + ] + .into_iter() + .all(|(key, expected)| { + d.get(key) + .and_then(serde_json::Value::as_u64) + .unwrap_or(if key == "pp" || key == "cp" { 1 } else { 0 }) + == u64::from(expected) + }) + }) + .collect(); + let roles = matching_identity + .iter() + .map(|d| { + d.get("worker_type") + .filter(|role| !role.is_null()) + .cloned() + .map(serde_json::from_value::) + .transpose() + .map_err(|error| invalid(format!("FPM deployment worker_type: {error}"))) + }) + .collect::, _>>()?; + if roles.iter().any(Option::is_none) && roles.iter().any(Option::is_some) { + return Err(invalid( + "ambiguous FPM deployment identity: cannot mix legacy shared and role-specific profiles", + )); + } + let mut matching = + matching_identity + .into_iter() + .zip(roles) + .filter_map(|(deployment, role)| { + role.is_none_or(|role| role == self.worker_type) + .then_some(deployment) + }); + let deployment = matching + .next() + .ok_or_else(|| invalid("no matching FPM deployment resources"))?; + if matching.next().is_some() { + return Err(invalid("duplicate FPM deployment resources")); + } + let resources: FpmResourceConfig = serde_json::from_value( + deployment + .get("resources") + .cloned() + .ok_or_else(|| invalid("FPM deployment requires resources"))?, + ) + .map_err(|error| invalid(format!("FPM resource profile: {error}")))?; + resources.validate()?; + Ok(Some(resources)) + } +} + +impl FpmResourceConfig { + fn estimate_budget( + &self, + request: &FpmCacheBudgetRequest, + context: u64, + ) -> Result { + self.validate()?; + if request.memory_fraction_kind != "of_total" { + return Err(invalid( + "vLLM FPM resources require memory_fraction_kind='of_total'", + )); + } + let fraction = request.memory_fraction_value; + if !fraction.is_finite() || !(0.0..=1.0).contains(&fraction) { + return Err(invalid( + "memory_fraction_value must be finite and in [0, 1]", + )); + } + if request + .tolerance_fraction + .is_some_and(|t| !t.is_finite() || !(0.0..1.0).contains(&t)) + { + return Err(invalid("tolerance_fraction must be finite and in [0, 1)")); + } + if request.total_gpu_capacity_bytes == 0 + || request.total_gpu_capacity_bytes > MAX_EXACT_BYTES + { + return Err(invalid( + "GPU capacity must be positive and no greater than 2**53", + )); + } + if request.max_num_tokens == 0 + || request.max_num_tokens > self.max_num_tokens + || request.max_batch_size == 0 + || request.max_batch_size > self.max_batch_size + { + return Err(invalid( + "FPM resource envelope exceeded: provide bounds for the requested positive rank-local scheduler envelope", + )); + } + let (total, memory_breakdown) = if let Some(runtime) = &self.runtime_memory { + if request.max_num_tokens != self.max_num_tokens + || request.max_batch_size != self.max_batch_size + { + return Err(invalid( + "runtime memory requires the exact recorded rank-local scheduler settings; collect new runtime memory for changed settings", + )); + } + if context > runtime.max_model_len { + return Err(invalid( + "context_length exceeds runtime memory max_model_len", + )); + } + if fraction != runtime.gpu_memory_utilization { + return Err(invalid( + "memory_fraction_value must match runtime memory gpu_memory_utilization", + )); + } + if request.cuda_graph_reserved_bytes != 0 { + return Err(invalid( + "runtime memory already includes graph reservations; cuda_graph_reserved_bytes must be 0", + )); + } + if (request.total_gpu_capacity_bytes as f64 * fraction).floor() + < runtime.kv_cache_bytes as f64 + { + return Err(invalid( + "GPU memory budget is smaller than the recorded runtime KV cache allocation", + )); + } + (runtime.kv_cache_bytes, None) + } else { + let non_kv = self + .non_kv_bytes()? + .checked_add(request.cuda_graph_reserved_bytes) + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| { + invalid("total FPM non-KV bytes including CUDA graphs must not exceed 2**53") + })?; + let available = + (request.total_gpu_capacity_bytes as f64 * fraction - non_kv as f64).floor(); + if available < 1.0 { + return Err(invalid( + "no KV budget after non-KV resources and CUDA graphs", + )); + } + ( + available as u64, + Some(MemoryBreakdown { + weights_bytes: self.weights_bytes.unwrap(), + activations_bytes: self.activations_bytes.unwrap(), + runtime_overhead_bytes: self.runtime_overhead_bytes.unwrap(), + comm_overhead_bytes: self.comm_overhead_bytes.unwrap(), + cuda_graph_reserved_bytes: request.cuda_graph_reserved_bytes, + }), + ) + }; + let token_count = |bytes| self.kv_bytes_per_token.map(|rate| bytes / rate); + let footprint = |tokens: u64, chunk: u32| { + let bytes = if self.cache_layout == FpmCacheLayout::Linear { + tokens.checked_mul(self.kv_bytes_per_token.unwrap()) + } else { + self.cache_groups.iter().try_fold(0u64, |sum, group| { + sum.checked_add(group.peak_request_bytes(tokens, chunk).ok()?) + }) + }; + bytes + .filter(|bytes| *bytes <= MAX_EXACT_BYTES) + .ok_or_else(|| invalid("cache request allocation exceeds 2**53 bytes")) + }; + let request_peak_cache_bytes = footprint(context, request.max_num_tokens)?; + let request_occupancy_cache_bytes = request + .request_occupancy_tokens + .map(|tokens| { + if tokens == 0 || tokens > context { + return Err(invalid( + "request_occupancy_tokens must be positive and within context_length", + )); + } + footprint(tokens, 1) + }) + .transpose()?; + Ok(FpmCacheBudget { + total_gpu_capacity_bytes: request.total_gpu_capacity_bytes, + total_kv_size_bytes: total, + kv_size_per_token_bytes: self.kv_bytes_per_token, + total_kv_size_tokens: token_count(total), + source: "profile".into(), + memory_breakdown, + tolerance_adjusted: request.tolerance_fraction.map(|tolerance| { + let bytes = (total as f64 * (1.0 - tolerance)).floor() as u64; + FpmCacheBudgetAdjusted { + tolerance_fraction: tolerance, + total_kv_size_bytes: bytes, + total_kv_size_tokens: token_count(bytes), + } + }), + resource_provenance: self.provenance.clone(), + cache_layout: self.cache_layout, + cache_groups: self.cache_groups.clone(), + request_peak_cache_bytes, + request_occupancy_cache_bytes, + }) + } +} + +impl ForwardPassPerfModel { + pub fn estimate_cache_budget( + &self, + request: &FpmCacheBudgetRequest, + ) -> Result { + self.provenance() + .ok_or_else(|| invalid("cache resources require canonical model provenance"))? + .config + .estimate_cache_budget(request) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn group(window: Option) -> FpmCacheGroup { + FpmCacheGroup { + name: "attention".into(), + kind: FpmCacheKind::Attention, + num_layers: 2, + block_size_tokens: 16, + page_size_bytes: 160, + sliding_window: window, + } + } + + fn resources() -> FpmResourceConfig { + FpmResourceConfig { + weights_bytes: Some(100), + activations_bytes: Some(20), + runtime_overhead_bytes: Some(30), + comm_overhead_bytes: Some(50), + runtime_memory: None, + kv_bytes_per_token: None, + cache_layout: FpmCacheLayout::Grouped, + cache_groups: vec![group(Some(32))], + max_num_tokens: 128, + max_batch_size: 8, + provenance: "Hand-derived rank-local byte fixture".into(), + } + } + + fn request() -> FpmCacheBudgetRequest { + FpmCacheBudgetRequest { + total_gpu_capacity_bytes: 1000, + memory_fraction_kind: "of_total".into(), + memory_fraction_value: 0.8, + max_num_tokens: 128, + max_batch_size: 8, + context_length: None, + request_occupancy_tokens: None, + cuda_graph_reserved_bytes: 20, + tolerance_fraction: Some(0.1), + } + } + + fn runtime_resources() -> FpmResourceConfig { + FpmResourceConfig { + weights_bytes: None, + activations_bytes: None, + runtime_overhead_bytes: None, + comm_overhead_bytes: None, + runtime_memory: Some(FpmRuntimeMemoryConfig { + kv_cache_bytes: 600, + gpu_memory_utilization: 0.8, + max_model_len: 4096, + provenance: "Synthetic initialized-worker cache allocation".into(), + }), + ..resources() + } + } + + fn role_config() -> ForwardPassPerfModelConfig { + let mut config = ForwardPassPerfModelConfig::new( + "test/model", + "test_gpu", + BackendKind::Vllm, + ForwardPassWorkerType::Prefill, + ); + config.backend_version = Some("0.27.0".into()); + let deployment = serde_json::json!({ + "system": "test_gpu", "backend": "vllm", "backend_version": "0.27.0", + "tp": 1, "dp": 1, "moe_tp": 1, "moe_ep": 1, + "worker_type": "prefill", "resources": runtime_resources(), + }); + let mut decode = deployment.clone(); + decode["worker_type"] = "decode".into(); + decode["resources"]["runtime_memory"]["kv_cache_bytes"] = 700.into(); + decode["resources"]["max_batch_size"] = 4.into(); + config.fpm_profile = Some( + serde_json::json!({ + "model": "test/model", "deployments": [deployment, decode], + }) + .as_object() + .unwrap() + .clone(), + ); + config + } + + #[test] + fn profile_resources_select_independent_roles_after_round_trip() { + let config = role_config(); + let mut reloaded: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&config).unwrap()).unwrap(); + let prefill = reloaded.fpm_resources().unwrap().unwrap(); + assert_eq!(prefill.runtime_memory.unwrap().kv_cache_bytes, 600); + assert_eq!(prefill.max_batch_size, 8); + reloaded.worker_type = ForwardPassWorkerType::Decode; + let decode = reloaded.fpm_resources().unwrap().unwrap(); + assert_eq!(decode.runtime_memory.unwrap().kv_cache_bytes, 700); + assert_eq!(decode.max_batch_size, 4); + reloaded.worker_type = ForwardPassWorkerType::Aggregated; + assert!( + reloaded + .fpm_resources() + .unwrap_err() + .to_string() + .contains("no matching") + ); + } + + #[test] + fn legacy_shared_resources_are_preserved_without_role_inference() { + let mut config = role_config(); + let deployments = config.fpm_profile.as_mut().unwrap()["deployments"] + .as_array_mut() + .unwrap(); + deployments.truncate(1); + deployments[0] + .as_object_mut() + .unwrap() + .remove("worker_type"); + for role in [ + ForwardPassWorkerType::Prefill, + ForwardPassWorkerType::Decode, + ForwardPassWorkerType::Aggregated, + ] { + config.worker_type = role; + assert_eq!( + config + .fpm_resources() + .unwrap() + .unwrap() + .runtime_memory + .unwrap() + .kv_cache_bytes, + 600 + ); + } + assert!( + config.fpm_profile.as_ref().unwrap()["deployments"][0] + .get("worker_type") + .is_none() + ); + } + + #[test] + fn resource_selection_rejects_ambiguous_shared_and_role_specific_profiles() { + let mut config = role_config(); + config.fpm_profile.as_mut().unwrap()["deployments"][0] + .as_object_mut() + .unwrap() + .remove("worker_type"); + assert!( + config + .fpm_resources() + .unwrap_err() + .to_string() + .contains("cannot mix legacy") + ); + } + + #[test] + fn pending_memory_validates_but_cannot_produce_a_cache_budget() { + let mut pending = resources(); + pending.activations_bytes = None; + pending.validate().unwrap(); + let saved = serde_json::to_value(&pending).unwrap(); + assert!(saved.get("activations_bytes").is_none()); + assert!(saved.get("runtime_memory").is_none()); + assert_eq!(saved["weights_bytes"], 100); + assert_eq!( + serde_json::from_value::(saved).unwrap(), + pending + ); + assert!( + pending + .estimate_budget(&request(), 4096) + .unwrap_err() + .to_string() + .contains("pending runtime profiling") + ); + } + + #[test] + fn runtime_memory_preserves_observed_capacity_without_a_breakdown() { + let mut runtime = runtime_resources(); + runtime.validate().unwrap(); + let mut q = request(); + q.cuda_graph_reserved_bytes = 0; + let grouped = runtime.estimate_budget(&q, 4096).unwrap(); + assert_eq!(grouped.total_kv_size_bytes, 600); + assert_eq!(grouped.total_kv_size_tokens, None); + assert_eq!(grouped.memory_breakdown, None); + assert_eq!(grouped.request_peak_cache_bytes, 1760); + + runtime.cache_layout = FpmCacheLayout::Linear; + runtime.cache_groups.clear(); + runtime.kv_bytes_per_token = Some(10); + q.total_gpu_capacity_bytes = 2000; + let linear = runtime.estimate_budget(&q, 4096).unwrap(); + // Runtime recorded 600 bytes: a larger device does not double capacity. + // Ten bytes/token gives 60 raw tokens and 54 with a 10% margin. + assert_eq!(linear.total_kv_size_bytes, 600); + assert_eq!(linear.total_kv_size_tokens, Some(60)); + let adjusted = linear.tolerance_adjusted.unwrap(); + assert_eq!(adjusted.total_kv_size_bytes, 540); + assert_eq!(adjusted.total_kv_size_tokens, Some(54)); + assert_eq!(linear.memory_breakdown, None); + let saved = serde_json::to_value(&runtime).unwrap(); + assert!(saved.get("weights_bytes").is_none()); + assert_eq!( + serde_json::from_value::(saved).unwrap(), + runtime + ); + } + + #[test] + fn steady_occupancy_keeps_runtime_scheduler_and_context_admission() { + let runtime = runtime_resources(); + let mut q = request(); + q.cuda_graph_reserved_bytes = 0; + q.request_occupancy_tokens = Some(1152); + let budget = runtime.estimate_budget(&q, 4096).unwrap(); + // A 32-token retained window can straddle three 16-token pages. + // A 128-token prefill chunk needs eleven pages. Each page is 160 bytes. + assert_eq!(budget.request_occupancy_cache_bytes, Some(480)); + assert_eq!(budget.request_peak_cache_bytes, 1760); + q.max_num_tokens = 1; + assert!(runtime.estimate_budget(&q, 4096).is_err()); + q.max_num_tokens = 128; + assert!(runtime.estimate_budget(&q, 8192).is_err()); + q.request_occupancy_tokens = Some(4097); + assert!(runtime.estimate_budget(&q, 4096).is_err()); + } + + #[test] + fn runtime_memory_rejects_changed_settings_and_invalid_evidence() { + let runtime = runtime_resources(); + let mut q = request(); + q.cuda_graph_reserved_bytes = 0; + for changed in [ + FpmCacheBudgetRequest { + max_num_tokens: 64, + ..q.clone() + }, + FpmCacheBudgetRequest { + max_batch_size: 4, + ..q.clone() + }, + FpmCacheBudgetRequest { + memory_fraction_value: 0.9, + ..q.clone() + }, + FpmCacheBudgetRequest { + cuda_graph_reserved_bytes: 1, + ..q.clone() + }, + FpmCacheBudgetRequest { + total_gpu_capacity_bytes: 599, + ..q.clone() + }, + FpmCacheBudgetRequest { + total_gpu_capacity_bytes: 749, + ..q.clone() + }, + ] { + assert!(runtime.estimate_budget(&changed, 4096).is_err()); + } + assert!(runtime.estimate_budget(&q, 4097).is_err()); + let mut mixed = runtime.clone(); + mixed.weights_bytes = Some(0); + assert!(mixed.validate().is_err()); + let mut invalid = runtime.runtime_memory.clone().unwrap(); + for bytes in [0, MAX_EXACT_BYTES + 1] { + invalid.kv_cache_bytes = bytes; + assert!(invalid.validate().is_err()); + } + invalid.kv_cache_bytes = MAX_EXACT_BYTES; + invalid.validate().unwrap(); + for fraction in [0.0, -0.1, 1.1, f64::NAN, f64::INFINITY] { + invalid.gpu_memory_utilization = fraction; + assert!(invalid.validate().is_err()); + } + invalid.gpu_memory_utilization = 1.0; + invalid.max_model_len = 0; + assert!(invalid.validate().is_err()); + invalid.max_model_len = 1; + invalid.provenance = " ".into(); + assert!(invalid.validate().is_err()); + } + + #[test] + fn grouped_page_boundaries_keep_full_history_and_evict_windows() { + let local = group(Some(32)); + // First 32 tokens occupy two physical pages. Token 33 retains parts of + // pages 0, 1, and 2; page 0 disappears only when its history is obsolete. + assert_eq!(local.bytes_for_forward(0, 32).unwrap(), 320); + assert_eq!(local.block_range(32, 33).unwrap(), 0..3); + assert_eq!(local.block_range(48, 49).unwrap(), 1..4); + assert_eq!(local.bytes_for_forward(64, 65).unwrap(), 480); + assert_eq!(group(None).bytes_for_forward(64, 65).unwrap(), 800); + // A 128-token prefill chunk keeps ten pages here, exceeding the window. + assert_eq!(local.block_range(512, 640).unwrap(), 30..40); + assert_eq!(local.bytes_for_forward(512, 640).unwrap(), 1600); + assert!(local.bytes_for_forward(33, 32).is_err()); + } + + #[test] + fn cache_budget_keeps_byte_tolerance_without_inventing_token_capacity() { + let result = resources().estimate_budget(&request(), 4096).unwrap(); + // 80% of 1000, minus 100+20+30+50 non-KV and 20 graph bytes. + assert_eq!(result.total_kv_size_bytes, 580); + assert_eq!( + result + .tolerance_adjusted + .as_ref() + .unwrap() + .total_kv_size_bytes, + 522 + ); + assert_eq!(result.total_kv_size_tokens, None); + assert_eq!(result.kv_size_per_token_bytes, None); + assert_eq!( + result.tolerance_adjusted.unwrap().total_kv_size_tokens, + None + ); + // At most 11 aligned pages cover 31 history + 128 scheduled tokens. + assert_eq!(result.request_peak_cache_bytes, 1760); + assert_eq!( + result.memory_breakdown.unwrap().cuda_graph_reserved_bytes, + 20 + ); + } + + #[test] + fn grouped_resources_validate_layout_and_exact_bytes() { + let original = resources(); + assert_eq!( + serde_json::from_str::(&serde_json::to_string(&original).unwrap()) + .unwrap(), + original + ); + let mut bad = original.clone(); + bad.kv_bytes_per_token = Some(1); + assert!(bad.validate().is_err()); + bad = original.clone(); + bad.cache_groups.push(bad.cache_groups[0].clone()); + assert!(bad.validate().is_err()); + bad = original.clone(); + bad.cache_groups[0].page_size_bytes = MAX_EXACT_BYTES; + assert!(bad.cache_groups[0].bytes_for_forward(0, 17).is_err()); + bad.weights_bytes = Some(u64::MAX); + assert!(bad.validate().is_err()); + let mut unknown = serde_json::to_value(original).unwrap(); + unknown["cache_groups"][0]["unknown_page_semantics"] = true.into(); + assert!(serde_json::from_value::(unknown).is_err()); + } + + #[test] + fn grouped_budget_rejects_invalid_fraction_envelope_and_capacity() { + let r = resources(); + for fraction in [f64::NAN, f64::INFINITY, -0.1, 1.1] { + let mut q = request(); + q.memory_fraction_value = fraction; + assert!(r.estimate_budget(&q, 4096).is_err()); + } + for tolerance in [f64::NAN, -0.1, 1.0] { + let mut q = request(); + q.tolerance_fraction = Some(tolerance); + assert!(r.estimate_budget(&q, 4096).is_err()); + } + let mut q = request(); + q.max_batch_size = 9; + assert!(r.estimate_budget(&q, 4096).is_err()); + q = request(); + q.total_gpu_capacity_bytes = 0; + assert!(r.estimate_budget(&q, 4096).is_err()); + q = request(); + q.memory_fraction_value = 0.1; + assert!(r.estimate_budget(&q, 4096).is_err()); + q = request(); + q.cuda_graph_reserved_bytes = MAX_EXACT_BYTES; + assert!(r.estimate_budget(&q, 4096).is_err()); + } + + #[test] + fn convolution_requires_an_explicit_retention_window() { + let mut conv = group(None); + conv.kind = FpmCacheKind::Convolution; + assert!(conv.validate().is_err()); + conv.sliding_window = Some(4); + conv.block_size_tokens = 4; + conv.page_size_bytes = 8; + // Token 9 uses pages containing 5..8 and 9..12, even for a four-token window. + assert_eq!(conv.block_range(8, 9).unwrap(), 1..3); + assert_eq!(conv.bytes_for_forward(8, 9).unwrap(), 16); + } +} diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index ea807e34b..c86b17e58 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -166,6 +166,239 @@ fn native_model(options: ForwardPassPerfOptions) -> ForwardPassPerfModel { ForwardPassPerfModel::from_engine(fixture_engine(), options) } +/// Synthetic measured timings with no SOL operations. The fixture's expected +/// answers are its literal measurements and raw linear midpoints. +fn direct_coverage_model(dir: &std::path::Path) -> ForwardPassPerfModel { + use crate::operators::fpm_forward::{FpmForwardOp, FpmInterpolation, FpmPhase}; + use crate::perf_database::fpm_forward::tests::{RowSpec, default_identity, write_pair}; + let mut rows = Vec::new(); + for (tokens, latency) in [(100, 10.0), (200, 20.0)] { + for (kv, bump) in [(0, 0.0), (200, 4.0)] { + rows.push(RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms: latency + bump, + ..Default::default() + }); + } + } + for (batch, kv, latency) in [ + (8, 64, 8.0), + (8, 128, 16.0), + (9, 64, 9.0), + (9, 128, 17.0), + (16, 64, 16.0), + (16, 128, 24.0), + (17, 64, 170.0), + (17, 128, 178.0), + ] { + rows.push(RowSpec { + batch_size: batch, + total_kv_read_tokens: kv, + latency_ms: latency, + kv_seed_regime: Some("real_kv"), + ..Default::default() + }); + } + rows.push(RowSpec { + batch_size: 8, + total_kv_read_tokens: 256, + latency_ms: 99.0, + kv_seed_regime: Some("fake_fallback"), + ..Default::default() + }); + write_pair(dir, &rows); + let mut db = PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.24.0").unwrap(); + db.set_fpm_forward_for_test(crate::perf_database::FpmForwardTable::new( + dir.to_path_buf(), + "b200_sxm", + "vllm", + "0.25.1", + )); + let op = |phase: FpmPhase| { + Op::FpmForward(FpmForwardOp { + name: format!("fpm_forward_{}", phase.as_str()), + phase, + model_path: "org/model-a".into(), + dcp_size: None, + match_identity: default_identity(4), + original_fmha_quant_mode: None, + weight_bytes: 0.0, + verify_width: 1, + interpolation: FpmInterpolation::Direct, + sol_ops: vec![], + }) + }; + let spec = EngineSpec::new( + fixture_engine_config(), + vec![op(FpmPhase::Prefill)], + vec![op(FpmPhase::Decode)], + ); + let engine = Engine::build(spec, Arc::new(db)).unwrap(); + let mut model = ForwardPassPerfModel::from_engine(Arc::new(engine), Default::default()); + model.query_coverage = Some(Default::default()); + model +} + +#[test] +fn direct_coverage_records_real_resolutions_and_mixed_baselines() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + assert_eq!(model.predict_prefill_latency(1, 100, 0).unwrap(), 10.0); + // At 150 new / 100 cached tokens, each axis is halfway between rows. + assert_eq!(model.predict_prefill_latency(1, 250, 100).unwrap(), 17.0); + assert_eq!(model.predict_decode_latency_total(8, 64).unwrap(), 8.0); + assert_eq!(model.predict_decode_latency_total(8, 96).unwrap(), 12.0); + let mixed = ForwardPassMetrics { + scheduled_requests: ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 100, + num_decode_requests: 12, + sum_decode_kv_tokens: 96, + ..Default::default() + }, + ..Default::default() + }; + // Same graph bracket rows 9/16 give decode=16, floor=12. Prefill=10. + assert_eq!( + model.estimate_forward_pass_time_ms(&[mixed]).unwrap(), + Some(14.0) + ); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.counting_unit, "native_lookup_resolutions"); + assert_eq!(report.queries.measured, 3); + assert_eq!(report.queries.interpolated, 4); + assert_eq!(report.prefill.measured, 2); + assert_eq!(report.prefill.interpolated, 1); + assert_eq!(report.decode.measured, 1); + assert_eq!(report.decode.interpolated, 2); + assert_eq!(report.mixed_decode_baseline.interpolated, 1); + assert_eq!(report.queries.unsupported, 0); + assert!(report.gaps.is_empty()); +} + +#[test] +fn direct_coverage_keeps_partial_failures_and_excludes_fake_rows() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + let mixed = ForwardPassMetrics { + scheduled_requests: ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 100, + num_decode_requests: 8, + sum_decode_kv_tokens: 256, + ..Default::default() + }, + ..Default::default() + }; + let error = model.estimate_forward_pass_time_ms(&[mixed]).unwrap_err(); + assert!(error.to_string().contains("genuine FPM points")); + assert!(model.predict_decode_latency_total(8, 256).is_err()); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.prefill.measured, 1); + assert_eq!(report.decode.unsupported, 2); + assert_eq!(report.mixed_decode_baseline, Default::default()); + assert_eq!(report.gaps.len(), 1); + let gap = &report.gaps[0]; + assert_eq!(gap.occurrences, 2); + assert_eq!(gap.phase, "decode"); + assert_eq!(gap.model_path, "org/model-a"); + assert_eq!(gap.cell_identity["tp"], "4"); + assert!(!gap.cell_ids.is_empty()); + assert_eq!(gap.coordinates["batch_size"], 8.0); + assert_eq!(gap.coordinates["total_kv_read_tokens"], 256.0); + assert_eq!(gap.reason, error.to_string()); + assert_eq!(model.fpm_decode_kv_ceiling().unwrap(), Some(128)); + assert_eq!(model.fpm_query_coverage().unwrap().unwrap(), report); +} + +#[test] +fn direct_coverage_is_bounded_and_owned_by_each_model() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + model.predict_prefill_latency(1, 100, 0).unwrap(); + let clone = model.clone(); + assert_eq!( + clone.fpm_query_coverage().unwrap().unwrap().queries, + Default::default() + ); + assert_eq!( + model + .fpm_query_coverage() + .unwrap() + .unwrap() + .queries + .measured, + 1 + ); + assert!(clone.predict_prefill_latency(3, 100, 0).is_err()); + assert_eq!( + model + .fpm_query_coverage() + .unwrap() + .unwrap() + .queries + .unsupported, + 0 + ); + for kv in 1000..1140 { + assert!(model.predict_decode_latency_total(8, kv).is_err()); + } + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.queries.unsupported, 140); + assert_eq!(report.gaps.len(), report.gap_limit); + assert_eq!(report.omitted_gap_queries, 12); + // A retained gap continues counting duplicates even after the cap. + assert!(model.predict_decode_latency_total(8, 1000).is_err()); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.gaps[0].occurrences, 2); + assert_eq!(report.omitted_gap_queries, 12); +} + +#[test] +fn direct_coverage_static_wrappers_preserve_engine_answers_and_errors() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + let engine = model.native_engine().unwrap(); + for (batch, isl, prefix) in [ + (1, 100, 0), + (1, 250, 100), + (1, 400, 200), + (0, 100, 0), + (1, 100, 100), + (3, 100, 0), + ] { + assert_eq!( + model + .predict_prefill_latency(batch, isl, prefix) + .map_err(|err| err.to_string()), + engine + .predict_prefill_latency(batch, isl, prefix) + .map_err(|err| err.to_string()), + ); + } + for (batch, kv) in [(0, 0), (8, 64), (8, 96), (12, 96), (8, 256)] { + assert_eq!( + model + .predict_decode_latency_total(batch, kv) + .map_err(|err| err.to_string()), + engine + .predict_decode_latency_total(batch, kv) + .map_err(|err| err.to_string()), + ); + } + assert_eq!( + model.fpm_decode_kv_ceiling().unwrap(), + engine.fpm_decode_kv_ceiling().unwrap() + ); + let mut disabled = model.clone(); + disabled.query_coverage = None; + assert_eq!(disabled.predict_prefill_latency(1, 250, 100).unwrap(), 17.0); + assert_eq!(disabled.fpm_query_coverage().unwrap(), None); +} + fn regression_model( worker_type: ForwardPassWorkerType, options: ForwardPassPerfOptions, diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index 3830f3742..20938c69e 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -310,6 +310,14 @@ fn estimate_kwargs<'py>( let (fraction_kind, fraction_value) = req.kv_cache_memory_fraction.to_wire(); let kwargs = PyDict::new(py); kwargs.set_item("backend_version", engine.backend_version.as_deref())?; + kwargs.set_item( + "worker_type", + engine + .extra + .get("worker_type") + .map(String::as_str) + .unwrap_or("aggregated"), + )?; kwargs.set_item("max_num_tokens", req.max_num_tokens)?; kwargs.set_item("max_batch_size", req.max_batch_size)?; kwargs.set_item("memory_fraction_kind", fraction_kind)?; @@ -372,6 +380,12 @@ fn estimate_from_dict( use pyo3::exceptions::PyValueError; use pyo3::types::PyAnyMethods; + if out.get_item("total_kv_size_tokens")?.is_none() { + return Err(PyValueError::new_err( + "grouped FPM caches have no scalar token capacity; use ForwardPassPerfModelConfig::estimate_cache_budget and a cache-group-aware allocator", + )); + } + let u64_at = |k: &str| -> pyo3::PyResult { out.get_item(k)?.extract::() }; let source = match out.get_item("source")?.extract::()?.as_str() { @@ -623,6 +637,21 @@ mod tests { } } + #[cfg(feature = "python")] + #[test] + fn scalar_memory_transport_rejects_grouped_capacity() { + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyDictMethods}; + + pyo3::prepare_freethreaded_python(); + Python::with_gil(|py| { + let out = PyDict::new(py); + out.set_item("total_kv_size_tokens", py.None()).unwrap(); + let error = estimate_from_dict(out.as_any()).unwrap_err(); + assert!(error.to_string().contains("no scalar token capacity")); + }); + } + #[cfg(feature = "python")] #[test] fn py_dict_forwarding_and_parsing_are_deterministic() { diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index f89fbb72f..810ded819 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -55,7 +55,12 @@ pub use fpm::{ }; pub use fpm::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, - OpLevelConfig, RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, + FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, +}; +pub use fpm::{ + FpmCacheBudget, FpmCacheBudgetAdjusted, FpmCacheBudgetRequest, FpmCacheGroup, FpmCacheKind, + FpmCacheLayout, FpmResourceConfig, FpmRuntimeMemoryConfig, }; // Forward-pass metrics telemetry types and schema version, plus the // crate-internal validation helper. Re-exported at the crate root so existing diff --git a/crates/core/src/perfmodel/operators/fpm_forward.rs b/crates/core/src/perfmodel/operators/fpm_forward.rs index 8832efdcd..dfc264480 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -41,6 +41,7 @@ use serde::{Deserialize, Serialize}; use crate::{DirectFpmQueryEvidence, FpmCoordinates, FpmMeasurementSupport, FpmQueryResolution}; use crate::common::error::AicError; +use crate::fpm::coverage::{DirectFpmResolution, FpmQueryCoverage, FpmQueryPurpose}; use crate::operators::op::{Op, RuntimeContext}; use crate::operators::{PerformanceResult, Source}; use crate::perf_database::PerfDatabase; @@ -161,6 +162,15 @@ impl FpmForwardOp { &self, db: &PerfDatabase, ctx: &RuntimeContext, + ) -> Result { + self.query_with_coverage(db, ctx, None) + } + + pub(crate) fn query_with_coverage( + &self, + db: &PerfDatabase, + ctx: &RuntimeContext, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { if self.verify_width == 0 || (self.phase == FpmPhase::Prefill && self.verify_width != 1) { return Err(data_err(format!( @@ -189,7 +199,6 @@ impl FpmForwardOp { ctx.beam_width ))); } - let cell = self.select_cell(db)?; let b = batch_size as f64; let coords: Vec = match self.phase { FpmPhase::Prefill => { @@ -208,7 +217,7 @@ impl FpmForwardOp { vec![b, b / w * s as f64] } }; - self.resolve(db, cell, &coords) + self.query_totals_with_coverage(db, &coords, coverage) } /// Resolve at RAW per-rank iteration totals — the table's native @@ -221,6 +230,15 @@ impl FpmForwardOp { &self, db: &PerfDatabase, coords: &[f64], + ) -> Result { + self.query_totals_with_coverage(db, coords, None) + } + + pub(crate) fn query_totals_with_coverage( + &self, + db: &PerfDatabase, + coords: &[f64], + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { let expected = match self.phase { FpmPhase::Prefill => 3, @@ -233,6 +251,14 @@ impl FpmForwardOp { coords ))); } + if let Some(coverage) = coverage { + return self.query_direct_with_coverage( + db, + coords, + FpmQueryPurpose::ForwardPass, + coverage, + ); + } let cell = self.select_cell(db)?; self.resolve(db, cell, coords) } @@ -272,6 +298,16 @@ impl FpmForwardOp { db: &PerfDatabase, batch_size: u32, total_kv: f64, + ) -> Result { + self.query_pass_baseline_with_coverage(db, batch_size, total_kv, None) + } + + pub(crate) fn query_pass_baseline_with_coverage( + &self, + db: &PerfDatabase, + batch_size: u32, + total_kv: f64, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { if self.phase != FpmPhase::Decode { return Err(data_err(format!( @@ -284,9 +320,19 @@ impl FpmForwardOp { "invalid FPM baseline query: batch_size={batch_size}" ))); } + if let Some(coverage) = coverage { + return self.query_direct_with_coverage( + db, + &[batch_size as f64, total_kv], + FpmQueryPurpose::MixedDecodeBaseline, + coverage, + ); + } let cell = self.select_cell(db)?; if self.interpolation == FpmInterpolation::Direct { - return self.resolve_direct_decode(cell, &[batch_size as f64, total_kv], true); + return self + .resolve_direct_decode(cell, &[batch_size as f64, total_kv], true) + .map(|(result, _)| result); } let Some(domain) = cell.decode_domain else { return Err(data_err(format!( @@ -371,7 +417,8 @@ impl FpmForwardOp { return match self.phase { FpmPhase::Prefill => self.resolve_direct_prefill(cell, coords), FpmPhase::Decode => self.resolve_direct_decode(cell, coords, false), - }; + } + .map(|(result, _)| result); } // Data-certified prefill batch clamp (mirrors Python _resolve): the // regime coordinate is the token TOTAL, which stays untouched — the @@ -504,11 +551,46 @@ impl FpmForwardOp { Ok(PerformanceResult::new(latency, Source::Silicon)) } + fn query_direct_with_coverage( + &self, + db: &PerfDatabase, + coords: &[f64], + purpose: FpmQueryPurpose, + coverage: &mut FpmQueryCoverage, + ) -> Result { + let mut cell_ids = &[][..]; + let result = (|| { + if self.interpolation != FpmInterpolation::Direct { + return Err(AicError::InvalidEngineConfig( + "query coverage requires direct FPM interpolation".into(), + )); + } + let cell = self.select_cell(db)?; + cell_ids = cell.cell_ids.as_slice(); + match self.phase { + FpmPhase::Prefill => self.resolve_direct_prefill(cell, coords), + FpmPhase::Decode => self.resolve_direct_decode( + cell, + coords, + purpose == FpmQueryPurpose::MixedDecodeBaseline, + ), + } + })(); + coverage.record( + self, + cell_ids, + coords, + purpose, + result.as_ref().map(|(_, kind)| *kind), + ); + result.map(|(result, _)| result) + } + fn resolve_direct_prefill( &self, cell: &FpmForwardCell, coords: &[f64], - ) -> Result { + ) -> Result<(PerformanceResult, DirectFpmResolution), AicError> { let (batch, tokens, kv) = (coords[0], coords[1], coords[2]); if !batch.is_finite() || batch.fract() != 0.0 || batch < 1.0 || batch > u32::MAX as f64 { return Err(self.direct_coverage_err( @@ -517,17 +599,15 @@ impl FpmForwardOp { "batch_size must be a positive integer", )); } - if let Some(value) = cell + if let Some((curve, value)) = cell .direct_prefill .get(&(batch as u32, kv as u32)) .filter(|_| kv == (kv as u32) as f64) - .and_then(|curve| direct_curve_value(curve, tokens)) + .and_then(|curve| direct_curve_value(curve, tokens).map(|value| (curve, value))) { - let support = direct_curve_support( - &cell.direct_prefill[&(batch as u32, kv as u32)], - tokens, - |token| fpm_coordinates(self.phase, &[batch, token, kv]), - ); + let support = direct_curve_support(curve, tokens, |token| { + fpm_coordinates(self.phase, &[batch, token, kv]) + }); let resolution = if support.len() == 1 { FpmQueryResolution::ExactLookup } else { @@ -591,7 +671,7 @@ impl FpmForwardOp { resolution: FpmQueryResolution, support: Vec, decode_baseline: bool, - ) -> Result { + ) -> Result<(PerformanceResult, DirectFpmResolution), AicError> { if !latency.is_finite() || latency <= 0.0 { return Err(self.direct_coverage_err( cell, @@ -599,6 +679,11 @@ impl FpmForwardOp { &format!("interpolation produced an invalid latency ({latency})"), )); } + let coverage_resolution = if resolution == FpmQueryResolution::ExactLookup { + DirectFpmResolution::Measured + } else { + DirectFpmResolution::Interpolated + }; let mut result = PerformanceResult::new(latency, Source::Silicon); result.fpm_queries = Some(Box::new(vec![DirectFpmQueryEvidence { phase: self.phase.as_str().into(), @@ -609,7 +694,7 @@ impl FpmForwardOp { latency_ms: latency, support, }])); - Ok(result) + Ok((result, coverage_resolution)) } fn direct_coverage_err(&self, cell: &FpmForwardCell, coords: &[f64], reason: &str) -> AicError { @@ -636,7 +721,7 @@ impl FpmForwardOp { cell: &FpmForwardCell, coords: &[f64], baseline: bool, - ) -> Result { + ) -> Result<(PerformanceResult, DirectFpmResolution), AicError> { let (lo_row, hi_row) = self.direct_decode_rows(cell, coords)?; let row_value = |row: u32| { let curve = &cell.direct_decode[&row]; diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 5aea8651e..5bf2de12d 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1121,6 +1121,7 @@ struct EngineBuildRequest { systems_path: Option, forward_model: Option, fpm_profile: Option, + worker_type: Option, fpm_interpolation: Option, cp_size: u32, decode_workload_distribution: Option, @@ -1180,6 +1181,7 @@ impl AicEngineBuilder { systems_path: None, forward_model: None, fpm_profile: None, + worker_type: None, fpm_interpolation: None, cp_size: 1, decode_workload_distribution: None, @@ -1576,6 +1578,9 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result "prefill", + crate::ForwardPassWorkerType::Decode => "decode", + crate::ForwardPassWorkerType::Aggregated => "aggregated", + } + .to_owned(), + ), fpm_profile: config .fpm_profile .as_ref() @@ -1875,6 +1888,7 @@ fn engine_build_request( systems_path: systems_path.map(str::to_owned), forward_model: config.forward_model.clone(), fpm_profile: config.extra.get("fpm_profile").cloned(), + worker_type: config.extra.get("worker_type").cloned(), fpm_interpolation: config.extra.get("fpm_interpolation").cloned(), cp_size: config.parallel.cp_size.unwrap_or(1), decode_workload_distribution: None, @@ -2009,6 +2023,17 @@ impl PyForwardPassPerfModel { serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } + /// Size declarative resources using canonical configuration, without timing + /// data or analytical graph construction. + #[staticmethod] + fn estimate_cache_budget(config_json: &str, budget_json: &str) -> PyResult { + let config = parse_forward_pass_config(config_json)?; + let request: crate::perfmodel::FpmCacheBudgetRequest = serde_json::from_str(budget_json) + .map_err(|e| PyValueError::new_err(format!("invalid cache budget: {e}")))?; + let estimate = config.estimate_cache_budget(&request).map_err(aic_to_py)?; + serde_json::to_string(&estimate).map_err(|e| PyValueError::new_err(e.to_string())) + } + /// Share typed FPM option validation with the compilation adapter. #[staticmethod] #[pyo3(signature = (options_json, fmha_quant_mode=None, comm_quant_mode=None, has_profile=false))] @@ -2077,8 +2102,20 @@ impl PyForwardPassPerfModel { "decode_workload_distribution cannot be migrated from a legacy EngineConfig; use ForwardPassPerfModelConfig.estimator_config.op_level".into(), ))); } - let legacy: EngineConfig = + let mut legacy: EngineConfig = serde_json::from_value(value).map_err(|e| PyValueError::new_err(e.to_string()))?; + if legacy + .extra + .get("worker_type") + .is_some_and(|saved| saved != worker_type) + { + return Err(PyValueError::new_err( + "saved FPM worker_type conflicts with requested worker_type", + )); + } + legacy + .extra + .insert("worker_type".into(), worker_type.to_owned()); let mut request = engine_build_request( &legacy, legacy.systems_path.as_ref().and_then(|path| path.to_str()), @@ -2165,7 +2202,7 @@ impl PyForwardPassPerfModel { serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } - /// Direct latency for the qualified homogeneous graph-prefill profile. + /// Uncorrected native prefill latency, including qualified graph-prefill shapes. #[pyo3(signature=(bs, isl, prefix=PrefillArgument::from(0)))] fn predict_prefill_latency( &self, @@ -2174,7 +2211,7 @@ impl PyForwardPassPerfModel { isl: PrefillArgument, prefix: PrefillArgument, ) -> PyResult { - if !(bs.exact && isl.exact && prefix.exact) { + if self.inner.has_prefill_graph_profile() && !(bs.exact && isl.exact && prefix.exact) { return Err(aic_to_py(crate::perf_database::prefill_graph::error( "arguments must be exact Python integers in the unsigned 32-bit range", ))); @@ -2257,6 +2294,31 @@ impl PyForwardPassPerfModel { .map_err(|e| PyValueError::new_err(format!("diagnostics serialize: {e}"))) } + /// Bounded direct-FPM native lookup evidence, or JSON null when disabled. + fn fpm_query_coverage(&self) -> PyResult { + let coverage = self.inner.fpm_query_coverage().map_err(aic_to_py)?; + serde_json::to_string(&coverage) + .map_err(|error| PyValueError::new_err(format!("FPM coverage serialize: {error}"))) + } + + fn predict_decode_latency_total( + &self, + py: Python<'_>, + batch_size: u32, + total_past_kv_tokens: u32, + ) -> PyResult { + py.allow_threads(|| { + self.inner + .predict_decode_latency_total(batch_size, total_past_kv_tokens) + }) + .map_err(aic_to_py) + } + + fn fpm_decode_kv_ceiling(&self, py: Python<'_>) -> PyResult> { + py.allow_threads(|| self.inner.fpm_decode_kv_ceiling()) + .map_err(aic_to_py) + } + /// Regression store labels, readiness and retained counts as JSON. /// Includes cold stores; native AIC models return an empty list. fn regression_store_diagnostics(&self) -> PyResult { diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 143a1a2c1..07cb6827e 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -453,7 +453,8 @@ enum AicPhaseProvider { struct AicTimingModel { engine: Py, - diagnostic_model: Option, + diagnostic_model: Option>, + fpm_model: Option>, decoder_replay: bool, phase_provider: AicPhaseProvider, use_fpm_decode_totals: bool, @@ -486,8 +487,10 @@ impl AicTimingModel { } config.validate_parallel_shape()?; config.resolved_memory_fraction()?; - let model = ForwardPassPerfModel::best_available(request) - .context("AIC timing provider could not construct the requested estimator")?; + let model = Arc::new( + ForwardPassPerfModel::best_available(request) + .context("AIC timing provider could not construct the requested estimator")?, + ); let provenance = model .provenance() .context("canonical estimator omitted construction provenance")?; @@ -520,20 +523,17 @@ impl AicTimingModel { "AIC regression estimator is not ready: offline replay requires trained observations or a native estimator" )?; let phase_provider = AicPhaseProvider::Native(Arc::clone(&native)); - let (engine, fpm_decode_kv_ceiling) = Python::with_gil(|py| -> PyResult<_> { - let engine = Py::new(py, crate::AicEngine::from_shared_engine(native))?.into_any(); - let ceiling = if use_fpm_decode_totals { - engine - .bind(py) - .call_method0("fpm_decode_kv_ceiling")? - .extract::>()? - } else { - None - }; - Ok((engine, ceiling)) + let engine = Python::with_gil(|py| -> PyResult<_> { + Ok(Py::new(py, crate::AicEngine::from_shared_engine(native))?.into_any()) })?; + let fpm_decode_kv_ceiling = if use_fpm_decode_totals { + model.fpm_decode_kv_ceiling()? + } else { + None + }; Ok(Self { engine, + fpm_model: use_fpm_decode_totals.then(|| Arc::clone(&model)), diagnostic_model: Some(model), decoder_replay: config.decoder_replay, phase_provider, @@ -785,6 +785,11 @@ impl TimingModel for AicTimingModel { model.static_prefill_detailed(batch_size, mean_isl, mean_prefix)?, ); } + if let Some(model) = &self.fpm_model { + return model + .predict_prefill_latency(batch_size, mean_isl, mean_prefix) + .context("AIC prefill prediction failed"); + } Python::with_gil(|py| { self.engine .bind(py) @@ -824,6 +829,11 @@ impl TimingModel for AicTimingModel { return self .record_fpm_estimate((false, batch_size, 0, total_past_kv_tokens), estimate); } + if let Some(model) = &self.fpm_model { + return model + .predict_decode_latency_total(batch_size, total_past_kv_tokens) + .context("AIC decode prediction failed"); + } return Python::with_gil(|py| { self.engine .bind(py) @@ -919,45 +929,70 @@ fn checked_u32(value: usize, name: &str) -> Result { u32::try_from(value).with_context(|| format!("{name} {value} exceeds AIC's u32 limit")) } +fn aic_capacity_kwargs<'py>( + py: Python<'py>, + config: &AicTimingConfig, + role: &ReplayRoleConfig, +) -> PyResult> { + let (memory_fraction_kind, memory_fraction_value) = config + .resolved_memory_fraction() + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + let kwargs = PyDict::new(py); + kwargs.set_item("backend_version", config.resolved_backend_version())?; + kwargs.set_item("scheduler_block_size", role.rank.block_size)?; + kwargs.set_item("max_num_tokens", role.rank.max_num_batched_tokens)?; + kwargs.set_item("max_batch_size", role.rank.max_num_seqs)?; + if config.fpm_profile.is_some() { + kwargs.set_item("context_length", role.rank.max_model_len)?; + kwargs.set_item( + "worker_type", + match config + .worker_type + .unwrap_or(ForwardPassWorkerType::Aggregated) + { + ForwardPassWorkerType::Prefill => "prefill", + ForwardPassWorkerType::Decode => "decode", + ForwardPassWorkerType::Aggregated => "aggregated", + }, + )?; + } + kwargs.set_item("memory_fraction_kind", memory_fraction_kind)?; + kwargs.set_item("memory_fraction_value", memory_fraction_value)?; + kwargs.set_item("tp_size", config.tp)?; + kwargs.set_item("pp_size", config.pp)?; + kwargs.set_item("attention_dp_size", config.attention_dp)?; + kwargs.set_item("moe_tp_size", config.moe_tp_size)?; + kwargs.set_item("moe_ep_size", config.moe_ep_size)?; + kwargs.set_item("gemm_quant_mode", config.gemm_dtype.as_deref())?; + kwargs.set_item("moe_quant_mode", config.moe_dtype.as_deref())?; + kwargs.set_item("fmha_quant_mode", config.fmha_dtype.as_deref())?; + kwargs.set_item("kvcache_quant_mode", config.kv_cache_dtype.as_deref())?; + kwargs.set_item("comm_quant_mode", config.comm_dtype.as_deref())?; + kwargs.set_item("moe_backend", config.moe_backend.as_deref())?; + kwargs.set_item("attention_backend", config.attention_backend.as_deref())?; + kwargs.set_item("enable_eplb", config.enable_eplb)?; + kwargs.set_item("wideep_num_slots", config.wideep_num_slots)?; + kwargs.set_item( + "fpm_profile", + config + .fpm_profile + .as_ref() + .map(|profile| serde_json::to_string(profile).expect("JSON profile")), + )?; + kwargs.set_item( + "cuda_graph_reserved_bytes", + config.cuda_graph_reserved_bytes, + )?; + // Capacity intentionally omits NextN until AIC's Eagle memory model no + // longer returns negative KV capacity. Timing compilation still uses it. + kwargs.set_item("systems_path", config.systems_path.as_deref())?; + Ok(kwargs) +} + fn estimate_aic_num_gpu_blocks(config: &AicTimingConfig, role: &ReplayRoleConfig) -> Result { - let (memory_fraction_kind, memory_fraction_value) = config.resolved_memory_fraction()?; Python::with_gil(|py| -> PyResult { let memory = PyModule::import(py, "aisimulate_core.sdk.memory")?; - let kwargs = PyDict::new(py); - kwargs.set_item("backend_version", config.resolved_backend_version())?; - kwargs.set_item("scheduler_block_size", role.rank.block_size)?; - kwargs.set_item("max_num_tokens", role.rank.max_num_batched_tokens)?; - kwargs.set_item("max_batch_size", role.rank.max_num_seqs)?; - kwargs.set_item("memory_fraction_kind", memory_fraction_kind)?; - kwargs.set_item("memory_fraction_value", memory_fraction_value)?; - kwargs.set_item("tp_size", config.tp)?; - kwargs.set_item("pp_size", config.pp)?; - kwargs.set_item("attention_dp_size", config.attention_dp)?; - kwargs.set_item("moe_tp_size", config.moe_tp_size)?; - kwargs.set_item("moe_ep_size", config.moe_ep_size)?; - kwargs.set_item("gemm_quant_mode", config.gemm_dtype.as_deref())?; - kwargs.set_item("moe_quant_mode", config.moe_dtype.as_deref())?; - kwargs.set_item("fmha_quant_mode", config.fmha_dtype.as_deref())?; - kwargs.set_item("kvcache_quant_mode", config.kv_cache_dtype.as_deref())?; - kwargs.set_item("comm_quant_mode", config.comm_dtype.as_deref())?; - kwargs.set_item("moe_backend", config.moe_backend.as_deref())?; - kwargs.set_item("attention_backend", config.attention_backend.as_deref())?; - kwargs.set_item("enable_eplb", config.enable_eplb)?; - kwargs.set_item("wideep_num_slots", config.wideep_num_slots)?; - kwargs.set_item( - "fpm_profile", - config - .fpm_profile - .as_ref() - .map(|profile| serde_json::to_string(profile).expect("JSON profile")), - )?; - kwargs.set_item( - "cuda_graph_reserved_bytes", - config.cuda_graph_reserved_bytes, - )?; - // Capacity intentionally omits NextN until AIC's Eagle memory model no - // longer returns negative KV capacity. Timing compilation still uses it. - kwargs.set_item("systems_path", config.systems_path.as_deref())?; + let kwargs = aic_capacity_kwargs(py, config, role)?; memory .getattr("estimate_num_gpu_blocks")? .call( @@ -973,6 +1008,31 @@ fn estimate_aic_num_gpu_blocks(config: &AicTimingConfig, role: &ReplayRoleConfig .map_err(|error| anyhow!("AIC KV-cache capacity estimation failed: {error}")) } +fn estimate_aic_grouped_budget( + config: &AicTimingConfig, + role: &ReplayRoleConfig, +) -> Result { + let serialized = Python::with_gil(|py| -> PyResult { + let memory = PyModule::import(py, "aisimulate_core.sdk.memory")?; + let kwargs = aic_capacity_kwargs(py, config, role)?; + kwargs.del_item("scheduler_block_size")?; + kwargs.set_item("context_length", role.rank.max_model_len)?; + let budget = memory.getattr("estimate_kv_cache")?.call( + ( + config.model.as_str(), + config.system.as_str(), + config.backend.as_str(), + ), + Some(&kwargs), + )?; + PyModule::import(py, "json")? + .call_method1("dumps", (budget,))? + .extract() + }) + .map_err(|error| anyhow!("AIC grouped-cache budget estimation failed: {error}"))?; + serde_json::from_str(&serialized).context("invalid AIC grouped-cache budget") +} + fn materialize_aic_capacity( config: &AicTimingConfig, role: &mut ReplayRoleConfig, @@ -1015,6 +1075,54 @@ fn materialize_aic_capacity( timing_depth as usize == engine_nextn, "AIC speculative depth={timing_depth} does not match engine aic_nextn={engine_nextn}" ); + let resources = if config.backend == "vllm" && config.pp == 1 { + config + .estimator_request( + config + .worker_type + .unwrap_or(ForwardPassWorkerType::Aggregated), + )? + .fpm_resources()? + } else { + None + }; + if let Some(resources) = &resources { + resources.require_memory()?; + ensure!( + resources.runtime_memory.is_none() || !capacity_is_explicit, + "runtime FPM memory cannot use fixed num_gpu_blocks; use the recorded cache allocation" + ); + } + if resources.as_ref().is_some_and(|resources| { + resources.cache_layout == crate::perfmodel::fpm::FpmCacheLayout::Grouped + }) { + ensure!( + !capacity_is_explicit, + "grouped FPM cache cannot use fixed num_gpu_blocks; use the profile byte budget" + ); + let budget = estimate_aic_grouped_budget(config, role)?; + ensure!( + role.rank.kv_cache_groups.is_empty() + || role.rank.kv_cache_groups == budget.cache_groups, + "engine cache groups conflict with the canonical FPM profile" + ); + ensure!( + role.rank + .kv_cache_capacity_bytes + .is_none_or(|bytes| bytes == budget.total_kv_size_bytes), + "engine cache byte capacity conflicts with the canonical FPM resource budget" + ); + role.rank.kv_cache_groups = budget.cache_groups; + role.rank.kv_cache_capacity_bytes = Some(budget.total_kv_size_bytes); + // Keep the historical positive field for EngineConfig compatibility. + // The grouped allocator exclusively consumes groups and byte capacity. + role.rank.validate()?; + return Ok(()); + } + ensure!( + role.rank.kv_cache_groups.is_empty() && role.rank.kv_cache_capacity_bytes.is_none(), + "grouped cache allocation requires matching canonical FPM profile resources" + ); if capacity_is_explicit || role.rank.state_cache.is_some() { return Ok(()); } @@ -1031,12 +1139,17 @@ fn materialize_aic_capacity( fn cap_role_capacity_to_fpm_decode_domain( role: &mut ReplayRoleConfig, decode_kv_ceiling: Option, - capacity_is_explicit: bool, + capacity_is_recorded: bool, ) -> Result<()> { + // A timing-table token ceiling is not a physical grouped-cache budget. + // Queries retain their logical contexts and report uncovered points normally. + if !role.rank.kv_cache_groups.is_empty() { + return Ok(()); + } let Some(decode_kv_ceiling) = decode_kv_ceiling else { return Ok(()); }; - if capacity_is_explicit || role.rank.state_cache.is_some() { + if capacity_is_recorded || role.rank.state_cache.is_some() { return Ok(()); } let covered_blocks = decode_kv_ceiling as usize / role.rank.block_size; @@ -1078,6 +1191,7 @@ fn resolve_role_timing( capacity_is_explicit: bool, worker_type: ForwardPassWorkerType, capture_performance_diagnostics: bool, + coverage_sources: &mut Vec<(ForwardPassWorkerType, Arc)>, ) -> Result>> { let TimingModelConfig::External { provider, config } = role.rank.timing_model.clone() else { return Ok(None); @@ -1089,7 +1203,20 @@ fn resolve_role_timing( ); let mut config: AicTimingConfig = serde_json::from_value(config).context("invalid AIC timing provider configuration")?; + let resources = if config.backend == "vllm" && config.pp == 1 { + config.estimator_request(worker_type)?.fpm_resources()? + } else { + None + }; + if let Some(resources) = &resources { + resources.require_memory()?; + } let mut timing = AicTimingModel::build(&mut config, worker_type)?; + if let Some(model) = &timing.fpm_model + && model.fpm_query_coverage()?.is_some() + { + coverage_sources.push((worker_type, Arc::clone(model))); + } if !capture_performance_diagnostics { timing.diagnostic_model = None; } @@ -1102,7 +1229,10 @@ fn resolve_role_timing( cap_role_capacity_to_fpm_decode_domain( role, timing.fpm_decode_kv_ceiling, - capacity_is_explicit, + capacity_is_explicit + || resources + .as_ref() + .is_some_and(|resources| resources.runtime_memory.is_some()), )?; let mut resolved = serde_json::to_value(config.estimator_request(worker_type)?)?; if let Some(object) = resolved.as_object_mut() { @@ -1233,16 +1363,36 @@ fn resolve_kv_capacity_concurrency( expected_tokens > 0, "KV load requires positive token lengths" ); - let per_rank_tokens = role - .rank - .num_gpu_blocks - .checked_mul(role.rank.block_size) - .context("KV capacity overflow")?; - let total_tokens = per_rank_tokens - .checked_mul(role.dp_size as usize) - .and_then(|value| value.checked_mul(replicas)) - .context("aggregate KV capacity overflow")?; - let capacity = total_tokens / expected_tokens; + let capacity = if !role.rank.kv_cache_groups.is_empty() { + let request_bytes = role + .rank + .kv_cache_groups + .iter() + .try_fold(0u64, |sum, group| { + sum.checked_add(group.peak_request_bytes(expected_tokens as u64, 1)?) + .context("grouped KV request bytes overflow") + })?; + let per_rank = role + .rank + .kv_cache_capacity_bytes + .context("grouped KV load requires byte capacity")? + / request_bytes; + usize::try_from(per_rank)? + .checked_mul(role.dp_size as usize) + .and_then(|value| value.checked_mul(replicas)) + .context("aggregate grouped KV concurrency overflow")? + } else { + let per_rank_tokens = role + .rank + .num_gpu_blocks + .checked_mul(role.rank.block_size) + .context("KV capacity overflow")?; + let total_tokens = per_rank_tokens + .checked_mul(role.dp_size as usize) + .and_then(|value| value.checked_mul(replicas)) + .context("aggregate KV capacity overflow")?; + total_tokens / expected_tokens + }; ensure!( capacity > 0, "candidate KV capacity cannot hold one request" @@ -1908,6 +2058,89 @@ fn infer_aic_timing_topology(engine: &mut serde_json::Value) -> Result<()> { Ok(()) } +#[derive(Debug)] +struct ReplayCoverageFailure { + error: anyhow::Error, + coverage: serde_json::Value, +} + +impl std::fmt::Display for ReplayCoverageFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "{}", self.error) + } +} + +impl std::error::Error for ReplayCoverageFailure { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.error.source() + } +} + +fn replay_failure( + error: anyhow::Error, + coverage: Result>, +) -> anyhow::Error { + match coverage { + Ok(Some(coverage)) => ReplayCoverageFailure { error, coverage }.into(), + Ok(None) => error, + Err(coverage_error) => error.context(format!( + "FPM coverage diagnostics failed: {coverage_error:#}" + )), + } +} + +fn replay_fpm_coverage( + sources: &[(ForwardPassWorkerType, Arc)], + report: Option<&crate::replay::ReplayReport>, +) -> Result> { + if sources.is_empty() { + return Ok(None); + } + let mut measured = 0u64; + let mut interpolated = 0u64; + let mut unsupported = 0u64; + let mut roles = Vec::with_capacity(sources.len()); + for (role, model) in sources { + let coverage = model + .fpm_query_coverage()? + .context("requested FPM query coverage is unavailable")?; + measured += coverage.queries.measured; + interpolated += coverage.queries.interpolated; + unsupported += coverage.queries.unsupported; + roles.push(serde_json::json!({"worker_type": role, "coverage": coverage})); + } + let replay_completed = report.is_some_and(|report| { + let counts = &report.request_counts; + counts.num_requests > 0 + && counts.completed_requests == counts.num_requests + && report.agentic_graph.as_ref().is_none_or(|graph| { + graph.node_count == counts.completed_requests + && report + .agentic_play_outcomes + .as_ref() + .is_some_and(|outcomes| { + outcomes.len() == graph.play_count + && outcomes.iter().all(|outcome| { + matches!( + outcome.status, + crate::replay::loadgen::AgenticPlayStatus::Completed + ) + }) + }) + }) + }); + let covered = replay_completed && measured + interpolated > 0 && unsupported == 0; + Ok(Some(serde_json::json!({ + "schema_version": 1, + "status": if covered { "covered" } else { "incomplete" }, + "replay_completed": replay_completed, + "counting_unit": "native_lookup_resolutions", + "queries": {"measured": measured, "interpolated": interpolated, "unsupported": unsupported}, + "roles": roles, + "accuracy": "not_assessed", + }))) +} + fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let (mut spec, mut traffic, capture_performance_diagnostics) = match serde_json::from_str(payload).context("invalid AISimulate execution ReplaySpec")? { @@ -1951,8 +2184,10 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { ReplayTopology::Disaggregated { .. } => 2, }; let mut power_sources = Vec::with_capacity(expected_power_sources); + let mut coverage_sources = Vec::new(); - let (mut report, artifacts, resolved_weka_timestamp_basis) = match spec.topology.clone() { + let replay_result = (|| -> Result<_> { + match spec.topology.clone() { ReplayTopology::Aggregated { .. } => { let mut role = aggregated_role(&engine_config); let capacity_is_explicit = role @@ -1963,6 +2198,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { capacity_is_explicit, ForwardPassWorkerType::Aggregated, capture_performance_diagnostics, + &mut coverage_sources, )?; if let Some(timing) = timing.as_ref() { power_sources.push(TimingPowerSource::new(Arc::clone(timing), &role.rank)); @@ -2061,12 +2297,14 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { prefill_capacity_is_explicit, ForwardPassWorkerType::Prefill, capture_performance_diagnostics, + &mut coverage_sources, )?; let decode_timing = resolve_role_timing( &mut decode, decode_capacity_is_explicit, ForwardPassWorkerType::Decode, capture_performance_diagnostics, + &mut coverage_sources, )?; if let Some(timing) = prefill_timing.as_ref() { power_sources.push(TimingPowerSource::new(Arc::clone(timing), &prefill.rank)); @@ -2129,8 +2367,19 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { ) .map(|(report, artifacts)| (report, artifacts, resolved_basis)) } - } - .context("AISimulate replay failed")?; + } + .map_err(Into::into) + })() + .context("AISimulate replay failed"); + let (mut report, artifacts, resolved_weka_timestamp_basis) = match replay_result { + Ok(result) => result, + Err(error) => { + return Err(replay_failure( + error, + replay_fpm_coverage(&coverage_sources, None), + )); + } + }; let (timing_evidence, unavailable_reason) = if power_sources.len() == expected_power_sources { ( replay_timing_evidence(&power_sources)?, @@ -2154,6 +2403,9 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { )?)); let mut report_json = serde_json::to_value(&report).context("serializing AISimulate replay report summary")?; + if let Some(coverage) = replay_fpm_coverage(&coverage_sources, Some(&report))? { + report_json["fpm_query_coverage"] = coverage; + } if capture_performance_diagnostics { let performance_evidence = replay_timing_evidence_with_fpm(&power_sources, true)?; report_json["performance_diagnostics"] = @@ -2229,7 +2481,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { } fn replay_python_error(error: anyhow::Error) -> PyErr { - if error.chain().any(|cause| { + let exception = if error.chain().any(|cause| { matches!( cause.downcast_ref::(), Some(crate::replay::ReplayError::ResourceLimited(_)) @@ -2238,7 +2490,15 @@ fn replay_python_error(error: anyhow::Error) -> PyErr { PyMemoryError::new_err(format!("{error:#}")) } else { PyRuntimeError::new_err(format!("{error:#}")) + }; + if let Some(failure) = error.downcast_ref::() { + Python::with_gil(|py| { + let _ = exception + .value(py) + .setattr("fpm_query_coverage", failure.coverage.to_string()); + }); } + exception } /// Execute one canonical serialized ReplaySpec and return serialized report JSON. @@ -2274,6 +2534,21 @@ mod tests { use super::*; + #[test] + fn replay_coverage_errors_preserve_the_original_chain() { + let original = || anyhow::anyhow!("original failure").context("AISimulate replay failed"); + let wrapped = replay_failure(original(), Ok(Some(serde_json::json!({})))); + assert_eq!( + format!("{wrapped:#}"), + "AISimulate replay failed: original failure" + ); + let missing = replay_failure(original(), Err(anyhow::anyhow!("coverage unavailable"))); + assert_eq!( + format!("{missing:#}"), + "FPM coverage diagnostics failed: coverage unavailable: AISimulate replay failed: original failure" + ); + } + #[test] fn replay_python_error_preserves_contextual_resource_failure() { pyo3::prepare_freethreaded_python(); @@ -2726,6 +3001,7 @@ mod tests { AicTimingModel { engine, diagnostic_model: None, + fpm_model: None, decoder_replay: false, phase_provider: AicPhaseProvider::Python, use_fpm_decode_totals, @@ -3075,6 +3351,51 @@ mod tests { } } + #[test] + fn fixed_capacity_does_not_require_vllm_pp1_profile_resources() { + for (backend, pp) in [("sglang", 1), ("vllm", 2)] { + let mut config = aic_config(); + config.backend = backend.into(); + config.pp = pp; + config.fpm_profile = Some(serde_json::Map::new()); + let mut role = aggregated_role(&ReplayEngineConfig::default()); + role.rank.backend = if backend == "sglang" { + Backend::Sglang + } else { + Backend::Vllm + }; + materialize_aic_capacity(&config, &mut role, true, |_, _| unreachable!()).unwrap(); + } + } + + #[test] + fn grouped_load_and_timing_domain_preserve_byte_capacity() { + let mut role = aggregated_role(&ReplayEngineConfig::default()); + role.dp_size = 3; + role.rank.kv_cache_groups = serde_json::from_value(serde_json::json!([{ + "name":"window", "kind":"attention", "num_layers":1, + "block_size_tokens":16, "page_size_bytes":160, "sliding_window":32 + }])) + .unwrap(); + role.rank.kv_cache_capacity_bytes = Some(1000); + cap_role_capacity_to_fpm_decode_domain(&mut role, Some(1), false).unwrap(); + assert_eq!(role.rank.kv_cache_capacity_bytes, Some(1000)); + let mut traffic: RuntimeTraffic = serde_json::from_value(serde_json::json!({ + "source_type":"synthetic", "load_type":"kv_capacity_fraction", + "kv_load_ratio":1.0, "isl":1152, "osl":4 + })) + .unwrap(); + // A 32-token window can straddle three 16-token pages (480 bytes). + // Each rank holds two requests; three DP ranks and two replicas hold 12. + assert_eq!( + resolve_kv_capacity_concurrency(&mut traffic, &role, 2).unwrap(), + Some(12) + ); + assert_eq!(traffic.concurrency, Some(12)); + role.rank.kv_cache_capacity_bytes = Some(479); + assert!(resolve_kv_capacity_concurrency(&mut traffic, &role, 2).is_err()); + } + #[test] fn ngram_timing_requires_matching_scheduler_depth_and_no_mtp() { let mut config = aic_config(); @@ -3946,10 +4267,10 @@ mod tests { }); let mut timing = python_timing_model(engine, false); timing.phase_provider = AicPhaseProvider::Native(Arc::clone(&native)); - timing.diagnostic_model = Some(ForwardPassPerfModel::from_engine( + timing.diagnostic_model = Some(Arc::new(ForwardPassPerfModel::from_engine( native, ForwardPassPerfOptions::default(), - )); + ))); for kind in ["diagnostics", "evidence"] { for mode in ["static_ctx", "static_gen"] { let error = timing diff --git a/crates/core/src/replay/components/engine.rs b/crates/core/src/replay/components/engine.rs index 55025f1e5..62390805e 100644 --- a/crates/core/src/replay/components/engine.rs +++ b/crates/core/src/replay/components/engine.rs @@ -75,11 +75,13 @@ struct RankMetricsState { } impl RankMetricsState { - fn new(dp_rank: u32, total_blocks: u64) -> Self { + fn new(dp_rank: u32, total_blocks: u64, kv_cache_capacity_bytes: Option) -> Self { Self { latest: Metrics { dp_rank, total_blocks, + kv_cache_capacity_bytes, + kv_cache_used_bytes: kv_cache_capacity_bytes.map(|_| 0), ..Metrics::default() }, interval: ReplaySchedulerIntervalMetrics::default(), @@ -115,6 +117,8 @@ impl RankMetricsState { active_blocks: self.latest.active_blocks, inactive_blocks: self.latest.inactive_blocks, total_blocks: self.latest.total_blocks, + kv_cache_used_bytes: self.latest.kv_cache_used_bytes, + kv_cache_capacity_bytes: self.latest.kv_cache_capacity_bytes, active_cache_usage: self.latest.cache_usage, physical_cache_usage: self.latest.physical_cache_usage, running_requests: self.latest.running_requests, @@ -229,7 +233,13 @@ where for worker in self.workers.iter_mut().filter_map(Option::as_mut) { worker.telemetry = Some( (0..dp_size) - .map(|dp_rank| RankMetricsState::new(dp_rank, total_blocks)) + .map(|dp_rank| { + RankMetricsState::new( + dp_rank, + total_blocks, + self.factory.kv_cache_capacity_bytes(), + ) + }) .collect(), ); } @@ -319,7 +329,13 @@ where } let telemetry = self.telemetry.as_ref().map(|_| { (0..self.factory.dp_size()) - .map(|dp_rank| RankMetricsState::new(dp_rank, self.factory.total_blocks())) + .map(|dp_rank| { + RankMetricsState::new( + dp_rank, + self.factory.total_blocks(), + self.factory.kv_cache_capacity_bytes(), + ) + }) .collect() }); self.workers.push(Some(LogicalWorker { @@ -1501,7 +1517,7 @@ mod tests { #[test] fn scheduler_cache_metrics_accumulate_and_drain_per_tick() { - let mut state = RankMetricsState::new(2, 100); + let mut state = RankMetricsState::new(2, 100, None); state .update( &Metrics { diff --git a/crates/core/src/replay/engine.rs b/crates/core/src/replay/engine.rs index d17aa4e28..9e59b4017 100644 --- a/crates/core/src/replay/engine.rs +++ b/crates/core/src/replay/engine.rs @@ -208,6 +208,7 @@ pub struct ReplayRoleFactory { tensor_parallel_size: u32, backend: Backend, total_blocks: u64, + kv_cache_capacity_bytes: Option, pub(crate) g3_config: Option, pub(crate) g3_block_bytes: usize, pub(crate) g3_tier: Option, @@ -289,6 +290,10 @@ impl ReplayRoleFactory { pub fn total_blocks(&self) -> u64 { self.total_blocks } + + pub(crate) fn kv_cache_capacity_bytes(&self) -> Option { + self.kv_cache_capacity_bytes + } } /// Resolves built-in or Runner-provided timing once, then creates role factories. @@ -351,9 +356,14 @@ impl ReplayEngineFactory { WorkerStage::Decode => self.decode_timing.as_ref().or(self.timing.as_ref()), }; let backend = role.rank.backend; - let total_blocks = u64::try_from(role.rank.num_gpu_blocks).map_err(|_| { - ReplayError::InvalidSpec("engine KV block count exceeds the metrics range".into()) - })?; + let kv_cache_capacity_bytes = role.rank.kv_cache_capacity_bytes; + let total_blocks = if role.rank.kv_cache_groups.is_empty() { + u64::try_from(role.rank.num_gpu_blocks).map_err(|_| { + ReplayError::InvalidSpec("engine KV block count exceeds the metrics range".into()) + })? + } else { + 0 + }; role.rank.validate().map_err(engine_error)?; let g3_config = role.rank.g3_offload.take(); let g3_block_bytes = role @@ -373,6 +383,7 @@ impl ReplayEngineFactory { tensor_parallel_size: role.tensor_parallel_size, backend, total_blocks, + kv_cache_capacity_bytes, g3_config, g3_block_bytes, g3_tier: None, diff --git a/crates/core/src/replay/telemetry.rs b/crates/core/src/replay/telemetry.rs index cd4e656a3..07be8d8db 100644 --- a/crates/core/src/replay/telemetry.rs +++ b/crates/core/src/replay/telemetry.rs @@ -29,14 +29,22 @@ pub enum ReplayTelemetrySampleKind { pub struct ReplaySchedulerMetricsSnapshot { pub worker_id: usize, pub dp_rank: u32, - /// Backend-native legacy occupancy. vLLM counts active references; SGLang - /// counts occupied page-pool blocks, including radix-resident pages. + /// Backend-native occupancy. vLLM counts active references (physical pages + /// across all groups for grouped caches); SGLang counts occupied page-pool + /// blocks, including radix-resident pages. pub active_blocks: u64, /// Reusable resident blocks excluded from `active_blocks` (vLLM only; /// SGLang reports zero because its legacy occupancy already includes them). pub inactive_blocks: u64, pub total_blocks: u64, - /// Legacy/backend-native `active_blocks / total_blocks` utilization. + /// Shared physical grouped-cache occupancy and capacity. Linear caches + /// retain the legacy block metrics and omit these fields. + #[serde(skip_serializing_if = "Option::is_none")] + pub kv_cache_used_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub kv_cache_capacity_bytes: Option, + /// Legacy/backend-native block utilization, or physical byte utilization + /// for grouped caches (which have no scalar total block capacity). pub active_cache_usage: f64, /// Physical resident utilization; equal to active utilization for SGLang. pub physical_cache_usage: f64, diff --git a/crates/core/tests/runtime_contract.rs b/crates/core/tests/runtime_contract.rs index 4a6262939..66009c185 100644 --- a/crates/core/tests/runtime_contract.rs +++ b/crates/core/tests/runtime_contract.rs @@ -100,6 +100,73 @@ impl ReplayTelemetryObserver for FailingTelemetryObserver { } } +#[test] +fn grouped_cache_replay_emits_physical_bytes_without_scalar_block_capacity() { + let samples = Arc::new(Mutex::new(Vec::new())); + let mut req = request("windowed", 10.0, 1152, 4); + req.dp_rank = Some(0); + let mut spec = aggregated_spec(Backend::Vllm, 1, 0.0, vec![req]); + spec.engine["dp_size"] = serde_json::json!(2); + let rank = &mut spec.engine["rank"]; + rank["enable_prefix_caching"] = false.into(); + rank["max_num_batched_tokens"] = 128.into(); + rank["kv_cache_capacity_bytes"] = 4000.into(); + rank["kv_cache_groups"] = serde_json::json!([ + {"name":"full", "kind":"attention", "num_layers":1, "block_size_tokens":64, "page_size_bytes":64}, + {"name":"window", "kind":"attention", "num_layers":2, "block_size_tokens":16, "page_size_bytes":16, "sliding_window":512}, + {"name":"conv", "kind":"convolution", "num_layers":3, "block_size_tokens":4, "page_size_bytes":8, "sliding_window":4} + ]); + let plain = Replayer::new(spec.clone(), ReplayEngineFactory::new()) + .unwrap() + .run() + .unwrap(); + let report = Replayer::new(spec, ReplayEngineFactory::new()) + .unwrap() + .with_telemetry_observer( + 1.0, + Box::new(RecordingTelemetryObserver { + samples: Arc::clone(&samples), + }), + ) + .unwrap() + .run() + .unwrap(); + assert_eq!(report.request_counts.completed_requests, 1); + assert_eq!( + report.request_counts.total_output_tokens, + plain.request_counts.total_output_tokens + ); + let samples = samples.lock().unwrap(); + let mut saw_resident_pages = false; + for sample in samples.iter() { + assert_eq!(sample.decode_scheduler_metrics.len(), 2); + for row in &sample.decode_scheduler_metrics { + assert_eq!((row.total_blocks, row.inactive_blocks), (0, 0)); + assert_eq!(row.kv_cache_capacity_bytes, Some(4000)); + let used = row.kv_cache_used_bytes.unwrap(); + assert_eq!(row.active_blocks > 0, used > 0); + assert!(used <= 4000); + assert_eq!(row.physical_cache_usage, used as f64 / 4000.0); + if row.dp_rank == 1 { + assert_eq!(used, 0); + } + saw_resident_pages |= used > 0; + let emitted = serde_json::to_value(row).unwrap(); + assert_eq!(emitted["kv_cache_used_bytes"], used); + assert_eq!(emitted["kv_cache_capacity_bytes"], 4000); + } + } + assert!(saw_resident_pages); + assert!( + samples + .last() + .unwrap() + .decode_scheduler_metrics + .iter() + .all(|row| row.kv_cache_used_bytes == Some(0)) + ); +} + #[test] fn replay_accept_length_counts_decode_work_before_output_truncation() { for backend in [Backend::Vllm, Backend::Sglang] { diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index 0b4306f46..1f0b05182 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -11,7 +11,7 @@ use aisimulate_core::{ AicEngine, AicEngineBuilder, AicError, BackendKind, DatabaseMode, EstimationMode, EstimatorConfig, ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassRegressionStoreDiagnostics, ForwardPassWorkerType, FpmInterpolationMethod, - KvCacheEstimateRequest, + FpmQueryCoverage, KvCacheEstimateRequest, }; /// Compile the ergonomic engine builder without starting embedded Python. @@ -80,6 +80,19 @@ pub fn profile_model( ForwardPassPerfModel::best_available(config) } +/// Coverage is owned by the model returned by the one construction API. +pub fn fpm_coverage(model: &ForwardPassPerfModel) -> Result, AicError> { + model.fpm_query_coverage() +} + +pub fn covered_prefill(model: &ForwardPassPerfModel) -> Result { + model.predict_prefill_latency(1, 128, 0) +} + +pub fn covered_decode(model: &ForwardPassPerfModel) -> Result { + model.predict_decode_latency_total(1, 128) +} + pub fn best_available_model( config: ForwardPassPerfModelConfig, ) -> Result { @@ -583,6 +596,11 @@ pub fn rebuild_replay_report_literals( } /// Detailed phase evidence is reachable through the canonical model. -pub fn operation_diagnostics(model: &ForwardPassPerfModel) -> Result, AicError> { +pub fn operation_diagnostics( + model: &ForwardPassPerfModel, +) -> Result< + Vec, + AicError, +> { model.static_phase_diagnostics(1, 128, 0, true) } diff --git a/docs/core-api.md b/docs/core-api.md index 6ee9a4738..cf8bb9870 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -151,6 +151,123 @@ because deserialization defaults it to zero. Rust source that constructs patterns must include the new field. This source migration is part of the next minor API update. +### FPM profile cache groups and byte budgets + +FPM resources support `cache_layout: linear` with a positive +`kv_bytes_per_token`, or `cache_layout: grouped` with a nonempty `cache_groups` +list and no scalar token rate. Existing linear profiles keep their serialized +shape. Python exports `FpmCacheGroup` from `aisimulate.fpm_profile` and +`aisimulate_core.fpm_profile`; Rust exports it from `aisimulate_core::perfmodel`. + +Memory can be pending, declared, or observed at runtime. The four legacy fields +`weights_bytes`, `activations_bytes`, `runtime_overhead_bytes`, and +`comm_overhead_bytes` are optional. A complete set retains the existing non-KV +budget calculation. A partial set is planning evidence only: missing values are +not zero, and cache sizing or replay fails with an instruction to finalize +runtime memory. Schema validation, collection planning, and direct timing queries +can use a pending profile. Omitted fields remain absent from serialized output. + +Alternatively, `resources.runtime_memory` records the selected deployment's +initialized cache allocation: + +```json +{ + "kv_cache_bytes": 1073741824, + "gpu_memory_utilization": 0.9, + "max_model_len": 32768, + "provenance": "Verified collection worker initialization; see saved evidence." +} +``` + +The byte count is a positive integer no greater than `2**53`; utilization is +finite and in `(0, 1]`; context and provenance are required. Runtime memory +cannot coexist with any of the four legacy non-KV fields. Python exposes +`FpmRuntimeMemoryProfile` and resource properties `memory_source` +(`pending`, `declared`, or `runtime`) and `memory_ready`; `require_memory()` +rejects pending resources. Rust exposes `FpmRuntimeMemoryConfig` and +`FpmResourceConfig::require_memory()`. Rust callers using resource struct literals +must wrap legacy byte values in `Some(...)` and supply `runtime_memory: None`. + +Each group has a unique `name`, `kind` (`attention` or `convolution`), positive +`num_layers`, `block_size_tokens`, and `page_size_bytes`, and an optional positive +`sliding_window`. An omitted or null window retains full history; convolution +groups require a window. `page_size_bytes` is the **rank-local aggregate for all +layers in the group**, including runtime padding. Do not multiply it by +`num_layers` again. Runtime block sizes and padding are deployment inputs; model +geometry alone does not establish them. See the +[grouped-profile review workflow](fpm-self-service.md#review-grouped-cache-resources). + +Use `RustForwardPassPerfModel.estimate_cache_budget(config, budget)` with the +same canonical `ForwardPassPerfModelConfig` used for timing. Rust exposes +`ForwardPassPerfModelConfig::estimate_cache_budget(&request)` and +`ForwardPassPerfModel::estimate_cache_budget(&request)`, taking +`FpmCacheBudgetRequest` and returning `FpmCacheBudget`. The config method checks +the exact profile deployment and scheduler envelope without constructing a +timing model, loading timings, or building an operation graph. Grouped planning +requires the native extension; it does not require GPU access or measured data. + +| Budget field | Meaning | +| --- | --- | +| `total_gpu_capacity_bytes` | Positive rank-local GPU capacity. | +| `memory_fraction_kind`, `memory_fraction_value` | `of_total` and the fraction of GPU memory available to the worker. | +| `max_num_tokens`, `max_batch_size` | Positive rank-local scheduler limits within a declared envelope; exact recorded values for runtime memory. | +| `context_length` | Optional positive request bound, at most the profile context and runtime `max_model_len` when present; defaults to the profile context. | +| `request_occupancy_tokens` | Optional positive logical request length within the selected context, for a separate steady decode footprint with one new token per forward pass. | +| `cuda_graph_reserved_bytes` | Separate nonnegative reservation for declared memory; must be zero for runtime memory. | +| `tolerance_fraction` | Optional fraction in `[0, 1)` deducted from the resulting cache byte budget. | + +The result includes `total_kv_size_bytes`, `memory_breakdown`, +`resource_provenance`, `cache_layout`, `cache_groups`, and +`request_peak_cache_bytes`. The latter is a conservative single-request bound +including block alignment and the configured prefill chunk, not a reservation +for every scheduler slot. A smaller context changes this peak bound, not the +declared non-cache resource bounds. `tolerance_adjusted`, when present, provides +the reduced byte budget. For grouped caches, `kv_size_per_token_bytes`, +`total_kv_size_tokens`, and the adjusted token capacity are null (`None` in Python). +There is no equivalent scalar token capacity. + +When `request_occupancy_tokens` is supplied, the result also includes +`request_occupancy_cache_bytes`. KV-relative synthetic loads use this native +footprint while retaining the actual scheduler settings and context for resource +admission. Observed linear pools retain their recorded capacity even when timing +coverage is smaller; uncovered queries report timing errors. + +Runtime memory uses `kv_cache_bytes` directly and returns +`memory_breakdown: null` (`Option::None` in Rust). It does not invent activation +or other non-KV components. The requested memory fraction must exactly match the +recorded utilization; the device budget must contain the observed cache pool. +A different device capacity argument never rescales the pool or qualifies a new +deployment. Changed scheduler settings require new runtime evidence. Explicit +scalar cache-capacity overrides are rejected for runtime profiles in prediction, +recommendation and replay; use the recorded allocation. A tolerance margin may +reduce it through the canonical budget API. + +Python `estimate_kv_cache(..., fpm_profile=..., context_length=...)` delegates +grouped and runtime-memory profiles to this native budget method. +`estimate_num_gpu_blocks` also accepts `context_length` for runtime validation; +it and the +legacy Rust scalar `KvCacheEstimate` transport reject grouped profiles; use the +groups and shared byte budget instead. Existing complete linear declarations +keep their budget calculation and serialized shape. + +Grouped native execution supports cold aggregated vLLM, PP1/CP1, HBM-only cache, +non-speculative decoding and `prefix_caching: false`. Prefix reuse, host/G3 +offload, disaggregation and fixed scalar cache capacity are rejected. Allocation +is atomic across all groups against one shared rank-local byte budget. A window +of `W` keeps the blocks covering the last `W - 1` computed tokens and all tokens +scheduled for the next forward. Expired completed pages are evicted before +subsequent forwards; temporary prefill pages remain charged through completion. +Full-attention groups retain full history. This physical retention never truncates +logical request progress or the context coordinates sent to FPM timing queries. + +The Rust Replay observer's `ReplaySchedulerMetricsSnapshot` reports optional +`kv_cache_used_bytes` and `kv_cache_capacity_bytes` for grouped caches; its +`active_cache_usage` and `physical_cache_usage` use byte occupancy. `active_blocks` +counts resident group pages, while `total_blocks` is zero because heterogeneous +pages have no scalar block capacity. Use the byte fields for capacity comparisons. +These fields are currently exposed through the Rust observer API, not the Python +JSON replay runtime. Linear snapshots retain their existing block metrics. + ## Choosing a forward-pass API Use `RustForwardPassPerfModel.best_available(config)` from Python or @@ -326,10 +443,14 @@ nested paths. The supported namespaces are: architecture without a registered class when a profile is supplied. Without a profile, auto retains SOL. Explicit SOL requires a registered analytical model; direct requires a profile and also supports registered architectures. +- `fpm_interpolation.collect_coverage`: `false` (default). Opt in to bounded + evidence from actual direct-FPM lookups, as described below. It requires + explicit `estimation_mode: fpm_interpolation`, resolved `method: direct`, and + `fallback_policy: deny`. It is omitted from normalized serialization when false. The top-level `fpm_profile` contains the complete profile dictionary: pinned model revision, architecture, context length, expert count, deployment precision -and topology, conservative resource bounds, and provenance. A profile requires +and topology, cache geometry, memory evidence, and provenance. A profile requires an explicit literal `backend_version` that matches its selected deployment; slot aliases and omitted versions are rejected. This profile schema does not declare recorded DCP, so combining `fpm_profile` with an explicit `dcp` is @@ -337,6 +458,17 @@ rejected; measured DCP profiles without `fpm_profile` retain their existing rout Profile/schema and precision conflicts fail before estimator fallback. Omitted precision fields are filled from the profile and preserved in the resolved canonical configuration. +Each deployment may declare `worker_type: prefill`, `decode`, or `aggregated`. +The canonical construction request selects only the matching role's precision, +scheduler envelope, cache geometry and memory evidence. Separate P/D deployments +may share hardware and topology while retaining different resources. Historical +deployments without this field keep their shared-resource behavior and omit the +field on serialization; they cannot coexist with role-specific deployments at +the same hardware/runtime/topology identity. The Python memory adapters accept +`worker_type` to select the same role, defaulting to `aggregated` for existing +callers. An absent or different explicit role is an error, not a fallback to +another role's capacity. + Runtime normalization and construction verify the profile architecture against checkpoint `config.json` from the local model path or pinned remote revision before interpolation selection. Missing or malformed architecture metadata and @@ -353,8 +485,9 @@ whole-forward native operations without SOL operations, and never constructs an analytical graph. Profile resource estimates and memory planning do not require timing data or a native timing model. -Direct readiness requires genuine measurements for each operation's phase; -prefill-only operations do not require decode rows, and vice versa. Construction +Direct readiness requires genuine measurements for the request's `worker_type`: +prefill requires prefill rows, decode requires decode rows, and aggregated +requires both phases. Querying an absent phase still fails explicitly. Construction continues to later systems roots when a required phase is unavailable. Query coverage still needs an exact point or supported interpolation. Cross-KV prefill uses the nearest same-batch lower and upper KV curves that both cover the @@ -398,6 +531,70 @@ return `None`. `tune_with_fpms()` preserves the established FPM observation contract. Native construction still uses Python model compilation; estimator selection, regression, correction, and latency computation are owned by Rust. +### Direct-FPM query coverage + +Enable coverage through the same canonical `best_available(config)` construction: + +```yaml +estimation_mode: fpm_interpolation +fallback_policy: deny +estimator_config: + fpm_interpolation: + method: direct + collect_coverage: true +``` + +The rest of the configuration must identify the exact deployment, including its +complete `fpm_profile`, literal backend version and data roots. CLI prediction +places these timing controls under `engine.workers..timing`. Coverage does +not change point selection, interpolation, correction or latency, and it does +not permit fallback or extrapolation after a missing query. + +The returned model exposes `fpm_query_coverage()` as +`Result, AicError>` in Rust and `dict[str, Any] | None` +in the Python SDK. The raw PyO3 method returns a JSON string, using `null` when +disabled. Evidence belongs to that model instance; cloning a Rust model starts +a fresh accumulator so candidate evaluations do not share counts. The public +Rust facade exports `FpmQueryCoverage`, `FpmQueryCoverageCounts`, `FpmQueryGap` +and `FpmQueryPurpose`. + +A snapshot has cumulative `queries`, `prefill`, `decode` and +`mixed_decode_baseline` counts, each with `measured`, `interpolated` and +`unsupported`. Classification comes from the native lookup that supplied or +rejected the timing, including mixed-pass decode-floor lookups. Counts describe +`native_lookup_resolutions`; they are not request or replay-iteration counts, +and an external timing-cache hit adds no lookup. The snapshot retains up to 128 +distinct gaps with phase, purpose, model/cell identity, query coordinates, reason +and occurrence count. `omitted_gap_queries` counts failed occurrences beyond +that storage limit; unsupported totals are not truncated. Snapshot reads do not +reset evidence, and lookup errors remain errors with partial evidence available. +The snapshot alone does not assess replay completion or prediction accuracy. + +For existing static phase consumers, the same returned model provides these +uncorrected native timings, in milliseconds: + +| Method | Coordinate contract | +| --- | --- | +| `predict_prefill_latency(batch_size, isl, prefix)` | Full input length and cached-prefix length per request; direct-FPM totals are `(batch_size, batch_size * (isl - prefix), batch_size * prefix)`. | +| `predict_decode_latency_total(batch_size, total_past_kv_tokens)` | Exact total past-KV across the decode batch, preserving the native FPM coordinate without a mean-context conversion. | +| `fpm_decode_kv_ceiling()` | Largest collected decode KV total, or `None`; reading this bound does not record a timing lookup or prove shape-specific coverage. | + +These methods preserve the existing engine's phase behavior and record actual +direct lookup evidence when enabled. Empty work does not invent a measured +query. Whole-iteration `estimate_forward_pass_time_ms()` remains the API for +scheduled per-rank FPM telemetry, including mixed work and its existing online +correction behavior. + +Ordinary replay uses this returned model for FPM timing and emits +`fpm_query_coverage` in its report when collection is enabled. The CLI also saves +`fpm-coverage.json`, including partial evidence when a native timing error stops +replay; the Python exception retains a `fpm_query_coverage` attribute for that +failure path. A passing replay coverage status requires completed requests, +nonempty resolved queries and no unsupported lookup. It is distinct from +operation/energy evidence, which whole-model FPM does not provide. See +[FPM replay validation](fpm-self-service.md#validate-fpm-query-coverage-with-agentx-replay) +for the stricter whole-corpus completion checks and saved onboarding artifacts. + ### Recursive regression and statistics rebuilding Regression maintains centered sufficient statistics for the retained samples @@ -600,6 +797,9 @@ The supported `aisimulate_core::perfmodel` Rust surface is grouped as follows: - KV-cache estimation: `estimate_kv_cache`, `KvCacheEstimateRequest`, `KvCacheEstimateOptions`, `KvCacheMemoryFraction`, and estimate/result/error types; +- FPM profile cache resources: `FpmCacheGroup`, `FpmCacheKind`, `FpmCacheLayout`, + `FpmResourceConfig`, `FpmRuntimeMemoryConfig`, `FpmCacheBudgetRequest`, `FpmCacheBudget`, and + `FpmCacheBudgetAdjusted`, through the canonical model/config budget method; - wire identity: `EngineConfig`, `ParallelMapping`, `QuantizationConfig`, `SpeculativeConfig`, `BackendKind`, `DatabaseMode`, and `DataType`; - schema gates: `ENGINE_CONFIG_SCHEMA_VERSION`, diff --git a/docs/fpm-self-service.md b/docs/fpm-self-service.md index d3681ea0e..ccd16e754 100644 --- a/docs/fpm-self-service.md +++ b/docs/fpm-self-service.md @@ -5,13 +5,316 @@ SPDX-License-Identifier: Apache-2.0 # FPM self-service -`aisimulate onboard` guides onboarding a new model for FPM simulation on your designated hardware platform. It records the model, runtime, target GPU system and allocation, plans one TP, DEP, or TEP worker, and produces ordinary `predict` and `recommend` configurations that read your collected FPM data. It builds on AISimulate's per-worker FPM support and packaged collector. +`aisimulate onboard` guides onboarding a new model for FPM simulation on your designated hardware platform. It records the model, runtime, target GPU system and interconnect, and lets you select one or more TP, DEP, or TEP worker configurations. Each configuration has its own resource profile, collection plan and minimum collection GPU requirement. You review the resource and collection limits, collect whole-forward timings through Dynamo self-benchmark, then validate their query coverage using ordinary trace replay. Each plan also produces ordinary `predict` and `recommend` configurations for its worker. -Planning works before the model has an AISimulate model class or measured FPM timings. With a supplied FPM profile, planning validates the declared deployment identity and estimates memory admission from your resource bounds. Runtime compatibility and data readiness remain **unchecked**, and accuracy is **not assessed**. This setup does not provision GPUs or run target preflight checks. +Collection limits and validation traffic are separate inputs. AISimulate sets runtime limits and a prefill CUDA graph policy. New requests leave capture configuration to the pinned runtime; an explicit capture extension is available when needed. Dynamo self-benchmark uses the initialized engine, image sampling defaults and runtime feasibility checks to generate and measure the exact grid. AgentX traces exercise the resulting FPM library through replay; their variable request lengths do not require a fixed input/output length or latency target during onboarding. -Ordinary FPM `predict` and `recommend` accept an inline `engine.fpm_profile` with model identity and rank-local resource bounds. Direct interpolation uses measured timings without constructing an op-level model. Automatic interpolation selection retains SOL for registered architectures. See [Choose the model execution route](#choose-the-model-execution-route) for selection and coverage rules. +Planning works before the model has an AISimulate model class or measured FPM timings. With a supplied FPM profile, planning validates the declared deployment identity. Fresh config-derived profiles leave memory pending until a runtime probe or runtime initialization during collection; no activation or non-KV memory bound is required. Runtime compatibility and data readiness remain **unchecked**, and accuracy is **not assessed**. This setup does not provision GPUs or run target preflight checks. -The Inkling pilot targets NVIDIA GB200 with TP, DEP, and TEP through the class-independent route. Its checkpoint, runtime, allocation, profile resource bounds, and strategy degrees must be pinned before collection. This guide does not establish Inkling readiness or GB200 accuracy. +Ordinary FPM `predict` and `recommend` accept an inline `engine.fpm_profile` with model identity, cache geometry and resolved rank-local memory evidence. Direct interpolation uses measured timings without constructing an op-level model. Automatic interpolation selection retains SOL for registered architectures. See [Choose the model execution route](#choose-the-model-execution-route) for selection and coverage rules. + +## Onboard with an agent + +For a request such as "Help me onboard my model for FPM simulation on my target GPUs," follow these six stages. Claude Code reaches them through the root `CLAUDE.md` import of `AGENTS.md`; other agents use the same `AGENTS.md` entry point. Create one [session checkpoint](#checkpoint-and-resume-an-onboarding-session) during stage 1 and update it after meaningful findings and decisions, including unfinished investigation and partial profile review. The sections below remain the detailed CLI reference. + +| Stage | Required result | +| --- | --- | +| 1. Inspect the model and target | Accessible config/profile and packaged hardware specification identified; supported metadata read and gaps recorded. | +| 2. Choose the worker and collection limits | Sourced joint precision options and collector limitations reviewed; checkpoint identity/revision, runtime, interconnect, precision combination(s), worker configurations, memory fraction, context, scheduler bounds and capture policy selected. Minimum collection GPUs are derived for each; actual target-runtime compatibility remains unchecked. | +| 3. Derive, review and save collection inputs | Each configuration's identity, precision, topology, cache geometry and runtime settings reviewed and accepted. Memory may remain pending; planning estimates are identified separately. | +| 4. Plan collection | Validated saved plan and collector preview for each configuration; execution resources, independent scheduling, prerequisites and recovery paths recorded. | +| 5. Collect, verify and finalize | Existing evidence inspected and bounded prefill/decode readiness established before full collection. Formal timings and initialized-worker memory evidence verified. `onboard finalize` writes a fresh simulation plan; review and accept its resolved profile separately. | +| 6. Validate replay and run predict/recommend | Cold aggregated trace replay checks direct-FPM query coverage; completed and incomplete results are distinguished. Ordinary prediction/recommendation examples remain available; accuracy is assessed separately. | + +At each transition or blocker, give a short update: **Stage N/6 — name; result or blocker; next action.** Ask only for missing information needed for the current stage, using facts and decisions already supplied. Continue independent authorized work while an answer is pending. Move on when the required evidence is available; stage transitions do not require another approval. Honor prior authorization, including authorization to execute collection, while preserving the profile acceptance step below. + +### 1. Inspect the model and target + +Start with only the missing **Hugging Face model ID** (`organization/model-name`) and **target GPU platform**. If neither is known, a first response can be: + +> Stage 1/6 — Inspect the model and target. What is your Hugging Face model ID (`organization/model-name`) and target GPU platform? I'll retrieve and inspect the configuration first. If your model isn't on Hugging Face, share a local config or checkpoint path instead. + +Omit facts already supplied from this request. Accept a supplied local config, complete FPM profile or checkpoint path without requiring a Hub ID; do not ask for both an ID and a config upfront. With a Hugging Face ID, use available authorized Hub access to retrieve only `config.json` at this stage. Honor any supplied revision and record the repository ID and resolved commit. If no revision was supplied, report the resolved commit for the deployment choice in stage 2. If the config cannot be retrieved, explain the specific gap and ask for a local config. The agent retrieves the file; AISimulate's CLI consumes a local config/profile and does not download it or the checkpoint. + +Inspect the accessible config/profile before asking for model kind, architecture, expert count or context limit; derive supported metadata and record its source. Ask about these fields only when the configuration leaves them unresolved. For CLI setup, `--model` identifies the selected Hugging Face repository or actual checkpoint path, `--model-config` names the local JSON file, and `--model-revision` pins the checkpoint separately. A config file hash is not a checkpoint revision. This route does not require an op-level model class or per-operation silicon data. If the hardware lacks a packaged system specification, report that integration gap. + +A natively multimodal checkpoint can be onboarded for its text decoder. Explain this scope during inspection and profile review: FPM timing models text-decoder execution and excludes encoder, projector and preprocessing latency. Config-derived estimates describe the decoder only. Observed cache capacity reflects every component actually loaded by the collection worker, including resident non-text components; do not subtract guessed encoder memory from that observation. This does not qualify full multimodal workloads. + +The agent handles checkout and environment checks: record the branch/commit, follow [development setup](../DEVELOPMENT.md#initial-setup), activate the environment, and inspect `aisimulate onboard --help` and `aisimulate onboard init --help`. Check later subcommands before using them. Report missing commands/options as a version mismatch, rather than asking the user to supply unsupported inputs. + +### 2. Choose the worker and collection limits + +Establish the intended serving role before reviewing runtime settings. For P/D disaggregation, create independently accepted `--worker-type prefill` and `--worker-type decode` configurations; their topology, graph settings, scheduler limits and memory budget may differ. For aggregated serving, use `--worker-type aggregated`: both probe launches receive the same accepted serving configuration, while actual per-batch graph execution may differ. Inspect runtime support before proposing modes; do not assume every model/backend supports prefill FULL graphs. See [serving roles](#serving-roles-and-cuda-graph-settings) for the CLI and multi-configuration format. Omitted roles in existing requests retain their historical shared-profile behavior and are not relabeled as independently measured P/D profiles. + +At stage 2 entry, reuse stage 1's checkpoint revision and target hardware, and resolve the literal framework version from supplied or available deployment metadata. Ask for the missing version pin only if it remains unknown, then investigate that version before proposing precision choices. Resolve any remaining checkpoint identity or interconnect gaps from available metadata before asking. A model label, config hash or example revision does not establish a checkpoint pin. An existing vLLM launch command or configuration is optional evidence: preserve explicit choices, but do not require one or begin with a questionnaire about every topology and dtype. Do not ask for total available GPUs, node allocation, GPUs per node or replica budgets during onboarding; those are choices for actual prediction or recommendation runs. + +Complete the [runtime investigation and options table](#investigate-runtime-constraints-and-precision-options) before asking for derivable precision or cache inputs. New config-based setup proposes `comm_quant_mode: half` with the current collector's identity source; this is not a claim that all NCCL tensors use FP16. The agent investigates and prepares supported inputs; the CLI does not import arbitrary launch arguments or automatically inspect a remote runtime. + +Show the [runtime and collection defaults](#runtime-and-collection-limits): 0.90 GPU memory utilization, profile scheduler bounds or 8,192 scheduled tokens and 256 sequences per rank, and `--prefill-cudagraph-policy runtime`. Fresh config/profile-based setup caps context at the smaller of the declared limit and the 256,000-token AgentX reference bound. Runtime graph policy leaves the effective mode/capture list and exact sample grid unresolved until engine initialization. Explain these starting policies and allow edits; they do not establish memory fit or timing coverage. Do not ask for fixed input/output lengths, concurrency, TTFT or TPOT as required collection inputs. Those flags customize optional synthetic validation examples. + +With a local model config, inspect the [read-only topology preview](#preview-and-choose-parallelism) for each considered precision combination using the actual identity, runtime, target and collection limits. Supply the precision/layout facts established by inspection, investigate remaining diagnostics, ask only for unresolved choices, and rerun the preview. Present the default and alternatives with their resource assumptions and unresolved fields. A candidate marked `estimated_fit` has a complete optional planning estimate; it is not a performance ranking, runtime qualification or a measurement. If no default is available, explain why and select an exact candidate; missing non-KV estimates do not prevent collection. Do not silently assume topology or precision. + +Help choose one or more TP configurations, or the relevant TP/DEP/TEP configurations for MoE, using the preview's exact flags or the [supported topology fields](#create-the-request). Honor explicit choices even when they lie outside the automatic shortlist. Explain each selected tuple; do not require the user to know every parallelism field upfront. Derive its minimum collection GPUs as attention TP times attention DP: TP4 requires four GPUs, while DEP8 requires eight. This requirement does not declare available capacity or establish runtime placement or compatibility. Use [directory output](#onboard-multiple-parallel-configurations) to prepare several configurations in one session; each profile and collection plan still selects one exact tuple. Collection runs can reuse the same GPUs, so do not add their requirements into an onboarding GPU budget. + +#### Investigate runtime constraints and precision options + +Perform this investigation proactively once the checkpoint, hardware and framework version are known. A local source checkout, user-supplied launch command, initialization logs and a prestarted server are optional evidence, not prerequisites. Use available local sources first and authorized official remote sources when needed. Do not download weight shards or launch GPU work merely to discover options. + +1. Inspect the pinned checkpoint's config, quantization sidecars and lightweight tensor/index metadata where available. Identify quantized and unquantized weight parts and any fixed cache/state constraints; a repository name or one weight-quantization label does not describe every tensor. Distinguish metadata-derived estimates from measured allocations. +2. Read the pinned official framework's model implementation, quantization/backend selection and attention/cache implementation. Trace model dtype validation, backend/kernel dispatch, cache dtype/block sizing and convolution or other state allocation, including version, GPU and topology conditions. Record the relevant flags, environment variables and defaults. For a supplied vendor/custom build, inspect its pinned image/build metadata and known patches; record uncertainty if they are unavailable. Unpatched release behavior alone does not establish that build's limits. General dtype enums or another version's documentation do not establish support for this model. If the pinned build lacks the model or a required backend, report that incompatibility explicitly. +3. Classify each finding as a **fixed checkpoint constraint**, **runtime default**, **configurable alternative**, or **unknown**. Keep model dtype, mixed weight storage, FMHA FPM identity and actual runtime compute/dispatch, attention KV dtype, convolution state and communication identity separate. Trace their dependencies to form plausible joint combinations; do not create a Cartesian product of independently advertised dtypes. Keep unknown FMHA/KV facts unresolved rather than inferring them from weight precision. Record missing evidence and access limits before asking only for facts or choices that investigation cannot resolve. + +Present a concise **options table with one row per joint combination**, including the default and evidence-backed alternatives. Do not invent alternatives to fill the table. Use these columns: + +| Column | Required content | +| --- | --- | +| Option and joint precision | Distinct weight parts/model dtype, FMHA identity and runtime compute/dispatch, attention KV and fixed state; label fixed, default, configurable or unknown facts. | +| Conditions and controls | Pinned version, GPU, backend and topology dependencies; required runtime flags/environment settings and defaults. | +| Sources and confidence | Checkpoint revision and metadata path; official source URL/path and tag/commit for each claim. Distinguish source-confirmed behavior, inference and unresolved or untested conditions. | +| Runtime support | Supported by the inspected version, conditional/unverified, or ruled out, with the reason. Source review does not establish successful target execution. | +| Collector compatibility | Compatible with the current profile collector's identity checks, blocked, or unresolved, independently of runtime support. | +| Resource effects and limits | Expected changes to weights, attention cache, fixed state, workspaces/padding and memory fit; retain unknown bounds. Accuracy remains unmeasured without matched evidence. | + +The current profile collector requires KV dtype to equal its checkpoint-inferred native dtype and checks GEMM, MoE, FMHA and communication identities against its resolved configuration. Those checks restrict collection even when the pinned runtime supports another combination. Mark any blocked alternative with its precise limitation; user selection or a schema-valid profile does not override serving dispatch. Do not relabel checkpoint data or use extra launch arguments to disguise an unsupported identity. See [profile collection constraints](#plan-preview-and-explicitly-execute). + +Treat runtime defaults as one proposal and let the user choose the precision combination(s) to cover, preserving prior explicit choices. Record the choice, rationale, sources and unresolved conditions in profile provenance. A blocked selection stays blocked until the required capability is available; do not substitute a default silently. For multiple selected combinations, use a separate fresh output root per precision combination and the existing `--parallel-configs` topology list within each. Duplicate resolved topologies within the same serving role are rejected even if their precision overrides differ. This is agent orchestration over the existing CLI, not a new precision-matrix command or schema. Recompute resources and review each exact precision/topology profile in stage 3 before accepting it. + +### 3. Derive, review and save the profile + +Use [a local model config](#start-from-a-local-model-config), or [a supplied profile](#provide-identity-and-resource-metadata), for the selected deployment. Both produce a class-independent direct-FPM plan. Derive supported cache geometry and runtime facts before asking for unresolved inputs. Explain each value's source and limitations. Do not ask users to guess activation, runtime, communication or aggregate non-KV bytes. Fresh config-derived profiles keep memory pending; optional estimates remain planning evidence until collection observes the actual cache allocation. Resolve effective weight, FMHA, communication and KV precision separately; a quantized checkpoint label does not determine all of them. Preserve replacements and their rationale in overrides/provenance. Review all effective values and assumptions using the appropriate flow below. + +For several configurations, read shared model, runtime and hardware inputs once, then derive and review each selected precision/topology profile independently. Do not transfer rank-local byte bounds or `cache_groups` between precision combinations or tuples. Put headless overrides under the corresponding `--parallel-configs` entry's `resource_overrides`; interactive edits apply only to the profile being reviewed. Shared `cache_block_sizes` can be reused only when valid for the same runtime, backend and precision, with page bytes derived for each tuple. + +If runtime-dependent geometry or memory remains unknown, follow [runtime probing](#resolve-cache-geometry-with-a-runtime-probe) before requiring a complete profile. Use matching verified evidence or execute a bounded probe for every selected configuration and each phase required by its role: prefill or decode for a single-role request, both for explicit aggregated and role-less historical requests. A missing bundled adapter triggers source investigation and campaign-local instrumentation; it does not end the investigation. + +For sliding-window or supported convolution state, follow [grouped cache review](#review-grouped-cache-resources). Derive layer geometry first, then resolve runtime block sizes for each chosen worker and review aggregate page bytes, including padding. A scalar bytes-per-token estimate cannot replace these groups. + +For agent sessions, use the headless flow below and [checkpoint each configuration's draft and acceptance](#checkpoint-and-resume-an-onboarding-session), even when a terminal is available. This preserves partial review across sessions. The standalone interactive flow remains available, but does not automatically save its prompts or partial acceptance to the session checkpoint. + +| Workflow | Review and save behavior | +| --- | --- | +| Checkpointed agent session, with or without a PTY | Supply identity/collection flags and overrides without `--interactive`; use `--parallel-configs` with `--output-dir` for several configurations of one selected precision combination. Missing required inputs exit 2 without saving; use the diagnostics to investigate first and ask only for still-unresolved facts or choices. A successful command writes immediately. Initially write to a separate path such as `draft-request.yaml` or a fresh `draft-onboarding/` directory and show every embedded profile, its sources and scope for user review. Save each configuration's draft, edits and explicit acceptance in the session checkpoint as review progresses. Apply edits in the inputs/overrides and repeat draft review; rerun unchanged accepted inputs to a fresh final path and verify that the profiles match the accepted drafts before planning. Draft names are only a convention; they have no special CLI status. Review a supplied `--fpm-profile` in the same way. | +| Standalone interactive CLI in a terminal | Run `aisimulate onboard init --model-config /path/to/config.json --interactive --output-dir onboarding`, or retain `--output support-request.yaml` for one request. Directory output accepts comma-separated candidate numbers. Inspect missing-field prompts and review the final profile, using `edit` to revise its values before entering `accept`. All selected profiles must be accepted before directory output is saved. Cancellation creates no new artifacts and preserves prior output; partial prompts or acceptance are not persisted in the session checkpoint. This final review is specific to `--model-config --interactive`; supplying `--fpm-profile` does not add it. | + +Do not pipe answers into `--interactive`: it requires a terminal. Scripted setup has no built-in acceptance prompt. Keep drafts separate from final requests and use a fresh draft directory for each revision. Preserve prior outputs when revising a deployment. + +Model identity, runtime and topology are not profile-review edit fields. If they change, return to stage 2, regenerate dependent estimates and review the new request; use new output paths for the changed deployment. A precision change also requires reconsidering the joint combination and collector compatibility, recomputing dependent resources and reviewing the complete revised profile. Explicit byte overrides can survive CLI edits; rederive any that depend on the changed precision rather than retaining stale bounds. + +### 4. Plan collection + +Follow [Plan, preview, and explicitly execute](#plan-preview-and-explicitly-execute) using a new output directory. For directory output, read `onboarding.json` and use each emitted `aisimulate onboard plan` command to create that configuration's `collection/` directory. Inspect each `support-plan.json`, embedded/saved profile, generated prediction/recommendation configs and `commands.json`. Run `onboard collect-fpm` without `--execute` to print the collector command; this does not run the collector's own plan or check the target runtime. Inspect the read-only collector plan when its input environment is available and identify any missing prerequisites. + +Explain each configuration's minimum collection GPUs, rank-local scheduler/resource envelope and [shared collection policy](#how-the-collection-grid-is-determined). AISimulate configures collection inputs and Dynamo generates the exact points; the plan does not derive a second grid from trace requests. In runtime graph mode, report the capture list and exact point count as unresolved until engine initialization. Each plan selects one parallel tuple and generates single-worker validation configs. Continue stages 5 and 6 separately for each configuration, preserving its data and result paths. Actual collection resources and placement must be checked in the collector environment before execution. The synthetic request count does not bound timing samples or collection duration. + +Prepare the per-configuration execution table in [Orchestrate independent collection campaigns](#orchestrate-independent-collection-campaigns), including the reason for any job dependency. Sharing GPU capacity does not make one configuration's data a prerequisite for another. + +### 5. Collect and verify data + +First inspect any existing timing data for a matching deployment. Reuse a verified matching formal pair when available; do not collect again merely to complete a stage or because historical native campaign artifacts are unavailable. For new or continuing collection, [check timing readiness](#check-timing-readiness-before-full-collection) against saved native artifacts for each configuration. Where evidence is missing, prepare a bounded smoke for each phase required by the selected role before a full grid; explicit aggregated and role-less historical requests require both prefill and decode. A completed runtime memory probe does not replace this timing check. + +For new collection, prepare the compatible pinned Dynamo/vLLM image, accessible checkpoint and GPU resources. Choose the existing collection executor at this stage: Kubernetes needs its namespace and deployment permissions; Slurm needs a caller-owned `sbatch`/`salloc` allocation with Pyxis and shared storage. Follow the [deployment options](#choose-the-collection-executor) below and the [FPM collection guide](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md). An already-running HTTP server is not required: the AISimulate collector generates the runtime and launches benchmark workers through the selected executor. The engine initializes the model and cache, resolves its graph configuration, and Dynamo self-benchmark generates and times the admitted points. Use the agreed scope and existing execution authorization, reporting concrete missing prerequisites when blocked. Preview and execute with the same deployment options. `--execute --smoke` collects diagnostic evidence without a formal FPM pair; full execution requires matching usable timing evidence for every selected cell. Inspect effective precision, graph mode/capture sizes, cache allocation/padding and supported benchmark seeding during bring-up; preserve checkpoints, initialization logs and raw evidence. + +Follow the [orchestration and recovery procedure](#orchestrate-independent-collection-campaigns) for every selected configuration. Track worker completion separately from publication and validation, and continue unrelated campaigns when one fails. + +For both reused and new data, [inspect the published pair](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md#4-validate-and-install-the-fpm-profile): verify hashes, schema and identities, including actual runtime, topology, precision and available prefill/decode cells. Keep it at the generated configs' local systems path. Record any historical checkpoint-revision uncertainty in provenance; a declared revision does not prove that old measurements used it. A successful preview or smoke run is not formal data, and a matching pair does not prove all simulated queries are covered. + +Run the standard [collection quality checks](#validate-collection-and-serving-accuracy) against the original collection directory: point validity and effective execution inspection, comparable full-grid repeatability, and withheld-coordinate interpolation. Review and save the editable policy before measurements. A published table is usable evidence even when a quality gate is incomplete or failed, but it is not qualified accuracy evidence. + +Then, for a pending profile, [finalize runtime memory](#finalize-runtime-memory) into a fresh directory. Supply the collection report through optional `--collection-report` to preserve its independently verified quality status in the resolved profile; omission records `not_assessed`. Review and accept the resolved profile separately, save its provenance in the session checkpoint, then use that directory in stage 6. Matching timing data alone cannot finalize memory when initialization evidence is missing. Exploratory prediction/recommendation remains available with resolved memory while accuracy is unqualified. + +### 6. Validate replay and run predict/recommend + +For heterogeneous P/D recommendation, candidate enumeration uses each role’s +hardware and explicit `prefill` or `decode` worker identity. + +For an independent prefill or decode request, stop the collection workflow at that role's verified quality and memory results. Its plan exports `worker.yaml`, the role-tagged FPM profile and an isolated systems directory; it does not create an executable full P/D prediction or recommendation. The other role must be accepted and qualified independently before composing an ordinary serving configuration. `onboard validate-fpm` currently handles aggregated replay only. Existing grouped-cache P/D handoff limitations concern subsequent replay, not collection: decode collection initializes representative state locally and does not require an actual P/D transfer. Keep this replay limitation in the checkpoint without blocking collection qualification. + +[Validate a local AgentX trace](#validate-fpm-query-coverage-with-agentx-replay) through `onboard validate-fpm`. It writes an inspectable ordinary prediction config, runs cold aggregated replay, and saves coverage evidence separately from the collection plan. Use the selected complete local corpus or one complete play for a first check, and state that selection. Keep `engine.systems_paths`, the target model projection and the direct-FPM profile intact. Report the command's exit status and exact missing coordinates. Missing timing stops replay; a partial report does not audit the rest of the trace. + +[Run the generated ordinary configurations](#run-the-generated-ordinary-configurations) when a synthetic prediction or recommendation check is useful. For deployment prediction or optimization, set the desired replicas and GPU budget in the ordinary runtime configs. Do not change precision labels or silently switch timing methods to obtain a result. Return to stage 5 to address missing data, or stage 2 if the selected deployment or collection limits change. A different validation trace alone does not invalidate collected timings. Report the simulation stage as incomplete while required queries fail. + +Complete the standard [matched serving validation](#validate-collection-and-serving-accuracy) for one representative compatible play. At handoff, provide the session checkpoint path and ensure it records the checkout revision, final request/profile, plan directory, policy, data pair/provenance, exact commands/exit statuses and all result paths. Distinguish estimated memory fit and CPU planning from actual target-runtime checks, formal data/coverage, completed simulations and measured accuracy. Missing or failed validation keeps accuracy unqualified; ordinary exploratory prediction and recommendation remain available. + +### Checkpoint and resume an onboarding session + +Use one `onboarding-checkpoint.json` for the whole session, including every selected precision/topology configuration. The agent creates it during stage 1, before all inputs are known, and saves after meaningful findings, user decisions, draft edits, acceptance and command results. Users do not need to request a save. These commands persist the supplied state; they do not observe conversations, investigate remote capabilities or automatically intercept other onboarding commands. The agent is responsible for invoking them. + +Keep the file outside fresh `init --output-dir` roots, for example beside `drafts/` and separate final precision directories under `model-onboarding/`. Those directories must still satisfy the existing fresh-output rules. Create the checkpoint with: + +```bash +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json +``` + +Calling the same command again inspects existing state without replacing it. Each saved revision has a generated revision number. Updates require `--expect-revision` with the number just read; a stale revision rejects the update. Reload and reconcile concurrent changes instead of retrying with an invented revision. Saves use an OS lock and atomic file replacement. The adjacent lock file is an implementation detail, not another user checkpoint; leave it in place. + +The JSON document uses `schema_version: "aisimulate-onboarding-checkpoint/v1"`. Its [schema](../python/aisimulate/src/aisimulate/support/checkpoint.py) holds: + +| Field | Purpose | +| --- | --- | +| `inputs` | Shared effective model, runtime, hardware and collection choices, including immutable revisions. Incomplete inputs are allowed. | +| `research`, `decisions`, `pending_questions` | Findings with sources/confidence, user rationale and unresolved questions. These are context, not executable instructions or acceptance. | +| `validation_inputs` | Trace selection and other validation-only choices, separate from collection identity. | +| `progress` | Agent-reported stage (`1`–`6`), status (`in_progress`, `blocked` or `complete`), blockers and next action. A saved status does not prove successful execution. | +| `configurations` | Records keyed by a stable user-chosen ID. Each has its own inputs, draft request, validation inputs, progress and artifact references; acceptance is generated by the CLI. | +| `artifacts` | Shared source-file references. Configuration records hold their own profiles, plans, data, results and collector-checkpoint references. | + +Both commands return a JSON report containing the full saved `state`, per-configuration `profile_accepted` and `effective_status`, `integrity_issues`, `archived_artifacts`, verification limits and `saved`. Read the current revision from `state.revision`. `effective_status` is `draft`, `accepted` or `needs_attention`; the separately reported stage/status is agent-supplied, not a verified completion result. An update can save successfully while reporting existing artifact integrity issues and exiting 2; check `saved` and the new revision before retrying. + +`--update FILE.json` applies a recursive JSON merge patch; `--update -` reads the patch from stdin. Objects merge, arrays replace and `null` removes a field where the schema permits it. Do not edit CLI-owned schema, revision, acceptance or hash fields. Put effective selections in `inputs` or the configuration's `draft_request`, not only in research or decision notes: dependency checks use the former. For example, after initial creation, save supplied facts and unfinished investigation with the current revision: + +```bash +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json \ + --expect-revision 1 --update - <<'JSON' +{ + "inputs": {"model": "organization/model-name", "gpu": "h200_sxm"}, + "research": {"runtime": {"status": "pending", "sources": []}}, + "pending_questions": ["Which pinned vLLM version will collection use?"], + "progress": {"stage": 1, "status": "in_progress", "next_action": "Retrieve and inspect the model config."} +} +JSON +``` + +Replace these illustrative inputs with the user's facts. The example revision assumes the preceding creation was the only save; always use the actual current revision. + +#### Preserve partial profile review + +Use [headless draft generation](#3-derive-review-and-save-the-profile) for checkpointed agent sessions. Save each complete generated request, including its embedded profile and provenance, as that configuration's `draft_request`. Incomplete draft requests may also be saved while resolving inputs, but cannot be accepted. For example, after generating two draft requests at `model-onboarding/draft-a.yaml` and `model-onboarding/draft-b.yaml`, prepare their update without changing the requests: + +```bash +python - <<'PY' +import json +from pathlib import Path +import yaml + +root = Path("model-onboarding") +update = {"configurations": {}} +for name in ("a", "b"): + update["configurations"][f"worker-{name}"] = { + "draft_request": yaml.safe_load((root / f"draft-{name}.yaml").read_text()), + "progress": {"stage": 3, "status": "in_progress", "next_action": "Review this exact draft request and its assumptions."}, + } +(root / "draft-update.json").write_text(json.dumps(update, indent=2) + "\n") +PY +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json \ + --expect-revision 2 --update ./model-onboarding/draft-update.json +``` + +The update file is disposable input; the checkpoint contains the drafts. Show each complete request/profile, source and assumption to the user. Save requested edits before presenting the revised values. Only after the user explicitly accepts `worker-a`'s exact values, record that acceptance: + +```bash +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json \ + --expect-revision 3 --accept-profile worker-a --update - <<'JSON' +{ + "configurations": { + "worker-a": { + "progress": {"next_action": "Publish the accepted draft to a fresh final path, verify equality, then plan collection."} + } + } +} +JSON +``` + +The CLI validates the complete request/profile and binds acceptance to its content and relevant shared/configuration inputs. `worker-b` remains a draft. Repeat `--accept-profile` to record multiple explicitly approved configurations together; never use it to infer approval from a saved file, stage label or another configuration's acceptance. Saving or resuming does not publish final requests or launch collection. Regenerate each accepted request to a fresh final path using unchanged reviewed inputs and verify equality before planning; checkpoint acceptance does not replace that comparison. Standalone `init --interactive` remains all-or-nothing and does not save partial prompt/review state automatically. + +#### Record artifacts and resume + +Artifacts are objects keyed by a descriptive name. Register only existing files, with paths relative to the checkpoint directory where possible. Each reference declares `kind` (`file`, `request`, `profile` or `collector_checkpoint`) and `scope` (`input`, `collection` or `validation`). Shared artifacts require `scope: input`; collection and validation artifacts belong to a configuration. For example, the following configuration fragment registers a final request after publication: + +```json +{ + "configurations": { + "worker-a": { + "artifacts": { + "final_request": {"path": "final-a/request.yaml", "kind": "request", "scope": "collection"} + } + } + } +} +``` + +Immutable files receive a SHA-256 snapshot when registered. `kind: request` and `kind: profile` also check canonical contents against the configuration's draft. Later unrelated saves, including resubmitting the same path, do not silently refresh a snapshot. References are current by default (`archived: false`). Keep old output files for inspection, explicitly set superseded references to `archived: true`, and register current outputs under new names/paths. Archives retain their original path, kind, scope and hashes. The report lists them under `archived_artifacts` as historical and unverified: they do not undergo current integrity or semantic checks, and missing archived files do not block resume. Current outputs still require the usual verification; do not archive unresolved current evidence merely to suppress a failure. + +For example, after accepting worker-a's revised draft and publishing its matching request and plan into `final-a-v2/`, archive the old request/plan/collector references while registering the replacements. If the checkpoint now reports revision 7, use the following patch; substitute the actual revision and reference names: + +```bash +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json \ + --expect-revision 7 --update - <<'JSON' +{ + "research": {"worker-a-revision": "Revised collection limits; earlier outputs retained for inspection."}, + "configurations": { + "worker-a": { + "artifacts": { + "final_request": {"archived": true}, + "plan": {"archived": true}, + "collector": {"archived": true}, + "final_request_v2": {"path": "final-a-v2/request.yaml", "kind": "request", "scope": "collection"}, + "plan_v2": {"path": "final-a-v2/collection/support-plan.json", "scope": "collection"} + } + } + } +} +JSON +``` + +Retiring an `input` reference removes it from the dependency context and invalidates affected acceptance and downstream progress; review and accept the revised inputs explicitly. Setting `archived: false` restores checks against the original snapshot/context, so a stale, modified or missing reference reports an issue again. Deliberately replacing a snapshot at the same path requires removing its reference in one save and registering it again after review; prefer new paths and archived references to preserve history. + +Record the existing collector checkpoint as `kind: collector_checkpoint` with `scope: collection`: its mutable cell progress stays in the collector's own file and is not duplicated or treated as immutable drift. The session checks that this file exists and is a readable JSON object; the collector verifies its schema and frozen-plan identity on actual collection resume. + +Resume in the same or another agent session using the single entry point: + +```bash +aisimulate onboard resume \ + --checkpoint ./model-onboarding/onboarding-checkpoint.json +``` + +Use the checkpoint's absolute path when resuming from another working directory. Its relative artifact paths resolve from the checkpoint directory. Resume is read-only: it returns saved context, effective per-configuration progress, integrity issues and next actions without accepting a profile, executing saved command strings or starting collection. Missing or modified immutable artifacts produce actionable issues and a nonzero exit status while retaining readable saved context. An unreadable, malformed or unsupported-schema checkpoint fails with a concise error. + +Relevant shared inputs affect every configuration; configuration input/draft changes affect only that configuration. Such changes invalidate affected acceptance and downstream progress. Research/rationale-only edits preserve acceptance. Changes under `validation_inputs` invalidate validation progress without discarding valid collection data. Archive superseded references explicitly when replacing stale outputs; leaving them current continues to block readiness even after replacements are registered. Do not reuse stale outputs as current or delete completed collection to make a status pass. + +Resume verifies checkpoint integrity and current referenced file snapshots; archived references are historical and unverified. It does not establish runtime compatibility, inspect every plan/data semantic or certify collection, coverage or accuracy from saved stage strings. Inspect the corresponding command results and ordinary plan/data/validation evidence before continuing. Existing collector frozen-plan identity checks remain authoritative when actual collection resumes through `onboard collect-fpm --resume`. + +#### Preserve collection deployment options + +At the collection stage, save the exact executor, pinned image, mounts, networking and other deployment options in the existing configuration's `inputs.collection_deployment`. Record them before registering collection artifacts so the checkpoint binds those artifacts to the selected deployment. For example, if the current checkpoint revision is 8: + +```bash +aisimulate onboard checkpoint \ + --file ./model-onboarding/onboarding-checkpoint.json \ + --expect-revision 8 --update - <<'JSON' +{ + "configurations": { + "worker-a": { + "inputs": { + "collection_deployment": { + "executor": "slurm", + "image": "/shared/images/pinned-runtime.sqsh", + "container_mount": ["/shared/models:/models:ro", "/shared/hf-cache:/root/.cache/huggingface"], + "cpus_per_task": 16, + "cpu_bind": "cores", + "transport": "ib" + } + }, + "progress": {"next_action": "Inspect the allocation, preview with the saved deployment options, then resume the matching campaign."} + } + } +} +JSON +``` + +Use the actual revision, configuration ID and deployment values. This object is saved agent context, using the existing generic `inputs` field; it is not a second collector configuration format. On `onboard resume`, recover these values and pass the corresponding `collect-fpm` flags explicitly, including every repeated `--container-mount`. Preserve any supplied `dynamo_version` and Kubernetes options in the same way. Save the reviewed command and its results as session context. Neither `resume` nor the generated generic commands automatically applies this object or executes a saved command string. + +As with other effective input changes, changing these values invalidates that configuration's acceptance and marks its existing collection references stale. Review the affected context, retain the old campaign, archive superseded references with a reason, and plan into a new directory. Record replacement artifacts and use the existing exact-profile acceptance workflow; do not rewrite or discard the collector's checkpoint to force a changed deployment to resume. An interrupted unchanged campaign uses its original deployment options and `collect-fpm --resume`. Acquiring a replacement Slurm allocation does not by itself change the saved deployment options; the collector still verifies the current allocation and frozen plan. + +### Resume from existing work + +If a session checkpoint exists, run `onboard resume` first, resolve integrity issues and continue each configuration's unfinished work without repeating shared intake or accepted review. If older work has no session checkpoint, inspect the saved request/profile, plan, data pair and results, then record their verified state and any explicit acceptance evidence in a new checkpoint. For directory output, `onboarding.json` locates configurations and next plan commands; it is an output index, not the session checkpoint. A saved file without evidence of acceptance still needs stage 3 review. + +Validate that artifacts still match the checkout's CLI, selected deployment and collection settings using the existing checks below. An accepted final request can start at stage 4, a valid plan at stage 5, and a verified matching data pair at stage 6. Preserve completed artifacts when blocked and save the specific missing input and next action. Changed deployment, resource profile or collection bounds require review and a new collection directory. Validation-only changes can reuse the verified collection plan; use a separate results directory for each replay. Collection's `--resume` continues its matching collector campaign, while `onboard resume` supplies the broader session context. ## Create the request @@ -21,7 +324,7 @@ Use an environment installed from this checkout; see [development setup](../DEVE aisimulate onboard init --interactive --output support-request.yaml ``` -Enter the actual model identifier or checkpoint path, pinned model revision, dense/MoE kind, pinned vLLM version, GPU system, allocation, and interconnect. Then choose a TP size and pilot workload. Supplied options skip their prompts. Enter accepts displayed defaults; invalid values can be corrected; Ctrl-C or end-of-input cancels without saving. Existing files require `--overwrite`. +Enter the actual model identifier or checkpoint path, pinned model revision, dense/MoE kind, pinned vLLM version, GPU system and interconnect. Choose a worker and review its context, scheduler and capture limits. Setup derives and displays the GPUs required for that worker. Supplied options skip their prompts. Enter accepts displayed defaults; invalid values can be corrected; Ctrl-C or end-of-input cancels without saving. Existing files require `--overwrite`. Both guided and scripted setup use onboarding. `--profile onboarding` is an optional spelling of the same behavior. For automation, supply the identity flags directly: @@ -31,16 +334,18 @@ aisimulate onboard init \ --model-revision YOUR_IMMUTABLE_REVISION \ --model-kind dense \ --framework-version YOUR_PINNED_VLLM_VERSION \ - --gpu h200_sxm --gpu-count 4 --interconnect nvswitch \ - --tensor-parallel 2 \ + --gpu h200_sxm --interconnect nvswitch \ + --tensor-parallel 2 --context-length YOUR_MODEL_CONTEXT_LIMIT \ --output support-request.yaml ``` -Replace the model and runtime placeholders with the actual inputs. The example does not identify an Inkling checkpoint or claim that four H200s can run your model. Framework support currently selects vLLM. Optional tokenizer, chat-template, and AISimulate revisions are recorded only when supplied. +Replace the model, runtime and hardware inputs with your deployment's values. Framework support currently selects vLLM. Optional tokenizer, chat-template, and AISimulate revisions are recorded only when supplied. + +With `--model-config`, omitted parallelism flags trigger model/hardware-aware suggestions. Headless setup selects only a fully assessed default; otherwise it exits without saving. Identifier-only and supplied-profile setup retain TP1 when parallelism is omitted. Review the default limits below and change them through `aisimulate onboard init --help` options or config-based profile review. -The default pilot uses 1,024 input tokens, 128 output tokens, concurrency 1, four requests, a 16,384-token context limit, TTFT target 1,000 ms, and TPOT target 100 ms. Scripted setup defaults to TP1 unless `--tensor-parallel` is supplied. These are planning defaults, not measured model capacity or latency. Change them with the corresponding flags shown by `aisimulate onboard init --help`. +The optional `predict/pilot.yaml` and `recommend/pilot.yaml` examples use 1,024 input tokens, 128 output tokens, concurrency 1, four requests, TTFT target 1,000 ms and TPOT target 100 ms. Their workload/SLA flags do not size memory, choose the collector grid or change the collection identity. The example token lengths must still fit the selected runtime context. Guided intake does not prompt for these synthetic workload or SLA fields. -One node is the default; GPUs per node then equals `--gpu-count`. For multiple nodes, supply both `--node-count` and `--gpus-per-node`; their product must equal the total allocation. Each selected worker must fit on one node. Advanced flags include `--request-count`, `--max-candidates`, `--objective`, `--seed`, and `--sm`. +Onboarding takes no GPU-pool or node-allocation inputs. The selected topology establishes a minimum collection requirement, not a reservation or a check that those GPUs are available. Advanced flags include `--request-count`, `--objective`, `--seed`, and `--sm`. Replica counts and optimization budgets remain configurable in ordinary `predict` and `recommend` inputs. `--tensor-parallel` always means attention TP. Use all MoE dimensions explicitly for DEP and TEP; replicas are independent workers: @@ -50,7 +355,274 @@ One node is the default; GPUs per node then equals `--gpu-count`. For multiple n | DEP8 | `--tensor-parallel 1 --attention-data-parallel 8 --moe-tensor-parallel 1 --moe-expert-parallel 8` | `(1, 1, 8, 1, 8, 1)` | | TEP8 | `--tensor-parallel 8 --attention-data-parallel 1 --moe-tensor-parallel 1 --moe-expert-parallel 8` | `(8, 1, 1, 1, 8, 1)` | -The first profile implementation supports vLLM decoder-only models with PP1, CP1 and linear KV storage. Recorded DCP is not represented by this profile schema and is rejected; existing registered-model DCP profiles remain a separate supported path. AFD, encoder pools, speculative decoding, nonlinear recurrent state, wide EP and EPLB need additional metadata/semantics and are rejected by this route. +The profile route supports vLLM text decoders, including the text portion of multimodal checkpoints, with PP1, CP1 and linear or grouped cache storage. Recorded DCP is not represented by this profile schema and is rejected; existing registered-model DCP profiles remain a separate supported path. Groups describe full attention, sliding-window attention and supported short-convolution state. Grouped execution currently requires cold aggregated workers, HBM-only cache, no speculative decoding and `prefix_caching: false`. Prefix reuse, host/G3 offload, disaggregation and scalar cache-capacity overrides are rejected for grouped profiles. AFD, encoder pools, unsupported recurrent state, wide EP and EPLB remain outside this route. + +### Runtime and collection limits + +These settings define the runtime envelope used for profile sizing, collection and generated configurations. They describe different constraints; none independently proves that a trace is covered. + +| Setting | Scope and initial value | +| --- | --- | +| `--context-length` | Maximum input plus output tokens for one request. Fresh config/profile-based setup uses `min(declared context, 256000)`; identifier-only setup uses 256,000 until reviewed. The saved field remains `search.context_length`. | +| `--max-num-tokens` | Scheduled token budget per attention-DP rank. Use the supplied profile's bound, or an initial policy of 8,192. Saved as `collection.max_num_tokens` when explicitly set. | +| `--max-batch-size` | Scheduler sequence bound per attention-DP rank. Use the supplied profile's bound, or an initial policy of 256. It is independent of validation concurrency and does not request every prefill batch up to the bound. Saved as `collection.max_batch_size` when explicitly set. | +| `--gpu-memory-utilization` | Finite fraction of total GPU memory in `(0, 1]`, initially `0.90`. Used consistently for topology/resource admission, collection and generated prediction/recommendation/replay memory budgets. Saved as `collection.gpu_memory_utilization`; it does not establish fit or reserve GPUs. | +| `--prefill-cudagraph-policy` | `runtime` for new initialization: leave prefill compilation settings to the pinned runtime and resolve actual captures at engine initialization. `explicit` selects a reviewed capture extension. Saved as `collection.prefill_cudagraph_policy`. | +| `--max-prefill-cudagraph-size` | Positive prefill capture limit for `explicit` policy. Supplying a number selects `explicit` unless a contradictory `runtime` policy was also supplied. Choosing `explicit` without a number uses 2,048. Saved as `collection.max_prefill_cudagraph_size`; match the target serving configuration. | + +For example, add `--gpu-memory-utilization 0.92 --prefill-cudagraph-policy runtime` to review a different memory fraction while retaining engine-selected graphs, or `--max-prefill-cudagraph-size 2048` for the explicit extension. Runtime policy with a numeric capture size is rejected. Old saved requests that omit `collection.prefill_cudagraph_policy` retain the prior explicit policy and 2,048-token default; they are not reinterpreted as runtime-policy requests. An omitted legacy memory fraction retains the prior collector-argument omission and the existing 0.90 simulation budget. Review changes in a new request/collection directory. + +### Serving roles and CUDA graph settings + +Use `--worker-type prefill`, `decode` or `aggregated` to bind new collection and memory evidence to its serving role. `prefill` and `decode` requests collect only their own phase and use separate output roots. Their configured graph modes, capture limits and usable KV capacities need not agree with each other. Within each role, observed settings and rank-local cache evidence must match that role's accepted launch; an unexpected runtime change still fails validation. Aggregated requests exercise both phases with one shared serving configuration and combine compatible memory observations conservatively. A batch selecting PIECEWISE during prefill and FULL during decode is not itself a configuration mismatch. + +The following optional settings describe vLLM's configuration for the selected role. They require an explicit `--worker-type`, are included in saved request and collection identities, and remain editable during the review/accept/save step. Enter `clear` when editing a generic graph field to remove that override: + +| Option | Meaning | +| --- | --- | +| `--cudagraph-mode` | Configured vLLM mode: `NONE`, `PIECEWISE`, `FULL`, `FULL_DECODE_ONLY` or `FULL_AND_PIECEWISE`. Runtime/backend support still needs verification. | +| `--cudagraph-capture-sizes 1 2 4 8 16` | Increasing, distinct positive capture sizes; use `--prefill-cudagraph-policy runtime` to avoid a competing legacy override. | +| `--max-cudagraph-capture-size 2048` | Maximum capture size for this serving configuration. In aggregated mode it is passed to both phase launches. When a capture list is also supplied, it must equal the largest listed size. Mode `NONE` rejects capture settings. | + +The ordinary `--max-num-tokens`, `--max-batch-size`, `--context-length` and `--gpu-memory-utilization` settings belong to the selected role. Choosing a smaller smoke workload does not reduce these accepted engine settings. The older prefill capture flags remain available: on explicit aggregated requests the same capture override applies to both launches; role-less historical requests retain their original behavior. Do not combine generic capture size/list options with the legacy explicit capture policy. + +For two roles on the same topology, use per-entry `worker_type` and `collection` in `--parallel-configs`. This illustrative example assumes those graph modes are supported; it does not assert measured memory fit: + +```yaml +- worker_type: prefill + tensor_parallel: 4 + collection: + prefill_cudagraph_policy: runtime + cudagraph_mode: PIECEWISE + max_num_tokens: 8192 + max_batch_size: 256 + gpu_memory_utilization: 0.90 +- worker_type: decode + tensor_parallel: 4 + collection: + prefill_cudagraph_policy: runtime + cudagraph_mode: FULL_DECODE_ONLY + max_num_tokens: 2048 + max_batch_size: 128 + gpu_memory_utilization: 0.92 +``` + +Pass this file with the existing model/config, pinned runtime, hardware and precision inputs to `onboard init --parallel-configs roles.yaml --output-dir ./onboarding`. Review and accept each role separately. Outputs such as `prefill-tp4-dp1-moe-tp1-moe-ep1/` and `decode-tp4-dp1-moe-tp1-moe-ep1/` have independent request identities, plans, memory observations and validation reports. Store them under separate configuration IDs in the session checkpoint. A role or graph-setting edit invalidates that configuration's previous acceptance and collection identity; it does not change the other role's accepted evidence. + +For role-less historical requests, an explicit capture extension configures prefill only; decode retains its runtime-selected captures. For these requests, both `onboard probe-runtime` and `onboard collect-fpm` previews warn that one shared memory profile requires matching observed graph settings across phases. The warning does not reject explicit capture policy or establish that it will fail. Use runtime policy, explicit captures verified to match decode, or supported eager execution in both phases when that matches the serving target. Actual phase/rank disagreement is rejected during memory import/finalization; requested settings alone cannot establish compatibility. + +Profile resource bounds must cover the selected scheduler envelope. Conflicting declarations fail rather than silently clipping the request. With config-based setup, editing scheduler bounds recomputes dependent estimates; review those estimates again. An explicit byte override remains your declared bound until you edit it. + +The 256,000-token context policy is motivated by the [AgentX reference subset](#validate-fpm-query-coverage-with-agentx-replay), whose request limit is 256,000, not 262,144. It is a starting point for review, not an assertion that every maximum-length request fits or that the FPM table covers it. The scheduler sequence limit does not reserve maximum context for all sequences simultaneously. For the normal onboarding route, vLLM resolves cache capacity during initialization. Finalization uses observed request-usable capacity and group geometry; replay applies actual request/cache occupancy. Existing complete declared profiles retain their legacy arithmetic. GPU memory utilization is a budgeting input, not a hard limit on activation/workspace allocations. + +## Start from a local model config + +Use a local Hugging Face-style JSON configuration to fill supported metadata and create the existing FPM profile: + +```bash +aisimulate onboard init \ + --model-config /models/your-pinned-checkpoint/config.json \ + --interactive --output support-request.yaml +``` + +Setup reads that file and an adjacent `hf_quant_config.json` ModelOpt sidecar, when present, without downloading a checkpoint, importing model code, constructing an analytical model, or launching GPU work. It records each source's SHA-256 separately, rejects conflicting inline/sidecar quantization declarations, and preserves explicit profile overrides. Weight quantization does not establish runtime FMHA or KV precision; only explicit KV metadata supplies a cache-precision fact. It displays source information and derived inputs, then asks for unresolved values. Missing model metadata is collected before runtime/collection options so the model's context limit can bound the selected envelope. Config identity hints skip their ordinary prompts; explicit CLI identity options take precedence and conflicts can be corrected. A pinned checkpoint revision, literal runtime version, GPU system and interconnect still need your input when absent. The config's SHA-256 records the local source; it is not a checkpoint revision. + +With no parallelism flags, guided setup asks for missing shared precision/layout facts, shows a small topology shortlist with resource reasons, and lets you choose one worker with `--output`, or several with [`--output-dir`](#onboard-multiple-parallel-configurations). Enter selects only the displayed single default, and only when all required profile inputs are resolved and its estimated bytes fit the budget. Otherwise a candidate selection is required before setup resolves the remaining precision and cache geometry. Any explicit parallelism flag bypasses suggestions and retains the existing topology validation. Each chosen topology then follows the same profile review, edit, accept and cancellation flow below. + +When present, `text_config` must be one nonempty decoder configuration object. Only its geometry is used; wrapper and encoder dimensions are never merged into it. Otherwise, setup reads the flat decoder fields even when vision or audio metadata is present. Nested decoder architecture/model-type declarations select supported decoder validation and resource estimates. The profile preserves the declared architecture used by collection: nested `architectures` when present, otherwise the wrapper architecture. Provenance records both identities when they differ. A wrapper architecture never establishes an unknown nested decoder's resource layout. Shared outer dtype and quantization metadata are inherited only when the text section does not declare that metadata. The checkpoint identity hint and SHA-256 remain those of the original document. `model_max_length` is accepted as a context-limit alias; conflicting context declarations fail explicitly. + +Scripted intake and interactive final review state that FPM models the text decoder only. Config-derived estimates exclude multimodal encoders, projectors, preprocessing and other non-text components. Observed runtime cache capacity accounts for all components actually loaded by the worker, while timing describes the text decoder only. The saved profile preserves that notice in provenance through planning and generated prediction/recommendation configurations. Known incompatible text-decoder cache layouts remain rejected, and unknown layouts require explicit resource accounting. + +Memory remains pending for the selected TP, DEP or TEP worker until collection. Prompts request only unresolved identity, precision and supported cache geometry. Optional legacy byte overrides describe a conservative per-rank estimate. Enter integer bytes or an explicit unit such as `70 GiB`, `512 MiB`, or `1.5 GB`; the conversion must produce a whole number of bytes. Invalid individual values are prompted again. Model revision, runtime version, and deployment conflicts return to the corresponding option while retaining accepted resource answers. During initial input collection, incompatible config facts or combined resource bounds fail with the reason and leave no request; correct the source or overrides and rerun setup. Ctrl-C or end-of-input exits 130 and leaves no partial request or replacement of an existing request. + +Once all required inputs are available, guided setup shows the complete profile, sources, and selected deployment for review. Choose `edit` to replace an estimate, a previous answer, or a value from `--resource-overrides`. Choose a field by name, enter its value, and review the updated profile. For example, these illustrative inputs edit GPU memory utilization and record why: + +```text +Review action (accept/edit/cancel): edit +Field to edit: gpu_memory_utilization +gpu_memory_utilization (fraction in (0, 1]): 0.9 +... updated profile and sources ... +Review action (accept/edit/cancel): edit +Field to edit: provenance +provenance: Reviewed runtime collection settings; memory remains pending initialization. +... updated profile and sources ... +Review action (accept/edit/cancel): accept +``` + +Changing an input recomputes dependent estimates: for example, changing `kv_cache_dtype` updates inferred linear `kv_bytes_per_token` or grouped attention page bytes, and changing `max_num_tokens` updates inferred activation bytes. The review also exposes `runtime_context_length`, `gpu_memory_utilization`, `prefill_cudagraph_policy` and `max_prefill_cudagraph_size` as editable collection settings; `context_length` edits the model profile's declared maximum. Editing the policy to `runtime` clears a saved explicit size; editing the size selects `explicit`. Other explicit values remain in place until you edit those fields themselves. If an edit makes a memory estimate unavailable, it remains pending; setup asks only for missing required identity, precision or cache geometry. Invalid individual answers can be corrected; an edit that conflicts with the config or complete request is rejected with the reason, and the previous profile is retained. Model identity and topology remain the declared deployment. + +Enter `accept` to accept the exact reviewed request; Enter alone does not accept it. You can make repeated edits before accepting. Single-file output saves after acceptance; directory output saves after every selected profile has been accepted. Choose `cancel` at topology selection or review, or press Ctrl-C or send end-of-input at any prompt, to exit 130 without creating artifacts or replacing existing output. Single-file output also preserves an existing request with `--overwrite`; directory output rejects `--overwrite`. This review step applies to `--model-config --interactive`; scripted setup never prompts. + +For automation, supply unresolved precision and cache geometry through a flat JSON or YAML file. This example describes a BF16 decoder whose standard geometry comes from its local config. Non-KV bytes are intentionally omitted; initialization during collection will establish memory: + +```yaml +# resource-overrides.yaml +gemm_quant_mode: bfloat16 +moe_quant_mode: bfloat16 +fmha_quant_mode: bfloat16 +comm_quant_mode: half +kv_cache_dtype: bfloat16 +cache_layout: linear +kv_bytes_per_token: 262144 +max_num_tokens: 8192 +max_batch_size: 256 +provenance: Illustrative file shape only; replace values and this explanation with their actual source. +``` + +```bash +aisimulate onboard init \ + --model-config /models/your-pinned-checkpoint/config.json \ + --resource-overrides resource-overrides.yaml \ + --model /models/your-pinned-checkpoint \ + --model-revision YOUR_IMMUTABLE_REVISION \ + --framework-version YOUR_PINNED_VLLM_VERSION \ + --gpu h200_sxm --interconnect nvswitch \ + --tensor-parallel 4 --output support-request.yaml +``` + +Replace the model, revision, runtime and hardware inputs with the worker you will actually run. `--model-config` and `--fpm-profile` are mutually exclusive; `--resource-overrides` requires `--model-config` and also works with guided setup. Flat per-rank byte overrides require an explicit topology because the same byte bound cannot be transferred or rescaled across candidate tuples. For multiple configurations, place them in each `--parallel-configs` entry's `resource_overrides` instead. Shared precision/layout overrides can be used for automatic suggestions. Scripted setup never reads terminal input. It exits 2 without writing a request when required inputs remain unresolved or no fully assessed automatic default exists. Validation first lists all missing or invalid target identity and collection options. Once that stage is valid, it lists unresolved profile fields, or candidate-specific gaps and exact topology flags, and asks for the necessary inputs or guided setup. + +The flat override fields are `architecture`, `context_length`, `num_experts`, every precision and resource field shown above, `cache_block_sizes`, `cache_groups`, `moe_backend`, `attention_backend`, and optional `provenance`. For grouped resources, replace `cache_layout: linear` and `kv_bytes_per_token` with the declarations described below. Unknown fields, duplicate fields, invalid types, unsupported values, and incompatible cache semantics are rejected. Overrides are recorded with per-field provenance rather than discarded. Profile `context_length` is the model's declared maximum; the CLI `--context-length` selects a runtime limit that cannot exceed it. Scheduler limits are per attention-DP rank, as described in [Runtime and collection limits](#runtime-and-collection-limits). + +Reusing an observed profile with `onboard init --fpm-profile` defaults to its recorded memory utilization and limits context to its observed `max_model_len`. An explicit smaller context is allowed; changing utilization or exceeding the observed context requires new runtime memory evidence. + +Derivation is deliberately limited. Unambiguous architecture, context, expert count and model kind can be read independently of memory support. The initial supported estimates are: + +| Quantity | Supported derivation and limits | +| --- | --- | +| Weight bytes | Vanilla full-attention Llama, Mistral, Qwen2, Qwen3 and Mixtral BF16 layouts with complete geometry and vocabulary divisible by 64 and TP. Estimates include supported biases, tied/untied embeddings, replicated norms, and the Mixtral router. Quantized storage, unknown padding, auxiliary prediction tensors, shared experts and unsupported custom layouts remain unestimated; collection resolves runtime capacity. | +| KV bytes per token | Full-attention MHA/GQA in those layouts plus Qwen3 MoE and MiniMax-M2, with an explicit supported cache dtype and static tensor storage. KV heads shard or replicate according to TP; attention DP does not divide a rank's cache. MLA/DSA storage and dynamic or per-token quantization scales need explicit accounting. | +| Grouped cache pages | Declared full/sliding attention layers with sufficient MHA/GQA geometry, selected TP, effective KV dtype and explicit runtime block sizes. Supported short-convolution state gets separate windowed groups. Derived page bytes are minimum packed tensor estimates over all group layers; runtime padding and nonstandard storage need explicit review or overrides. | +| Activation bytes | The existing backend estimate for those full-attention layouts, 16-bit FMHA, and TP1/2/4/8, within the declared scheduler envelope. It has a 70 MiB floor and does not cover DSA workspaces or speculative decoding. | +| Runtime overhead | The selected packaged hardware specification's per-rank `misc.other_mem` estimate when present. It excludes the separate CUDA graph reservation. | +| Communication overhead | The exact `misc.nccl_mem` TP entry for a pure-TP worker when present. DEP and TEP estimates remain unavailable; total GPU count does not establish their communication storage. These estimates are optional. | + +Source notes identify assumptions and packaged hardware hashes. Weight precision does not establish runtime FMHA or KV precision. Config intake proposes the current collector communication identity `half` with its source. Other schema-valid communication overrides remain recorded for review, with a warning that the current collector rejects them for new collection. Recognized quantization declarations can establish checkpoint/FPM precision identities or an explicit KV dtype independently of the unknown quantized storage size; a checkpoint label does not mean every tensor has that dtype. Configs may supply useful identity facts while weight or activation estimates remain unavailable. Only unsupported cache semantics require additional geometry; non-KV estimates never become required questions. These estimates do not certify runtime memory fit. Unknown layouts must supply the remaining supported semantics explicitly; known incompatible layouts are rejected. + +The saved request embeds the complete profile and its provenance. The original model config and override file are no longer needed by `onboard plan`, generated prediction and recommendation configs, or request replay. You can inspect or supply the complete profile directly as described next. + +### Review grouped cache resources + +Config intake recognizes unambiguous full/sliding layer declarations, including `sliding_window`/`layer_types` and supported `sliding_window_size`/`local_layer_ids` aliases. It also recognizes the supported `use_sconv` and convolution-kernel metadata. It derives available layer/head geometry and records the source. Runtime attention block sizes after page unification, effective KV dtype and allocation padding are not established by the model config. Unknown resource assumptions remain explicit; unsupported recurrent semantics and contradictory layer/window declarations fail. + +After selecting precision and topology, investigate the pinned cache/backend implementation and available allocation metadata for block sizes after page unification; ask only for values that remain unresolved. The model config's omission does not imply that the user must supply them. Supply `cache_block_sizes` using the generated group names listed by intake. For example, a mixed full/sliding model might require this JSON answer, or the equivalent mapping in a `--resource-overrides` file: + +```json +{"full_attention": 16, "sliding_attention": 16} +``` + +These sizes are illustrative; use the target runtime's tokens per block after page unification. Supported convolution groups use the pinned kernel width as their default block size, which can also be overridden by group name. Partial answers retain the supplied entries and ask for unresolved attention groups. These declarations describe the cache being modeled; verify that they match the actual serving allocation. Keep config-derived packed minima labeled as estimates until runtime evidence establishes the allocation, including padding; revise the profile and its provenance if bring-up reveals different values. + +Review the resulting `cache_groups` and their geometry/provenance. Each `FpmCacheGroup` records `name`, `kind` (`attention` or `convolution`), `num_layers`, `block_size_tokens`, `page_size_bytes`, and optional `sliding_window`. **`page_size_bytes` is the rank-local aggregate cost of one logical block across every layer in that semantic group, including applicable runtime padding.** Do not multiply it by `num_layers` again or divide it by attention DP. An omitted or null window retains full history; a positive window limits retained history, and convolution groups require one. + +This aggregate is distinct from the runtime's padded physical page size for an individual layer or shared allocation slot. Even when the runtime unifies physical page sizes, semantic groups with different layer counts can have different aggregate bytes. Do not require equal `page_size_bytes` across all profile groups. Runtime observation import checks actual tensor/storage strides, sharing, padded pages and pool slots before deriving the profile's group costs and usable capacity. + +Config-only page formulas can omit runtime padding and are minimum estimates. In the standard `memory_pending` flow, these estimates do not authorize simulation: probing or collection must resolve geometry and capacity first. The legacy complete declared-memory path trusts the reviewed declaration, so its group bytes must conservatively cover the actual padded allocation; supplying non-KV bounds does not verify those page costs. Record the source and any unresolved padding explicitly. + +To include runtime padding or other known storage, choose `edit`, select `cache_groups`, and enter a complete JSON list of group objects with integer byte values. Update `provenance` with the source and remaining assumptions, then review the revised profile before `accept`. A headless `--resource-overrides` file can supply the same complete list instead of `cache_block_sizes`. For example, this illustrates the cache portion of an override for two full-attention layers and six windowed layers; replace all values with the reviewed layout: + +```yaml +cache_layout: grouped +cache_groups: + - name: full_attention + kind: attention + num_layers: 2 + block_size_tokens: 16 + page_size_bytes: 131072 + - name: sliding_attention + kind: attention + num_layers: 6 + block_size_tokens: 16 + page_size_bytes: 393216 + sliding_window: 512 +provenance: Illustrative cache shape only; replace with the runtime layout, padding and source. +``` + +Include the remaining precision and cache geometry requested by intake; non-KV memory stays pending. Explicit group overrides must preserve the config's known attention/convolution layer counts and retention windows. Do not supply `kv_bytes_per_token` with grouped resources or use a linear override to discard declared windowed/convolution state. Saved profiles preserve the full groups and their provenance through planning, prediction, recommendation and replay. + +The native allocator shares one rank-local byte budget across groups and admits their allocations atomically. Full-attention groups retain history; windowed groups retain blocks covering the last `window - 1` computed tokens plus all tokens scheduled for the next forward. Expired completed pages are evicted before subsequent forwards. Temporary prefill pages remain charged through completion, so a long prefill chunk can need more memory than the retained window alone. Logical request progress and FPM query context remain unchanged: a long-context request still needs timing coverage for that logical context even when its retained cache has stopped growing. + +### Preview and choose parallelism + +`--suggest-parallel` emits one JSON report to stdout without creating a request, plan, timing data or output directory. It requires `--model-config` and the actual target identity; config hints can supply the model label and kind, but checkpoint revision and runtime version are never invented. Use your intended runtime/collection limits and shared precision declarations for meaningful resource accounting: + +```bash +aisimulate onboard init \ + --model-config /models/your-pinned-checkpoint/config.json \ + --model /models/your-pinned-checkpoint \ + --model-revision YOUR_IMMUTABLE_REVISION \ + --framework-version YOUR_PINNED_VLLM_VERSION \ + --gpu h200_sxm --interconnect nvswitch \ + --suggest-parallel +``` + +The report includes hardware/config provenance, exact `cli_flags`, required GPUs, resolved fields and sources, `missing` inputs, rejected candidates and a nullable `default`. Without explicit runtime precision, the example can report `needs_inputs` candidates and no default; use those diagnostics to collect shared facts in a flat `--resource-overrides` file and rerun. The preview does not validate or write `--output`, even if that request path already exists. It cannot be combined with `--interactive`, profile options, explicit topology flags, `--output-dir` or `--parallel-configs`. To save a request, remove `--suggest-parallel`; use a fully assessed automatic default or select a candidate by adding its exact flags. To save several configurations, use directory output below. Rank-local byte overrides require an explicit choice. A pending-input candidate can be selected in guided setup, which asks for its missing bounds before final review. + +Dense decoders consider TP; MoE decoders consider pure TP, DEP and TEP. The shortlist contains at most two choices per family and six distinct tuples overall, preferring the smallest estimated fit and the next one. Widths are powers of two within the packaged node domain. A declared NVLink/NVSwitch fabric can use a packaged fast rack domain when its inter-node bandwidth matches or exceeds its intra-node bandwidth; a declared `none` interconnect permits only one GPU automatically. This is hardware architecture, not an available GPU allocation. Explicit topology choices may exceed the shortlist. + +The memory precheck compares complete per-rank resource estimates against the selected `--gpu-memory-utilization` fraction of packaged GPU memory, initially 90%. Linear cache reserves one full runtime context plus one cached token per rank for planner headroom. Grouped cache uses the native conservative per-request peak at the selected context and scheduled-token bound, including block alignment and transient prefill pages. Grouped planning requires this checkout's compiled AISimulate extension. Neither path multiplies context by the maximum scheduler sequence count or assumes balanced attention-DP routing. Concurrent requests share the cache capacity left after non-cache resources. Missing precision, weights, cache, activations or reservations prevent a default; a known lower bound can reject an oversized candidate but cannot prove fit. Planning estimates for DEP/TEP communication storage, many quantized or custom decoder layouts, and large TP activation envelopes remain unresolved unless supplied. These optional byte estimates are not required to collect observed runtime memory. CUDA graph reservations and non-text components remain outside these estimates. Only known geometry constraints are checked; generated-plan admission and actual serving-runtime checks remain authoritative. The shortlist does not establish timing coverage, optimal performance or measured accuracy. + +### Onboard multiple parallel configurations + +Use `--model-config` with `--output-dir` to prepare separate requests and profiles in one session. Guided setup reads shared inputs once, then accepts one or more comma-separated candidate numbers: + +```bash +aisimulate onboard init \ + --model-config /models/your-pinned-checkpoint/config.json \ + --interactive --output-dir ./onboarding +``` + +For example, enter `1,2` to select the first two displayed candidates. Blank input selects only the assessed single default, when one exists. Invalid, repeated or out-of-range numbers are prompted again. Setup derives cache geometry separately for each selected tuple; memory remains pending until its collection worker initializes. Collection settings start from the shared CLI flags and can be changed for each profile in guided review. Review and explicitly `accept` each profile; edits to one profile do not change another. Cancellation, EOF or an invalid later profile leaves no new output, even after an earlier profile was accepted. + +For headless setup or explicit choices outside the shortlist, supply `--parallel-configs` as a nonempty JSON or YAML list. Each entry requires a positive integer `tensor_parallel`; optional topology fields are `attention_data_parallel`, `moe_tensor_parallel` and `moe_expert_parallel`. Use these fields to express the TP, DEP and TEP tuples in the [topology table](#create-the-request). PP and CP remain 1. Unknown fields, invalid types, unsupported topologies and duplicate resolved tuples within the same serving role are rejected, including duplicate tuples with different precision overrides. Optional `worker_type` and `collection` select per-role launch settings as described in [serving roles](#serving-roles-and-cuda-graph-settings). For multiple selected precision combinations, use separate fresh output roots as described in [stage 2](#investigate-runtime-constraints-and-precision-options); this does not expand collector support. + +These example files describe two dense TP choices. Precision and cache geometry come from shared runtime facts and the model config; memory is observed separately for each tuple: + +```yaml +# shared-overrides.yaml +gemm_quant_mode: bfloat16 +moe_quant_mode: bfloat16 +fmha_quant_mode: bfloat16 +comm_quant_mode: half +kv_cache_dtype: bfloat16 +cache_layout: linear +``` + +```yaml +# parallel-configs.yaml +- tensor_parallel: 2 +- tensor_parallel: 4 +``` + +```bash +aisimulate onboard init \ + --model-config /models/your-pinned-checkpoint/config.json \ + --model /models/your-pinned-checkpoint \ + --model-revision YOUR_IMMUTABLE_REVISION \ + --framework-version YOUR_PINNED_VLLM_VERSION \ + --gpu h200_sxm --interconnect nvswitch \ + --resource-overrides shared-overrides.yaml \ + --parallel-configs parallel-configs.yaml --output-dir ./onboarding +``` + +Each entry's `resource_overrides` uses the existing [flat override fields](#start-from-a-local-model-config). Shared `--resource-overrides` may contain precision/layout metadata and `cache_block_sizes` valid for the selected runtime, backend and precision. For multiple choices, shared per-rank byte overrides and complete `cache_groups` are rejected: put them in the corresponding entry or enter them during its guided review. Cache page bytes are derived separately for each precision and tuple, even when block sizes are shared. Do not copy or rescale a reviewed byte bound from another configuration. + +`--parallel-configs` requires both `--model-config` and `--output-dir`, and conflicts with explicit topology flags and `--suggest-parallel`. Adding `--interactive` skips candidate selection but still reviews every supplied configuration. Explicit topology flags with `--output-dir` and no configuration list produce one configuration. Existing `--output FILE` behavior and single-file request/profile schemas are unchanged; `--output` and `--output-dir` are mutually exclusive. + +Choose a fresh output root or an existing empty directory. Directory output rejects `--overwrite`, nonempty roots and symlink/path conflicts, and does not replace collection results. After every profile validates and, when interactive, is accepted, setup publishes: + +```text +onboarding/ + onboarding.json + tp2-dp1-moe-tp1-moe-ep1/ + request.yaml + fpm-profile.json + tp4-dp1-moe-tp1-moe-ep1/ + request.yaml + fpm-profile.json +``` + +Each deterministic directory names its resolved tuple; each request embeds one deployment and matches its exported profile. The `configurations` list in `onboarding.json` records each topology, minimum collection GPUs, request/profile paths, intended `collection/` directory and exact next `aisimulate onboard plan` command in `plan_command`. Run those commands to create the ordinary plans, then use each plan's existing `collect-fpm` preview, collection and `validate-fpm` commands with separate result directories. Setup does not launch collection. The selected configurations can be collected in separate runs on the same GPUs; their GPU requirements are not summed. Once saved, the requests and profiles can be planned without the original config or override files. ## Provide identity and resource metadata @@ -68,11 +640,11 @@ For a model without an analytical class, create a JSON or YAML FPM profile and p | Deployment runtime hints | `moe_backend`, `attention_backend` (default `auto`), `enable_wideep`, `enable_eplb` (currently both `false`). | | Deployment `resources` | Conservative bounds described below, specific to this topology. | -Each `resources` record requires integer `weights_bytes`, `activations_bytes`, `runtime_overhead_bytes`, `comm_overhead_bytes`, positive `kv_bytes_per_token`, `cache_layout: linear`, positive `max_num_tokens` and `max_batch_size`, and a `provenance` explanation. Obtain these from checkpoint tensor/storage metadata and serving-runtime accounting or conservative explicit estimates. A profile is an input declaration; AISimulate does not certify that its values describe your runtime. +A legacy declared-memory `resources` record supplies all four integer fields: `weights_bytes`, `activations_bytes`, `runtime_overhead_bytes` and `comm_overhead_bytes`. They are optional for the pending runtime-memory workflow. Every profile requires positive `max_num_tokens` and `max_batch_size`, and a `provenance` explanation. Choose either `cache_layout: linear` with positive `kv_bytes_per_token`, or `cache_layout: grouped` with a nonempty `cache_groups` list and no scalar token rate, as described in [grouped cache review](#review-grouped-cache-resources). Existing linear profiles retain their serialized shape. Supply a complete legacy declaration only when those bounds are already known and intended; new onboarding observes runtime cache capacity instead. A profile is an input declaration; AISimulate does not certify that its values describe your runtime. All resource bytes are **per rank**, and must bound every rank. Scheduler envelope limits are also per attention-DP rank; they are not the summed batch and token totals used by the worker's FPM timing query. Do not copy TP resource values into DEP/TEP merely because GPU counts match. Include all non-KV storage, including quantization overheads. Overhead fields exclude the separate `kv_cache.capacity.cuda_graph_reserved_bytes` reservation. Profiles carry explicit quantization identity; `text_only` and `unrecorded_quant_modes` overrides are rejected on this route. Requests above the declared scheduler envelope fail explicitly. Collection validates both prefill and explicit decode batch limits against that envelope before launching workers. -Automatic KV admission subtracts these non-KV bounds and CUDA graph reservation from the configured fraction of GPU memory, then divides by the declared KV bytes per token. `context_length: max` uses the profile limit; `bytes_per_token: auto` uses the exact deployment's cache geometry. These operations need a hardware specification and profile, but no FPM timing files or model graph. +For complete legacy declarations, automatic KV admission subtracts these non-KV bounds and CUDA graph reservation from the configured fraction of GPU memory. Linear cache divides the remainder by its declared bytes per token; grouped cache passes the shared byte budget and group pages to the native allocator. `context_length: max` uses the profile limit. Keep `bytes_per_token: auto` for profile resources; grouped profiles reject explicit scalar rates and block-count capacities. These operations need a hardware specification and profile, but no FPM timing files or model graph. The [canonical cache-budget API](core-api.md#fpm-profile-cache-groups-and-byte-budgets) returns grouped `total_kv_size_bytes` and `request_peak_cache_bytes`; scalar `total_kv_size_tokens` and `kv_size_per_token_bytes` are null. `estimate_num_gpu_blocks` rejects grouped profiles. The declared checkpoint revision is preserved for review and replay. Existing FPM tables may not pin a checkpoint revision; supplying one in a profile does not prove that the historical measurements came from that exact revision. Record that limitation in `provenance` when validating existing data. @@ -89,15 +661,30 @@ aisimulate onboard collect-fpm \ --output-dir ./aisimulate-support ``` -The first command saves the request, `support-plan.json`, `commands.json`, `predict/pilot.yaml`, `recommend/pilot.yaml`, and a local `systems/` directory. When supplied, the profile is also saved as `fpm-model-profile.json`, included in the collector command, and embedded in prediction/recommendation configs. The plan records the CPU resource estimate. The second command prints the collector invocation without launching it. Generated command vectors and printed next commands use absolute output paths and preserve spaces or shell punctuation. Use a separate output directory for each request. `--overwrite` can repair missing generated files for the identical saved request, including an interrupted write before `support-plan.json` was created. Repair requires an intact matching `request.yaml`, rejects changed inputs, and preserves existing collected data. +The first command saves the request, `support-plan.json`, `commands.json`, `predict/pilot.yaml`, `recommend/pilot.yaml`, and a local `systems/` directory. When supplied, the profile is also saved as `fpm-model-profile.json`, included in the collector command, and embedded in prediction/recommendation configs. The plan records pending memory or the explicitly declared resource estimate, plus effective collection limits and their sources. The second command prints the collector invocation without launching it. Generated command vectors and printed next commands use absolute output paths and preserve spaces or shell punctuation. Use a separate output directory for each deployment or collection envelope. + +The v2 plan's `request_id` identifies collection inputs: deployment identity, exact profile and collection settings. Its separate `validation_id` also includes synthetic workload/SLA, recommendation objective and seed. To change only those validation inputs, copy the request to a new file, edit that copy, then run `onboard plan --config UPDATED_REQUEST --output-dir EXISTING_PLAN --overwrite`. The command first validates the original saved request and every existing generated file against the saved plan, then refreshes validation examples and their hashes. It can also repair missing generated files for the same collection. If writing was interrupted before `support-plan.json` was created, repair requires an intact identical `request.yaml` and checks every existing generated input before writing. It preserves collected data, checkpoints and prior result directories. Changed collection/profile inputs or modified existing generated files are rejected; do not edit files inside the plan to bypass those checks. + +If an AISimulate update changes generated guidance, recreate the plan in a new output directory: validation compares generated files byte for byte and does not migrate existing plans. A validation trace is supplied separately to `validate-fpm`, so selecting another trace does not require regenerating the collection plan. -If a software update changes generated guidance, recreate the plan in a new output directory: repair compares generated files byte for byte and does not migrate older plans. Guidance changes do not change the saved-request identity checks used by collection. +The plan reports `fpm.collection_gpus_required`, derived as attention TP times attention DP. Collection uses only that worker width and the matching `tp`, `pure_tp`, `dep`, or `tep` preset. The generated prediction and recommendation configs both use one worker. Recommendation has one exact preset and one trial; its `optimization.constraints.max_candidate_gpus` equals the selected worker's requirement so even workers wider than the runtime default are representable. That derived cap describes this validation example and does not declare the user's total available GPUs. There is no replica expansion or deployment optimization during onboarding. -The guided plan uses one selected parallel tuple. By default it evaluates a single worker, so recommendation is not a broad deployment search. `--max-candidates 2` additionally considers the largest count of identical workers that fits the allocation, when that differs from one worker. Each choice gets an independent recommendation config pinned to that replica count with a one-trial budget. The single worker keeps `recommend/pilot.yaml`; the second choice uses `recommend/replicas-N.yaml`, where `N` is its replica count. The plan reports the actual candidate count and lists both config and result paths. Collection uses the matching `tp`, `pure_tp`, `dep`, or `tep` preset and exact worker GPU count. +Saved `aisimulate-support-request/v1` requests are read as v2 while preserving their declared context and profile resource bounds; an omitted legacy context retains 16,384. Legacy validation concurrency is not reused as a collection sequence limit. Review these retained and newly independent settings before use. Old `aisimulate-support-plan/v1` directories must be regenerated in a **new directory** because their collection commands depended on the synthetic validation workload. Do not overwrite old timings or checkpoints to reuse the old plan identity. -Before execution, prepare the real checkpoint and the pinned runtime using the existing [FPM collection guide](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md). The packaged collector invokes a Generator-resolved Dynamo/vLLM deployment and needs the corresponding GPU resources, deployment configuration, permissions, and model access. Invoking its command locally does not create that environment. `commands.json` publishes the guarded `aisimulate onboard collect-fpm --execute` command for collection, alongside a read-only collector planning command. +Some earlier v1 drafts also included `identity.gpu_count`, `identity.node_count`, `identity.gpus_per_node` and `search.max_candidates`; those obsolete fields remain rejected. Copy the request, remove them and review the retained topology/profile before regeneration. Plan summaries use `fpm.collection_gpus_required` and `search.candidates[].required_gpus`. The old `onboard init` allocation flags (`--gpu-count`, `--node-count`, `--gpus-per-node`, `--max-candidates`) are not accepted. Ordinary `predict` and `recommend` replica and GPU-budget APIs are unchanged. -After those prerequisites are ready, explicitly launch collection from that environment: +Before execution, prepare the real checkpoint and compatible pinned runtime image using the existing [FPM collection guide](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md). The packaged collector generates the runtime, creates the benchmark workloads and launches Dynamo/vLLM workers. It needs GPU resources, deployment configuration, model/image access and the selected executor's prerequisites below; it does not require a prestarted serving process or HTTP endpoint. The command launches workers within that prepared environment rather than provisioning the cluster itself. `commands.json` publishes the guarded `aisimulate onboard collect-fpm --execute` command for collection, alongside a read-only collector planning command. Apply the selected deployment options to either command; the generated examples retain Kubernetes defaults. + +For a source checkout installed with `uv sync`, activate its environment and expose the collector source package before running generated collector commands. From the repository root: + +```bash +source python/aisimulate/.venv/bin/activate +export PYTHONPATH="$PWD/python/aisimulate${PYTHONPATH:+:$PYTHONPATH}" +``` + +Alternatively, run the collector from the `python/aisimulate` source directory with that environment active. The installed release wheel includes the collector and needs no source-path adjustment. These steps make the command importable; they do not provide the target deployment or GPUs. + +After those prerequisites are ready, complete the [readiness procedure](#check-timing-readiness-before-full-collection) before explicitly launching full collection from that environment: ```bash aisimulate onboard collect-fpm \ @@ -115,23 +702,610 @@ python3 -c "import importlib.metadata; print(importlib.metadata.version('vllm')) Use that complete string for both the request's `--framework-version` and the collection profile's deployment `backend_version`. On a mismatch, the error reports the expected and observed strings, and the observed version remains in `collector-provenance.json`. Recover by using the pinned runtime, or by creating a new collection profile pinned to the observed full version and regenerating the request and collection plan in a new output directory. -For a profile-based campaign, the collector validates its resolved topology and precision against the supplied profile before execution. It does not relabel an FP8 cell as BF16 or change checkpoint quantization to satisfy the profile. A mismatch reports the conflicting field and requires a matching profile/runtime or a supported collector configuration. Resource admission uses the supplied bounds without constructing an analytical model. Profile contents participate in the collector's frozen-plan identity. +For a profile-based campaign, the collector validates its resolved topology and precision against the supplied profile before execution. It does not relabel an FP8 cell as BF16 or change checkpoint quantization to satisfy the profile. A mismatch reports the conflicting field and requires a matching profile/runtime or a supported collector configuration. Pending profiles defer memory admission to the initialized runtime; complete declared profiles retain CPU admission without constructing an analytical model. Profile contents participate in the collector's frozen-plan identity. New profile-based collection currently uses checkpoint-native weight, FMHA and KV precision without consulting op-level timing tables. The requested KV dtype must match that checkpoint-native dtype. An older FPM profile can still be valid for prediction while being unsuitable for new collection: the historical MiniMax/H200 BF16 FMHA fallback cell differs from its checkpoint-native FP8 inference. Use the historical identity to query those timings and a matching native profile for new collection; the collector rejects that mismatch. Publish the new FP8 campaign into a clean, separate dataset using a new onboarding output directory. Existing cell IDs omit FMHA precision, and publication retains the first published run for a cell ID. If the historical BF16 dataset already holds that cell ID, publication skips the new FP8 run. A collection profile must select one literal runtime version for the target hardware/backend. Successful formal collection and resume verify that the published data's pod-reported runtime version exactly matches `framework_version`, including suffixes such as `+cu128`. A mismatch exits 1 and preserves collection artifacts. Use the declared runtime or create a new request and plan for the observed version; generated configs are never silently retargeted. Diagnostic smoke runs do not publish or verify formal data. -Set deployment options directly on `onboard collect-fpm`: `--dynamo-version VERSION`, `--image IMAGE`, `--namespace NAME`, `--model-cache NAME[:MOUNT[:SUBPATH]]`, `--transport nvlink|ib|efa`, and `--image-pull-secret NAME`. The mount, when supplied, is an absolute container path. Prefer an immutable image digest. Supply the same options when previewing, executing, and resuming; deployment settings are part of the collector's frozen-plan identity, so changed settings require a new output directory. Arbitrary collector arguments and engine overrides are not accepted by this command. +### Check timing readiness before full collection + +Apply this procedure independently to every selected precision/topology configuration before new full collection. The readiness gate governs collection; it does not require recollecting compatible historical formal data that can already be used by ordinary prediction/recommendation. Historical data reuse retains the separate identity, memory, query-coverage and accuracy checks, including recorded provenance gaps. First inspect any saved native campaign artifacts: + +```bash +aisimulate onboard collect-fpm \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support --check-readiness +``` + +This CPU-only command prints a JSON assessment without launching workers, writing a report file or updating the session checkpoint. It exits 0 when ready and 1 when incomplete or blocked; invalid CLI inputs or onboarding plans retain exit 2. Do not combine it with `--execute`, `--smoke`, `--limit` or `--resume`. It validates the saved plan, current attempts and native evidence rather than trusting saved success labels. Existing formal attempts take precedence over smoke evidence for the same cell, including failed formal attempts; an older successful smoke cannot hide a later failure. Missing evidence is incomplete, not successful. This check reports the archived deployment; full execution additionally compares it with the requested runtime/launch plan. A saved report does not transfer readiness to another image, runtime, precision, topology or collection envelope. + +When there is no reusable complete evidence and collection is authorized, use the same reviewed deployment options for a bounded smoke: + +```bash +aisimulate onboard collect-fpm \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support --execute --smoke +``` + +Guided smoke defaults to **all cells in this configuration's plan**: both phases for aggregated/legacy requests, or the selected prefill/decode phase for an independent role. Dynamo generates a bounded native sample grid; smoke does not reduce the reviewed context, scheduler, graph or memory settings. An older smoke that omitted an explicit prefill capture configuration cannot qualify the formal launch: inspect its archived arguments, preserve the evidence and requalify with the matching settings in a new campaign. `--limit N` explicitly limits cells and may omit a required phase, leaving readiness incomplete. Smoke and formal collection use separate collector checkpoints and artifact directories, and smoke never publishes the formal data pair. Use the same executor/image/mount/transport options as the intended formal run; for an existing smoke, add `--resume` after inspecting its checkpoint. + +Execution saves `fpm-readiness.json` in the onboarding output directory and returns nonzero for an execution failure or incomplete/blocked readiness, including a limited smoke that omitted required cells. Inspect the per-cell status, blockers and next action, together with these evidence categories: + +| Evidence | Interpretation | +| --- | --- | +| Frozen runtime/launch identity and observed backend | Verify the exact planned deployment and current attempt. Requested settings do not independently attest every runtime component. | +| Native KV warm-up metadata | Inspect `enabled`, `warm_eligible`, `skip_reason` and per-rank diagnostics. Missing or contradictory required protocol evidence cannot qualify the cell. | +| Native and formal seed-regime counts | Formal conversion applies the selected strategy's protocol, including approved `skip:moe_tp_balanced_by_construction` and the narrowly verified dense full-attention skip described below. Inspect native evidence and the publication's classification; do not decide usability from a raw stamp alone. | +| Direct-eligible point counts | Formal `fake_fallback` rows are excluded. Zero eligible decode points blocks full collection and direct prediction. Mixed usable/fallback data covers only the reported samples and is not complete query coverage. | +| Available execution-worker observations | Inspect observed attention groups, resolved cache settings and graph configuration. Contradictions in these checked observations are blockers. Missing optional observations remain unqualified without requiring completed memory sizing or serving validation before timing collection. Referenced `fpm-memory-*.json` files are unvalidated sources retained for separate runtime import/finalization checks; their presence does not qualify allocation or memory. | +| CPU allocation and affinity | For new Slurm campaigns, inspect `execution.cpu_affinity` and the per-node launcher, per-rank worker and scheduler sidecars. The actual step must expose the requested CPU pool; local schedulers must have sufficient distinct CPUs available, and observed threads must stay within that pool. Missing evidence is incomplete; contradictions or insufficient capacity block full collection. A shared pool does not mean dedicated per-rank pinning. | + +Zero-KV prefill can retain a native `real_prefix` path stamp while its formal seed regime is `n/a`; the stamp alone does not establish that cached tokens were read. Keep the native provenance and use the validated phase/token counts and formal regime when interpreting readiness. + +Dense decode can be synthetic by design only when the producer reports enabled warm-up with `warm_eligible: false` and `skip_reason: dense_model_content_insensitive`, and the pinned native evidence corroborates a dense full-attention model and its exact runtime/topology. The supported case is single-rank or attention TP with PP1, DP1, MoE TP1/EP1, and prefill and decode context parallelism both equal to 1. Initialized cache groups must be causal full attention, with no sliding window, attention chunking, convolution or recurrent state; a `FullAttentionSpec` type name alone is insufficient. Matching model and runtime identities and every expected native rank are required. Consistent qualifying points publish as `skip:dense_model_content_insensitive` under the explicit `dense_full_attention_v1` metadata policy. This retains their measured samples and raw provenance; it does not establish real KV content or serving accuracy. Other `fake_fallback` points remain ineligible for direct timing. + +Readiness counts describe the current reaggregation of verified native artifacts. They do not relabel an existing formal pair. Historical publications without the dense classification marker retain their recorded regimes during verification, finalization and validation; any revised publication needs a fresh output tree while preserving the original pair. + +`ready_for_full_collection` permits the next collection step only for the verified sampled scope. It does not establish memory readiness, interpolation coverage, repeatability or workload accuracy. Full execution checks matching usable smoke or formal evidence for every current cell before launching new GPU work; it does not treat a completed memory probe, process exit or existing report file as sufficient. Continue the existing collection-quality, memory-finalization, replay and matched-serving steps afterward. + +If no usable decode exists, inspect the pinned runtime, exact warm-up flags, eligibility/skip reason and worker logs before deciding why. `fake_fallback` alone does not prove unsupported convolution: TP eligibility, prefix settings, capacity and transient execution failures can also matter. Preserve raw results and completed prefill. Record a deterministic capability failure as blocked and stop unchanged automatic retries. Unsupported convolution or other recurrent-state seeding needs a separate Dynamo/runtime implementation and bounded qualification of the resulting runtime; disabling required warm-up or relabeling synthetic state does not provide support. Continue independent configurations under their own checks. + +One coordinator saves the report path, command/exit status, blockers and next action in the existing session checkpoint, using per-configuration `artifacts` and `progress` plus `research` for the investigation. For a read-only assessment, retain its emitted JSON in the research record or a separately saved evidence file; do not register a nonexistent report. Reload the checkpoint revision first and follow the [writer coordination rule](#orchestrate-independent-collection-campaigns); this assessment does not save conversational progress automatically. When execution updates a report, refresh its registered artifact reference. Preserve superseded evidence and decisions rather than editing a report to declare readiness. + +Use the [recovery procedure](#orchestrate-independent-collection-campaigns) for compatible saved artifacts. Guided formal resume may proceed without ready timing only when its checkpoint proves that no worker launch is needed. That exception permits artifact recovery; it does not qualify unusable timings or authorize a full GPU retry. Changed collection/runtime inputs require a new reviewed campaign, and AISimulate updates that change generated plan guidance retain the existing new-directory requirement. + +### Choose the collection executor + +Set deployment options directly on `onboard collect-fpm`: + +| Option | Kubernetes (default) | Slurm (`--executor slurm`) | +| --- | --- | --- | +| `--image IMAGE` | Overrides the worker image. | Required Pyxis image or accessible SquashFS image path. | +| `--container-mount SRC[:DST[:FLAGS]]` | Rejected; use the model-cache PVC option. | Repeat for checkpoint, cache or other required mounts; spelling and order are preserved. | +| `--cpus-per-task N` | Rejected. | Positive CPUs for each node's one-task worker/scheduler pool; new campaigns default to 16. | +| `--cpu-bind cores\|none` | Rejected. | Slurm binding for that node task; new campaigns default to `cores`. `none` retains the allowed mask imposed by the allocation/container. | +| `--namespace`, `--model-cache NAME[:MOUNT[:SUBPATH]]`, `--image-pull-secret` | Existing namespace, model PVC and registry-secret settings. A supplied mount is an absolute container path. | Rejected because they configure Kubernetes resources. | +| `--dynamo-version VERSION` | Pinned template version. | Same template-version selection. | +| `--transport nvlink\|ib\|efa` | GPU networking transport. | GPU networking transport; it does not select the executor. | + +Prefer an immutable image digest or a pinned image file. Supply the same options when previewing, executing and resuming; deployment settings are part of the collector's frozen-plan identity, so changed settings require a new output directory. [Save these options in the session checkpoint](#preserve-collection-deployment-options). Arbitrary collector arguments and engine overrides are not accepted by this command. + +For Slurm, run within an existing allocation obtained through your cluster's ordinary `salloc` or `sbatch` process. The collector uses `srun` with Pyxis/Enroot to launch its benchmark workers; it does not submit an allocation or choose job dependencies. The allocated node count must match the generated collector plan, with enough GPUs per node for the selected worker. Follow [campaign orchestration](#orchestrate-independent-collection-campaigns) when several configurations share capacity. These are collection-time requirements, not questions about the user's total GPU pool during model intake. Support for the executor does not qualify every model/runtime/topology combination. + +Request matching CPU resources in the allocation, for example `salloc --nodes=2 --ntasks-per-node=1 --gpus-per-node=4 --cpus-per-task=16` with your cluster's account, partition and time options. The collector starts one task per node with explicit `--cpus-per-task=16 --cpu-bind=cores`; local engine schedulers, GPU workers and supporting threads share this pool. The initial 16-CPU choice follows the existing Slurm generator default and is editable, not a measured optimum. It does not assign one dedicated core to each rank or choose a NUMA layout. Increase it when the local DP topology or measured host load requires more capacity; a request smaller than the number of local DP schedulers is rejected. CPU counts and observed logical CPU masks can differ with SMT. Slurm CPU IDs are used to count allocation capacity, not assumed to match Linux affinity numbering. + +The collector checks allocated CPU counts rather than full-node CPU capacity. Inside the same container step that will start the engine, it records `fpm-cpu-launcher.json` before model initialization and fails if the effective allowed mask is unavailable or too small. Worker observations after graph warm-up and scheduler initialization observations record actual process and thread masks. Inspect these during smoke, especially for DEP; do not infer isolation from the batch allocation size or from a separate diagnostic step. An allocation conflict fails promptly instead of waiting indefinitely for a node task. + +Keep the campaign directory on storage shared at the same absolute path across the allocated nodes. Make the pinned image and checkpoint accessible, and mount the model/cache paths required inside the container. The collector adds its own staged runtime and result mounts, preserves initialization and native-result evidence, and cleans up its own named job steps without canceling the caller's allocation. Commas and control characters are not allowed in individual mount strings; use repeated flags for multiple mounts. For example, preview an existing plan in shared storage: + +```bash +aisimulate onboard collect-fpm \ + --config /shared/onboarding/request.yaml \ + --output-dir /shared/onboarding/collection \ + --executor slurm \ + --cpus-per-task 16 --cpu-bind cores \ + --image /shared/images/pinned-runtime.sqsh \ + --container-mount /shared/models:/models:ro \ + --container-mount /shared/hf-cache:/root/.cache/huggingface \ + --transport ib +``` + +Use your cluster's actual paths and networking. Add `--execute` to this same command when its prepared allocation is available; add `--resume` as well when continuing its existing campaign. Both executors use the same native-grid collection, checkpoint, formal publication and memory-finalization workflow. Slurm does not introduce another timing grid or automatically launch an HTTP serving endpoint. + +Use the same CPU settings for runtime probes, smoke, formal collection and representative repeats. Save both values with the other deployment inputs in the checkpoint; previewed collector commands use `--fpm-slurm-cpus-per-task` and `--fpm-slurm-cpu-bind`. Changing the CPU policy requires a fresh campaign and matching smoke; preserve earlier data and reports instead of changing their saved identity. Legacy plans without a CPU policy remain readable for compatible CPU-only recovery and memory import, with CPU affinity explicitly unverified. They cannot launch new workers under silently added defaults. New runtime instrumentation must emit the worker and scheduler CPU evidence before its smoke can qualify a new Slurm campaign. + +For the [readiness smoke](#check-timing-readiness-before-full-collection), add `--execute --smoke`; `--limit N` also requires `--smoke` and may leave required cells unchecked. Diagnostic smoke and limited runs do not publish formal FPM data. Existing campaign data, raw artifacts, or checkpoints require explicit `--resume` and a readable matching collector checkpoint; otherwise choose a new output directory. A custom `--checkpoint-dir`, if needed, must remain inside the plan's `fpm-checkpoint/` directory. Selecting an empty checkpoint directory does not allow reuse of existing campaign artifacts. Smoke and formal campaigns have separate checkpoints and artifact directories, so an existing smoke run does not prevent the first formal run, or vice versa. The collector verifies the resumed checkpoint's frozen-plan identity. + +Planning and collection reject concurrent onboarding operations against the same plan directory. The persistent `.support.lock` file uses an OS advisory lock; ownership is released when the process exits, including after an abrupt termination. Leave the file in place. Request validation, saved onboarding-plan checks, and collector input resolution exit 2. Failures after collector execution starts, including a frozen checkpoint identity mismatch, exit 1 with a concise message. Interruption exits 130. + +### Orchestrate independent collection campaigns + +This is the onboarding agent's procedure over the existing CLI and scheduler. The CLI does not schedule a collection campaign across configurations. Use this planning and outcome tracking for runtime probes, formal collection and validation, but choose concurrency according to the command's actual prerequisites and checkpoint writes. Formal `collect-fpm` jobs do not update the session checkpoint; probe/import commands do. Schedule collection-quality checks after their native collection and formal publication, replay after memory resolution, and matched serving from the required collection/replay evidence. + +Before submitting jobs, save an execution table keyed by the session checkpoint's configuration IDs. For each configuration, record: + +- The reviewed inputs and plan/preview paths: accepted request/profile for formal collection, or draft/model config and probe preview while geometry is unresolved. Keep incomplete profiles unaccepted. Include pinned checkpoint/runtime/image, precision/topology, exact previewed command and deployment options. +- Its own collection, collector-checkpoint, artifact and validation paths; generated node count, minimum GPUs and GPUs per node; required phases and remaining access or runtime prerequisites. +- Intended concurrency, authorized resource/time limits, any actual data dependencies and the reason for each scheduling dependency. The exact grid and duration remain unresolved until the runtime supplies evidence; do not invent a second grid. + +Submit independent formal collection configurations independently by default and let the scheduler place them within their resource requests. For example, on a four-GPU-per-node system, **if the generated plans confirm this placement**, four selected configurations have the following resource needs: + +| Configuration | Collection allocation | Data dependency on another configuration | +| --- | --- | --- | +| TP4 | One node, four GPUs | None | +| DEP4 | One node, four GPUs | None | +| TP8 | Two nodes, four GPUs per node | None | +| DEP8 | Two nodes, four GPUs per node | None | + +Keep output roots separate and size each allocation to the plan it will execute. The current Slurm runner requires exactly the generated node count; a two-node allocation cannot simply be reused for a one-node plan. Parallel collection is optional; sequential runs can reuse capacity. + +The caller owns Slurm submission dependencies. The collector launches workers inside the supplied allocation and does not add `sbatch` dependencies between profiles. Choose dependencies according to their actual purpose: -For a diagnostic run, add `--execute --smoke`; `--limit N` also requires `--smoke`. Diagnostic smoke and limited runs do not publish formal FPM data. Existing campaign data, raw artifacts, or checkpoints require explicit `--resume` and a readable matching collector checkpoint; otherwise choose a new output directory. A custom `--checkpoint-dir`, if needed, must remain inside the plan's `fpm-checkpoint/` directory. Selecting an empty checkpoint directory does not allow reuse of existing campaign artifacts. Smoke and formal campaigns have separate checkpoints and artifact directories, so an existing smoke run does not prevent the first formal run, or vice versa. The collector verifies the resumed checkpoint's frozen-plan identity. +| Purpose | Submission policy | +| --- | --- | +| Independent configurations with no sequencing requirement | No dependency. | +| Intentionally sequence independent jobs to bound concurrent resource use or coordinate shared checkpoint writers | `afterany:JOB_ID`, which waits for termination regardless of success. | +| Consume an artifact produced by a prerequisite job | `afterok:JOB_ID`, which requires successful completion; the consumer must still verify the required artifact. | + +These are [Slurm dependency semantics](https://slurm.schedmd.com/sbatch.html#OPT_dependency). A dependency does not reserve the same physical GPUs. In the example, DEP4 does not need TP4 publication, and DEP8 does not need TP8 publication. Do not use `afterok` between those pairs merely to sequence capacity reuse. Preserve nonzero collection/post-processing exits; do not mask failures to release dependents. Pause other configurations only when evidence identifies a shared blocker, such as an inaccessible checkpoint or incompatible common image. + +Track each configuration through distinct outcomes: + +1. **Execution:** scheduler job ID/state and batch exit, per-rank worker/step exits, initialization logs and native outputs. A successful submission is not execution; completed GPU steps are not a completed campaign. Inspect [Slurm accounting](https://slurm.schedmd.com/sacct.html), including job and step records, alongside collector evidence. +2. **Collection and publication:** validated cell status, cleanup outcome, runtime compatibility and a verified formal data/metadata pair. Record a post-processing failure here even if every GPU step succeeded. +3. **Memory and validation:** separately record resolved memory and profile acceptance, collection quality gates, replay coverage and matched serving accuracy. Do not infer any of these from a successful job or a published table. + +Keep the execution table, scheduling rationale, exact commands, job history, stage outcomes and recovery decisions under the existing checkpoint's top-level `research`, indexed by configuration ID. Use each configuration's `progress` for its current stage, blocker and next action, and `artifacts` for evidence references. Register its mutable collector checkpoint as `kind: collector_checkpoint`, `scope: collection`; that file remains authoritative for individual cells. Keep effective deployment options in `inputs.collection_deployment` as described above, but keep volatile job IDs and scheduler state out of accepted inputs. One coordinator records progress with revision checks. Independent formal collection jobs write their separate campaign directories. + +**Serialize session-checkpoint writers for each invocation's full duration.** `onboard probe-runtime`, including preview, and `onboard import-observations` read a revision before doing their work and save the shared checkpoint afterward. The agent must coordinate them with every other writer of that file: do not run another probe/import, coordinator checkpoint update or profile acceptance while one is in progress, including during its GPU work. There is no automatic lock covering that full interval. Save planned commands and output paths before submission; retain new job IDs, notes and results in the ordinary job logs/outputs until they can be merged safely. After the invocation exits, reload the checkpoint and reconcile its current revision before the next write. This coordination does not make one configuration's success a prerequisite for another: once a failed invocation exits, an independent probe can proceed unless it has a shared blocker. Keep its output root separate. + +When a campaign fails: + +1. Inspect the batch and worker states, logs, collector checkpoint and artifact validity to identify the failed stage. Check whether any job or owned worker step is still active before resubmitting; prevent duplicate execution and preserve the failed attempt's evidence. +2. Continue independent configurations. A failed `afterok` dependency may leave a dependent pending or cancelled under site policy; requeueing the predecessor does not repair that failed dependency. Inspect the dependent and cancel or replace only the campaign's own job as needed before resubmitting it independently, or with `afterany` if sequencing is still required. Do not rerun successful TP measurements just to release an unrelated DEP campaign. +3. Use the original command's documented resume path and retain the same reviewed execution inputs and roots while repairing compatible post-processing. For formal collection, `onboard collect-fpm --resume --execute` can recover completed native artifacts, skip verified passed cells and retry publication. It may still launch missing or unfinished cells, so inspect the checkpoint, [timing readiness](#check-timing-readiness-before-full-collection) and resource requirements first; resume is not a guaranteed CPU-only operation. The readiness gate permits recovery without ready timing only when the matching checkpoint proves no worker launch is needed. Cells still marked `failed` or `cleanup_failed` are not rerun by guided resume. +4. If failed formal-collection cells need another GPU attempt, follow the lower-level collector's [recovery procedure](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md#recovery-and-cleanup), archive the old raw/log evidence and preserve the full emitted command's frozen arguments and paths. `--resume-retry-failed` belongs to `python -m collector.fpm_forward`, not `onboard collect-fpm`. After that retry, return through guided `collect-fpm --resume --execute` for its runtime-compatibility checks before finalization. If only memory finalization failed, retry `onboard finalize` with intact evidence and a fresh resolved output directory. Changes to reviewed runtime settings require a new reviewed campaign; never edit cell status or bypass ownership/artifact checks to declare success. + +If a probe/import fails at its checkpoint save because the revision changed, preserve its produced reports and observations. Reload and reconcile the latest checkpoint before following the [probe resume/import procedure](#preview-execute-import-and-review) under the writer rule above. Complete compatible observations can be imported into a fresh output directory without another GPU run. With unchanged inputs, probe `--resume --execute` reuses a verified complete attempt; incomplete or invalid evidence can still require GPU work. Do not repeat completed measurements merely to repair checkpoint bookkeeping. + +### How the collection grid is determined + +AISimulate and Dynamo jointly determine what is collected. For formal collection, their responsibilities are: + +| Part of collection policy | Responsibility | +| --- | --- | +| Worker deployment | The AISimulate collector generates the runtime and launches benchmark workers through Kubernetes or an existing Slurm/Pyxis allocation. | +| Runtime bounds | AISimulate forwards the reviewed per-request context, scheduled-token and sequence bounds, and configured GPU memory fraction, to prefill and decode workers. | +| Runtime capture policy | With `runtime`, AISimulate emits no prefill compilation override or capture-derived new-token sample cap. Capture sizes and the corresponding new-token axis/counts remain unresolved before engine initialization; the scheduler-derived KV-read sample cap remains bounded. | +| Explicit capture policy | With `explicit`, AISimulate constructs the capture-size list and overrides the prefill engine's compilation configuration. It derives prefill new-token and KV-read sample caps from captures and runtime bounds. These caps are inputs to Dynamo, not an exact point list. | +| Remaining sampling defaults | The collector inherits formal prefill batch-sample and decode-sample caps from the deployed Dynamo image. Inspect that image's defaults; a scheduler sequence bound does not request every prefill batch up to the bound. | +| Engine initialization | vLLM loads the model and allocates cache/workspaces, then resolves supported graph mode and capture sizes for the selected runtime configuration. Inspect effective values in initialization logs and raw runtime artifacts. | +| Exact points and timings | Dynamo self-benchmark combines those inputs with the initialized engine's capture sizes, KV capacity and feasibility checks to generate, admit and time the actual points. | + +An explicit prefill compilation override changes the engine being measured, so its capture configuration must match the serving target. Runtime policy also requires checking the actual resolved mode and captures: the same requested settings can resolve differently across model/runtime versions or scheduler bounds. The deployed image's inherited sampling defaults can change the collected set. Inspect the generated configuration, initialized engine and image behavior; runtime bounds alone do not describe the full sampling policy. + +The onboarding `support-plan.json` records the reviewed bounds and policy in `fpm.runtime_limits` and summarizes grid generation in `fpm.sampling`. For runtime policy, the collector plan's `point_generation.prefill_sampling` records `cudagraph_policy: runtime` and leaves capture sizes, capture counts and capture-dependent new-token fields null. Explicit policy expands the selected capture list and sample caps. The exact point count remains runtime-determined in either mode. The `point_generation` owner remains Dynamo because Dynamo generates the actual runtime points; that ownership does not mean all collection policy comes from Dynamo. + +The onboarding command supplies `--fpm-max-model-len`, `--fpm-max-num-batched-tokens`, `--fpm-max-num-seqs` and the selected `--fpm-prefill-cudagraph-policy`; new requests also supply `--fpm-gpu-memory-utilization`. A supplied numeric capture limit is forwarded as `--fpm-max-prefill-cudagraph-size`; explicit policy without one retains the collector's 2,048-token default. `--fpm-max-prefill-isl` is a total new-token budget, not a per-request context limit. DEP does not multiply these rank-local settings by GPU count. Onboarding does not build a separate trace-derived grid or promise a sample count or collection duration. + +Native artifact completion and skip checks account for the generated grid. A complete grid does not guarantee that AgentX replay or another direct-FPM caller can resolve every requested query. Validate those query shapes separately, using the published measurements and supported interpolation. + +Direct collector callers can set the same shared runtime flags. Omitting `--fpm-prefill-cudagraph-policy` retains the standalone collector's explicit 2,048-token capture default; select `runtime` deliberately to defer captures. Runtime policy with a numeric capture size is rejected. Omitting `--fpm-gpu-memory-utilization` preserves the existing runtime/default behavior; an explicit value must be finite and in `(0, 1]`. With a supplied FPM profile, omitted scheduler/context limits use that profile's bounds; without one, an omitted context retains the existing vLLM `-1` auto-fit behavior. Explicit limits must be positive and cannot exceed the supplied profile or contradict the prefill sampling bounds. Inspect the generated command and collector plan before committing GPU time. Successful formal collection publishes the FPM Parquet file and metadata pair into the plan's local systems data directory; diagnostic success alone does not provide that pair. Validate the worker's actual FPM query shapes afterward, including summed queries for DEP. + +New collection campaigns request KV preparation in both phases by default: `DYN_BENCH_PREFILL_REAL_SEED=on` for cached prefill and `DYN_BENCH_KV_WARMUP=on` for decode. Cached prefill computes the required prefix before timing the forward pass; prefill with no existing KV has no prefix to seed. Requested prefill seeding requires a compatible runtime and is checked before model loading. Explicit deployment environment overrides remain explicit, subject to the existing collection protocol checks. These defaults are part of the collection identity; preserve old campaigns and use a fresh campaign when changing the policy. + +Decode warm-up still depends on the deployed runtime's eligibility and available cache capacity. An enabled switch does not prove real KV was used for every point: inspect raw `kv_seed_regime`, warm-up eligibility, skip reasons and fallback evidence. The current formal table labels ordinary prefill seeding `n/a`; retain the raw artifacts to distinguish `real_prefix` from `fake_prefix`. No successful collection establishes prediction accuracy by itself. + +Benchmark prefix seeding and replay cache reuse have different purposes. The collector may retain prefixes to prepare timing points outside the measured forward pass. Cold replay starts from an empty cache, and grouped replay retains `prefix_caching: false` for cross-request reuse. Do not copy that replay flag into benchmark launches. Smoke sampling is diagnostic and uses small caps; it does not establish the formal runtime capture list or point count. + +## Resolve cache geometry with a runtime probe + +Use `onboard probe-runtime` when the selected runtime's block sizes, physical page costs, padding or capacity remain unresolved. It launches the existing collector before a complete simulation profile exists. A profile, timing table, guessed geometry or an already-running server is not a prerequisite. Select the model/runtime, hardware, joint precision, explicit worker topology, context, scheduler limits, memory fraction and capture policy first. The probe bounds native benchmark sampling while preserving those reviewed runtime settings; normal model initialization and graph warm-up still run. + +Both probe preview/execution and `onboard import-observations` automatically update the session checkpoint after their work. Follow the [campaign writer rule](#orchestrate-independent-collection-campaigns): serialize each full invocation with all other writers of that checkpoint, defer coordinator saves while it runs, then reload its current revision. Separate output directories alone do not prevent conflicting saves. Keep produced evidence if a save fails and reconcile the checkpoint before compatible resume or import. + +An onboarding agent should investigate and author a campaign-local adapter when no compatible bundled adapter exists. Inspect the exact pinned source and relevant vendor patches, document the field mappings, and test them through the runtime's extension classes. Use the existing authorized campaign scope to preview and execute; ask only for missing access, unresolved user choices, or a resource/runtime constraint requiring a decision. Neither an upstream adapter merge nor a runtime downgrade is required to use a local bundle. A runtime with unsupported cache semantics still needs an explicit limitation or a separately supported representation. + +### Save the launch inputs in the existing checkpoint + +`--configuration` selects checkpoint keys; the name never implies topology. The command reads each configuration's `draft_request.identity/search/collection` and fills absent facts from shared or per-configuration `inputs`. Conflicting facts are rejected. An incomplete draft is allowed. For inputs without a profile, supply these fields: + +| Checkpoint input | Required meaning | +| --- | --- | +| Identity fields | `model`, `model_revision`, `model_kind`, `framework`, literal `framework_version`, packaged `gpu`, and `interconnect`; optional `sm`. | +| Search fields | Explicit `tensor_parallel`, `attention_data_parallel`, `moe_tensor_parallel`, `moe_expert_parallel`, and `context_length`. This route currently represents PP1/CP1. | +| `model_config` | Local `config.json` path, or `{ "path": "...", "sha256": "..." }`. Relative paths are based on the checkpoint directory. The runtime must load matching config bytes from its actual checkpoint path or local HF cache. A discovered adjacent `hf_quant_config.json` is pinned separately in `model_config.source_files`; its contents are not merged into or written over `config.json`. | +| `precision` | `gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `kv_cache_dtype`, and `comm_quant_mode`; optional selected backends and `enable_wideep`/`enable_eplb`. A complete existing profile can supply this exact precision identity. | +| Collection fields | `max_num_tokens`, `max_batch_size`, `gpu_memory_utilization`, `prefill_cudagraph_policy`, and an explicit capture size only with `explicit` policy. Missing scheduler/fraction values use the documented 8,192/256/0.90 defaults. Fresh inputs use runtime capture policy; saved legacy drafts retain their policy. | +| `collection_deployment` | Existing executor/image/model-cache/mount/transport options. `container_mount` is a list. Explicit deployment flags override and save the selected configurations' deployment choices. | +| `resource_overrides` | Per-configuration metadata assumptions, if required for the model. Cache geometry can remain absent at probe entry. | + +Save source research and outstanding questions after each meaningful finding with `onboard checkpoint`, using its current revision. Register the local bundle files and source notes as input artifacts; do not put independent configurations' byte bounds into shared inputs. Probe and import also register their results, but they do not observe or save the agent's investigation automatically. + +### Choose or author the instrumentation bundle + +Omit `--instrumentation` for the audited bundled vLLM 0.27.0 adapter. For the inspected vLLM 0.28.0/Dynamo build, export the complete [campaign-local example](../python/aisimulate/collector/fpm_forward/runtime/vllm-0.28.0.example.json) into a fresh directory: + +```bash +python - <<'PY' +from pathlib import Path +from collector.fpm_forward import bundled_instrumentation +from collector.fpm_forward.runtime_instrumentation import load_instrumentation, freeze_instrumentation + +assets = Path(bundled_instrumentation.__file__).parent / "runtime" +bundle = load_instrumentation(assets / "vllm-0.28.0.example.json", "0.28.0") +frozen = freeze_instrumentation(bundle, Path("model-onboarding/instrumentation")) +print(frozen.manifest_path) +print(frozen.sha256) +PY +``` + +Run the export with Python from the installed AISimulate environment. In a source checkout, set `PYTHONPATH=python/aisimulate` when running it from the repository root. The example contains runnable worker/scheduler adapters, source notes and exact installed-source hashes. Its [source audit](../python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.28.0.md) identifies the inspected immutable image. Inspect your image's actual source before using it. A mismatch produces unresolved evidence with the observed hashes; investigate the patch and update its mapping/tests before creating a new bundle. Source inspection and CPU fakes establish implementation readiness. Only a successful live probe and independent import establish evidence for your selected campaign. + +For another pin, create a fresh bundle using the same public manifest fields. This abbreviated template illustrates the structure; replace the revision and hashes with inspected values and declare every imported adapter file: + +```yaml +schema_version: aisimulate-runtime-instrumentation/v1 +runtime: + framework: vllm + version: "" + source_revision: "" + source_files: + vllm/v1/worker/gpu_worker.py: "" + # Include cache allocation, retention, backend and scheduler implementations. +files: [observer.py, source-notes.md] +source_notes: source-notes.md +worker_class: observer.ObservedWorker +scheduler_class: observer.ObservedInstrumentedScheduler +observation_schema: aisimulate-runtime-observation/v1 +``` + +Paths are normalized relative paths within the bundle. The manifest and every declared file contribute to its identity. `source_notes` and nonempty installed `source_files` are required for complete import. The runner freezes the bundle for each attempt and injects the classes through the native worker/scheduler extension mechanism. Loading a manifest, previewing launches and importing saved observations never execute the supplied Python. + +Use the [runnable hooks](../python/aisimulate/collector/fpm_forward/runtime/instrumentation/hooks.py) and [field contract](../python/aisimulate/collector/fpm_forward/runtime/instrumentation/README.md) as the starting template. Preserve each native `super` call, its return value and exceptions. Record profiled cache allowance after native profiling, successful allocation after initialization, and physical views only after warm-up/capture returns. The scheduler records its native initial pool and free queue. Observe outside timed forwards; never change the runtime's graph, precision, scheduler or forward behavior to make the observation pass. + +New launches verify the frozen bundle's bytes and import its declared modules from their exact paths, including in spawned workers. For a campaign-local adapter, include `instrumentation_binding: observed_binding()` in every worker and scheduler observation after the native lifecycle hook, importing `observed_binding` from the collector's `fpm_runtime_instrumentation` entrypoint. The bundled observer demonstrates this call. Do not copy the launch's declared binding as evidence: missing receipts or contradictory imported sources leave memory import incomplete. Saved campaigns without this binding remain readable with `import_binding: unverified_legacy`; resuming them does not retroactively verify their imports or rewrite their bundles. + +The runner supplies `AISIMULATE_RUNTIME_CONTEXT`, `AISIMULATE_RUNTIME_INSTRUMENTATION` and `AISIMULATE_RUNTIME_OBSERVATION_DIR`. The context binds attempt, configuration, phase, bundle, exact launch facts and expected ranks. Worker files use `schema_version: aisimulate-runtime-observation/v1`, the same binding, `kind: worker`, DP/TP/PP coordinates, actual runtime/config/hardware evidence, a lifecycle record, raw cache data and `unresolved_fields`. Schedulers use `kind: scheduler`, a DP coordinate and null TP/PP coordinates. They report initial free blocks, null/permanent reservation IDs and group-to-pool mapping. + +A worker's physical view has this shape (illustrative values, not model geometry): + +```json +{ + "storage_id": "storage-0", + "storage_offset_bytes": 0, + "shape": [100, 8, 16, 256], + "stride_bytes": [65536, 8192, 512, 2], + "element_size_bytes": 2, + "block_axis": 0, + "kernel_block_size_tokens": 16 +} +``` + +Retain real shapes and byte strides, including KV-plane-first layouts. Obtain the block axis and kernel block size from the initialized backend; do not guess an axis because its length happens to equal the pool block count. Record backing storages once and every allocation's full storage size, declared `shared_by`, byte offset and packed block stride. Read the runtime spec's padded page size rather than reconstructing an older formula. Keep full/sliding/convolution retention distinct, report unknown specs, and require vLLM 0.28's `extra_retained_tokens` to be explicitly zero for this representation. Hash the actual loaded config file, not the host-side planning copy; config identity does not attest checkpoint weight contents. + +### Preview, execute, import and review + +Use the same deployment options for preview, execution and resume. These commands show two already selected configurations; omit or add configuration keys to match the campaign. The checkpoint stays outside the fresh output directories. + +```bash +aisimulate onboard probe-runtime \ + --checkpoint ./model-onboarding/onboarding-checkpoint.json \ + --configuration worker-a --configuration worker-b \ + --instrumentation ./model-onboarding/instrumentation/manifest.json \ + --executor slurm --image '' \ + --output-dir ./model-onboarding/probe-001 + +# After reviewing the preview, execute within the authorized campaign. +aisimulate onboard probe-runtime \ + --checkpoint ./model-onboarding/onboarding-checkpoint.json \ + --configuration worker-a --configuration worker-b \ + --instrumentation ./model-onboarding/instrumentation/manifest.json \ + --executor slurm --image '' \ + --output-dir ./model-onboarding/probe-001 --execute + +aisimulate onboard import-observations \ + --checkpoint ./model-onboarding/onboarding-checkpoint.json \ + --configuration worker-a --configuration worker-b \ + --observations ./model-onboarding/probe-001/observations.json \ + --output-dir ./model-onboarding/runtime-drafts-001 +``` -Planning and collection reject concurrent onboarding operations. The persistent `.support.lock` file uses an OS advisory lock; ownership is released when the process exits, including after an abrupt termination. Leave the file in place. Request validation, saved onboarding-plan checks, and collector input resolution exit 2. Failures after collector execution starts, including a frozen checkpoint identity mismatch, exit 1 with a concise message. Interruption exits 130. +Use the [existing Slurm or Kubernetes deployment prerequisites](#choose-the-collection-executor), including required mounts and a caller-owned Slurm allocation. Every selected configuration runs the phases required by its role, covering all worker ranks and one scheduler per DP rank; explicit aggregated and role-less historical requests run both prefill and decode. Sampling caps limit diagnostic points; they do not shrink the reviewed context, scheduler limits, memory fraction or capture settings. Use `--resume --execute` with the same inputs to continue an interrupted probe. Instrumentation or runtime-setting changes require a fresh compatible attempt/output; retain failed evidence. -The collector narrows initial prefill sampling with the pilot's input-token and concurrency bounds. With an FPM profile, both the sample batch size and total new-token axis are capped by the same rank-local scheduler limits used in the generated prediction and recommendation configs. DEP does not multiply these limits by the GPU count; coverage of the worker's summed FPM queries still requires validation. Decode uses the collector's existing sampling profile within the declared resource envelope; a four-request synthetic pilot does not imply four timing samples or a short decode campaign. Inspect the generated command and collector plan before committing GPU time. Successful formal collection publishes the FPM Parquet file and metadata pair into the plan's local systems data directory; diagnostic success alone does not provide that pair. +`observations.json` is a runner-owned index of immutable launch manifests, frozen bundles and raw artifacts. CPU import rechecks their hashes, exact identity/build/precision/settings, ranks and lifecycle, actual tensor/storage accounting, retention and pool reservations. It compares compatible phases/ranks and retains their individual values before taking the limiting usable capacity. An observer's `passed` or `status` field cannot authorize a profile. + +Import writes a validation result for every selected configuration and fresh `request.yaml`/`fpm-model-profile.json` drafts for complete configurations. Partial success is preserved and the command exits nonzero for any incomplete configuration. It registers evidence and drafts in the same checkpoint without accepting them or overwriting prior accepted profiles. Present each complete draft for the existing [review/edit/accept workflow](#preserve-partial-profile-review). Preserve observed values and provenance when recording user overrides. Runtime-affecting edits require a fresh compatible probe. Bundle/evidence changes invalidate dependent acceptance and collection references while preserving history. + +After acceptance, continue ordinary planning and formal timing collection. Observed memory/layout does not establish timing availability, query coverage or serving accuracy. Formal collection automatically inherits the imported profile’s verified bundle and launch evidence; no extra public instrumentation flag is needed on `collect-fpm`. Its deployment options must match the accepted probe. After collection, `runtime-compatibility.json` checks the actual runtime, group geometry and usable capacity against that accepted bound. The existing timing publication and accuracy stages still apply. + +| Diagnostic | Investigate and retain | +| --- | --- | +| Missing or duplicate ranks | Actual launched worker/DP counts, process failures and per-rank logs. Never fill a missing rank with another rank's record. | +| Missing warm-up/capture evidence | Native initialization failure, incomplete lifecycle hooks, or an unobserved capture outcome. Requested graph settings alone are insufficient. | +| Source/config hash mismatch | Installed vendor patches, selected backend source, actual model-cache mount and loaded revision. Preserve both hashes and the old attempt. | +| Storage/alias/padding mismatch | Actual untyped storage, byte strides/offsets, backend block axis and runtime `shared_by` declarations. Do not flatten groups or count shared storage twice. | +| Retention/pool mismatch | Unsupported spec, additional retained tokens, reservations, multiple pools, watermark, compressed storage or offload. A positive byte count does not make these semantics representable. | +| Phase/rank graph disagreement | Effective settings and native graph resolution across all ranks and required phases of the same request. Keep the raw records and rerun with the intended serving policy. Separate prefill and decode requests may intentionally use different graph settings. | + +## Finalize runtime memory + +Fresh config-based onboarding does not ask for activation, runtime, communication or aggregate non-KV memory bytes. Optional derived values appear as planning estimates in the preview and provenance. They do not complete the saved memory profile. Supplying all four non-KV byte overrides explicitly retains the complete declared-profile route; a partial override remains recorded but does not establish simulation readiness. Existing complete profiles remain supported. + +Collection observes initialized worker cache allocation and scheduler pool capacity without an extra GPU run. For pending profiles on audited vLLM 0.27.0, both phases use synchronous scheduling because the current prefill benchmark requires it. Review this fixed collection policy in stage 2; it is not an assertion about vLLM serving defaults. The initial observer supports the audited vLLM 0.27.0 full/sliding-attention cache interfaces, including supported convolution storage represented by those interfaces. Historical memory artifacts retain this strict version guard. The [runtime probe and observation import](#resolve-cache-geometry-with-a-runtime-probe) route also accepts source-audited campaign-local adapters for other exact runtime builds; its evidence is checked independently before producing a profile. Unknown layouts report unresolved memory; they do not automatically qualify for finalization. Do not substitute guessed bytes to hide that diagnostic. + +The supported memory evidence requires automatic HBM-only allocation without weight offload, synchronous scheduling, complete graph settings, and an audited unquantized, FP8 or ModelOpt NVFP4 precision mapping. Mixed precision such as NVFP4 routed experts with BF16 attention and shared experts is supported. Unknown precision mappings remain unresolved. Worker and scheduler settings must agree, with recorded worker graph downgrades allowed only for the audited runtime behavior. + +After successful formal collection: + +```bash +aisimulate onboard finalize \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --collection-report ./collection-quality/collection-validation.json \ + --resolved-output-dir ./aisimulate-support-resolved +``` + +`--collection-report` is optional. When supplied, finalization independently recomputes the existing report's mandatory collection gates under its frozen policy, verifies report/policy/source hashes, and binds it to the exact original request and collection directory. The finalization manifest and exported profile provenance carry a `collection_quality` record: + +| `collection_quality.status` | Meaning | +| --- | --- | +| `passed` | All mandatory collection gates pass for the report's evaluated scope and frozen policy. | +| `failed` | At least one mandatory gate failed; memory can still be finalized for exploratory use. | +| `incomplete` | Required collection evidence or assessment remains incomplete; memory can still be finalized if its own evidence is complete. | +| `not_assessed` | No collection report was supplied. A completed memory profile does not imply collection quality. | + +Every state retains `accuracy: not_assessed`; collection quality and matched serving accuracy are separate. Assessed records retain the report/policy identities and individual gate statuses. An unrelated, stale, modified or contradictory report is rejected rather than copied as a trusted status. Keep the referenced evidence available: saved-plan, runtime-profile and replay verification recheck finalized quality, including profiles finalized directly from collection without an earlier runtime probe. Historical finalizations without this marker retain their original representation and remain unassessed; they are not silently qualified. Failed or incomplete quality does not itself disable ordinary exploratory prediction/recommendation. + +The command verifies the immutable onboarding and collector plans, passed cell attempts, native timing artifacts, runtime resource evidence, and formal Parquet/metadata commit. It reconstructs the formal timing rows from their native evidence before copying the pair. It requires compatible cache geometry and resolved memory-relevant launch settings, including CUDA graph configuration, across cells and required phases within the same request. Explicit aggregated and role-less historical requests share one profile across both phases; separate prefill and decode requests retain their own graph settings and memory capacities. For compatible observations it uses the smallest observed request-usable capacity across ranks and cells within the selected role. Source hashes, original graph evidence and the capacity rule are retained in `finalization.json` and profile provenance. + +Saved effective collection options must match the reviewed request, including capture policy and explicit capture sizes. Finalization also rejects an observed immutable model revision or loaded-config commit that contradicts the requested immutable revision. Local PVC model paths and missing revision metadata remain valid; this check does not prove weight contents. + +The fresh directory contains `request.yaml`, `fpm-model-profile.json`, ordinary prediction/recommendation configs and a verified copy of the formal data pair. The original collection is unchanged. Runtime cache groups preserve allocation padding and shared-pool slot costs, including null/permanent reservations, instead of deriving an invented non-KV breakdown. Runtime memory is tied to the recorded scheduler, context, precision, topology and memory policy; changing these settings requires new compatible memory evidence. + +If formal observations show less usable cache capacity than the accepted runtime probe, preserve the original collection and import its complete formal runtime observation index through `onboard import-observations`. Review and explicitly accept the resulting capacity revision in the same configuration's checkpoint. Then finalize with both requests: + +```bash +aisimulate onboard finalize \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --memory-config ./reviewed-capacity-revision/request.yaml \ + --resolved-output-dir ./aisimulate-support-resolved +``` + +`--memory-config` permits only a lower cache-capacity value, equal to the verified formal minimum, and updated provenance. The complete model, precision, topology, cache geometry, runtime, graph and collection settings must remain unchanged. The revision must derive from the exact successful formal cell attempts and be explicitly accepted in its referenced checkpoint; guessed bounds, capacity overrides, unrelated probes and changed execution settings are rejected. The command retains the original collection identity and binds the revision request by path, SHA-256 and request identity in the new finalization provenance. Keep the revision request and its evidence immutable and available for subsequent verification. This correction reuses the measured timing pair and requires no formal GPU recollection. Continue `validate-collection` against the original request and collection directory. + +Add `--collection-report ./collection-quality/collection-validation.json` to that command to retain quality on a capacity revision too. The report remains bound to the original timing request and collection, while `--memory-config` identifies the separately accepted memory revision. + +Review the resolved profile and assumptions with the user. Finalization does not accept it in the session checkpoint. Save its request as a new draft, record the new artifact paths and source relationship, and obtain acceptance of the exact resolved profile with the existing checkpoint workflow. Preserve the original collection references. Only then continue simulation using `./aisimulate-support-resolved/request.yaml` and `--output-dir ./aisimulate-support-resolved`. A pending profile produces a clear error before prediction/recommendation cache sizing or replay validation. + +Resuming `onboarding-checkpoint.json` restores the onboarding conversation and accepted inputs; it does not migrate collection artifacts. Current collection and finalization use a schema-11 collection plan and schema-7 formal FPM publication. The collector can still verify schema-10 plan hashes and read their matching native timing and memory evidence, preserving historical cell and attempt identities. Historical schema-6 formal publications cannot be finalized by this collector, and automatic migration is unsupported. Keep those artifacts unchanged and use fresh output directories for any new collection. + +## Validate collection and serving accuracy + +The standard procedure has four distinct checks. Collection validity verifies each native point and its published row, plus observed attention groups, initialized graph configuration and KV initialization regime. Repeatability checks independent repetitions of the complete native grid under the same frozen execution protocol. Bounded subsets remain diagnostics until their agreement with full-grid execution is established. Withheld-coordinate validation checks native direct interpolation against measurements excluded from its input table. Matched serving checks the selected workload separately. Fixed weight/KV precision, a complete grid, a measured-point lookup or successful replay alone cannot establish accuracy. + +Use separate validation directories for each precision/topology campaign. The original request, collection plan, raw results and formal pair remain unchanged. Save the following JSON as `validation-policy.json`, or supply equivalent YAML; omitted fields use these defaults and unknown fields are rejected: + +```json +{ + "schema_version": "aisimulate-onboarding-validation-policy/v2", + "repeatability": { + "comparison_mode": "full_grid", + "samples": 5, + "max_points_per_cell": 12, + "max_attempts_per_sample": 2, + "max_cv": 0.05, + "max_source_relative_difference": 0.05 + }, + "interpolation": { + "max_points_per_phase": 16, + "seed": 42, + "max_p95_relative_error": 0.20, + "max_unsupported_fraction": 0.0 + }, + "serving": { + "max_latency_p95_relative_error": 0.20, + "max_throughput_relative_error": 0.20, + "max_dispatch_delay_ms": 5.0, + "max_dispatch_delay_fraction": 0.01 + } +} +``` + +| Property | Meaning | +| --- | --- | +| `repeatability.comparison_mode` | `full_grid` repeats the source native sweep with fresh engines. `bounded` selects a diagnostic subset and cannot qualify the full-grid source. | +| `repeatability.samples` | Fresh, independently launched measurements per coordinate, at least 2. The source measurement is retained separately. | +| `repeatability.max_points_per_cell` | In `bounded` mode only, maximum representative coordinates per phase cell, at least 1. A budget that drops required observed strata is rejected. | +| `repeatability.max_attempts_per_sample` | Total attempts permitted per independent sample, including the first launch. Failed attempts and valid slow measurements remain recorded. | +| `repeatability.max_cv` | Maximum sample standard deviation divided by sample mean across all valid independent estimates in a comparable population. Adjacent internal decode steps are not independent samples. | +| `repeatability.max_source_relative_difference` | Maximum `abs(fresh median - source) / source` at each coordinate. Applies only when source and fresh full-grid evidence are comparable; distinct from CV. | +| `interpolation.max_points_per_phase` | Maximum withheld coordinates for each phase, at least 1; selection retains boundary anchors and a viable retained measurement table. | +| `interpolation.seed` | Nonnegative deterministic selection seed. | +| `interpolation.max_p95_relative_error` | Maximum p95 absolute relative forward-time error, separately for prefill and decode. | +| `interpolation.max_unsupported_fraction` | Maximum unsupported fraction among selected holdouts, from 0 to 1. Unsupported queries are never counted as successful numerical predictions. | +| `serving.max_latency_p95_relative_error` | Maximum p95 absolute relative TTFT and per-request mean inter-token-latency errors, assessed separately; also applies to supplied forward samples. | +| `serving.max_throughput_relative_error` | Maximum absolute relative error of total output tokens divided by elapsed time from first dispatch to last response, including inter-turn delays. | +| `serving.max_dispatch_delay_ms` | Absolute dispatch-timing tolerance in milliseconds, default 5.0. | +| `serving.max_dispatch_delay_fraction` | Relative dispatch-timing tolerance, default 0.01 of the required inter-turn gap. Allowed jitter is the larger of this bound and the absolute bound. This does not relax trace timing-policy compatibility. | + +Ratios are dimensionless: `0.20` means 20%. Error is `abs(predicted - measured) / measured`; percentiles use linear empirical quantiles. Timing reports identify their units. These initial criteria are editable and are not statistical confidence bounds or universal accuracy guarantees. + +### Stage 5: collection quality + +Freeze the policy and native grid, inspect source validity/execution and run the CPU holdout evaluator: + +```bash +aisimulate onboard validate-collection \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --validation-output-dir ./collection-quality \ + --policy ./validation-policy.json +``` + +This writes `policy.json`, `campaign.json`, `repeatability-selection.json`, `holdout/` and `collection-validation.json`. It does not launch GPU work. Its initial overall status normally remains `incomplete` until repeats are collected. Failed numerical checks return nonzero even during preparation. To execute the frozen repeats with the archived collector deployment, use the same Kubernetes environment or caller-owned Slurm allocation as the source campaign: + +```bash +aisimulate onboard validate-collection \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --validation-output-dir ./collection-quality \ + --resume --execute +``` + +Each full-grid repetition launches a fresh engine through the existing executor with the source collection options and separate raw artifacts, attempt IDs and runtime run IDs. Dynamo still determines and admits the native points. Repetitions preserve the source seed, warm-up, KV preparation, runtime/graph settings and allocation policy; changing those settings requires a fresh source campaign. Global warm-up and KV preparation remain outside measured forward time. The per-launch estimate remains the maximum across ranks, followed by a median across independent launches. The derived aggregate is separate evidence and never overwrites the formal table. + +Comparability is checked before numeric stability: inspect the per-rank injected-prompt hashes, runtime pins, expected graph/capture and KV regimes, preparation settings and native sweep history. Compare the same DP rank across launches; different ranks may legitimately use different prompts. Missing or mismatched evidence is reported explicitly. Prompt hashes do not observe sampled continuation tokens, KV/recurrent-state tensors or prove execution-history equivalence. A valid slow measurement remains part of its population; latency alone is not a reason to declare it incomparable. + +Fresh assessments freeze `observation_evidence_version: 2` separately from the validation-policy version. Requested scheduling records distinguish an explicit generated flag, a runtime default and an unreported launch. An omitted async-scheduling flag does not imply either effective Boolean value. Inspect initialized native scheduler and worker observations; missing scheduling remains unestablished, contradictory observations reject, and synchronous/asynchronous runs cannot be pooled as comparable. + +Evidence v2 also checks the producer's per-sample graph mode and token padding, process-local forward indices and timing intervals, and the warm-up/preparation ledger before comparing dispatch, timing basis and recorded preparation. Missing producer execution evidence leaves fresh comparability unestablished and cannot qualify the source table; matching requested graphs or stable latencies do not replace it. Absolute timestamps and latency differences are not equality criteria. These observations describe scheduler wall time, not CUDA-event GPU execution time, GC activity or CPU stall causes. They do not establish the complete graph descriptor, capture-versus-replay, attention kernel, equal KV tensors, equivalent graph/kernel/cache warming, or decode real-KV chain warm-up history. Retain those limits and all valid slow samples. + +Exclude a workload prediction comparison from the onboarding review scope for DP imbalance only when actual per-rank workload evidence, including expert-load evidence where relevant, establishes imbalanced work in that comparison. DP topology or rank latency spread alone is insufficient. Record the excluded comparison, supporting evidence and resulting qualification scope in the checkpoint and review, preserving all raw samples. The CLI has no imbalance-exclusion feature; this review decision leaves measurement gates active, retains valid slow samples in their populations and does not relax CV thresholds or automatically qualify failed or incomplete results. + +Inspect individual attempts and raw internal samples, the median, sample count, mean, minimum/maximum, sample standard deviation and CV in `repeatability-assessment.json`. CV uses every valid independent estimate in the comparable population, including slow samples. Adjacent decode steps and consolidated provenance duplicates never count as independent launches. The source-inclusive CV is retained as a diagnostic, not blindly pooled into the fresh population. Qualification of the published source additionally requires comparable full-grid evidence and source-versus-median agreement within `max_source_relative_difference`; stable new measurements cannot qualify an incomparable or outlying historical source. Keep the default 5% CV bound while investigating instability. + +For a bounded diagnostic, set `comparison_mode` to `bounded` in a new policy and use a fresh output directory. Its subset covers small/large batches, short/long KV contexts and selected graph boundaries, but executes a different sweep history. It cannot pass the full-grid quality gate even if its own measurements are stable. Preserve both populations and establish their empirical agreement before considering a different validation protocol. Failed attempts remain recorded; `--retry-failed` requires explicit resume and stays within the frozen `max_attempts_per_sample` budget. A successful slow sample is never replaced by a faster retry. + +The source must retain a schema-v11 collector plan, successful cell attempts, formal publication and archived `generator-overrides.json`. New GPU launches require the source collector revision; a different checkout must not silently remeasure old data under new behavior. Historical artifacts remain readable, but missing required observations, measurement-protocol metadata or deployment evidence prevent a fresh assessment from qualifying them. Saved assessments without `observation_evidence_version` retain evidence v1 and its earlier, narrower observation scope on resume, recheck and compatible repeat reuse; they do not acquire the v2 checks or new evidence retroactively. Explicit validation policy v1 also retains its bounded-subset meaning and cannot qualify under policy v2. Preserve frozen historical outputs and use a fresh validation directory for a new protocol. Expected point modes and initialized graph configuration alone remain distinct from observed per-sample dispatch. + +For a source with a frozen Slurm CPU policy, repeats retain that exact node-task CPU request and binding. Review CPU affinity evidence for each fresh attempt alongside the timing spread; adequate requested CPUs alone do not establish the worker and scheduler masks. A changed CPU policy is a new collection experiment, not a repeat of the old launch. + +The holdout evaluator removes every physical row at each selected coordinate, preserves envelope anchors, and uses the canonical direct estimator with denied fallback and no analytical model class. It also retains the adjacent measured points on both sides of an execution-mode change along each batch, token or KV curve. For example, if native evidence labels 512 tokens `PIECEWISE` and 513 tokens `NONE`, both measurements stay in the training table. The boundary comes from the full validated native grid and resolved graph configuration for the same collection attempt, not a fixed token limit, a timing jump or the small repeatability subset. Inspect `capture_boundaries` in the holdout plan/report for retained anchors, source hashes and incomplete evidence. Missing or unknown mode metadata leaves boundary assessment incomplete; contradictory rank/configuration evidence is rejected. Expected point modes are not an observed per-call dispatch trace. + +This is an onboarding holdout-selection rule. It does not modify interpolation in ordinary predict, recommend or replay, prove that every interpolation bracket is within one execution regime, or assess behavior when a required boundary measurement is missing. Changes within one graph mode, including capture-bucket or kernel transitions, remain unclassified. Numerical failures and unsupported counts remain separate for prefill/decode. Sparse tables without viable holdouts remain incomplete. References remain the original single-sample timings: repeatability evidence helps diagnose noise but does not subtract measurement noise from interpolation error. + +To change thresholds, edit a copy of the policy and create a new assessment without launching measurements: + +```bash +aisimulate onboard validate-collection \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --validation-output-dir ./collection-quality-reassessed \ + --policy ./revised-validation-policy.json \ + --repeatability-dir ./collection-quality/repeatability +``` + +Raw samples and source collection are revalidated before reuse. A changed repeat count, comparison mode, bounded subset budget or attempt budget requires a fresh measurement campaign; a changed holdout selection can reuse the formal table and rerun CPU evaluation. `--resume` requires unchanged policy, source evidence and holdout selection version. Threshold-only reassessment preserves the original numerical result and all raw attempts; changing a threshold is an explicit policy decision, not a retry strategy. Assessments created before execution-boundary retention require a fresh validation output directory; keep their original failures and reports. The revised selection evaluates a narrower scope, so a changed result is not proof that core interpolation or serving accuracy improved. Preserve prior assessments rather than editing their saved policy or status fields. + +### Stage 6: matched serving + +After memory finalization, run [ordinary replay coverage](#validate-fpm-query-coverage-with-agentx-replay) for the selected local play. Keep that replay report separate from collection quality. + +Before preparation, install [AIPerf at the accepted immutable revision](https://github.com/ai-dynamo/aiperf/tree/7db2ba37a62aa80c882bc90eaf61cc8073e2387b) in a separate environment from AISimulate: + +```bash +uv venv --python 3.13 ./aiperf-validation-env +uv pip install --python ./aiperf-validation-env/bin/python \ + 'aiperf @ git+https://github.com/ai-dynamo/aiperf.git@7db2ba37a62aa80c882bc90eaf61cc8073e2387b' +./aiperf-validation-env/bin/python -m aiperf --version +``` + +This revision supports Python 3.11–3.13 and installs its tokenizer dependencies, including Transformers, SentencePiece, tiktoken and protobuf. Supply the matching checkpoint's local tokenizer files and chat template with `--tokenizer`; preparation loads them locally without remote code. The installation must retain the exact Git repository and commit in its PEP 610 `direct_url.json` metadata, which preparation and execution verify. A package-index installation with the same version number does not establish that source identity. + +Prepare a matched serving recipe using that environment: + +```bash +aisimulate onboard validate-serving --action prepare \ + --collection-report ./collection-quality/collection-validation.json \ + --replay-report ./aisimulate-validation/agentx/validation.json \ + --validation-output-dir ./serving-prepared \ + --endpoint http://127.0.0.1:8000 \ + --tokenizer /path/to/pinned-tokenizer \ + --aiperf-python ./aiperf-validation-env/bin/python +``` + +Preparation rechecks collection, repeatability, holdout and replay evidence. It derives expected execution from the source collection; it never substitutes requested defaults for unreported attention/graph/runtime observations. Differences across phase/rank settings remain unresolved rather than choosing one arbitrarily. It accepts the original request or a memory-resolved request with verified finalization provenance and identical formal data. Missing prediction/coverage hashes in an older replay report require rerunning `validate-fpm`; the old report remains a coverage diagnostic. + +Collection's phase-specific prefix-caching flag is retained as evidence but excluded from this uniformity comparison because it implements benchmark KV preparation. The actual serving flag must independently match the replay's prefix-reuse policy; this exception does not waive the serving cache check. + +This iteration requires **one complete single-stream Weka play**. Additionally, all predecessor recorded API times must be zero/absent, or all inter-turn idle gaps must be zero, so the pinned AIPerf load semantics match the existing replay. Plays with both positive predecessor API times and positive idle gaps, branching plays and multi-play corpora remain explicitly incomplete. Select a naturally compatible complete play; do not flatten dependencies or edit timestamps to force a pass. Broader workload coverage requires additional matched validation. + +Trace token lengths alone do not include the target checkpoint's chat-template overhead. Preparation uses the pinned AIPerf environment and local target tokenizer to freeze actual chat payloads, tokenize them with that template, and create a separate explicit `agentic_mooncake` trace with exact target token identities and one-token blocks. This preserves prefix equality at the engine's cache-block boundaries even when the source Weka blocks are larger. It preserves the original play's request identities, serial dependencies, arrival bounds, inter-turn gaps and requested output lengths. It then runs ordinary strict direct-FPM prediction on this derived trace with the same engine/profile/formal table. `matched_workload` in the recipe records these payloads, tokenization and derived prediction/coverage artifacts; the original Weka trace and its broader replay report remain unchanged. A new target shape without FPM coverage leaves preparation incomplete. + +`serving/recipe.json` contains the pinned AIPerf revision and inspectable command, plus a `serving_observer` section with copied module hashes, environment variables and `--worker-cls` arguments. Start a fresh ordinary serving worker using the source model/runtime/topology/precision/graph settings and this observation setup. Make its evidence directory accessible to the validation process. The caller manages this server; preparing a recipe does not launch it, and Dynamo self-benchmark does not act as the serving benchmark. The observer checks supported actual worker ranks, attention groups, graphs and initialized runtime settings. A source-pinned image digest comes from the frozen launch; it is not a separate container attestation. + +Run traffic explicitly with the pinned AIPerf installation: + +```bash +aisimulate onboard validate-serving --action run \ + --recipe ./serving-prepared/serving/recipe.json \ + --aiperf-python ./aiperf-validation-env/bin/python +``` + +The producer verifies its pin, checks the server's `/tokenize` responses against the frozen target token IDs/template, sends the frozen payloads through the native Mooncake raw-payload loader, and preserves native per-request records and execution receipts. A different server tokenizer/template cannot pass by reporting the same total token count. The endpoint must expose the matching vLLM tokenization API. Fill the generated `observed-execution.template.json` from actual serving initialization and scope evidence, preserving source paths/hashes; expected values are not observations. Then assess in another fresh directory: + +```bash +aisimulate onboard validate-serving --action assess \ + --collection-report ./collection-quality/collection-validation.json \ + --replay-report ./aisimulate-validation/agentx/validation.json \ + --recipe ./serving-prepared/serving/recipe.json \ + --execution-evidence ./serving-prepared/serving/observed-execution.json \ + --validation-output-dir ./onboarding-assessment +``` + +Add `--forward-evidence PATH` when matched instrumented forward timings are available. Their predictions are computed by the native direct estimator; they are independent from TTFT, TPOT and throughput metrics. Without them, forward validation is `unavailable`, and any successful claim is limited to the reported client metrics. Missing required runtime/request evidence remains incomplete; numerical or execution mismatches fail. Supplying forward evidence that fails also prevents qualification. Serving-only threshold changes can reuse the recipe and raw measurements by passing an edited `--policy` to a fresh assessment. + +The final `validation.json` reports collection, replay and serving gates independently. Only `status: passed` with `accuracy: qualified_for_evaluated_scope` qualifies the evaluated coordinates/configuration/play. The command independently recomputes assessments and verifies source hashes rather than trusting saved `passed` flags. It makes no claim for untested concurrency, mixed batches, other plays or other precision/topology cells. Ordinary prediction/recommendation remains available while accuracy is unqualified. + +### Save validation progress in the session checkpoint + +Use the existing single checkpoint and its revision checks. Add validation-only settings under the relevant configuration's `validation_inputs`, and artifacts with `scope: validation`; do not change collection `inputs` for threshold or trace edits. For example, merge this patch after the referenced reports exist: + +```json +{ + "configurations": { + "tp4": { + "validation_inputs": { + "policy": "/absolute/path/collection-quality/policy.json", + "repeatability": { + "comparison_mode": "full_grid", + "samples": 5, + "max_attempts_per_sample": 2, + "max_cv": 0.05, + "max_source_relative_difference": 0.05 + }, + "trace": "/absolute/path/selected-play.jsonl" + }, + "artifacts": { + "validation_policy": {"path": "/absolute/path/collection-quality/policy.json", "scope": "validation"}, + "collection_quality": {"path": "/absolute/path/collection-quality/collection-validation.json", "scope": "validation"}, + "onboarding_validation": {"path": "/absolute/path/onboarding-assessment/validation.json", "scope": "validation"} + } + } + } +} +``` + +Apply it with `aisimulate onboard checkpoint --file PATH --expect-revision CURRENT --update PATCH.json`. Archive superseded validation references with a reason and register replacement paths, as in the [checkpoint workflow](#checkpoint-and-resume-an-onboarding-session). Policy/trace edits invalidate assessments, not accepted profiles or valid collected data. Save the frozen repeatability policy, source protocol/seed/preparation and selected coordinates before execution; the checkpoint records context and does not execute or apply these settings. After resuming a partial repeatability run, replace its changed report reference before recording completion. Saved stage labels never substitute for rechecking machine evidence. + +## Validate FPM query coverage with AgentX replay + +After verifying a matching formal FPM pair and resolving memory, run the existing cold aggregated replay path against a local Weka JSON or JSONL file. For pending collection profiles, use the fresh directory from `onboard finalize`: + +```bash +aisimulate onboard validate-fpm \ + --config ./aisimulate-support-resolved/request.yaml \ + --output-dir ./aisimulate-support-resolved \ + --trace /path/to/traces.jsonl \ + --validation-output-dir ./aisimulate-validation/agentx +``` + +The output directory must be separate from, and neither inside nor above, the collection directory. Use a fresh directory or `--overwrite` to replace prior validation outputs. The command verifies the saved collection plan, writes `predict.yaml`, and invokes ordinary `aisimulate predict` with per-request evidence. It does not download traces or run GPU collection. The reviewed target model, topology, runtime, profile, scheduler limits and systems paths remain fixed; source model labels in a trace do not replace the configured target model. + +The current validation scope is **cold aggregated replay, one client lane, HBM-only cache and no speculative decoding**. The cache starts cold. Linear profiles can accumulate normal prefix reuse; grouped profiles retain `prefix_caching: false` and require vLLM with PP1/CP1. Nested timestamps are interpreted relative to the root play (`nested_timestamp_basis: absolute`). The command replays the complete supplied file; selecting one complete play is useful for a first check, but does not establish coverage of a larger corpus. Preserve the full play and its dependencies when preparing a subset. Ordinary prediction's host-memory checks still apply; a larger corpus can require a larger host, and a preflight rejection remains incomplete validation. The current path builds on existing vLLM/SGLang replay support; onboarding collection remains vLLM. Integrating seeded cache snapshots, explicit warmup and expanded replay modes from the merged AgentX changes, including [#207](https://github.com/ai-dynamo/aisimulate/pull/207) and [#235](https://github.com/ai-dynamo/aisimulate/pull/235), remains follow-up work for onboarding; those modes are not prerequisites for this workflow. + +The reference corpus is [semianalysisai/cc-traces-weka-062126-256k](https://huggingface.co/datasets/semianalysisai/cc-traces-weka-062126-256k/tree/8fecd2fc56694469f758f0afbbb6335ad3043740). To download that pinned revision separately, with Hugging Face access available: + +```bash +hf download semianalysisai/cc-traces-weka-062126-256k traces.jsonl \ + --type dataset \ + --revision 8fecd2fc56694469f758f0afbbb6335ad3043740 \ + --local-dir ./agentx-reference +``` + +The expected `traces.jsonl` SHA-256 is `e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef`. Record the file hash and selection when using a derived subset. No trace corpus is bundled or fetched automatically by AISimulate. + +### Read the coverage result + +| Artifact | Meaning | +| --- | --- | +| `predict.yaml` | Inspectable ordinary prediction input with strict direct FPM and native coverage collection enabled. | +| `validation.json` | Overall `covered` or `incomplete` result, issues, collection identity, trace/config/FPM file hashes and replay scope. Accuracy remains `not_assessed`. | +| `prediction/fpm-coverage.json` | Native lookup counts and bounded missing-query details, also saved when a timing lookup fails after replay starts. | +| `prediction/prediction.json` | Ordinary prediction report when produced, including per-request completion, AgentX graph and target-model projection. | + +`measured` means the native lookup resolved from an exact measured point; `interpolated` means it used supported direct interpolation; `unsupported` means the native query failed. Counts are **native lookup resolutions**, not requests, replay iterations or tokens. Reusing a cached timing does not make another lookup. Each role records prefill, decode and any mixed-pass decode-baseline counts. Missing entries identify phase, purpose, cell identity, coordinates, reason and occurrence count. Each model retains at most 128 distinct gaps and reports omitted gap occurrences; its total unsupported count remains complete. + +An exit status of 0 and overall `validation.json` status `covered` require a nonempty completed replay, all selected requests and plays completed with their requested output, and at least one supported native query with no unsupported queries. An early failure, failed or aborted request, empty query set or interruption cannot pass. Inspect `validation.json` first: an individual native snapshot does not certify replay completion. A missing query stops ordinary prediction and returns nonzero while retaining available evidence. That partial report describes only the attempted portion; it does not identify every missing point in the remainder of the corpus. + +For a missing decode query, inspect its exact batch size and total past-KV coordinate together. A high maximum KV value elsewhere in a table does not prove coverage for batch size 1 at that context, or for another batch/capture region. Prefill likewise needs the requested batch and prompt/KV support. Use these gaps to review the collection envelope and matching data, preserving strict direct interpolation and denied fallback. Successful coverage establishes that the selected replay could obtain timings; it does not establish predictive accuracy. Assess accuracy separately with matched silicon measurements. ## Run the generated ordinary configurations -After the selected model execution route is available and formal data collection is complete: +The generated synthetic configurations are optional small prediction/recommendation examples after the selected model route and formal data are available: ```bash aisimulate predict \ @@ -143,53 +1317,21 @@ aisimulate recommend \ --output-dir ./aisimulate-support/recommend-results/pilot ``` -With two candidates, run both replica counts. This loop derives the candidate names from the validated saved request and constructs fixed `aisimulate recommend` commands: - -```bash -python3 - <<'PY' -import shlex -import subprocess -from pathlib import Path - -from aisimulate.support.plan import check_plan -from aisimulate.support.schema import SupportRequest - -root = Path("./aisimulate-support").resolve() -request = SupportRequest.from_yaml(root / "request.yaml") -check_plan(request, root) -max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.worker_gpus) -replicas = list(dict.fromkeys((1, max_replicas)))[:request.search.max_candidates] -statuses = [] -for count in replicas: - name = "pilot" if count == 1 else f"replicas-{count}" - command = [ - "aisimulate", "recommend", - "--config", str(root / f"recommend/{name}.yaml"), - "--output-dir", str(root / f"recommend-results/{name}"), - "--format", "json", - ] - result = subprocess.run(command, check=False) - statuses.append(result.returncode) - print(f"exit {result.returncode}: {shlex.join(command)}", flush=True) -raise SystemExit(1 if any(statuses) else 0) -PY -``` - -The loop reports each command's exit status and attempts all commands, including when the pilot finds no feasible candidate. It exits 1 after all attempts if any command returned a nonzero status, otherwise 0. Successful outputs remain usable even when the loop exits 1; inspect each result before comparing candidates. - -Run either the individual recommendation command or the loop against fresh result directories. The loop also handles a one-candidate plan and does not execute entries from `commands.json`. Results remain separate under `recommend-results/pilot` and, when present, `recommend-results/replicas-N`; compare their objective and latency results for the same workload. This plan does not produce a combined ranking or search additional TP sizes or scheduler settings. +Use fresh result directories and inspect each command's exit status and results. These examples validate one selected worker and workload. For an actual deployment run, copy the generated config and set `engine.workers.aggregated.parallelism.replicas` in prediction, or the replica values in recommendation's `parallelism.preset` and `optimization.constraints.max_candidate_gpus` to your intended deployment budget. Keep the exact worker topology and profile consistent with the measured data. Those runtime choices do not require changing the onboarding request or recollecting the same worker's timings. The generated configurations select `engine.workers.aggregated.timing.estimation_mode: fpm_interpolation` with `fallback_policy: deny`. With a profile, they also set `estimator_config.fpm_interpolation.method: direct`. `engine.systems_paths` contains the plan's absolute local systems directory, which supplies hardware and collected FPM data. Recommendation preserves the resolved root, interpolation method, and complete profile in exported prediction configs. Moving the plan to another machine requires updating absolute paths or regenerating it there. Existing configurations using the single-root `engine.systems_path` input are normalized to `engine.systems_paths` when saved. Both spellings resolve relative local paths against the working directory when the configuration is loaded, so saving and reloading from another directory preserves the selected roots. The canonical list's `default` entry continues to select packaged data. A worker’s `timing.systems_paths` overrides engine-level roots; performance metadata records those effective roots. Empty or whitespace-only roots are rejected by the shared root type in both prediction and recommendation configurations. These configurations target AISimulate's standalone `predict` and `recommend` commands. Dynamo's replay Planner AIC session adapter does not currently forward custom systems roots; using that downstream path requires a separate adapter update. +Grouped configurations also set `engine.workers.aggregated.kv_cache.prefix_caching: false`; preserve it in prediction, recommendation and `validate-fpm` replay. The same reviewed groups and byte budget reach native allocation in each path. Changing cache layout does not introduce a new collection grid or change the FPM timing-table coordinates: use the [shared AISimulate/Dynamo collection policy](#how-the-collection-grid-is-determined) and validate the resulting logical-context queries. + Ordinary `predict` and `recommend` retain their existing defaults when no profile is supplied. Generated configs can be edited through the public schema. A recommendation with a profile enumerates only its declared deployment tuples that fit the GPU budget; it preserves the complete inline profile and interpolation choice in exported prediction configs. Default scheduler search ranges may exceed a profile's envelope, so pin or bound those domains explicitly. A successful simulation is not an accuracy result. Compare its output with an independent run of the same model, runtime, topology, and workload to assess accuracy. ## Choose the model execution route Hand off the saved request, pinned model configuration, and plan. Both routes need a canonical checkpoint identity, effective precision and topology, correct weight and KV-cache accounting, and matching whole-forward FPM measurements. Collected timings alone do not establish memory fit. -- **Registered-model/SOL route:** reuse a compatible analytical class or follow [How to Add a New Model](../python/aisimulate/docs/add_a_new_model.md) when choosing to add one. Verify its operation graph, memory/cache accounting, and native FPM SOL execution. +- **Registered-model/SOL route:** reuse a compatible analytical class or follow the optional [model-integration procedure and CPU checks](../python/aisimulate/docs/fpm/model-integration.md#2-registered-model-route-reuse-or-implement-the-model-description) when choosing to add one. Verify its operation graph, memory/cache accounting, and native FPM SOL execution. - **Class-independent direct route:** supply the identity/resource profile and configure the worker as shown below. The guided planner selects this route whenever a profile is supplied. No operation graph is constructed for resources, timing, or recommendation candidates. ```yaml @@ -215,3 +1357,19 @@ Saved legacy interpolation controls merge with explicitly supplied correction, r Per-operation silicon profiling described in the model guide is not required by either FPM route. The workflow collects whole-forward timings, then verifies prediction and recommendation for the exact target deployment. Report timing coverage and interpolation error separately; successful simulation alone does not establish measured accuracy. Accuracy gates and interpolation corrections remain deferred until new measurements are collected with the provenance from [Dynamo #15110](https://github.com/ai-dynamo/dynamo/pull/15110). Those records support investigation; they do not by themselves establish the cause of earlier interpolation errors. If a reproducible CUDA graph boundary causes the discrepancy, interpolation must respect that boundary. No graph-aware partitioning or accuracy policy is introduced here. + +### Profile boundary checks + +Scripted onboarding requires `--context-length` unless a model config or FPM +profile supplies the context limit; 256000 is a cap, not an assumed model limit. +Aggregated collection inherits the profile's shared scheduler token limit. If +selected deployments have different limits, supply a shared +`--fpm-max-num-batched-tokens` within every deployment's envelope. + +Grouped-cache replay requires a runner that advertises grouped-cache support; +use `--stack engine`. The guard selects the requested deployment for both canonical +and flat engine arguments, so unrelated grouped entries do not block linear replay. +Recommendation skips deployment modes for which the +profile has no matching role deployments. Finalized collection quality remains +verified when checking a plan; rendering collection arguments without a runtime +probe does not repeat that verification. diff --git a/docs/sweeper/traffic.md b/docs/sweeper/traffic.md index f1ba2270e..5e65ea4bf 100644 --- a/docs/sweeper/traffic.md +++ b/docs/sweeper/traffic.md @@ -108,8 +108,11 @@ temporary KV pressure or request retraction. Batching combinations that leave no are reported to Vizier as infeasible before replay. Every resulting candidate records `kv_load_ratio`, the derived `concurrency`, -`kv_load_concurrency_capacity`, `kv_load_capacity_tokens`, and per-role -`*_kv_capacity_tokens` for traceability. +`kv_load_concurrency_capacity` for traceability. Linear caches also record +`kv_load_capacity_tokens` and per-role `*_kv_capacity_tokens`. These token fields +are absent for grouped caches, which instead record `kv_load_capacity_bytes`, +`kv_load_request_cache_bytes`. Consumers must treat token-capacity fields as optional +and use the byte-based fields for grouped layouts. ## `num_request_ratio` (synthetic length scales with the load) diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index fa727a448..07aa9ab91 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -157,6 +157,29 @@ https://gitlab-master.nvidia.com/dl/sglang/sglang/-/blob/02c5a855aceb968c310e6fb ## vLLM +The grouped-cache config derivation and metadata regression tests in +`src/aisimulate/support/config_profile.py`, +`tests/unit/test_support_config_profile.py`, and +`tests/unit/test_support_topology.py` are adapted and modified from vLLM's +Inkling cache layout, decoder state construction, and configuration defaults at +commit `98dff2a81d747d1dba01a47f939f48c3526d4206`: + +- [`vllm/models/inkling/nvidia/sconv_swa_attn.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py) +- [`vllm/models/inkling/nvidia/model.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/model.py) +- [`vllm/models/inkling/configs.py`](https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/configs.py) + +Copyright contributors to the vLLM project. + +This material is licensed under Apache-2.0; the upstream license at the +identified revision is available at: +https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/LICENSE + +AISimulate modifies these contracts for metadata-only rank-local aggregate +cache sizing and validation, without constructing runtime tensors or a model. +The derived tests use original synthetic geometries and do not copy a +checkpoint fixture. The derived files retain the upstream copyright and +identify the modified source material. + The repository-root files `crates/core/src/engine/scheduler/vllm/core.rs`, `crates/core/src/engine/kv_manager/{state_cache_manager,vllm_backend}.rs`, `crates/core/src/engine/cache/vllm_block_pool.rs`, diff --git a/python/aisimulate/collector/fpm_forward/bundled_instrumentation.py b/python/aisimulate/collector/fpm_forward/bundled_instrumentation.py new file mode 100644 index 000000000..07255887d --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/bundled_instrumentation.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Select the audited bundle; other pins use an explicit campaign-local manifest.""" + +from pathlib import Path + +from .runtime_instrumentation import InstrumentationBundle, load_instrumentation + + +def bundled_instrumentation(version: str) -> InstrumentationBundle | None: + if version != "0.27.0": + return None + return load_instrumentation(Path(__file__).parent / "runtime" / "vllm-0.27.0.json", version) diff --git a/python/aisimulate/collector/fpm_forward/cli.py b/python/aisimulate/collector/fpm_forward/cli.py index 7096f5933..0f226eaf8 100644 --- a/python/aisimulate/collector/fpm_forward/cli.py +++ b/python/aisimulate/collector/fpm_forward/cli.py @@ -14,7 +14,7 @@ from collector.model_cases import build_collection_case_plan from .config import add_fpm_arguments, add_fpm_generator_arguments -from .entry import resolve_inputs, resolve_run_inputs, run_resolved +from .entry import _load_generator_overrides, resolve_inputs, resolve_run_inputs, run_resolved _INPUT_ERRORS = (OSError, RuntimeError, TypeError, ValueError) @@ -28,7 +28,7 @@ def _parser() -> argparse.ArgumentParser: parser.add_argument("--model-path", default=None) parser.add_argument("--model-architecture", default=None) parser.add_argument("--model-cases", default=None, help="Optional model cases YAML path.") - parser.add_argument("--gpu", required=True, help="Target AIC system, for example b200_sxm.") + parser.add_argument("--gpu", default=None, help="Target AIC system, required for a new campaign.") parser.add_argument("--sm", type=int, default=None, help="Optional explicit SM version for case planning.") parser.add_argument("--plan-only", action="store_true", help="Print the frozen FPM plan and exit.") parser.add_argument("--smoke", action="store_true", help="Run the minimal smoke sampling profile.") @@ -40,6 +40,29 @@ def _parser() -> argparse.ArgumentParser: help="Retry failed cells while resuming; requires --resume.", ) parser.add_argument("--checkpoint-dir", default=".collector_checkpoint") + validation = parser.add_argument_group("Repeatability validation") + validation.add_argument("--repeatability-source-campaign", default=None, help="Existing frozen campaign directory.") + validation.add_argument( + "--repeatability-source-checkpoint", default=None, help="Its passed fpm_forward.json checkpoint." + ) + validation.add_argument( + "--repeatability-output-dir", default=None, help="Fresh validation output; original data is preserved." + ) + validation.add_argument("--repeatability-samples", type=int, default=5) + validation.add_argument( + "--repeatability-comparison-mode", + choices=("full_grid", "bounded"), + default="full_grid", + help="Repeat the complete native grid; bounded subsets are diagnostic only.", + ) + validation.add_argument( + "--repeatability-max-points", type=int, default=12, help="Representative point budget in bounded mode." + ) + validation.add_argument( + "--repeatability-max-attempts", type=int, default=2, help="Attempt budget per fresh sample." + ) + validation.add_argument("--repeatability-cv-threshold", type=float, default=0.05) + validation.add_argument("--repeatability-source-agreement-threshold", type=float, default=0.05) add_fpm_arguments(parser) add_fpm_generator_arguments(parser) return parser @@ -48,12 +71,61 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: parser = _parser() args = parser.parse_args(argv) - if not (args.model_path or args.model_architecture or args.model_cases): + if not (args.model_path or args.model_architecture or args.model_cases or args.repeatability_source_campaign): parser.error("FPM requires --model-path, --model-architecture, or --model-cases") if args.resume_retry_failed and not args.resume: parser.error("--resume-retry-failed requires --resume") + repeatability_paths = ( + args.repeatability_source_campaign, + args.repeatability_source_checkpoint, + args.repeatability_output_dir, + ) + repeatability = any(repeatability_paths) + if repeatability and not all(repeatability_paths): + parser.error("repeatability requires source campaign, source checkpoint and a separate output directory") + if repeatability and (args.smoke or args.limit is not None): + parser.error("repeatability cannot use --smoke or --limit") + if not repeatability and not args.gpu: + parser.error("FPM requires --gpu for a new campaign") try: + if repeatability: + from .repeatability import ( + freeze_repeatability_plan, + load_repeatability_deployment, + load_repeatability_source, + run_repeatability, + ) + + plan = load_repeatability_source(args.repeatability_source_campaign) + generator_overrides = _load_generator_overrides(args) + if not generator_overrides and not args.plan_only: + generator_overrides = load_repeatability_deployment(args.repeatability_source_campaign) + kwargs = { + "samples": args.repeatability_samples, + "max_points_per_cell": args.repeatability_max_points, + "cv_threshold": args.repeatability_cv_threshold, + "comparison_mode": args.repeatability_comparison_mode, + "max_attempts_per_sample": args.repeatability_max_attempts, + "source_agreement_threshold": args.repeatability_source_agreement_threshold, + } + if args.plan_only: + result = freeze_repeatability_plan( + plan, args.repeatability_source_campaign, args.repeatability_source_checkpoint, **kwargs + ) + else: + result = run_repeatability( + plan, + generator_overrides=generator_overrides, + source_campaign_dir=args.repeatability_source_campaign, + source_checkpoint_path=args.repeatability_source_checkpoint, + output_dir=args.repeatability_output_dir, + resume=args.resume, + retry_failed=args.resume_retry_failed, + **kwargs, + ) + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 if args.plan_only or result["status"] == "passed" else 1 case_plan = build_collection_case_plan( backend=args.backend, model_path=args.model_path, diff --git a/python/aisimulate/collector/fpm_forward/config.py b/python/aisimulate/collector/fpm_forward/config.py index 8d62210f7..d5411d364 100644 --- a/python/aisimulate/collector/fpm_forward/config.py +++ b/python/aisimulate/collector/fpm_forward/config.py @@ -6,25 +6,105 @@ from __future__ import annotations import argparse +import copy import hashlib import json +import math from dataclasses import dataclass from pathlib import Path +from typing import Any FPM_FORWARD_OP = "fpm_forward" FPM_WARMUP_ITERATIONS = 5 FPM_MEASUREMENT_REPEATS = 1 FPM_MAX_PREFILL_ISL = 8192 -# Default preserves the extended-capture collection policy. Use -# --fpm-max-prefill-cudagraph-size (e.g. 512) to align the collected machine -# with a deployment that runs vLLM's stock capture defaults. +# Legacy standalone campaigns retain their extended-capture policy. Use +# --fpm-prefill-cudagraph-policy runtime to leave graph selection to vLLM. FPM_MAX_PREFILL_CUDAGRAPH_SIZE = 2048 VLLM_AUTO_FIT_MAX_MODEL_LEN = -1 +# Shared node-task pool, matching generator/config/deployment_config.yaml's +# SlurmConfig default. This is an editable starting point, not per-rank pinning. +FPM_SLURM_CPUS_PER_TASK = 16 +FPM_SLURM_CPU_BIND = "cores" + + +def resolve_slurm_cpu_policy( + cpus_per_task: int | None, + cpu_bind: str | None, + *, + resume: bool = False, + checkpoint_dir: str | Path | None = None, + artifact_root: str | Path | None = None, + smoke: bool = False, +) -> tuple[int | None, str | None]: + """Freeze new defaults, or retain a verified archived policy on resume. + + An old plan with no CPU policy stays unverified. The execution transport + rejects new workers for it, while compatible post-processing remains usable. + """ + if (cpus_per_task is None or cpu_bind is None) and resume and checkpoint_dir and artifact_root: + checkpoint_path = Path(checkpoint_dir).expanduser() / ( + "fpm_forward_smoke.json" if smoke else "fpm_forward.json" + ) + if checkpoint_path.exists(): + from .runtime_memory import validate_saved_plan + + checkpoint = json.loads(checkpoint_path.read_text()) + digest = checkpoint.get("plan_sha256") + if not isinstance(digest, str) or len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("Slurm CPU policy requires a valid archived checkpoint plan hash") + campaign = Path(artifact_root).expanduser() / digest[:16] + if smoke: + campaign /= "smoke" + saved = json.loads((campaign / "collection-plan.json").read_text()) + validate_saved_plan(saved) + if saved["sha256"] != digest: + raise ValueError("archived Slurm CPU policy does not match the checkpoint") + options = saved["options"] + return ( + options.get("slurm_cpus_per_task") if cpus_per_task is None else cpus_per_task, + options.get("slurm_cpu_bind") if cpu_bind is None else cpu_bind, + ) + return ( + FPM_SLURM_CPUS_PER_TASK if cpus_per_task is None else cpus_per_task, + FPM_SLURM_CPU_BIND if cpu_bind is None else cpu_bind, + ) + + +# Dynamo 1.5.0: components/src/dynamo/vllm/instrumented_scheduler.py, +# _kvwarm_flag_on / _bench_realseed_on at b83b1d9304ebfc624709ac46db32b1b6f1ff1615: +# https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/instrumented_scheduler.py +# Decode uses the runtime's eligibility gate; cached prefill has a separate +# real-prefix seeding switch. Dynamo retains ownership of point generation +# and feasibility. +FPM_KV_WARMUP_DEFAULTS = {"DYN_BENCH_KV_WARMUP": "on", "DYN_BENCH_PREFILL_REAL_SEED": "on"} PARALLEL_AXES = ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp") PARALLEL_PRESETS = ("auto", "tp", "tep", "dep", "pure_tp") +def with_kv_warmup_defaults(overrides: dict[str, Any]) -> dict[str, Any]: + """Resolve declared warm-up switches without replacing explicit choices.""" + + resolved = copy.deepcopy(overrides) + k8s = resolved.setdefault("K8sConfig", {}) + environment = k8s.setdefault("extra_env", []) + if environment is None: + environment = k8s["extra_env"] = [] + declared = set() + for item in environment: + name = item["name"] + if name in FPM_KV_WARMUP_DEFAULTS: + value = item.get("value") + if not isinstance(value, str) or value.lower() not in {"on", "1", "true", "off", "0", "false"}: + raise ValueError(f"{name} must be a literal on/off, 1/0 or true/false string") + declared.add(name) + environment.extend( + {"name": name, "value": value} for name, value in FPM_KV_WARMUP_DEFAULTS.items() if name not in declared + ) + return resolved + + def _positive_int(value: str) -> int: parsed = int(value) if parsed < 1: @@ -32,6 +112,13 @@ def _positive_int(value: str) -> int: return parsed +def _positive_or_auto_int(value: str) -> int: + parsed = int(value) + if parsed != VLLM_AUTO_FIT_MAX_MODEL_LEN and parsed <= 0: + raise argparse.ArgumentTypeError("must be positive or -1 for runtime auto-fit") + return parsed + + def _at_least_two_int(value: str) -> int: parsed = int(value) if parsed < 2: @@ -46,6 +133,13 @@ def _nonnegative_int(value: str) -> int: return parsed +def _memory_fraction(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or not 0 < parsed <= 1: + raise argparse.ArgumentTypeError("value must be finite and in (0, 1]") + return parsed + + def _optional_size_list(values: list[int] | None) -> tuple[int, ...] | None: if values is None: return None @@ -154,11 +248,12 @@ class PrefillSamplingProfile: max_isl: int max_batch_size: int | None max_total_prefill_tokens: int - cudagraph_capture_sizes: tuple[int, ...] - max_cudagraph_capture_size: int - new_token_axis_points: tuple[int, ...] - max_new_token_samples: int + cudagraph_capture_sizes: tuple[int, ...] | None + max_cudagraph_capture_size: int | None + new_token_axis_points: tuple[int, ...] | None + max_new_token_samples: int | None max_kv_read_token_samples: int + cudagraph_policy: str = "explicit" @classmethod def build( @@ -166,17 +261,26 @@ def build( *, max_isl: int, max_batch_size: int | None, - max_cudagraph_capture_size: int = FPM_MAX_PREFILL_CUDAGRAPH_SIZE, + max_cudagraph_capture_size: int | None = None, + cudagraph_policy: str = "explicit", ) -> PrefillSamplingProfile: if max_isl < 2: raise ValueError("FPM max prefill ISL must be at least 2") if max_batch_size is not None and max_batch_size < 1: raise ValueError("FPM max prefill batch size must be positive") - if max_cudagraph_capture_size < 1: - raise ValueError("FPM max prefill CUDA-graph capture size must be positive") - - capture_sizes = _prefill_cudagraph_capture_sizes(max_isl, max_cudagraph_capture_size) - new_token_points = _dynamo_cudagraph_axis_points(capture_sizes, max_isl) + if cudagraph_policy not in {"explicit", "runtime"}: + raise ValueError("FPM prefill CUDA-graph policy must be runtime or explicit") + if cudagraph_policy == "runtime" and max_cudagraph_capture_size is not None: + raise ValueError("runtime prefill CUDA-graph policy cannot specify a capture size") + capture_sizes = None + new_token_points = None + if cudagraph_policy == "explicit": + if max_cudagraph_capture_size is None: + max_cudagraph_capture_size = FPM_MAX_PREFILL_CUDAGRAPH_SIZE + if type(max_cudagraph_capture_size) is not int or max_cudagraph_capture_size < 1: + raise ValueError("FPM max prefill CUDA-graph capture size must be positive") + capture_sizes = _prefill_cudagraph_capture_sizes(max_isl, max_cudagraph_capture_size) + new_token_points = _dynamo_cudagraph_axis_points(capture_sizes, max_isl) batch_upper_bound = min(max_isl, max_batch_size or max_isl) # PR11509's KV axis contains zero, a batch-aligned minimum, powers of # two in block units, and an exact maximum. This upper bound prevents @@ -188,25 +292,48 @@ def build( max_batch_size=max_batch_size, max_total_prefill_tokens=max_isl, cudagraph_capture_sizes=capture_sizes, - max_cudagraph_capture_size=capture_sizes[-1], + max_cudagraph_capture_size=capture_sizes[-1] if capture_sizes is not None else None, new_token_axis_points=new_token_points, - max_new_token_samples=max(2, len(new_token_points)), + max_new_token_samples=max(2, len(new_token_points)) if new_token_points is not None else None, max_kv_read_token_samples=max_kv_read_samples, + cudagraph_policy=cudagraph_policy, ) def to_dict(self) -> dict[str, object]: - return { + payload = { "max_isl": self.max_isl, "max_batch_size": self.max_batch_size, "max_total_prefill_tokens": self.max_total_prefill_tokens, - "cudagraph_capture_sizes": list(self.cudagraph_capture_sizes), - "cudagraph_capture_size_count": len(self.cudagraph_capture_sizes), + "cudagraph_capture_sizes": list(self.cudagraph_capture_sizes) + if self.cudagraph_capture_sizes is not None + else None, + "cudagraph_capture_size_count": len(self.cudagraph_capture_sizes) + if self.cudagraph_capture_sizes is not None + else None, "max_cudagraph_capture_size": self.max_cudagraph_capture_size, - "new_token_axis_points": list(self.new_token_axis_points), - "new_token_axis_point_count": len(self.new_token_axis_points), + "new_token_axis_points": list(self.new_token_axis_points) + if self.new_token_axis_points is not None + else None, + "new_token_axis_point_count": len(self.new_token_axis_points) + if self.new_token_axis_points is not None + else None, "prefill_max_new_token_samples": self.max_new_token_samples, "prefill_max_kv_read_token_samples": self.max_kv_read_token_samples, } + # Keep legacy explicit-capture plan identity unchanged. Runtime-selected + # graphs and their candidate axes cannot be known before initialization. + if self.cudagraph_policy == "runtime": + payload["cudagraph_policy"] = "runtime" + return payload + + +def collection_phases(worker_type: str | None) -> tuple[str, ...]: + """Preserve historical phase pairs unless an exact serving role is selected.""" + if worker_type in (None, "aggregated"): + return ("prefill", "decode") + if worker_type in ("prefill", "decode"): + return (worker_type,) + raise ValueError("worker_type must be prefill, decode or aggregated") @dataclass(frozen=True, slots=True) @@ -242,10 +369,19 @@ class FPMCollectionOptions: # vLLM 0.24 auto-fits this limit after model/CUDA-graph profiling, and # Dynamo records the resolved value in ``limits.max_model_len``. vllm_max_model_len: int = VLLM_AUTO_FIT_MAX_MODEL_LEN - max_prefill_isl: int = FPM_MAX_PREFILL_ISL + max_num_batched_tokens: int | None = None + max_num_seqs: int | None = None + max_prefill_isl: int | None = None max_prefill_batch_size: int | None = None + max_prefill_cudagraph_size: int | None = None + prefill_cudagraph_policy: str = "explicit" + worker_type: str | None = None + cudagraph_mode: str | None = None + cudagraph_capture_sizes: tuple[int, ...] | None = None + max_cudagraph_capture_size: int | None = None + gpu_memory_utilization: float | None = None + max_decode_batch_size: int | None = None - max_prefill_cudagraph_size: int = FPM_MAX_PREFILL_CUDAGRAPH_SIZE decoder_replay: bool = False enforce_eager: bool = False benchmark_points_json: str | None = None @@ -253,13 +389,164 @@ class FPMCollectionOptions: executor: str = "kubernetes" slurm_container_image: str = "" slurm_container_mounts: tuple[str, ...] = () + # None preserves legacy saved-plan identity; fresh CLI plans resolve both. + slurm_cpus_per_task: int | None = None + slurm_cpu_bind: str | None = None + + def __post_init__(self) -> None: + collection_phases(self.worker_type) + graph_fields = (self.cudagraph_mode, self.cudagraph_capture_sizes, self.max_cudagraph_capture_size) + if any(value is not None for value in graph_fields) and self.worker_type is None: + raise ValueError("CUDA graph settings require an explicit --fpm-worker-type") + if self.cudagraph_mode is not None and self.cudagraph_mode not in { + "NONE", + "PIECEWISE", + "FULL", + "FULL_DECODE_ONLY", + "FULL_AND_PIECEWISE", + }: + raise ValueError("unsupported --fpm-cudagraph-mode") + if self.cudagraph_capture_sizes is not None: + sizes = self.cudagraph_capture_sizes + if ( + not sizes + or any(type(value) is not int or value < 1 for value in sizes) + or tuple(sorted(set(sizes))) != sizes + ): + raise ValueError("--fpm-cudagraph-capture-sizes must be sorted unique positive integers") + if self.max_cudagraph_capture_size is not None and self.max_cudagraph_capture_size != max(sizes): + raise ValueError("--fpm-max-cudagraph-capture-size must equal the largest capture size") + if self.max_cudagraph_capture_size is not None and ( + type(self.max_cudagraph_capture_size) is not int or self.max_cudagraph_capture_size < 1 + ): + raise ValueError("--fpm-max-cudagraph-capture-size must be positive") + if (self.cudagraph_capture_sizes is not None or self.max_cudagraph_capture_size is not None) and ( + self.prefill_cudagraph_policy != "runtime" or self.max_prefill_cudagraph_size is not None + ): + raise ValueError("generic CUDA graph captures require --fpm-prefill-cudagraph-policy runtime") + if self.cudagraph_mode == "NONE" and ( + self.cudagraph_capture_sizes is not None + or self.max_cudagraph_capture_size is not None + or self.prefill_cudagraph_policy != "runtime" + ): + raise ValueError("CUDA graph mode NONE cannot declare capture sizes") + maximum = self.max_cudagraph_capture_size or max(self.cudagraph_capture_sizes or (0,)) + if self.max_num_batched_tokens is not None and maximum > self.max_num_batched_tokens: + raise ValueError("CUDA graph capture size exceeds max_num_batched_tokens") + if (self.slurm_cpus_per_task is None) != (self.slurm_cpu_bind is None): + raise ValueError("Slurm CPU policy requires both cpus-per-task and cpu-bind") + if self.slurm_cpus_per_task is not None: + if self.executor != "slurm": + raise ValueError("Slurm CPU options require --fpm-executor slurm") + if type(self.slurm_cpus_per_task) is not int or self.slurm_cpus_per_task < 1: + raise ValueError("--fpm-slurm-cpus-per-task must be a positive integer") + if self.slurm_cpu_bind not in {"cores", "none"}: + raise ValueError("--fpm-slurm-cpu-bind must be cores or none") + if self.prefill_cudagraph_policy not in {"runtime", "explicit"}: + raise ValueError("--fpm-prefill-cudagraph-policy must be runtime or explicit") + if self.prefill_cudagraph_policy == "runtime" and self.max_prefill_cudagraph_size is not None: + raise ValueError("--fpm-prefill-cudagraph-policy runtime cannot use --fpm-max-prefill-cudagraph-size") + if self.prefill_cudagraph_policy == "explicit" and self.max_prefill_cudagraph_size is None: + object.__setattr__(self, "max_prefill_cudagraph_size", FPM_MAX_PREFILL_CUDAGRAPH_SIZE) + if self.gpu_memory_utilization is not None and ( + type(self.gpu_memory_utilization) not in {int, float} + or not math.isfinite(self.gpu_memory_utilization) + or not 0 < self.gpu_memory_utilization <= 1 + ): + raise ValueError("--fpm-gpu-memory-utilization must be finite and in (0, 1]") + for name, value in ( + ( + "max-model-len", + None if self.vllm_max_model_len == VLLM_AUTO_FIT_MAX_MODEL_LEN else self.vllm_max_model_len, + ), + ("max-num-batched-tokens", self.max_num_batched_tokens), + ("max-num-seqs", self.max_num_seqs), + ("max-decode-batch-size", self.max_decode_batch_size), + ("max-prefill-cudagraph-size", self.max_prefill_cudagraph_size), + ): + if value is not None and (type(value) is not int or value < 1): + raise ValueError(f"--fpm-{name} must be a positive integer") + if ( + self.max_prefill_isl is not None + and self.max_num_batched_tokens is not None + and self.max_prefill_isl > self.max_num_batched_tokens + ): + raise ValueError("--fpm-max-prefill-isl exceeds --fpm-max-num-batched-tokens") + if ( + self.max_prefill_batch_size is not None + and self.max_num_seqs is not None + and self.max_prefill_batch_size > self.max_num_seqs + ): + raise ValueError("--fpm-max-prefill-batch-size exceeds --fpm-max-num-seqs") + if ( + self.max_decode_batch_size is not None + and self.max_num_seqs is not None + and self.max_decode_batch_size > self.max_num_seqs + ): + raise ValueError("--fpm-max-decode-batch-size exceeds --fpm-max-num-seqs") + self.validate_scheduler_limits() + + def validate_scheduler_limits( + self, *, profile_max_num_tokens: int | None = None, profile_max_batch_size: int | None = None + ) -> None: + """Check known effective phase limits without inventing runtime defaults.""" + shared_tokens = self.max_num_batched_tokens or profile_max_num_tokens + shared_sequences = self.max_num_seqs or profile_max_batch_size + if self.worker_type == "aggregated" and ( + (shared_tokens is not None and self.max_prefill_isl not in (None, shared_tokens)) + or ( + shared_sequences is not None + and ( + self.max_prefill_batch_size not in (None, shared_sequences) + or self.max_decode_batch_size not in (None, shared_sequences) + ) + ) + ): + raise ValueError("aggregated serving requires shared prefill and decode scheduler limits") + for phase, tokens, sequences in ( + ("decode", shared_tokens, self.max_decode_batch_size or shared_sequences), + ( + "prefill", + self.max_prefill_isl or shared_tokens or FPM_MAX_PREFILL_ISL, + self.max_prefill_batch_size or shared_sequences, + ), + ): + # vLLM v0.25.1 vllm/config/scheduler.py:286 rejects token budgets below + # max_num_seqs: every scheduled sequence needs at least one token. + if tokens is not None and sequences is not None and tokens < sequences: + raise ValueError( + f"FPM {phase} max_num_batched_tokens ({tokens}) must be at least max_num_seqs ({sequences}); " + "edit the shared/prefill limits or profile bounds" + ) + + @property + def workload_kinds(self) -> tuple[str, ...]: + return collection_phases(self.worker_type) + + @property + def compilation_config(self) -> dict[str, object]: + result = {} + if self.cudagraph_mode is not None: + result["cudagraph_mode"] = self.cudagraph_mode + if self.cudagraph_capture_sizes is not None: + result["cudagraph_capture_sizes"] = list(self.cudagraph_capture_sizes) + if self.max_cudagraph_capture_size is not None: + result["max_cudagraph_capture_size"] = self.max_cudagraph_capture_size + elif self.cudagraph_capture_sizes is not None: + result["max_cudagraph_capture_size"] = max(self.cudagraph_capture_sizes) + return result @property def prefill_sampling(self) -> PrefillSamplingProfile: return PrefillSamplingProfile.build( - max_isl=self.max_prefill_isl, - max_batch_size=self.max_prefill_batch_size, + max_isl=self.max_prefill_isl + if self.max_prefill_isl is not None + else (self.max_num_batched_tokens or FPM_MAX_PREFILL_ISL), + max_batch_size=self.max_prefill_batch_size + if self.max_prefill_batch_size is not None + else self.max_num_seqs, max_cudagraph_capture_size=self.max_prefill_cudagraph_size, + cudagraph_policy=self.prefill_cudagraph_policy, ) @classmethod @@ -301,10 +588,21 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: executor = getattr(args, "fpm_executor", None) or "kubernetes" image = getattr(args, "fpm_slurm_container_image", None) or "" mounts = tuple(getattr(args, "fpm_slurm_container_mount", None) or ()) + cpus = getattr(args, "fpm_slurm_cpus_per_task", None) + cpu_bind = getattr(args, "fpm_slurm_cpu_bind", None) if executor == "slurm" and not image: raise ValueError("--fpm-executor slurm requires --fpm-slurm-container-image") - if executor != "slurm" and (image or mounts): - raise ValueError("Slurm container options require --fpm-executor slurm") + if executor != "slurm" and (image or mounts or cpus is not None or cpu_bind is not None): + raise ValueError("Slurm container/CPU options require --fpm-executor slurm") + if executor == "slurm": + cpus, cpu_bind = resolve_slurm_cpu_policy( + cpus, + cpu_bind, + resume=bool(getattr(args, "resume", False)), + checkpoint_dir=getattr(args, "checkpoint_dir", None), + artifact_root=getattr(args, "fpm_artifact_root", None) or "fpm_forward_artifacts", + smoke=bool(getattr(args, "smoke", False)), + ) model_len = getattr(args, "fpm_max_model_len", None) if model_len is not None and model_len != -1 and model_len < 1: @@ -343,22 +641,38 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: if getattr(args, "fpm_warmup_iterations", None) is None else args.fpm_warmup_iterations ), - vllm_max_model_len=model_len if model_len is not None else VLLM_AUTO_FIT_MAX_MODEL_LEN, + vllm_max_model_len=( + VLLM_AUTO_FIT_MAX_MODEL_LEN + if getattr(args, "fpm_max_model_len", None) is None + else args.fpm_max_model_len + ), + max_num_batched_tokens=getattr(args, "fpm_max_num_batched_tokens", None), + max_num_seqs=getattr(args, "fpm_max_num_seqs", None), max_decode_batch_size=getattr(args, "fpm_max_decode_batch_size", None), - max_prefill_isl=getattr(args, "fpm_max_prefill_isl", None) or FPM_MAX_PREFILL_ISL, + max_prefill_isl=getattr(args, "fpm_max_prefill_isl", None), max_prefill_batch_size=getattr(args, "fpm_max_prefill_batch_size", None), - max_prefill_cudagraph_size=( - getattr(args, "fpm_max_prefill_cudagraph_size", None) or FPM_MAX_PREFILL_CUDAGRAPH_SIZE + max_prefill_cudagraph_size=getattr(args, "fpm_max_prefill_cudagraph_size", None), + prefill_cudagraph_policy=getattr(args, "fpm_prefill_cudagraph_policy", None) or "explicit", + worker_type=getattr(args, "fpm_worker_type", None), + cudagraph_mode=getattr(args, "fpm_cudagraph_mode", None), + cudagraph_capture_sizes=( + tuple(args.fpm_cudagraph_capture_sizes) + if getattr(args, "fpm_cudagraph_capture_sizes", None) is not None + else None ), + max_cudagraph_capture_size=getattr(args, "fpm_max_cudagraph_capture_size", None), + gpu_memory_utilization=getattr(args, "fpm_gpu_memory_utilization", None), decoder_replay=bool(getattr(args, "fpm_decoder_replay", False)), enforce_eager=bool(getattr(args, "fpm_enforce_eager", False)), executor=executor, slurm_container_image=image, slurm_container_mounts=mounts, + slurm_cpus_per_task=cpus, + slurm_cpu_bind=cpu_bind, ) def to_dict(self) -> dict[str, object]: - payload = { + result = { "decoder_replay": self.decoder_replay, "executor": self.executor, "slurm_container_image": self.slurm_container_image, @@ -387,14 +701,33 @@ def to_dict(self) -> dict[str, object]: "point_source": "dynamo_native_self_benchmark", "prefill_sampling": self.prefill_sampling.to_dict(), } + # Preserve the existing frozen-plan representation when the new shared + # runtime limits are absent. + if self.worker_type is not None: + result["worker_type"] = self.worker_type + if self.cudagraph_mode is not None: + result["cudagraph_mode"] = self.cudagraph_mode + if self.cudagraph_capture_sizes is not None: + result["cudagraph_capture_sizes"] = list(self.cudagraph_capture_sizes) + if self.max_cudagraph_capture_size is not None: + result["max_cudagraph_capture_size"] = self.max_cudagraph_capture_size + if self.slurm_cpus_per_task is not None: + result["slurm_cpus_per_task"] = self.slurm_cpus_per_task + result["slurm_cpu_bind"] = self.slurm_cpu_bind + if self.max_num_batched_tokens is not None: + result["max_num_batched_tokens"] = self.max_num_batched_tokens + if self.max_num_seqs is not None: + result["max_num_seqs"] = self.max_num_seqs + if self.gpu_memory_utilization is not None: + result["gpu_memory_utilization"] = self.gpu_memory_utilization if self.enforce_eager: - payload["enforce_eager"] = True + result["enforce_eager"] = True if self.benchmark_points_json is not None: - payload["benchmark_points"] = { + result["benchmark_points"] = { "payload": json.loads(self.benchmark_points_json), "sha256": self.benchmark_points_sha256, } - return payload + return result def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: @@ -404,12 +737,43 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: "FPM forward collection", "Whole-model forward-pass planning, execution, and publication.", ) + group.add_argument( + "--fpm-worker-type", + choices=("prefill", "decode", "aggregated"), + default=None, + help="Collect one serving role, or both phases of an aggregated engine.", + ) + group.add_argument( + "--fpm-cudagraph-mode", + choices=("NONE", "PIECEWISE", "FULL", "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"), + default=None, + ) + group.add_argument("--fpm-cudagraph-capture-sizes", type=int, nargs="+", default=None) + group.add_argument("--fpm-max-cudagraph-capture-size", type=int, default=None) group.add_argument( "--fpm-model-profile", default=None, metavar="PATH", help="JSON or YAML FPM model profile with exact deployment identities and rank-local resource bounds.", ) + group.add_argument( + "--fpm-runtime-instrumentation", + default=None, + metavar="PATH", + help="Frozen instrumentation manifest for verifying collection against accepted runtime observations.", + ) + group.add_argument( + "--fpm-runtime-launch", + default=None, + metavar="PATH", + help="JSON launch facts from the accepted runtime probe; requires runtime instrumentation and configuration.", + ) + group.add_argument( + "--fpm-runtime-configuration", + default=None, + metavar="KEY", + help="Checkpoint configuration key for runtime observations; requires instrumentation and launch facts.", + ) group.add_argument( "--fpm-enforce-eager", action="store_true", @@ -429,7 +793,10 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: help="Use true bounded decoder replay; requires verified runtime support.", ) group.add_argument( - "--fpm-max-model-len", type=int, default=None, help="vLLM context limit; positive or -1 for runtime auto-fit." + "--fpm-max-model-len", + type=_positive_or_auto_int, + default=None, + help="Maximum context length for both phases; defaults to the profile context or vLLM auto-fit (-1).", ) group.add_argument( "--fpm-max-decode-batch-size", @@ -517,6 +884,18 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: f"(default: {FPM_WARMUP_ITERATIONS})." ), ) + group.add_argument( + "--fpm-max-num-batched-tokens", + type=_positive_int, + default=None, + help="Shared vLLM scheduled-token limit for prefill and decode; may be narrowed for prefill.", + ) + group.add_argument( + "--fpm-max-num-seqs", + type=_positive_int, + default=None, + help="Shared vLLM sequence limit for prefill and decode; may be narrowed for prefill.", + ) group.add_argument( "--fpm-max-prefill-isl", dest="fpm_max_prefill_isl", @@ -525,7 +904,8 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: help=( "Maximum total scheduled prefill new-token axis; the batch=1 points also " "cover this per-request new-token length " - f"(default: {FPM_MAX_PREFILL_ISL})." + "(defaults to the shared token limit, otherwise " + f"{FPM_MAX_PREFILL_ISL} within supplied profile bounds)." ), ) group.add_argument( @@ -533,7 +913,22 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: dest="fpm_max_prefill_batch_size", type=_positive_int, default=None, - help="Optional prefill max_num_seqs override; defaults to the Dynamo/vLLM runtime value.", + help="Optional narrower prefill max_num_seqs; defaults to shared/profile bounds or the Dynamo/vLLM value.", + ) + group.add_argument( + "--fpm-gpu-memory-utilization", + type=_memory_fraction, + default=None, + help="vLLM fraction of total GPU memory for weights, runtime and KV cache; omitted uses the runtime default.", + ) + group.add_argument( + "--fpm-prefill-cudagraph-policy", + choices=("runtime", "explicit"), + default=None, + help=( + "Prefill graph policy: runtime leaves graph selection to vLLM; " + "explicit retains the bounded capture axis (default)." + ), ) group.add_argument( "--fpm-max-prefill-cudagraph-size", @@ -541,9 +936,8 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: type=_positive_int, default=None, help=( - "Largest prefill CUDA-graph capture size in the collected axis; set 512 to " - "mirror vLLM's stock capture defaults " - f"(default: {FPM_MAX_PREFILL_CUDAGRAPH_SIZE})." + "Largest prefill CUDA-graph capture size with the explicit policy; " + f"defaults to {FPM_MAX_PREFILL_CUDAGRAPH_SIZE}. Cannot be combined with the runtime policy." ), ) group.add_argument( @@ -607,6 +1001,18 @@ def add_fpm_generator_arguments(parser: argparse.ArgumentParser) -> None: default=None, help="Pyxis SOURCE:TARGET mount; repeat for checkpoint/runtime/cache paths.", ) + group.add_argument( + "--fpm-slurm-cpus-per-task", + type=_positive_int, + default=None, + help="CPUs for each node's one-task Slurm step (new plans: 16); shared by local workers and schedulers.", + ) + group.add_argument( + "--fpm-slurm-cpu-bind", + choices=["cores", "none"], + default=None, + help="Slurm node-task CPU binding (new plans: cores); no dedicated per-rank pinning.", + ) group.add_argument( "--generator-config", default=None, @@ -658,6 +1064,13 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_kv_cache_dtypes", "fpm_model_config", "fpm_model_profile", + "fpm_runtime_instrumentation", + "fpm_runtime_launch", + "fpm_runtime_configuration", + "fpm_worker_type", + "fpm_cudagraph_mode", + "fpm_cudagraph_capture_sizes", + "fpm_max_cudagraph_capture_size", "fpm_tp_sizes", "fpm_pp_sizes", "fpm_dp_sizes", @@ -671,11 +1084,15 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_enable_wideep", "fpm_enable_eplb", "fpm_warmup_iterations", + "fpm_max_model_len", + "fpm_max_num_batched_tokens", + "fpm_max_num_seqs", + "fpm_gpu_memory_utilization", + "fpm_prefill_cudagraph_policy", "fpm_max_prefill_isl", "fpm_max_prefill_batch_size", "fpm_max_prefill_cudagraph_size", "fpm_artifact_root", - "fpm_max_model_len", "fpm_max_decode_batch_size", "fpm_decoder_replay", "fpm_enforce_eager", @@ -683,6 +1100,8 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_executor", "fpm_slurm_container_image", "fpm_slurm_container_mount", + "fpm_slurm_cpus_per_task", + "fpm_slurm_cpu_bind", "fpm_database_root", "fpm_publish_partial", # Deployment-only Generator inputs are registered unconditionally on diff --git a/python/aisimulate/collector/fpm_forward/cpu_affinity.py b/python/aisimulate/collector/fpm_forward/cpu_affinity.py new file mode 100644 index 000000000..707942a13 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/cpu_affinity.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate initialization-time CPU evidence separately from memory and timings.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any + +from .runtime_instrumentation import read_json + + +def _integers(value: Any, label: str, *, minimum: int, empty: bool = False) -> list[int]: + if ( + not isinstance(value, list) + or (not value and not empty) + or any(type(item) is not int or item < minimum for item in value) + or value != sorted(set(value)) + ): + raise ValueError(f"{label} must contain sorted, unique integers >= {minimum}") + return value + + +def _snapshot(record: dict[str, Any]) -> None: + if not isinstance(record.get("hostname"), str) or not record["hostname"].strip(): + raise ValueError("CPU observation has no hostname") + for name in ("pid", "observer_thread_id"): + if type(record.get(name)) is not int or record[name] < 1: + raise ValueError(f"CPU observation has invalid {name}") + status = record.get("status") + if status not in {"observed", "partial", "unavailable"}: + raise ValueError("CPU observation has invalid status") + main = _integers( + record.get("main_thread_allowed_cpus"), "main thread CPU mask", minimum=0, empty=status == "unavailable" + ) + before = _integers(record.get("thread_ids_before"), "thread IDs before observation", minimum=1, empty=True) + after = _integers(record.get("thread_ids_after"), "thread IDs after observation", minimum=1, empty=True) + threads = record.get("threads") + if not isinstance(threads, list): + raise ValueError("CPU observation lacks thread masks") + seen = {} + for thread in threads: + if not isinstance(thread, dict): + raise ValueError("CPU thread observation must be an object") + tid = thread.get("tid") + if type(tid) is not int or tid < 1 or tid in seen or tid not in before: + raise ValueError("CPU thread observation has an invalid, duplicate or unlisted thread ID") + seen[tid] = _integers(thread.get("allowed_cpus"), "thread CPU mask", minimum=0) + errors, thread_errors = record.get("errors"), record.get("thread_errors") + if not isinstance(errors, list) or any(not isinstance(error, str) or not error for error in errors): + raise ValueError("CPU observation errors must be a list of messages") + if not isinstance(thread_errors, list) or any( + not isinstance(item, dict) + or type(item.get("tid")) is not int + or item["tid"] not in before + or not isinstance(item.get("error"), str) + or not item["error"] + for item in thread_errors + ): + raise ValueError("CPU thread errors must identify an observed thread and error") + if status == "observed" and ( + before != after + or set(seen) != set(before) + or seen.get(record["pid"]) != main + or record["observer_thread_id"] not in seen + or errors + or thread_errors + ): + raise ValueError("CPU observation claims complete thread evidence with missing or inconsistent masks") + + +def _can_schedule_distinct(masks: list[list[int]]) -> bool: + """Check main-thread capacity without requiring disjoint affinity masks.""" + assigned: dict[int, int] = {} + + def assign(index: int, seen: set[int]) -> bool: + for cpu in masks[index]: + if cpu in seen: + continue + seen.add(cpu) + previous = assigned.get(cpu) + if previous is None or assign(previous, seen): + assigned[cpu] = index + return True + return False + + return all(assign(index, set()) for index in range(len(masks))) + + +def inspect_cpu_affinity( + cell, raw_root: Path, collection, *, plan, expected_nodes: int | None = None +) -> dict[str, Any]: + """Inspect only this cell's attempt; old artifacts remain CPU-unqualified. + + Main-thread masks and thread snapshots describe specific initialization + hooks. They do not prove CPU utilization, exclusive cores, NUMA locality or + affinity throughout later timed forwards. Shared multi-core pools are valid. + """ + requested = getattr(plan.options, "slurm_cpus_per_task", None) + binding = getattr(plan.options, "slurm_cpu_bind", None) + required = plan.options.executor == "slurm" and (requested is not None or binding is not None) + missing, failures, observations = [], [], [] + if required and (type(requested) is not int or requested < 1 or binding not in {"cores", "none"}): + failures.append("frozen CPU policy requires positive cpus_per_task and cpu_bind=cores|none") + elif not required: + missing.append("frozen CPU policy is unavailable; historical observations do not qualify CPU allocation") + if required and expected_nodes is None: + missing.append("generated launch node count is unavailable") + elif expected_nodes is not None and (type(expected_nodes) is not int or expected_nodes < 1): + failures.append("generated launch node count must be a positive integer") + expected_nodes = None + expected_workers = { + (dp, tp, pp) + for dp in range(cell.topology.dp) + for tp in range(cell.topology.tp) + for pp in range(cell.topology.pp) + } + workers, schedulers, launchers = {}, {}, {} + processes = set() + measurements = { + "launcher": "before_engine_start", + "worker": "after_warmup", + "scheduler": "scheduler_initialized", + } + for path in sorted(raw_root.glob("**/fpm-cpu-*.json")): + try: + raw = path.read_bytes() + record = read_json(path) + source = {"path": str(path.resolve()), "sha256": hashlib.sha256(raw).hexdigest(), "size_bytes": len(raw)} + item = {**record, "source": source} + observations.append(item) + if ( + record.get("schema_name") != "aisimulate_fpm_cpu_affinity" + or type(record.get("schema_version")) is not int + or record["schema_version"] != 1 + ): + raise ValueError("CPU observation has an unsupported schema") + provenance = record.get("collector_provenance") + if not isinstance(provenance, dict) or any( + provenance.get(key) != value + for key, value in { + "plan_sha256": plan.sha256, + "cell_id": cell.cell_id, + "attempt_id": collection.collector_attempt_id, + }.items() + ): + raise ValueError("CPU evidence has a different plan, cell or attempt") + kind = record.get("kind") + if kind not in measurements or record.get("measurement") != measurements[kind]: + raise ValueError("CPU observation has an unsupported role or lifecycle hook") + _snapshot(record) + process = (kind, record["hostname"], record["pid"]) + if process in processes: + raise ValueError("CPU observation has a duplicate process identity for its role") + processes.add(process) + rank = tuple(record.get(name) for name in ("dp_rank", "tp_rank", "pp_rank")) + if kind == "worker": + if any(type(value) is not int for value in rank) or rank not in expected_workers or rank in workers: + raise ValueError("CPU worker rank is unexpected or duplicate") + workers[rank] = item + suffix = f"-dp{rank[0]}-tp{rank[1]}-pp{rank[2]}" + elif kind == "scheduler": + dp, tp, pp = rank + if ( + type(dp) is not int + or not 0 <= dp < cell.topology.dp + or tp is not None + or pp is not None + or dp in schedulers + ): + raise ValueError("CPU scheduler rank is unexpected or duplicate") + schedulers[dp] = item + suffix = f"-dp{dp}" + else: + if any(value is not None for value in rank) or record["hostname"] in launchers: + raise ValueError("CPU launcher has rank coordinates or a duplicate hostname") + launchers[record["hostname"]] = item + suffix = "" + if path.name != f"fpm-cpu-{kind}{suffix}.json": + raise ValueError("CPU observation filename disagrees with its role/rank") + if record["status"] != "observed": + missing.append(f"{kind} {rank} on {record['hostname']}: {record['status']} CPU affinity snapshot") + except (OSError, TypeError, ValueError, KeyError) as error: + failures.append(f"{path.name}: {error}") + if workers.keys() != expected_workers: + missing.append(f"worker CPU observations missing for ranks {sorted(expected_workers - workers.keys())}") + if schedulers.keys() != set(range(cell.topology.dp)): + missing.append( + "scheduler CPU observations missing for DP ranks " + f"{sorted(set(range(cell.topology.dp)) - schedulers.keys())}" + ) + for dp, scheduler in schedulers.items(): + worker = workers.get((dp, 0, 0)) + if worker is not None and worker["hostname"] != scheduler["hostname"]: + failures.append(f"DP{dp} scheduler host differs from its TP0/PP0 worker") + observed_hosts = {item["hostname"] for item in [*workers.values(), *schedulers.values()]} + if expected_nodes is not None and len(workers) == len(expected_workers): + worker_hosts = {item["hostname"] for item in workers.values()} + if len(worker_hosts) != expected_nodes or ( + len(launchers) == len(worker_hosts) and launchers.keys() != worker_hosts + ): + failures.append("observed worker/launcher hosts differ from the generated launch node count") + if len(launchers) > expected_nodes: + failures.append("launcher observations exceed the generated node count") + nodes = [] + for hostname in sorted(observed_hosts | launchers.keys()): + local_workers = [item for item in workers.values() if item["hostname"] == hostname] + local_schedulers = [item for item in schedulers.values() if item["hostname"] == hostname] + launcher = launchers.get(hostname) + node = { + "hostname": hostname, + "worker_ranks": [ + {name: item[name] for name in ("dp_rank", "tp_rank", "pp_rank")} for item in local_workers + ], + "scheduler_dp_ranks": [item["dp_rank"] for item in local_schedulers], + "launcher_allowed_cpus": launcher["main_thread_allowed_cpus"] if launcher else None, + } + nodes.append(node) + if launcher is None: + missing.append(f"launcher CPU pool observation missing for {hostname}") + # A rank on a host that has no launcher while all expected pools are + # present is also caught by the empty worker pool check below. + continue + if not local_workers and len(workers) == len(expected_workers): + failures.append(f"launcher {hostname} has no worker CPU observations on its host") + pool = set(launcher["main_thread_allowed_cpus"]) + if required: + if ( + type(launcher.get("requested_cpus_per_task")) is not int + or launcher["requested_cpus_per_task"] != requested + or launcher.get("cpu_bind") != binding + ): + failures.append(f"launcher {hostname} CPU policy differs from the frozen requested policy") + if pool and type(requested) is int and len(pool) < requested: + failures.append( + f"launcher {hostname} allows fewer logical CPUs than requested cpus_per_task={requested}" + ) + if type(requested) is int and requested < len(local_schedulers): + failures.append(f"launcher {hostname} requested CPU count is below its local DP scheduler count") + local_gpus = launcher.get("local_gpu_count") + if type(local_gpus) is not int or local_gpus < 1: + failures.append(f"launcher {hostname} has no valid local GPU count") + elif local_gpus != len(local_workers) and len(workers) == len(expected_workers): + failures.append(f"launcher {hostname} local GPU count differs from observed worker rank count") + if expected_nodes is not None and local_gpus != len(expected_workers) / expected_nodes: + failures.append(f"launcher {hostname} local GPU count differs from the generated launch geometry") + for item in [launcher, *local_workers, *local_schedulers]: + if pool and not set(item["main_thread_allowed_cpus"]) <= pool: + failures.append(f"{item['kind']} on {hostname} has a main-thread mask outside the launcher CPU pool") + for thread in item["threads"]: + if pool and not set(thread["allowed_cpus"]) <= pool: + failures.append( + f"{item['kind']} on {hostname} has thread {thread['tid']} outside the launcher CPU pool" + ) + masks = [item["main_thread_allowed_cpus"] for item in local_schedulers] + if masks and all(masks): + node["scheduler_main_threads_have_distinct_cpu_capacity"] = _can_schedule_distinct(masks) + if not node["scheduler_main_threads_have_distinct_cpu_capacity"]: + failures.append(f"DP scheduler main threads on {hostname} cannot run on distinct logical CPUs") + if not launchers: + missing.append("launcher CPU pool observations are missing") + return { + "status": "failed" if failures else "incomplete" if missing else "qualified", + "policy_required": required, + "policy": {"executor": plan.options.executor, "cpus_per_task": requested, "cpu_bind": binding}, + "expected_nodes": expected_nodes, + "missing_evidence": sorted(set(missing)), + "failures": sorted(set(failures)), + "observations": observations, + "nodes": nodes, + "scope": "Linux logical CPU masks at launcher startup, worker warm-up completion and scheduler initialization; " + "per-thread snapshots, not utilization, exclusive-core, NUMA or timed-forward affinity qualification", + } diff --git a/python/aisimulate/collector/fpm_forward/database.py b/python/aisimulate/collector/fpm_forward/database.py index 920d9a3cc..631ac3011 100644 --- a/python/aisimulate/collector/fpm_forward/database.py +++ b/python/aisimulate/collector/fpm_forward/database.py @@ -22,7 +22,7 @@ from aisimulate.fpm_contract import FPM_RESOLVED_CONFIG_GLOB from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY -from .native_artifact import validate_native_collection +from .native_artifact import NativeCollection, _rank_artifacts, select_native_measurements, validate_native_collection from .planner import FPMCell, FPMCollectionPlan, backend_identity_columns logger = logging.getLogger(__name__) @@ -97,14 +97,123 @@ def _validate_backend_markers(cell: FPMCell, cell_dir: Path) -> None: raise ValueError(f"backend marker mismatch in {path}: {mismatches}") +def _full_attention_group(group: object) -> bool: + if not isinstance(group, dict): + return False + if group.get("type") == "FullAttentionSpec": + # FullAttentionSpec also represents sliding/chunked attention when + # vLLM disables its hybrid allocator; the type name is insufficient. + # https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/kv_cache_interface.py#L235 + return ( + "specs" not in group + and "sliding_window" in group + and group["sliding_window"] is None + and "attention_chunk_size" in group + and group["attention_chunk_size"] is None + and group.get("non_causal") is False + ) + if group.get("type") == "UniformTypeKVCacheSpecs": + specs = group.get("specs") + return ( + isinstance(specs, list) + and bool(specs) + and all( + isinstance(spec, dict) and spec.get("type") == "FullAttentionSpec" and _full_attention_group(spec) + for spec in specs + ) + ) + return False + + +def _dense_synthetic_kv_eligible( + plan: FPMCollectionPlan, cell: FPMCell, raw_root: Path, collection: NativeCollection +) -> bool: + """Corroborate the producer's deliberate dense skip with initialized cache evidence.""" + meta = collection.kvwarm_meta or {} + if ( + meta.get("enabled") is not True + or meta.get("warm_eligible") is not False + or meta.get("skip_reason") != "dense_model_content_insensitive" + or plan.capability.is_moe is not False + or plan.capability.attention_source != "dense_attention" + or plan.capability.model_config is None + or cell.parallel_strategy not in {"tp", "single"} + or any(getattr(cell.topology, name) != 1 for name in ("pp", "dp", "moe_tp", "moe_ep", "cp")) + ): + return False + source = plan.capability.model_config.payload + text = source.get("text_config") + # Dynamo checks both configs for experts before considering recurrent state: + # https://github.com/ai-dynamo/dynamo/blob/92cf45e2e6707ac32b1fee275f949270fdd61200/components/src/dynamo/vllm/instrumented_scheduler.py#L5585 + # Its skip reason alone therefore does not establish dense full attention. + if any( + config.get(key) + for config in (source, text) + if isinstance(config, dict) + for key in ("num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts") + ): + return False + model_type, architectures = source.get("model_type"), source.get("architectures") + if not isinstance(model_type, str) or not model_type or not isinstance(architectures, list) or not architectures: + return False + expected_parallel = { + "tensor_parallel_size": cell.topology.tp, + "pipeline_parallel_size": cell.topology.pp, + "data_parallel_size": cell.topology.dp, + "prefill_context_parallel_size": cell.topology.cp, + "decode_context_parallel_size": cell.topology.cp, + "enable_expert_parallel": False, + } + artifacts = _rank_artifacts(raw_root) + for _path, payload in artifacts: + engine = payload.get("engine") + if not isinstance(engine, dict) or engine.get("capture_error") is not None: + return False + model, parallel = engine.get("model"), engine.get("parallel") + versions, cache = engine.get("versions"), engine.get("kv_cache") + if not all(isinstance(block, dict) for block in (model, parallel, versions, cache)): + return False + if ( + model.get("model_type") != model_type + or model.get("architectures") != architectures + or versions.get("vllm") != collection.backend_version + or any( + type(parallel.get(key)) is not type(value) or parallel.get(key) != value + for key, value in expected_parallel.items() + ) + or parallel.get("data_parallel_rank") != payload["dp"]["rank"] + ): + return False + if model.get("model") != plan.model_path: + from .execution_evidence import _effective_launch + + try: + launch_model = _effective_launch(plan, raw_root.parent).get("--model") + except (OSError, TypeError, ValueError, KeyError): + return False + if model.get("model") != launch_model: + return False + # Initialized group evidence, including vLLM's uniform-spec wrapper: + # https://github.com/ai-dynamo/dynamo/blob/92cf45e2e6707ac32b1fee275f949270fdd61200/components/src/dynamo/vllm/instrumented_scheduler.py#L291 + groups = cache.get("groups") + if not isinstance(groups, list) or not groups or not all(_full_attention_group(group) for group in groups): + return False + return bool(artifacts) + + def aggregate_cell( plan: FPMCollectionPlan, cell: FPMCell, cell_dir: Path, *, expected_attempt_id: str, + dense_synthetic_kv: bool = True, ) -> list[dict[str, Any]]: - """Validate native rank artifacts and take max-rank latency per grid point.""" + """Validate native rank artifacts and take max-rank latency per grid point. + + Revalidating an existing publication uses its explicit dense-KV policy; + an older commit record retains the historical fake-fallback classification. + """ if not expected_attempt_id: raise ValueError(f"cannot aggregate {cell.cell_id} without an expected Collector attempt identity") @@ -118,45 +227,7 @@ def aggregate_cell( backend_version = collection.backend_version capability = plan.capability - # The steady-state decode policy clamps every per-sequence context below - # the measurable minimum up to it, so several requested plan points can - # collapse onto one achieved physical coordinate (e.g. batch=3 with - # requested total-kv 3/4/5/6 all measure total-kv 6). Repeated samples of - # one coordinate would violate the database's unique-key contract, so keep - # exactly one per coordinate: the native (unclamped) sample when present, - # otherwise the clamped sample with the lowest benchmark_id. Two native - # samples on one coordinate remain a hard error — the grid itself - # guarantees native coordinates are unique. - grouped: dict[tuple[str, int, int, int], list[Any]] = {} - for measurement in collection.points: - point = measurement.point - key = ( - str(point["point_type"]), - int(point["batch_size"]), - int(point["total_prefill_tokens"]), - int(point["total_kv_read_tokens"]), - ) - grouped.setdefault(key, []).append(measurement) - selected: list[Any] = [] - dropped_clamped = 0 - for key, measurements in grouped.items(): - natives = [m for m in measurements if "context_clamped" not in (m.point.get("sample_reasons") or ())] - if len(natives) > 1: - raise ValueError( - f"conflicting FPM measurements for physical coordinate {key} in " - f"{cell.cell_id}: {len(natives)} unclamped samples share one key" - ) - if natives: - selected.append(natives[0]) - else: - selected.append(min(measurements, key=lambda m: int(m.point["benchmark_id"]))) - dropped_clamped += len(measurements) - 1 - if dropped_clamped: - logger.info( - "FPM %s: consolidated %d context-clamped duplicate sample(s) onto their achieved physical coordinates", - cell.cell_id, - dropped_clamped, - ) + selected = select_native_measurements(collection, cell_id=cell.cell_id) # KV seed regime (schema v6 additive column): a per-row RECORD of the # engine's measurement protocol for this point, consumed by the modeling @@ -169,6 +240,13 @@ def aggregate_cell( # legitimate topology regime. kvwarm_meta = collection.kvwarm_meta approved_skip_reasons = {"moe_tp_balanced_by_construction"} + dense_synthetic = ( + dense_synthetic_kv + and cell.workload_kind == "decode" + and kvwarm_meta is not None + and kvwarm_meta.get("skip_reason") == "dense_model_content_insensitive" + and _dense_synthetic_kv_eligible(plan, cell, cell_dir / "raw", collection) + ) def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: if phase == "prefill" and (not cell.execution_identity[0] or int(point["total_kv_read_tokens"]) == 0): @@ -184,6 +262,8 @@ def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: raise ValueError(f"native point reports real KV under warm-up skip_reason={skip_reason!r}") return "real_kv" if "kvwarm_fake_fallback" in reasons: + if dense_synthetic: + return "skip:dense_model_content_insensitive" return "fake_fallback" if skip_reason: return "fake_fallback" @@ -277,6 +357,30 @@ def _sha256(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() +def published_dense_synthetic_cells(metadata: dict[str, Any]) -> frozenset[str]: + """Read the explicit classification policy without upgrading historical rows.""" + marker = metadata.get("kv_seed_classification") + if marker is None: + return frozenset() + if not isinstance(marker, dict) or marker.get("policy") != "dense_full_attention_v1": + raise ValueError("unsupported FPM KV seed classification policy") + cells = marker.get("cells") + if ( + not isinstance(cells, list) + or not cells + or any(not isinstance(cell, str) or not cell for cell in cells) + or len(cells) != len(set(cells)) + ): + raise ValueError("FPM KV seed classification requires distinct nonempty cell IDs") + return frozenset(cells) + + +def _validate_kv_seed_classification(metadata: dict[str, Any], rows: list[dict[str, Any]]) -> None: + promoted = {row["cell_id"] for row in rows if row.get("kv_seed_regime") == "skip:dense_model_content_insensitive"} + if published_dense_synthetic_cells(metadata) != promoted: + raise ValueError("FPM KV seed classification policy differs from published rows") + + def validate_formal_database_commit( parquet_path: Path, metadata_path: Path, @@ -302,6 +406,12 @@ def validate_formal_database_commit( payload = json.loads(metadata_path.read_text()) if not isinstance(payload, dict): raise TypeError(f"FPM database commit record must be a mapping: {metadata_path}") + if payload.get("schema_name") == "aic_fpm_forward_perf" and payload.get("schema_version") == 6: + raise ValueError( + "historical schema-6 FPM publications cannot be finalized by this collector; schema 7 is required. " + "Preserve the existing artifacts. Automatic migration is unsupported; " + "use a fresh output directory for any new collection." + ) expected = { "schema_name": "aic_fpm_forward_perf", "schema_version": 7, @@ -334,6 +444,7 @@ def validate_formal_database_commit( f"parquet={parquet.metadata.num_rows}, metadata={row_count!r}" ) rows = parquet.read().to_pylist() + _validate_kv_seed_classification(payload, rows) identities = _run_identities_by_cell(rows, source="committed") grouped: dict[str, list[dict[str, Any]]] = {} for row in rows: @@ -599,6 +710,7 @@ def write_formal_database( f"contract (missing columns: {missing}); publish to a clean destination: {parquet_path}" ) existing_rows = table.to_pylist() + _validate_kv_seed_classification(committed, existing_rows) if committed.get("schema_version") == 6: for row in existing_rows: row.update(zip(EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, strict=True)) @@ -678,6 +790,15 @@ def write_formal_database( "backend": plan.backend, "backend_version": version, } + dense_cells = sorted( + { + row["cell_id"] + for row in merged + if row.get("kv_seed_regime") == "skip:dense_model_content_insensitive" + } + ) + if dense_cells: + metadata["kv_seed_classification"] = {"policy": "dense_full_attention_v1", "cells": dense_cells} temporary_metadata.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n") os.replace(temporary, parquet_path) # Metadata is the commit record: readers must validate its parquet diff --git a/python/aisimulate/collector/fpm_forward/entry.py b/python/aisimulate/collector/fpm_forward/entry.py index 3abbe1b3c..60ce81a08 100644 --- a/python/aisimulate/collector/fpm_forward/entry.py +++ b/python/aisimulate/collector/fpm_forward/entry.py @@ -150,6 +150,24 @@ def resolve_inputs(args: argparse.Namespace, case_plan) -> ResolvedFPMInputs: if not isinstance(payload, dict): raise ValueError("--fpm-model-profile must contain a JSON or YAML mapping") fpm_profile = load_fpm_profile(payload) + runtime_paths = tuple( + getattr(args, name, None) + for name in ("fpm_runtime_instrumentation", "fpm_runtime_launch", "fpm_runtime_configuration") + ) + runtime_inputs = {} + if any(value is not None for value in runtime_paths): + if not all(isinstance(value, str) and value for value in runtime_paths): + raise ValueError( + "--fpm-runtime-instrumentation, --fpm-runtime-launch and " + "--fpm-runtime-configuration are required together" + ) + from .runtime_instrumentation import read_json + + runtime_inputs = { + "runtime_instrumentation": runtime_paths[0], + "runtime_launch": read_json(Path(runtime_paths[1]).expanduser()), + "runtime_configuration": runtime_paths[2], + } plan = build_collection_plan( backend=args.backend, model_path=case_plan.model_path, @@ -162,6 +180,7 @@ def resolve_inputs(args: argparse.Namespace, case_plan) -> ResolvedFPMInputs: fpm_profile=fpm_profile, collector_config={}, generator_overrides=generator_overrides, + **runtime_inputs, ) return plan, generator_overrides diff --git a/python/aisimulate/collector/fpm_forward/execution_evidence.py b/python/aisimulate/collector/fpm_forward/execution_evidence.py new file mode 100644 index 000000000..2df117278 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/execution_evidence.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inspect observed execution separately from requested benchmark settings.""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +import shlex +from collections import Counter +from pathlib import Path, PurePosixPath +from typing import Any + +import yaml + +from .config import with_kv_warmup_defaults +from .cpu_affinity import inspect_cpu_affinity +from .native_artifact import NativeCollection, _rank_artifacts +from .planner import FPMCell, FPMCollectionPlan, _canonical_hash +from .runtime.fpm_memory_observer import EXECUTION_SUPPORTED_VERSIONS +from .runtime_memory import _dtype, _validate_graph, _validate_precision, _validate_revision + +_GRAPH_FIELDS = {"mode", "cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"} +_CONFIG_TYPES = { + "model_config": { + "model": (str,), + "revision": (str, type(None)), + "dtype": (str,), + "quantization": (str, type(None)), + "max_model_len": (int,), + "enforce_eager": (bool,), + }, + "cache_config": { + "cache_dtype": (str,), + "gpu_memory_utilization": (int, float), + "enable_prefix_caching": (bool,), + "kv_cache_memory_bytes": (int, type(None)), + "num_gpu_blocks_override": (int, type(None)), + }, + "scheduler_config": {"max_num_batched_tokens": (int,), "max_num_seqs": (int,), "async_scheduling": (bool,)}, + "parallel_config": { + "tensor_parallel_size": (int,), + "pipeline_parallel_size": (int,), + "data_parallel_size": (int,), + "data_parallel_size_local": (int,), + "data_parallel_external_lb": (bool,), + "enable_expert_parallel": (bool,), + "enable_eplb": (bool,), + "decode_context_parallel_size": (int,), + "prefill_context_parallel_size": (int,), + }, + "kernel_config": {"moe_backend": (str,)}, +} + + +class _MissingExecutionEvidence(ValueError): + pass + + +def _requested_scheduling(flags: dict[str, Any] | None) -> dict[str, Any]: + if flags is None: + return {"source": "unreported", "async_scheduling": None} + choices = [ + name == "--async-scheduling" for name in flags if name in {"--async-scheduling", "--no-async-scheduling"} + ] + return { + "source": "generated_flag" if choices else "runtime_default", + "async_scheduling": choices[-1] if choices else None, + } + + +def _effective_launch(plan: FPMCollectionPlan, cell_dir: Path) -> dict[str, Any]: + """Read the generator's literal argv array; never evaluate archived shell.""" + try: + for name in ("fpm_env.sh", "collector-runtime-env.sh"): + (cell_dir / name).read_bytes() + deployment = json.loads((cell_dir.parent.parent / "generator-overrides.json").read_text()) + request = json.loads((cell_dir / "generator-request.json").read_text()) + script = (cell_dir / "run.sh").read_text() + except FileNotFoundError as error: + raise _MissingExecutionEvidence(f"frozen generated launch is missing: {error.filename}") from error + if not isinstance(deployment, dict): + raise ValueError("archived deployment inputs must be an object") + mount = deployment.get("K8sConfig", {}) + if not isinstance(mount, dict): + raise ValueError("archived K8sConfig must be an object") + if _canonical_hash(with_kv_warmup_defaults(deployment)) != plan.generator_config_sha256: + raise ValueError("archived deployment inputs differ from the source plan") + expected_model = plan.model_path + if mount.get("k8s_model_path_in_pvc"): + if not mount.get("k8s_pvc_mount_path"): + raise ValueError("checkpoint mount has no explicit container mount path") + expected_model = str(PurePosixPath(mount["k8s_pvc_mount_path"]) / mount["k8s_model_path_in_pvc"]) + if not isinstance(request, dict): + raise ValueError("generated request must be an object") + service = request.get("ServiceConfig") + if not isinstance(service, dict): + raise ValueError("generated ServiceConfig must be an object") + if service.get("model_path") != expected_model: + raise ValueError("generated model path differs from the plan and its explicit checkpoint mount") + commands = re.findall(r"^engine_command=\((.*)\)$", script, re.MULTILINE) + if len(commands) != 1: + raise _MissingExecutionEvidence("frozen launch does not contain one literal generated engine_command array") + argv = shlex.split(commands[0]) + if argv[:3] != ["python3", "-m", "dynamo.vllm"]: + raise _MissingExecutionEvidence("frozen launch uses an unsupported engine command") + # The current generator emits one shell-quoted token per argument. Read + # only these literal tokens, including argparse's --option=value form. + flags = {} + for index, token in enumerate(argv): + if not token.startswith("--"): + continue + name, separator, value = token.partition("=") + flags.pop(name, None) + flags[name] = ( + value + if separator + else argv[index + 1] + if index + 1 < len(argv) and not argv[index + 1].startswith("--") + else True + ) + actual_model = flags.get("--model") + if isinstance(actual_model, str) and ("$" in actual_model or "`" in actual_model): + raise _MissingExecutionEvidence("generated model path is not a literal runtime setting") + if actual_model != expected_model: + raise ValueError("generated engine model differs from the frozen plan and checkpoint mount") + return flags + + +def _matches_pinned_hub_snapshot(expected: str, model: dict[str, Any], revisions: set[Any]) -> bool: + """Recognize a resolved Hub cache identity, not an arbitrary local alias. + + Offline vLLM can replace the Hub ID with its cached snapshot directory. + Require its full repository identity and immutable commit to agree with + both the source pins and the loaded config; this is not a weight hash. + """ + if ( + re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9._-]*/[A-Za-z0-9_][A-Za-z0-9._-]*", expected) is None + or "--" in expected + or ".." in expected + or any(part.endswith((".", "-", ".git")) for part in expected.split("/")) + ): + return False + observed = model["model"] + path = PurePosixPath(observed) + if path.anchor != "/" or str(path) != observed or ".." in path.parts or "\\" in observed or "\x00" in observed: + return False + if len(path.parts) < 4 or path.parts[-3:-1] != (f"models--{expected.replace('/', '--')}", "snapshots"): + return False + commit = path.name + if re.fullmatch(r"[0-9a-fA-F]{40}", commit) is None or not revisions: + return False + if any( + not isinstance(revision, str) + or re.fullmatch(r"[0-9a-fA-F]{40}", revision) is None + or revision.lower() != commit.lower() + for revision in revisions + ): + return False + loaded = model.get("loaded_config_commit_hash") + return isinstance(loaded, str) and loaded.lower() == commit.lower() + + +def _inspect_config(payload, cell, plan, flags, missing, failures): + config = payload.get("resolved_config") + if not isinstance(config, dict): + missing.append("resolved runtime configuration is unreported") + config = {} + valid = {} + for section, fields in _CONFIG_TYPES.items(): + values = config.get(section) + if not isinstance(values, dict): + missing.append(f"runtime {section} is unreported") + continue + valid[section] = {} + for name, types in fields.items(): + label = f"{section}.{name}" + if name not in values: + missing.append(f"runtime {label} is unreported") + elif type(values[name]) not in types or (isinstance(values[name], str) and not values[name]): + failures.append(f"runtime {label} has an invalid type or empty value") + elif type(values[name]) in (int, float) and (not math.isfinite(values[name]) or values[name] <= 0): + failures.append(f"runtime {label} must be positive and finite") + else: + valid[section][name] = values[name] + model, cache, parallel = (valid.get(key, {}) for key in ("model_config", "cache_config", "parallel_config")) + if cache.get("gpu_memory_utilization", 1) > 1: + failures.append("runtime GPU memory utilization must not exceed one") + if parallel.get("data_parallel_size_local", 1) > parallel.get("data_parallel_size", 1): + failures.append("runtime local DP size exceeds global DP size") + expected = { + "model_config": {"enforce_eager": plan.options.enforce_eager}, + "parallel_config": { + "tensor_parallel_size": cell.topology.tp, + "pipeline_parallel_size": cell.topology.pp, + "data_parallel_size": cell.topology.dp, + "decode_context_parallel_size": cell.topology.cp, + "prefill_context_parallel_size": 1, + "enable_expert_parallel": cell.topology.moe_ep > 1, + "enable_eplb": bool(cell.backend_policy.aic_fields.get("enable_eplb", False)), + }, + } + if plan.options.vllm_max_model_len > 0: + expected["model_config"]["max_model_len"] = plan.options.vllm_max_model_len + if plan.options.gpu_memory_utilization is not None: + expected["cache_config"] = {"gpu_memory_utilization": plan.options.gpu_memory_utilization} + for section, fields in expected.items(): + for name, value in fields.items(): + if name in valid.get(section, {}) and valid[section][name] != value: + failures.append(f"runtime {section}.{name} differs from the frozen plan") + expected_revision = plan.fpm_profile.model_revision if plan.fpm_profile is not None else None + revisions = {value for value in (expected_revision, (flags or {}).get("--revision")) if value is not None} + launch_fields = { + "--model": ("model_config", "model", str), + "--max-model-len": ("model_config", "max_model_len", int), + "--gpu-memory-utilization": ("cache_config", "gpu_memory_utilization", float), + "--kv-cache-memory-bytes": ("cache_config", "kv_cache_memory_bytes", int), + "--num-gpu-blocks-override": ("cache_config", "num_gpu_blocks_override", int), + "--max-num-batched-tokens": ("scheduler_config", "max_num_batched_tokens", int), + "--max-num-seqs": ("scheduler_config", "max_num_seqs", int), + "--tensor-parallel-size": ("parallel_config", "tensor_parallel_size", int), + "--pipeline-parallel-size": ("parallel_config", "pipeline_parallel_size", int), + "--data-parallel-size": ("parallel_config", "data_parallel_size", int), + "--decode-context-parallel-size": ("parallel_config", "decode_context_parallel_size", int), + "--prefill-context-parallel-size": ("parallel_config", "prefill_context_parallel_size", int), + "--moe-backend": ("kernel_config", "moe_backend", str), + } + if flags is not None: + for name in ("kv_cache_memory_bytes", "num_gpu_blocks_override"): + if f"--{name.replace('_', '-')}" not in flags and cache.get(name) is not None: + failures.append(f"runtime cache_config.{name} differs from automatic generated cache sizing") + for flag, (section, name, convert) in launch_fields.items(): + if flag not in flags or name not in valid.get(section, {}): + continue + raw = flags[flag] + if type(raw) is not str or "$" in raw or "`" in raw: + missing.append(f"generated {flag} is not a literal runtime setting") + continue + try: + value = convert(raw) + except ValueError: + failures.append(f"generated {flag} is not a valid runtime setting") + continue + if (flag == "--max-model-len" and value == -1) or (flag == "--moe-backend" and value == "auto"): + continue + if valid[section][name] != value: + if flag == "--model" and _matches_pinned_hub_snapshot(value, config["model_config"], revisions): + continue + failures.append(f"runtime {section}.{name} differs from generated {flag}") + for flag, section, name in ( + ("async-scheduling", "scheduler_config", "async_scheduling"), + ("enable-prefix-caching", "cache_config", "enable_prefix_caching"), + ("enforce-eager", "model_config", "enforce_eager"), + ): + choices = [key == f"--{flag}" for key in flags if key in {f"--{flag}", f"--no-{flag}"}] + if choices and name in valid.get(section, {}) and valid[section][name] != choices[-1]: + failures.append(f"runtime {section}.{name} differs from generated flags") + for flag, actual in (("--dtype", model.get("dtype")), ("--kv-cache-dtype", cache.get("cache_dtype"))): + if flag in flags and flags[flag] != "auto" and actual is not None and _dtype(actual) != _dtype(flags[flag]): + failures.append(f"runtime precision differs from generated {flag}") + if "--quantization" in flags and "quantization" in model and model["quantization"] != flags["--quantization"]: + failures.append("runtime quantization differs from generated --quantization") + if not revisions: + missing.append("source has no pinned model revision") + raw_model = config.get("model_config") + loaded = raw_model.get("loaded_config_commit_hash") if isinstance(raw_model, dict) else None + if loaded is not None and (not isinstance(loaded, str) or re.fullmatch(r"[0-9a-fA-F]{40}", loaded) is None): + failures.append("runtime loaded model config commit hash is malformed") + for revision in revisions: + observed = model.get("revision") + immutable = isinstance(revision, str) and re.fullmatch(r"[0-9a-fA-F]{40}", revision) is not None + if observed == revision or (immutable and isinstance(loaded, str) and loaded.lower() == revision.lower()): + pass + elif observed is None: + missing.append(f"runtime model revision has no evidence for {revision!r}") + else: + failures.append("runtime model revision differs from the pinned source revision") + if immutable and isinstance(config.get("model_config"), dict): + try: + _validate_revision(config, revision) + except ValueError as error: + failures.append(str(error)) + if "dtype" in model and "cache_dtype" in cache: + actual = model["dtype"] if cache["cache_dtype"] == "auto" else cache["cache_dtype"] + if _dtype(actual) != cell.kv_cache_dtype: + failures.append("runtime KV precision differs from the planned identity") + method = model.get("quantization") + quant = config.get("quantization_config") + if "quantization_config" not in config: + missing.append("runtime quantization_config is unreported") + elif "quantization" in model and "dtype" in model: + required = {"type"} + if method in {"modelopt", "modelopt_fp4"}: + required.add("quant_method") + elif method == "fp8": + required.update({"activation_scheme", "weight_block_size"}) + if method is not None and quant is None: + failures.append("runtime quantization method contradicts its absent quantization configuration") + elif method is not None and (not isinstance(quant, dict) or not required.issubset(quant)): + missing.append("runtime quantization family configuration is incomplete") + elif method not in {None, "modelopt", "modelopt_fp4", "fp8"}: + missing.append(f"runtime quantization mapping is unaudited: {method}") + else: + try: + _validate_precision(config, cell) + except ValueError as error: + failures.append(str(error)) + graphs = [] + for label, graph in ( + ("graph_config", payload.get("graph_config")), + ("compilation_config", config.get("compilation_config")), + ): + if not isinstance(graph, dict) or not _GRAPH_FIELDS.issubset(graph): + missing.append(f"runtime {label} is incomplete") + continue + if "enforce_eager" not in model or "max_num_batched_tokens" not in valid.get("scheduler_config", {}): + continue + try: + graphs.append( + _validate_graph( + graph, + eager=model["enforce_eager"], + tokens=valid["scheduler_config"]["max_num_batched_tokens"], + allow_unused_captures=True, + ) + ) + except ValueError as error: + failures.append(str(error)) + if len(graphs) == 2 and graphs[0] != graphs[1]: + failures.append("worker graph_config contradicts its resolved_config.compilation_config") + + +def file_evidence(path: Path) -> dict[str, Any]: + raw = path.read_bytes() + return {"path": str(path.resolve()), "sha256": hashlib.sha256(raw).hexdigest(), "size_bytes": len(raw)} + + +def inspect_execution_evidence( + cell: FPMCell, raw_root: Path, collection: NativeCollection, *, plan: FPMCollectionPlan, evidence_version: int = 2 +) -> dict[str, Any]: + """Require actual worker observations; native point expectations stay labelled. + + The caller first validates native point/rank/timing integrity. Missing + execution observations leave qualification incomplete without invalidating + historical timing artifacts or manufacturing a requested backend identity. + """ + if type(evidence_version) is not int or evidence_version not in {1, 2}: + raise ValueError("unsupported execution evidence version") + missing = [] + failures = [] + flags = None + try: + flags = _effective_launch(plan, raw_root.parent) + except _MissingExecutionEvidence as error: + missing.append(str(error)) + except (ValueError, TypeError) as error: + failures.append(str(error)) + if collection.backend_version != plan.capability.aic_database_version: + failures.append("observed backend version differs from the frozen plan") + if collection.backend_version not in EXECUTION_SUPPORTED_VERSIONS: + missing.append(f"execution configuration inspection is unaudited for {collection.backend_version}") + workers = [] + ranks = set() + scheduling = {"requested": _requested_scheduling(flags), "observed_workers": []} + for path in sorted(raw_root.glob("**/fpm-execution-worker-*.json")): + payload = json.loads(path.read_text()) + if not isinstance(payload, dict): + raise ValueError(f"runtime execution evidence must be an object: {path}") + provenance = payload.get("collector_provenance", {}) + if not isinstance(provenance, dict): + raise ValueError(f"runtime execution collector_provenance must be an object: {path}") + if ( + payload.get("schema_name") != "aisimulate_fpm_runtime_execution" + or payload.get("schema_version") != 1 + or provenance.get("plan_sha256") != plan.sha256 + or provenance.get("cell_id") != cell.cell_id + or provenance.get("attempt_id") != collection.collector_attempt_id + or payload.get("backend_version") != collection.backend_version + ): + raise ValueError(f"runtime execution evidence has a different plan, cell, attempt or version: {path}") + rank = tuple(payload.get(key) for key in ("dp_rank", "tp_rank", "pp_rank")) + if any(type(value) is not int or value < 0 for value in rank) or rank in ranks: + raise ValueError(f"runtime execution worker ranks are invalid or duplicated: {path}") + ranks.add(rank) + item = {**payload, "source": file_evidence(path)} + workers.append(item) + if payload.get("status") != "observed": + missing.append(f"worker {rank}: {payload.get('error', 'execution observation unresolved')}") + continue + groups = payload.get("attention_groups") + if ( + not isinstance(groups, list) + or not groups + or any( + not isinstance(group, dict) + or not isinstance(group.get("backend_class"), str) + or not group["backend_class"] + or not isinstance(group.get("layer_names"), list) + or not group["layer_names"] + or any(not isinstance(name, str) or not name for name in group["layer_names"]) + for group in groups + ) + ): + failures.append(f"runtime attention group observation is malformed: {path}") + _inspect_config(payload, cell, plan, flags, missing, failures) + config = payload.get("resolved_config") + scheduler = config.get("scheduler_config") if isinstance(config, dict) else None + scheduling["observed_workers"].append( + { + "rank": list(rank), + "async_scheduling": scheduler.get("async_scheduling") if isinstance(scheduler, dict) else None, + "source": item["source"], + } + ) + expected = { + (dp, tp, pp) + for dp in range(cell.topology.dp) + for tp in range(cell.topology.tp) + for pp in range(cell.topology.pp) + } + if ranks - expected: + raise ValueError("runtime execution observations contain unexpected worker ranks") + if ranks != expected: + missing.append(f"worker execution observations missing for ranks {sorted(expected - ranks)}") + modes = { + item["async_scheduling"] for item in scheduling["observed_workers"] if type(item["async_scheduling"]) is bool + } + if evidence_version == 2 and len(modes) > 1: + failures.append("workers disagree on effective async_scheduling") + scheduling["effective_async_scheduling"] = next(iter(modes)) if len(modes) == 1 else None + scheduling["status"] = ( + "mismatch" + if len(modes) > 1 + else "observed" + if len(scheduling["observed_workers"]) == len(expected) + and all(type(item["async_scheduling"]) is bool for item in scheduling["observed_workers"]) + else "unestablished" + ) + native_graphs = [] + for path, payload in _rank_artifacts(raw_root): + native_graphs.append( + {"dp_rank": payload["dp"]["rank"], "config": payload.get("cudagraph"), "source": file_evidence(path)} + ) + if not isinstance(payload.get("cudagraph"), dict): + missing.append(f"native graph configuration unreported for {path.name}") + regimes = Counter(measurement.kv_seed_regime or "unreported" for measurement in collection.points) + if "unreported" in regimes: + missing.append("per-point KV initialization regime is unreported") + cpu_nodes, cpu_manifest, cpu_geometry_missing, cpu_geometry_failure = None, None, None, None + if getattr(plan.options, "slurm_cpus_per_task", None) is not None: + from .runner import FPM_MANIFEST_FILENAME, _expected_nodes + + manifest = raw_root.parent / FPM_MANIFEST_FILENAME + try: + cpu_manifest = file_evidence(manifest) + cpu_nodes = _expected_nodes(manifest) + except OSError as error: + cpu_geometry_missing = str(error) + except (TypeError, ValueError, KeyError, yaml.YAMLError) as error: + cpu_geometry_failure = str(error) + cpu = inspect_cpu_affinity(cell, raw_root, collection, plan=plan, expected_nodes=cpu_nodes) + cpu["launch_geometry_source"] = cpu_manifest + if cpu_geometry_missing is not None: + cpu["missing_evidence"].append(f"generated launch geometry: {cpu_geometry_missing}") + if cpu_geometry_failure is not None: + cpu["failures"].append(f"generated launch geometry: {cpu_geometry_failure}") + cpu["status"] = "failed" + if cpu["policy_required"]: + missing.extend(f"CPU affinity: {message}" for message in cpu["missing_evidence"]) + failures.extend(f"CPU affinity: {message}" for message in cpu["failures"]) + dispatch = "unreported" + if evidence_version == 2: + from .measurement_evidence import extract_measurement_evidence + + measurements = extract_measurement_evidence(raw_root, collection) + missing.extend(f"measurement evidence: {reason}" for reason in measurements["reasons"]) + for point in measurements["points"].values(): + missing.extend(f"measurement evidence: {reason}" for reason in point["reasons"]) + if measurements["warmup_evidence"]: + dispatch = ( + "observed" if "observed_per_call_graph_dispatch" not in measurements["unobserved"] else "incomplete" + ) + return { + "status": "failed" if failures else "incomplete" if missing else "qualified", + "missing_evidence": sorted(set(missing)), + "failures": sorted(set(failures)), + "observed_workers": workers, + "native_graph_config": native_graphs, + "kv_seed_regime_counts": dict(regimes), + "cpu_affinity": cpu, + **({"scheduling": scheduling} if evidence_version == 2 else {}), + "per_point_dispatch": dispatch, + "scope": "initialized attention backends and resolved graph configuration; " + + ( + "per-sample native dispatch observations" + if dispatch == "observed" + else "no complete per-point dispatch trace" + if evidence_version == 2 + else "no per-point dispatch trace" + ), + } diff --git a/python/aisimulate/collector/fpm_forward/measurement_evidence.py b/python/aisimulate/collector/fpm_forward/measurement_evidence.py new file mode 100644 index 000000000..3529a8a11 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/measurement_evidence.py @@ -0,0 +1,699 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare recorded benchmark conditions without using latency as evidence. + +The consumer contract is Dynamo's measurement_protocol/benchmark_measurement +schema 1 (ai-dynamo/dynamo, commit 92cf45e2e6707ac32b1fee275f949270fdd61200). +Prompt hashes observe admission inputs, not generated tokens or cache tensors. +Effective worker configuration is checked separately by execution validation. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any + +from .native_artifact import NativeCollection, _rank_artifacts, select_native_measurements + +_UNOBSERVED = [ + "sampled_continuation_token_ids", + "kv_cache_tensors", + "recurrent_state_tensors", + "execution_history_equivalence", + "observed_per_call_graph_dispatch", +] +_GRAPH_FIELDS = ("expected_cudagraph_mode", "expected_capture_size", "padding_tokens") +_RUNTIME_FIELDS = ("limits", "cudagraph", "recurrent_state") +_EXECUTION_CONTRACT = { + "schema_version": 1, + "sample_field": "benchmark_sample", + "warmup_field": "warmup_evidence", + "forward_index_scope": "rank_process", + "timing_clock": "process_local_monotonic", + "cudagraph_source": "ModelRunnerOutput.cudagraph_stats", + "warmup_scope": "request_completion_and_existing_shape_validation", +} + + +def _hash(value: Any) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +def _coordinates(point: dict[str, Any]) -> dict[str, Any]: + result = {key: point[key] for key in ("batch_size", "total_kv_read_tokens")} + if point["point_type"] == "prefill": + result["total_prefill_tokens"] = point["total_prefill_tokens"] + result.update({key: point[key] for key in ("partition", "rows") if point.get(key) is not None}) + return result + + +def _content_key(point: dict[str, Any]) -> str: + # These fields are the producer's schema-1 content identity, distinct from + # AISimulate's phase-cell-local coordinate key used by repeatability. + identity = {key: point[key] for key in ("point_type", "batch_size")} + identity.update({key: point.get(key, 0) for key in ("total_prefill_tokens", "total_kv_read_tokens")}) + identity.update({key: point.get(key) for key in ("partition", "rows")}) + return _hash(identity) + + +def _integer(value: Any, minimum: int = 0) -> bool: + return type(value) is int and value >= minimum + + +def _digest(value: Any) -> bool: + return isinstance(value, str) and len(value) == 64 and all(char in "0123456789abcdef" for char in value) + + +def _regime_gaps(point: dict[str, Any], regime: str | None) -> list[str]: + reasons = [f"missing {key}" for key in _GRAPH_FIELDS if key not in point] + if not isinstance(point.get("expected_cudagraph_mode"), str) or not point["expected_cudagraph_mode"]: + reasons.append("missing expected CUDA graph mode") + if point.get("expected_cudagraph_mode") not in {None, "NONE"}: + if not _integer(point.get("expected_capture_size"), 1): + reasons.append("missing expected CUDA graph capture size") + if not _integer(point.get("padding_tokens")): + reasons.append("missing expected CUDA graph padding") + if not regime: + reasons.append("missing KV initialization regime") + return reasons + + +def _validate_protocol(protocol: Any) -> None: + if not isinstance(protocol, dict): + raise ValueError("measurement_protocol must be a mapping") + fixed = { + "schema_version": 1, + "content_identity": "coordinate_rank_slot_v1", + "prompt_hash_encoding": "uint32_le", + "independent_repetitions": 1, + "timing_metric": "scheduler_wall_time", + "input_evidence_scope": "injected_prompt_token_ids", + } + if any(protocol.get(key) != value for key, value in fixed.items()) or not all( + _integer(protocol.get(key), 1) for key in ("schema_version", "independent_repetitions") + ): + raise ValueError("measurement_protocol has unsupported identity, timing, or repetition semantics") + if not all(isinstance(protocol.get(key), str) for key in ("content_seed", "synthetic_pool_tag")) or protocol.get( + "synthetic_content" + ) not in {"random", "sharegpt", "sharegpt_chain", "zeros"}: + raise ValueError("measurement_protocol has malformed content settings") + unobserved = protocol.get("unobserved") + if not isinstance(unobserved, list) or not all(isinstance(item, str) and item for item in unobserved): + raise ValueError("measurement_protocol must describe unobserved inputs") + preparation = protocol.get("preparation") + if ( + not isinstance(preparation, dict) + or not all( + _integer(preparation.get(key)) for key in ("warmup_iterations", "giant_kv_threshold", "giant_kv_repeats") + ) + or not all(type(preparation.get(key)) is bool for key in ("prefill_real_seed", "decode_real_kv_warmup")) + ): + raise ValueError("measurement_protocol has malformed preparation settings") + + +def _validate_prompts(prompts: Any) -> bool: + if not isinstance(prompts, dict) or prompts.get("status") not in {"recorded", "unavailable"}: + raise ValueError("benchmark_measurement has malformed prompt evidence") + requests = prompts.get("requests") + if prompts["status"] == "unavailable": + if prompts.get("sha256") is not None or requests != []: + raise ValueError("unavailable prompt evidence contains observed hashes") + return False + if not isinstance(requests, list) or not requests: + raise ValueError("recorded prompt evidence has no requests") + for index, request in enumerate(requests): + if ( + not isinstance(request, dict) + or not _integer(request.get("slot")) + or request["slot"] != index + or not _integer(request.get("num_tokens")) + or not _digest(request.get("sha256")) + ): + raise ValueError("recorded prompt evidence has malformed request hashes") + if prompts.get("sha256") != _hash(requests): + raise ValueError("recorded prompt evidence aggregate hash disagrees with requests") + return True + + +def _validate_estimate(measurement: dict[str, Any], fpm: dict[str, Any], rank: int, benchmark_id: int) -> None: + raw = measurement.get("raw_fpms") + expected = measurement.get("expected_internal_samples") + estimate = measurement.get("estimate") + if ( + not _integer(expected, 1) + or not isinstance(raw, list) + or not raw + or len(raw) > expected + or not isinstance(estimate, dict) + ): + raise ValueError("benchmark_measurement has malformed internal sample evidence") + if not _integer(fpm.get("counter_id")) or fpm["counter_id"] != benchmark_id: + raise ValueError("benchmark_measurement retained counter does not match its point") + for sample in raw: + if ( + not isinstance(sample, dict) + or type(sample.get("dp_rank")) is not int + or sample["dp_rank"] != rank + or not _integer(sample.get("counter_id")) + or sample["counter_id"] != benchmark_id + or isinstance(sample.get("wall_time"), bool) + or not isinstance(sample.get("wall_time"), int | float) + or not math.isfinite(sample["wall_time"]) + or sample["wall_time"] <= 0 + ): + raise ValueError("benchmark_measurement has invalid raw counter/rank/timing evidence") + method = estimate.get("method") + if method == "single_step" and len(raw) == 1: + indices = [0] + elif method == "last_step" and expected > 1 and len(raw) == expected: + indices = [len(raw) - 1] + elif method == "adjacent_upper_median" and expected > 2 and len(raw) > 1: + indices = list(range(1, len(raw))) + else: + raise ValueError("benchmark_measurement has inconsistent reduction method") + if estimate.get("raw_sample_indices") != indices or any( + type(index) is not int for index in estimate["raw_sample_indices"] + ): + raise ValueError("benchmark_measurement has inconsistent raw sample indices") + walls = sorted(raw[index]["wall_time"] for index in indices) + if not math.isclose(walls[len(walls) // 2], fpm["wall_time"], rel_tol=1e-9, abs_tol=1e-12): + raise ValueError("benchmark_measurement reduction disagrees with retained timing") + + +def _read_measurement( + row: dict[str, Any], *, rank: int, grid_digest: str, position: int, protocol: dict[str, Any] | None +) -> tuple[dict[str, Any] | None, list[str]]: + fpm = row["fpms"][0] + if protocol is None: + if "benchmark_measurement" in fpm: + raise ValueError("benchmark_measurement has no measurement_protocol") + return None, ["missing measurement_protocol and benchmark_measurement"] + measurement = fpm.get("benchmark_measurement") + if ( + not isinstance(measurement, dict) + or type(measurement.get("schema_version")) is not int + or measurement["schema_version"] != 1 + or type(measurement.get("dp_rank")) is not int + or measurement["dp_rank"] != rank + or measurement.get("point_key") != _content_key(row["point"]) + ): + raise ValueError("benchmark_measurement has missing or inconsistent rank/point identity") + preparation = measurement.get("preparation") + if ( + not isinstance(preparation, dict) + or preparation.get("grid_digest") != grid_digest + or not _integer(preparation.get("completed_points_before")) + or preparation["completed_points_before"] != position + or preparation.get("kv_seed_regime") != row.get("kv_seed_regime") + ): + raise ValueError("benchmark_measurement has inconsistent grid/order/KV preparation") + recorded = _validate_prompts(measurement.get("prompts")) + _validate_estimate(measurement, fpm, rank, row["point"]["benchmark_id"]) + return measurement, [] if recorded else [f"rank {rank}: injected prompt hashes unavailable"] + + +def _scheduling_observations(raw_root: Path, artifacts, collection) -> tuple[dict[str, Any], list[str]]: + """Read resolved scheduler/worker state, never derive it from launch defaults.""" + observations = {str(rank): [] for rank, *_ in collection.rank_timings} + for path, payload in artifacts: + engine = payload.get("engine") + if not isinstance(engine, dict) or engine.get("capture_error"): + continue + scheduler = engine.get("scheduler") + if isinstance(scheduler, dict) and "async_scheduling" in scheduler: + observations[str(payload["dp"]["rank"])].append( + {"async_scheduling": scheduler["async_scheduling"], "source": str(path)} + ) + for path in sorted(raw_root.glob("**/fpm-execution-worker-*.json")): + payload = json.loads(path.read_text()) + if not isinstance(payload, dict) or payload.get("status") != "observed": + continue + provenance = payload.get("collector_provenance", {}) + rank = str(payload.get("dp_rank")) + if ( + payload.get("schema_name") != "aisimulate_fpm_runtime_execution" + or payload.get("schema_version") != 1 + or not isinstance(provenance, dict) + or provenance.get("attempt_id") != collection.collector_attempt_id + or payload.get("backend_version") != collection.backend_version + or rank not in observations + ): + raise ValueError("scheduling observation has inconsistent worker provenance") + config = payload.get("resolved_config") + scheduler = config.get("scheduler_config") if isinstance(config, dict) else None + if isinstance(scheduler, dict) and "async_scheduling" in scheduler: + observations[rank].append({"async_scheduling": scheduler["async_scheduling"], "source": str(path)}) + result, missing = {}, [] + for rank, items in observations.items(): + if any(item["async_scheduling"] is not None and type(item["async_scheduling"]) is not bool for item in items): + raise ValueError("observed async_scheduling must be Boolean") + modes = {item["async_scheduling"] for item in items if item["async_scheduling"] is not None} + if len(modes) > 1: + raise ValueError(f"rank {rank}: native scheduler and workers disagree on async_scheduling") + result[rank] = {"async_scheduling": next(iter(modes)) if modes else None, "observations": items} + if not modes: + missing.append(f"rank {rank}: effective async_scheduling is unreported") + if len({item["async_scheduling"] for item in result.values()} - {None}) > 1: + raise ValueError("native ranks disagree on effective async_scheduling") + return result, missing + + +def _warmup_history(payload: dict[str, Any], protocol: dict[str, Any] | None): + if protocol is None or "execution_evidence" not in protocol: + return None, [] + if protocol["execution_evidence"] != _EXECUTION_CONTRACT: + raise ValueError("unsupported measurement execution_evidence contract") + history = payload.get("warmup_evidence") + if ( + not isinstance(history, dict) + or history.get("status") not in {"recorded", "unavailable"} + or not isinstance(history.get("records"), list) + or (history["status"] == "unavailable" and history["records"]) + ): + raise ValueError("invalid warmup_evidence status or records") + prefixes = [_hash([])] + digest = hashlib.sha256() + previous_completed = 0 + previous_start = 0 + for record in history["records"]: + if ( + not isinstance(record, dict) + or record.get("kind") not in {"global", "eager_shape", "real_prefix_seed", "real_prefix_shape"} + or record.get("status") not in {"running", "completed", "failed"} + or not isinstance(record.get("validation"), dict) + or record["validation"].get("status") not in {"not_performed", "passed", "failed"} + or not _integer(record.get("completed_points_before")) + ): + raise ValueError("malformed warmup_evidence record") + start, end = record.get("forward_index_start"), record.get("forward_index_end") + count = record.get("observed_forward_count") + if ( + not _integer(count) + or any(value is not None and not _integer(value) for value in (start, end)) + or (start is not None and end is not None and (end < start or count != end - start)) + or (record["status"] == "running") != (end is None) + or (start is None and (record["kind"] != "eager_shape" or record["status"] != "failed" or count != 0)) + or (record["validation"]["status"] == "passed" and record["status"] != "completed") + or (record["validation"]["status"] == "failed" and record["status"] != "failed") + ): + raise ValueError("warmup_evidence has contradictory completion or forward indices") + completed = record["completed_points_before"] + if ( + not previous_completed <= completed <= len(payload["results"]) + or (record["kind"] == "global" and completed != 0) + or (start is not None and start < previous_start) + ): + raise ValueError("warmup_evidence has impossible preparation order") + previous_completed = completed + if start is not None: + previous_start = start + digest.update(_hash(record).encode()) + prefixes.append(digest.hexdigest()) + return history, prefixes + + +def _observed_execution(measurement, fpm, history, prefixes, previous_forward): + """Validate within one rank/process; absolute clocks are not compared across runs.""" + if "benchmark_sample" in fpm: + raise ValueError("retained estimate cannot claim a single benchmark_sample") + observed, reasons = [], [] + preceding_forward = previous_forward + for index, raw in enumerate(measurement["raw_fpms"]): + sample = raw.get("benchmark_sample") + if ( + not isinstance(sample, dict) + or not _integer(sample.get("sample_index")) + or sample["sample_index"] != index + or not _integer(sample.get("forward_index")) + or sample["forward_index"] <= previous_forward + ): + raise ValueError("benchmark_sample has invalid sample or process-local forward index") + previous_forward = sample["forward_index"] + timing, graph = sample.get("timing"), sample.get("cudagraph") + if not isinstance(timing, dict) or timing.get("basis") not in {"schedule_to_output", "inter_output"}: + raise ValueError("benchmark_sample has invalid timing basis") + start, end = timing.get("start_monotonic"), timing.get("end_monotonic") + if ( + type(end) not in {int, float} + or not math.isfinite(end) + or ( + start is not None + and ( + type(start) not in {int, float} + or not math.isfinite(start) + or start > end + or not math.isclose(end - start, raw["wall_time"], rel_tol=1e-9, abs_tol=1e-9) + ) + ) + ): + raise ValueError("benchmark_sample contradicts its local timing interval") + if start is None: + reasons.append("sample timing start is unavailable") + if not isinstance(graph, dict) or graph.get("status") not in {"observed", "unavailable"}: + raise ValueError("benchmark_sample has invalid graph observation status") + values = [ + graph.get(name) for name in ("runtime_mode", "num_unpadded_tokens", "num_padded_tokens", "num_paddings") + ] + if graph["status"] == "unavailable": + if any(value is not None for value in values): + raise ValueError("unavailable graph observation contains observed values") + reasons.append("per-sample CUDA graph dispatch is unavailable") + elif ( + values[0] not in {"NONE", "PIECEWISE", "FULL"} + or not all(_integer(value) for value in values[1:]) + or values[2] - values[1] != values[3] + ): + raise ValueError("benchmark_sample has invalid observed graph dispatch or padding") + observed.append( + { + "sample_index": index, + "forward_index": previous_forward, + "timing_basis": timing["basis"], + "cudagraph": graph, + } + ) + count = measurement["preparation"].get("warmup_records_before") + prefix = None + if history["status"] == "unavailable": + if count is not None: + raise ValueError("unavailable warmup history claims an observed record count") + reasons.append("warmup history is unavailable") + elif not _integer(count) or count >= len(prefixes): + raise ValueError("measurement has invalid warmup_records_before") + else: + prefix = prefixes[count] + first_forward = observed[0]["forward_index"] + completed = measurement["preparation"]["completed_points_before"] + if count != sum(record["completed_points_before"] <= completed for record in history["records"]): + raise ValueError("measurement warmup_records_before contradicts preparation order") + for record in history["records"][:count]: + start, end = record["forward_index_start"], record["forward_index_end"] + if (start is not None and start > first_forward) or (end is not None and end > first_forward): + raise ValueError("recorded preparation ends after its measurement") + if record["completed_points_before"] == completed and any( + value is not None and value <= preceding_forward for value in (start, end) + ): + raise ValueError("recorded warmup preparation precedes its completed point count") + if record["status"] == "running": + prefix = None + reasons.append("referenced warmup preparation is still running") + return {"samples": observed, "warmup_prefix_sha256": prefix}, reasons, previous_forward + + +def extract_measurement_evidence( + raw_root: Path, collection: NativeCollection, *, evidence_version: int = 2 +) -> dict[str, Any]: + """Read schema-1 evidence for the publication-selected native samples. + + Call after native validation. Missing legacy evidence is explicitly + unestablished; malformed or contradictory new evidence is rejected. + Raw adjacent steps are retained as evidence of one launch, never counted as + independent observations. Duplicate consolidation uses publication's rule. + """ + if type(evidence_version) is not int or evidence_version not in {1, 2}: + raise ValueError("unsupported measurement evidence version") + artifacts = _rank_artifacts(raw_root) + if not artifacts: + raise ValueError("measurement evidence has no native rank artifacts") + protocol = artifacts[0][1].get("measurement_protocol") + present = "measurement_protocol" in artifacts[0][1] + if present: + _validate_protocol(protocol) + expected_ranks = {rank for rank, *_ in collection.rank_timings} + measurements = {item.point["benchmark_id"]: item for item in collection.points} + selected = select_native_measurements(collection, cell_id="measurement evidence") + selected_ids = {item.point["benchmark_id"] for item in selected} + points = { + _hash(_coordinates(item.point)): { + "point": item.point, + "kv_seed_regime": item.kv_seed_regime, + "regime": {key: item.point.get(key) for key in _GRAPH_FIELDS}, + "wall_time_seconds": max(wall for _, wall in item.rank_wall_times), + "ranks": {}, + "reasons": _regime_gaps(item.point, item.kv_seed_regime), + } + for item in selected + } + result: dict[str, Any] = { + "schema_version": evidence_version, + "backend_version": collection.backend_version, + "collector_attempt_id": collection.collector_attempt_id, + "runtime_run_id": collection.runtime_run_id, + "runtime_grid_digest": collection.runtime_grid_digest, + "measurement_protocol": protocol, + "runtime_identity": {}, + "benchmark_config": {}, + "points": points, + "raw_files": [], + "reasons": [] if present else ["legacy artifact has no measurement protocol"], + "unobserved": sorted(set(_UNOBSERVED + (protocol or {}).get("unobserved", []))), + } + if evidence_version == 2: + result["scheduling"], scheduling_gaps = _scheduling_observations(raw_root, artifacts, collection) + result["reasons"].extend(scheduling_gaps) + result["warmup_evidence"] = {} + result["unobserved"] = sorted( + set( + result["unobserved"] + + [ + "complete_cuda_graph_descriptor", + "capture_versus_replay", + "gpu_execution_time", + "warmup_graph_kernel_and_cache_equivalence", + "decode_real_kv_chain_warmup_history", + ] + ) + ) + seen_ranks = set() + execution_order = None + for path, payload in artifacts: + if ("measurement_protocol" in payload) != present or payload.get("measurement_protocol") != protocol: + raise ValueError("native ranks disagree on measurement_protocol (including mixed legacy/new evidence)") + rank = payload["dp"]["rank"] + if not _integer(rank) or rank not in expected_ranks or rank in seen_ranks: + raise ValueError("measurement evidence has inconsistent DP ranks") + seen_ranks.add(rank) + history, prefixes = _warmup_history(payload, protocol) if evidence_version == 2 else (None, []) + previous_forward = -1 + if history is not None: + result["warmup_evidence"][str(rank)] = history + elif evidence_version == 2: + result["reasons"].append(f"rank {rank}: per-sample execution and warmup evidence is unreported") + if (payload.get("run_id"), payload.get("grid_digest")) != ( + collection.runtime_run_id, + collection.runtime_grid_digest, + ): + raise ValueError("measurement evidence differs from validated native run/grid identity") + runtime = {key: payload.get(key) for key in _RUNTIME_FIELDS} + result["runtime_identity"][str(rank)] = runtime + config = payload.get("config") + result["benchmark_config"][str(rank)] = ( + {key: value for key, value in config.items() if key != "output_path"} if isinstance(config, dict) else None + ) + for key, value in runtime.items(): + if not isinstance(value, dict) or not value: + result["reasons"].append(f"rank {rank}: missing native {key} evidence") + if not isinstance(config, dict) or not config: + result["reasons"].append(f"rank {rank}: missing benchmark config") + rows = payload["results"] + if len(rows) != len(measurements) or {row["point"]["benchmark_id"] for row in rows} != set(measurements): + raise ValueError("measurement evidence differs from validated native point set") + order = [row["point"]["benchmark_id"] for row in rows] + if execution_order is None: + execution_order = order + elif present and order != execution_order: + raise ValueError("native ranks disagree on recorded measurement order") + prefix = hashlib.sha256() + prefix_missing = False + for position, row in enumerate(rows): + benchmark_id = row["point"]["benchmark_id"] + source = measurements[benchmark_id] + if ( + row["point"] != source.point + or row.get("kv_seed_regime") != source.kv_seed_regime + or row["fpms"][0]["wall_time"] != dict(source.rank_wall_times)[rank] + ): + raise ValueError("measurement evidence differs from validated native sample") + measurement, reasons = _read_measurement( + row, rank=rank, grid_digest=collection.runtime_grid_digest, position=position, protocol=protocol + ) + prompt_missing = bool(reasons) + if history is not None: + execution, gaps, previous_forward = _observed_execution( + measurement, row["fpms"][0], history, prefixes, previous_forward + ) + measurement = {**measurement, "observed_execution": execution} + reasons.extend(f"rank {rank}: {reason}" for reason in gaps) + if benchmark_id in selected_ids: + point = points[_hash(_coordinates(row["point"]))] + point["ranks"][str(rank)] = ( + {**measurement, "execution_prefix_sha256": prefix.hexdigest()} if measurement is not None else None + ) + point["reasons"].extend(reasons) + if prefix_missing: + point["reasons"].append(f"rank {rank}: preceding prompt evidence incomplete") + prefix_missing |= prompt_missing + preparation = measurement["preparation"] if measurement else None + if evidence_version == 2 and preparation is not None: + preparation = {key: value for key, value in preparation.items() if key != "warmup_records_before"} + prefix.update( + _hash( + { + "point": row["point"], + "prompts": measurement["prompts"] if measurement else None, + "preparation": preparation, + } + ).encode() + ) + result["raw_files"].append({"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}) + if seen_ranks != expected_ranks: + raise ValueError("measurement evidence is missing native ranks") + for point in points.values(): + point["status"] = "unestablished" if point["reasons"] else "recorded" + if ( + evidence_version == 2 + and protocol is not None + and "execution_evidence" in protocol + and all( + sample["cudagraph"]["status"] == "observed" + for point in points.values() + for measurement in point["ranks"].values() + for sample in measurement["observed_execution"]["samples"] + ) + ): + result["unobserved"].remove("observed_per_call_graph_dispatch") + result["status"] = ( + "unestablished" if result["reasons"] or any(point["reasons"] for point in points.values()) else "recorded" + ) + return result + + +def compare_measurements( + reference: dict[str, Any], candidate: dict[str, Any], *, point_key: str, same_context: bool = True +) -> dict[str, Any]: + """Assess observed identity first, keeping cross-grid history unestablished. + + Comparable means matching recorded conditions for this observation scope; + it does not prove identical cache tensors, per-call dispatch, or stability. + Latency never determines comparability and is deliberately not inspected. + """ + missing = [] + mismatches = [] + extended = reference.get("schema_version") == 2 and candidate.get("schema_version") == 2 + for label, evidence in (("reference", reference), ("candidate", candidate)): + missing.extend(f"{label}: {reason}" for reason in evidence.get("reasons", [])) + if evidence.get("schema_version") not in {1, 2} or evidence.get("measurement_protocol") is None: + missing.append(f"{label}: missing supported measurement evidence") + if not evidence.get("runtime_run_id"): + missing.append(f"{label}: missing independent runtime run identity") + for key in ("backend_version", "measurement_protocol", "runtime_identity"): + left_value, right_value = reference.get(key), candidate.get(key) + if extended and key == "measurement_protocol": + left_value, right_value = ( + {name: value for name, value in item.items() if name not in {"execution_evidence", "unobserved"}} + if isinstance(item, dict) + else item + for item in (left_value, right_value) + ) + if left_value is not None and right_value is not None and left_value != right_value: + mismatches.append(f"{key} differs") + if reference.get("schema_version") == 2 or candidate.get("schema_version") == 2: + for rank in sorted(set(reference.get("scheduling", {})) | set(candidate.get("scheduling", {}))): + values = [ + item.get("scheduling", {}).get(rank, {}).get("async_scheduling") for item in (reference, candidate) + ] + if any(type(value) is not bool for value in values): + missing.append(f"rank {rank}: effective async_scheduling comparison is unestablished") + elif values[0] != values[1]: + mismatches.append(f"rank {rank}: effective async_scheduling differs") + if reference.get("runtime_run_id") and reference.get("runtime_run_id") == candidate.get("runtime_run_id"): + mismatches.append("runtime_run_id is not an independent launch") + left = reference.get("points", {}).get(point_key) + right = candidate.get("points", {}).get(point_key) + if left is None or right is None: + missing.append("point is missing from reference or candidate") + else: + for label, point in (("reference", left), ("candidate", right)): + missing.extend(f"{label}: {reason}" for reason in point.get("reasons", [])) + if _content_key(left["point"]) != _content_key(right["point"]): + mismatches.append("point shape/partition differs") + for key in ("kv_seed_regime", "regime"): + if left.get(key) != right.get(key): + mismatches.append(f"{key} differs") + if set(left["ranks"]) != set(right["ranks"]): + mismatches.append("DP rank set differs") + for rank in sorted(set(left["ranks"]) & set(right["ranks"])): + lhs, rhs = left["ranks"][rank], right["ranks"][rank] + if lhs is None or rhs is None: + missing.append(f"rank {rank}: missing measurement evidence") + continue + for key in ("point_key", "prompts", "expected_internal_samples", "estimate"): + if lhs[key] != rhs[key]: + mismatches.append(f"rank {rank}: {key} differs") + observed_left, observed_right = lhs.get("observed_execution"), rhs.get("observed_execution") + if observed_left is not None and observed_right is not None: + # Relative forward order belongs to matched-sweep preparation; + # absolute process-local timestamps never establish equality. + fields = ( + ("sample_index", "timing_basis", "forward_index") + if same_context + else ("sample_index", "timing_basis") + ) + left_samples, right_samples = observed_left["samples"], observed_right["samples"] + if len(left_samples) != len(right_samples): + mismatches.append(f"rank {rank}: observed sample count differs") + for a, b in zip(left_samples, right_samples, strict=False): + if any(a[name] != b[name] for name in fields): + mismatches.append(f"rank {rank}: observed sample order or timing basis differs") + if a["cudagraph"]["status"] == b["cudagraph"]["status"] == "observed": + if a["cudagraph"] != b["cudagraph"]: + mismatches.append(f"rank {rank}: observed graph dispatch differs") + else: + missing.append(f"rank {rank}: observed graph dispatch comparison is unestablished") + if same_context: + if observed_left["warmup_prefix_sha256"] is None or observed_right["warmup_prefix_sha256"] is None: + missing.append(f"rank {rank}: observed warmup comparison is unestablished") + elif observed_left["warmup_prefix_sha256"] != observed_right["warmup_prefix_sha256"]: + mismatches.append(f"rank {rank}: observed warmup preparation differs") + elif extended or (observed_left is None) != (observed_right is None): + missing.append(f"rank {rank}: per-sample execution comparison is unestablished") + if same_context: + for key in ("preparation", "execution_prefix_sha256"): + left_value, right_value = lhs[key], rhs[key] + if extended and key == "preparation": + left_value, right_value = ( + {name: value for name, value in item.items() if name != "warmup_records_before"} + for item in (left_value, right_value) + ) + if left_value != right_value: + mismatches.append(f"rank {rank}: {key} differs") + if same_context: + for key in ("runtime_grid_digest", "benchmark_config"): + if reference.get(key) is not None and candidate.get(key) is not None and reference[key] != candidate[key]: + mismatches.append(f"{key} differs") + identity_status = "mismatch" if mismatches else "unestablished" if missing else "comparable" + reasons = sorted({*mismatches, *missing}) + if not same_context: + reasons.append("cross-context execution history equivalence has not been established") + return { + "status": "unestablished" if not same_context and identity_status == "comparable" else identity_status, + "identity_status": identity_status, + "reasons": reasons, + "unobserved": sorted( + set( + ( + _UNOBSERVED + if reference.get("schema_version") != 2 or candidate.get("schema_version") != 2 + else _UNOBSERVED[:-1] + ) + + reference.get("unobserved", []) + + candidate.get("unobserved", []) + ) + ), + "scope": "recorded measurement protocol and injected prompts; worker runtime validation remains required", + "context": "matched_native_sweep" if same_context else "cross_context_unproven", + } diff --git a/python/aisimulate/collector/fpm_forward/memory_admission.py b/python/aisimulate/collector/fpm_forward/memory_admission.py index df741d2c1..c80b79178 100644 --- a/python/aisimulate/collector/fpm_forward/memory_admission.py +++ b/python/aisimulate/collector/fpm_forward/memory_admission.py @@ -9,7 +9,7 @@ import math from dataclasses import dataclass -from aisimulate.sdk.memory import KVCacheEstimator +from aisimulate.sdk.memory import KVCacheEstimator, kv_cache_budget_bytes from aisimulate_core.sdk.errors import PerfDataNotAvailableError from aisimulate_core.sdk.fpm_profile import FpmModelProfile from aisimulate_core.sdk.perf_database import load_system_spec @@ -28,15 +28,17 @@ class DTypeMemoryEstimate: gpu_capacity_bytes: int | None reason: str provenance: str | None = None + gpu_memory_budget_bytes: int | None = None def to_dict(self) -> dict[str, object]: + budget = self.gpu_capacity_bytes if self.gpu_memory_budget_bytes is None else self.gpu_memory_budget_bytes payload = { "kv_cache_dtype": self.kv_cache_dtype, "disposition": self.disposition, "estimated_non_kv_bytes": self.estimated_non_kv_bytes, "gpu_capacity_bytes": self.gpu_capacity_bytes, "headroom_bytes": ( - self.gpu_capacity_bytes - self.estimated_non_kv_bytes + budget - self.estimated_non_kv_bytes if self.estimated_non_kv_bytes is not None and self.gpu_capacity_bytes is not None else None ), @@ -44,6 +46,8 @@ def to_dict(self) -> dict[str, object]: } if self.provenance is not None: payload["provenance"] = self.provenance + if self.gpu_memory_budget_bytes is not None: + payload["gpu_memory_budget_bytes"] = self.gpu_memory_budget_bytes return payload @@ -93,6 +97,8 @@ def _estimate_dtype( max_new_tokens: int, fpm_profile: FpmModelProfile | None = None, max_batch_size: int | None = None, + gpu_memory_utilization: float | None = None, + worker_type: str | None = None, ) -> DTypeMemoryEstimate: # Planner-owned capability data fails closed: resolve_model_capability # guarantees an fmha mapping for every resolved KV dtype, so a missing @@ -116,6 +122,7 @@ def _estimate_dtype( moe_tp_size=topology.moe_tp, moe_ep_size=topology.moe_ep, cp_size=topology.cp, + worker_type=worker_type or "aggregated", ) resources = deployment.resources resources.validate_envelope( @@ -125,19 +132,40 @@ def _estimate_dtype( capacity = math.floor(float(load_system_spec(system)["gpu"]["mem_capacity"])) if capacity <= 0: raise ValueError(f"invalid GPU memory capacity for {system!r}: {capacity}") - rejected = resources.non_kv_bytes >= capacity + budget = math.floor( + kv_cache_budget_bytes( + capacity=capacity, + non_kv=0, + fraction=gpu_memory_utilization if gpu_memory_utilization is not None else 1.0, + of_free=False, + ) + ) + if resources.memory_source != "declared": + return DTypeMemoryEstimate( + kv_cache_dtype=kv_cache_dtype, + disposition="unknown", + estimated_non_kv_bytes=None, + gpu_capacity_bytes=capacity, + reason="memory capacity is determined by this collection worker's runtime initialization; " + "no non-KV byte declaration is required", + provenance=resources.provenance, + gpu_memory_budget_bytes=budget if gpu_memory_utilization is not None else None, + ) + rejected = resources.non_kv_bytes >= budget + capacity_label = "configured GPU memory budget" if gpu_memory_utilization is not None else "GPU capacity" return DTypeMemoryEstimate( kv_cache_dtype=kv_cache_dtype, disposition="rejected" if rejected else "admitted", estimated_non_kv_bytes=resources.non_kv_bytes, gpu_capacity_bytes=capacity, reason=( - "declared rank-local non-KV resource bound is not below GPU capacity" + f"declared rank-local non-KV resource bound is not below {capacity_label}" if rejected - else "declared rank-local non-KV resource bound is below GPU capacity; " + else f"declared rank-local non-KV resource bound is below {capacity_label}; " "runtime profiling still determines KV capacity and CUDA-graph memory" ), provenance=resources.provenance, + gpu_memory_budget_bytes=budget if gpu_memory_utilization is not None else None, ) try: breakdown = KVCacheEstimator.from_request( @@ -201,17 +229,27 @@ def _estimate_dtype( reason=f"AIC memory estimate unavailable: {type(error).__name__}: {error}", ) - rejected = non_kv >= capacity + budget = math.floor( + kv_cache_budget_bytes( + capacity=capacity, + non_kv=0, + fraction=gpu_memory_utilization if gpu_memory_utilization is not None else 1.0, + of_free=False, + ) + ) + rejected = non_kv >= budget + capacity_label = "configured GPU memory budget" if gpu_memory_utilization is not None else "GPU capacity" return DTypeMemoryEstimate( kv_cache_dtype=kv_cache_dtype, disposition="rejected" if rejected else "admitted", estimated_non_kv_bytes=non_kv, gpu_capacity_bytes=capacity, reason=( - "AIC configured max-new-token non-KV memory is not below GPU capacity" + f"AIC configured max-new-token non-KV memory is not below {capacity_label}" if rejected - else "AIC configured max-new-token non-KV memory is below GPU capacity" + else f"AIC configured max-new-token non-KV memory is below {capacity_label}" ), + gpu_memory_budget_bytes=budget if gpu_memory_utilization is not None else None, ) @@ -225,12 +263,15 @@ def filter_memory_infeasible_topologies( max_new_tokens: int, fpm_profile: FpmModelProfile | None = None, max_batch_size: int | None = None, + gpu_memory_utilization: float | None = None, + worker_type: str | None = None, ) -> tuple[tuple[ParallelTopology, ...], tuple[TopologyMemoryDecision, ...]]: """Drop topologies that cannot fit the configured max-new-token envelope. This is intentionally a one-sided generation-time filter. A topology is rejected only when every requested KV dtype has a successful estimate - and all estimates exceed rank-local physical capacity. Without a profile, + and all estimates exceed the declared share of rank-local capacity (physical + capacity when no GPU memory fraction is declared). Without a profile, unknown AIC estimates remain runnable. Supplied profiles must explicitly cover the requested deployment and envelope; validation errors propagate. """ @@ -254,6 +295,7 @@ def filter_memory_infeasible_topologies( moe_tp_size=topology.moe_tp, moe_ep_size=topology.moe_ep, cp_size=topology.cp, + worker_type=worker_type or "aggregated", ).resources if fpm_profile is not None else None @@ -269,6 +311,8 @@ def filter_memory_infeasible_topologies( max_new_tokens=max_new_tokens, fpm_profile=fpm_profile, max_batch_size=max_batch_size, + gpu_memory_utilization=gpu_memory_utilization, + worker_type=worker_type, ) for kv_cache_dtype in capability.dtype.kv_cache_dtypes ) @@ -283,7 +327,7 @@ def filter_memory_infeasible_topologies( admitted.append(topology) else: disposition = "rejected" - reason = "all requested KV dtypes exceed rank-local GPU capacity at configured max new tokens" + reason = "all requested KV dtypes exceed rank-local GPU memory budget at configured max new tokens" rejected_capacity.append((topology, estimates)) decisions.append( TopologyMemoryDecision( @@ -292,7 +336,11 @@ def filter_memory_infeasible_topologies( max_new_tokens=max_new_tokens, estimates=estimates, reason=reason, - source="fpm_profile_declared" if fpm_profile is not None else "aic_native_configured_max_new_tokens", + source=( + f"fpm_profile_{resources.memory_source}" + if resources is not None + else "aic_native_configured_max_new_tokens" + ), max_batch_size=resources.max_batch_size if resources is not None else 1, profile_max_num_tokens=resources.max_num_tokens if resources is not None else None, ) @@ -310,14 +358,15 @@ def filter_memory_infeasible_topologies( if best.estimated_non_kv_bytes is None: details.append(f"{topology.to_dict()}={best.reason}") else: + budget = ( + best.gpu_capacity_bytes if best.gpu_memory_budget_bytes is None else best.gpu_memory_budget_bytes + ) details.append( - f"{topology.to_dict()}=" - f"{best.estimated_non_kv_bytes / 2**30:.2f}/" - f"{best.gpu_capacity_bytes / 2**30:.2f} GiB" + f"{topology.to_dict()}={best.estimated_non_kv_bytes / 2**30:.2f}/{budget / 2**30:.2f} GiB" ) logger.warning( "fpm_forward: dropped %d/%d topologies (%s configured max-new-token non-KV memory " - "exceeds GPU capacity, system=%s, max_new_tokens=%d): %s", + "exceeds GPU memory budget, system=%s, max_new_tokens=%d): %s", len(rejected_capacity), len(topologies), "FPM profile" if fpm_profile is not None else "AIC", diff --git a/python/aisimulate/collector/fpm_forward/model_capability.py b/python/aisimulate/collector/fpm_forward/model_capability.py index 68839456f..c0583faa5 100644 --- a/python/aisimulate/collector/fpm_forward/model_capability.py +++ b/python/aisimulate/collector/fpm_forward/model_capability.py @@ -11,6 +11,7 @@ from pathlib import Path from typing import Any +from aisimulate.sdk.common import MULTIMODAL_TEXT_CONFIG_KEY from aisimulate.sdk.memory import NaiveKVCacheEstimator from aisimulate.sdk.utils import ( HuggingFaceDownloadError, @@ -20,6 +21,7 @@ _get_model_config_path, _parse_hf_config_json, ) +from aisimulate_core.sdk.fpm_model_metadata import attach_quantization_sidecar, inherit_decoder_precision _ATTENTION_SOURCE_OPS = { "dsa_module": frozenset({"dsa_context_module", "dsa_generation_module"}), @@ -64,13 +66,26 @@ def payload(self) -> dict[str, Any]: @property def effective_payload(self) -> dict[str, Any]: - """Return the text-model view while preserving top-level quant metadata.""" + """Return decoder metadata, inheriting only identity and shared precision.""" payload = self.payload text_config = payload.get("text_config") if isinstance(text_config, dict): - payload = {**payload, **text_config} - return payload + effective, _ = inherit_decoder_precision(payload, text_config) + if effective.get("architectures") is None and payload.get("architectures") is not None: + effective["architectures"] = payload["architectures"] + else: + effective = payload + sidecar = effective.get("hf_quant_config") + if sidecar is None and isinstance(text_config, dict) and isinstance(payload.get("hf_quant_config"), dict): + try: + return attach_quantization_sidecar(effective, payload["hf_quant_config"]) + except ValueError: + # Inline wrapper precision can describe non-decoder components. + # Preserve decoder scope on a conflicting wrapper declaration; + # actual adjacent sidecars are checked before freezing the source. + return effective + return attach_quantization_sidecar(effective, sidecar) if isinstance(sidecar, dict) else effective def parsed_payload(self) -> dict[str, Any]: """Return the frozen payload in the SDK's parsed model-parameter shape. @@ -83,8 +98,16 @@ def parsed_payload(self) -> dict[str, Any]: """ raw_config = _attach_inferred_quant_fields(self.effective_payload) + parser_config = raw_config + architecture = (raw_config.get("architectures") or [None])[0] + text_key = MULTIMODAL_TEXT_CONFIG_KEY.get(architecture) + source = self.payload + if text_key and isinstance(source.get(text_key), dict): + # Registered wrapper parsers require their original encoder and + # processor context, but read model geometry from the decoder. + parser_config = {**source, "architectures": [architecture], text_key: raw_config} try: - parsed = _parse_hf_config_json(raw_config) + parsed = _parse_hf_config_json(parser_config) except ValueError: # Bootstrap templates intentionally admit real HF architectures # that the native AIC parser does not know yet. Normalize their @@ -131,6 +154,20 @@ def _load_json_mapping(path: Path, *, description: str) -> tuple[dict[str, Any], return payload, path +def _attach_modelopt_sidecar(payload: dict[str, Any], sidecar: dict[str, Any]) -> dict[str, Any]: + text_config = payload.get("text_config") + if payload.get("hf_quant_config") is not None and payload["hf_quant_config"] != sidecar: + raise ValueError("conflicting inline hf_quant_config and adjacent hf_quant_config.json") + if isinstance(text_config, dict): + decoder, _ = inherit_decoder_precision(payload, text_config) + attach_quantization_sidecar(decoder, sidecar) # Reject contradictory source declarations before planning. + else: + attach_quantization_sidecar(payload, sidecar) + # Preserve the established frozen source/hash for already supported inputs. + # Null handling and inferred scalar fields belong to the effective view. + return _attach_hf_quant_config(payload, sidecar) + + def _attach_sibling_quant_config(payload: dict[str, Any], config_path: Path, *, cached: bool) -> dict[str, Any]: suffix = "_hf_quant_config.json" if cached else "hf_quant_config.json" quant_path = ( @@ -141,7 +178,7 @@ def _attach_sibling_quant_config(payload: dict[str, Any], config_path: Path, *, if not quant_path.is_file(): return payload quant_config, _ = _load_json_mapping(quant_path, description="model quantization config") - return _attach_hf_quant_config(payload, quant_config) + return _attach_modelopt_sidecar(payload, quant_config) def _local_config_path(raw_path: str, *, explicit: bool) -> Path | None: @@ -191,7 +228,7 @@ def _download_model_config(model_path: str) -> ResolvedModelConfig: if quant_config is not None: if not isinstance(quant_config, dict): raise TypeError(f"Hugging Face hf_quant_config.json must be a mapping for model {model_path!r}") - payload = _attach_hf_quant_config(payload, quant_config) + payload = _attach_modelopt_sidecar(payload, quant_config) return ResolvedModelConfig( payload, source_kind="huggingface", diff --git a/python/aisimulate/collector/fpm_forward/native_artifact.py b/python/aisimulate/collector/fpm_forward/native_artifact.py index 1794d2344..402ba4037 100644 --- a/python/aisimulate/collector/fpm_forward/native_artifact.py +++ b/python/aisimulate/collector/fpm_forward/native_artifact.py @@ -7,6 +7,7 @@ import hashlib import json +import logging import math import re from dataclasses import dataclass @@ -22,6 +23,8 @@ from .planner import FPMCell from .types import KVWARM_STRATEGIES +logger = logging.getLogger(__name__) + COLLECTOR_PROVENANCE_FILENAME = "collector-provenance.json" @@ -29,6 +32,9 @@ class NativePointMeasurement: point: dict[str, Any] rank_wall_times: tuple[tuple[int, float], ...] + # Keep native row provenance for diagnostics, without reinterpreting the + # formal database's historical prefill kv_seed_regime contract. + kv_seed_regime: str | None = None # Distinguishes "no artifact seen yet" from a legitimately absent (legacy) @@ -36,6 +42,36 @@ class NativePointMeasurement: _KVWARM_UNSEEN = object() +def _validate_seed_regime(row: dict[str, Any], path: Path) -> str | None: + regime = row.get("kv_seed_regime") + if regime is None: + return None # Historical artifacts did not report a row-level regime. + if not isinstance(regime, str) or not regime: + raise ValueError(f"native kv_seed_regime must be a non-empty string: {path}") + # Dynamo's _kvwarm_seed_regime reports these point-injection stamps: + # https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/instrumented_scheduler.py#L4721 + stamps = { + "real_prefix": "prefill_real_seed", + "fake_prefix": "prefill_fake_prefix", + "real_kv": "kvwarm_real_kv", + "fake_fallback": "kvwarm_fake_fallback", + } + reasons = set(row["point"].get("sample_reasons") or []) + observed_stamps = reasons.intersection(stamps.values()) + if regime in stamps and observed_stamps != {stamps[regime]}: + raise ValueError(f"native kv_seed_regime disagrees with point injection stamps: {path}") + if regime in {"real_prefix", "fake_prefix"} and row["point"]["point_type"] != "prefill": + raise ValueError(f"native prefix seed regime requires prefill: {path}") + # Real-seed staging stamps the path before checking whether any KV needs + # seeding; an uncached prefill can therefore retain real_prefix provenance. + # https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/instrumented_scheduler.py#L4298 + if regime == "fake_prefix" and row["point"]["total_kv_read_tokens"] <= 0: + raise ValueError(f"native fake prefix seed regime requires cached prefill: {path}") + if regime not in stamps and observed_stamps: + raise ValueError(f"native kv_seed_regime disagrees with point injection stamps: {path}") + return regime + + def _validate_kvwarm_contract(cell: FPMCell, kvwarm: object, path: Path) -> dict[str, Any] | None: """Validate that engine protocol evidence matches the rendered strategy.""" @@ -86,6 +122,94 @@ class NativeCollection: input_provenance: dict[str, Any] | None = None +def _zero_kv_prefill_sample(measurements: list[NativePointMeasurement]) -> NativePointMeasurement | None: + """Prefer the unique ordinary sample over equivalent real-prefix path samples.""" + ordinary = [m for m in measurements if m.kv_seed_regime == "not_applicable"] + if len(ordinary) != 1: + return None + sample = ordinary[0] + if sample.point["point_type"] != "prefill" or sample.point["total_kv_read_tokens"] != 0: + return None + expected = {key: value for key, value in sample.point.items() if key != "benchmark_id"} + expected["sample_reasons"] = list(expected.get("sample_reasons") or []) + for measurement in measurements: + if measurement is sample: + continue + if measurement.kv_seed_regime != "real_prefix": + return None + point = {key: value for key, value in measurement.point.items() if key != "benchmark_id"} + point["sample_reasons"] = [ + reason for reason in (point.get("sample_reasons") or []) if reason != "prefill_real_seed" + ] + if point != expected: + return None + return sample + + +def select_native_measurements(collection: NativeCollection, *, cell_id: str) -> list[NativePointMeasurement]: + """Select the formal per-coordinate samples from one validated attempt. + + Used by publication and by repeatability source selection, so the original + reference timing is identical. Never combine independent repeat attempts. + """ + # The steady-state decode policy clamps every per-sequence context below + # the measurable minimum up to it, so several requested plan points can + # collapse onto one achieved physical coordinate (e.g. batch=3 with + # requested total-kv 3/4/5/6 all measure total-kv 6). Repeated samples of + # one coordinate would violate the database's unique-key contract, so keep + # exactly one per coordinate: the native (unclamped) sample when present, + # otherwise the clamped sample with the lowest benchmark_id. + # Real-prefix staging can also emit zero-KV prefill samples beside the + # ordinary sample of the same work. Only this provenance difference is + # consolidatable: keep the unique ordinary sample with identical remaining + # point metadata. Other native collisions remain a hard error. Raw samples + # stay intact; neither reduction averages timings nor selects by latency. + grouped: dict[tuple[str, int, int, int], list[NativePointMeasurement]] = {} + for measurement in collection.points: + point = measurement.point + key = ( + str(point["point_type"]), + int(point["batch_size"]), + int(point["total_prefill_tokens"]), + int(point["total_kv_read_tokens"]), + ) + grouped.setdefault(key, []).append(measurement) + selected: list[NativePointMeasurement] = [] + dropped_clamped = 0 + dropped_zero_kv = 0 + for key, measurements in grouped.items(): + natives = [m for m in measurements if "context_clamped" not in (m.point.get("sample_reasons") or ())] + if len(natives) > 1: + sample = _zero_kv_prefill_sample(measurements) + if sample is not None: + selected.append(sample) + dropped_zero_kv += len(measurements) - 1 + continue + raise ValueError( + f"conflicting FPM measurements for physical coordinate {key} in " + f"{cell_id}: {len(natives)} unclamped samples share one key" + ) + if natives: + selected.append(natives[0]) + else: + selected.append(min(measurements, key=lambda m: int(m.point["benchmark_id"]))) + dropped_clamped += len(measurements) - 1 + if dropped_clamped: + logger.info( + "FPM %s: consolidated %d context-clamped duplicate sample(s) onto their achieved physical coordinates", + cell_id, + dropped_clamped, + ) + if dropped_zero_kv: + logger.info( + "FPM %s: consolidated %d zero-KV prefill provenance duplicate sample(s) onto their ordinary samples", + cell_id, + dropped_zero_kv, + ) + + return selected + + def _validate_execution_provenance(cell: FPMCell, payload: dict[str, Any], path: Path) -> dict[str, Any] | None: """Config-bound curves require engine evidence for execution and real input.""" if not cell.execution_identity[0]: @@ -381,6 +505,7 @@ def validate_native_collection( kvwarm_seen: object = _KVWARM_UNSEEN input_provenance: dict[str, Any] | None = None local_fpms: dict[tuple[int, int], dict[str, Any]] = {} + seed_regimes: dict[int, str | None] = {} rank_timings: list[tuple[int, float, float]] = [] for path, payload in rank_payloads: @@ -502,6 +627,17 @@ def validate_native_collection( if not isinstance(fpms, list) or len(fpms) != 1: raise ValueError(f"native point must contain exactly one local FPM: {path}") _validate_fpm(point, fpms[0], rank=rank) + seed_regime = _validate_seed_regime(row, path) + if ( + seed_regime == "real_kv" + and cell.workload_kind == "decode" + and kvwarm is not None + and (not kvwarm["enabled"] or not kvwarm["warm_eligible"]) + ): + raise ValueError(f"native real KV seed regime disagrees with warm-up eligibility: {path}") + if benchmark_id in seed_regimes and seed_regimes[benchmark_id] != seed_regime: + raise ValueError(f"native DP ranks disagree on per-point KV seed regime: {path}") + seed_regimes[benchmark_id] = seed_regime local_fpms[(benchmark_id, rank)] = fpms[0] points.append(point) # The engine writes results in execution order, and KV warm-up's @@ -568,7 +704,11 @@ def validate_native_collection( ): raise ValueError(f"native iteration wall_time mismatch for benchmark_id={benchmark_id}") measured_iteration_seconds += group_wall_time - measurements.append(NativePointMeasurement(point=dict(point), rank_wall_times=tuple(wall_times))) + measurements.append( + NativePointMeasurement( + point=dict(point), rank_wall_times=tuple(wall_times), kv_seed_regime=seed_regimes[benchmark_id] + ) + ) for rank, _elapsed, measured in rank_timings: if not math.isclose(measured, measured_iteration_seconds, rel_tol=1e-9, abs_tol=1e-12): diff --git a/python/aisimulate/collector/fpm_forward/planner.py b/python/aisimulate/collector/fpm_forward/planner.py index 3e3deb194..e083d34de 100644 --- a/python/aisimulate/collector/fpm_forward/planner.py +++ b/python/aisimulate/collector/fpm_forward/planner.py @@ -15,7 +15,7 @@ import logging import os import subprocess -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path, PurePosixPath from typing import Any @@ -23,13 +23,47 @@ from aisimulate_core.sdk.fpm_profile import FpmDeploymentProfile, FpmModelProfile, load_fpm_profile from .capabilities import ModelCapabilityProfile, ResolvedDTypeProfile, resolve_model_capability -from .config import PARALLEL_AXES, FPMCollectionOptions +from .config import ( + FPM_MAX_PREFILL_ISL, + PARALLEL_AXES, + VLLM_AUTO_FIT_MAX_MODEL_LEN, + FPMCollectionOptions, + with_kv_warmup_defaults, +) from .memory_admission import TopologyMemoryDecision, filter_memory_infeasible_topologies +from .runtime.fpm_memory_observer import SUPPORTED_VERSION as MEMORY_OBSERVER_VERSION from .topology import enumerate_fpm_topologies, topology_strategy from .types import ParallelTopology logger = logging.getLogger(__name__) + +def _runtime_memory_policy( + profile: FpmModelProfile | None, backend_version: str, runtime_observation: dict[str, Any] | None = None +) -> dict[str, object] | None: + if runtime_observation is not None: + return { + "source": "runtime_instrumentation", + "observation": "enabled", + "selected_vllm_version": backend_version, + "bundle_sha256": runtime_observation["bundle_sha256"], + "async_scheduling": False, + } + if profile is None or not any( + deployment.resources.memory_source == "pending" for deployment in profile.deployments + ): + return None + available = backend_version == MEMORY_OBSERVER_VERSION + return { + "source": "vllm_initialization", + "supported_vllm_version": MEMORY_OBSERVER_VERSION, + "selected_vllm_version": backend_version, + "observation": "enabled" if available else "unavailable_for_runtime", + "evidence_schema_version": 1, + "async_scheduling": False if available else None, + } + + _INSTALLED_DISTRIBUTION = "aisimulate" _INSTALLED_PAYLOAD_ROOTS = frozenset(("aisimulate_core", "aisimulate", "collector")) _INSTALLED_PLANNER_PATH = PurePosixPath("collector/fpm_forward/planner.py") @@ -39,6 +73,9 @@ PurePosixPath("collector/fpm_forward/runner.py"), PurePosixPath("collector/fpm_forward/runtime/fpm_exec.sh"), PurePosixPath("collector/fpm_forward/runtime/preflight.py"), + PurePosixPath("collector/fpm_forward/runtime/fpm_memory_observer.py"), + PurePosixPath("collector/fpm_forward/runtime/fpm_memory_worker.py"), + PurePosixPath("collector/fpm_forward/runtime/fpm_memory_scheduler.py"), ) ) @@ -491,6 +528,27 @@ class FPMCollectionPlan: cells: tuple[FPMCell, ...] sha256: str _fpm_profile_json: str | None = field(default=None, repr=False) + _runtime_observation_json: str | None = field(default=None, repr=False) + + @property + def runtime_instrumentation(self): + if self._runtime_observation_json is None: + return None + from .runtime_instrumentation import load_instrumentation + + reference = json.loads(self._runtime_observation_json) + bundle = load_instrumentation(reference["manifest"], expected_version=self.capability.aic_database_version) + if bundle.sha256 != reference["bundle_sha256"]: + raise ValueError("formal collection instrumentation changed after planning") + return bundle + + @property + def runtime_launch(self) -> dict[str, Any] | None: + return json.loads(self._runtime_observation_json)["launch"] if self._runtime_observation_json else None + + @property + def runtime_configuration(self) -> str | None: + return json.loads(self._runtime_observation_json)["configuration"] if self._runtime_observation_json else None @property def fpm_profile(self) -> FpmModelProfile | None: @@ -512,9 +570,11 @@ def deployment_profile(self, cell: FPMCell) -> FpmDeploymentProfile | None: moe_tp_size=cell.topology.moe_tp, moe_ep_size=cell.topology.moe_ep, cp_size=cell.topology.cp, + worker_type=self.options.worker_type or "aggregated", ) def to_dict(self) -> dict[str, object]: + prefill_sampling = self.options.prefill_sampling.to_dict() explicit_points = ( json.loads(self.options.benchmark_points_json) if self.options.benchmark_points_json is not None else None ) @@ -542,7 +602,7 @@ def to_dict(self) -> dict[str, object]: "partition_policy": "balanced_v1", "point_admission": "dynamo_live_scheduler", "precondition": "vllm_engine_initialized", - "prefill_sampling": self.options.prefill_sampling.to_dict(), + "prefill_sampling": prefill_sampling, "planned_point_count": ( sum(len(explicit_points.get(phase, [])) for phase in ("prefill", "decode")) if explicit_points is not None @@ -570,14 +630,22 @@ def to_dict(self) -> dict[str, object]: ), "backend_policies": len(self.backend_policies), "cells": len(self.cells), - "prefill_cudagraph_capture_sizes": len(self.options.prefill_sampling.cudagraph_capture_sizes), - "prefill_new_token_axis_points": len(self.options.prefill_sampling.new_token_axis_points), + "prefill_cudagraph_capture_sizes": prefill_sampling["cudagraph_capture_size_count"], + "prefill_new_token_axis_points": prefill_sampling["new_token_axis_point_count"], "points": "runtime-determined", }, "sha256": self.sha256, } if self._fpm_profile_json is not None: payload["fpm_profile"] = json.loads(self._fpm_profile_json) + runtime_observation = json.loads(self._runtime_observation_json) if self._runtime_observation_json else None + memory_policy = _runtime_memory_policy( + self.fpm_profile, self.capability.aic_database_version, runtime_observation + ) + if memory_policy is not None: + payload["runtime_memory_policy"] = memory_policy + if runtime_observation is not None: + payload["runtime_observation"] = runtime_observation return payload @@ -623,6 +691,9 @@ def build_collection_plan( fpm_profile: FpmModelProfile | dict[str, Any] | None = None, collector_config: dict[str, Any] | None = None, generator_overrides: dict[str, Any] | None = None, + runtime_instrumentation=None, + runtime_launch: dict[str, Any] | None = None, + runtime_configuration: str = "collection", ) -> FPMCollectionPlan: if backend != "vllm": raise ValueError("FPM Generator V1 currently supports only backend=vllm") @@ -645,7 +716,9 @@ def build_collection_plan( f"{system}/{backend}; found {sorted(versions)}" ) collector_config["aic_database_version"] = next(iter(versions)) - generator_config_sha256 = _canonical_hash(generator_overrides or {}) + # Freeze resolved warm-up defaults as well as explicit deployment inputs; + # changing the default must not resume a campaign under its old plan hash. + generator_config_sha256 = _canonical_hash(with_kv_warmup_defaults(generator_overrides or {})) capability = resolve_model_capability( backend=backend, model_path=model_path, @@ -687,7 +760,7 @@ def build_collection_plan( ) policies = _backend_policies(options, collector_config, backend=backend) if profile is not None: - _validate_profile_identities( + deployments = _validate_profile_identities( profile, capability, candidate_topologies, @@ -695,8 +768,50 @@ def build_collection_plan( model_path, system, backend, - max_decode_batch_size=options.max_decode_batch_size, + worker_type=options.worker_type, + ) + if options.vllm_max_model_len > profile.context_length: + raise ValueError( + f"--fpm-max-model-len={options.vllm_max_model_len} exceeds " + f"FPM profile context_length={profile.context_length}" + ) + max_prefill_isl = options.max_prefill_isl + if max_prefill_isl is None: + max_prefill_isl = options.max_num_batched_tokens + if max_prefill_isl is None: + token_limits = {deployment.resources.max_num_tokens for deployment in deployments} + if options.worker_type == "aggregated": + if len(token_limits) != 1: + raise ValueError( + "aggregated deployments have different token limits; " + "set an explicit shared --fpm-max-num-batched-tokens" + ) + max_prefill_isl = next(iter(token_limits)) + else: + max_prefill_isl = min(FPM_MAX_PREFILL_ISL, *token_limits) + options = replace( + options, + vllm_max_model_len=( + profile.context_length + if options.vllm_max_model_len == VLLM_AUTO_FIT_MAX_MODEL_LEN + else options.vllm_max_model_len + ), + max_prefill_isl=max_prefill_isl, ) + # Validate the shared decode envelope as well as narrower prefill + # controls before memory admission can queue or drop any deployment. + for deployment in deployments: + resources = deployment.resources + options.validate_scheduler_limits( + profile_max_num_tokens=resources.max_num_tokens, profile_max_batch_size=resources.max_batch_size + ) + resources.validate_envelope( + max_num_tokens=options.max_num_batched_tokens or max_prefill_isl, + max_batch_size=max( + options.max_decode_batch_size or options.max_num_seqs or resources.max_batch_size, + options.max_prefill_batch_size or options.max_num_seqs or resources.max_batch_size, + ), + ) topologies, topology_memory_admission = filter_memory_infeasible_topologies( backend=backend, model_path=model_path, @@ -705,7 +820,9 @@ def build_collection_plan( topologies=candidate_topologies, max_new_tokens=options.prefill_sampling.max_total_prefill_tokens, fpm_profile=profile, - max_batch_size=options.max_prefill_batch_size, + max_batch_size=options.prefill_sampling.max_batch_size, + gpu_memory_utilization=options.gpu_memory_utilization, + worker_type=options.worker_type, ) weight_quantization = capability.dtype.gemm_quant_mode runnable_dtype_pairs = { @@ -762,7 +879,7 @@ def build_collection_plan( comm_quant_mode=capability.dtype.comm_quant_mode, fmha_resolution=capability.dtype.fmha_resolution_by_kv_dtype[kv_cache_dtype], ) - for phase in ("prefill", "decode") + for phase in options.workload_kinds for topology in topologies for kv_cache_dtype in capability.dtype.kv_cache_dtypes if (topology, kv_cache_dtype) in runnable_dtype_pairs @@ -788,6 +905,45 @@ def build_collection_plan( if profile is not None: canonical["fpm_profile"] = profile.model_dump(mode="json") profile_json = json.dumps(canonical["fpm_profile"], sort_keys=True, separators=(",", ":")) + runtime_observation_json = None + if runtime_instrumentation is not None: + from .runtime_instrumentation import load_instrumentation + from .runtime_probe import normalize_probe_launch, validate_collection_probe_launch + + if runtime_launch is None: + raise ValueError("formal runtime instrumentation requires the accepted probe launch facts") + bundle = ( + load_instrumentation(runtime_instrumentation, expected_version=capability.aic_database_version) + if isinstance(runtime_instrumentation, (str, Path)) + else runtime_instrumentation + ) + launch = normalize_probe_launch(runtime_launch) + validate_collection_probe_launch( + launch, + bundle, + model_path=model_path, + system=system, + backend=backend, + backend_version=capability.aic_database_version, + cells=cells, + options=options, + profile=profile, + generator_overrides=generator_overrides or {}, + ) + canonical["runtime_observation"] = { + "manifest": str(bundle.manifest_path), + "bundle_sha256": bundle.sha256, + "launch": launch, + "configuration": runtime_configuration, + } + runtime_observation_json = json.dumps(canonical["runtime_observation"], sort_keys=True) + elif runtime_launch is not None: + raise ValueError("accepted probe launch facts require runtime instrumentation for formal collection") + memory_policy = _runtime_memory_policy( + profile, capability.aic_database_version, canonical.get("runtime_observation") + ) + if memory_policy is not None: + canonical["runtime_memory_policy"] = memory_policy return FPMCollectionPlan( backend=backend, model_path=model_path, @@ -803,6 +959,7 @@ def build_collection_plan( cells=cells, sha256=_canonical_hash(canonical), _fpm_profile_json=profile_json, + _runtime_observation_json=runtime_observation_json, ) @@ -815,8 +972,8 @@ def _validate_profile_identities( system: str, backend: str, *, - max_decode_batch_size: int | None = None, -) -> None: + worker_type: str | None = None, +) -> tuple[FpmDeploymentProfile, ...]: """Require a resource declaration for every requested cell, before admission.""" if profile.architecture != capability.architecture: raise ValueError( @@ -835,6 +992,7 @@ def _validate_profile_identities( "enable_wideep", "enable_eplb", ) + deployments = [] for topology in topologies: deployment = profile.select( model=model_path, @@ -847,9 +1005,9 @@ def _validate_profile_identities( moe_tp_size=topology.moe_tp, moe_ep_size=topology.moe_ep, cp_size=topology.cp, + worker_type=worker_type or "aggregated", ) - if max_decode_batch_size is not None: - deployment.resources.validate_envelope(max_num_tokens=1, max_batch_size=max_decode_batch_size) + deployments.append(deployment) for kv_dtype in capability.dtype.kv_cache_dtypes: for policy in policies: resolved = [ @@ -873,3 +1031,4 @@ def _validate_profile_identities( f"FPM collection profile identity mismatch: {conflicts}; supply a profile matching " "the collection configuration. The profile does not override serving dispatch." ) + return tuple(deployments) diff --git a/python/aisimulate/collector/fpm_forward/repeatability.py b/python/aisimulate/collector/fpm_forward/repeatability.py new file mode 100644 index 000000000..b54f43153 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/repeatability.py @@ -0,0 +1,1218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Independent native sweep repetitions with auditable measurement evidence. + +AISimulate selects coordinates already measured by the source campaign. Dynamo +owns their admission and execution through its existing explicit-point API: +https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/benchmark_points.py +No runtime scheduler implementation is copied or replaced here. +""" + +from __future__ import annotations + +import copy +import json +import math +import statistics +from collections import Counter +from dataclasses import fields, replace +from pathlib import Path +from typing import Any + +from aisimulate.fpm_contract import FPM_MANIFEST_FILENAME + +from . import planner +from .capabilities import ModelCapabilityProfile, ResolvedDTypeProfile +from .config import FPMCollectionOptions, with_kv_warmup_defaults +from .execution_evidence import file_evidence, inspect_execution_evidence +from .measurement_evidence import compare_measurements, extract_measurement_evidence +from .memory_admission import DTypeMemoryEstimate, TopologyMemoryDecision +from .model_capability import ResolvedModelConfig +from .native_artifact import ( + NativeCollection, + NativePointMeasurement, + _rank_artifacts, + select_native_measurements, + validate_native_collection, +) +from .planner import BackendPolicy, FPMCollectionPlan, _canonical_hash, _hash_stable_admission +from .runner import ( + CHECKPOINT_SCHEMA, + FPM_RECOVERABLE_STATUSES, + _atomic_json, + _cell_runner, + _file_manifest, + _recover_completed_attempt, + _validate_points_receipts, + run_collection, +) +from .runtime_memory import cell_from_dict, validate_saved_plan +from .types import ParallelTopology + +PLAN_FILENAME = "repeatability-plan.json" +REPORT_FILENAME = "repeatability.json" +AGGREGATE_FILENAME = "repeatability-aggregate.json" +_LAUNCH_FILES = ("generator-request.json", "run.sh", "fpm_env.sh", "collector-runtime-env.sh") + + +class _IncompleteSample(ValueError): + pass + + +def load_repeatability_deployment(source_campaign_dir: str | Path) -> dict[str, Any]: + """Load the archived effective deployment inputs and verify their plan hash.""" + root = Path(source_campaign_dir).expanduser().resolve() + path = root / "generator-overrides.json" + if not path.is_file(): + raise ValueError( + "source campaign lacks archived generator-overrides.json; supply the original explicit " + "deployment inputs and verify their generator_config_sha256, or collect a new source campaign" + ) + overrides = json.loads(path.read_text()) + saved = json.loads((root / "collection-plan.json").read_text()) + validate_saved_plan(saved) + if not isinstance(overrides, dict): + raise ValueError("archived deployment inputs must be an object") + if not isinstance(overrides.get("K8sConfig", {}), dict): + raise ValueError("archived K8sConfig must be an object") + if _canonical_hash(with_kv_warmup_defaults(overrides)) != saved["generator_config_sha256"]: + raise ValueError("archived deployment inputs differ from the source plan") + return overrides + + +def load_repeatability_source(source_campaign_dir: str | Path) -> FPMCollectionPlan: + """Load an exact archived runnable plan without resolving models or hardware. + + Deployment-only Generator inputs are supplied separately to run_repeatability + and checked against the source's frozen generator_config_sha256. This loader + never upgrades historical schemas or applies new planning defaults silently. + """ + path = Path(source_campaign_dir).expanduser().resolve() / "collection-plan.json" + saved = json.loads(path.read_text()) + validate_saved_plan(saved) + if saved["schema_version"] != 11: + raise ValueError( + "repeatability execution requires a schema-v11 source campaign; historical data stays readable" + ) + raw_options = saved["options"] + option_names = {item.name for item in fields(FPMCollectionOptions)} + options = {name: raw_options[name] for name in option_names if name in raw_options} + for name, value in list(options.items()): + if isinstance(value, list): + options[name] = tuple(value) + sampling = raw_options["prefill_sampling"] + options.update( + warmup_iterations=raw_options["global_warmup_iterations"], + max_prefill_isl=sampling["max_isl"], + max_prefill_batch_size=sampling["max_batch_size"], + max_prefill_cudagraph_size=sampling["max_cudagraph_capture_size"], + prefill_cudagraph_policy=sampling.get("cudagraph_policy", "explicit"), + ) + if points := raw_options.get("benchmark_points"): + options.update( + benchmark_points_json=json.dumps(points["payload"], sort_keys=True, separators=(",", ":")), + benchmark_points_sha256=points["sha256"], + ) + capability = copy.deepcopy(saved["capability"]) + config = capability.pop("model_config") + dtype = capability.pop("dtype") + dtype["kv_cache_dtypes"] = tuple(dtype["kv_cache_dtypes"]) + resolved_dtype = ResolvedDTypeProfile(**dtype) + resolved_capability = ModelCapabilityProfile( + **capability, + dtype=resolved_dtype, + model_config=ResolvedModelConfig( + config["payload"], source_kind=config["source_kind"], source_reference=str(path) + ), + ) + admissions = [] + for item in saved["topology_memory_admission"]: + envelope = item["activation_envelope"] + estimates = [] + for estimate in item["estimates"]: + estimates.append( + DTypeMemoryEstimate(**{key: value for key, value in estimate.items() if key != "headroom_bytes"}) + ) + admissions.append( + TopologyMemoryDecision( + topology=ParallelTopology(**item["topology"]), + disposition=item["disposition"], + source=item["source"], + reason=item["reason"], + estimates=tuple(estimates), + max_new_tokens=envelope.get("requested_prefill_tokens", envelope.get("max_new_tokens")), + max_batch_size=envelope["max_batch_size"], + profile_max_num_tokens=envelope.get("max_num_tokens"), + ) + ) + raw_dtype = dict(saved["dtype_profile"]) + raw_dtype["kv_cache_dtypes"] = tuple(raw_dtype["kv_cache_dtypes"]) + plan = FPMCollectionPlan( + **{ + name: saved[name] + for name in ("backend", "model_path", "system", "aic_revision", "generator_config_sha256", "sha256") + }, + options=FPMCollectionOptions(**options), + capability=resolved_capability, + dtype_profile=ResolvedDTypeProfile(**raw_dtype), + topologies=tuple( + ParallelTopology(**{key: value for key, value in item.items() if key != "strategy"}) + for item in saved["topologies"] + ), + topology_memory_admission=tuple(admissions), + backend_policies=tuple(BackendPolicy(**item) for item in saved["backend_policies"]), + cells=tuple(cell_from_dict(item) for item in saved["cells"]), + _fpm_profile_json=json.dumps(saved["fpm_profile"]) if "fpm_profile" in saved else None, + _runtime_observation_json=json.dumps(saved["runtime_observation"]) if "runtime_observation" in saved else None, + ) + if plan.to_dict() != saved: + raise ValueError("saved repeatability source cannot be reconstructed exactly by this collector version") + return plan + + +def _coordinates(point: dict[str, Any]) -> dict[str, Any]: + coordinates = {key: point[key] for key in ("batch_size", "total_kv_read_tokens")} + if point["point_type"] == "prefill": + coordinates["total_prefill_tokens"] = point["total_prefill_tokens"] + for key in ("partition", "rows"): + if point.get(key) is not None: + coordinates[key] = point[key] + return coordinates + + +def _point_key(point: dict[str, Any]) -> str: + return _canonical_hash(_coordinates(point)) + + +def _regime(measurement: NativePointMeasurement) -> tuple[Any, ...]: + point = measurement.point + return ( + measurement.kv_seed_regime, + point.get("expected_cudagraph_mode"), + point.get("expected_capture_size"), + point.get("padding_tokens"), + ) + + +def _select_points(collection: NativeCollection, limit: int | None, *, cell_id: str) -> list[dict[str, Any]]: + """Cover extremes and regime/boundary representatives with a fixed budget.""" + points = sorted( + select_native_measurements(collection, cell_id=cell_id), + key=lambda item: ( + item.point["batch_size"], + item.point["total_prefill_tokens"], + item.point["total_kv_read_tokens"], + _point_key(item.point), + ), + ) + if len({_point_key(item.point) for item in points}) != len(points): + raise ValueError("repeatability requires unique native coordinates; duplicate regimes cannot be blended") + roles: list[set[str]] = [set() for _ in points] + axes = { + "batch": [item.point["batch_size"] for item in points], + "kv_per_request": [item.point["total_kv_read_tokens"] / item.point["batch_size"] for item in points], + "new_tokens": [item.point["total_prefill_tokens"] for item in points], + } + for name, values in axes.items(): + for index, value in enumerate(values): + for label, edge in (("minimum", min(values)), ("maximum", max(values))): + if value == edge: + roles[index].add(f"{name}_{label}") + for index, item in enumerate(points): + roles[index].add(f"regime:{item.kv_seed_regime or 'unreported'}:{item.point.get('expected_cudagraph_mode')}") + # Select the first and last observed capture boundary with neighbours that + # actually exist in this grid. We do not invent new coordinates or promise + # every capture size is represented by a bounded subset. + token_axis = [item.point["total_prefill_tokens"] or item.point["batch_size"] for item in points] + captures = sorted({value for item in points if (value := item.point.get("expected_capture_size")) is not None}) + for boundary in sorted(set(captures[:1] + captures[-1:])): + for label, values in ( + ("below", [value for value in token_axis if value < boundary]), + ("at", [value for value in token_axis if value == boundary]), + ("above", [value for value in token_axis if value > boundary]), + ): + if not values: + continue + target = max(values) if label == "below" else min(values) + for index, value in enumerate(token_axis): + if value == target: + roles[index].add(f"capture_{boundary}_{label}") + uncovered = set().union(*roles) if limit is not None else set() + selected = set() if limit is not None else set(range(len(points))) + while uncovered: + index = max(range(len(points)), key=lambda index: (len(roles[index] & uncovered), -index)) + if len(selected) == limit: + raise ValueError( + f"repeatability max_points_per_cell={limit} cannot cover observed regimes/extremes; " + f"increase the bound (uncovered: {sorted(uncovered)})" + ) + selected.add(index) + uncovered -= roles[index] + # Fill remaining slots evenly across the sorted native grid. Every sample + # remains a measured point, including when the source grid itself is small. + target_count = min(limit, len(points)) if limit is not None else len(points) + candidates = [round(index * (len(points) - 1) / max(1, target_count - 1)) for index in range(target_count)] + for index in [*candidates, *range(len(points))]: + if len(selected) >= target_count: + break + selected.add(index) + return [ + { + "key": _point_key(points[index].point), + "coordinates": _coordinates(points[index].point), + "source_point": points[index].point, + "source_wall_time_seconds": max(value for _, value in points[index].rank_wall_times), + "source_rank_wall_times": [list(value) for value in points[index].rank_wall_times], + "kv_seed_regime": points[index].kv_seed_regime, + "selection_reasons": sorted(roles[index]) if limit is not None else ["full_native_sweep"], + } + for index in sorted(selected) + ] + + +def freeze_repeatability_plan( + source_plan: FPMCollectionPlan, + source_campaign_dir: str | Path, + source_checkpoint_path: str | Path, + *, + samples: int = 5, + max_points_per_cell: int = 12, + cv_threshold: float = 0.05, + comparison_mode: str = "full_grid", + max_attempts_per_sample: int = 2, + source_agreement_threshold: float = 0.05, + observation_evidence_version: int = 2, +) -> dict[str, Any]: + """Inspect validated source artifacts and return a deterministic frozen plan.""" + if type(observation_evidence_version) is not int or observation_evidence_version not in {1, 2}: + raise ValueError("unsupported repeatability observation evidence version") + if type(samples) is not int or samples < 2: + raise ValueError("repeatability samples must be an integer >= 2") + if type(max_points_per_cell) is not int or max_points_per_cell < 1: + raise ValueError("repeatability max_points_per_cell must be a positive integer") + if type(cv_threshold) not in (int, float) or not math.isfinite(cv_threshold) or cv_threshold < 0: + raise ValueError("repeatability cv_threshold must be finite and nonnegative") + if comparison_mode not in {"full_grid", "bounded"}: + raise ValueError("repeatability comparison_mode must be full_grid or bounded") + if type(max_attempts_per_sample) is not int or max_attempts_per_sample < 1: + raise ValueError("repeatability max_attempts_per_sample must be a positive integer") + if ( + type(source_agreement_threshold) not in (int, float) + or not math.isfinite(source_agreement_threshold) + or source_agreement_threshold < 0 + ): + raise ValueError("repeatability source_agreement_threshold must be finite and nonnegative") + root = Path(source_campaign_dir).expanduser().resolve() + checkpoint_path = Path(source_checkpoint_path).expanduser().resolve() + saved = json.loads((root / "collection-plan.json").read_text()) + validate_saved_plan(saved) + if saved != source_plan.to_dict(): + raise ValueError("repeatability source plan differs from the saved campaign; reuse its exact frozen inputs") + checkpoint = json.loads(checkpoint_path.read_text()) + if checkpoint.get("schema") != CHECKPOINT_SCHEMA or checkpoint.get("plan_sha256") != source_plan.sha256: + raise ValueError("repeatability source checkpoint does not match its campaign") + cells = [] + for cell in source_plan.cells: + entry = checkpoint.get("cells", {}).get(cell.cell_id, {}) + if entry.get("status") != "passed" or not entry.get("attempt_id"): + raise ValueError(f"repeatability requires a passed source cell: {cell.cell_id}") + cell_dir = root / "cells" / cell.cell_id + collection = validate_native_collection( + cell, + cell_dir / "raw", + expected_plan_sha256=source_plan.sha256, + expected_attempt_id=entry["attempt_id"], + ) + _validate_points_receipts(source_plan, cell, cell_dir / "raw", entry["attempt_id"]) + selected = _select_points( + collection, max_points_per_cell if comparison_mode == "bounded" else None, cell_id=cell.cell_id + ) + capture_sizes = sorted( + {value for item in collection.points if (value := item.point.get("expected_capture_size")) is not None} + ) + selected_boundaries = ( + sorted(set(capture_sizes[:1] + capture_sizes[-1:])) if comparison_mode == "bounded" else capture_sizes + ) + manifest = {"schema_version": 3, "prefill": [], "decode": []} + manifest[cell.workload_kind] = [item["coordinates"] for item in selected] + cells.append( + { + "cell_id": cell.cell_id, + "phase": cell.workload_kind, + "source_attempt_id": collection.collector_attempt_id, + "source_runtime_run_id": collection.runtime_run_id, + "source_measurement": extract_measurement_evidence( + cell_dir / "raw", collection, evidence_version=observation_evidence_version + ), + "source_point_count": len(collection.points), + "points": selected, + "selection_coverage": { + "covered_strata": sorted({reason for point in selected for reason in point["selection_reasons"]}), + "uncovered_strata": [], + "observed_expected_capture_sizes": capture_sizes, + "selected_capture_boundaries": selected_boundaries, + "unselected_capture_boundaries": sorted(set(capture_sizes) - set(selected_boundaries)), + }, + "benchmark_points": manifest, + "source_files": _file_manifest(cell_dir / "raw"), + "source_launch": { + name: file_evidence(cell_dir / name) if (cell_dir / name).is_file() else None + for name in _LAUNCH_FILES + }, + "execution": inspect_execution_evidence( + cell, cell_dir / "raw", collection, plan=source_plan, evidence_version=observation_evidence_version + ), + } + ) + payload = { + **({"observation_evidence_version": 2} if observation_evidence_version == 2 else {}), + "schema_name": "aisimulate_fpm_repeatability_plan", + "schema_version": 2, + "source_plan_sha256": source_plan.sha256, + "source_campaign_dir": str(root), + "source_checkpoint": file_evidence(checkpoint_path), + "source_deployment": file_evidence(root / "generator-overrides.json") + if (root / "generator-overrides.json").is_file() + else None, + "policy": { + "samples": samples, + "max_points_per_cell": max_points_per_cell, + "cv_threshold": cv_threshold, + "comparison_mode": comparison_mode, + "max_attempts_per_sample": max_attempts_per_sample, + "source_agreement_threshold": source_agreement_threshold, + }, + "sampling": "independent_native_full_sweep_launches" + if comparison_mode == "full_grid" + else "independent_native_subset_launches", + "cells": cells, + } + return {**payload, "sha256": _canonical_hash(payload)} + + +def _subset_plan( + source: FPMCollectionPlan, selected: dict[str, Any], *, comparison_mode: str = "bounded" +) -> FPMCollectionPlan: + canonical = json.dumps(selected["benchmark_points"], sort_keys=True, separators=(",", ":")) + options = replace( + source.options, + benchmark_points_json=canonical, + benchmark_points_sha256=_canonical_hash(selected["benchmark_points"]), + ) + if comparison_mode == "full_grid": + options = source.options + plan = replace( + source, options=options, cells=tuple(cell for cell in source.cells if cell.cell_id == selected["cell_id"]) + ) + # Use the ordinary collection-plan hash contract so existing native readers + # can independently validate every sample campaign. + payload = { + "backend": plan.backend, + "model_path": plan.model_path, + "system": plan.system, + "aic_revision": plan.aic_revision, + "generator_config_sha256": plan.generator_config_sha256, + "options": options.to_dict(), + "capability": plan.capability.to_dict(), + "dtype_profile": plan.dtype_profile.to_dict(), + "point_generation": "dynamo_native_self_benchmark", + "topology_memory_admission": [_hash_stable_admission(item) for item in plan.topology_memory_admission], + "topologies": [item.to_dict() for item in plan.topologies], + "policies": [item.to_dict() for item in plan.backend_policies], + "cells": [cell.to_dict() for cell in plan.cells], + } + serialized = plan.to_dict() + for key in ("fpm_profile", "runtime_memory_policy", "runtime_observation"): + if key in serialized: + payload[key] = serialized[key] + return replace(plan, sha256=_canonical_hash(payload)) + + +def _worker_identity(evidence): + return sorted( + [ + { + key: worker[key] + for key in ( + "dp_rank", + "tp_rank", + "pp_rank", + "backend_version", + "attention_groups", + "graph_config", + "resolved_config", + ) + } + for worker in evidence["observed_workers"] + ], + key=lambda worker: (worker["dp_rank"], worker["tp_rank"], worker["pp_rank"]), + ) + + +def _sample_evidence(plan: FPMCollectionPlan, selected: dict[str, Any], directory: Path) -> dict[str, Any]: + checkpoint = json.loads((directory / "checkpoint" / "fpm_forward.json").read_text()) + cell = plan.cells[0] + entry = checkpoint.get("cells", {}).get(cell.cell_id, {}) + if entry.get("status") not in {"passed", "cleanup_failed", *FPM_RECOVERABLE_STATUSES} or not entry.get( + "attempt_id" + ): + raise _IncompleteSample("repeatability sample collector checkpoint is not complete") + if checkpoint.get("schema") != CHECKPOINT_SCHEMA or checkpoint.get("plan_sha256") != plan.sha256: + raise ValueError("repeatability sample collector checkpoint identity changed") + root = directory / "artifacts" / plan.sha256[:16] / "cells" / cell.cell_id + rank_payloads = _rank_artifacts(root / "raw") + if not rank_payloads or len(rank_payloads) < cell.topology.dp: + raise _IncompleteSample("repeatability sample is missing native rank artifacts") + collection = validate_native_collection( + cell, + root / "raw", + expected_plan_sha256=plan.sha256, + expected_attempt_id=entry["attempt_id"], + ) + _validate_points_receipts(plan, cell, root / "raw", entry["attempt_id"]) + canonical = select_native_measurements(collection, cell_id=cell.cell_id) + actual = {_point_key(item.point): item for item in canonical} + expected = {item["key"]: item for item in selected["points"]} + if len(actual) != len(canonical) or set(actual) != set(expected): + raise ValueError("repeatability runtime did not measure exactly the frozen coordinates") + for key, item in actual.items(): + original = expected[key] + original_measurement = NativePointMeasurement( + point=original["source_point"], + rank_wall_times=(), + kv_seed_regime=original["kv_seed_regime"], + ) + if _regime(item) != _regime(original_measurement): + raise ValueError(f"repeatability execution/seed regime changed for point {key}") + evidence_version = selected["source_measurement"]["schema_version"] + execution = inspect_execution_evidence(cell, root / "raw", collection, plan=plan, evidence_version=evidence_version) + if selected["execution"]["status"] == execution["status"] == "qualified" and _worker_identity( + execution + ) != _worker_identity(selected["execution"]): + raise ValueError("repeatability observed attention backend, graph or runtime configuration changed") + return { + "collector_status": entry["status"], + "collector_error": {key: entry[key] for key in ("error_type", "error") if key in entry}, + "cleanup_error": entry.get("cleanup_error"), + "attempt_id": collection.collector_attempt_id, + "runtime_run_id": collection.runtime_run_id, + "runtime_grid_digest": collection.runtime_grid_digest, + "measurement": extract_measurement_evidence(root / "raw", collection, evidence_version=evidence_version), + "points": { + key: { + "wall_time_seconds": max(value for _, value in item.rank_wall_times), + "rank_wall_times": [list(value) for value in item.rank_wall_times], + "point": item.point, + "kv_seed_regime": item.kv_seed_regime, + } + for key, item in actual.items() + }, + "raw_files": _file_manifest(root / "raw"), + "launch": {name: file_evidence(root / name) if (root / name).is_file() else None for name in _LAUNCH_FILES}, + "execution": execution, + } + + +def _record_sample_evidence(plan, selected, directory, frozen, report, attempt): + evidence = _sample_evidence(plan, selected, directory) + used_attempts = {cell["source_attempt_id"] for cell in frozen["cells"]} + used_runs = {cell["source_runtime_run_id"] for cell in frozen["cells"]} + for entries in report["samples"].values(): + for entry in entries: + for prior in entry["attempts"]: + if prior.get("evidence"): + used_attempts.add(prior["evidence"]["attempt_id"]) + used_runs.add(prior["evidence"]["runtime_run_id"]) + if evidence["runtime_run_id"] in used_runs: + raise ValueError("repeatability reused a runtime run instead of an independent launch") + if evidence["attempt_id"] in used_attempts: + raise ValueError("repeatability reused a collector attempt instead of an independent launch") + attempt["evidence"] = evidence + if evidence["execution"]["status"] == "failed": + attempt.update( + status="failed", + failure_kind="validation_failed", + error="runtime execution evidence contradicts the frozen source", + ) + elif evidence["collector_status"] == "cleanup_failed" or evidence["cleanup_error"] is not None: + attempt.update(status="failed", failure_kind="cleanup_failed", error=evidence["cleanup_error"]) + elif evidence["collector_status"] in FPM_RECOVERABLE_STATUSES: + attempt.update( + status="failed", + failure_kind="postprocessing_failed", + error="complete native measurement retained; CPU-only post-processing recovery is required before " + "qualification or further launches; GPU replacement is disabled", + ) + elif not attempt.get("errors"): + attempt["status"] = "passed" + + +def _retry_sample_cleanup(plan, selected, directory, attempt): + """Retry teardown alone, preserving the original measurement and checkpoint.""" + cell_root = directory / "artifacts" / plan.sha256[:16] / "cells" / selected["cell_id"] + cleanup = {"status": "running"} + attempt.setdefault("cleanup_attempts", []).append(cleanup) + try: + manifest = cell_root / FPM_MANIFEST_FILENAME + if not manifest.is_file(): + raise ValueError("repeatability sample has no manifest to verify cleanup") + _cell_runner(plan, plan.cells[0], manifest, cell_root).cleanup() + if _sample_evidence(plan, selected, directory) != attempt["evidence"]: + raise ValueError("repeatability sample evidence changed during cleanup") + cleanup["status"] = "passed" + if attempt["evidence"]["collector_status"] in FPM_RECOVERABLE_STATUSES: + attempt.update(status="failed", failure_kind="postprocessing_failed") + else: + attempt["status"] = "passed" + except (KeyboardInterrupt, SystemExit): + cleanup["status"] = "interrupted" + raise + except Exception as error: + cleanup.update(status="failed", error=f"{type(error).__name__}: {error}") + finally: + attempt.setdefault("cleanup_file_history", []).append(attempt.get("files", {})) + attempt["files"] = _file_manifest(directory) + + +def _retry_sample_postprocessing(plan, selected, directory, attempt): + """Validate salvaged artifacts without relaunching or rewriting their checkpoint.""" + recovery = {"status": "running"} + attempt.setdefault("postprocessing_attempts", []).append(recovery) + try: + checkpoint = json.loads((directory / "checkpoint" / "fpm_forward.json").read_text()) + entry = dict(checkpoint["cells"][selected["cell_id"]]) + if attempt.get("cleanup_attempts") and attempt["cleanup_attempts"][-1]["status"] == "passed": + entry.pop("cleanup_error", None) + recovered = _recover_completed_attempt(plan, plan.cells[0], directory / "artifacts" / plan.sha256[:16], entry) + if recovered is None: + raise ValueError( + "CPU-only recovery could not validate the retained observation; inspect collector artifacts " + "and teardown before retrying recovery; the original GPU measurement cannot be replaced" + ) + if _sample_evidence(plan, selected, directory) != attempt["evidence"]: + raise ValueError("repeatability sample evidence changed during post-processing recovery") + recovery.update(status="passed", artifact_recovery=recovered["artifact_recovery"]) + attempt["status"] = "passed" + except (KeyboardInterrupt, SystemExit): + recovery["status"] = "interrupted" + raise + except Exception as error: + recovery.update(status="failed", error=f"{type(error).__name__}: {error}") + finally: + attempt.setdefault("postprocessing_file_history", []).append(attempt.get("files", {})) + attempt["files"] = _file_manifest(directory) + + +def _comparison_status(comparisons: list[dict[str, Any]]) -> str: + if any(item["status"] == "mismatch" for item in comparisons): + return "mismatch" + if not comparisons or any(item["status"] != "comparable" for item in comparisons): + return "unestablished" + return "comparable" + + +def _summarize(frozen: dict[str, Any], report: dict[str, Any]) -> None: + """Keep independent observations, identity eligibility and stability separate.""" + points = [] + required = frozen["policy"]["samples"] + full_grid = frozen["policy"]["comparison_mode"] == "full_grid" + execution_failed = any(cell["execution"]["status"] == "failed" for cell in frozen["cells"]) + execution_complete = True + validation_failed = False + for cell in frozen["cells"]: + slots = report["samples"][cell["cell_id"]] + latest = [sample["attempts"][-1] for sample in slots if sample["attempts"]] + validation_failed |= any(item.get("failure_kind") == "validation_failed" for item in latest) + # A valid observation is never superseded. Preserve even a completed + # measurement whose collection subsequently failed during teardown. + measured = [attempt for slot in slots for attempt in slot["attempts"] if attempt.get("evidence")] + execution_failed |= any(item["evidence"]["execution"]["status"] == "failed" for item in measured) + execution_complete &= len(measured) == required and all( + item["status"] == "passed" and item["evidence"]["execution"]["status"] == "qualified" for item in measured + ) + reference = measured[0]["evidence"].get("measurement", {}) if measured else {} + for point in cell["points"]: + key = point["key"] + comparisons = [] + if measured: + first_point = reference.get("points", {}).get(key, {}) + first_recorded = ( + not reference.get("reasons") + and first_point.get("status") == "recorded" + and measured[0]["evidence"]["execution"]["status"] == "qualified" + ) + comparisons.append( + { + "status": "mismatch" + if measured[0]["evidence"]["execution"]["status"] == "failed" + else "comparable" + if first_recorded + else "unestablished", + "reasons": [*reference.get("reasons", []), *first_point.get("reasons", [])], + "scope": "reference observation; independence checked by launch identities", + } + ) + for item in measured[1:]: + comparison = compare_measurements(reference, item["evidence"].get("measurement", {}), point_key=key) + if item["evidence"]["execution"]["status"] != "qualified": + comparison = { + **comparison, + "status": "mismatch" + if item["evidence"]["execution"]["status"] == "failed" + else "unestablished", + "reasons": [*comparison["reasons"], "worker execution is not qualified"], + } + if ( + measured[0]["evidence"]["execution"]["status"] == "qualified" + and item["evidence"]["execution"]["status"] == "qualified" + and _worker_identity(measured[0]["evidence"]["execution"]) + != _worker_identity(item["evidence"]["execution"]) + ): + comparison = { + **comparison, + "status": "mismatch", + "reasons": [*comparison["reasons"], "observed worker runtime configuration differs"], + } + comparisons.append(comparison) + comparable = _comparison_status(comparisons) + source_comparison = compare_measurements( + cell["source_measurement"], reference, point_key=key, same_context=full_grid + ) + if cell["execution"]["status"] != "qualified": + source_comparison = { + **source_comparison, + "status": "mismatch" if cell["execution"]["status"] == "failed" else "unestablished", + "reasons": [*source_comparison["reasons"], "source worker execution is not qualified"], + } + raw_values = [item["evidence"]["points"][key]["wall_time_seconds"] for item in measured] + # Raw timings remain visible even when identity cannot be established; + # they must not be pooled into a qualified numerical population. + values = raw_values if comparable == "comparable" else [] + mean = statistics.mean(values) if values else None + median = statistics.median(values) if values else None + stddev = statistics.stdev(values) if len(values) >= 2 else None + cv = stddev / mean if stddev is not None else None + inclusive = [point["source_wall_time_seconds"], *values] + inclusive_stddev = statistics.stdev(inclusive) if values else None + inclusive_cv = inclusive_stddev / statistics.mean(inclusive) if inclusive_stddev is not None else None + difference = ( + abs(median / point["source_wall_time_seconds"] - 1) + if median is not None and source_comparison["status"] == "comparable" + else None + ) + status = ( + "mismatch" + if comparable == "mismatch" + else "unestablished" + if comparable != "comparable" + else "incomplete" + if len(values) != required + else "passed" + if cv is not None and cv <= frozen["policy"]["cv_threshold"] + else "unstable" + ) + points.append( + { + "cell_id": cell["cell_id"], + "phase": cell["phase"], + "key": key, + "coordinates": point["coordinates"], + "source_wall_time_seconds": point["source_wall_time_seconds"], + "raw_samples_seconds": raw_values, + "samples_seconds": values, + "sample_count": len(values), + "raw_sample_count": len(raw_values), + "observations": [ + { + "attempt_id": item["evidence"]["attempt_id"], + "runtime_run_id": item["evidence"]["runtime_run_id"], + "directory": item["directory"], + "status": item["status"], + **item["evidence"]["points"][key], + } + for item in measured + ], + "comparability": {"status": comparable, "comparisons": comparisons}, + "source_comparison": source_comparison, + "mean_seconds": mean, + "median_seconds": median, + "minimum_seconds": min(values) if values else None, + "maximum_seconds": max(values) if values else None, + "sample_cv": cv, + "sample_stddev_seconds": stddev, + "source_inclusive_cv": inclusive_cv, + "source_inclusive_stddev_seconds": inclusive_stddev, + "source_inclusive_sample_count": len(inclusive), + "source_inclusive_usage": "diagnostic_only; source is not an independent fresh repetition", + "source_relative_difference": difference, + "source_agreement": ( + "unestablished" + if difference is None or len(values) != required + else "passed" + if difference <= frozen["policy"]["source_agreement_threshold"] + else "failed" + ), + "status": status, + } + ) + counts = Counter(item["status"] for item in points) + repeatability = ( + "unstable" + if counts["unstable"] + else "mismatch" + if counts["mismatch"] + else "unestablished" + if counts["unestablished"] + else "incomplete" + if counts["incomplete"] + else "passed" + ) + new_status = ( + "unstable" + if repeatability == "unstable" + else "failed" + if execution_failed or validation_failed or counts["mismatch"] + else "qualified" + if repeatability == "passed" and execution_complete + else "unestablished" + if repeatability == "unestablished" + else "incomplete" + ) + source_status = ( + "failed" + if any( + point["source_comparison"]["status"] == "mismatch" or point["source_agreement"] == "failed" + for point in points + ) + else "qualified" + if full_grid + and new_status == "qualified" + and all( + point["source_comparison"]["status"] == "comparable" and point["source_agreement"] == "passed" + for point in points + ) + else "unestablished" + ) + report.update( + points=points, + repeatability={ + "status": repeatability, + "cv_threshold": frozen["policy"]["cv_threshold"], + "point_counts": dict(counts), + "measurement_count": "independent fresh launches; adjacent internal steps are not repetitions", + "time_unit": "seconds", + "standard_deviation_denominator": "n - 1", + }, + new_population={"status": new_status, "scope": "full_native_sweep" if full_grid else "bounded_diagnostic"}, + source_qualification={ + "status": source_status, + "source_agreement_threshold": frozen["policy"]["source_agreement_threshold"], + "agreement_metric": "abs(fresh_median / source_wall_time - 1)", + "scope": "original_source" if full_grid else "bounded_subset_cannot_qualify_full_source", + }, + execution={ + "status": "failed" + if execution_failed or validation_failed + else "qualified" + if execution_complete + else "incomplete" + }, + status=( + "failed" + if new_status in {"failed", "unstable"} or source_status == "failed" + else "passed" + if new_status == source_status == "qualified" + else "incomplete" + ), + ) + + +def _save_report(root: Path, report: dict[str, Any]) -> None: + """Publish an auditable derived estimate without replacing formal source data.""" + aggregate = { + "schema_name": "aisimulate_fpm_repeatability_aggregate", + "schema_version": 1, + "plan_sha256": report["plan_sha256"], + "status": report["new_population"]["status"], + "scope": report["new_population"]["scope"], + "method": "median_of_independent_launch_max_rank_estimates", + "formal_source_replaced": False, + "points": [ + { + key: point[key] + for key in ( + "cell_id", + "phase", + "key", + "coordinates", + "median_seconds", + "samples_seconds", + "raw_samples_seconds", + "sample_cv", + "status", + "comparability", + "observations", + ) + } + for point in report["points"] + ], + } + # The durable report must keep referencing its previous immutable aggregate + # until the new report is committed. An interruption cannot orphan its hash. + aggregate_path = root / f"{Path(AGGREGATE_FILENAME).stem}-{_canonical_hash(aggregate)}.json" + if aggregate_path.exists(): + if json.loads(aggregate_path.read_text()) != aggregate: + raise ValueError("repeatability aggregate generation changed") + else: + _atomic_json(aggregate_path, aggregate) + report["aggregate"] = file_evidence(aggregate_path) + _atomic_json(root / REPORT_FILENAME, report) + + +def assess_repeatability( + frozen_plan: dict[str, Any], + report: dict[str, Any], + *, + cv_threshold: float, + source_agreement_threshold: float | None = None, +) -> dict[str, Any]: + """Reassess preserved samples without scheduling any additional collection. + + The returned assessment records the original frozen plan and the changed + criterion. A caller must retain it at a new report path, preserving history. + Source and raw sample evidence are revalidated before any reassessment. + """ + if type(cv_threshold) not in (int, float) or not math.isfinite(cv_threshold) or cv_threshold < 0: + raise ValueError("repeatability cv_threshold must be finite and nonnegative") + if report.get("plan_sha256") != frozen_plan.get("sha256"): + raise ValueError("repeatability report does not match its frozen plan") + payload = {key: value for key, value in frozen_plan.items() if key != "sha256"} + if _canonical_hash(payload) != frozen_plan.get("sha256"): + raise ValueError("repeatability frozen plan hash changed") + if frozen_plan.get("schema_version") == 1: + return _assess_legacy(frozen_plan, report, cv_threshold=cv_threshold) + if frozen_plan.get("schema_version") != 2 or report.get("schema_version") != 2: + raise ValueError("unsupported repeatability evidence schema") + source = load_repeatability_source(frozen_plan["source_campaign_dir"]) + fresh = freeze_repeatability_plan( + source, + frozen_plan["source_campaign_dir"], + frozen_plan["source_checkpoint"]["path"], + observation_evidence_version=frozen_plan.get("observation_evidence_version", 1), + **frozen_plan["policy"], + ) + if fresh != frozen_plan: + raise ValueError("repeatability source artifacts changed before reassessment") + _validate_saved_samples(source, frozen_plan, report, Path(report["output_dir"]).resolve()) + result = copy.deepcopy(report) + assessment_plan = copy.deepcopy(frozen_plan) + assessment_plan["policy"]["cv_threshold"] = cv_threshold + if source_agreement_threshold is not None: + if ( + type(source_agreement_threshold) not in (int, float) + or not math.isfinite(source_agreement_threshold) + or source_agreement_threshold < 0 + ): + raise ValueError("repeatability source_agreement_threshold must be finite and nonnegative") + assessment_plan["policy"]["source_agreement_threshold"] = source_agreement_threshold + result["assessment"] = { + "source_report_sha256": _canonical_hash(report), + "cv_threshold": cv_threshold, + "source_agreement_threshold": assessment_plan["policy"]["source_agreement_threshold"], + } + _summarize(assessment_plan, result) + return result + + +def _assess_legacy(frozen, report, *, cv_threshold): + """Keep historical bounded reports inspectable without a protocol upgrade.""" + root = Path(frozen["source_campaign_dir"]).resolve() + if file_evidence(Path(frozen["source_checkpoint"]["path"])) != frozen["source_checkpoint"]: + raise ValueError("repeatability historical source checkpoint changed") + for cell in frozen["cells"]: + if _file_manifest(root / "cells" / cell["cell_id"] / "raw") != cell["source_files"]: + raise ValueError("repeatability historical source artifacts changed") + output = Path(report["output_dir"]).resolve() + for samples in report["samples"].values(): + for sample in samples: + for attempt in sample["attempts"]: + path = (output / attempt["directory"]).resolve() + if not path.is_relative_to(output) or _file_manifest(path) != attempt.get("files"): + raise ValueError("repeatability historical sample artifacts changed") + result = copy.deepcopy(report) + result["historical_assessment"] = { + key: report[key] for key in ("status", "repeatability", "execution") if key in report + } + result["assessment"] = { + "source_report_sha256": _canonical_hash(report), + "cv_threshold": cv_threshold, + "reason": "legacy bounded measurements lack a frozen comparable protocol", + } + result["status"] = "incomplete" + result["source_qualification"] = {"status": "unestablished"} + result["new_population"] = {"status": "unestablished", "scope": "legacy_diagnostic"} + return result + + +def _validate_saved_samples(source, frozen, report, root): + if report.get("schema_version") != 2: + raise ValueError("legacy repeatability evidence cannot resume execution; preserve it for offline inspection") + if report.get("aggregate") and file_evidence(Path(report["aggregate"]["path"])) != report["aggregate"]: + raise ValueError("repeatability aggregate artifact changed") + if report.get("output_dir") != str(root): + raise ValueError("repeatability report output directory changed") + if report.get("policy") != frozen["policy"]: + raise ValueError("repeatability saved measurement policy changed") + expected_cells = {cell["cell_id"] for cell in frozen["cells"]} + if set(report.get("samples", {})) != expected_cells: + raise ValueError("repeatability saved sample cells do not match the plan") + seen_attempts = {cell["source_attempt_id"] for cell in frozen["cells"]} + seen_runs = {cell["source_runtime_run_id"] for cell in frozen["cells"]} + for cell in frozen["cells"]: + samples = report["samples"][cell["cell_id"]] + if len(samples) != frozen["policy"]["samples"]: + raise ValueError("repeatability saved sample count changed") + plan = _subset_plan(source, cell, comparison_mode=frozen["policy"]["comparison_mode"]) + for index, sample in enumerate(samples, start=1): + if sample.get("sample_index") != index or not isinstance(sample.get("attempts"), list): + raise ValueError("repeatability saved sample indices are invalid") + if len(sample["attempts"]) > frozen["policy"]["max_attempts_per_sample"]: + raise ValueError("repeatability saved attempts exceed the frozen retry budget") + for attempt_index, attempt in enumerate(sample["attempts"], start=1): + relative = Path("samples") / cell["cell_id"] / f"sample-{index:02d}" / f"attempt-{attempt_index:02d}" + directory = root / relative + if attempt.get("directory") != str(relative) or not directory.resolve().is_relative_to(root): + raise ValueError("repeatability saved attempt directory is invalid") + if attempt.get("status") not in {"running", "interrupted", "failed", "passed"}: + raise ValueError("repeatability saved attempt status is invalid") + if attempt.get("files") is not None and _file_manifest(directory) != attempt["files"]: + raise ValueError("repeatability attempt files changed") + if attempt["status"] != "passed" and not attempt.get("evidence"): + continue + if attempt_index != len(sample["attempts"]): + raise ValueError("repeatability valid measurement was superseded") + evidence = _sample_evidence(plan, cell, directory) + if evidence != attempt.get("evidence"): + raise ValueError("saved repeatability sample evidence changed") + if ( + attempt["status"] == "passed" + and (evidence["collector_status"] == "cleanup_failed" or evidence["cleanup_error"] is not None) + and not (attempt.get("cleanup_attempts") and attempt["cleanup_attempts"][-1]["status"] == "passed") + ): + raise ValueError("repeatability sample cleanup is not resolved") + if ( + attempt["status"] == "passed" + and evidence["collector_status"] in FPM_RECOVERABLE_STATUSES + and not ( + attempt.get("postprocessing_attempts") + and attempt["postprocessing_attempts"][-1]["status"] == "passed" + ) + ): + raise ValueError("repeatability sample post-processing is not resolved") + if evidence["attempt_id"] in seen_attempts: + raise ValueError("repeatability reused an attempt instead of an independent launch") + if evidence["runtime_run_id"] in seen_runs: + raise ValueError("repeatability reused a runtime run instead of an independent launch") + seen_attempts.add(evidence["attempt_id"]) + seen_runs.add(evidence["runtime_run_id"]) + + +def run_repeatability( + source_plan: FPMCollectionPlan, + *, + generator_overrides: dict[str, Any], + source_campaign_dir: str | Path, + source_checkpoint_path: str | Path, + output_dir: str | Path, + resume: bool = False, + retry_failed: bool = False, + samples: int = 5, + max_points_per_cell: int = 12, + cv_threshold: float = 0.05, + comparison_mode: str = "full_grid", + max_attempts_per_sample: int = 2, + source_agreement_threshold: float = 0.05, + observation_evidence_version: int | None = None, +) -> dict[str, Any]: + """Repeat the original sweep, or collect an explicitly diagnostic bounded subset. + + The caller owns execution authorization and an allocation for the selected + Kubernetes or Slurm transport. No serving data or formal table is replaced. + """ + if retry_failed and not resume: + raise ValueError("repeatability retry_failed requires resume") + if _canonical_hash(with_kv_warmup_defaults(generator_overrides)) != source_plan.generator_config_sha256: + raise ValueError("repeatability deployment inputs differ from the original launch") + root = Path(output_dir).expanduser().resolve() + previous = json.loads((root / PLAN_FILENAME).read_text()) if resume and (root / PLAN_FILENAME).exists() else None + if observation_evidence_version is None: + observation_evidence_version = previous.get("observation_evidence_version", 1) if previous is not None else 2 + frozen = freeze_repeatability_plan( + source_plan, + source_campaign_dir, + source_checkpoint_path, + samples=samples, + max_points_per_cell=max_points_per_cell, + cv_threshold=cv_threshold, + comparison_mode=comparison_mode, + max_attempts_per_sample=max_attempts_per_sample, + source_agreement_threshold=source_agreement_threshold, + observation_evidence_version=observation_evidence_version, + ) + source = Path(source_campaign_dir).expanduser().resolve() + if root.is_relative_to(source) or source.is_relative_to(root): + raise ValueError("repeatability output must be separate from the source campaign") + if root.exists() and any(root.iterdir()): + if not resume: + raise ValueError("repeatability output exists; use resume or a fresh directory") + if json.loads((root / PLAN_FILENAME).read_text()) != frozen: + raise ValueError("repeatability frozen plan or source artifacts changed") + report = json.loads((root / REPORT_FILENAME).read_text()) + if report.get("plan_sha256") != frozen["sha256"]: + raise ValueError("repeatability report does not match its frozen plan") + _validate_saved_samples(source_plan, frozen, report, root) + _summarize(frozen, report) + else: + if resume: + raise ValueError("repeatability resume requires existing validation evidence") + root.mkdir(parents=True, exist_ok=True) + _atomic_json(root / PLAN_FILENAME, frozen) + report = { + "schema_name": "aisimulate_fpm_repeatability", + "schema_version": 2, + "plan_sha256": frozen["sha256"], + "source_plan_sha256": source_plan.sha256, + "output_dir": str(root), + "policy": frozen["policy"], + "samples": { + cell["cell_id"]: [{"sample_index": index + 1, "attempts": []} for index in range(samples)] + for cell in frozen["cells"] + }, + } + _summarize(frozen, report) + _save_report(root, report) + if any(cell["execution"]["status"] == "failed" for cell in frozen["cells"]): + # A stable repeat set cannot qualify a measured wrong-source campaign. + # Retain its assessment without spending another benchmark launch. + _atomic_json(root / REPORT_FILENAME, report) + return report + for selected in frozen["cells"]: + plan = _subset_plan(source_plan, selected, comparison_mode=comparison_mode) + for sample in report["samples"][selected["cell_id"]]: + previous = sample["attempts"][-1] if sample["attempts"] else None + if previous is not None: + previous_dir = root / previous["directory"] + if previous["status"] != "passed" and not previous.get("evidence"): + # Collection can finish before the final report is committed. + # Recover its verified observation without launching a replacement. + try: + _record_sample_evidence(plan, selected, previous_dir, frozen, report, previous) + except (FileNotFoundError, _IncompleteSample): + pass + except ValueError as error: + previous.update(status="failed", failure_kind="validation_failed", error=str(error)) + previous["files"] = _file_manifest(previous_dir) + _summarize(frozen, report) + _save_report(root, report) + if previous["status"] == "passed": + continue + if previous.get("evidence"): + if retry_failed and previous.get("failure_kind") == "cleanup_failed": + try: + _retry_sample_cleanup(plan, selected, previous_dir, previous) + finally: + _summarize(frozen, report) + _save_report(root, report) + if previous["status"] == "passed": + continue + if retry_failed and previous.get("failure_kind") == "postprocessing_failed": + try: + _retry_sample_postprocessing(plan, selected, previous_dir, previous) + finally: + _summarize(frozen, report) + _save_report(root, report) + if previous["status"] == "passed": + continue + return report + if not retry_failed: + return report + if len(sample["attempts"]) >= max_attempts_per_sample: + raise ValueError("repeatability frozen retry budget exhausted; retain attempts and investigate") + # A failed teardown must be resolved before any next launch. + manifest = ( + previous_dir + / "artifacts" + / plan.sha256[:16] + / "cells" + / selected["cell_id"] + / FPM_MANIFEST_FILENAME + ) + if manifest.exists(): + try: + _cell_runner(plan, plan.cells[0], manifest, manifest.parent).cleanup() + finally: + # Teardown failures may append transport diagnostics; + # retain the old snapshot and record the appended files. + previous.setdefault("cleanup_file_history", []).append(previous.get("files", {})) + previous["files"] = _file_manifest(previous_dir) + _atomic_json(root / REPORT_FILENAME, report) + relative = ( + Path("samples") + / selected["cell_id"] + / f"sample-{sample['sample_index']:02d}" + / f"attempt-{len(sample['attempts']) + 1:02d}" + ) + if planner._git_revision() != source_plan.aic_revision: + raise ValueError( + "repeatability launch requires the source collector revision; the current checkout/code " + "is unqualified for this campaign. Use the matching revision or collect a new source campaign. " + "Existing samples remain available for offline assessment." + ) + directory = root / relative + directory.mkdir(parents=True, exist_ok=False) + attempt = {"directory": str(relative), "status": "running"} + sample["attempts"].append(attempt) + _atomic_json(root / REPORT_FILENAME, report) + try: + validating = False + errors = run_collection( + plan, + generator_overrides=generator_overrides, + checkpoint_dir=str(directory / "checkpoint"), + artifact_root=str(directory / "artifacts"), + resume=False, + retry_failed=False, + publish_database=False, + ) + if errors: + attempt.update(status="failed", errors=errors, failure_kind="collection_error") + validating = True + _record_sample_evidence(plan, selected, directory, frozen, report, attempt) + except (KeyboardInterrupt, SystemExit): + attempt.update(status="interrupted", failure_kind="interrupted") + raise + except Exception as error: + kind = ( + "missing_evidence" + if isinstance(error, (FileNotFoundError, _IncompleteSample)) + else "validation_failed" + if validating and isinstance(error, ValueError) + else "collection_error" + ) + attempt.update(status="failed", failure_kind=kind, error=f"{type(error).__name__}: {error}") + finally: + attempt["files"] = _file_manifest(directory) + _summarize(frozen, report) + _save_report(root, report) + if ( + attempt["status"] != "passed" + or report["repeatability"]["status"] == "mismatch" + or report["source_qualification"]["status"] == "failed" + ): + # Keep failures visible. Do not launch the remaining subset + # repetitions into an uninvestigated runtime/cleanup failure. + return report + _summarize(frozen, report) + _save_report(root, report) + return report diff --git a/python/aisimulate/collector/fpm_forward/runner.py b/python/aisimulate/collector/fpm_forward/runner.py index 4c3e1f273..258084f5e 100644 --- a/python/aisimulate/collector/fpm_forward/runner.py +++ b/python/aisimulate/collector/fpm_forward/runner.py @@ -22,6 +22,7 @@ import time import uuid import zlib +from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import contextmanager, suppress from dataclasses import replace @@ -45,8 +46,11 @@ fpm_workload_node_count, ) +from .config import FPM_KV_WARMUP_DEFAULTS, with_kv_warmup_defaults from .native_artifact import COLLECTOR_PROVENANCE_FILENAME, validate_native_collection from .planner import FPMCell, FPMCollectionPlan +from .runtime.fpm_memory_observer import EXECUTION_SUPPORTED_VERSIONS +from .runtime.fpm_memory_observer import SUPPORTED_VERSION as MEMORY_OBSERVER_VERSION from .types import KVWARM_STRATEGIES logger = logging.getLogger(__name__) @@ -86,8 +90,10 @@ # insensitive to both flags (measured 100-108 ms across all four combinations # at 8192 new tokens, M2.7 tp4+EP). _FPM_VLLM_PREFILL_ARGS = ("--no-async-scheduling",) -# Decode always keeps async scheduling on (vLLM's default) -- the -# steady-state second step models a production decode iteration, and +# Ordinary decode leaves scheduling to the runtime default. Memory observation +# explicitly requests synchronous scheduling in both phases to establish one +# compatible memory profile; execution evidence records the effective choice. +# The steady-state second step models a production decode iteration, and # production overlaps scheduler CPU work with the GPU. Against real traffic # at (256, 2.1M KV): async 26.5 ms (1.03x of the 25.8 ms measured), sync # 31.3 ms (1.21x). @@ -119,6 +125,29 @@ REMOTE_EXIT_MARKER = "__FPM_REMOTE_EXIT_CODE__=" REMOTE_FILES_MARKER = "__FPM_REMOTE_FILES__=" REMOTE_WORKDIR = "/tmp/fpm-bench" +_MEMORY_OBSERVER_FILES = ("fpm_memory_observer.py", "fpm_memory_worker.py", "fpm_memory_scheduler.py") + + +def _observe_runtime_memory(plan: FPMCollectionPlan, cell: FPMCell) -> bool: + if getattr(plan, "runtime_instrumentation", None) is not None: + return True + if getattr(plan, "fpm_profile", None) is None: + return False + deployment = plan.deployment_profile(cell) + return ( + deployment is not None + and deployment.resources.memory_source == "pending" + and deployment.backend_version == MEMORY_OBSERVER_VERSION + ) + + +def _observe_runtime_execution(plan: FPMCollectionPlan, cell: FPMCell) -> bool: + if getattr(plan, "runtime_instrumentation", None) is not None: + return False + return ( + getattr(getattr(plan, "capability", None), "aic_database_version", None) in EXECUTION_SUPPORTED_VERSIONS + and not cell.execution_identity[0] + ) def _utc_now() -> str: @@ -1203,30 +1232,48 @@ def _cell_generator_overrides( if not payload[cell.workload_kind]: raise ValueError(f"benchmark-points manifest has no {cell.workload_kind} points for this cell") scheduler_args.extend(["--benchmark-points-file", f"{REMOTE_WORKDIR}/{POINTS_FILENAME}"]) + max_num_tokens = getattr(plan.options, "max_num_batched_tokens", None) + max_num_seqs = getattr(plan.options, "max_num_seqs", None) + compilation_config = dict(getattr(plan.options, "compilation_config", {})) + if getattr(plan, "fpm_profile", None) is not None: + deployment_profile = plan.deployment_profile(cell) + assert deployment_profile is not None + resources = deployment_profile.resources + max_num_tokens = resources.max_num_tokens if max_num_tokens is None else max_num_tokens + max_num_seqs = resources.max_batch_size if max_num_seqs is None else max_num_seqs if cell.workload_kind == "decode" and getattr(plan.options, "max_decode_batch_size", None): - scheduler_args.extend(["--max-num-seqs", str(plan.options.max_decode_batch_size)]) - if cell.workload_kind == "prefill" and not smoke: + max_num_seqs = plan.options.max_decode_batch_size + if cell.workload_kind == "prefill": profile = plan.options.prefill_sampling - compilation_config = { - "cudagraph_capture_sizes": list(profile.cudagraph_capture_sizes), - "max_cudagraph_capture_size": profile.max_cudagraph_capture_size, - } + if not smoke or getattr(plan.options, "max_prefill_isl", None) is not None: + max_num_tokens = profile.max_total_prefill_tokens + if profile.max_batch_size is not None: + max_num_seqs = profile.max_batch_size + # A declared serving role owns one engine configuration. Aggregated probes + # exercise both workloads without changing that configuration. Untagged + # historical plans retain their original prefill-only override. + if cell.workload_kind == "prefill" or getattr(plan.options, "worker_type", None) is not None: + captures = plan.options.prefill_sampling + if captures.cudagraph_capture_sizes is not None and not enforce_eager: + compilation_config.update( + cudagraph_capture_sizes=list(captures.cudagraph_capture_sizes), + max_cudagraph_capture_size=captures.max_cudagraph_capture_size, + ) + if compilation_config and not enforce_eager: + scheduler_args.extend( + ["--compilation-config", json.dumps(compilation_config, sort_keys=True, separators=(",", ":"))] + ) + if cell.workload_kind == "prefill" and not smoke: + if profile.cudagraph_capture_sizes is not None: + scheduler_args.extend(["--prefill-max-new-token-samples", str(profile.max_new_token_samples)]) + # This cap depends only on scheduler bounds, not the runtime's graph + # choices. The new-token axis remains Dynamo-owned in runtime mode. scheduler_args.extend( [ - "--max-num-batched-tokens", - str(profile.max_total_prefill_tokens), - "--prefill-max-new-token-samples", - str(profile.max_new_token_samples), "--prefill-max-kv-read-token-samples", str(profile.max_kv_read_token_samples), ] ) - if not enforce_eager: - scheduler_args.extend( - ["--compilation-config", json.dumps(compilation_config, sort_keys=True, separators=(",", ":"))] - ) - if profile.max_batch_size is not None: - scheduler_args.extend(["--max-num-seqs", str(profile.max_batch_size)]) elif smoke: if cell.workload_kind == "prefill": scheduler_args.extend( @@ -1248,16 +1295,31 @@ def _cell_generator_overrides( "2", ] ) - if getattr(plan, "fpm_profile", None) is not None: - deployment_profile = plan.deployment_profile(cell) - assert deployment_profile is not None - resources = deployment_profile.resources - # Native point generation sees the declared scheduler bounds before - # any cases are queued. Never discard runtime cases to fit a profile. - if "--max-num-batched-tokens" not in scheduler_args: - scheduler_args.extend(["--max-num-batched-tokens", str(resources.max_num_tokens)]) - if "--max-num-seqs" not in scheduler_args: - scheduler_args.extend(["--max-num-seqs", str(resources.max_batch_size)]) + probe_sampling = getattr(plan, "probe_sampling", None) + if probe_sampling is not None: + sampling_flags = ( + { + "--prefill-max-new-token-samples": probe_sampling["max_new_token_samples"], + "--prefill-max-kv-read-token-samples": probe_sampling["max_kv_read_token_samples"], + "--prefix-max-batch-size-samples": probe_sampling["max_batch_size_samples"], + } + if cell.workload_kind == "prefill" + else { + "--decode-max-kv-read-token-samples": probe_sampling["max_kv_read_token_samples"], + "--decode-max-batch-size-samples": probe_sampling["max_batch_size_samples"], + } + ) + for flag, value in sampling_flags.items(): + if flag in scheduler_args: + scheduler_args[scheduler_args.index(flag) + 1] = str(value) + else: + scheduler_args.extend([flag, str(value)]) + # Native point generation sees the runtime bounds before any cases are + # queued. Each flag is emitted once, after resolving the phase's bounds. + if max_num_tokens is not None: + scheduler_args.extend(["--max-num-batched-tokens", str(max_num_tokens)]) + if max_num_seqs is not None: + scheduler_args.extend(["--max-num-seqs", str(max_num_seqs)]) model_args = [] architecture = getattr(getattr(plan, "capability", None), "architecture", None) if architecture == "DeepseekV41ForCausalLM": @@ -1266,9 +1328,9 @@ def _cell_generator_overrides( model_args.extend( ["--language-model-only", "--tokenizer-mode=deepseek_v41", '--engram-config={"cpu_offload":false}'] ) - if cell.workload_kind == "decode" or smoke: - model_args.extend(["--max-num-batched-tokens", str(plan.options.max_prefill_isl)]) - if cell.workload_kind == "prefill" and smoke and plan.options.max_prefill_batch_size: + if (cell.workload_kind == "decode" or smoke) and max_num_tokens is None: + model_args.extend(["--max-num-batched-tokens", str(plan.options.prefill_sampling.max_total_prefill_tokens)]) + if cell.workload_kind == "prefill" and smoke and plan.options.max_prefill_batch_size and max_num_seqs is None: model_args.extend(["--max-num-seqs", str(plan.options.max_prefill_batch_size)]) if architecture == "GlmMoeDsaForCausalLM": # This is the serving path validated by the pinned GLM-5.2 vLLM image. @@ -1279,6 +1341,50 @@ def _cell_generator_overrides( {"name": FPM_ENGINE_BENCHMARK_OUTPUT_ENV, "value": f"{FPM_RESULTS_DIR}/benchmark.json"}, {"name": FPM_RUN_ID_ENV, "value": cell.cell_id}, ] + observe_memory = _observe_runtime_memory(plan, cell) + observe_execution = _observe_runtime_execution(plan, cell) + instrumentation = getattr(plan, "runtime_instrumentation", None) + if instrumentation is not None: + # Pass the selected pin to native model construction for both probes and + # later formal collection. Local checkpoints otherwise resolve revision + # to None; loaded config bytes remain independently verified on import. + model_args.extend(["--revision", plan.runtime_launch["identity"]["model_revision"]]) + if observe_memory: + if cell.workload_kind == "decode": + # Use the prefill collection requirement for both phases when + # deriving one serving memory profile. Existing declared profiles + # retain their established phase-specific scheduling policies. + model_args.append("--no-async-scheduling") + model_args.extend( + [ + "--worker-cls", + "fpm_runtime_instrumentation.ObservedWorker" + if instrumentation is not None + else "fpm_memory_worker.FpmResourceWorker", + "--scheduler-cls", + "fpm_runtime_instrumentation.ObservedInstrumentedScheduler" + if instrumentation is not None + else "fpm_memory_scheduler.FpmResourceInstrumentedScheduler", + ] + ) + elif observe_execution: + model_args.extend( + [ + "--worker-cls", + "fpm_memory_worker.FpmExecutionWorker", + "--scheduler-cls", + "fpm_memory_scheduler.FpmExecutionInstrumentedScheduler", + ] + ) + if instrumentation is not None: + env.extend( + {"name": name, "value": value} + for name, value in { + "AISIMULATE_RUNTIME_CONTEXT": f"{REMOTE_WORKDIR}/runtime-probe-context.json", + "AISIMULATE_RUNTIME_INSTRUMENTATION": f"{REMOTE_WORKDIR}/runtime-instrumentation/manifest.json", + "AISIMULATE_RUNTIME_OBSERVATION_DIR": FPM_RESULTS_DIR, + }.items() + ) if architecture == "DeepseekV41ForCausalLM": from aisimulate_core.sdk.deepseek_v41 import MODEL_REVISION @@ -1316,6 +1422,11 @@ def _cell_generator_overrides( } }, } + gpu_memory_utilization = getattr(plan.options, "gpu_memory_utilization", None) + if gpu_memory_utilization is not None: + # Use the existing backend mapping to --gpu-memory-utilization. The + # Generator preserves explicitly supplied values of this common key. + generated["params"]["agg"]["kv_cache_free_gpu_memory_fraction"] = gpu_memory_utilization policy = cell.backend_policy.generator_overrides merged = _deep_merge(_deep_merge(base, generated), policy) @@ -1363,8 +1474,30 @@ def _cell_generator_overrides( raise ValueError(f"{READINESS_TIMEOUT_ENV} must be an integer from 1 through 3600") resolved_env[READINESS_TIMEOUT_ENV] = {"name": READINESS_TIMEOUT_ENV, "value": str(readiness)} merged.setdefault("K8sConfig", {})["extra_env"] = list(resolved_env.values()) + merged = with_kv_warmup_defaults(merged) policy_args = ((policy.get("params") or {}).get("agg") or {}).get("extra_cli_args") or [] + policy_flags = {str(argument).split("=", 1)[0].split(" ", 1)[0] for argument in policy_args} + if instrumentation is not None and "--revision" in policy_flags: + raise ValueError("backend policy cannot override the runtime probe's selected model revision") + if (observe_memory or observe_execution) and "--worker-cls" in policy_flags: + raise ValueError("backend policy cannot replace the runtime execution/memory observer worker class") + if (observe_memory or observe_execution) and "--scheduler-cls" in policy_flags: + raise ValueError("backend policy cannot replace runtime execution/memory observer scheduler class") + if observe_memory and policy_flags & {"--kv-cache-memory-bytes", "--num-gpu-blocks-override"}: + raise ValueError("backend policy cannot replace runtime memory observer classes or automatic cache sizing") + if observe_memory and any( + str(argument).split("=", 1)[0].split(" ", 1)[0] == "--async-scheduling" for argument in policy_args + ): + raise ValueError("runtime memory collection requires synchronous scheduling in both phases") + if gpu_memory_utilization is not None and any( + str(argument).split("=", 1)[0].split(" ", 1)[0] == "--gpu-memory-utilization" for argument in policy_args + ): + raise ValueError("backend policy cannot override --fpm-gpu-memory-utilization with extra CLI arguments") + if cell.workload_kind == "prefill" and profile.cudagraph_policy == "runtime": + graph_options = ("--compilation-config", "--cudagraph-capture-sizes", "--max-cudagraph-capture-size") + if any(str(argument).split("=", 1)[0] in graph_options for argument in policy_args): + raise ValueError("runtime prefill CUDA-graph policy cannot use backend capture overrides") if any(str(arg).split("=", 1)[0] in {"--enforce-eager", "--no-enforce-eager"} for arg in policy_args): raise ValueError("eager execution must be supplied through --fpm-enforce-eager") if any(str(arg).split("=", 1)[0] == "--benchmark-points-file" for arg in policy_args): @@ -1400,6 +1533,11 @@ def _cell_generator_overrides( resolved_args.extend(["--benchmark-timeout", str(DEFAULT_BENCHMARK_TIMEOUT_SECONDS)]) resolved_args.extend(scheduler_args) merged_agg = merged.setdefault("params", {}).setdefault("agg", {}) + if ( + gpu_memory_utilization is not None + and merged_agg.get("kv_cache_free_gpu_memory_fraction") != gpu_memory_utilization + ): + raise ValueError("backend policy cannot change --fpm-gpu-memory-utilization") merged_agg.update({"extra_cli_args": resolved_args}) return merged @@ -1415,7 +1553,7 @@ def _configured_sampling_metadata( cell: FPMCell, *, smoke: bool, -) -> dict[str, int | str]: +) -> dict[str, int | str | None]: canonical = _frozen_points(plan) if canonical is not None: if smoke: @@ -1430,13 +1568,23 @@ def _configured_sampling_metadata( # arguments that conflict with this resolved protocol. return {"decode_prefix_caching": _decode_prefix_caching_mode(cell)} if smoke: - return {"prefill_max_new_token_samples": 2} + payload: dict[str, int | str | None] = {"prefill_max_new_token_samples": 2} + if plan.options.prefill_sampling.cudagraph_policy == "runtime": + payload["prefill_cudagraph_policy"] = "runtime" + return payload profile = plan.options.prefill_sampling - return { - "prefill_cudagraph_capture_size_count": len(profile.cudagraph_capture_sizes), - "prefill_requested_new_token_axis_count": len(profile.new_token_axis_points), + payload = { + "prefill_cudagraph_capture_size_count": len(profile.cudagraph_capture_sizes) + if profile.cudagraph_capture_sizes is not None + else None, + "prefill_requested_new_token_axis_count": len(profile.new_token_axis_points) + if profile.new_token_axis_points is not None + else None, "prefill_max_new_token_samples": profile.max_new_token_samples, } + if profile.cudagraph_policy == "runtime": + payload["prefill_cudagraph_policy"] = "runtime" + return payload def _write_runtime_environment(cell_dir: Path, overrides: dict[str, Any]) -> None: @@ -1444,11 +1592,15 @@ def _write_runtime_environment(cell_dir: Path, overrides: dict[str, Any]) -> Non # this Collector-owned file because Slurm does not start a Kubernetes Pod # with extra_env, and run.sh's exports happen after the preflight process. names = { + *FPM_KV_WARMUP_DEFAULTS, READINESS_TIMEOUT_ENV, "PYTHONPATH", "DYN_FPM_DSV41_REAL_KV", "DYN_FPM_INPUT_TEXT", "DYN_FPM_TOKENIZER_REVISION", + "AISIMULATE_RUNTIME_CONTEXT", + "AISIMULATE_RUNTIME_INSTRUMENTATION", + "AISIMULATE_RUNTIME_OBSERVATION_DIR", } lines = ["# Generated Collector startup environment; engine settings remain in run.sh."] for item in overrides["K8sConfig"]["extra_env"]: @@ -1544,7 +1696,7 @@ def _runtime_collection_summary( *, expected_plan_sha256: str | None = None, expected_attempt_id: str | None = None, -) -> dict[str, int]: +) -> dict[str, Any]: """Return auditable unique-axis counts from a validated native grid.""" collection = validate_native_collection( @@ -1558,7 +1710,12 @@ def _runtime_collection_summary( "measured_point_count": len(points), "measured_batch_size_axis_count": len({int(point["batch_size"]) for point in points}), "measured_kv_read_axis_count": len({int(point["total_kv_read_tokens"]) for point in points}), + "native_kv_seed_regime_counts": dict( + Counter(measurement.kv_seed_regime or "unreported" for measurement in collection.points) + ), } + if collection.kvwarm_meta is not None: + summary["native_kvwarm"] = collection.kvwarm_meta if cell.workload_kind == "prefill": summary["measured_new_token_axis_count"] = len({int(point["total_prefill_tokens"]) for point in points}) return summary @@ -1618,6 +1775,9 @@ def _cell_runner(plan: FPMCollectionPlan, cell: FPMCell, manifest: Path, cell_di image=plan.options.slurm_container_image, mounts=plan.options.slurm_container_mounts, total_gpus=cell.topology.total_gpus, + cpus_per_task=plan.options.slurm_cpus_per_task, + cpu_bind=plan.options.slurm_cpu_bind, + attention_tp=cell.topology.tp, ) if executor != "kubernetes": raise ValueError(f"unknown FPM executor {executor!r}") @@ -1819,6 +1979,7 @@ def run_collection( cell_limit: int | None = None, database_root: str | None = None, publish_partial: bool = False, + publish_database: bool = True, ) -> list[dict[str, object]]: """Render and run every cell, always tearing down owned resources.""" @@ -1834,6 +1995,7 @@ def run_collection( cell_limit=cell_limit, database_root=database_root, publish_partial=publish_partial, + publish_database=publish_database, ) @@ -1849,6 +2011,7 @@ def _run_collection_impl( cell_limit: int | None = None, database_root: str | None = None, publish_partial: bool = False, + publish_database: bool = True, ) -> list[dict[str, object]]: if _frozen_points(plan) is not None and smoke: raise ValueError("--fpm-benchmark-points-file cannot be combined with --smoke") @@ -1858,11 +2021,55 @@ def _run_collection_impl( root /= "smoke" root.mkdir(parents=True, exist_ok=True) _atomic_json(root / "collection-plan.json", plan.to_dict()) + _atomic_json(root / "generator-overrides.json", with_kv_warmup_defaults(generator_overrides)) checkpoint_name = "fpm_forward_smoke.json" if smoke else "fpm_forward.json" checkpoint_path = Path(checkpoint_dir).expanduser().resolve() / checkpoint_name checkpoint = _load_checkpoint(checkpoint_path, plan, resume) + instrumentation = getattr(plan, "runtime_instrumentation", None) + observation_attempt_id = None + if instrumentation is not None: + from .runtime_probe import prepare_collection_observations + + observation_attempt_id = checkpoint.setdefault("runtime_observation_attempt_id", uuid.uuid4().hex) + instrumentation = prepare_collection_observations(plan, root, observation_attempt_id) + _atomic_json(checkpoint_path, checkpoint) errors: list[dict[str, object]] = [] run_attempts: list[dict[str, Any]] = [] + observation_archive_failed: set[str] = set() + + def archive_observations(cell, record) -> bool: + if instrumentation is None: + return True + from .runtime_probe import record_collection_observations + + try: + observation_index = record_collection_observations( + plan, + root, + observation_attempt_id, + cell, + record["attempt_id"], + root / "cells" / cell.cell_id, + record["status"], + ) + checkpoint["runtime_observations"] = str(observation_index) + record.pop("observation_error", None) + return True + except Exception as error: + record["status"] = "failed" + record["observation_error"] = str(error) + observation_archive_failed.add(cell.cell_id) + errors.append( + { + "module": "fpm_forward", + "cell_id": cell.cell_id, + "error_type": type(error).__name__, + "error_message": str(error), + "classification": "runtime_observation_archive_failed", + } + ) + return False + runtime_exec = Path(__file__).resolve().parent / "runtime" / "fpm_exec.sh" runtime_preflight = Path(__file__).resolve().parent / "runtime" / "preflight.py" target_cells = plan.cells[: (cell_limit or (1 if smoke else len(plan.cells)))] @@ -1870,7 +2077,8 @@ def _run_collection_impl( formal_database_terminal = False database_entry = checkpoint.get("database") if ( - resume + publish_database + and resume and not smoke and isinstance(database_entry, dict) and database_entry.get("status") == "passed" @@ -1910,8 +2118,11 @@ def _run_collection_impl( recovered = _recover_completed_attempt(plan, cell, root, entry) if recovered is None: continue + archived = archive_observations(cell, recovered) checkpoint["cells"][cell.cell_id] = recovered checkpoint_changed = True + if not archived: + continue logger.info( "Recovered completed FPM cell %s from strictly validated artifacts for attempt %s", cell.cell_id, @@ -1925,6 +2136,12 @@ def _run_collection_impl( entry = checkpoint["cells"].get(cell.cell_id) if not isinstance(entry, dict) or entry.get("status") != "passed": continue + if instrumentation is not None: + # Native recovery and older checkpoints can precede observation + # archival. Repair the index before this passed cell is skipped. + checkpoint_changed = True + if not archive_observations(cell, entry): + continue # A validated terminal database no longer depends on retained raw # receipts. Unpublished cells still need them before publication. if not formal_database_terminal and _frozen_points(plan) is not None: @@ -1968,11 +2185,21 @@ def _run_collection_impl( for cell in target_cells: previous = checkpoint["cells"].get(cell.cell_id, {}) + if cell.cell_id in observation_archive_failed: + continue if resume and previous.get("status") == "passed": continue if resume and previous.get("status") in {"failed", "cleanup_failed"} and not retry_failed: continue + if getattr(plan.options, "executor", "kubernetes") == "slurm" and plan.options.slurm_cpus_per_task is None: + # Do not discard retained raw/log evidence on an old campaign in + # an attempt to launch with unrecorded CPU defaults. + raise ValueError( + "saved Slurm campaign has no frozen CPU policy; preserve its artifacts for CPU-only recovery " + "and use a fresh campaign and smoke for new workers" + ) + cell_dir = root / "cells" / cell.cell_id if getattr(plan.options, "executor", "kubernetes") == "slurm": abandoned_manifest = cell_dir / FPM_MANIFEST_FILENAME @@ -2002,6 +2229,10 @@ def _run_collection_impl( if cell_dir.exists() and not resume: shutil.rmtree(cell_dir) cell_dir.mkdir(parents=True, exist_ok=True) + # Preserve abruptly interrupted attempts before replacing raw files. + if instrumentation is not None and previous.get("attempt_id") and not archive_observations(cell, previous): + _atomic_json(checkpoint_path, checkpoint) + continue for stale_dir in (cell_dir / "raw", cell_dir / "logs"): if stale_dir.exists(): shutil.rmtree(stale_dir) @@ -2060,6 +2291,16 @@ def _run_collection_impl( pods = resource.wait_ready(_expected_nodes(manifest)) phase_marks["schedule_s"] = round(time.monotonic() - mark, 3) mark = time.monotonic() + instrumentation_files = [] + if instrumentation is not None: + from .runtime_probe import launch_context, stage_runtime_instrumentation + + context = launch_context( + plan, cell, configuration=plan.runtime_configuration, attempt_id=observation_attempt_id + ) + context["collector_attempt_id"] = attempt_id + context["cell_id"] = cell.cell_id + instrumentation_files = stage_runtime_instrumentation(instrumentation, cell_dir, context) resource.stage( pods, [ @@ -2068,6 +2309,19 @@ def _run_collection_impl( runtime_env, runtime_exec, runtime_preflight, + *[ + runtime_exec.parent / filename + for filename in _MEMORY_OBSERVER_FILES + if ( + instrumentation is None + and (_observe_runtime_memory(plan, cell) or _observe_runtime_execution(plan, cell)) + ) + or ( + getattr(plan.options, "executor", "kubernetes") == "slurm" + and filename == "fpm_memory_observer.py" + ) + ], + *instrumentation_files, *_stage_points_file(plan, cell_dir), *( [ @@ -2166,6 +2420,8 @@ def _run_collection_impl( record["collector_phase_seconds"] = dict(phase_marks) record["completed_at"] = _utc_now() record["duration_seconds"] = round(time.monotonic() - cell_started, 3) + if instrumentation is not None: + archive_observations(cell, record) run_attempts.append(_run_manifest_attempt(cell, record)) _atomic_json(checkpoint_path, checkpoint) # R16 §3 run manifest: the machine-readable timing map for the speed-up @@ -2191,6 +2447,17 @@ def _run_collection_impl( if formal_database_terminal: return errors all_passed = all(checkpoint["cells"].get(cell.cell_id, {}).get("status") == "passed" for cell in target_cells) + if not publish_database: + if not all_passed and not errors: + errors.append( + { + "module": "fpm_forward", + "error_type": "IncompleteCampaign", + "error_message": "validation collection has incomplete cells", + "classification": "campaign_incomplete", + } + ) + return errors # Formal publication eligibility must agree with completion: a deliberate # partial run (cell_limit below the frozen plan) can pass every targeted # cell yet cannot publish, and deserves the honest campaign_incomplete diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh b/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh index 366a609c3..bc3737c5e 100644 --- a/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh @@ -20,12 +20,40 @@ source "${workdir}/fpm_env.sh" # This is staged by the Collector from the same resolved deployment settings # used by run.sh. It carries startup configuration across both transports. source "${workdir}/collector-runtime-env.sh" +# This runs inside the SAME Slurm step/container as the engine, before model +# initialization. Separate provenance steps cannot establish the engine's mask. +if [[ -n "${FPM_SLURM_CPUS_PER_TASK:-}" ]]; then + python3 - "${workdir}/fpm_memory_observer.py" <<'PY' +import importlib.util +import os +from pathlib import Path +import sys + +spec = importlib.util.spec_from_file_location("fpm_launcher_cpu_observer", sys.argv[1]) +observer = importlib.util.module_from_spec(spec) +spec.loader.exec_module(observer) +requested = int(os.environ["FPM_SLURM_CPUS_PER_TASK"]) +report = observer.observe_cpu( + "launcher", requested_cpus_per_task=requested, cpu_bind=os.environ["FPM_SLURM_CPU_BIND"], + local_gpu_count=int(os.environ["FPM_LOCAL_GPU_COUNT"]), directory=Path("/results"), +) +if report["status"] != "observed" or len(report["main_thread_allowed_cpus"]) < requested: + raise SystemExit("Slurm engine CPU affinity is unavailable or smaller than requested; inspect fpm-cpu-launcher.json") +PY +fi if [[ ! "${FPM_READINESS_TIMEOUT_SECONDS:-}" =~ ^[1-9][0-9]*$ ]] || (( ${#FPM_READINESS_TIMEOUT_SECONDS} > 4 || FPM_READINESS_TIMEOUT_SECONDS > 3600 )); then echo "FPM_READINESS_TIMEOUT_SECONDS must be an integer from 1 through 3600" >&2 exit 2 fi +if [[ -f "${workdir}/fpm_memory_worker.py" ]]; then + # The wrappers are staged alongside this script and delegate all execution + # to the image's installed vLLM/Dynamo classes. Supported execution or + # memory observation stages them into the new pod. Avoid an empty path entry. + export PYTHONPATH="${workdir}${PYTHONPATH:+:${PYTHONPATH}}" +fi + etcd_endpoint="http://${FPM_MASTER_ADDR}:2379" engine_pid="" etcd_pid="" diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_observer.py b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_observer.py new file mode 100644 index 000000000..c1b2df9f8 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_observer.py @@ -0,0 +1,641 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Observe initialized vLLM memory without changing execution or profiling. + +Original collector adapters against vLLM 0.27.0, commit +4bdc8a788d2e2ce9165d552b3d4d8b72604626bf: vllm/v1/worker/gpu_worker.py +(determine_available_memory, initialize_from_config, compile_or_warm_up_model), +vllm/v1/worker/gpu_model_runner.py (initialize_kv_cache_tensors), and +vllm/v1/core/{kv_cache_manager,block_pool}.py. Cache specs are read, not rebuilt. +https://github.com/vllm-project/vllm/tree/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1 + +Unsupported observation APIs write unresolved evidence. Exceptions from the +runtime itself are never caught by these observation helpers or wrappers. +""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import logging +import os +import socket +import threading +from enum import Enum +from pathlib import Path +from typing import Any + +SCHEMA_NAME = "aisimulate_fpm_runtime_memory" +SCHEMA_VERSION = 1 +SUPPORTED_VERSION = "0.27.0" +EXECUTION_SUPPORTED_VERSIONS = ("0.27.0", "0.28.0") +RESULTS_DIR = Path("/results") +_SPEC_MODULE = "vllm.v1.kv_cache_interface." +_CONV_CLASS = "vllm.models.inkling.nvidia.sconv_swa_attn.InklingConvState" +_CONFIG_FIELDS = { + "model_config": ("model", "revision", "dtype", "quantization", "max_model_len", "enforce_eager"), + "cache_config": ( + "cache_dtype", + "gpu_memory_utilization", + "kv_cache_memory_bytes", + "enable_prefix_caching", + "num_gpu_blocks_override", + ), + "scheduler_config": ("max_num_batched_tokens", "max_num_seqs", "async_scheduling"), + "parallel_config": ( + "tensor_parallel_size", + "pipeline_parallel_size", + "data_parallel_size", + "data_parallel_size_local", + "data_parallel_external_lb", + "enable_expert_parallel", + "enable_eplb", + "decode_context_parallel_size", + "prefill_context_parallel_size", + ), + "compilation_config": ("mode", "cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"), + "kernel_config": ("moe_backend",), +} + + +def cpu_affinity_snapshot(*, task_directory: Path = Path("/proc/self/task")) -> dict[str, Any]: + """Observe Linux thread affinity once, without binding or running a model. + + Linux affinity is per thread. Querying PID explicitly observes the main + thread even when a runtime hook executes on another Python/native thread. + The task scan is not atomic; changing or unreadable threads stay partial. + CPU IDs are logical CPUs, not a physical-core or NUMA placement claim. + """ + pid = os.getpid() + payload: dict[str, Any] = { + "hostname": socket.gethostname(), + "pid": pid, + "observer_thread_id": threading.get_native_id(), + "main_thread_allowed_cpus": [], + "threads": [], + "thread_ids_before": [], + "thread_ids_after": [], + "thread_errors": [], + "errors": [], + "status": "unavailable", + } + try: + payload["main_thread_allowed_cpus"] = sorted(os.sched_getaffinity(pid)) + if not payload["main_thread_allowed_cpus"]: + raise ValueError("main thread has an empty CPU affinity mask") + except (AttributeError, OSError, ValueError) as error: + payload["errors"].append(f"main thread affinity: {type(error).__name__}: {error}") + return payload + try: + before = payload["thread_ids_before"] = sorted( + int(path.name) for path in task_directory.iterdir() if path.name.isdecimal() + ) + for tid in before: + try: + cpus = sorted(os.sched_getaffinity(tid)) + if not cpus: + raise ValueError("thread has an empty CPU affinity mask") + payload["threads"].append({"tid": tid, "allowed_cpus": cpus}) + except (OSError, ValueError) as error: + payload["thread_errors"].append({"tid": tid, "error": f"{type(error).__name__}: {error}"}) + after = payload["thread_ids_after"] = sorted( + int(path.name) for path in task_directory.iterdir() if path.name.isdecimal() + ) + if before != after: + payload["errors"].append("thread list changed during observation") + if payload["observer_thread_id"] not in before: + payload["errors"].append("observation thread is missing from the task snapshot") + main = [thread for thread in payload["threads"] if thread["tid"] == pid] + if not main or main[0]["allowed_cpus"] != payload["main_thread_allowed_cpus"]: + payload["errors"].append("main thread was missing or its affinity changed during observation") + except OSError as error: + payload["errors"].append(f"thread enumeration: {type(error).__name__}: {error}") + payload["status"] = "partial" if payload["errors"] or payload["thread_errors"] else "observed" + return payload + + +def observe_cpu( + kind: str, + *, + dp_rank: int | None = None, + tp_rank: int | None = None, + pp_rank: int | None = None, + requested_cpus_per_task: int | None = None, + cpu_bind: str | None = None, + local_gpu_count: int | None = None, + directory: Path | None = None, +) -> dict[str, Any]: + """Save independent CPU evidence; unavailable observation never changes memory. + + The launcher uses the same container and srun step as the engine, before + engine startup. Its caller must inspect this saved result before proceeding. + This helper records affinity; it does not set it or claim it remains fixed. + """ + measurements = { + "launcher": "before_engine_start", + "worker": "after_warmup", + "scheduler": "scheduler_initialized", + } + if kind not in measurements: + raise ValueError("CPU observation requires launcher, worker or scheduler role") + ranks = {"dp_rank": dp_rank, "tp_rank": tp_rank, "pp_rank": pp_rank} + required = set(ranks) if kind == "worker" else {"dp_rank"} if kind == "scheduler" else set() + if any( + (type(value) is not int or value < 0) if key in required else value is not None for key, value in ranks.items() + ): + raise ValueError("CPU observation has invalid role/rank coordinates") + if directory is None: + directory, binding = _execution_destination() + else: + source = directory / "collector-provenance.json" + raw = source.read_bytes() + provenance = json.loads(raw) + if not isinstance(provenance, dict): + raise ValueError("CPU observation collector provenance must be an object") + binding = { + "collector_provenance": provenance, + "execution_provenance": provenance, + "provenance_source": { + "path": str(source), + "sha256": hashlib.sha256(raw).hexdigest(), + "size_bytes": len(raw), + }, + } + payload = { + "schema_name": "aisimulate_fpm_cpu_affinity", + "schema_version": 1, + "kind": kind, + "measurement": measurements[kind], + **ranks, + **binding, + **cpu_affinity_snapshot(), + } + if kind == "launcher": + payload.update( + requested_cpus_per_task=requested_cpus_per_task, cpu_bind=cpu_bind, local_gpu_count=local_gpu_count + ) + suffix = ( + f"-dp{dp_rank}-tp{tp_rank}-pp{pp_rank}" if kind == "worker" else f"-dp{dp_rank}" if kind == "scheduler" else "" + ) + path = directory / f"fpm-cpu-{kind}{suffix}.json" + try: + _write_execution_file(path, payload) + except FileExistsError: + raise RuntimeError(f"refusing duplicate runtime CPU observation: {path}") from None + return payload + + +def _class_name(value: object) -> str: + return f"{type(value).__module__}.{type(value).__name__}" + + +def _json_value(value: Any) -> Any: + if isinstance(value, Enum): + return value.name + if value is None or type(value) in (str, int, float, bool): + return value + if isinstance(value, (list, tuple)): + return [_json_value(item) for item in value] + return str(value) + + +def compilation_config(config: Any) -> dict[str, Any]: + """Snapshot the graph settings before the worker resolves backend support.""" + compilation = config.compilation_config + return {key: _json_value(getattr(compilation, key)) for key in _CONFIG_FIELDS["compilation_config"]} + + +def _write_execution_file(path: Path, payload: dict[str, Any]) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + with temporary.open("x", encoding="utf-8") as handle: + handle.write(json.dumps(payload, sort_keys=True, indent=2) + "\n") + handle.flush() + os.fsync(handle.fileno()) + # Publish atomically without replacing an earlier observation, including + # when two worker processes accidentally report the same rank. + os.link(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _execution_destination() -> tuple[Path, dict[str, Any]]: + output_value = os.environ.get("FPM_EXECUTION_OUTPUT_DIR") + provenance_value = os.environ.get("FPM_EXECUTION_PROVENANCE_FILE") + custom = output_value is not None or provenance_value is not None + if custom: + if not output_value or not provenance_value: + raise ValueError("FPM_EXECUTION_OUTPUT_DIR and FPM_EXECUTION_PROVENANCE_FILE must be supplied together") + for name, value in (("output directory", output_value), ("provenance file", provenance_value)): + if not Path(value).is_absolute() or any(char in value for char in ("\n", "\r", "\x00")): + raise ValueError(f"execution observation {name} must be an absolute path without control characters") + directory, source = Path(output_value), Path(provenance_value) + else: + directory, source = RESULTS_DIR, RESULTS_DIR / "collector-provenance.json" + raw = source.read_bytes() + provenance = json.loads(raw) + if not isinstance(provenance, dict): + raise ValueError("execution provenance must be a JSON object") + binding = { + "execution_provenance": provenance, + "provenance_source": {"path": str(source), "sha256": hashlib.sha256(raw).hexdigest(), "size_bytes": len(raw)}, + } + if not custom: + binding["collector_provenance"] = provenance + return directory, binding + if ( + provenance.get("purpose") != "matched_serving_accuracy" + or not isinstance(provenance.get("run_id"), str) + or not provenance["run_id"].strip() + or "collector_attempt_id" in provenance + or provenance.get("schema_name") == "aic_fpm_collector_provenance" + ): + raise ValueError("custom execution provenance requires a serving run_id and purpose=matched_serving_accuracy") + directory.mkdir(parents=True, exist_ok=True) + marker = directory / ".execution-run.json" + if (directory / "collector-provenance.json").exists() or ( + not marker.exists() and any(directory.glob("fpm-execution-worker-*.json")) + ): + raise ValueError("serving execution output requires a fresh directory, separate from collector attempts") + try: + _write_execution_file(marker, binding) + except FileExistsError: + if json.loads(marker.read_bytes()) != binding: + raise ValueError("execution observation directory already belongs to a different serving run") from None + return directory, binding + + +def observe_execution(worker: Any, *, dp_rank: int, tp_rank: int, pp_rank: int) -> None: + """Record selected backend objects after initialization, outside timed forwards. + + This original adapter reads AttentionGroup.backend/layer_names populated by + GPUModelRunner.initialize_attn_backend, including runtime backend wrappers: + https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu_model_runner.py#L7020 + https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/utils.py#L217 + The same read-only interfaces were verified for vLLM 0.28.0 at: + https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu_model_runner.py#L7155 + https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/utils.py#L241 + It does not infer a backend from requested flags or trace per-point dispatch. + + Ordinary serving can use FpmExecutionWorker with absolute + FPM_EXECUTION_OUTPUT_DIR and FPM_EXECUTION_PROVENANCE_FILE paths. Both are + required together; the provenance identifies a fresh matched-serving run. + Collector launches retain their default /results output and provenance. + """ + directory, binding = _execution_destination() + payload = { + "schema_name": "aisimulate_fpm_runtime_execution", + "schema_version": 1, + **binding, + "backend_version": importlib.metadata.version("vllm"), + "dp_rank": dp_rank, + "tp_rank": tp_rank, + "pp_rank": pp_rank, + "per_point_dispatch": "unreported", + } + try: + if payload["backend_version"] not in EXECUTION_SUPPORTED_VERSIONS: + raise ValueError(f"execution observation requires audited vLLM {EXECUTION_SUPPORTED_VERSIONS}") + groups = [] + for cache_groups in worker.model_runner.attn_groups: + for group in cache_groups: + names = list(group.layer_names) + if not names or any(not isinstance(name, str) or not name for name in names): + raise ValueError("runtime attention group has no explicit layer names") + backend = group.backend + if not isinstance(backend, type): + raise ValueError("runtime attention group backend is not a class") + groups.append( + { + "kv_cache_group_id": group.kv_cache_group_id, + "backend_class": f"{backend.__module__}.{backend.__qualname__}", + "layer_names": names, + } + ) + if not groups: + raise ValueError("runtime reported no initialized attention groups") + payload.update( + attention_groups=groups, + graph_config=compilation_config(worker.vllm_config), + resolved_config=execution_config(worker.vllm_config), + status="observed", + ) + except Exception as error: + payload.update(status="unresolved", error=f"{type(error).__name__}: {error}") + logging.getLogger(__name__).warning("FPM execution evidence is incomplete: %s", payload["error"]) + path = directory / f"fpm-execution-worker-dp{dp_rank}-tp{tp_rank}-pp{pp_rank}.json" + if path.exists(): + raise RuntimeError(f"refusing duplicate runtime execution observation: {path}") + try: + _write_execution_file(path, payload) + except FileExistsError: + raise RuntimeError(f"refusing duplicate runtime execution observation: {path}") from None + + +def execution_config(config: Any) -> dict[str, Any]: + """Read execution identity without requiring the memory-finalization APIs.""" + result = { + section: { + key: _json_value(getattr(getattr(config, section), key)) + for key in fields + if hasattr(getattr(config, section), key) + } + for section, fields in _CONFIG_FIELDS.items() + if getattr(config, section, None) is not None + } + hf_config = getattr(config.model_config, "hf_config", None) + if hf_config is not None and getattr(hf_config, "_commit_hash", None) is not None: + result["model_config"]["loaded_config_commit_hash"] = hf_config._commit_hash + quant = config.quant_config + result["quantization_config"] = ( + { + "type": _class_name(quant), + **{ + key: _json_value(getattr(quant, key)) + for key in ("quant_method", "activation_scheme", "weight_block_size") + if hasattr(quant, key) + }, + } + if quant is not None + else None + ) + return result + + +def resolved_config(config: Any) -> dict[str, Any]: + result = execution_config(config) + # vLLM 0.27.0 config/offload.py owns weight offload; cache_config does not. + offload = config.offload_config + result["offload_config"] = { + "offload_backend": offload.offload_backend, + "uva": {"cpu_offload_gb": offload.uva.cpu_offload_gb}, + "prefetch": { + key: getattr(offload.prefetch, key) + for key in ("offload_group_size", "offload_num_in_group", "offload_prefetch_step") + }, + } + return result + + +def _offloader_type() -> str: + # vLLM 0.27.0 model_executor/offloader/base.py:get_offloader is read-only. + from vllm.model_executor.offloader import get_offloader + + return _class_name(get_offloader()) + + +def _positive(value: Any, label: str) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{label} must be a positive integer, got {value!r}") + return value + + +def cache_groups(cache_config: Any) -> list[dict[str, Any]]: + groups = [] + all_names: set[str] = set() + for group in cache_config.kv_cache_groups: + spec = group.kv_cache_spec + spec_type = _class_name(spec) + if spec_type not in {_SPEC_MODULE + "FullAttentionSpec", _SPEC_MODULE + "SlidingWindowSpec"}: + raise ValueError(f"unsupported runtime cache spec: {spec_type}") + if getattr(group, "is_eagle_group", False): + raise ValueError("speculative cache groups are unsupported") + names = list(group.layer_names) + if not names or any(not isinstance(name, str) or not name for name in names): + raise ValueError("runtime cache groups require explicit layer names") + if len(names) != len(set(names)) or all_names.intersection(names): + raise ValueError("runtime cache group layer aliases are unsupported") + all_names.update(names) + window = getattr(spec, "sliding_window", None) + if spec_type.endswith(".SlidingWindowSpec"): + _positive(window, "sliding_window") + elif window is not None or getattr(spec, "attention_chunk_size", None) is not None: + raise ValueError("full-attention cache with alternate retention is unsupported") + groups.append( + { + "layer_names": names, + "spec_type": spec_type, + "block_size_tokens": _positive(spec.block_size, "cache block size"), + "spec_page_size_bytes": _positive(spec.page_size_bytes, "cache spec page size"), + "sliding_window": window, + "dtype": str(spec.dtype), + } + ) + if not groups: + raise ValueError("runtime emitted no cache groups") + return groups + + +def _storage(tensor: Any) -> tuple[tuple[str, int], int]: + storage = tensor.untyped_storage() + device = str(storage.device) + if not device.startswith("cuda:"): + raise ValueError(f"cache storage is not on a CUDA device: {device}") + return (device, _positive(storage.data_ptr(), "storage pointer")), _positive(storage.nbytes(), "storage bytes") + + +def _tensor_storages(tensors: Any) -> dict[tuple[str, int], int]: + result: dict[tuple[str, int], int] = {} + + def visit(value: Any) -> None: + if isinstance(value, dict): + for item in value.values(): + visit(item) + elif isinstance(value, (list, tuple)): + for item in value: + visit(item) + else: + key, size = _storage(value) + if key in result and result[key] != size: + raise ValueError("aliased cache storage reports inconsistent sizes") + result[key] = size + + visit(tensors) + if not result: + raise ValueError("runtime emitted no physical cache storage") + if len({device for device, _pointer in result}) != 1: + raise ValueError("rank cache spans multiple devices") + return result + + +def worker_memory(worker: Any) -> dict[str, Any]: + if not getattr(worker, "_fpm_cache_initialized", False): + raise ValueError("worker cache initialization was not observed") + runner = worker.model_runner + if getattr(runner, "shared_kv_cache_layers", None): + raise ValueError("cross-layer KV sharing is not yet supported by memory finalization") + config = runner.kv_cache_config + groups = cache_groups(config) + storages = _tensor_storages(runner.kv_caches) + context = worker.vllm_config.compilation_config.static_forward_context + layers = {name: context[name] for group in groups for name in group["layer_names"]} + layer_storages = {name: _storage(layer.kv_cache) for name, layer in layers.items()} + if dict(layer_storages.values()) != storages: + raise ValueError("cache group layers do not account for every model-runner storage") + + # Every observed alias must be declared by the runtime allocation plan. + # Packed tensors explicitly share one backing allocation through block_stride. + expected: dict[tuple[str, int], int] = {} + covered: set[str] = set() + packed_key = None + allocations = [] + for tensor in config.kv_cache_tensors: + names = list(tensor.shared_by) + if not names or covered.intersection(names) or any(name not in layers for name in names): + raise ValueError("runtime tensor allocation has missing or duplicate cache layers") + keys = {layer_storages[name][0] for name in names} + if len(keys) != 1: + raise ValueError("runtime-declared shared tensor uses different physical storages") + key = next(iter(keys)) + size = _positive(tensor.size, "configured cache tensor bytes") + stride = getattr(tensor, "block_stride", 0) + if type(stride) is not int or stride < 0: + raise ValueError("invalid packed cache block stride") + if stride: + if packed_key is None: + if key in expected: + raise ValueError("packed cache aliases an unpacked allocation") + packed_key = key + elif key != packed_key: + raise ValueError("packed cache tensors use multiple backing allocations") + elif key in expected: + raise ValueError("undeclared cache tensor alias") + if key in expected and expected[key] != size: + raise ValueError("packed allocation sizes disagree") + expected[key] = size + covered.update(names) + allocations.append( + {"size": size, "shared_by": names, "block_stride": stride, "offset": getattr(tensor, "offset", 0)} + ) + if covered != set(layers) or expected != storages: + raise ValueError("runtime allocation plan disagrees with physical cache storage") + for group in groups: + classes = {name: _class_name(layers[name]) for name in group["layer_names"]} + conv = {value == _CONV_CLASS for value in classes.values()} + if len(conv) != 1: + raise ValueError("mixed convolution and attention cache group is unsupported") + group["kind"] = "convolution" if conv == {True} else "attention" + group["layer_classes"] = classes + if group["kind"] == "convolution" and group["sliding_window"] != 4: + raise ValueError("unsupported Inkling convolution retention") + diagnostics = { + name: _json_value(getattr(worker, name)) + for name in ( + "available_kv_cache_memory_bytes", + "peak_activation_memory", + "non_torch_memory", + "total_consumed", + "requested_memory", + "cudagraph_memory_estimate", + ) + if hasattr(worker, name) + } + return { + "num_blocks": _positive(config.num_blocks, "cache block count"), + "available_cache_bytes": _positive(worker._fpm_available_cache_bytes, "profiled cache budget"), + "allocated_cache_bytes": sum(storages.values()), + "storages": [ + {"device": key[0], "pointer": key[1], "size_bytes": size} for key, size in sorted(storages.items()) + ], + "tensor_allocations": allocations, + "groups": groups, + "diagnostics": diagnostics, + } + + +def scheduler_memory(scheduler: Any, config: Any) -> dict[str, Any]: + manager = scheduler.kv_cache_manager + pool = manager.block_pool + managers = manager.coordinator.single_type_managers + if not managers or any(item.block_pool is not pool for item in managers): + raise ValueError("runtime cache groups must share one block pool") + num_blocks = _positive(config.num_blocks, "scheduler cache block count") + if pool.num_gpu_blocks != num_blocks: + raise ValueError("scheduler and configured block pool sizes disagree") + free = _positive(pool.get_num_free_blocks(), "initial free cache blocks") + if free >= num_blocks or getattr(pool.null_block, "is_null", False) is not True: + raise ValueError("scheduler null block reservation is not observable") + if type(manager.watermark_blocks) is not int or manager.watermark_blocks != 0: + raise ValueError("nonzero scheduler cache watermark is unsupported") + return { + "num_blocks": num_blocks, + "initial_free_blocks": free, + "reserved_blocks": num_blocks - free, + "null_block_id": pool.null_block.block_id, + "watermark_blocks": manager.watermark_blocks, + "pool_count": 1, + "groups": cache_groups(config), + } + + +def observe( + kind: str, + owner: Any, + *, + dp_rank: int, + tp_rank: int | None = None, + pp_rank: int | None = None, + cache_config: Any = None, +) -> None: + """Write atomic evidence after successful initialization, never run a model.""" + provenance = json.loads((RESULTS_DIR / "collector-provenance.json").read_text()) + payload: dict[str, Any] = { + "schema_name": SCHEMA_NAME, + "schema_version": SCHEMA_VERSION, + "kind": kind, + "collector_provenance": provenance, + "backend_version": importlib.metadata.version("vllm"), + "dp_rank": dp_rank, + "tp_rank": tp_rank, + "pp_rank": pp_rank, + } + try: + if payload["backend_version"] != SUPPORTED_VERSION: + raise ValueError( + f"memory observation currently supports vLLM {SUPPORTED_VERSION}; found {payload['backend_version']}" + ) + if ( + getattr(owner.vllm_config, "speculative_config", None) is not None + or getattr(owner.vllm_config, "kv_transfer_config", None) is not None + ): + raise ValueError("runtime memory observation supports non-speculative local cache only") + payload["resolved_config"] = resolved_config(owner.vllm_config) + offload = payload["resolved_config"]["offload_config"] + if kind == "worker": + payload["effective_offloader"] = _offloader_type() + initial = getattr(owner, "_fpm_initial_compilation_config", None) + if initial is not None: + payload["initial_compilation_config"] = initial + cache = owner.vllm_config.cache_config + if ( + getattr(cache, "kv_cache_memory_bytes", None) is not None + or getattr(cache, "num_gpu_blocks_override", None) is not None + or offload["offload_backend"] != "auto" + or offload["uva"]["cpu_offload_gb"] != 0 + or offload["prefetch"]["offload_group_size"] != 0 + or ( + kind == "worker" + and payload["effective_offloader"] != "vllm.model_executor.offloader.base.NoopOffloader" + ) + ): + raise ValueError("manual KV allocation or offloaded cache is unsupported by runtime memory finalization") + payload["cache"] = worker_memory(owner) if kind == "worker" else scheduler_memory(owner, cache_config) + payload["status"] = "resolved" + except Exception as error: + payload.update(status="unresolved", error=f"{type(error).__name__}: {error}") + logging.getLogger(__name__).warning("FPM runtime memory remains unresolved: %s", payload["error"]) + suffix = f"-tp{tp_rank}-pp{pp_rank}" if kind == "worker" else "" + path = RESULTS_DIR / f"fpm-memory-{kind}-dp{dp_rank}{suffix}.json" + if path.exists(): + raise RuntimeError(f"refusing duplicate runtime memory observation: {path}") + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + with temporary.open("x", encoding="utf-8") as handle: + handle.write(json.dumps(payload, sort_keys=True, indent=2) + "\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_scheduler.py b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_scheduler.py new file mode 100644 index 000000000..cbbc78302 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_scheduler.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Observe the existing Dynamo scheduler after normal initialization. + +Constructor and DP rank API audited at Dynamo commit +41882ae9b07232eed4850fb1daf8c958abb2556a, components/src/dynamo/vllm/instrumented_scheduler.py. +https://github.com/ai-dynamo/dynamo/blob/41882ae9b07232eed4850fb1daf8c958abb2556a/components/src/dynamo/vllm/instrumented_scheduler.py +""" + +from fpm_memory_observer import observe, observe_cpu + +try: + from dynamo.vllm.instrumented_scheduler import InstrumentedScheduler +except ImportError as error: + raise RuntimeError( + "FPM memory observation requires an image with Dynamo's native InstrumentedScheduler; " + "see the pinned API contract in this module" + ) from error + + +class FpmExecutionInstrumentedScheduler(InstrumentedScheduler): + def __init__(self, vllm_config, kv_cache_config, *args, **kwargs): + super().__init__(vllm_config, kv_cache_config, *args, **kwargs) + observe_cpu("scheduler", dp_rank=self._fpm_dp_rank) + + +class FpmResourceInstrumentedScheduler(FpmExecutionInstrumentedScheduler): + def __init__(self, vllm_config, kv_cache_config, *args, **kwargs): + super().__init__(vllm_config, kv_cache_config, *args, **kwargs) + observe("scheduler", self, dp_rank=self._fpm_dp_rank, cache_config=kv_cache_config) diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_worker.py b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_worker.py new file mode 100644 index 000000000..998cb3ca4 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_memory_worker.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read-only worker hooks; the audited API contract is in fpm_memory_observer.""" + +import logging + +from fpm_memory_observer import compilation_config, observe, observe_cpu, observe_execution + +try: + from vllm.distributed import get_pp_group, get_tp_group + from vllm.v1.worker.gpu_worker import Worker +except ImportError as error: + raise RuntimeError( + "FPM observation requires a compatible vLLM V1 GPU-worker image; " + "execution observation supports vLLM 0.27.0/0.28.0, memory supports 0.27.0" + ) from error + + +class FpmExecutionWorker(Worker): + def compile_or_warm_up_model(self): + result = super().compile_or_warm_up_model() + parallel = self.vllm_config.parallel_config + dp_rank = getattr(parallel, "data_parallel_index", None) + if dp_rank is None: + dp_rank = parallel.data_parallel_rank + observe_cpu( + "worker", dp_rank=dp_rank, tp_rank=get_tp_group().rank_in_group, pp_rank=get_pp_group().rank_in_group + ) + observe_execution( + self, dp_rank=dp_rank, tp_rank=get_tp_group().rank_in_group, pp_rank=get_pp_group().rank_in_group + ) + return result + + +class FpmResourceWorker(FpmExecutionWorker): + def determine_available_memory(self): + # vLLM 0.27.0 profiles a minimal cache and resolves graph mode here, + # before initialize_from_config. Retain the first snapshot on repeats. + if not hasattr(self, "_fpm_initial_compilation_config"): + try: + self._fpm_initial_compilation_config = compilation_config(self.vllm_config) + except (AttributeError, TypeError) as error: + self._fpm_initial_compilation_config = None + logging.getLogger(__name__).warning("FPM initial graph configuration is unavailable: %s", error) + result = super().determine_available_memory() + self._fpm_available_cache_bytes = result + return result + + def initialize_from_config(self, kv_cache_config): + result = super().initialize_from_config(kv_cache_config) + self._fpm_cache_initialized = True + return result + + def compile_or_warm_up_model(self): + result = super().compile_or_warm_up_model() + parallel = self.vllm_config.parallel_config + dp_rank = getattr(parallel, "data_parallel_index", None) + if dp_rank is None: + dp_rank = parallel.data_parallel_rank + observe( + "worker", self, dp_rank=dp_rank, tp_rank=get_tp_group().rank_in_group, pp_rank=get_pp_group().rank_in_group + ) + return result diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_runtime_instrumentation.py b/python/aisimulate/collector/fpm_forward/runtime/fpm_runtime_instrumentation.py new file mode 100644 index 000000000..6bcba845f --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_runtime_instrumentation.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Import selected runtime adapters from verified frozen source, including on spawn. + +Only manifest-declared modules use this loader. Runtime dependencies retain normal +Python imports; this verifies provenance, and is not a sandbox for custom code. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.abc +import importlib.machinery +import importlib.util +import json +import os +import socket +import sys +from pathlib import Path, PurePosixPath + +BINDING_SCHEMA = "aisimulate-runtime-instrumentation-imports/v1" +_bundle = None + + +def _digest(raw): + return hashlib.sha256(raw).hexdigest() + + +def _json(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +class _Bundle(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def __init__(self): + context = json.loads(Path(os.environ["AISIMULATE_RUNTIME_CONTEXT"]).read_text()) + self.manifest_path = Path(os.environ["AISIMULATE_RUNTIME_INSTRUMENTATION"]).resolve() + self.root = self.manifest_path.parent + self.directory = Path(os.environ["AISIMULATE_RUNTIME_OBSERVATION_DIR"]) + self.path = self.directory / f"runtime-instrumentation-imports-{socket.gethostname()}-{os.getpid()}.json" + self.record = { + "schema_version": BINDING_SCHEMA, + **{key: context[key] for key in ("attempt_id", "configuration", "phase")}, + "hostname": socket.gethostname(), + "pid": os.getpid(), + "root": str(self.root), + "status": "verifying", + "expected_bundle_sha256": context["bundle_sha256"], + "loader_sha256": _digest(Path(__file__).read_bytes()), + "imports": {}, + "classes": {}, + } + try: + self.manifest = json.loads(self.manifest_path.read_text()) + self.files = self._files() + self.record.update(files=self.files, bundle_sha256=self._identity(self.files)) + if self.record["bundle_sha256"] != context["bundle_sha256"]: + raise ValueError("runtime instrumentation bundle hash mismatch") + if context.get("instrumentation_binding") != { + "schema_version": BINDING_SCHEMA, + "loader_sha256": self.record["loader_sha256"], + }: + raise ValueError("runtime instrumentation loader identity mismatch") + self.modules = {} + for name in self.files: + if name.endswith(".py"): + parts = list(PurePosixPath(name).with_suffix("").parts) + if parts[-1] == "__init__": + parts.pop() + self.modules[".".join(parts)] = name + for name in self.modules: + if name in sys.modules: + raise ValueError(f"instrumentation module already imported outside verified bundle loading: {name}") + sys.meta_path.insert(0, self) + except Exception as error: + self.fail(error) + raise + + def _files(self): + files = {} + for name in self.manifest["files"]: + path = PurePosixPath(name) + if path.is_absolute() or ".." in path.parts or path.as_posix() != name: + raise ValueError("unsafe runtime instrumentation path") + source = self.root / name + if not source.resolve().is_relative_to(self.root) or source.is_symlink(): + raise ValueError("runtime instrumentation source escaped the frozen bundle") + files[name] = _digest(source.read_bytes()) + return files + + def _identity(self, files): + return _digest(_json({"manifest": self.manifest, "files": files}).encode()) + + def write(self): + if self.record["pid"] != os.getpid(): + # Forked workers inherit the verified modules, not the parent's + # process identity or receipt file. Do not claim a fresh import. + self.record["inherited_from_pid"] = self.record["pid"] + self.record.update(pid=os.getpid(), hostname=socket.gethostname()) + self.path = self.directory / f"runtime-instrumentation-imports-{socket.gethostname()}-{os.getpid()}.json" + self.path.parent.mkdir(parents=True, exist_ok=True) + temporary = self.path.with_suffix(".tmp") + temporary.write_text(_json(self.record) + "\n") + temporary.replace(self.path) + + def fail(self, error): + self.record.update(status="failed", error=f"{type(error).__name__}: {error}") + self.write() + + def find_spec(self, fullname, path=None, target=None): + name = self.modules.get(fullname) + if name is None: + return None + source = self.root / name + package = [] if source.name == "__init__.py" else None + return importlib.util.spec_from_file_location(fullname, source, loader=self, submodule_search_locations=package) + + def create_module(self, spec): + return None + + def get_resource_reader(self, fullname): + # Use Python's ordinary package-data interface; source execution still + # passes through exec_module and all declared files remain hash-checked. + source = self.root / self.modules[fullname] + return importlib.machinery.SourceFileLoader(fullname, str(source)).get_resource_reader(fullname) + + def exec_module(self, module): + name = self.modules[module.__name__] + source = self.root / name + try: + raw = source.read_bytes() + digest = _digest(raw) + if digest != self.files[name]: + raise ValueError(f"runtime instrumentation source changed: {name}") + # Execute the bytes that were hashed, never an unrelated cached pyc. + exec(compile(raw, str(source), "exec"), module.__dict__) + self.record["imports"][module.__name__] = {"file": name, "path": str(source), "sha256": digest} + self.write() + except Exception as error: + self.fail(error) + raise + + def checked_record(self): + try: + if self.record["status"] == "failed": + raise ValueError(self.record["error"]) + if json.loads(self.manifest_path.read_text()) != self.manifest or self._files() != self.files: + raise ValueError("runtime instrumentation bundle changed after import") + for name, evidence in self.record["imports"].items(): + module = sys.modules.get(name) + if module is None or Path(module.__file__).resolve() != Path(evidence["path"]): + raise ValueError(f"runtime instrumentation imported module origin changed: {name}") + self.write() + return json.loads(_json(self.record)) + except Exception as error: + self.fail(error) + raise + + +def observed_binding(): + """Read actual imports after hooks initialize; declarations alone cannot pass.""" + if _bundle is None: + raise ValueError("runtime instrumentation did not load through its verified entrypoint") + return _bundle.checked_record() + + +def __getattr__(name): + global _bundle + field = {"ObservedWorker": "worker_class", "ObservedInstrumentedScheduler": "scheduler_class"}.get(name) + if field is None: + raise AttributeError(name) + if _bundle is None: + _bundle = _Bundle() + try: + module_name, _, class_name = _bundle.manifest[field].rpartition(".") + module = importlib.import_module(module_name) + actual = getattr(module, class_name) + if not isinstance(actual, type) or actual.__module__ not in _bundle.record["imports"]: + raise ValueError("instrumentation class must be defined by an imported frozen module") + # An empty subclass preserves native methods and isinstance checks. Its + # stable module/name also sends spawned-process unpickling through here. + selected = type(name, (actual,), {"__module__": __name__}) + _bundle.record["classes"][field] = _bundle.manifest[field] + _bundle.record["status"] = "verified" + _bundle.checked_record() + globals()[name] = selected + return selected + except Exception as error: + _bundle.fail(error) + raise diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/README.md b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/README.md new file mode 100644 index 000000000..a72edef7d --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/README.md @@ -0,0 +1,57 @@ + + +# Runtime observation adapters + +These are original AISimulate adapters written against inspected runtime interfaces. They contain no copied vLLM or Dynamo implementation. The immutable upstream revisions and installation-relative file SHA-256 values are recorded in each manifest. The version-specific source note is part of the hashed bundle. vLLM and Dynamo are Apache-2.0 projects; these references document the interfaces used, without redistributing their source. + +`vllm-0.27.0.json` is the bundled selection. `vllm-0.28.0.example.json` is an explicit campaign-local example for its inspected runtime build. Freeze the latter with `load_instrumentation` and `freeze_instrumentation` into a fresh campaign directory, as shown in the [user guide](../../../../../../docs/fpm-self-service.md#resolve-cache-geometry-with-a-runtime-probe). Other runtime pins need their own source audit, manifest and adapter changes; changing the version string alone is insufficient. An example bundle is implementation evidence, not qualification of every model, backend, GPU or topology on that version. + +## Runtime interface + +The collector stages the frozen bundle and a collector-owned import entrypoint, and supplies these environment variables: + +| Variable | Contents | +| --- | --- | +| `AISIMULATE_RUNTIME_CONTEXT` | Absolute path to the runner-owned `aisimulate-runtime-probe-launch/v1` JSON document: `attempt_id`, `configuration`, `phase`, `bundle_sha256`, exact `launch`, and `expected_ranks`. Sampling is a separate field. | +| `AISIMULATE_RUNTIME_INSTRUMENTATION` | Absolute path to the frozen JSON `manifest.json`. | +| `AISIMULATE_RUNTIME_OBSERVATION_DIR` | The current process's results directory; the collector uses `/results`. | + +Worker and scheduler classes inherit the native vLLM worker and Dynamo `InstrumentedScheduler`. The worker preserves return values and exceptions from memory profiling, cache initialization and warm-up. The scheduler calls the native constructor first. Neither intercepts timed forwards, selects kernels, changes scheduler settings, or changes sampling. Python modules are imported only by those runtime processes; bundle preview/freezing and observation import do not import them. + +Both adapters export `ObservedInstrumentedScheduler`. The runner passes `fpm_runtime_instrumentation.ObservedInstrumentedScheduler` to `--scheduler-cls`; both pinned Dynamo argument consumers require its string to contain `InstrumentedScheduler` before importing the class. The entrypoint verifies the complete frozen bundle, including source notes, then imports only requested modules from their declared paths. It rejects modules previously loaded outside that verified import path. Its method-free subclass retains native inheritance and routes spawned-process unpickling through the same verification. Runtime dependencies outside the manifest keep normal Python import behavior; this is provenance checking, not a sandbox for campaign code. + +The entrypoint emits `runtime-instrumentation-imports--.json` with actual file hashes, loaded module paths, classes and failures. The separate built-in Slurm launcher CPU observer is outside this bundle and is not attributed to its hash. New contexts contain `instrumentation_binding` with the collector loader's hash. Every worker/scheduler observation for such a context must include `instrumentation_binding` returned by `from fpm_runtime_instrumentation import observed_binding`; call `observed_binding()` after the native lifecycle hook, never copy the context marker as observation. The bundled adapter does this automatically. The CPU importer checks observed imports against the frozen manifest and context without executing campaign code. Archived contexts without this marker remain readable with `import_binding: unverified_legacy`; they are not retroactively verified or rewritten. + +## Observation meanings + +| Field | Source, units and lifecycle | +| --- | --- | +| `runtime.version` | Installed `importlib.metadata.version("vllm")`, checked against both class version and manifest. | +| `runtime.source_revision`, `runtime.source_files` | Pinned interpretation revision plus SHA-256 of actual installed files, including the Dynamo scheduler and inspected backends. All must match. A supplied generated vLLM commit ID is checked and retained separately in `source_identity`; file hashes do not attest the entire source tree. | +| `launch`, `identity`, attempt/configuration/phase/bundle | Runner context copied verbatim. These declarations are independently checked against actual resolved configuration and hardware on import. The separate `instrumentation_binding` records the verified runtime bundle bytes and modules actually imported. | +| `resolved_config` | Actual initialized vLLM model/cache/scheduler/parallel/compilation/kernel/quantization/offload fields, including speculative and transfer configuration. No launch value substitutes for an unobserved field. | +| `model_config_sha256`, `loaded_model_config` | SHA-256 of the actual `hf_config_path or model` directory's `config.json`, or the local HF cache's `config.json` at `hf_config._commit_hash`. No network request is made. Missing files, non-HF formats and version-selected alternative configuration files remain unresolved. The launch passes the reviewed `--revision` to native model construction, including for local directories. Import requires the observed revision and independently checks the loaded config bytes; this binds configuration and the declared source revision, not weight bytes. | +| `model_config_source_files` | When the launch declares adjacent model-config sidecars, SHA-256 values are read from the actual config directory. Missing files remain unresolved; expected hashes are never echoed. The primary config bytes remain unchanged. | +| `hardware` | Worker device properties: CUDA major/minor combined as SM, actual device name and physical total bytes. The requested GPU label is retained separately for independent family/SM checking. | +| `cache.available_cache_bytes` | The byte allowance returned by native `determine_available_memory`, preserved after `super` succeeds. It is not post-warm-up free HBM. | +| `raw_cache_config` | Unnormalized spec fields and allocation declarations retained independently, including when a spec or layout cannot be represented. | +| `cache.num_blocks`, `groups` | Initialized `KVCacheConfig`: exact layer membership, raw spec class, logical block tokens, spec page bytes, dtype and retention. Spec page bytes come from the runtime property, including padding. | +| `cache.storages` | Unique `(CUDA device, untyped storage pointer)` with actual backing-storage bytes and a stable rank-local storage ID. Shared backing storage is counted once. | +| `cache.tensor_allocations` | Runtime `size/shared_by/block_stride/offset`, with the observed storage ID. Size is the entire backing storage; stride and offset are bytes. Zero block stride denotes an unpacked allocation. Runtime declarations must explain all aliases. | +| `cache.layer_tensors` | Every layer's actual tensor shape, byte strides, byte storage offset and element width. Block axis comes from its initialized backend's `get_kv_cache_block_dim`, never from matching a shape value. Kernel block size comes from the exact initialized runner: V1 uses `_kernel_block_sizes`, V2 uses `kernel_block_sizes`. The observer records the actual runner class and sizes, requires that implementation's source hash, and rejects unknown classes or missing initialized state. | +| `cache.pool_id` | Stable logical ID per DP worker/scheduler pair. Actual scheduler manager identity proves every group shares the same pool; process pointers are not compared across processes. | +| `cache.initial_free_blocks`, reservations | Initial block-pool count and free queue after native scheduler construction; set difference gives reserved IDs. These audited constructors reserve only the null block. Additional reservations remain unresolved until their permanence has a source mapping. | +| `lifecycle` | Workers record only after native warm-up returns successfully and cache initialization was observed. `capture_completed` means the runtime's selected capture policy completed, including its intentional eager/no-capture path. Actual graph policy and sizes remain in `resolved_config`. Schedulers record after native initialization. | + +The supported representation is one shared HBM block pool with full/sliding retention, zero watermark, no speculative decoding, transfer or offload, and no compressed storage blocks. The existing audited `InklingConvState`/window-four mapping is preserved by the historical helper; any model-specific backend also needs its own declared source hash before this new bundle can complete. Unknown classes and retention are not renamed into supported specs. Actual prefix-caching configuration is retained; `semantics.prefix_reuse=false` does not authorize cross-request replay reuse. + +Both logical K/V-plane-first and packed views retain their physical strides. The CPU importer independently checks affine byte bounds, pool/kernel block ratio, packed slot intervals, padding, declared aliases, reservations and compatible rank/phase settings. An adapter-supplied status cannot authorize import. + +Independent collection errors are kept in `unresolved_fields` and `error` with available raw data. Duplicate records fail instead of overwriting earlier evidence. Native initialization/warm-up exceptions propagate unchanged. Formal collection also emits the unchanged `fpm-execution-worker-*.json` evidence for existing execution-quality consumers. + +## Checks for a new campaign-local adapter + +Check the frozen manifest's scheduler string through the pinned native benchmark argument consumer before direct class construction, and verify the exported class inherits the native scheduler. Exercise the runtime extension classes against small runtime-interface fakes before using GPUs. Cover actual backend block axes (including repeated dimension lengths), packed offsets, padded pages, declared aliases, source/version mismatches, actual config-byte hashes, absent fields, retention changes and native exceptions. Then use the normal probe command on the selected hardware and import both phases for every selected configuration. Preserve failed attempts and corroborating initialization logs. A source audit and fake-runtime test cannot replace that live campaign evidence. diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/__init__.py b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/__init__.py new file mode 100644 index 000000000..52a7a9daf --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/hooks.py b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/hooks.py new file mode 100644 index 000000000..6e43a38c6 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/hooks.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lifecycle hooks shared by the two source-audited adapters.""" + +import logging + +from fpm_memory_observer import compilation_config +from vllm.distributed import get_pp_group, get_tp_group + +from .observer import observe + + +class WorkerObservation: + def determine_available_memory(self): + if not hasattr(self, "_fpm_initial_compilation_config"): + try: + self._fpm_initial_compilation_config = compilation_config(self.vllm_config) + except (AttributeError, TypeError) as error: + logging.getLogger(__name__).warning("Initial graph configuration unavailable: %s", error) + result = super().determine_available_memory() + self._fpm_available_cache_bytes = result + return result + + def initialize_from_config(self, kv_cache_config): + result = super().initialize_from_config(kv_cache_config) + self._fpm_cache_initialized = True + return result + + def compile_or_warm_up_model(self): + result = super().compile_or_warm_up_model() + self._observation_warmup_completed = True + parallel = self.vllm_config.parallel_config + dp_rank = getattr(parallel, "data_parallel_index", None) + if dp_rank is None: + dp_rank = parallel.data_parallel_rank + observe( + "worker", + self, + version=self.observation_version, + dp_rank=dp_rank, + tp_rank=get_tp_group().rank_in_group, + pp_rank=get_pp_group().rank_in_group, + ) + return result + + +class SchedulerObservation: + def __init__(self, vllm_config, kv_cache_config, *args, **kwargs): + super().__init__(vllm_config, kv_cache_config, *args, **kwargs) + observe( + "scheduler", self, version=self.observation_version, dp_rank=self._fpm_dp_rank, cache_config=kv_cache_config + ) diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/observer.py b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/observer.py new file mode 100644 index 000000000..333149d91 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/observer.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Original read-only adapters; see the bundled version-specific source notes. + +The runner imports this file only inside the selected runtime. Inspection and +CPU import use the manifest, never this module. No model or timing code is run. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.metadata +import json +import logging +import os +from pathlib import Path +from typing import Any + +import fpm_memory_observer as legacy + + +def _json(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): _json(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_json(item) for item in value] + return legacy._json_value(value) + + +def _installed_file(name: str) -> Path: + package, *parts = name.split("/") + module = importlib.import_module(package) + paths = [Path(root).joinpath(*parts) for root in module.__path__] + found = [path for path in paths if path.is_file()] + if len(found) != 1: + raise ValueError(f"runtime source path must resolve uniquely: {name}") + return found[0] + + +def _runtime(payload: dict, manifest: dict, version: str) -> None: + actual_version = importlib.metadata.version("vllm") + build = importlib.import_module("vllm._version") + commit = getattr(build, "__commit_id__", None) + pinned = manifest["runtime"] + observed = payload["runtime"] = { + "framework": "vllm", + "version": actual_version, + "source_revision": pinned["source_revision"], + "source_files": {}, + "image": payload["launch"]["deployment"]["image"], + } + payload["source_identity"] = {"vllm_commit_id": commit, "source_revision_basis": "declared source and file hashes"} + errors = [] + if actual_version != version or pinned["version"] != version: + errors.append(f"adapter requires vLLM {version}; installed={actual_version}, manifest={pinned['version']}") + # Wheels may omit a VCS identifier. The exact inspected source files remain + # mandatory; when a wheel supplies an identifier, it must also agree. + if commit is not None: + short = str(commit).removeprefix("g") + if len(short) < 7 or not pinned["source_revision"].startswith(short): + errors.append(f"installed vLLM commit {commit!r} contradicts source revision") + for name, expected in pinned["source_files"].items(): + try: + digest = hashlib.sha256(_installed_file(name).read_bytes()).hexdigest() + observed["source_files"][name] = digest + if digest != expected: + errors.append(f"runtime source hash differs: {name}") + except (OSError, ValueError, AttributeError, ImportError) as error: + errors.append(f"runtime source {name}: {error}") + if errors: + raise ValueError("; ".join(errors)) + + +def _loaded_config(owner: Any, payload: dict) -> None: + model = owner.vllm_config.model_config + config_format = getattr(model, "config_format", "auto") + if config_format not in ("auto", "hf"): + raise ValueError(f"loaded config format requires another source mapping: {config_format}") + source = getattr(model, "hf_config_path", None) or model.model + candidate = Path(source) + if candidate.is_dir(): + if config_format == "auto" and (candidate / "params.json").exists(): + raise ValueError("automatic HF/Mistral config selection requires an explicit loaded-file mapping") + candidate = candidate / "config.json" + elif not candidate.is_file(): + # Resolve only the already-loaded immutable HF revision. Never download + # another config during observation or hash a caller-supplied host copy. + commit = getattr(getattr(model, "hf_config", None), "_commit_hash", None) + if not isinstance(commit, str) or len(commit) != 40: + raise ValueError("loaded remote config has no immutable cache commit") + from huggingface_hub import try_to_load_from_cache + + if config_format == "auto" and isinstance(try_to_load_from_cache(source, "params.json", revision=commit), str): + raise ValueError("automatic cached HF/Mistral config selection is ambiguous") + cached = try_to_load_from_cache(source, "config.json", revision=commit) + if not isinstance(cached, str): + raise ValueError("actual loaded config.json is absent from the local HF cache") + candidate = Path(cached) + if candidate.name != "config.json": + raise ValueError("this adapter requires an actual loaded HF config.json") + raw = candidate.read_bytes() + document = json.loads(raw) + if document.get("configuration_files"): + raise ValueError("version-selected HF configuration files require a separate audited mapping") + digest = hashlib.sha256(raw).hexdigest() + payload["model_config_sha256"] = digest + payload["loaded_model_config"] = {"path": str(candidate.resolve()), "sha256": digest} + expected_sources = payload["launch"].get("model_config", {}).get("source_files") + if expected_sources is not None: + observed = payload["model_config_source_files"] = {} + for name in expected_sources: + if ( + not isinstance(name, str) + or name in ("", ".", "..") + or Path(name).name != name + or any(char in name for char in ("\\", ":", "\x00", "\n", "\r")) + ): + raise ValueError("model config sidecars must be adjacent relative filenames") + observed[name] = hashlib.sha256((candidate.parent / name).read_bytes()).hexdigest() + + +def _resolved(owner: Any) -> dict: + config = owner.vllm_config + value = legacy.resolved_config(config) + for key in ("speculative_config", "kv_transfer_config"): + value[key] = _json(getattr(config, key, None)) + for key in ("hf_config_path", "hf_overrides", "config_format"): + if hasattr(config.model_config, key): + value["model_config"][key] = _json(getattr(config.model_config, key)) + return value + + +def _raw_cache(config: Any) -> dict: + fields = ( + "block_size", + "storage_block_size", + "page_size_bytes", + "page_size_padded", + "dtype", + "sliding_window", + "extra_retained_tokens", + "attention_chunk_size", + "num_kv_heads", + "head_size", + "kv_quant_mode", + ) + return { + "num_blocks": config.num_blocks, + "groups": [ + { + "layer_names": list(group.layer_names), + "spec_type": legacy._class_name(group.kv_cache_spec), + **{ + key: _json(getattr(group.kv_cache_spec, key)) for key in fields if hasattr(group.kv_cache_spec, key) + }, + } + for group in config.kv_cache_groups + ], + "tensor_allocations": [ + { + key: _json(getattr(tensor, key)) + for key in ("size", "shared_by", "block_stride", "offset") + if hasattr(tensor, key) + } + for tensor in config.kv_cache_tensors + ], + } + + +def _groups(cache: dict, config: Any, version: str, pool_id: str) -> None: + errors = [] + for raw, group in zip(config.kv_cache_groups, cache["groups"], strict=True): + spec = raw.kv_cache_spec + # 0.27 lacks this field. 0.28 introduces it on SlidingWindowSpec; its + # absence there is an incompatible runtime API, not a default of zero. + if version == "0.28.0" and group["spec_type"].endswith(".SlidingWindowSpec"): + extra = getattr(spec, "extra_retained_tokens", None) + else: + extra = getattr(spec, "extra_retained_tokens", 0) + chunk = getattr(spec, "attention_chunk_size", None) + group.update(pool_id=pool_id, extra_retained_tokens=extra, attention_chunk_size=chunk) + if type(extra) is not int or extra != 0 or chunk is not None: + errors.append(f"unsupported cache retention for {group['layer_names']}: extra={extra!r}, chunk={chunk!r}") + if getattr(spec, "storage_block_size", spec.block_size) != spec.block_size: + errors.append(f"compressed cache storage is unsupported: {group['layer_names']}") + if errors: + raise ValueError("; ".join(errors)) + + +def _layer_views(owner: Any, cache: dict, source_files: dict) -> None: + runner = owner.model_runner + runner_class = legacy._class_name(runner) + # Worker.init_device selects either native runner. Each initialize_kv_cache + # stores the sizes it passes to the tensor allocator under a different name; + # missing state must not fall back to another runner's field or spec sizes. + size_fields = { + "vllm.v1.worker.gpu_model_runner.GPUModelRunner": "_kernel_block_sizes", + "vllm.v1.worker.gpu.model_runner.GPUModelRunner": "kernel_block_sizes", + } + if runner_class not in size_fields: + raise ValueError(f"model runner has no audited cache mapping: {runner_class}") + runner_source = type(runner).__module__.replace(".", "/") + ".py" + if runner_source not in source_files: + raise ValueError(f"selected model runner source has not been audited: {runner_source}") + kernel_block_sizes = getattr(runner, size_fields[runner_class]) + if not isinstance(kernel_block_sizes, list) or len(kernel_block_sizes) != len(cache["groups"]): + raise ValueError("initialized kernel block sizes do not cover every cache group") + kernel_block_sizes = [legacy._positive(size, "kernel block size") for size in kernel_block_sizes] + cache["model_runner"] = {"class": runner_class, "kernel_block_sizes": kernel_block_sizes} + layers = owner.vllm_config.compilation_config.static_forward_context + storage_ids = {} + for index, storage in enumerate(cache["storages"]): + storage["storage_id"] = f"storage-{index}" + storage_ids[(storage["device"], storage["pointer"])] = storage["storage_id"] + for allocation in cache["tensor_allocations"]: + key, _ = legacy._storage(layers[allocation["shared_by"][0]].kv_cache) + allocation["storage_id"] = storage_ids[key] + views = cache["layer_tensors"] = {} + groups = [] + backends = set() + for attention_groups in runner.attn_groups: + for group in attention_groups: + spec = group.kv_cache_spec + kernel = kernel_block_sizes[group.kv_cache_group_id] + if getattr(spec, "storage_block_size", spec.block_size) != spec.block_size: + raise ValueError("compressed kernel cache blocks are unsupported") + # Both audited runners use auto for an unquantized group's shape. + # Read KVQuantMode by its public enum name; do not infer from dtype. + mode = getattr(spec, "kv_quant_mode", None) + if mode is None: + raise ValueError("cache spec has no observed KV quantization mode") + dtype = ( + "auto" + if getattr(mode, "name", None) == "NONE" + else (getattr(spec, "cache_dtype_str", None) or owner.vllm_config.cache_config.cache_dtype) + ) + backend = group.backend + backend_source = backend.__module__.replace(".", "/") + ".py" + if backend_source not in source_files: + raise ValueError(f"selected backend source has not been audited: {backend_source}") + axis = backend.get_kv_cache_block_dim(kernel, spec.num_kv_heads, spec.head_size, cache_dtype_str=dtype) + if type(axis) is not int or axis < 0: + raise ValueError("backend returned an invalid cache block dimension") + backends.add(backend.get_name()) + groups.append( + { + "kv_cache_group_id": group.kv_cache_group_id, + "backend_class": f"{backend.__module__}.{backend.__qualname__}", + "layer_names": list(group.layer_names), + } + ) + for name in group.layer_names: + if name in views: + raise ValueError(f"duplicate attention group layer: {name}") + tensor = layers[name].kv_cache + key, _ = legacy._storage(tensor) + width = legacy._positive(tensor.element_size(), "tensor element size") + shape = list(tensor.shape) + if axis >= len(shape): + raise ValueError("backend block axis is outside the physical tensor shape") + views[name] = { + "storage_id": storage_ids[key], + "storage_offset_bytes": tensor.storage_offset() * width, + "shape": shape, + "stride_bytes": [stride * width for stride in tensor.stride()], + "element_size_bytes": width, + "block_axis": axis, + "kernel_block_size_tokens": kernel, + } + expected = {name for group in cache["groups"] for name in group["layer_names"]} + if set(views) != expected: + raise ValueError("attention backend views do not cover every physical cache layer") + cache["attention_groups"] = groups + cache["attention_backends"] = sorted(backends) + + +def _reservations(owner: Any, cache: dict) -> None: + pool = owner.kv_cache_manager.block_pool + free = [block.block_id for block in pool.free_block_queue.get_all_free_blocks()] + ids = [block.block_id for block in pool.blocks] + if len(set(ids)) != cache["num_blocks"] or set(ids) != set(range(cache["num_blocks"])): + raise ValueError("block pool does not expose every unique block ID") + if len(set(free)) != len(free) or len(free) != cache["initial_free_blocks"] or not set(free) <= set(ids): + raise ValueError("free block queue disagrees with the observed initial count") + reserved = sorted(set(ids) - set(free)) + cache["permanent_reserved_block_ids"] = reserved + cache["group_pool_ids"] = [cache["pool_id"]] * len(cache["groups"]) + # Audited BlockPool constructors reserve only null_block. Do not classify + # a live request allocation as a permanent reservation from its ref count. + if reserved != [pool.null_block.block_id]: + raise ValueError("additional initialized reservations lack an audited permanent-reservation mapping") + + +def _semantics(owner: Any, payload: dict) -> None: + config = owner.vllm_config + offload = payload["resolved_config"]["offload_config"] + semantics = payload["cache"]["semantics"] = { + "allocation": "shared_block_pool", + "storage": "hbm", + "retention": "full_or_window", + "prefix_reuse": False, + "offload": ( + offload["offload_backend"] != "auto" + or offload["uva"]["cpu_offload_gb"] != 0 + or offload["prefetch"]["offload_group_size"] != 0 + or payload.get("effective_offloader") != "vllm.model_executor.offloader.base.NoopOffloader" + ), + "speculative": getattr(config, "speculative_config", None) is not None, + } + if ( + semantics["offload"] + or semantics["speculative"] + or getattr(config, "kv_transfer_config", None) is not None + or getattr(config.cache_config, "kv_cache_memory_bytes", None) is not None + or getattr(config.cache_config, "num_gpu_blocks_override", None) is not None + ): + raise ValueError("offload, speculative, transferred or manually overridden cache is unsupported") + + +def _hardware(owner: Any, gpu: str) -> dict: + import torch + + properties = torch.cuda.get_device_properties(owner.device) + return { + "gpu": gpu, + "sm": properties.major * 10 + properties.minor, + "device_name": properties.name, + "total_memory_bytes": properties.total_memory, + } + + +def observe( + kind: str, + owner: Any, + *, + version: str, + dp_rank: int, + tp_rank: int | None = None, + pp_rank: int | None = None, + cache_config: Any = None, +) -> None: + """Publish raw independent observations after a successful runtime hook.""" + context = json.loads(Path(os.environ["AISIMULATE_RUNTIME_CONTEXT"]).read_text()) + manifest = json.loads(Path(os.environ["AISIMULATE_RUNTIME_INSTRUMENTATION"]).read_text()) + directory = Path(os.environ["AISIMULATE_RUNTIME_OBSERVATION_DIR"]) + payload = { + "schema_version": "aisimulate-runtime-observation/v1", + **{key: context[key] for key in ("attempt_id", "configuration", "phase", "bundle_sha256", "launch")}, + "identity": context["launch"]["identity"], + "kind": kind, + "dp_rank": dp_rank, + "tp_rank": tp_rank, + "pp_rank": pp_rank, + "unresolved_fields": [], + "lifecycle": { + "measurement": "after_warmup", + "cache_initialized": bool(getattr(owner, "_fpm_cache_initialized", False)), + "warmup_completed": bool(getattr(owner, "_observation_warmup_completed", False)), + "capture_completed": bool(getattr(owner, "_observation_warmup_completed", False)), + } + if kind == "worker" + else {"measurement": "scheduler_initialized", "cache_initialized": True}, + } + errors = [] + + # CPU readiness is independent of physical-memory import. A missing or + # partial CPU snapshot must not relabel otherwise valid cache observations. + try: + cpu = legacy.observe_cpu(kind, dp_rank=dp_rank, tp_rank=tp_rank, pp_rank=pp_rank, directory=directory) + payload["cpu_affinity_observation"] = {"status": cpu["status"], "hostname": cpu["hostname"], "pid": cpu["pid"]} + except Exception as error: + payload["cpu_affinity_observation"] = {"status": "unavailable", "error": f"{type(error).__name__}: {error}"} + logging.getLogger(__name__).warning("CPU affinity observation unavailable: %s", error) + + def capture(label, action): + try: + action() + except Exception as error: + payload["unresolved_fields"].append(label) + errors.append(f"{label}: {type(error).__name__}: {error}") + + if "instrumentation_binding" in context: + from fpm_runtime_instrumentation import observed_binding + + capture("instrumentation_binding", lambda: payload.update(instrumentation_binding=observed_binding())) + capture("runtime", lambda: _runtime(payload, manifest, version)) + capture("model_config_sha256", lambda: _loaded_config(owner, payload)) + capture("resolved_config", lambda: payload.update(resolved_config=_resolved(owner))) + pool_id = f"dp{dp_rank}-shared-pool" + if kind == "worker": + capture("hardware", lambda: payload.update(hardware=_hardware(owner, payload["identity"]["gpu"]))) + capture("effective_offloader", lambda: payload.update(effective_offloader=legacy._offloader_type())) + if hasattr(owner, "_fpm_initial_compilation_config"): + payload["initial_compilation_config"] = owner._fpm_initial_compilation_config + capture("cache", lambda: payload.update(cache=legacy.worker_memory(owner))) + config = getattr(getattr(owner, "model_runner", None), "kv_cache_config", None) + else: + config = cache_config + capture("cache", lambda: payload.update(cache=legacy.scheduler_memory(owner, config))) + capture("raw_cache_config", lambda: payload.update(raw_cache_config=_raw_cache(config))) + if "cache" in payload: + cache = payload["cache"] + cache["pool_id"] = pool_id + capture("cache.groups", lambda: _groups(cache, config, version, pool_id)) + if kind == "worker": + capture("cache.semantics", lambda: _semantics(owner, payload)) + capture("cache.layer_tensors", lambda: _layer_views(owner, cache, manifest["runtime"]["source_files"])) + if "attention_groups" in cache: + payload["attention_groups"] = cache.pop("attention_groups") + names = cache.pop("attention_backends") + payload["attention_backend"] = names[0] if len(names) == 1 else names + else: + capture("cache.permanent_reserved_block_ids", lambda: _reservations(owner, cache)) + if kind == "worker" and ("collector_attempt_id" in context or (directory / "collector-provenance.json").exists()): + # Retain the existing formal-collection execution contract unchanged. + previous = legacy.RESULTS_DIR + try: + legacy.RESULTS_DIR = directory + capture( + "execution_evidence", + lambda: legacy.observe_execution(owner, dp_rank=dp_rank, tp_rank=tp_rank, pp_rank=pp_rank), + ) + finally: + legacy.RESULTS_DIR = previous + if errors: + payload["error"] = "; ".join(errors) + logging.getLogger(__name__).warning("Runtime observation incomplete: %s", payload["error"]) + directory.mkdir(parents=True, exist_ok=True) + suffix = f"-tp{tp_rank}-pp{pp_rank}" if kind == "worker" else "" + path = directory / f"runtime-observation-{kind}-dp{dp_rank}{suffix}.json" + legacy._write_execution_file(path, payload) diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.27.0.md b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.27.0.md new file mode 100644 index 000000000..4c386ad4b --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.27.0.md @@ -0,0 +1,26 @@ + + +# vLLM 0.27.0 source mapping + +Read [the common field and lifecycle contract](README.md) with this version-specific audit. The manifest pins vLLM revision `4bdc8a788d2e2ce9165d552b3d4d8b72604626bf` and Dynamo scheduler revision `41882ae9b07232eed4850fb1daf8c958abb2556a` through exact installation-file hashes. The source notes and adapter files are hashed into each campaign bundle. + +| Interpretation | Immutable source | +| --- | --- | +| Native profile allowance, initialization and post-warm-up boundary | [GPU worker](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu_worker.py): `determine_available_memory`, `initialize_from_config`, `compile_or_warm_up_model`. | +| Actual backing allocation, packed slot metadata and resolved kernel blocks | [GPU model runner](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu_model_runner.py): `_allocate_kv_cache_tensors`, `_reshape_kv_cache_tensors`, `initialize_kv_cache`. | +| V2 runner selection and initialized kernel blocks | [GPU worker](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu_worker.py): `init_device` selects the runner; [V2 model runner](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu/model_runner.py): `initialize_kv_cache` records `kernel_block_sizes` returned by `init_attn_backend` and passes that same list to `init_kv_cache`. V1 records `_kernel_block_sizes` instead. Both implementation files are source-hash bound. | +| KV-plane/packed/padded affine views and kernel block subdivision | [GPU attention utilities](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/gpu/attn_utils.py): `_reshape_attention_kv_cache`, `_reshape_kv_cache`. | +| Backend-provided block axis and layout | [backend interface](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/attention/backend.py): `get_kv_cache_block_dim`; [FlashAttention](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/attention/backends/flash_attn.py) and [FlashInfer](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/attention/backends/flashinfer.py): `get_kv_cache_shape`, `get_kv_cache_stride_order`. Physical tensor strides retain the selected HND/NHD layout. | +| Attention group/kernel-block preparation and binding views to layers | [worker utilities](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/worker/utils.py): `AttentionGroup`, `prepare_kernel_block_sizes`, `bind_kv_cache`. | +| Raw spec classes, dtype, page bytes and retention | [KV cache interface](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/kv_cache_interface.py): `FullAttentionSpec`, `SlidingWindowSpec`, allocation tensor fields. | +| Shared pool, null reservation, free queue and retention manager | [cache manager](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/core/kv_cache_manager.py), [block pool](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/core/block_pool.py), [cache utilities](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/core/kv_cache_utils.py), [single-type managers](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/v1/core/single_type_kv_cache_manager.py). The scheduler observation precedes request scheduling. | +| Loaded HF configuration source | [model configuration](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/config/model.py): `get_config(hf_config_path or model, ..., revision, ...)`; [configuration parser](https://github.com/vllm-project/vllm/blob/4bdc8a788d2e2ce9165d552b3d4d8b72604626bf/vllm/transformers_utils/config.py): `HFConfigParser.parse`. | +| Scheduler constructor, DP rank and native self-benchmark | [Dynamo instrumented scheduler](https://github.com/ai-dynamo/dynamo/blob/41882ae9b07232eed4850fb1daf8c958abb2556a/components/src/dynamo/vllm/instrumented_scheduler.py): native initialization and `_fpm_dp_rank`. | +| Benchmark scheduler argument acceptance | [Dynamo argument consumer](https://github.com/ai-dynamo/dynamo/blob/41882ae9b07232eed4850fb1daf8c958abb2556a/components/src/dynamo/vllm/args.py#L354): `update_engine_config_with_dynamo` requires `InstrumentedScheduler` in the scheduler-class string before import. The exported class is `ObservedInstrumentedScheduler`, a native subclass. | + +`SlidingWindowSpec` has no `extra_retained_tokens` field at this revision. Its absence maps explicitly to zero. An unexpected nonzero value remains unsupported. + +Both pins reject attention chunk retention, speculative groups, unknown cache spec classes and unsupported alias layouts. Runtime source checks include the selected backend module; selecting another backend requires auditing and declaring that implementation rather than claiming that these inspected backends cover it. A vendor patch changing a declared file produces unresolved evidence with the actual hash retained. Update source notes, mapping tests and the manifest in a fresh campaign attempt after inspecting such a patch. diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.28.0.md b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.28.0.md new file mode 100644 index 000000000..7fcb2d8c2 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm-0.28.0.md @@ -0,0 +1,28 @@ + + +# vLLM 0.28.0 source mapping + +Read [the common field and lifecycle contract](README.md) with this version-specific audit. The manifest pins vLLM revision `2cf0a6915ce544dc493a0990f2ea38d81601128a` and Dynamo scheduler revision `b83b1d9304ebfc624709ac46db32b1b6f1ff1615` through exact installation-file hashes. The source notes and adapter files are hashed into each campaign bundle. + +| Interpretation | Immutable source | +| --- | --- | +| Native profile allowance, initialization and post-warm-up boundary | [GPU worker](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu_worker.py): `determine_available_memory`, `initialize_from_config`, `compile_or_warm_up_model`. | +| Actual backing allocation, packed slot metadata and resolved kernel blocks | [GPU model runner](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu_model_runner.py): `_allocate_kv_cache_tensors`, `_reshape_kv_cache_tensors`, `initialize_kv_cache`. | +| V2 runner selection and initialized kernel blocks | [GPU worker](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu_worker.py): `init_device` selects the runner; [V2 model runner](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu/model_runner.py): `initialize_kv_cache` records `kernel_block_sizes` returned by `init_attn_backend` and passes that same list to `init_kv_cache`. V1 records `_kernel_block_sizes` instead. Both implementation files are source-hash bound. | +| KV-plane/packed/padded affine views and kernel block subdivision | [GPU attention utilities](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/gpu/attn_utils.py): `_reshape_attention_kv_cache`, `_reshape_kv_cache`. | +| Backend-provided block axis and layout | [backend interface](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/attention/backend.py): `get_kv_cache_block_dim`; [FlashAttention](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/attention/backends/flash_attn.py) and [FlashInfer](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/attention/backends/flashinfer.py): `get_kv_cache_shape`, `get_kv_cache_stride_order`. Physical tensor strides retain the selected HND/NHD layout. | +| Attention group/kernel-block preparation and binding views to layers | [worker utilities](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/worker/utils.py): `AttentionGroup`, `prepare_kernel_block_sizes`, `bind_kv_cache`. | +| Raw spec classes, dtype, page bytes and retention | [KV cache interface](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/kv_cache_interface.py): `FullAttentionSpec`, `SlidingWindowSpec`, allocation tensor fields. | +| Shared pool, null reservation, free queue and retention manager | [cache manager](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/core/kv_cache_manager.py), [block pool](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/core/block_pool.py), [cache utilities](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/core/kv_cache_utils.py), [single-type managers](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/v1/core/single_type_kv_cache_manager.py). The scheduler observation precedes request scheduling. | +| Loaded HF configuration source | [model configuration](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/config/model.py): `get_config(hf_config_path or model, ..., revision, ...)`; [configuration parser](https://github.com/vllm-project/vllm/blob/2cf0a6915ce544dc493a0990f2ea38d81601128a/vllm/transformers_utils/config.py): `HFConfigParser.parse`. | +| Scheduler constructor, DP rank and native self-benchmark | [Dynamo instrumented scheduler](https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/instrumented_scheduler.py): native initialization and `_fpm_dp_rank`. | +| Benchmark scheduler argument acceptance | [Dynamo argument consumer](https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/args.py#L393): `update_engine_config_with_dynamo` requires `InstrumentedScheduler` in the scheduler-class string before import. The exported class is `ObservedInstrumentedScheduler`, a native subclass. | + +`SlidingWindowSpec.extra_retained_tokens` is present at this revision and extends retained history for speculative prefill. The adapter requires the actual integer value zero; an absent or nonzero value remains unresolved. Page sizing changed from 0.27, so this adapter reads the actual `page_size_bytes` property and rejects compressed `storage_block_size` instead of applying an older formula. + +Both pins reject attention chunk retention, speculative groups, unknown cache spec classes and unsupported alias layouts. Runtime source checks include the selected backend module; selecting another backend requires auditing and declaring that implementation rather than claiming that these inspected backends cover it. A vendor patch changing a declared file produces unresolved evidence with the actual hash retained. Update source notes, mapping tests and the manifest in a fresh campaign attempt after inspecting such a patch. + +This example was source-audited against the ARM64 Dynamo 1.5.0 image `nvcr.io/nvidia/ai-dynamo/vllm-runtime@sha256:93506f88fc21e74aa2611fedf1ac8077040dab509e1861a62bb24390aea463c0`. The relevant installed source files matched the pinned upstream sources. The generated vLLM build identifier was `g2cf0a6915`. This identifies the inspected source build; live verification remains a per-campaign requirement. diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_027.py b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_027.py new file mode 100644 index 000000000..fe2360728 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_027.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM 0.27.0 adapter; source mapping and limitations are in vllm-0.27.0.md.""" + +from dynamo.vllm.instrumented_scheduler import InstrumentedScheduler +from vllm.v1.worker.gpu_worker import Worker + +from .hooks import SchedulerObservation, WorkerObservation + + +class ObservedWorker(WorkerObservation, Worker): + observation_version = "0.27.0" + + +class ObservedInstrumentedScheduler(SchedulerObservation, InstrumentedScheduler): + observation_version = "0.27.0" diff --git a/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_028.py b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_028.py new file mode 100644 index 000000000..1baffcf92 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/instrumentation/vllm_028.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM 0.28.0 adapter; source mapping and limitations are in vllm-0.28.0.md.""" + +from dynamo.vllm.instrumented_scheduler import InstrumentedScheduler +from vllm.v1.worker.gpu_worker import Worker + +from .hooks import SchedulerObservation, WorkerObservation + + +class ObservedWorker(WorkerObservation, Worker): + observation_version = "0.28.0" + + +class ObservedInstrumentedScheduler(SchedulerObservation, InstrumentedScheduler): + observation_version = "0.28.0" diff --git a/python/aisimulate/collector/fpm_forward/runtime/preflight.py b/python/aisimulate/collector/fpm_forward/runtime/preflight.py index 7ca48c756..5238b59f8 100644 --- a/python/aisimulate/collector/fpm_forward/runtime/preflight.py +++ b/python/aisimulate/collector/fpm_forward/runtime/preflight.py @@ -40,13 +40,18 @@ def _write_audit(audit: dict) -> None: def main() -> None: + phase = os.environ.get("FPM_BENCHMARK_MODE") + real_prefill_requested = phase in {"prefill", "agg"} and os.environ.get( + "DYN_BENCH_PREFILL_REAL_SEED", "off" + ).lower() in {"on", "1", "true"} + dsv41_adapter = os.environ.get("DYN_FPM_DSV41_REAL_KV") == "1" # The audit artifact exists precisely to document a rejected image, so an # image whose runtime module is missing entirely (pre-PR11509) must still # produce it before this process fails the pod. try: from dynamo.vllm.instrumented_scheduler import BenchmarkPoint, InstrumentedScheduler - if os.environ.get("DYN_FPM_DSV41_REAL_KV") == "1": + if dsv41_adapter: from dsv41_scheduler import DeepseekV41RealKVScheduler if InstrumentedScheduler is not DeepseekV41RealKVScheduler: @@ -76,13 +81,21 @@ def main() -> None: fields = set(getattr(BenchmarkPoint, "__dataclass_fields__", {})) missing_fields = sorted(GRAPH_AWARE_FIELDS - fields) - missing_methods = sorted(name for name in GRAPH_AWARE_METHODS if not hasattr(InstrumentedScheduler, name)) + required_methods = set(GRAPH_AWARE_METHODS) + if real_prefill_requested and not dsv41_adapter: + # Native cached-prefill seeding is separate from decode warm-up: + # https://github.com/ai-dynamo/dynamo/blob/b83b1d9304ebfc624709ac46db32b1b6f1ff1615/components/src/dynamo/vllm/instrumented_scheduler.py#L3326 + # The source-checked V4.1 adapter owns its own real-KV prefill path. + required_methods.add("_bench_realseed_on") + missing_methods = sorted(name for name in required_methods if not hasattr(InstrumentedScheduler, name)) audit = { "schema_version": 1, "runtime_contract": "dynamo_pr11509_native_schema_v2_kvwarm_v1", "benchmark_point_fields": sorted(fields), "missing_fields": missing_fields, "missing_methods": missing_methods, + "prefill_real_seed_requested": real_prefill_requested, + "prefill_real_seed_implementation": "dsv41_adapter" if dsv41_adapter else "dynamo_native", "status": "passed" if not missing_fields and not missing_methods else "failed", } _write_audit(audit) diff --git a/python/aisimulate/collector/fpm_forward/runtime/vllm-0.27.0.json b/python/aisimulate/collector/fpm_forward/runtime/vllm-0.27.0.json new file mode 100644 index 000000000..35c26868a --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/vllm-0.27.0.json @@ -0,0 +1,41 @@ +{ + "schema_version": "aisimulate-runtime-instrumentation/v1", + "runtime": { + "framework": "vllm", + "version": "0.27.0", + "source_revision": "4bdc8a788d2e2ce9165d552b3d4d8b72604626bf", + "source_files": { + "vllm/v1/worker/gpu_worker.py": "50134ca3a147f5470e556437cbd40a13ea05b098efbac6553d5ea2f3e5f7ff57", + "vllm/v1/worker/gpu_model_runner.py": "993f2c926012241bcb8cc7f568c8e82968e4d81ca125641c831ecf501aac02ac", + "vllm/v1/core/kv_cache_utils.py": "0ab5385925f660395612e3f69030cad3b18c2e0bd6fb61d3ff43072a6cefb74f", + "vllm/v1/core/sched/scheduler.py": "c67bda2886b52865ddafabaae7d797c359e930752f374421a33e537d94a5f45a", + "vllm/v1/core/block_pool.py": "51cad2fd425128a0ff433ca4685acfc022e40659153ea5d7180fc586e1eebb6c", + "vllm/v1/kv_cache_interface.py": "bfbd62f02d900faf043ce0aeb46de0a10d7ce531351fced789805a6b817a7f80", + "vllm/version.py": "bd7a42776b630b611e2f5fa4698bc97e1f4f2583c35ccb246ec2e7bd43c83313", + "vllm/v1/worker/gpu/attn_utils.py": "5f889b6694eb83cdfa1f2d099e646e2c985d5213a68168edc821038de58946a0", + "vllm/v1/core/single_type_kv_cache_manager.py": "8c9030040cbc3b7176e228e10ff6491c3cf6c3e95bdfb573532fcf75f5385ebd", + "vllm/v1/attention/backends/flash_attn.py": "03017560e89dffc5453b6525bf76940e143b38939022ff5f74f872119fe12e48", + "vllm/v1/core/kv_cache_manager.py": "70f7f608c0963af155540630a5633483e5c19c0c0280dfd621760d5db2a419a6", + "vllm/v1/attention/backends/flashinfer.py": "758acb584792328cb20b05373ec8cab93aee76b48562baf1b40f8f4a55c5d054", + "vllm/v1/attention/backend.py": "14803680d6f0f3fcf04dd48def6b961c5dae70c061287610c8f0e1bbc82cb471", + "vllm/config/model.py": "d14c60ce1b2a7d9c47983ef0df0c10c4f8083ced5bc45935e8ababfe59f22736", + "vllm/transformers_utils/config.py": "17b57bfe69ba076cc5e01e3fefdc86e5023a1f85a759a99c16bef49569828dff", + "dynamo/vllm/instrumented_scheduler.py": "443255313134f32331c6b40cfc90fe72064f207e089c8d95867f8963c22e20a4", + "vllm/v1/worker/utils.py": "8bdd2c60ee6443092cc0ea60e62343fd1ae6ad9a94888925bd7a83d5b4ae110d", + "vllm/v1/worker/gpu/model_runner.py": "60fb531a045d6879da25f57cdb95584cc5272e4ef33b69e6b9c2ff3acbd78f8a" + } + }, + "files": [ + "fpm_memory_observer.py", + "instrumentation/__init__.py", + "instrumentation/observer.py", + "instrumentation/hooks.py", + "instrumentation/vllm_027.py", + "instrumentation/vllm-0.27.0.md", + "instrumentation/README.md" + ], + "worker_class": "instrumentation.vllm_027.ObservedWorker", + "scheduler_class": "instrumentation.vllm_027.ObservedInstrumentedScheduler", + "observation_schema": "aisimulate-runtime-observation/v1", + "source_notes": "instrumentation/vllm-0.27.0.md" +} diff --git a/python/aisimulate/collector/fpm_forward/runtime/vllm-0.28.0.example.json b/python/aisimulate/collector/fpm_forward/runtime/vllm-0.28.0.example.json new file mode 100644 index 000000000..434df428e --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/vllm-0.28.0.example.json @@ -0,0 +1,42 @@ +{ + "schema_version": "aisimulate-runtime-instrumentation/v1", + "runtime": { + "framework": "vllm", + "version": "0.28.0", + "source_revision": "2cf0a6915ce544dc493a0990f2ea38d81601128a", + "source_files": { + "vllm/v1/worker/gpu_worker.py": "5e8faf3e00649289c81c2917dd4557a80b0b351c33295d8e0c4b204a304e6552", + "vllm/v1/worker/gpu_model_runner.py": "7f2890eefca1efe25565bf1c7e5906a87948ae922610a7aaac620b28b46f26aa", + "vllm/v1/core/kv_cache_utils.py": "088f2201bee86fade694e78141b6e99a5cd0cdd23c5c7ab3526dd119f76e4aec", + "vllm/v1/core/sched/scheduler.py": "4b79ab760c3864af6838d555c01335deea3aac97c9be2be415fa0e5dab32f07e", + "vllm/v1/core/block_pool.py": "ddee56dccb2208411b3a035918e917ce8f56a9858471e9ca12b420d5d79bc69c", + "vllm/v1/kv_cache_interface.py": "932e9175796706506c34c3014cc3c69b146e199426dbbacf19f24056f7332ec3", + "vllm/version.py": "bd7a42776b630b611e2f5fa4698bc97e1f4f2583c35ccb246ec2e7bd43c83313", + "vllm/v1/worker/gpu/attn_utils.py": "af27f25caf705b34201d1938b81ad212a65ed9178b19e89a517811f390f16eab", + "vllm/v1/core/single_type_kv_cache_manager.py": "dc6d2e64cc4cd8cbb3f4bd77d2d035038658d75b6e52f9ace0cf637a4897f71f", + "vllm/v1/attention/backends/flash_attn.py": "44f4f0b560f0b09575d52a9e7a44f064bb600796079edd548aab2e21b4f31064", + "vllm/v1/core/kv_cache_manager.py": "2d20c3d98845cfd8d88a2f66b8fc6402ea1fcfbee16879d2364a4a9e45b8fa47", + "vllm/v1/attention/backends/flashinfer.py": "5ce66c9df9ffa6a3f5f1e480b99831ac2627a3a8086a4b9a797c58bdfd9ed18f", + "vllm/v1/attention/backend.py": "d77652795155fc743c3d8233b064cf855090d72aca485a48732b4fb3a1b736db", + "vllm/config/model.py": "17c687232886184f0390f38fc1c2c8ae078eaf24ebd1960a6b0c6a0669a35a98", + "vllm/transformers_utils/config.py": "869c46fc5c78aa33b86df168c56a6e9e52a1e22ed5becccabce153c8839c6a30", + "dynamo/vllm/instrumented_scheduler.py": "89547cf5ccf65a5ddd2dab015d8fc35759a472cd53bc441918981eef48749c28", + "vllm/_version.py": "63904896f9449a84bac95c17567eb1dea3592783ac3aa13916f5928b9458e9ac", + "vllm/v1/worker/utils.py": "e9d4381d903f7345c9b2fa7e9facfdfe7de7c36c2488bf7e28398183282b7a07", + "vllm/v1/worker/gpu/model_runner.py": "fb08762e3a025fdc55ab8d2605a96a57d1fad7efae0fbb605cc4fe2973028aca" + } + }, + "files": [ + "fpm_memory_observer.py", + "instrumentation/__init__.py", + "instrumentation/observer.py", + "instrumentation/hooks.py", + "instrumentation/vllm_028.py", + "instrumentation/vllm-0.28.0.md", + "instrumentation/README.md" + ], + "worker_class": "instrumentation.vllm_028.ObservedWorker", + "scheduler_class": "instrumentation.vllm_028.ObservedInstrumentedScheduler", + "observation_schema": "aisimulate-runtime-observation/v1", + "source_notes": "instrumentation/vllm-0.28.0.md" +} diff --git a/python/aisimulate/collector/fpm_forward/runtime_instrumentation.py b/python/aisimulate/collector/fpm_forward/runtime_instrumentation.py new file mode 100644 index 000000000..f820466c1 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime_instrumentation.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inspect and freeze campaign-local observer files without importing code.""" + +from __future__ import annotations + +import hashlib +import json +import re +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any + +import yaml + +INSTRUMENTATION_SCHEMA = "aisimulate-runtime-instrumentation/v1" +OBSERVATION_SCHEMA = "aisimulate-runtime-observation/v1" +INSTRUMENTATION_LOADER = "fpm_runtime_instrumentation.py" + + +def runtime_binding() -> dict[str, str]: + """Identify the collector loader without importing runtime or adapter code.""" + return { + "schema_version": "aisimulate-runtime-instrumentation-imports/v1", + "loader_sha256": sha256_bytes((Path(__file__).parent / "runtime" / INSTRUMENTATION_LOADER).read_bytes()), + } + + +def canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def validate_sha256(value: Any, label: str) -> str: + if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value): + raise ValueError(f"{label} requires a lowercase SHA-256") + return value + + +def relative_path(value: Any) -> str: + if ( + not isinstance(value, str) + or not value + or any(ord(char) < 32 for char in value) + or "\\" in value + or ":" in value + or PurePosixPath(value).is_absolute() + or any(part in {"", ".", ".."} for part in value.split("/")) + ): + raise ValueError(f"artifact paths must be normalized relative paths: {value!r}") + return value + + +def contained_file(root: Path, value: Any) -> Path: + name = relative_path(value) + current = root + for part in PurePosixPath(name).parts: + current = current / part + if current.is_symlink(): + raise ValueError(f"artifact path contains a symlink: {name}") + if not current.is_file() or not current.resolve().is_relative_to(root.resolve()): + raise ValueError(f"artifact is not a contained regular file: {name}") + return current + + +def _unique_pairs(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result = {} + for key, value in pairs: + if not isinstance(key, str) or key in result: + raise ValueError(f"duplicate or non-string object key: {key!r}") + result[key] = value + return result + + +class _UniqueSafeLoader(yaml.SafeLoader): + pass + + +def _yaml_mapping(loader: _UniqueSafeLoader, node: yaml.MappingNode) -> dict[str, Any]: + return _unique_pairs([(loader.construct_object(k), loader.construct_object(v)) for k, v in node.value]) + + +_UniqueSafeLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _yaml_mapping) + + +def read_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8"), object_pairs_hook=_unique_pairs) + canonical_json(value) + except (OSError, UnicodeError, TypeError, RecursionError, ValueError) as error: + raise ValueError(f"invalid JSON artifact {path}: {error}") from error + if not isinstance(value, dict): + raise ValueError(f"JSON artifact must be an object: {path}") + return value + + +@dataclass(frozen=True, slots=True) +class InstrumentationBundle: + root: Path + manifest_path: Path + _manifest_json: str + _files_json: str + sha256: str + + @property + def manifest(self) -> dict[str, Any]: + return json.loads(self._manifest_json) + + @property + def files(self) -> dict[str, str]: + return json.loads(self._files_json) + + +def load_instrumentation(path: str | Path, expected_version: str | None = None) -> InstrumentationBundle: + """Parse a safe manifest and hash its complete declared bundle, without imports.""" + path = Path(path).absolute() + if path.is_symlink(): + raise ValueError("instrumentation manifest cannot be a symlink") + try: + manifest = yaml.load(path.read_text(encoding="utf-8"), Loader=_UniqueSafeLoader) + canonical_json(manifest) + except (OSError, UnicodeError, yaml.YAMLError, TypeError, RecursionError, ValueError) as error: + raise ValueError(f"invalid instrumentation manifest: {error}") from error + required = {"schema_version", "runtime", "files", "worker_class", "scheduler_class", "observation_schema"} + if ( + not isinstance(manifest, dict) + or not required <= manifest.keys() + or manifest.keys() - required - {"source_notes"} + ): + raise ValueError("instrumentation manifest has missing or unknown fields") + if manifest["schema_version"] != INSTRUMENTATION_SCHEMA or manifest["observation_schema"] != OBSERVATION_SCHEMA: + raise ValueError("unsupported instrumentation or observation schema") + runtime = manifest["runtime"] + if ( + not isinstance(runtime, dict) + or not {"framework", "version", "source_revision"} <= runtime.keys() + or runtime.keys() - {"framework", "version", "source_revision", "source_files"} + or runtime["framework"] != "vllm" + or not isinstance(runtime["version"], str) + or not runtime["version"].strip() + or not isinstance(runtime["source_revision"], str) + or not re.fullmatch(r"[0-9a-f]{40}|[0-9a-f]{64}", runtime["source_revision"]) + ): + raise ValueError("instrumentation requires a pinned vLLM version and exact source revision") + if expected_version is not None and runtime["version"] != expected_version: + raise ValueError("instrumentation runtime version differs from the requested version") + if "source_files" in runtime: + if not isinstance(runtime["source_files"], dict) or not runtime["source_files"]: + raise ValueError("runtime source_files must contain installation-relative paths and hashes") + for name, digest in runtime["source_files"].items(): + relative_path(name) + validate_sha256(digest, "runtime source file") + names = manifest["files"] + if not isinstance(names, list) or not names: + raise ValueError("instrumentation requires declared files") + for name in names: + relative_path(name) + if len({name.casefold() for name in names}) != len(names) or "manifest.json" in names: + raise ValueError("instrumentation file paths must be unique and cannot use reserved manifest.json") + files = {name: sha256_bytes(contained_file(path.parent, name).read_bytes()) for name in names} + if INSTRUMENTATION_LOADER in files or any(name.startswith("fpm_runtime_instrumentation/") for name in files): + raise ValueError("instrumentation cannot replace the collector's runtime loader") + modules: dict[str, str] = {} + for name in names: + if not name.endswith(".py"): + continue + parts = list(PurePosixPath(name).with_suffix("").parts) + if parts[-1] == "__init__": + parts.pop() + module = ".".join(parts) + if not module or any(not part.isidentifier() for part in parts): + raise ValueError(f"instrumentation module path is ambiguous: {name}") + if module in modules: + raise ValueError(f"instrumentation module is ambiguous: {module}") + modules[module] = name + for field in ("worker_class", "scheduler_class"): + value = manifest[field] + parts = value.split(".") if isinstance(value, str) else [] + if len(parts) < 2 or any(not part.isidentifier() for part in parts) or ".".join(parts[:-1]) not in modules: + raise ValueError(f"{field} must reference an unambiguous declared Python module and class") + for depth in range(1, len(parts) - 1): + if "/".join(parts[:depth]) + "/__init__.py" not in files: + raise ValueError(f"{field} package must have a declared __init__.py") + if "source_notes" in manifest and ( + not isinstance(manifest["source_notes"], str) or manifest["source_notes"] not in files + ): + raise ValueError("source_notes must name a declared bundle file") + identity = sha256_bytes(canonical_json({"manifest": manifest, "files": files}).encode()) + return InstrumentationBundle(path.parent, path, canonical_json(manifest), canonical_json(files), identity) + + +def freeze_instrumentation(bundle: InstrumentationBundle, target: str | Path) -> InstrumentationBundle: + """Copy verified content to a fresh attempt; never rewrite an existing bundle.""" + target = Path(target).absolute() + if target.exists() or target.is_symlink(): + raise ValueError("instrumentation freezing requires a fresh directory") + if load_instrumentation(bundle.manifest_path).sha256 != bundle.sha256: + raise ValueError("instrumentation changed after it was loaded") + target.parent.mkdir(parents=True, exist_ok=True) + temporary = Path(tempfile.mkdtemp(prefix=".instrumentation-", dir=target.parent)) + try: + for name, digest in bundle.files.items(): + raw = contained_file(bundle.root, name).read_bytes() + if sha256_bytes(raw) != digest: + raise ValueError(f"instrumentation file changed while freezing: {name}") + destination = temporary / name + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(raw) + (temporary / "manifest.json").write_text(bundle._manifest_json + "\n", encoding="utf-8") + if load_instrumentation(temporary / "manifest.json").sha256 != bundle.sha256: + raise ValueError("instrumentation bundle changed while freezing") + temporary.rename(target) + finally: + if temporary.exists(): + shutil.rmtree(temporary) + return load_instrumentation(target / "manifest.json") diff --git a/python/aisimulate/collector/fpm_forward/runtime_memory.py b/python/aisimulate/collector/fpm_forward/runtime_memory.py new file mode 100644 index 000000000..9a6c94665 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime_memory.py @@ -0,0 +1,736 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate runtime initialization evidence and resolve physical cache capacity.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY + +from .model_capability import ResolvedModelConfig +from .native_artifact import COLLECTOR_PROVENANCE_FILENAME, validate_native_collection +from .planner import BackendPolicy, FPMCell, _canonical_hash +from .runtime.fpm_memory_observer import SCHEMA_NAME, SCHEMA_VERSION, SUPPORTED_VERSION +from .types import ParallelTopology + +_CELL_EXECUTION_FIELDS = frozenset({"execution_identity", "input_text_sha256"}) + + +def _positive(value: Any, label: str) -> int: + if type(value) is not int or not 0 < value <= 2**53: + raise ValueError(f"{label} must be a positive integer no greater than 2**53") + return value + + +def cell_from_dict(payload: dict[str, Any]) -> FPMCell: + """Read an exact saved cell without resolving current models or capabilities.""" + if not isinstance(payload, dict): + raise ValueError("saved FPM cell must be an object") + execution_fields = _CELL_EXECUTION_FIELDS.intersection(payload) + if execution_fields and execution_fields != _CELL_EXECUTION_FIELDS: + raise ValueError("saved FPM cell has an incomplete execution identity") + execution = LEGACY_EXECUTION_IDENTITY + input_text_sha256 = "" + if execution_fields: + declared = payload["execution_identity"] + input_text_sha256 = payload["input_text_sha256"] + if ( + not isinstance(declared, dict) + or set(declared) != set(EXECUTION_COLUMNS) + or any(not isinstance(value, str) for value in declared.values()) + or not isinstance(input_text_sha256, str) + ): + raise ValueError("saved FPM cell has an invalid execution identity") + execution = tuple(declared[name] for name in EXECUTION_COLUMNS) + try: + topology = ParallelTopology(**payload["topology"]) + for name, value in topology.to_dict().items(): + _positive(value, f"topology.{name}") + dtypes = payload["resolved_dtypes"] + policy = payload["backend_policy"] + cell = FPMCell( + cell_id=payload["cell_id"], + workload_kind=payload["workload_kind"], + topology=topology, + weight_quantization=payload["weight_quantization"], + kv_cache_dtype=payload["kv_cache_dtype"], + backend_policy=BackendPolicy(**policy), + parallel_strategy=payload["parallel_strategy"], + gemm_quant_mode=dtypes["gemm_quant_mode"], + moe_quant_mode=dtypes["moe_quant_mode"], + fmha_quant_mode=dtypes["fmha_quant_mode"], + comm_quant_mode=dtypes["comm_quant_mode"], + fmha_resolution=dtypes["fmha_resolution"], + execution_identity=execution, + input_text_sha256=input_text_sha256, + ) + except (KeyError, TypeError) as error: + raise ValueError(f"invalid saved FPM cell: {error}") from error + serialized = cell.to_dict() + if not execution_fields: + # Old cells retain their exact stored shape and identity when hashed. + for name in _CELL_EXECUTION_FIELDS: + serialized.pop(name) + if serialized != payload or cell.workload_kind not in {"prefill", "decode"}: + raise ValueError("saved FPM cell does not match the native collection contract") + if not isinstance(cell.cell_id, str) or not cell.cell_id: + raise ValueError("saved FPM cell requires a nonempty identity") + return cell + + +def validate_saved_plan(payload: dict[str, Any]) -> None: + """Verify current and historical producer hashes without rewriting identity.""" + if ( + not isinstance(payload, dict) + or payload.get("schema_name") != "aic_fpm_collection_plan" + or payload.get("schema_version") not in (10, 11) + ): + raise ValueError("runtime memory finalization requires a schema-v10 or schema-v11 collection plan") + try: + canonical = { + key: payload[key] + for key in ( + "backend", + "model_path", + "system", + "aic_revision", + "generator_config_sha256", + "options", + "capability", + "dtype_profile", + "cells", + ) + } + canonical["point_generation"] = "dynamo_native_self_benchmark" + canonical["policies"] = payload["backend_policies"] + canonical["topologies"] = [ + {name: item[name] for name in ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp")} + for item in payload["topologies"] + ] + admissions = copy.deepcopy(payload["topology_memory_admission"]) + for admission in admissions: + admission.pop("reason", None) + for estimate in admission.get("estimates", []): + estimate.pop("reason", None) + canonical["topology_memory_admission"] = admissions + if "fpm_profile" in payload: + canonical["fpm_profile"] = payload["fpm_profile"] + if "runtime_memory_policy" in payload: + canonical["runtime_memory_policy"] = payload["runtime_memory_policy"] + if "runtime_observation" in payload: + canonical["runtime_observation"] = payload["runtime_observation"] + for item in payload["cells"]: + present = _CELL_EXECUTION_FIELDS.intersection(item) + required = _CELL_EXECUTION_FIELDS if payload["schema_version"] == 11 else frozenset() + if present != required: + raise ValueError("saved collection plan cell execution identity does not match its schema version") + cells = [cell_from_dict(item) for item in payload["cells"]] + if not cells or len({cell.cell_id for cell in cells}) != len(cells): + raise ValueError("saved collection plan requires unique cells") + except (KeyError, TypeError, AttributeError) as error: + raise ValueError(f"invalid saved collection plan: {error}") from error + if _canonical_hash(canonical) != payload.get("sha256"): + raise ValueError("saved collection plan SHA-256 does not match its immutable inputs") + + +@dataclass(frozen=True, slots=True) +class _SavedOptions: + warmup_iterations: int + benchmark_points_json: str | None + + +@dataclass(frozen=True, slots=True) +class _SavedCapability: + support_level: str + template_id: str | None + template_version: int | None + aic_database_version: str + is_moe: bool | None + attention_source: str | None + model_config: ResolvedModelConfig | None + + +@dataclass(frozen=True, slots=True) +class SavedCollectionIdentity: + """The verified plan fields used by native aggregation and commit validation.""" + + backend: str + model_path: str + system: str + sha256: str + generator_config_sha256: str + options: _SavedOptions + capability: _SavedCapability + cells: tuple[FPMCell, ...] + + +def saved_plan_identity(payload: dict[str, Any]) -> SavedCollectionIdentity: + validate_saved_plan(payload) + capability = payload["capability"] + config = capability.get("model_config") + model_config = None + if isinstance(config, dict) and isinstance(config.get("payload"), dict) and config["payload"]: + model_config = ResolvedModelConfig( + config["payload"], source_kind=config.get("source_kind", "saved"), source_reference="saved collection plan" + ) + points = payload["options"].get("benchmark_points") + return SavedCollectionIdentity( + backend=payload["backend"], + model_path=payload["model_path"], + system=payload["system"], + sha256=payload["sha256"], + generator_config_sha256=payload["generator_config_sha256"], + options=_SavedOptions( + warmup_iterations=payload["options"]["global_warmup_iterations"], + benchmark_points_json=json.dumps(points["payload"]) if points is not None else None, + ), + cells=tuple(cell_from_dict(item) for item in payload["cells"]), + capability=_SavedCapability( + **{ + key: capability[key] + for key in ( + "support_level", + "template_id", + "template_version", + "aic_database_version", + ) + }, + is_moe=capability.get("is_moe") if type(capability.get("is_moe")) is bool else None, + attention_source=capability.get("attention_source"), + model_config=model_config, + ), + ) + + +def _dtype(value: Any) -> str: + return { + "torch.bfloat16": "bfloat16", + "torch.float16": "half", + "float16": "half", + "fp8_e4m3": "fp8", + "fp8_e4m3fn": "fp8", + }.get(value, value) + + +def _resolved_cache_dtype(config: dict[str, Any]) -> str: + value = config["cache_config"].get("cache_dtype") + if value == "auto": + value = config["model_config"].get("dtype") + return _dtype(value) + + +def _validate_precision(config: dict[str, Any], cell: FPMCell) -> None: + """Check audited quantization families without assigning every layer a dtype. + + vLLM 4bdc8a788d2e2ce9165d552b3d4d8b72604626bf: + model_executor/layers/quantization/{fp8,modelopt}.py config classes and + ModelOptQuantConfigBase.get_quant_method preserve unquantized exclusions. + """ + model = config["model_config"] + base = _dtype(model.get("dtype")) + method = model.get("quantization") + quant = config.get("quantization_config") + if "quantization" not in model or "quantization_config" not in config or base not in {"bfloat16", "half"}: + raise ValueError("runtime memory weight precision evidence is incomplete or unsupported") + gemm, moe = None, None + if method is None and quant is None: + gemm = moe = base + elif isinstance(quant, dict): + prefix = "vllm.model_executor.layers.quantization." + if method == "modelopt_fp4" and quant.get("type") == prefix + "modelopt.ModelOptNvFp4Config": + if quant.get("quant_method") == "NVFP4": + gemm = moe = "nvfp4" + elif method == "fp8" and quant.get("type") == prefix + "fp8.Fp8Config": + block = quant.get("weight_block_size") + scheme = quant.get("activation_scheme") + if block == [128, 128] and scheme == "dynamic": + gemm = moe = "fp8_block" + elif block is None and scheme in {"static", "dynamic"}: + gemm, moe = ("fp8_static" if scheme == "static" else "fp8"), "fp8" + elif ( + method == "modelopt" + and quant.get("type") == prefix + "modelopt.ModelOptFp8Config" + and quant.get("quant_method") == "FP8" + ): + gemm, moe = "fp8_static", "fp8" + if gemm is None: + raise ValueError("runtime memory weight precision mapping is unsupported or contradicts quantization evidence") + # A mixed checkpoint may leave attention/shared experts in the base dtype. + # The global quantization method establishes a family, not every layer's mode. + if cell.gemm_quant_mode not in {base, gemm} or cell.moe_quant_mode not in {None, base, moe}: + raise ValueError("runtime memory weight precision differs from the planned GEMM/MoE identity") + + +def _validate_graph( + config: dict[str, Any], *, eager: Any, tokens: int, allow_unused_captures: bool = False +) -> dict[str, Any]: + # vLLM 0.27.0 config/vllm.py:_set_cudagraph_sizes and + # config/compilation.py:resolve_cudagraph_mode_and_sizes. A worker may retain + # configured sizes after disabling graphs for its initialized attention backend. + fields = {"mode", "cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"} + if not isinstance(config, dict) or not fields.issubset(config) or type(eager) is not bool: + raise ValueError("runtime memory evidence lacks complete graph configuration") + result = copy.deepcopy(config) + modes = ("NONE", "STOCK_TORCH_COMPILE", "DYNAMO_TRACE_ONCE", "VLLM_COMPILE") + mode = result["mode"] + if type(mode) is int and 0 <= mode < len(modes): + mode = modes[mode] + elif isinstance(mode, str) and mode in {"0", "1", "2", "3"}: + mode = modes[int(mode)] + if not isinstance(mode, str) or mode not in modes: + raise ValueError("runtime memory graph compilation mode is unsupported") + result["mode"] = mode + graph = result["cudagraph_mode"] + sizes, maximum = result["cudagraph_capture_sizes"], result["max_cudagraph_capture_size"] + if ( + not isinstance(graph, str) + or graph not in {"NONE", "PIECEWISE", "FULL", "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"} + or not isinstance(sizes, list) + or any(type(size) is not int or not 0 < size <= tokens for size in sizes) + or sizes != sorted(set(sizes)) + or type(maximum) is not int + or maximum != max(sizes, default=0) + or (graph != "NONE" and not sizes) + or (graph == "NONE" and sizes and not allow_unused_captures) + or (eager and (mode != "NONE" or graph != "NONE" or sizes)) + ): + raise ValueError("runtime memory graph configuration is incomplete or inconsistent") + return result + + +def _validate_no_offload(payload: dict[str, Any], config: dict[str, Any]) -> None: + offload = config.get("offload_config") + if ( + not isinstance(offload, dict) + or not isinstance(offload.get("uva"), dict) + or not isinstance(offload.get("prefetch"), dict) + ): + raise ValueError("runtime memory evidence lacks offload configuration") + uva, prefetch = offload["uva"], offload["prefetch"] + amount = uva.get("cpu_offload_gb") + if ( + offload.get("offload_backend") != "auto" + or type(amount) not in {int, float} + or amount != 0 + or type(prefetch.get("offload_group_size")) is not int + or prefetch["offload_group_size"] != 0 + or type(prefetch.get("offload_num_in_group")) is not int + or prefetch["offload_num_in_group"] < 1 + or type(prefetch.get("offload_prefetch_step")) is not int + or prefetch["offload_prefetch_step"] < 0 + or config["cache_config"].get("kv_cache_memory_bytes") is not None + or config["cache_config"].get("num_gpu_blocks_override") is not None + or ( + payload.get("kind") == "worker" + and payload.get("effective_offloader") != "vllm.model_executor.offloader.base.NoopOffloader" + ) + ): + raise ValueError("runtime memory requires the supported automatic HBM-only no-offload policy") + + +def validate_requested_graph(payload: dict[str, Any], config: dict[str, Any], expected: dict[str, Any]) -> None: + """Compare accepted engine settings before independently checked backend downgrades.""" + if not expected: + return + initial = _validate_graph( + payload.get("initial_compilation_config", config["compilation_config"]) + if payload.get("kind") == "worker" + else config["compilation_config"], + eager=config["model_config"]["enforce_eager"], + tokens=config["scheduler_config"]["max_num_batched_tokens"], + ) + for name, value in expected.items(): + if initial.get(name) != value: + raise ValueError(f"initialized {name} differs from the reviewed serving role") + + +def _validate_scheduler_config(worker: dict[str, Any], scheduler: dict[str, Any]) -> None: + actual, scheduled = (copy.deepcopy(item["resolved_config"]) for item in (worker, scheduler)) + for config, is_worker in ((actual, True), (scheduled, False)): + config["compilation_config"] = _validate_graph( + config["compilation_config"], + eager=config["model_config"]["enforce_eager"], + tokens=config["scheduler_config"]["max_num_batched_tokens"], + allow_unused_captures=is_worker and "initial_compilation_config" in worker, + ) + if actual == scheduled: + return + initial = worker.get("initial_compilation_config") + if not isinstance(initial, dict): + raise ValueError("runtime memory launch configurations differ between scheduler and worker") + initial = _validate_graph( + initial, + eager=actual["model_config"]["enforce_eager"], + tokens=actual["scheduler_config"]["max_num_batched_tokens"], + ) + original = {**actual, "compilation_config": initial} + graph = actual["compilation_config"] + # Only backend-dependent graph downgrades in the audited non-speculative + # MRV1 resolver are process-local. Compilation/capture sizes remain unchanged. + fallbacks = { + "FULL": {"FULL_AND_PIECEWISE", "FULL_DECODE_ONLY"}, + "FULL_AND_PIECEWISE": {"PIECEWISE", "NONE"}, + "FULL_DECODE_ONLY": {"PIECEWISE", "NONE"}, + } + if ( + original != scheduled + or graph["cudagraph_mode"] not in fallbacks.get(initial["cudagraph_mode"], set()) + or {**graph, "cudagraph_mode": initial["cudagraph_mode"]} != initial + ): + raise ValueError("runtime memory launch configurations differ between scheduler and worker") + + +def _validate_config( + payload: dict[str, Any], cell: FPMCell, *, context: int, tokens: int, sequences: int, fraction: float +) -> dict[str, Any]: + config = copy.deepcopy(payload.get("resolved_config")) + if not isinstance(config, dict): + raise ValueError("runtime memory evidence lacks resolved launch configuration") + expected = { + "model_config": {"max_model_len": context}, + "cache_config": {"gpu_memory_utilization": fraction}, + "scheduler_config": {"max_num_batched_tokens": tokens, "max_num_seqs": sequences, "async_scheduling": False}, + "parallel_config": { + "tensor_parallel_size": cell.topology.tp, + "pipeline_parallel_size": cell.topology.pp, + "data_parallel_size": cell.topology.dp, + "decode_context_parallel_size": cell.topology.cp, + }, + } + for section, fields in expected.items(): + actual = config.get(section) + if not isinstance(actual, dict): + raise ValueError(f"runtime memory evidence lacks {section}") + for key, value in fields.items(): + if type(actual.get(key)) is not type(value) or actual.get(key) != value: + raise ValueError( + f"runtime memory launch mismatch: {section}.{key}={actual.get(key)!r}, expected {value!r}" + ) + parallel = config["parallel_config"] + if parallel.get("prefill_context_parallel_size", 1) != 1: + raise ValueError("runtime memory with prefill context parallelism is unsupported") + if parallel.get("enable_expert_parallel") != (cell.topology.moe_ep > 1): + raise ValueError("runtime memory expert parallel setting differs from the planned topology") + if parallel.get("enable_eplb", False) is not False: + raise ValueError("runtime memory with expert load balancing is unsupported") + expected_backend = cell.backend_policy.aic_fields.get("moe_backend") + if ( + expected_backend not in (None, "auto") + and (config.get("kernel_config") or {}).get("moe_backend") != expected_backend + ): + raise ValueError("runtime memory MoE backend differs from the planned identity") + if _resolved_cache_dtype(config) != cell.kv_cache_dtype: + raise ValueError("runtime memory KV precision differs from the planned identity") + _validate_precision(config, cell) + _validate_no_offload(payload, config) + config["compilation_config"] = _validate_graph( + config.get("compilation_config"), + eager=config["model_config"].get("enforce_eager"), + tokens=tokens, + allow_unused_captures=payload.get("kind") == "worker" and "initial_compilation_config" in payload, + ) + return config + + +def _validate_revision(config: dict[str, Any], expected: str | None) -> None: + # A local PVC can omit revision or retain a symbolic ref. Only explicit + # immutable commits establish contradictions; this is not a weight hash. + if expected is None or re.fullmatch(r"[0-9a-fA-F]{40}", expected) is None: + return + model = config["model_config"] + for key in ("revision", "loaded_config_commit_hash"): + observed = model.get(key) + if ( + isinstance(observed, str) + and re.fullmatch(r"[0-9a-fA-F]{40}", observed) + and observed.lower() != expected.lower() + ): + raise ValueError(f"runtime memory model {key} contradicts the reviewed immutable revision") + + +def _normalized_groups(cache: dict[str, Any], scheduler_cache: dict[str, Any], page_bytes: int) -> list[dict[str, Any]]: + raw_groups = cache.get("groups") + scheduler_groups = scheduler_cache.get("groups") + if not isinstance(raw_groups, list) or not raw_groups or not isinstance(scheduler_groups, list): + raise ValueError("runtime memory evidence lacks cache groups") + if len(raw_groups) != len(scheduler_groups): + raise ValueError("worker and scheduler cache group counts disagree") + groups = [] + layer_names: set[str] = set() + for index, (group, scheduled) in enumerate(zip(raw_groups, scheduler_groups, strict=True)): + if not isinstance(group, dict) or not isinstance(scheduled, dict): + raise ValueError("runtime memory cache groups must be objects") + if {key: value for key, value in group.items() if key not in {"kind", "layer_classes"}} != scheduled: + raise ValueError("worker and scheduler cache group geometry disagrees") + names = group.get("layer_names") + if ( + not isinstance(names, list) + or not names + or any(not isinstance(name, str) or not name for name in names) + or len(set(names)) != len(names) + or layer_names.intersection(names) + ): + raise ValueError("runtime memory cache layer names must be nonempty and unique") + layer_names.update(names) + spec = group.get("spec_type") + if spec not in {"vllm.v1.kv_cache_interface.FullAttentionSpec", "vllm.v1.kv_cache_interface.SlidingWindowSpec"}: + raise ValueError(f"unsupported runtime cache spec: {spec!r}") + window = group.get("sliding_window") + if spec.endswith(".SlidingWindowSpec"): + _positive(window, "sliding window") + elif window is not None: + raise ValueError("full-attention cache cannot declare sliding retention") + kind = group.get("kind") + if kind not in {"attention", "convolution"} or (kind == "convolution" and window != 4): + raise ValueError("runtime memory cache kind is unsupported") + classes = group.get("layer_classes") + if ( + not isinstance(classes, dict) + or set(classes) != set(names) + or any(not isinstance(value, str) or not value for value in classes.values()) + ): + raise ValueError("runtime memory cache lacks complete layer class evidence") + convolution_class = "vllm.models.inkling.nvidia.sconv_swa_attn.InklingConvState" + if any((value == convolution_class) != (kind == "convolution") for value in classes.values()): + raise ValueError("runtime memory cache kind contradicts actual layer classes") + groups.append( + { + "name": f"runtime_group_{index}", + "kind": kind, + "num_layers": len(names), + "block_size_tokens": _positive(group.get("block_size_tokens"), "cache block size"), + "page_size_bytes": page_bytes, + "sliding_window": window, + } + ) + return groups + + +def validate_shared_pool_accounting(cache: dict[str, Any], scheduled: dict[str, Any]) -> dict[str, int]: + """Recompute the physical cost and usable count of one shared HBM block pool. + + Used by strict historical finalization and the independent observation + importer. Runtime-specific view/source/retention checks remain with callers. + """ + count = _positive(cache.get("num_blocks"), "worker cache block count") + free = _positive(scheduled.get("initial_free_blocks"), "initial free block count") + if ( + type(scheduled.get("num_blocks")) is not int + or scheduled["num_blocks"] != count + or scheduled.get("pool_count") != 1 + or type(scheduled.get("pool_count")) is not int + or scheduled.get("watermark_blocks") != 0 + or type(scheduled.get("watermark_blocks")) is not int + or free >= count + or scheduled.get("reserved_blocks") != count - free + or type(scheduled.get("reserved_blocks")) is not int + or type(scheduled.get("null_block_id")) is not int + or not 0 <= scheduled["null_block_id"] < count + ): + raise ValueError("scheduler block pool capacity/reservations are unsupported or inconsistent") + allocated = _positive(cache.get("allocated_cache_bytes"), "physical cache bytes") + available = _positive(cache.get("available_cache_bytes"), "profiled cache budget") + storages = cache.get("storages") + if not isinstance(storages, list) or not storages: + raise ValueError("runtime memory evidence lacks physical storage accounting") + keys = [(item.get("device"), item.get("pointer")) for item in storages] + if ( + len(keys) != len(set(keys)) + or len({device for device, _pointer in keys}) != 1 + or any( + not isinstance(device, str) or not device.startswith("cuda:") or type(pointer) is not int or pointer <= 0 + for device, pointer in keys + ) + or sum(_positive(item.get("size_bytes"), "storage bytes") for item in storages) != allocated + ): + raise ValueError("physical runtime cache storage accounting is inconsistent") + if allocated > available or allocated % count: + raise ValueError("physical cache allocation exceeds its budget or has a non-integral pool page size") + return { + "num_blocks": count, + "initial_free_blocks": free, + "allocated_cache_bytes": allocated, + "available_cache_bytes": available, + "page_size_bytes": allocated // count, + } + + +def resolve_runtime_resources( + cell: FPMCell, + raw_root: Path, + *, + expected_plan_sha256: str, + expected_attempt_id: str, + expected_backend_version: str, + expected_context_length: int, + expected_max_num_tokens: int, + expected_max_batch_size: int, + expected_gpu_memory_utilization: float, + expected_model_revision: str | None = None, + expected_compilation_config: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Require complete matching native timings and initialized rank/pool evidence. + + Every group draws pages from one scheduler block pool. Therefore each page + is charged the complete physical cost of one pool block, including layer + padding and shared buffers. Free pool blocks exclude permanent reservations. + No legacy activation/overhead decomposition is inferred from cache capacity. + """ + for name, value in ( + ("context", expected_context_length), + ("tokens", expected_max_num_tokens), + ("sequences", expected_max_batch_size), + ): + _positive(value, name) + fraction = expected_gpu_memory_utilization + if ( + isinstance(fraction, bool) + or not isinstance(fraction, (int, float)) + or not math.isfinite(fraction) + or not 0 < fraction <= 1 + ): + raise ValueError("runtime memory GPU memory utilization must be in (0, 1]") + if not expected_plan_sha256 or not expected_attempt_id: + raise ValueError("runtime memory resolution requires exact plan and attempt identities") + if expected_backend_version != SUPPORTED_VERSION: + raise ValueError(f"runtime memory conversion currently supports vLLM {SUPPORTED_VERSION}") + collection = validate_native_collection( + cell, raw_root, expected_plan_sha256=expected_plan_sha256, expected_attempt_id=expected_attempt_id + ) + if collection.backend_version != expected_backend_version: + raise ValueError("native timing runtime version differs from the memory profile") + if cell.topology.pp != 1 or cell.topology.cp != 1: + raise ValueError("runtime memory supports PP1 and CP1") + workers: dict[tuple[int, int], dict[str, Any]] = {} + schedulers: dict[int, dict[str, Any]] = {} + files = [] + canonical_config = None + for path in sorted(raw_root.glob("**/fpm-memory-*.json")): + if len(path.relative_to(raw_root).parts) != 2: + raise ValueError(f"runtime memory artifact must belong to one collected pod: {path}") + payload = json.loads(path.read_text()) + if ( + not isinstance(payload, dict) + or payload.get("schema_name") != SCHEMA_NAME + or payload.get("schema_version") != SCHEMA_VERSION + ): + raise ValueError(f"unsupported runtime memory evidence schema: {path}") + provenance_path = path.parent / COLLECTOR_PROVENANCE_FILENAME + provenance = json.loads(provenance_path.read_text()) + expected_provenance = { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell.cell_id, + "plan_sha256": expected_plan_sha256, + "attempt_id": expected_attempt_id, + "runtime": {"backend": "vllm", "backend_version": expected_backend_version}, + } + if provenance != expected_provenance or payload.get("collector_provenance") != provenance: + raise ValueError(f"runtime memory evidence belongs to a different plan, attempt, cell, or runtime: {path}") + if payload.get("backend_version") != expected_backend_version: + raise ValueError(f"runtime memory backend version mismatch: {path}") + if payload.get("status") != "resolved": + raise ValueError(f"runtime memory remains unresolved: {path}: {payload.get('error')}") + config = _validate_config( + payload, + cell, + context=expected_context_length, + tokens=expected_max_num_tokens, + sequences=expected_max_batch_size, + fraction=float(fraction), + ) + _validate_revision(config, expected_model_revision) + validate_requested_graph(payload, config, expected_compilation_config or {}) + dp = payload.get("dp_rank") + if type(dp) is not int or not 0 <= dp < cell.topology.dp: + raise ValueError(f"runtime memory has an invalid DP rank: {path}") + if payload.get("kind") == "worker": + if canonical_config is None: + canonical_config = config + elif config != canonical_config: + raise ValueError("runtime memory launch configurations differ across worker ranks") + tp, pp = payload.get("tp_rank"), payload.get("pp_rank") + if type(tp) is not int or not 0 <= tp < cell.topology.tp or type(pp) is not int or pp != 0: + raise ValueError(f"runtime memory has an invalid TP/PP rank: {path}") + key = (dp, tp) + if key in workers: + raise ValueError(f"duplicate runtime memory worker rank: {key}") + workers[key] = payload + elif payload.get("kind") == "scheduler": + if dp in schedulers or payload.get("tp_rank") is not None or payload.get("pp_rank") is not None: + raise ValueError(f"duplicate or invalid runtime memory scheduler rank: {dp}") + schedulers[dp] = payload + else: + raise ValueError(f"unknown runtime memory observation kind: {path}") + files.append( + { + "path": str(path.relative_to(raw_root)), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "evidence": payload, + } + ) + if set(workers) != {(dp, tp) for dp in range(cell.topology.dp) for tp in range(cell.topology.tp)}: + raise ValueError("runtime memory worker rank evidence is incomplete") + if set(schedulers) != set(range(cell.topology.dp)): + raise ValueError("runtime memory scheduler rank evidence is incomplete") + canonical_groups = None + canonical_geometry = None + capacities = [] + for (dp, tp), payload in sorted(workers.items()): + _validate_scheduler_config(payload, schedulers[dp]) + cache = payload.get("cache") + scheduled = schedulers[dp].get("cache") + if not isinstance(cache, dict) or not isinstance(scheduled, dict): + raise ValueError("runtime memory evidence lacks cache allocation") + accounting = validate_shared_pool_accounting(cache, scheduled) + count = accounting["num_blocks"] + free = accounting["initial_free_blocks"] + allocated = accounting["allocated_cache_bytes"] + groups = _normalized_groups(cache, scheduled, allocated // count) + geometry = [{key: value for key, value in group.items() if key != "layer_classes"} for group in cache["groups"]] + if canonical_groups is None: + canonical_groups, canonical_geometry = groups, geometry + elif groups != canonical_groups or geometry != canonical_geometry: + raise ValueError("cache group geometry or physical page cost differs across ranks") + capacities.append({"dp_rank": dp, "tp_rank": tp, "kv_cache_bytes": free * (allocated // count)}) + runtime_settings = copy.deepcopy(canonical_config) + runtime_settings["cache_config"].pop("enable_prefix_caching", None) + provenance = { + "source": "vllm_initialized_cache", + "observer_schema_version": SCHEMA_VERSION, + "plan_sha256": expected_plan_sha256, + "attempt_id": expected_attempt_id, + "cell_id": cell.cell_id, + "runtime_run_id": collection.runtime_run_id, + "runtime_grid_digest": collection.runtime_grid_digest, + "backend_version": expected_backend_version, + "resolved_config": canonical_config, + "runtime_settings": runtime_settings, + "rank_capacities": capacities, + "artifacts": files, + "scope": "physical memory allocated by the complete collection worker; " + "text-only timing does not subtract other loaded components", + } + return { + "max_num_tokens": expected_max_num_tokens, + "max_batch_size": expected_max_batch_size, + "cache_layout": "grouped", + "cache_groups": canonical_groups, + "runtime_memory": { + "kv_cache_bytes": min(item["kv_cache_bytes"] for item in capacities), + "gpu_memory_utilization": float(fraction), + "max_model_len": expected_context_length, + "provenance": json.dumps(provenance, sort_keys=True, separators=(",", ":")), + }, + "provenance": "Observed initialized vLLM shared cache pool after successful warmup; " + "minimum usable capacity across all ranks, with padding and permanent reservations included.", + } diff --git a/python/aisimulate/collector/fpm_forward/runtime_observations.py b/python/aisimulate/collector/fpm_forward/runtime_observations.py new file mode 100644 index 000000000..3bc830a4a --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime_observations.py @@ -0,0 +1,787 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only validation of saved runtime observations, independent of timing data. + +The supported allocation representation is one HBM pool shared by full-history +or sliding-window groups. Runtime mappings are evidence: the importer rechecks +raw geometry, configuration, physical views, and pool accounting itself. +""" + +from __future__ import annotations + +import copy +import math +import re +from importlib.resources import files +from pathlib import Path, PurePosixPath +from typing import Any + +from aisimulate_core.fpm_profile import FpmResourceProfile +from aisimulate_core.sdk.perf_database import load_system_spec + +from . import runtime_memory +from .config import PrefillSamplingProfile, collection_phases +from .planner import BackendPolicy, FPMCell +from .runtime_instrumentation import ( + OBSERVATION_SCHEMA, + InstrumentationBundle, + canonical_json, + contained_file, + load_instrumentation, + read_json, + sha256_bytes, + validate_sha256, +) +from .types import ParallelTopology + +INDEX_SCHEMA = "aisimulate-runtime-observations/v1" +LAUNCH_SCHEMA = "aisimulate-runtime-probe-launch/v1" +_SEMANTICS = { + "allocation": "shared_block_pool", + "storage": "hbm", + "retention": "full_or_window", + "prefix_reuse": False, + "offload": False, + "speculative": False, +} + + +def _same(actual: Any, expected: Any, label: str) -> None: + if canonical_json(actual) != canonical_json(expected): + raise ValueError(f"runtime observation {label} mismatch") + + +def _text(value: Any, label: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"runtime observation requires {label}") + return value + + +def _nonnegative(value: Any, label: str) -> int: + if type(value) is not int or not 0 <= value <= 2**53: + raise ValueError(f"{label} must be a nonnegative integer no greater than 2**53") + return value + + +def _positive(value: Any, label: str) -> int: + return runtime_memory._positive(value, label) + + +def expected_ranks(topology: dict[str, int]) -> dict[str, list[dict[str, int]]]: + """Derive rank coverage from explicit topology, never configuration labels.""" + if not isinstance(topology, dict) or set(topology) != {"tp", "pp", "dp", "moe_tp", "moe_ep", "cp"}: + raise ValueError("runtime observations require complete explicit topology") + for key, value in topology.items(): + _positive(value, f"topology.{key}") + if topology["pp"] != 1 or topology["cp"] != 1: + raise ValueError("runtime observation allocation supports PP1 and CP1") + return { + "workers": [ + {"dp_rank": dp, "tp_rank": tp, "pp_rank": 0} for dp in range(topology["dp"]) for tp in range(topology["tp"]) + ], + "schedulers": [{"dp_rank": dp} for dp in range(topology["dp"])], + } + + +def _launch_cell(launch: dict[str, Any], phase: str) -> FPMCell: + identity, collection, precision = (launch[key] for key in ("identity", "collection", "precision")) + topology = launch["topology"] + kind = identity.get("model_kind") + if kind == "moe": + if topology["moe_tp"] * topology["moe_ep"] != topology["tp"] * topology["dp"]: + raise ValueError("runtime observation attention and MoE topology sizes disagree") + elif kind != "dense" or any(topology[name] != 1 for name in ("dp", "moe_tp", "moe_ep")): + raise ValueError("dense runtime observations require DP1 and MoE TP1/EP1") + if identity.get("framework") != "vllm": + raise ValueError("runtime observation currently requires a vLLM source mapping") + for field in ("model", "model_revision", "framework_version", "gpu", "interconnect"): + _text(identity.get(field), f"identity.{field}") + for field in ("max_model_len", "max_num_batched_tokens", "max_num_seqs"): + _positive(collection.get(field), f"collection.{field}") + fraction = collection.get("gpu_memory_utilization") + if type(fraction) not in (int, float) or not math.isfinite(fraction) or not 0 < fraction <= 1: + raise ValueError("GPU memory utilization must be finite and in (0, 1]") + policy, size = collection.get("prefill_cudagraph_policy"), collection.get("max_prefill_cudagraph_size") + if policy not in {"runtime", "explicit"} or (policy == "runtime" and size is not None): + raise ValueError("runtime observation has an inconsistent prefill graph policy") + if policy == "explicit": + _positive(size, "explicit prefill graph size") + if collection.get("async_scheduling", False) is not False: + raise ValueError("runtime observation requires synchronous scheduling") + if precision.get("enable_eplb", False) is not False or precision.get("enable_wideep", False) is not False: + raise ValueError("runtime observation does not support EPLB or wide-EP allocations") + if precision.get("comm_quant_mode") != "half": + raise ValueError("runtime observation requires the collector half communication identity") + for field in ("gemm_quant_mode", "fmha_quant_mode", "kvcache_quant_mode"): + _text(precision.get(field), f"precision.{field}") + validate_sha256(launch["model_config"]["sha256"], "model configuration") + _text(launch["deployment"].get("image"), "exact runtime image") + return FPMCell( + cell_id="runtime-observation", + workload_kind=phase, + topology=ParallelTopology(**launch["topology"]), + weight_quantization=precision["gemm_quant_mode"], + kv_cache_dtype=precision["kvcache_quant_mode"], + backend_policy=BackendPolicy("runtime-observation", {}, {}, aic_fields=precision), + gemm_quant_mode=precision["gemm_quant_mode"], + moe_quant_mode=precision.get("moe_quant_mode"), + fmha_quant_mode=precision["fmha_quant_mode"], + comm_quant_mode=precision["comm_quant_mode"], + ) + + +def _artifact(root: Path, reference: dict[str, Any]) -> tuple[Path, str]: + path = contained_file(root, reference["path"]) + digest = validate_sha256(reference["sha256"], "runtime artifact") + if sha256_bytes(path.read_bytes()) != digest: + raise ValueError(f"runtime artifact hash mismatch: {reference['path']}") + return path, digest + + +def _validate_hardware(hardware: dict[str, Any], identity: dict[str, Any]) -> None: + system = identity["gpu"] + family = system.split("_", 1)[0] + if system == "rtx_pro_6000_server": + family = "rtx pro 6000" + elif not re.fullmatch(r"(?:a100|h100|h200|b200|b300|gh200|gb200|gb300|l40s|l4)(?:_sxm|_pcie)?", system): + raise ValueError("runtime hardware lacks a supported device-family mapping") + spec = load_system_spec(system, systems_paths=str(files("aisimulate_core") / "systems")) + expected_sm = spec.get("gpu", {}).get("sm_version") + actual_name = _text(hardware.get("device_name"), "hardware device name") + if ( + hardware.get("gpu") != system + or hardware.get("sm") != expected_sm + or not re.search(r"(? str: + deployment = launch["deployment"] + model_cache = deployment.get("model_cache") + if not model_cache: + return launch["identity"]["model"] + if deployment.get("executor") != "kubernetes": + raise ValueError("model_cache checkpoint mapping requires the Kubernetes executor") + parts = _text(model_cache, "deployment.model_cache").split(":") + if len(parts) > 3 or not parts[0]: + raise ValueError("model_cache must be NAME[:MOUNT[:SUBPATH]] with a non-empty PVC name") + # Match entry._load_generator_overrides and runner._cell_generator_overrides: + # a cache-only PVC preserves the public model; mount + subpath select the + # actual --model path while the served model name keeps its public identity. + mount = parts[1] if len(parts) > 1 else "" + subpath = parts[2] if len(parts) > 2 else "" + if not mount and not subpath: + return launch["identity"]["model"] + if not mount or not subpath: + raise ValueError("model_cache checkpoint mapping requires both mount and subpath") + relative = PurePosixPath(subpath) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("model_cache checkpoint subpath must be relative without '..'") + return str(PurePosixPath(mount) / relative) + + +def _matches_model_path(observed: Any, launch: dict[str, Any]) -> bool: + expected = _deployed_model_path(launch) + if observed == expected: + return True + revision = launch["identity"]["model_revision"] + # vLLM 0.27.0 engine/arg_utils.py:EngineArgs.__post_init__ resolves HF + # identifiers to local snapshots in offline mode. Bind that spelling to + # the exact repo and immutable revision; never alias explicit local paths. + if ( + not isinstance(observed, str) + or expected != launch["identity"]["model"] + or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*(?:/[A-Za-z0-9][A-Za-z0-9_.-]*)?", expected) is None + or "--" in expected + or ".." in expected + or re.fullmatch(r"[0-9a-fA-F]{40}", revision) is None + ): + return False + path = PurePosixPath(observed) + return ( + path.is_absolute() + and ".." not in path.parts + and path.parts[-3:] == ("models--" + expected.replace("/", "--"), "snapshots", revision) + ) + + +def _validate_record( + record: dict[str, Any], + *, + launch: dict[str, Any], + bundle: InstrumentationBundle, + binding: dict[str, str], + cell: FPMCell, +) -> dict[str, Any]: + if record.get("schema_version") != OBSERVATION_SCHEMA: + raise ValueError("unsupported runtime observation schema") + for name, value in binding.items(): + _same(record.get(name), value, name) + _same(record.get("identity"), launch["identity"], "model/hardware identity") + _same(record.get("launch"), launch, "effective launch settings") + validate_sha256(record.get("model_config_sha256"), "observed model configuration") + _same(record["model_config_sha256"], launch["model_config"]["sha256"], "loaded model configuration hash") + source_files = launch["model_config"].get("source_files", {}) + if not isinstance(source_files, dict) or source_files.keys() - {"hf_quant_config.json"}: + raise ValueError("model configuration source files must name the adjacent hf_quant_config.json") + for name, digest in source_files.items(): + validate_sha256(digest, f"model configuration source {name}") + _same(record.get("model_config_source_files", {}), source_files, "loaded model configuration source hashes") + _same( + record.get("runtime"), + {**bundle.manifest["runtime"], "image": launch["deployment"]["image"]}, + "runtime build/source", + ) + if record.get("unresolved_fields") != [] or record.get("error"): + raise ValueError( + f"runtime observation has unresolved fields: {record.get('unresolved_fields')}: {record.get('error')}" + ) + kind = record.get("kind") + if kind not in {"worker", "scheduler"}: + raise ValueError("runtime observation requires worker or scheduler role") + lifecycle = record.get("lifecycle", {}) + if lifecycle.get("cache_initialized") is not True: + raise ValueError("cache initialization lifecycle evidence is missing") + if kind == "worker": + if ( + lifecycle.get("measurement") != "after_warmup" + or lifecycle.get("warmup_completed") is not True + or lifecycle.get("capture_completed") is not True + ): + raise ValueError("worker warmup and graph capture completion evidence is missing") + hardware = record.get("hardware", {}) + _validate_hardware(hardware, launch["identity"]) + total = _positive(hardware.get("total_memory_bytes"), "hardware memory bytes") + if _positive(record["cache"]["available_cache_bytes"], "profiled cache allowance") > total: + raise ValueError("profiled cache allowance exceeds physical device memory") + elif lifecycle.get("measurement") != "scheduler_initialized": + raise ValueError("scheduler initialization lifecycle evidence is missing") + collection = launch["collection"] + config = runtime_memory._validate_config( + record, + cell, + context=collection["max_model_len"], + tokens=collection["max_num_batched_tokens"], + sequences=collection["max_num_seqs"], + fraction=float(collection["gpu_memory_utilization"]), + ) + if launch["identity"]["model_kind"] == "moe": + parallel = config["parallel_config"] + enabled = parallel["enable_expert_parallel"] + if type(enabled) is not bool: + raise ValueError("observed expert topology requires a Boolean expert-parallel setting") + # vLLM .27/.28 fused_moe/config.py:FusedMoEParallelConfig.make flattens + # attention TP and DP (supported PCP1). EP owns whole experts; partial + # simultaneous expert TP/EP is not represented by this source mapping. + width = parallel["tensor_parallel_size"] * parallel["data_parallel_size"] + expert_axes = (1, width) if enabled else (width, 1) + if (cell.topology.moe_tp, cell.topology.moe_ep) != expert_axes: + raise ValueError("observed expert topology differs from the selected MoE TP/EP axes") + if config.get("speculative_config") is not None or config.get("kv_transfer_config") is not None: + raise ValueError("speculative or transferred cache semantics are unsupported") + model = config["model_config"] + if not _matches_model_path(model.get("model"), launch): + raise ValueError("observed model path differs from the selected deployment") + revision = launch["identity"]["model_revision"] + if revision not in (model.get("revision"), model.get("loaded_config_commit_hash")): + raise ValueError("observed model revision differs from the selected model") + runtime_memory._validate_revision(config, revision) + if ( + cell.fmha_quant_mode not in {"half", "bfloat16"} + or runtime_memory._dtype(model.get("dtype")) != cell.fmha_quant_mode + ): + raise ValueError("FMHA precision lacks a supported observed compute-dtype mapping") + # The raw launch config is independent corroboration of the envelope. + if model.get("enforce_eager") is not collection.get("enforce_eager", False): + raise ValueError("runtime graph eager setting differs from the reviewed launch") + if ( + (binding["phase"] == "prefill" or launch.get("worker_type") is not None) + and collection["prefill_cudagraph_policy"] == "explicit" + and not model["enforce_eager"] + ): + sampling = PrefillSamplingProfile.build( + max_isl=collection["max_num_batched_tokens"], + max_batch_size=collection["max_num_seqs"], + max_cudagraph_capture_size=collection["max_prefill_cudagraph_size"], + ) + # Historical launches apply this to prefill only; explicit serving + # roles apply their accepted configuration to every exercised phase. + # Backend-dependent + # worker mode downgrades are checked against the scheduler separately; + # they do not authorize a different initial capture list or maximum. + initial = runtime_memory._validate_graph( + record.get("initial_compilation_config", config["compilation_config"]) + if kind == "worker" + else config["compilation_config"], + eager=model["enforce_eager"], + tokens=collection["max_num_batched_tokens"], + ) + if ( + initial["cudagraph_capture_sizes"] != list(sampling.cudagraph_capture_sizes) + or initial["max_cudagraph_capture_size"] != sampling.max_cudagraph_capture_size + ): + raise ValueError("initialized explicit prefill graph captures differ from the reviewed launch") + graph_settings = { + name: collection[name] + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size") + if collection.get(name) is not None + } + if graph_settings: + if launch.get("worker_type") is None: + raise ValueError("reviewed CUDA graph settings require an explicit serving role") + if "cudagraph_capture_sizes" in graph_settings: + graph_settings.setdefault("max_cudagraph_capture_size", max(graph_settings["cudagraph_capture_sizes"])) + runtime_memory.validate_requested_graph(record, config, graph_settings) + requested_backend = launch["precision"].get("attention_backend", "auto") + if requested_backend not in (None, "auto") and record.get("attention_backend") != requested_backend: + raise ValueError("runtime observation lacks the selected attention backend evidence") + return config + + +def _validate_import_binding(record: dict, expected: dict, bundle: InstrumentationBundle) -> None: + """Recheck actual runtime imports without importing any campaign code.""" + evidence = record.get("instrumentation_binding") + if not isinstance(evidence, dict) or evidence.get("status") != "verified" or evidence.get("error"): + raise ValueError("runtime observation lacks verified instrumentation imports") + if ( + not isinstance(expected, dict) + or set(expected) != {"schema_version", "loader_sha256"} + or expected.get("schema_version") != "aisimulate-runtime-instrumentation-imports/v1" + ): + raise ValueError("unsupported runtime instrumentation import binding") + validate_sha256(expected.get("loader_sha256"), "runtime instrumentation loader") + for name, value in expected.items(): + _same(evidence.get(name), value, f"instrumentation import {name}") + for name in ("attempt_id", "configuration", "phase", "bundle_sha256"): + _same(evidence.get(name), record.get(name), f"instrumentation import {name}") + _same(evidence.get("expected_bundle_sha256"), bundle.sha256, "instrumentation expected bundle") + _same(evidence.get("files"), bundle.files, "runtime instrumentation file hashes") + _positive(evidence.get("pid"), "instrumentation import process ID") + _text(evidence.get("hostname"), "instrumentation import hostname") + cpu = record.get("cpu_affinity_observation") + if isinstance(cpu, dict): + for name in ("pid", "hostname"): + if cpu.get(name) is not None: + _same(evidence[name], cpu[name], f"instrumentation import process {name}") + root = PurePosixPath(_text(evidence.get("root"), "runtime instrumentation root")) + if not root.is_absolute() or ".." in root.parts: + raise ValueError("runtime instrumentation root must be an absolute normalized path") + imports = evidence.get("imports") + if not isinstance(imports, dict) or not imports: + raise ValueError("runtime instrumentation import evidence is empty") + for module, imported in imports.items(): + if not isinstance(module, str) or not isinstance(imported, dict): + raise ValueError("runtime instrumentation module evidence must be an object") + filename = imported.get("file") + if ( + not isinstance(filename, str) + or filename not in bundle.files + or filename + not in ( + module.replace(".", "/") + ".py", + module.replace(".", "/") + "/__init__.py", + ) + ): + raise ValueError("runtime instrumentation module is not a declared bundle source") + _same(imported.get("sha256"), bundle.files[filename], "imported instrumentation source") + path = PurePosixPath(_text(imported.get("path"), "imported instrumentation path")) + if path != root / filename: + raise ValueError("runtime instrumentation module origin contradicts its declared source") + field = "worker_class" if record["kind"] == "worker" else "scheduler_class" + classes = evidence.get("classes") + if not isinstance(classes, dict): + raise ValueError("runtime instrumentation classes must be an object") + _same(classes.get(field), bundle.manifest[field], "imported instrumentation class") + if bundle.manifest[field].rsplit(".", 1)[0] not in imports: + raise ValueError("runtime instrumentation class module was not observed") + + +def _validate_alias_ownership(views: list[tuple[dict[str, Any], int]], count: int) -> None: + """Prove disjoint pool-block ownership for the union of aliased views.""" + ownership, block_span = None, 0 + for view, ratio in views: + axis = view["block_axis"] + strides = view["stride_bytes"] + step = ratio * strides[axis] + within = view["element_size_bytes"] + (ratio - 1) * strides[axis] + planes: list[tuple[int, int]] = [] + for stride, dimension in sorted( + (stride, dimension) + for index, (stride, dimension) in enumerate(zip(strides, view["shape"], strict=True)) + if index != axis and dimension > 1 + ): + if stride < step: + within += (dimension - 1) * stride + elif planes and stride == planes[-1][0] * planes[-1][1]: + previous_stride, previous_dimension = planes.pop() + planes.append((previous_stride, previous_dimension * dimension)) + else: + planes.append((stride, dimension)) + # One pool block may contain several kernel blocks, and K/V planes may + # enclose the block axis. Normalize those dimensions, including physical + # permutations, before comparing which pool ID owns a byte. Matching + # tensor-local bounds alone cannot establish this cross-layer property. + mapping = (step, planes) + if ownership is not None and mapping != ownership: + raise ValueError("aliased tensors have incompatible pool block ownership") + ownership = mapping + block_span = max(block_span, within) + step, planes = ownership + if block_span > step: + raise ValueError("aliased pool block byte intervals overlap") + span = (count - 1) * step + block_span + for stride, dimension in planes: + if stride < span: + raise ValueError("aliased pool block planes overlap") + span += (dimension - 1) * stride + + +def _physical_views( + cache: dict[str, Any], scheduled: dict[str, Any], precision: str +) -> tuple[list[dict[str, Any]], dict[str, int]]: + _same(cache.get("semantics"), _SEMANTICS, "supported allocation semantics") + accounting = runtime_memory.validate_shared_pool_accounting(cache, scheduled) + count, page = accounting["num_blocks"], accounting["page_size_bytes"] + pool = _text(cache.get("pool_id"), "pool identity") + if scheduled.get("pool_id") != pool or scheduled.get("group_pool_ids") != [pool] * len(cache["groups"]): + raise ValueError("worker and scheduler group-to-pool mapping disagrees") + reserved = scheduled.get("permanent_reserved_block_ids") + if ( + not isinstance(reserved, list) + or any(type(value) is not int or not 0 <= value < count for value in reserved) + or len(set(reserved)) != len(reserved) + or len(reserved) != scheduled["reserved_blocks"] + or scheduled["null_block_id"] not in reserved + ): + raise ValueError("scheduler permanent reservation IDs do not account for usable blocks") + groups = runtime_memory._normalized_groups(cache, scheduled, page) + layer_groups = {} + for group_id, group in enumerate(cache["groups"]): + if group.get("pool_id") != pool: + raise ValueError("cache group references another pool") + if ( + group.get("extra_retained_tokens") != 0 + or type(group.get("extra_retained_tokens")) is not int + or group.get("attention_chunk_size") is not None + ): + raise ValueError("unsupported or unresolved extra/chunk cache retention") + storage_dtype = runtime_memory._dtype(group.get("dtype")) + # vLLM 4bdc8a788d2e2ce9165d552b3d4d8b72604626bf + # utils/torch_utils.py:STR_DTYPE_TO_TORCH_DTYPE stores FP8 as uint8, + # also used by other packed formats. Only the independently validated + # resolved cache precision establishes FP8; retain raw storage evidence. + if storage_dtype != precision and not (storage_dtype == "torch.uint8" and precision == "fp8"): + raise ValueError("cache group dtype contradicts requested KV precision") + _positive(group.get("spec_page_size_bytes"), "cache spec page bytes") + for name in group["layer_names"]: + layer_groups[name] = (group_id, group) + storages = {} + for storage in cache["storages"]: + key = _text(storage.get("storage_id"), "backing storage identity") + if key in storages or storage["size_bytes"] % count: + raise ValueError("duplicate storage identity or non-integral storage slot size") + storages[key] = storage + tensors = cache.get("layer_tensors") + allocations = cache.get("tensor_allocations") + if ( + not isinstance(tensors, dict) + or set(tensors) != set(layer_groups) + or not isinstance(allocations, list) + or not allocations + ): + raise ValueError("physical layer tensor/allocation coverage is incomplete") + covered, used_storage = set(), set() + intervals: dict[str, list[tuple[int, int]]] = {} + packing: dict[str, int] = {} + for allocation in allocations: + storage_id = allocation.get("storage_id") + if ( + storage_id not in storages + or allocation.get("size") != storages[storage_id]["size_bytes"] + or type(allocation.get("size")) is not int + ): + raise ValueError("allocation plan differs from actual backing storage") + names = allocation.get("shared_by") + if ( + not isinstance(names, list) + or not names + or any(not isinstance(name, str) or name not in layer_groups for name in names) + or len(set(names)) != len(names) + or covered.intersection(names) + ): + raise ValueError("allocation plan has missing, duplicate or unknown layers") + if len({layer_groups[name][0] for name in names}) != len(names): + raise ValueError("layers in one cache group cannot alias the same allocation slot") + stride = _nonnegative(allocation.get("block_stride"), "allocation block stride") + offset = _nonnegative(allocation.get("offset"), "allocation byte offset") + slot = storages[storage_id]["size_bytes"] // count + if (stride and stride != slot) or (not stride and (offset or storage_id in used_storage)): + raise ValueError("allocation has an unsupported packed stride or undeclared storage alias") + if storage_id in packing and packing[storage_id] != stride: + raise ValueError("packed and unpacked allocations alias one storage") + packing[storage_id] = stride + width = max(layer_groups[name][1]["spec_page_size_bytes"] for name in names) + if offset + width > slot: + raise ValueError("allocation page bytes exceed the physical storage slot") + for lower, upper in intervals.setdefault(storage_id, []): + if offset < upper and lower < offset + width: + raise ValueError("packed allocation page intervals overlap") + intervals[storage_id].append((offset, offset + width)) + views = [] + for name in names: + view = tensors[name] + group = layer_groups[name][1] + observed_offset = _nonnegative(view.get("storage_offset_bytes"), "tensor storage byte offset") + if view.get("storage_id") != storage_id or observed_offset != offset: + raise ValueError("layer tensor aliases or offsets contradict the allocation plan") + shape, strides = view.get("shape"), view.get("stride_bytes") + item_size = _positive(view.get("element_size_bytes"), "tensor element bytes") + if not isinstance(shape, list) or not shape or not isinstance(strides, list) or len(shape) != len(strides): + raise ValueError("physical tensor shape and strides are incomplete") + for dimension in shape: + _positive(dimension, "tensor dimension") + for value in strides: + _nonnegative(value, "tensor byte stride") + if value % item_size: + raise ValueError("tensor byte strides must align to its element size") + axis = _nonnegative(view.get("block_axis"), "tensor block axis") + kernel = _positive(view.get("kernel_block_size_tokens"), "kernel block tokens") + if ( + axis >= len(shape) + or group["block_size_tokens"] % kernel + or shape[axis] != count * (group["block_size_tokens"] // kernel) + ): + raise ValueError("kernel tensor block coverage contradicts pool block geometry") + ratio = group["block_size_tokens"] // kernel + if stride and (ratio != 1 or strides[axis] != stride): + raise ValueError("packed kernel tensor block stride is unsupported") + # Sufficient affine non-overlap check for the dense/permuted/strided + # views produced by the audited runtime; no enumeration of blocks. + span = item_size + for byte_stride, dimension in sorted(zip(strides, shape, strict=True)): + if dimension > 1: + if byte_stride < span: + raise ValueError("physical tensor dimensions overlap") + span += (dimension - 1) * byte_stride + if offset + span > storages[storage_id]["size_bytes"]: + raise ValueError("physical tensor byte bounds exceed its backing storage") + logical = math.prod(shape) * item_size + if logical > count * group["spec_page_size_bytes"]: + raise ValueError("tensor logical bytes exceed source cache spec page bytes") + if stride: + within_page = item_size + sum( + (d - 1) * s for index, (d, s) in enumerate(zip(shape, strides, strict=True)) if index != axis + ) + if within_page > group["spec_page_size_bytes"]: + raise ValueError("packed tensor exceeds its declared per-slot page interval") + expected_item_size = {"bfloat16": 2, "half": 2, "fp8": 1}.get(precision) + if expected_item_size is not None and item_size != expected_item_size: + raise ValueError("tensor element size contradicts cache dtype") + views.append((view, ratio)) + _validate_alias_ownership(views, count) + covered.update(names) + used_storage.add(storage_id) + if covered != set(layer_groups) or used_storage != set(storages): + raise ValueError("allocation plan does not account for every layer and backing storage") + return groups, accounting + + +def _configuration(root: Path, label: str, saved: dict[str, Any], launch: dict[str, Any]) -> dict[str, Any]: + _same(saved.get("launch"), launch, "reviewed configuration launch") + ranks = expected_ranks(launch["topology"]) + attempts = saved.get("attempts") + if not isinstance(attempts, list) or not attempts: + raise ValueError("configuration has no runtime probe attempts") + ids = [_text(attempt.get("attempt_id"), "attempt identity") for attempt in attempts] + if len(ids) != len(set(ids)): + raise ValueError("duplicate runtime probe attempt identity") + active = saved.get("active_attempt_id") + if active not in ids: + raise ValueError("configuration requires an explicit active attempt identity") + attempt = attempts[ids.index(active)] + bundle_path = contained_file(root, attempt["bundle"]["manifest"]) + bundle = load_instrumentation(bundle_path, expected_version=launch["identity"]["framework_version"]) + _same(bundle.sha256, attempt["bundle"].get("sha256"), "instrumentation bundle hash") + if not bundle.manifest.get("source_notes") or not bundle.manifest["runtime"].get("source_files"): + raise ValueError("runtime source mapping requires hashed source_notes and runtime.source_files") + phases = attempt.get("phases", {}) + required_phases = collection_phases(launch.get("worker_type")) + if set(phases) != set(required_phases): + raise ValueError("runtime evidence requires every phase of the selected serving role") + canonical_settings, canonical_geometry, canonical_hardware, normalized = None, None, None, None + capacities, artifacts, launch_artifacts, runtime_artifacts, import_bindings = [], [], [], [], [] + used_paths: set[str] = set() + for phase in required_phases: + result = phases[phase] + binding = {"attempt_id": active, "configuration": label, "phase": phase, "bundle_sha256": bundle.sha256} + launch_path, digest = _artifact(root, result["launch_manifest"]) + context = read_json(launch_path) + if context.get("schema_version") != LAUNCH_SCHEMA: + raise ValueError("unsupported runtime launch manifest schema") + for name, value in {**binding, "launch": launch, "expected_ranks": ranks}.items(): + _same(context.get(name), value, f"launch manifest {name}") + launch_artifacts.append({**result["launch_manifest"], "sha256": digest}) + import_binding = context.get("instrumentation_binding") + import_bindings.append(import_binding is not None) + workers, schedulers = {}, {} + cell = _launch_cell(launch, phase) + references = result.get("artifacts") + if not isinstance(references, list): + raise ValueError("runtime phase is missing its artifact manifest") + for reference in references: + path, digest = _artifact(root, reference) + if str(path.resolve()) in used_paths: + raise ValueError("duplicate runtime artifact path across ranks or phases") + used_paths.add(str(path.resolve())) + if reference.get("kind") != "observation": + runtime_artifacts.append(copy.deepcopy(reference)) + continue + record = read_json(path) + config = _validate_record(record, launch=launch, bundle=bundle, binding=binding, cell=cell) + if import_binding is not None: + _validate_import_binding(record, import_binding, bundle) + elif "instrumentation_binding" in record: + raise ValueError("observed instrumentation imports lack their frozen launch binding") + dp = record.get("dp_rank") + if type(dp) is not int or not 0 <= dp < cell.topology.dp: + raise ValueError("invalid observed DP rank") + if record["kind"] == "worker": + tp, pp = record.get("tp_rank"), record.get("pp_rank") + if type(tp) is not int or not 0 <= tp < cell.topology.tp or type(pp) is not int or pp != 0: + raise ValueError("invalid observed TP/PP rank") + if (dp, tp) in workers: + raise ValueError("duplicate observed worker rank") + workers[(dp, tp)] = record + if canonical_hardware is None: + canonical_hardware = {key: record["hardware"][key] for key in ("gpu", "device_name", "sm")} + else: + _same( + {key: record["hardware"][key] for key in ("gpu", "device_name", "sm")}, + canonical_hardware, + "hardware across phases/ranks", + ) + settings = copy.deepcopy(config) + settings["cache_config"].pop("enable_prefix_caching", None) + if canonical_settings is None: + canonical_settings = settings + else: + if settings["compilation_config"] != canonical_settings["compilation_config"]: + raise ValueError( + "effective runtime/graph settings across phases/ranks differ: a shared memory " + "profile requires matching observed compilation and CUDA-graph captures. " + "use separate prefill/decode serving-role configurations when the intended engines differ" + ) + _same(settings, canonical_settings, "effective runtime/graph settings across phases/ranks") + else: + if dp in schedulers or record.get("tp_rank") is not None or record.get("pp_rank") is not None: + raise ValueError("duplicate or invalid observed scheduler rank") + schedulers[dp] = record + artifacts.append({"path": reference["path"], "sha256": digest, "evidence": record}) + if set(workers) != {(r["dp_rank"], r["tp_rank"]) for r in ranks["workers"]} or set(schedulers) != { + r["dp_rank"] for r in ranks["schedulers"] + }: + raise ValueError(f"{phase} worker/scheduler rank evidence is incomplete") + for (dp, tp), worker in sorted(workers.items()): + scheduler = schedulers[dp] + runtime_memory._validate_scheduler_config(worker, scheduler) + groups, accounting = _physical_views( + worker["cache"], scheduler["cache"], runtime_memory._resolved_cache_dtype(worker["resolved_config"]) + ) + geometry = [ + {key: value for key, value in group.items() if key != "pool_id"} for group in worker["cache"]["groups"] + ] + if canonical_geometry is None: + canonical_geometry, normalized = geometry, groups + else: + _same(geometry, canonical_geometry, "cache geometry across phases/ranks") + _same(groups, normalized, "physical page cost across phases/ranks") + capacities.append( + { + "phase": phase, + "dp_rank": dp, + "tp_rank": tp, + **accounting, + "kv_cache_bytes": accounting["initial_free_blocks"] * accounting["page_size_bytes"], + } + ) + provenance = { + "source": "runtime_observations", + "observation_schema": OBSERVATION_SCHEMA, + "configuration": label, + "attempt_id": active, + "bundle_sha256": bundle.sha256, + "runtime": bundle.manifest["runtime"], + "launch": launch, + "runtime_settings": canonical_settings, + "rank_capacities": capacities, + "launch_artifacts": launch_artifacts, + "artifacts": artifacts, + "runtime_artifacts": runtime_artifacts, + "instrumentation": { + "manifest": str(bundle_path.relative_to(root)), + "files": bundle.files, + "import_binding": "verified" if all(import_bindings) else "unverified_legacy", + }, + "scope": "Physical memory of every loaded worker component. The loaded model configuration hash " + "does not attest weight file contents. Source and synthetic checks alone do not establish " + "live runtime qualification.", + } + collection = launch["collection"] + resources = { + "cache_layout": "grouped", + "cache_groups": normalized, + "max_num_tokens": collection["max_num_batched_tokens"], + "max_batch_size": collection["max_num_seqs"], + "runtime_memory": { + "kv_cache_bytes": min(value["kv_cache_bytes"] for value in capacities), + "gpu_memory_utilization": float(collection["gpu_memory_utilization"]), + "max_model_len": collection["max_model_len"], + "provenance": canonical_json(provenance), + }, + "provenance": "Validated runtime shared-pool observations; minimum usable capacity across compatible phases " + "and ranks, including padding and permanent reservations.", + } + FpmResourceProfile.model_validate(resources) + return {"status": "complete", "resources": resources, "diagnostics": [], "provenance": provenance} + + +def validate_observations( + index_path: str | Path, expected_configurations: dict[str, dict[str, Any]] +) -> dict[str, dict[str, Any]]: + """Validate every requested configuration while preserving independent success. + + Only AISimulate code runs. No observer verdict is authoritative, no timing + rows are required, and historical attempts never substitute for the active + attempt. Returned resources remain drafts for the existing review workflow. + """ + index_path = Path(index_path).absolute() + if index_path.is_symlink(): + raise ValueError("runtime observation index cannot be a symlink") + index = read_json(index_path) + if index.get("schema_version") != INDEX_SCHEMA or not isinstance(index.get("configurations"), dict): + raise ValueError("unsupported runtime observation index schema") + if not isinstance(expected_configurations, dict) or not expected_configurations: + raise ValueError("runtime import requires selected configurations with reviewed launch facts") + results = {} + for label, launch in expected_configurations.items(): + try: + if label not in index["configurations"]: + raise ValueError("selected configuration has no runtime observation evidence") + results[label] = _configuration(index_path.parent, label, index["configurations"][label], launch) + results[label]["provenance"]["observations_index"] = { + "path": str(index_path), + "sha256": sha256_bytes(index_path.read_bytes()), + } + results[label]["resources"]["runtime_memory"]["provenance"] = canonical_json(results[label]["provenance"]) + except (ValueError, TypeError, KeyError, AttributeError, IndexError, OSError) as error: + results[label] = { + "status": "incomplete", + "resources": None, + "diagnostics": [f"{label}: {error}"], + "provenance": {"observations_index": str(index_path)}, + } + return results diff --git a/python/aisimulate/collector/fpm_forward/runtime_probe.py b/python/aisimulate/collector/fpm_forward/runtime_probe.py new file mode 100644 index 000000000..510ff6559 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime_probe.py @@ -0,0 +1,1034 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Launch runtime observations before simulation cache geometry is known.""" + +from __future__ import annotations + +import argparse +import copy +import fcntl +import hashlib +import json +import math +import re +import shutil +import uuid +import zipfile +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from aisimulate.fpm_advisory import prefill_graph_advisory + +from . import runner +from .config import FPM_WARMUP_ITERATIONS, FPMCollectionOptions, collection_phases +from .entry import _load_generator_overrides +from .model_capability import ResolvedModelConfig, load_model_config +from .planner import FPMCell, _backend_policies, _canonical_hash +from .runtime_instrumentation import contained_file, read_json, validate_sha256 +from .topology import enumerate_fpm_topologies, topology_strategy +from .types import ParallelTopology + +INDEX_SCHEMA = "aisimulate-runtime-observations/v1" +LAUNCH_SCHEMA = "aisimulate-runtime-probe-launch/v1" +CONTEXT_FILENAME = "runtime-probe-context.json" +BUNDLE_FILENAME = "runtime-instrumentation.zip" +PHASES = ("prefill", "decode") +# Dynamo b83b1d9304ebfc624709ac46db32b1b6f1ff1615, backend_args.py: +# _validate_benchmark_sampling requires two endpoints on each uniform axis. +DEFAULT_SAMPLING = {"max_new_token_samples": 2, "max_kv_read_token_samples": 2, "max_batch_size_samples": 2} + + +@dataclass(frozen=True, slots=True) +class _RuntimeModel: + architecture: str + aic_database_version: str + model_config: ResolvedModelConfig + + +@dataclass(frozen=True, slots=True) +class RuntimeProbePlan: + """The renderer's launch inputs, with no performance or resource profile.""" + + configuration: str + _launch_json: str + options: FPMCollectionOptions + capability: _RuntimeModel + cells: tuple[FPMCell, ...] + runtime_instrumentation: Any + sha256: str + + @property + def launch(self) -> dict[str, Any]: + return json.loads(self._launch_json) + + @property + def model_path(self) -> str: + return self.launch["identity"]["model"] + + @property + def system(self) -> str: + return self.launch["identity"]["gpu"] + + @property + def backend(self) -> str: + return self.launch["identity"]["framework"] + + @property + def fpm_profile(self) -> None: + return None + + @property + def runtime_launch(self) -> dict[str, Any]: + return self.launch + + @property + def probe_sampling(self) -> dict[str, int]: + return dict(DEFAULT_SAMPLING) + + +def _required_mapping(payload: dict[str, Any], field: str) -> dict[str, Any]: + value = payload.get(field) + if not isinstance(value, dict): + raise ValueError(f"runtime probe requires {field} launch facts") + return value + + +def normalize_probe_launch(value: dict[str, Any]) -> dict[str, Any]: + """Freeze explicit stage-2 launch facts and documented sampling defaults.""" + + launch = json.loads(json.dumps(value, allow_nan=False)) + collection_phases(launch.get("worker_type")) + identity = _required_mapping(launch, "identity") + for name in ("model", "model_revision", "model_kind", "framework", "framework_version", "gpu", "interconnect"): + item = identity.get(name) + if not isinstance(item, str) or not item.strip() or item != item.strip() or any(ord(c) < 32 for c in item): + raise ValueError(f"runtime probe requires identity.{name}") + if identity["framework"] != "vllm" or identity["model_kind"] not in {"dense", "moe"}: + raise ValueError("runtime probe requires framework=vllm and explicit model_kind=dense|moe") + if any( + identity[name].lower() in {"main", "master", "latest", "head"} + for name in ("model_revision", "framework_version") + ): + raise ValueError("runtime probe requires pinned model and framework revisions") + topology = _required_mapping(launch, "topology") + for name in ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp"): + if type(topology.get(name)) is not int or topology[name] < 1: + raise ValueError(f"runtime probe requires positive topology.{name}") + if topology["pp"] != 1 or topology["cp"] != 1: + raise ValueError("runtime observation currently supports PP1 and CP1") + precision = _required_mapping(launch, "precision") + for name in ("gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "kvcache_quant_mode", "comm_quant_mode"): + if not isinstance(precision.get(name), str) or not precision[name] or precision[name] == "auto": + raise ValueError(f"runtime probe requires explicit precision.{name}") + if precision["comm_quant_mode"] != "half": + raise ValueError("runtime probe requires the collector comm_quant_mode=half identity") + for name in ("moe_backend", "attention_backend"): + precision.setdefault(name, "auto") + if not isinstance(precision[name], str) or not precision[name]: + raise ValueError(f"precision.{name} must be a nonempty string") + for name in ("enable_wideep", "enable_eplb"): + precision.setdefault(name, False) + if type(precision[name]) is not bool: + raise ValueError(f"precision.{name} must be boolean") + collection = _required_mapping(launch, "collection") + for name in ("max_model_len", "max_num_batched_tokens", "max_num_seqs"): + if type(collection.get(name)) is not int or collection[name] < 1: + raise ValueError(f"runtime probe requires positive collection.{name}") + fraction = collection.get("gpu_memory_utilization") + if type(fraction) not in {int, float} or not math.isfinite(fraction) or not 0 < fraction <= 1: + raise ValueError("runtime probe requires collection.gpu_memory_utilization in (0, 1]") + policy = collection.get("prefill_cudagraph_policy") + if policy not in {"runtime", "explicit"}: + raise ValueError("runtime probe requires collection.prefill_cudagraph_policy=runtime|explicit") + if collection.get("max_prefill_cudagraph_size") is None: + collection["max_prefill_cudagraph_size"] = None if policy == "runtime" else 2048 + collection.setdefault("async_scheduling", False) + if collection["async_scheduling"] is not False: + raise ValueError("the native prefill observation protocol requires async_scheduling=false in both phases") + collection.setdefault("enforce_eager", False) + if type(collection["enforce_eager"]) is not bool: + raise ValueError("collection.enforce_eager must be boolean") + collection.setdefault("warmup_iterations", FPM_WARMUP_ITERATIONS) + if type(collection["warmup_iterations"]) is not int or collection["warmup_iterations"] < 0: + raise ValueError("collection.warmup_iterations must be a nonnegative integer") + config = _required_mapping(launch, "model_config") + if not isinstance(config.get("path"), str) or not config["path"]: + raise ValueError("runtime probe requires model_config.path and model_config.sha256") + validate_sha256(config.get("sha256"), "model configuration") + sources = config.get("source_files", {}) + if not isinstance(sources, dict) or sources.keys() - {"hf_quant_config.json"}: + raise ValueError("model_config.source_files supports the adjacent hf_quant_config.json only") + for name, digest in sources.items(): + validate_sha256(digest, f"model configuration source {name}") + deployment = _required_mapping(launch, "deployment") + deployment.setdefault("executor", "kubernetes") + deployment.setdefault("container_mount", []) + if deployment["executor"] not in {"kubernetes", "slurm"}: + raise ValueError("runtime probe executor must be kubernetes or slurm") + image = deployment.get("image") + if not isinstance(image, str) or not image or any(c.isspace() or ord(c) < 32 for c in image): + raise ValueError("runtime probe requires deployment.image for the pinned runtime") + mounts = deployment["container_mount"] + if not isinstance(mounts, list) or any( + not isinstance(m, str) or not m or any(ord(c) < 32 or c == "," for c in m) for m in mounts + ): + raise ValueError("deployment.container_mount must contain valid Slurm mount strings") + if deployment["executor"] == "slurm" and any( + deployment.get(key) is not None for key in ("namespace", "model_cache", "image_pull_secret") + ): + raise ValueError("Slurm runtime probes reject Kubernetes namespace, model_cache and image_pull_secret") + if deployment["executor"] != "slurm" and mounts: + raise ValueError("container_mount requires the Slurm executor") + cpus, binding = deployment.get("cpus_per_task"), deployment.get("cpu_bind") + if cpus is not None or binding is not None: + if deployment["executor"] != "slurm": + raise ValueError("CPU deployment options require the Slurm executor") + if type(cpus) is not int or cpus < 1 or binding not in {"cores", "none"}: + raise ValueError("Slurm CPU policy requires positive cpus_per_task and cpu_bind cores or none") + return launch + + +def _model_config_bytes(config: dict[str, Any]) -> dict[str, bytes]: + """Bind the raw config and every adjacent file consumed by the local loader.""" + + path = Path(config["path"]).expanduser().absolute() + sources = config.get("source_files", {}) + adjacent = path.parent / "hf_quant_config.json" + if (adjacent.exists() or adjacent.is_symlink()) and adjacent.name not in sources: + raise ValueError("adjacent hf_quant_config.json must be recorded in model_config.source_files") + contents = {"model-config.json": path.read_bytes()} + if hashlib.sha256(contents["model-config.json"]).hexdigest() != config["sha256"]: + raise ValueError("runtime probe model config does not match its recorded SHA-256") + for name, digest in sources.items(): + # HF snapshots use adjacent symlinks into a shared blob directory. The + # logical name is constrained above; its loaded bytes remain hash-bound. + source = path.parent / name + if not source.is_file(): + raise ValueError(f"model configuration source is not a regular file: {name}") + contents[name] = source.read_bytes() + if hashlib.sha256(contents[name]).hexdigest() != digest: + raise ValueError(f"model configuration source {name} does not match its recorded SHA-256") + return contents + + +def probe_generator_overrides(launch: dict[str, Any]) -> dict[str, Any]: + """Translate the existing deployment options through their shared parser.""" + + deployment = launch["deployment"] + image = deployment["image"] + args = argparse.Namespace( + generator_config=None, + generator_set=[f"K8sConfig.k8s_image={json.dumps(image)}"], + generator_dynamo_version=deployment.get("dynamo_version"), + generated_config_version=None, + namespace=deployment.get("namespace"), + transport=deployment.get("transport"), + image_pull_secret=deployment.get("image_pull_secret"), + model_cache=deployment.get("model_cache"), + ) + return _load_generator_overrides(args) + + +def build_runtime_probe_plan(configuration: str, facts: dict[str, Any], bundle: Any) -> RuntimeProbePlan: + launch = normalize_probe_launch(facts) + identity, precision, collection, deployment = ( + launch[key] for key in ("identity", "precision", "collection", "deployment") + ) + if bundle.manifest["runtime"]["version"] != identity["framework_version"]: + raise ValueError("instrumentation runtime version does not match the requested pin") + _model_config_bytes(launch["model_config"]) + model = load_model_config(identity["model"], explicit_config_path=launch["model_config"]["path"]) + # The shared loader reads the adjacent sidecar; reject changes during loading + # without replacing the raw primary-config hash with a merged payload hash. + _model_config_bytes(launch["model_config"]) + config = model.effective_payload + architecture = (config.get("architectures") or [None])[0] + if not isinstance(architecture, str) or not architecture: + raise ValueError("runtime probe model config requires an architecture") + from aisimulate.sdk.models.helpers import _infer_quant_modes_from_raw_config + from aisimulate.sdk.utils import _attach_inferred_quant_fields + + inferred = _infer_quant_modes_from_raw_config(_attach_inferred_quant_fields(config), architecture) + for name in ("gemm_quant_mode", "moe_quant_mode"): + actual = getattr(inferred.get(name), "name", inferred.get(name)) + if actual is None: + actual = {"bfloat16": "bfloat16", "float16": "half", "half": "half"}.get( + config.get("dtype") or config.get("torch_dtype") + ) + if actual != precision[name]: + raise ValueError(f"precision.{name}={precision[name]!r} does not match checkpoint-native {actual!r}") + topology = ParallelTopology(**launch["topology"]) + is_moe = identity["model_kind"] == "moe" + strategy = topology_strategy(topology, is_moe=is_moe) + preset = strategy if strategy != "single" else ("tep" if is_moe else "tp") + options = FPMCollectionOptions( + max_gpus=topology.total_gpus, + gpu_counts=(topology.total_gpus,), + parallel_presets=(preset,), + parallel_axes=(), + moe_backend=precision["moe_backend"], + attention_backend=precision["attention_backend"], + enable_wideep=str(precision["enable_wideep"]).lower(), + enable_eplb=str(precision["enable_eplb"]).lower(), + weight_quantizations=(precision["gemm_quant_mode"],), + kv_cache_dtypes=(precision["kvcache_quant_mode"],), + warmup_iterations=collection["warmup_iterations"], + vllm_max_model_len=collection["max_model_len"], + max_num_batched_tokens=collection["max_num_batched_tokens"], + max_num_seqs=collection["max_num_seqs"], + prefill_cudagraph_policy=collection["prefill_cudagraph_policy"], + worker_type=launch.get("worker_type"), + cudagraph_mode=collection.get("cudagraph_mode"), + cudagraph_capture_sizes=( + tuple(collection["cudagraph_capture_sizes"]) + if collection.get("cudagraph_capture_sizes") is not None + else None + ), + max_cudagraph_capture_size=collection.get("max_cudagraph_capture_size"), + max_prefill_cudagraph_size=collection["max_prefill_cudagraph_size"], + gpu_memory_utilization=collection["gpu_memory_utilization"], + enforce_eager=collection["enforce_eager"], + executor=deployment["executor"], + slurm_container_image=deployment["image"] if deployment["executor"] == "slurm" else "", + slurm_container_mounts=tuple(deployment["container_mount"]), + slurm_cpus_per_task=deployment.get("cpus_per_task"), + slurm_cpu_bind=deployment.get("cpu_bind"), + ) + admitted = enumerate_fpm_topologies(backend="vllm", is_moe=is_moe, options=options, allow_pure_tp=True) + if topology not in admitted: + raise ValueError(f"collector does not admit exact topology {topology.to_dict()}") + policy = _backend_policies(options, {}, backend="vllm")[0] + sha256 = _canonical_hash({"launch": launch, "bundle_sha256": bundle.sha256}) + cells = tuple( + FPMCell( + cell_id="probe-" + _canonical_hash({"configuration": configuration, "phase": phase, "plan": sha256})[:20], + workload_kind=phase, + topology=topology, + weight_quantization=precision["gemm_quant_mode"], + kv_cache_dtype=precision["kvcache_quant_mode"], + backend_policy=policy, + gemm_quant_mode=precision["gemm_quant_mode"], + moe_quant_mode=precision["moe_quant_mode"], + fmha_quant_mode=precision["fmha_quant_mode"], + comm_quant_mode=precision["comm_quant_mode"], + parallel_strategy=strategy, + fmha_resolution="runtime_probe_requested", + ) + for phase in options.workload_kinds + ) + return RuntimeProbePlan( + configuration, + json.dumps(launch, sort_keys=True), + options, + _RuntimeModel(architecture, identity["framework_version"], model), + cells, + bundle, + sha256, + ) + + +def validate_collection_probe_launch( + launch, bundle, *, model_path, system, backend, backend_version, cells, options, profile, generator_overrides +): + """Bind optional formal instrumentation to the actual collection inputs.""" + + expected = build_runtime_probe_plan("collection", launch, bundle) + identity = launch["identity"] + if (model_path, system, backend, backend_version) != ( + identity["model"], + identity["gpu"], + identity["framework"], + identity["framework_version"], + ): + raise ValueError("formal collection identity differs from the accepted runtime probe") + if profile is None or profile.model_revision != identity["model_revision"]: + raise ValueError("formal collection model revision differs from the accepted runtime probe") + expected_cell = expected.cells[0] + if {cell.workload_kind for cell in cells} != set(expected.options.workload_kinds): + raise ValueError("formal collection phases differ from the accepted serving role") + for cell in cells: + fields = ( + "topology", + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "comm_quant_mode", + "kv_cache_dtype", + ) + if ( + any(getattr(cell, key) != getattr(expected_cell, key) for key in fields) + or cell.backend_policy.to_dict() != expected_cell.backend_policy.to_dict() + ): + raise ValueError("formal collection topology or precision differs from the accepted runtime probe") + expected_options = expected.options + for name in ( + "vllm_max_model_len", + "prefill_cudagraph_policy", + "max_prefill_cudagraph_size", + "gpu_memory_utilization", + "enforce_eager", + "warmup_iterations", + "executor", + "slurm_container_image", + "slurm_container_mounts", + "slurm_cpus_per_task", + "slurm_cpu_bind", + "worker_type", + "cudagraph_mode", + "cudagraph_capture_sizes", + "max_cudagraph_capture_size", + ): + if getattr(options, name) != getattr(expected_options, name): + raise ValueError(f"formal collection {name} differs from the accepted runtime probe") + deployment = profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=backend_version, + tp_size=expected_cell.topology.tp, + pp_size=expected_cell.topology.pp, + attention_dp_size=expected_cell.topology.dp, + moe_tp_size=expected_cell.topology.moe_tp, + moe_ep_size=expected_cell.topology.moe_ep, + cp_size=expected_cell.topology.cp, + worker_type=options.worker_type or "aggregated", + ) + limits = { + "decode": ( + options.max_num_batched_tokens or deployment.resources.max_num_tokens, + options.max_decode_batch_size or options.max_num_seqs or deployment.resources.max_batch_size, + ), + "prefill": ( + options.prefill_sampling.max_total_prefill_tokens, + options.prefill_sampling.max_batch_size or options.max_num_seqs or deployment.resources.max_batch_size, + ), + } + required_limits = (launch["collection"]["max_num_batched_tokens"], launch["collection"]["max_num_seqs"]) + if any(limits[phase] != required_limits for phase in options.workload_kinds): + raise ValueError("formal phase scheduler limits differ from the accepted runtime probe") + wanted = probe_generator_overrides(launch) + actual = copy.deepcopy(generator_overrides) + if options.executor == "slurm": + wanted.get("K8sConfig", {}).pop("k8s_image", None) + actual.get("K8sConfig", {}).pop("k8s_image", None) + if runner.with_kv_warmup_defaults(wanted) != runner.with_kv_warmup_defaults(actual): + raise ValueError("formal collection deployment settings differ from the accepted runtime probe") + + +def prepare_collection_observations(plan, root: Path, parent_attempt_id: str): + """Freeze one parent attempt shared by phase-specific collector attempts.""" + + from .runtime_instrumentation import freeze_instrumentation, load_instrumentation + + destination = root / "runtime-observations" / parent_attempt_id / "instrumentation" + bundle = plan.runtime_instrumentation + if destination.exists(): + frozen = load_instrumentation(destination / "manifest.json") + if frozen.sha256 != bundle.sha256: + raise ValueError("formal collection frozen instrumentation hash mismatch") + return frozen + destination.parent.mkdir(parents=True, exist_ok=True) + return freeze_instrumentation(bundle, destination) + + +def record_collection_observations( + plan, root: Path, parent_attempt_id: str, cell, cell_attempt_id: str, cell_dir: Path, status: str +) -> Path: + """Archive a cell's evidence before the ordinary collector can retry it.""" + + index_path = root / "runtime-observations.json" + index = ( + json.loads(index_path.read_text()) + if index_path.exists() + else {"schema_version": INDEX_SCHEMA, "configurations": {}} + ) + configuration = plan.runtime_configuration + entry = index["configurations"].setdefault( + configuration, {"launch": plan.runtime_launch, "active_attempt_id": parent_attempt_id, "attempts": []} + ) + if entry["launch"] != plan.runtime_launch or entry["active_attempt_id"] != parent_attempt_id: + raise ValueError("formal collection observation attempt identity mismatch") + if not entry["attempts"]: + entry["attempts"].append( + { + "attempt_id": parent_attempt_id, + "status": "incomplete", + "bundle": { + "manifest": f"runtime-observations/{parent_attempt_id}/instrumentation/manifest.json", + "sha256": plan.runtime_instrumentation.sha256, + }, + "phases": {}, + "phase_attempts": {}, + } + ) + attempt = entry["attempts"][-1] + if attempt["attempt_id"] != parent_attempt_id or attempt["bundle"]["sha256"] != plan.runtime_instrumentation.sha256: + raise ValueError("formal collection frozen observation identity mismatch") + + def verify_context(path: Path) -> None: + saved = read_json(path) + expected = launch_context( + plan, cell, configuration=configuration, attempt_id=parent_attempt_id, saved_context=saved + ) + expected.update(collector_attempt_id=cell_attempt_id, cell_id=cell.cell_id) + if saved != expected: + raise ValueError("formal collection saved runtime context does not match its attempt") + + history = attempt["phase_attempts"].setdefault(cell.workload_kind, []) + for previous in reversed(history): + if previous["collector_attempt_id"] != cell_attempt_id or previous["status"] != status or status != "passed": + continue + if previous.get("launch_manifest") is not None: + verify_context(_verified_artifact(root, previous["launch_manifest"])) + elif status == "passed": + raise ValueError("formal collection passed attempt requires its saved runtime context") + for artifact in previous["artifacts"]: + _verified_artifact(root, artifact) + phase = previous + break + else: + context = cell_dir / CONTEXT_FILENAME + if context.is_file(): + verify_context(context) + elif status == "passed": + raise ValueError("formal collection passed attempt requires its saved runtime context") + snapshot = root / "runtime-observations" / parent_attempt_id / cell.workload_kind / cell_attempt_id + # A host can die after copying files and before committing the index. + # Preserve that unindexed directory and publish a new complete snapshot. + if snapshot.exists(): + snapshot = snapshot.with_name(snapshot.name + "-" + uuid.uuid4().hex) + staging = snapshot.with_name(".staging-" + uuid.uuid4().hex) + staging.mkdir(parents=True) + for name in ("raw", "logs"): + if (cell_dir / name).exists(): + shutil.copytree(cell_dir / name, staging / name, symlinks=True) + for path in cell_dir.iterdir(): + if path.is_file(): + if path.is_symlink(): + raise ValueError("formal collection launch files cannot be symlinks") + shutil.copy2(path, staging / path.name) + # Validate containment before exposing the snapshot to its index. + observation_artifacts(root, staging) + if (staging / CONTEXT_FILENAME).exists(): + verify_context(staging / CONTEXT_FILENAME) + staging.rename(snapshot) + phase = { + "status": status, + "collector_attempt_id": cell_attempt_id, + "cell_id": cell.cell_id, + "artifacts": observation_artifacts(root, snapshot), + } + if context.is_file(): + phase["launch_manifest"] = { + "path": str((snapshot / CONTEXT_FILENAME).relative_to(root)), + **runner._file_metadata(snapshot / CONTEXT_FILENAME), + } + history.append(phase) + attempt["phases"][cell.workload_kind] = phase + attempt["status"] = ( + "captured" + if set(attempt["phases"]) == set(collection_phases(plan.runtime_launch.get("worker_type"))) + and all(item["status"] == "passed" for item in attempt["phases"].values()) + else "incomplete" + ) + runner._atomic_json(index_path, index) + return index_path + + +def launch_context( + plan, cell: FPMCell, *, configuration: str, attempt_id: str, saved_context: dict[str, Any] | None = None +) -> dict[str, Any]: + from .runtime_instrumentation import runtime_binding + + topology = cell.topology + context = { + "schema_version": LAUNCH_SCHEMA, + "configuration": configuration, + "attempt_id": attempt_id, + "phase": cell.workload_kind, + "bundle_sha256": plan.runtime_instrumentation.sha256, + "instrumentation_binding": runtime_binding(), + "launch": plan.runtime_launch, + "sampling": getattr(plan, "probe_sampling", None), + "expected_ranks": { + "workers": [ + {"dp_rank": dp, "tp_rank": tp, "pp_rank": pp} + for dp in range(topology.dp) + for tp in range(topology.tp) + for pp in range(topology.pp) + ], + "schedulers": [{"dp_rank": dp} for dp in range(topology.dp)], + }, + } + if saved_context is not None: + # Recovery compares the archived attempt's identity. A loader upgrade + # does not invalidate its ownership or authorize relabeling old imports. + if "instrumentation_binding" in saved_context: + context["instrumentation_binding"] = saved_context["instrumentation_binding"] + else: + context.pop("instrumentation_binding") + return context + + +def stage_runtime_instrumentation(bundle, cell_dir: Path, context: dict[str, Any]) -> list[Path]: + """Package data without importing it; activation occurs only in the worker.""" + + from .runtime_instrumentation import INSTRUMENTATION_LOADER, runtime_binding + + if context.get("instrumentation_binding") != runtime_binding(): + raise ValueError("runtime instrumentation requires the current verified loader binding") + for name, payload in _model_config_bytes(context["launch"]["model_config"]).items(): + (cell_dir / name).write_bytes(payload) + archive = cell_dir / BUNDLE_FILENAME + with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_STORED) as target: + target.writestr("manifest.json", json.dumps(bundle.manifest, sort_keys=True)) + for relative, expected in bundle.files.items(): + payload = (bundle.root / relative).read_bytes() + if hashlib.sha256(payload).hexdigest() != expected: + raise ValueError(f"instrumentation changed before staging: {relative}") + target.writestr(relative, payload) + context_path = cell_dir / CONTEXT_FILENAME + runner._atomic_json(context_path, context) + archive_sha256 = hashlib.sha256(archive.read_bytes()).hexdigest() + # Archive bytes and entry names are verified before extraction. The host + # never imports campaign Python modules, including during preview. + setup = f""" +python3 - <<'AISIMULATE_INSTRUMENTATION' +import hashlib, pathlib, zipfile +archive = pathlib.Path('/tmp/fpm-bench/{BUNDLE_FILENAME}') +if hashlib.sha256(archive.read_bytes()).hexdigest() != {archive_sha256!r}: + raise RuntimeError('runtime instrumentation archive hash mismatch') +root = pathlib.Path('/tmp/fpm-bench/runtime-instrumentation') +with zipfile.ZipFile(archive) as bundle: + for name in bundle.namelist(): + path = pathlib.PurePosixPath(name) + if path.is_absolute() or '..' in path.parts or path.as_posix() != name: + raise RuntimeError('unsafe runtime instrumentation path') + bundle.extractall(root) +AISIMULATE_INSTRUMENTATION +export PYTHONPATH="/tmp/fpm-bench/runtime-instrumentation:/tmp/fpm-bench${{PYTHONPATH:+:${{PYTHONPATH}}}}" +""" + with (cell_dir / runner.RUNTIME_ENV_FILENAME).open("a") as handle: + handle.write(setup) + return [archive, context_path, Path(__file__).parent / "runtime" / INSTRUMENTATION_LOADER] + + +def observation_artifacts(root: Path, phase_dir: Path) -> list[dict[str, Any]]: + artifacts = [] + for directory in (phase_dir, phase_dir / "raw", phase_dir / "logs"): + paths = directory.iterdir() if directory == phase_dir else directory.rglob("*") + for path in sorted(paths) if directory.exists() else (): + if not path.is_file(): + continue + if path.name == CONTEXT_FILENAME and path.parent == phase_dir: + continue + if path.is_symlink() or not path.resolve().is_relative_to(root): + raise ValueError("runtime artifacts must be regular files within the probe output") + kind = "runtime-artifact" + if path.suffix == ".json": + try: + payload = json.loads(path.read_text()) + except (ValueError, UnicodeError): + payload = None + if isinstance(payload, dict) and payload.get("schema_version") == "aisimulate-runtime-observation/v1": + kind = "observation" + artifacts.append({"path": str(path.relative_to(root)), "kind": kind, **runner._file_metadata(path)}) + return artifacts + + +def _snapshot_probe_artifacts(root: Path, phase_dir: Path) -> list[dict[str, Any]]: + """Keep imported evidence outside the executor's potentially live result mount.""" + + references = observation_artifacts(root, phase_dir) + snapshot = phase_dir / ("evidence-" + uuid.uuid4().hex) + snapshot.mkdir() + artifacts = [] + for reference in references: + source = root / reference["path"] + target = snapshot / source.relative_to(phase_dir) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target) + artifacts.append( + {"path": str(target.relative_to(root)), "kind": reference["kind"], **runner._file_metadata(target)} + ) + return artifacts + + +def _recorded_attempt_valid(root: Path, attempt: dict[str, Any], bundle, phases: tuple[str, ...] = PHASES) -> bool: + from .runtime_instrumentation import load_instrumentation + + if attempt.get("status") != "completed" or attempt.get("bundle", {}).get("sha256") != bundle.sha256: + return False + frozen = load_instrumentation(root / attempt["bundle"]["manifest"]) + if frozen.sha256 != bundle.sha256 or set(attempt.get("phases", {})) != set(phases): + raise ValueError("saved probe attempt does not match its frozen instrumentation") + for phase in attempt["phases"].values(): + for artifact in [phase["launch_manifest"], *phase["artifacts"]]: + path = root / artifact["path"] + if ( + not path.resolve().is_relative_to(root) + or hashlib.sha256(path.read_bytes()).hexdigest() != artifact["sha256"] + ): + raise ValueError(f"saved probe artifact changed: {artifact['path']}") + return True + + +def _verified_artifact(root: Path, reference: dict[str, Any]) -> Path: + path = contained_file(root, reference["path"]) + digest = validate_sha256(reference.get("sha256"), "saved artifact") + if hashlib.sha256(path.read_bytes()).hexdigest() != digest: + raise ValueError(f"saved artifact SHA-256 mismatch: {reference['path']}") + return path + + +def resolve_probe_cpu_policy( + configuration: str, deployment: dict[str, Any], output_dir: str | Path, *, resume: bool +) -> dict[str, Any]: + """Resolve omitted CPU fields before launch validation, retaining old absence.""" + from .config import resolve_slurm_cpu_policy + + deployment = copy.deepcopy(deployment) + if deployment.get("executor") != "slurm": + return deployment + root = Path(output_dir).expanduser().resolve() + index_path = root / "observations.json" + if resume and index_path.exists(): + index = read_json(index_path) + if index.get("schema_version") != INDEX_SCHEMA or not isinstance(index.get("configurations"), dict): + raise ValueError("runtime probe observation index has an unsupported schema") + existing = index["configurations"].get(configuration) + if existing is not None: + if not isinstance(existing, dict): + raise ValueError("saved runtime probe configuration must be an object") + saved = _required_mapping(_required_mapping(existing, "launch"), "deployment") + inherited = { + field: saved[field] + for field in ("cpus_per_task", "cpu_bind") + if deployment.get(field) is None and field in saved + } + if inherited: + attempts = existing.get("attempts", []) + if not isinstance(attempts, list) or any(not isinstance(attempt, dict) for attempt in attempts): + raise ValueError("saved runtime probe attempts must be a list of objects") + phases = _required_mapping(attempts[-1], "phases") if attempts else {} + if any(not isinstance(phase, dict) for phase in phases.values()): + raise ValueError("saved runtime probe phases must be objects") + references = [phase["launch_manifest"] for phase in phases.values() if "launch_manifest" in phase] + if not references or any(not isinstance(reference, dict) for reference in references): + raise ValueError("saved CPU policy has no SHA-bound runtime launch context") + for reference in references: + context = read_json(_verified_artifact(root, reference)) + if context.get("launch") != existing["launch"]: + raise ValueError("saved CPU policy differs from its archived runtime launch context") + deployment.update(inherited) + return deployment + cpus, binding = resolve_slurm_cpu_policy(deployment.get("cpus_per_task"), deployment.get("cpu_bind")) + deployment.update(cpus_per_task=cpus, cpu_bind=binding) + return deployment + + +def _probe_configuration_slug(configuration: str) -> str: + return re.sub(r"[^a-zA-Z0-9_-]", "-", configuration)[:40] + "-" + _canonical_hash(configuration)[:8] + + +def _resource_manifest_reference(root: Path, plan, cell, directory: Path) -> dict[str, Any]: + manifest = contained_file(root, str((directory / runner.FPM_MANIFEST_FILENAME).relative_to(root))) + documents = runner._manifest_documents(manifest) + runner._workload_document(documents) + ownership = { + "aiconfigurator.nvidia.com/owned-by": "fpm-forward-collector", + "aiconfigurator.nvidia.com/plan": plan.sha256[:16], + runner.FPM_CELL_LABEL: cell.cell_id, + } + for document in documents: + labels = (document.get("metadata") or {}).get("labels") or {} + if any(labels.get(key) != value for key, value in ownership.items()): + raise ValueError("saved runtime resource does not match collector ownership") + return { + "path": str(manifest.relative_to(root)), + **runner._file_metadata(manifest), + "ownership": ownership, + "resources": [list(runner._resource_identity(document)) for document in documents], + } + + +def _cleanup_previous_attempt(root: Path, plan, attempt: dict[str, Any]) -> None: + """Verify saved launch ownership before any recovery command is allowed.""" + + from .runtime_instrumentation import load_instrumentation + + attempt_id = attempt["attempt_id"] + if not isinstance(attempt_id, str) or not re.fullmatch(r"[0-9a-f]{32}", attempt_id): + raise ValueError("invalid saved runtime probe attempt identity") + directory = root / "attempts" / _probe_configuration_slug(plan.configuration) / attempt_id + manifest_path = contained_file(root, attempt["bundle"]["manifest"]) + if manifest_path != directory / "instrumentation" / "manifest.json": + raise ValueError("saved runtime bundle path does not belong to this attempt") + bundle = load_instrumentation(manifest_path) + if bundle.sha256 != attempt["bundle"]["sha256"]: + raise ValueError("saved runtime instrumentation SHA-256 mismatch") + previous_plan = build_runtime_probe_plan(plan.configuration, plan.launch, bundle) + verified = [] + for cell in previous_plan.cells: + phase = attempt.get("phases", {}).get(cell.workload_kind, {}) + if phase.get("status") not in {"running", "interrupted", "cleanup_failed"}: + continue + reference = phase.get("launch_manifest") + if reference is None: + raise ValueError("interrupted runtime probe has no saved launch manifest") + context_path = _verified_artifact(root, reference) + phase_dir = directory / cell.workload_kind + if context_path != phase_dir / CONTEXT_FILENAME: + raise ValueError("saved runtime launch path does not belong to this attempt/phase") + saved = read_json(context_path) + expected = launch_context( + previous_plan, cell, configuration=plan.configuration, attempt_id=attempt_id, saved_context=saved + ) + if saved != expected: + raise ValueError("saved runtime launch context does not match its attempt identity") + resource_reference = phase.get("resource_manifest") + if not isinstance(resource_reference, dict): + raise ValueError("interrupted runtime probe has no saved resource manifest") + manifest = _verified_artifact(root, resource_reference) + if resource_reference != _resource_manifest_reference(root, previous_plan, cell, phase_dir): + raise ValueError("saved runtime resource manifest identity mismatch") + verified.append((cell, phase, phase_dir, manifest)) + for cell, phase, phase_dir, manifest in verified: + resource = runner._cell_runner(previous_plan, cell, manifest, phase_dir) + runner._salvage_artifacts(resource, cell.cell_id) + resource.cleanup() + phase["recovery_artifacts"] = _snapshot_probe_artifacts(root, phase_dir) + + +@contextmanager +def _probe_lock(root: Path): + with (root / ".runtime-probe.lock").open("a") as handle: + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + raise ValueError("another runtime probe owns this output directory") from error + try: + yield + finally: + fcntl.flock(handle, fcntl.LOCK_UN) + + +def probe_runtime( + configurations: dict[str, dict[str, Any]], + *, + instrumentation=None, + output_dir: str | Path, + execute: bool = False, + resume: bool = False, +) -> dict[str, Any]: + """Preview or capture all selected exact tuples; preserve every attempt.""" + + from .runtime_instrumentation import freeze_instrumentation, load_instrumentation + + if not isinstance(configurations, dict) or not configurations: + raise ValueError("runtime probe requires at least one selected configuration") + if instrumentation is None: + from .bundled_instrumentation import bundled_instrumentation + + versions = {item.get("identity", {}).get("framework_version") for item in configurations.values()} + if len(versions) != 1: + raise ValueError("bundled instrumentation selection requires one pinned runtime version") + version = next(iter(versions)) + instrumentation = bundled_instrumentation(version) + if instrumentation is None: + raise ValueError( + f"no bundled runtime observer for vLLM {version}; provide campaign-local --instrumentation" + ) + bundle = load_instrumentation(instrumentation) if isinstance(instrumentation, (str, Path)) else instrumentation + root = Path(output_dir).expanduser().resolve() + if root.exists() and any(root.iterdir()) and not (root / "probe-plan.json").exists(): + raise ValueError("runtime probe output must be fresh or contain this probe's saved plan") + root.mkdir(parents=True, exist_ok=True) + results: dict[str, Any] = {} + with _probe_lock(root), runner._sigterm_as_interrupt(): + index_path = root / "observations.json" + if index_path.exists() and execute and not resume: + raise ValueError("runtime probe observations already exist; use --resume or a fresh output directory") + index = ( + json.loads(index_path.read_text()) + if index_path.exists() + else {"schema_version": INDEX_SCHEMA, "configurations": {}} + ) + if index.get("schema_version") != INDEX_SCHEMA or not isinstance(index.get("configurations"), dict): + raise ValueError("runtime probe observation index has an unsupported schema") + plans = {} + for configuration, facts in configurations.items(): + try: + if not isinstance(configuration, str) or not configuration: + raise ValueError("configuration labels must be nonempty strings") + if facts.get("deployment", {}).get("executor") == "slurm": + facts = copy.deepcopy(facts) + facts["deployment"] = resolve_probe_cpu_policy( + configuration, facts["deployment"], root, resume=resume + ) + plan = build_runtime_probe_plan(configuration, facts, bundle) + existing = index["configurations"].get(configuration) + if existing is not None and existing["launch"] != plan.launch: + raise ValueError("launch settings changed; use a fresh probe output directory") + plans[configuration] = plan + except (OSError, TypeError, ValueError) as error: + results[configuration] = {"status": "failed", "diagnostics": [str(error)]} + runner._atomic_json( + root / "probe-plan.json", + { + "schema_version": "aisimulate-runtime-probe-plan/v1", + "configurations": {name: plan.launch for name, plan in plans.items()}, + "bundle_sha256": bundle.sha256, + }, + ) + for configuration, plan in plans.items(): + slug = _probe_configuration_slug(configuration) + result = results[configuration] = { + "status": "preview" if not execute else "completed", + "minimum_gpus": plan.cells[0].topology.total_gpus, + "launch": plan.launch, + "phases": {}, + "diagnostics": [], + } + collection = plan.launch["collection"] + advisory = ( + prefill_graph_advisory( + collection["prefill_cudagraph_policy"], enforce_eager=collection["enforce_eager"] + ) + if plan.launch.get("worker_type") is None + else None + ) + if advisory: + result["diagnostics"].append(advisory) + entry = index["configurations"].setdefault(configuration, {"launch": plan.launch, "attempts": []}) + previous = entry["attempts"][-1] if entry["attempts"] else None + if execute and resume and previous is not None: + try: + if _recorded_attempt_valid(root, previous, bundle, plan.options.workload_kinds): + result.update( + status="completed", + resumed=True, + attempt_id=previous["attempt_id"], + phases=copy.deepcopy(previous["phases"]), + ) + continue + except (OSError, TypeError, KeyError, ValueError) as error: + result["diagnostics"].append(f"previous attempt cannot be reused: {error}") + if execute and plan.options.executor == "slurm" and plan.options.slurm_cpus_per_task is None: + result["status"] = "failed" + result["diagnostics"].append( + "saved Slurm campaign has no frozen CPU policy; compatible recorded observations remain " + "readable, but new workers require a fresh campaign and smoke with explicit CPU settings" + ) + continue + if execute and resume and previous is not None: + try: + _cleanup_previous_attempt(root, plan, previous) + except Exception as error: + result["status"] = "failed" + result["diagnostics"].append(f"previous attempt cleanup failed: {error}") + continue + attempt_id = uuid.uuid4().hex + attempt_dir = root / ("attempts" if execute else "previews") / slug / attempt_id + attempt_dir.mkdir(parents=True) + frozen = freeze_instrumentation(bundle, attempt_dir / "instrumentation") + attempt = { + "attempt_id": attempt_id, + "status": "running", + "started_at": runner._utc_now(), + "bundle": {"manifest": str(frozen.manifest_path.relative_to(root)), "sha256": frozen.sha256}, + "phases": {}, + } + if execute: + entry["attempts"].append(attempt) + entry["active_attempt_id"] = attempt_id + runner._atomic_json(index_path, index) + for cell in plan.cells: + phase_dir = attempt_dir / cell.workload_kind + phase_dir.mkdir() + phase = attempt["phases"][cell.workload_kind] = {"status": "running", "artifacts": []} + resource = None + try: + runner._render_cell(plan, cell, phase_dir, probe_generator_overrides(plan.launch)) + context = launch_context(plan, cell, configuration=configuration, attempt_id=attempt_id) + extras = stage_runtime_instrumentation(frozen, phase_dir, context) + phase["launch_manifest"] = { + "path": str((phase_dir / CONTEXT_FILENAME).relative_to(root)), + **runner._file_metadata(phase_dir / CONTEXT_FILENAME), + } + phase["resource_manifest"] = _resource_manifest_reference(root, plan, cell, phase_dir) + phase["generator_request"] = str(phase_dir / "generator-request.json") + if execute: + runner._atomic_json(index_path, index) + manifest = phase_dir / runner.FPM_MANIFEST_FILENAME + resource = runner._cell_runner(plan, cell, manifest, phase_dir) + resource.cleanup() + resource.apply() + pods = resource.wait_ready(runner._expected_nodes(manifest)) + runtime = Path(runner.__file__).parent / "runtime" + resource.stage( + pods, + [ + phase_dir / runner.FPM_RUN_SCRIPT_FILENAME, + phase_dir / runner.FPM_ENV_FILENAME, + phase_dir / runner.RUNTIME_ENV_FILENAME, + runtime / "fpm_exec.sh", + runtime / "preflight.py", + *([runtime / "fpm_memory_observer.py"] if plan.options.executor == "slurm" else []), + *extras, + ], + ) + resource.prepare_attempt( + pods, + cell_id=cell.cell_id, + plan_sha256=plan.sha256, + attempt_id=attempt_id, + expected_backend_version=plan.capability.aic_database_version, + ) + resource.execute(pods) + resource.collect(pods, require_benchmark=False) + phase["status"] = "completed" if execute else "preview" + except KeyboardInterrupt: + phase["status"] = "interrupted" + result["status"] = attempt["status"] = "interrupted" + if resource is not None: + runner._salvage_artifacts(resource, cell.cell_id) + raise + except Exception as error: + phase.update(status="failed", error=str(error), error_type=type(error).__name__) + result["status"] = "failed" + result["diagnostics"].append(f"{cell.workload_kind}: {error}") + if resource is not None: + runner._salvage_artifacts(resource, cell.cell_id) + finally: + if resource is not None: + try: + resource.cleanup() + except Exception as error: + phase.update(status="cleanup_failed", cleanup_error=str(error)) + result["status"] = "failed" + result["diagnostics"].append(f"{cell.workload_kind} cleanup: {error}") + try: + phase["artifacts"] = _snapshot_probe_artifacts(root, phase_dir) + except (OSError, ValueError) as error: + phase.update(status="failed", capture_error=str(error)) + result["status"] = "failed" + result["diagnostics"].append(f"{cell.workload_kind} evidence capture: {error}") + result["phases"][cell.workload_kind] = copy.deepcopy(phase) + if execute: + attempt["status"] = result["status"] if result["status"] != "completed" else "running" + runner._atomic_json(index_path, index) + attempt.update(status=result["status"], completed_at=runner._utc_now()) + result["attempt_id"] = attempt_id + if execute: + runner._atomic_json(index_path, index) + failed = sum(result["status"] == "failed" for result in results.values()) + status = "failed" if failed == len(results) else "partial" if failed else "completed" if execute else "preview" + return { + "status": status, + "observations_path": str(root / "observations.json") if execute else None, + "configurations": results, + } diff --git a/python/aisimulate/collector/fpm_forward/slurm.py b/python/aisimulate/collector/fpm_forward/slurm.py index aea47d688..de2b718ac 100644 --- a/python/aisimulate/collector/fpm_forward/slurm.py +++ b/python/aisimulate/collector/fpm_forward/slurm.py @@ -26,7 +26,18 @@ class SlurmCellRunner: - def __init__(self, manifest: Path, cell_dir: Path, *, image: str, mounts: tuple[str, ...], total_gpus: int): + def __init__( + self, + manifest: Path, + cell_dir: Path, + *, + image: str, + mounts: tuple[str, ...], + total_gpus: int, + cpus_per_task: int | None = None, + cpu_bind: str | None = None, + attention_tp: int = 1, + ): from .runner import _expected_nodes self.cell_dir = cell_dir.resolve() @@ -41,6 +52,15 @@ def __init__(self, manifest: Path, cell_dir: Path, *, image: str, mounts: tuple[ self.job_id = os.environ.get("SLURM_JOB_ID", "") if not re.fullmatch(r"[0-9]+", self.job_id): raise ValueError("Slurm FPM must run inside an existing sbatch/salloc allocation") + if (cpus_per_task is not None or cpu_bind is not None) and ( + type(cpus_per_task) is not int or cpus_per_task < 1 or cpu_bind not in {"cores", "none"} + ): + raise ValueError("Slurm CPU policy requires positive cpus_per_task and cpu_bind cores or none") + local_dp = max(1, self.gpus_per_node // attention_tp) + if cpus_per_task is not None and cpus_per_task < local_dp: + raise ValueError(f"Slurm CPU pool {cpus_per_task} is smaller than {local_dp} local DP schedulers") + self.cpus_per_task = cpus_per_task + self.cpu_bind = cpu_bind self.step_name = f"fpm-{hashlib.sha256(str(self.cell_dir).encode()).hexdigest()[:20]}" # Keep ownership outside the replaceable cell payload so a fresh # invocation can tear down an abandoned allocation's named steps. @@ -81,12 +101,20 @@ def _command(self, args: list[str], *, timeout: float = 60, check: bool = True): raise def apply(self) -> None: + self._require_cpu_policy() for executable in ("srun", "scontrol", "squeue", "scancel"): if not shutil.which(executable): raise RuntimeError(f"Slurm FPM requires {executable}") self.owner_path.parent.mkdir(parents=True, exist_ok=True) self.owner_path.write_text(json.dumps({"job_id": self.job_id, "step_name": self.step_name}) + "\n") + def _require_cpu_policy(self) -> None: + if self.cpus_per_task is None or self.cpu_bind is None: + raise ValueError( + "saved Slurm campaign has no frozen CPU policy; CPU-only recovery remains available, " + "but new workers require a fresh campaign and smoke with explicit CPU settings" + ) + def wait_ready(self, expected_nodes: int, timeout_seconds: float = 900) -> list[str]: if expected_nodes != self.node_count: raise ValueError("Slurm expected node count disagrees with the generated manifest") @@ -118,12 +146,88 @@ def remaining() -> float: remaining() if len(hosts) != self.node_count or len(set(hosts)) != self.node_count: raise ValueError(f"FPM Slurm cell requires exactly {self.node_count} allocated nodes, got {hosts}") + self._qualify_cpu_allocation(hosts, timeout=remaining()) self.hosts = hosts return self.pods() if last_state not in {"PENDING", "CONFIGURING", "SUSPENDED"}: raise RuntimeError(f"Slurm allocation {self.job_id} cannot become ready from {last_state}") time.sleep(min(1.0, remaining())) + def _qualify_cpu_allocation(self, hosts: list[str], *, timeout: float) -> None: + """Use allocated CPU counts, never node-wide hardware capacity.""" + self._require_cpu_policy() + counts = [] + raw = os.environ.get("SLURM_JOB_CPUS_PER_NODE") + source = "SLURM_JOB_CPUS_PER_NODE" + if raw is not None: + for part in raw.split(","): + match = re.fullmatch(r"([1-9][0-9]*)(?:\(x([1-9][0-9]*)\))?", part) + if match is None: + raise ValueError("invalid allocated CPU counts in SLURM_JOB_CPUS_PER_NODE") + repeats = int(match.group(2) or 1) + if repeats > len(hosts): + raise ValueError("allocated CPU counts do not match Slurm nodes") + counts.extend([int(match.group(1))] * repeats) + else: + source = "scontrol show job --details CPU_IDs" + deadline = time.monotonic() + timeout + + def remaining() -> float: + budget = deadline - time.monotonic() + if budget <= 0: + raise TimeoutError("Slurm CPU allocation inspection exceeded the readiness deadline") + return budget + + raw = self._command( + ["scontrol", "show", "job", self.job_id, "--details"], + timeout=remaining(), + ).stdout + if re.search(r"(?:^|\s)JobId=" + re.escape(self.job_id) + r"(?:\s|$)", raw) is None: + raise ValueError("Slurm CPU allocation response does not identify the owned job") + # Identically allocated nodes may share one Nodes=host[1-2] entry. + # Expand that node set and bind counts by name, never row position. + masks = re.findall(r"(?:^|\s)Nodes=(\S+)\s+CPU_IDs=([0-9,-]+)(?:\s|$)", raw) + allocated = {} + for nodelist, mask in masks: + ids = set() + for segment in mask.split(","): + match = re.fullmatch(r"([0-9]+)(?:-([0-9]+))?", segment) + if match is None: + raise ValueError("invalid Slurm allocated CPU_IDs") + start, end = int(match.group(1)), int(match.group(2) or match.group(1)) + if end < start or end - start > 1048576: + raise ValueError("invalid Slurm allocated CPU_IDs range") + ids.update(range(start, end + 1)) + grouped_hosts = self._command( + ["scontrol", "show", "hostnames", nodelist], + timeout=remaining(), + ).stdout.split() + if not grouped_hosts or any(host in allocated or host not in hosts for host in grouped_hosts): + raise ValueError( + "Slurm allocated CPU node groups are missing, duplicated or outside the allocation" + ) + allocated.update({host: len(ids) for host in grouped_hosts}) + if set(allocated) != set(hosts): + raise ValueError("could not establish allocated CPUs for every Slurm node") + counts = [allocated[host] for host in hosts] + if len(counts) != len(hosts): + raise ValueError("could not establish allocated CPUs for every Slurm node") + report = { + "job_id": self.job_id, + "cpus_per_task": self.cpus_per_task, + "cpu_bind": self.cpu_bind, + "source": source, + "allocated_cpus_per_node": dict(zip(hosts, counts, strict=True)), + } + from .runner import _atomic_json + + _atomic_json(self.cell_dir / "slurm-cpu-allocation.json", report) + if any(count < self.cpus_per_task for count in counts): + raise ValueError( + f"Slurm allocation has insufficient CPUs for --cpus-per-task={self.cpus_per_task}: " + f"{report['allocated_cpus_per_node']}; request matching sbatch/salloc CPUs" + ) + def pods(self, *, include_terminating: bool = True) -> list[str]: del include_terminating return [f"node{rank:04d}" for rank in range(len(self.hosts))] @@ -137,6 +241,7 @@ def stage(self, pods: list[str], files: list[Path]) -> None: (self.cell_dir / "raw" / unit).mkdir(parents=True, exist_ok=True) def _exec(self, unit: str, command: list[str], *, timeout: int): + self._require_cpu_policy() rank = self.pods().index(unit) mounts = [ *self.mounts, @@ -155,6 +260,9 @@ def _exec(self, unit: str, command: list[str], *, timeout: int): "--nodes=1", "--ntasks=1", "--ntasks-per-node=1", + f"--cpus-per-task={self.cpus_per_task}", + f"--cpu-bind={self.cpu_bind}", + "--immediate=10", "--exclusive", "--exact", f"--nodelist={self.hosts[rank]}", @@ -163,9 +271,13 @@ def _exec(self, unit: str, command: list[str], *, timeout: int): f"--container-mounts={','.join(mounts)}", "--container-writable", "--container-workdir=/tmp/fpm-bench", - "env", + # Avoid Slurm resolving a bare entrypoint through inaccessible PATH entries. + "/usr/bin/env", f"FPM_NODE_RANK={rank}", f"FPM_MASTER_ADDR={self.hosts[0]}", + f"FPM_SLURM_CPUS_PER_TASK={self.cpus_per_task}", + f"FPM_SLURM_CPU_BIND={self.cpu_bind}", + f"FPM_LOCAL_GPU_COUNT={self.gpus_per_node}", *command, ], timeout=timeout, diff --git a/python/aisimulate/docs/add_a_new_model.md b/python/aisimulate/docs/add_a_new_model.md index b2e13314c..a2dfdbc39 100644 --- a/python/aisimulate/docs/add_a_new_model.md +++ b/python/aisimulate/docs/add_a_new_model.md @@ -11,6 +11,12 @@ All paths below are relative to the repository root unless linked otherwise. The `aisimulate_core` source namespace is retained inside the `aisimulate` wheel; installing or rebuilding a separate AIConfigurator package is unnecessary. +For whole-forward FPM onboarding, start with the [self-service config/profile workflow](../../../docs/fpm-self-service.md#onboard-with-an-agent), +which supports direct interpolation without an analytical model class. The +[FPM model-integration guide](fpm/model-integration.md) records shared metadata +requirements and the optional registered-model/SOL procedure to use alongside +this guide. Neither FPM route requires per-operation GPU timing collection. + ## Choose the smallest extension | What changed? | Start here | @@ -103,7 +109,7 @@ python/aisimulate/src/aisimulate_core/systems/data////< Do not relabel another backend version's data as newly measured. Update query version/support metadata only when the new cell meets its collection and consumer contracts. For whole-forward profiles, use the separate -[Self-benchmarking and FPM onboarding guide](fpm/self-benchmarking-and-onboarding.md). +[FPM self-benchmarking and onboarding guide](fpm/self-benchmarking-and-onboarding.md). ## 4. Validate the intended public path diff --git a/python/aisimulate/docs/fpm/README.md b/python/aisimulate/docs/fpm/README.md index fbeb0abea..bb16278b9 100644 --- a/python/aisimulate/docs/fpm/README.md +++ b/python/aisimulate/docs/fpm/README.md @@ -5,11 +5,10 @@ SPDX-License-Identifier: Apache-2.0 # Forward-pass models -Start with the [self-benchmark and FPM onboarding guide](self-benchmarking-and-onboarding.md). -Its seven steps cover support checks, measurement planning, collection, profile -validation, canonical model construction, Replay integration, and accuracy -validation. It explains which engine-iteration differences self-collection can -capture and which behaviors, including PP, need separate simulator support. +Start with the [FPM self-benchmarking and onboarding guide](self-benchmarking-and-onboarding.md) to collect +whole-forward measurements, publish a performance-data pair, load it through +the SDK, and run an AISimulate prediction. The guide includes prerequisites, +commands, expected artifacts, acceptance checks, and recovery steps. Self-benchmark collection currently supports vLLM configurations that pass model/runtime validation; new architectures can require benchmark adaptation. diff --git a/python/aisimulate/docs/fpm/aic-fpm-modeling-plan.md b/python/aisimulate/docs/fpm/aic-fpm-modeling-plan.md index 98c0cbb0c..0783364f4 100644 --- a/python/aisimulate/docs/fpm/aic-fpm-modeling-plan.md +++ b/python/aisimulate/docs/fpm/aic-fpm-modeling-plan.md @@ -18,7 +18,7 @@ consumes the **as-built** contract, not the planned one: | Topic | Original plan | As built (authoritative) | |---|---|---| | Physical axes | per-request `(B,S,P)` / `(B,L)` | per-DP-rank iteration totals `(batch_size, total_prefill_tokens, total_kv_read_tokens)`, `partition_policy=balanced_v1` | -| Point grid | collector-owned Halton/maximin, frozen holdout IDs | Dynamo PR11509 native self-benchmark owns the grid; runtime-admitted, no holdout set | +| Point grid | collector-owned Halton/maximin, frozen holdout IDs | AISimulate sets runtime limits and runtime/explicit capture policy; runtime policy defers captures to engine initialization, while explicit policy supplies the capture extension and derived sample caps. Dynamo native self-benchmark combines these inputs with initialized engine state, image sampling defaults and feasibility checks to generate the exact grid; no holdout set | | Repeats | 5 repeats, median, CV/quarantine | single sample per point (`dynamo_native_single_sample_v1`); DP max is baked into `latency_ms` | | Publication | candidate dir + human promotion | collector writes the formal pair directly (locked, atomic, conflict-fail merge) | | Sidecar | predictor ID, block size, per-axis min/max | schema v6: hashes, run identities, counts — **no predictor/domain fields** | diff --git a/python/aisimulate/docs/fpm/model-integration.md b/python/aisimulate/docs/fpm/model-integration.md new file mode 100644 index 000000000..faf632fb4 --- /dev/null +++ b/python/aisimulate/docs/fpm/model-integration.md @@ -0,0 +1,335 @@ + + +# Prepare model metadata and choose an FPM execution route + +FPM onboarding needs pinned model and deployment metadata, correct memory and +KV-cache accounting, and matching whole-forward timing data. The current +config/profile route supports ordinary FPM `predict` and `recommend` without +an op-level model class. Registered models can also use analytical SOL estimates +for timing transfer. Neither route requires per-operation GPU timing collection. + +Start with the [FPM self-service workflow](../../../../docs/fpm-self-service.md#onboard-with-an-agent) +to inspect a model, choose a worker, derive and review its resource profile, +plan and collect timings, and validate replay. This guide records the metadata +shared by both routes, then describes the **optional registered-model/SOL +integration procedure** in sections 2–5. Its CPU example uses the bundled +**Qwen3-0.6B configuration**; it does not establish coverage for another model. + +## Class-independent direct FPM + +Supply a local model config to `aisimulate onboard init --model-config`, or +provide an existing `--fpm-profile`. Config-based setup derives supported +resource estimates and identifies missing bounds; review and edit the exact +values, effective precisions and provenance before accepting the profile. +Follow the self-service guide's [terminal or headless review flow](../../../../docs/fpm-self-service.md#3-derive-review-and-save-the-profile). +Profile resources must bound every rank of the selected topology. A profile is +a declaration, not proof of runtime compatibility or measured memory fit. +For full/sliding attention and supported convolution state, follow +[grouped cache review](../../../../docs/fpm-self-service.md#review-grouped-cache-resources): +runtime block sizes remain explicit inputs, and each group page includes every +group layer plus runtime padding on one rank. Use the +[canonical byte-budget API](../../../../docs/core-api.md#fpm-profile-cache-groups-and-byte-budgets) +for grouped resources; a scalar token capacity cannot represent window eviction +or transient prefill pages. + +An existing launch configuration is optional evidence. The onboarding agent +proposes supported runtime settings from pinned checkpoint/runtime metadata and +available sidecars, asks only for unresolved facts, and presents the result for +review. The CLI does not import arbitrary launch arguments or inspect a remote +runtime automatically. Config intake proposes the current collector communication +identity `half`; unknown FMHA and KV precision remain explicit inputs. Packed +cache-page estimates remain distinct from runtime allocations including padding. + +The saved request embeds the profile, and generated ordinary configurations use +`engine.fpm_profile`, `estimation_mode: fpm_interpolation`, +`estimator_config.fpm_interpolation.method: direct`, and `fallback_policy: deny`. +Direct interpolation uses measured whole-forward timings without constructing +an operation graph for timing or resources. Unsupported metadata and uncovered +queries fail explicitly. See [execution-route selection and interpolation rules](../../../../docs/fpm-self-service.md#choose-the-model-execution-route) +for exact-point, curve and two-sided interpolation coverage. + +Each generated profile and plan selects one exact TP, DEP or TEP worker. To onboard +several configurations in one session, use the self-service guide's +[directory output](../../../../docs/fpm-self-service.md#onboard-multiple-parallel-configurations). +It reuses shared intake and creates a separate reviewed profile and request +for each tuple; rank-local byte bounds and cache groups are never transferred +between configurations. Follow each emitted plan command and collect or resume +each configuration independently. Follow the +[shared collection policy](../../../../docs/fpm-self-service.md#how-the-collection-grid-is-determined): +AISimulate sets runtime limits and the reviewed capture policy. New onboarding +uses `--prefill-cudagraph-policy runtime`, which leaves prefill compilation to +the pinned engine; explicit capture extension remains available. The collector +launches benchmark workers, and Dynamo uses initialized engine state, image +sampling defaults and feasibility checks to generate and time the exact grid. +Capture sizes and exact counts remain unresolved before engine initialization +in runtime mode. A complete generated grid does +not establish direct-FPM query coverage. Validation traffic is supplied +separately after a formal timing pair is verified. The +current [AgentX coverage check](../../../../docs/fpm-self-service.md#validate-fpm-query-coverage-with-agentx-replay) +uses cold aggregated replay, one client lane, HBM-only cache and no speculative +decoding. Missing timing stops replay and retains partial evidence; that evidence +cannot certify the rest of the trace. Coverage and measured accuracy are +separate results. + +## 1. Record the intended deployment + +Keep these inputs with the integration issue and eventual collection artifacts: + +| Input | What to inspect and record | +| --- | --- | +| Checkpoint | Canonical model ID, immutable revision, local `config.json` and optional quantization/processor metadata, and file hashes. AISimulate's config parser does not take a revision argument; use metadata from the pinned checkpoint in a local directory when an exact revision is needed. | +| Architecture | Layer types/counts, attention and KV heads, head dimension, FFN/expert dimensions, routing, shared experts, and any model-specific configuration fields. Compare these with the serving implementation. | +| Memory | Weight storage precision and replication/sharding; cache layout, quantization, sliding windows/compression, and fixed recurrent or decode state. | +| Runtime | GPU/system specification, backend and exact version, image digest, attention/MoE kernel choices, and effective GEMM/MoE/FMHA/KV/communication precision. Checkpoint weight precision alone does not specify all these values. | +| Worker topology | Exact `(TP, PP, attention DP, MoE TP, MoE EP, CP)` tuple. Minimum collection GPUs are attention TP times attention DP. Total GPU allocation, node reservations and replica budgets are not onboarding intake requirements. | +| Runtime and collection bounds | Per-request context, per-attention-DP-rank scheduled-token and sequence limits, GPU memory fraction (new onboarding starts at 0.90), and runtime or explicit prefill CUDA graph policy. Review these independently of replay traffic; record actual captures after initialization. These settings do not prove memory fit or timing coverage. | +| Validation traffic | Select a local trace when validating the collected pair. Fixed input/output lengths, concurrency, TTFT and TPOT are optional synthetic-example inputs, not collection requirements. | + +For example, MoE TP4 is `(4, 1, 1, 4, 1, 1)`, DEP8 is +`(1, 1, 8, 1, 8, 1)`, and TEP8 is `(8, 1, 1, 1, 8, 1)`. Equal GPU counts do +not make their resource bounds or timing cells interchangeable. Use the +[topology flags and collection limits](../../../../docs/fpm-self-service.md#create-the-request) +for the intended worker, and verify actual collection resources before execution. +The sequence limit does not reserve maximum context for every sequence. + +The profile-based collection workflow targets vLLM text decoders with +`PP=CP=1` and linear or grouped cache storage. Grouped prediction and replay +currently require cold aggregated execution, HBM-only cache, no speculative +decoding and `prefix_caching: false`; generated grouped configurations preserve +that setting. A multimodal checkpoint can supply an +unambiguous `text_config` or flat decoder fields. Text-decoder timing and +config-derived estimates exclude encoders, projectors, preprocessing and other +non-text components. Observed runtime cache capacity accounts for every component +actually loaded by the worker; do not subtract guessed encoder allocations. +Preserve that scope in profile provenance. Unsupported cache semantics, encoder +pools, AFD, speculative decoding, wide EP and EPLB remain outside this route. A +registered class existing for another mode is not evidence of support for the +intended FPM deployment. + +## 2. Registered-model route: reuse or implement the model description + +Follow [How to add a new model](../add_a_new_model.md) for the registry and native +operation contracts. For this registered-model/SOL route, review these concrete responsibilities: + +| Responsibility | Source and required result | +| --- | --- | +| Parse the checkpoint | [`sdk/utils.py`](../../src/aisimulate_core/sdk/utils.py), especially `get_model_config_from_model_path()`. Confirm the parsed geometry and architecture-specific `extra_params`; add parsing only for fields the existing parser cannot represent correctly. | +| Select a family | [`sdk/common.py`](../../src/aisimulate_core/sdk/common.py), `ARCHITECTURE_TO_MODEL_FAMILY`. Reuse a family only when its operation pipeline, precision handling, and cache behavior match the model. A new architecture name alone does not require a dedicated class. | +| Construct a new family when needed | The [registry examples](../../src/aisimulate_core/sdk/models/README.md#adding-a-new-model) show `@register_model`, `create()`, and reusable blocks. Implement the actual `context_ops` and `generation_ops`; an empty placeholder class is insufficient. | +| Describe analytical work | Each operation needs correct dimensions, quantization, layer/repetition scaling, and per-rank parallelism. Preserve the context/generation attention and `logits_gemm` naming contracts. Inspect communication, overlap, and phase-specific behavior as well as GEMMs. | +| Describe memory | Operation `get_weights()` values provide the weight inventory. Model `get_kvcache_bytes_per_sequence()` and its inverse `get_kvcache_max_tokens()` must represent the real cache curve. Override both for non-linear/windowed/compressed layouts; do not extrapolate a one-token slope across a window boundary. | +| Support FPM analytical execution | Every retained operation must be supported by the Rust [`fpm_sol` evaluator](../../../../crates/core/src/perfmodel/operators/fpm_sol.rs) for the proposed deployment and shapes. Adding an ordinary native operation or a `SOL_FULL` diagnostic does not automatically add this support. | + +Construct performance models through +`RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig)` and query +the returned model, as specified by the [canonical API](../../../../docs/core-api.md#choosing-a-forward-pass-api). +For registered FPM interpolation, select `estimation_mode: fpm_interpolation`, +`estimator_config.fpm_interpolation.method: sol`, and `fallback_policy: deny`. +The constructor requires a compatible registered class and a matching FPM pair. +Internally, the ordinary phase operations are retained as `sol_ops` inside +`FPMForwardOp`; their context weight inventory becomes `weight_bytes`. The model +class and cache methods remain available. Model authors should not duplicate +this phase replacement or implement a Python latency/SOL formula. + +## 3. Run CPU checks before collecting timings + +From the repository root, activate the installed [development environment](../../../../DEVELOPMENT.md) +and choose a fresh evidence directory: + +```bash +source python/aisimulate/.venv/bin/activate +export FPM_CHECK_DIR="$(mktemp -d)" +``` + +The following example needs no model weights, GPU, or matching Qwen FPM pair. It +uses bundled model metadata and the H200 system/backend catalog. The memory +entry point still requires a resolvable catalog/database layout, even though +memory sizing does not query measured operation latencies. The example pins a +queryable backend version; select a version allowed by your installed +[`query_versions.yaml`](../../src/aisimulate_core/systems/query_versions.yaml) +when adapting it. + +```bash +python - <<'PY' +import hashlib +import json +import math +import os +from pathlib import Path + +from aisimulate_core.sdk import ( + ForwardPassPerfModelConfig, + ModelConfig, + RustForwardPassPerfModel, + estimate_kv_cache, +) +from aisimulate_core.sdk.engine import build_ops_json +from aisimulate_core.sdk.models import get_model +from aisimulate_core.sdk.utils import get_model_config_from_model_path + +out = Path(os.environ["FPM_CHECK_DIR"]) +model_path, system, backend, version = "Qwen/Qwen3-0.6B", "h200_sxm", "vllm", "0.24.0" +topology = dict(tp_size=1, pp_size=1, attention_dp_size=1, moe_tp_size=1, moe_ep_size=1) +options = ModelConfig(**topology) +info = get_model_config_from_model_path(model_path) +model = get_model(model_path, options, backend) + +phase_ops = {"prefill": model.context_ops, "decode": model.generation_ops} +original = {phase: json.loads(build_ops_json(ops)) for phase, ops in phase_ops.items()} +assert all(original.values()) +weights = sum(op.get_weights() for op in model.context_ops) +assert weights > 0 + +kv_bytes = {s: model.get_kvcache_bytes_per_sequence(s) for s in (1, 128, 8192)} +for s, budget in kv_bytes.items(): + assert budget > 0 + capacity = model.get_kvcache_max_tokens(budget) + assert model.get_kvcache_bytes_per_sequence(capacity) <= budget + assert model.get_kvcache_bytes_per_sequence(capacity + 1) > budget + +precision = {k: getattr(model.config, k).name for k in ( + "gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "kvcache_quant_mode", "comm_quant_mode" +)} +memory = estimate_kv_cache( + model_path, system, backend, version, **topology, **precision, + max_num_tokens=2048, max_batch_size=8, + memory_fraction_kind="of_total", memory_fraction_value=0.9, + allow_naive_fallback=False, +) +assert memory["source"] == "native" and memory["total_kv_size_tokens"] > 0 + +# These are op-level analytical diagnostics, not FPM interpolation. +sol = {} +for phase in phase_ops: + config = ForwardPassPerfModelConfig( + model=model_path, system=system, backend=backend, backend_version=version, + worker_type=phase, tp=options.tp_size, pp=options.pp_size, + attention_dp=options.attention_dp_size, + moe_tp_size=options.moe_tp_size, moe_ep_size=options.moe_ep_size, + **precision, estimation_mode="op_level", database_mode="SOL", fallback_policy="deny", + ) + performance_model = RustForwardPassPerfModel.best_available(config) + rows = performance_model.static_phase_diagnostics( + batch_size=1, context_length=128, prefill=phase == "prefill", + ) + assert rows and all(row["details"]["sol"] is not None for row in rows) + assert all(math.isfinite(v) and v >= 0 for row in rows for v in row["details"]["sol"].values()) + assert sum(row["details"]["sol"]["latency_ms"] for row in rows) > 0 + sol[phase] = rows + +report = { + "model_path": model_path, "class": type(model).__name__, "family": model.model_family, + "system": system, "backend": backend, "backend_version": version, "topology": topology, + "geometry": {k: info[k] for k in ("architecture", "layers", "n", "n_kv", "d", "hidden_size")}, + "resolved_config_sha256": hashlib.sha256( + json.dumps(info["raw_config"], sort_keys=True).encode() + ).hexdigest(), + "quantization": precision, + "op_counts": {phase: len(ops) for phase, ops in phase_ops.items()}, + "weights_bytes_per_rank": weights, "kv_bytes_per_sequence_per_rank": kv_bytes, + "memory": memory, "op_level_diagnostics": sol, +} +for name, value in (("integration.json", report), ("op-level.json", original)): + (out / name).write_text(json.dumps(value, indent=2) + "\n") +print(json.dumps(report, indent=2)) +print(f"Saved integration evidence to {out}") +PY +``` + +**Expected for this example:** `LLAMAModel` / `LLAMA`, 28 layers, 16 attention +heads, 8 KV heads, head dimension 128, and 14 operation entries per phase. +The modeled weight inventory is 1,503,133,696 bytes per rank. BF16 KV costs +114,688 bytes per token; an 8,192-token sequence uses 939,524,096 bytes per rank. +The diagnostic produces finite nonnegative per-operation SOL values, including +zero-cost communication at TP=PP=1. Prefill processes 128 tokens; decode uses +128 past tokens and one new token. These are analytical estimates, not measured +latency or peak runtime memory. They do not exercise registered FPM timing +transfer. `resolved_config_sha256` hashes parsed metadata, not checkpoint weights. + +For your integration, replace the example model and deployment inputs together. +Compare the emitted operation dimensions and scale factors with the serving +architecture, and add independent expected-value tests for those facts. Exercise +each supported topology/precision and the relevant cache/window boundaries; +three positive sequence lengths and a successful import are not sufficient +model-specific evidence. See the existing [model configuration and memory tests](../../tests/unit/sdk/models/test_model_config.py) +for patterns. Include an unsupported architecture/topology case that must fail +explicitly, rather than falling back to another family or memory estimator. + +## 4. Verify the FPM execution contract + +For a registered-model integration, run its focused CPU tests together with the +existing [FPM graph and routing](../../tests/unit/sdk/test_fpm_forward.py), +[profile and canonical constructor](../../tests/unit/sdk/test_fpm_profile.py), +[operation serialization](../../tests/unit/sdk/test_opspec_coverage.py) and +[single-oracle](../../tests/cross_package/test_single_oracle_contract.py) +contract suites. Native [FPM SOL](../../../../crates/core/src/perfmodel/operators/fpm_sol.rs) +and [FPM forward](../../../../crates/core/src/perfmodel/operators/fpm_forward.rs) +tests cover analytical and transfer behavior. Follow the repository's +[test instructions](../../../../DEVELOPMENT.md#running-tests); on macOS, use +`-p no:timeout` as noted in the [repository guidance](../../../../AGENTS.md#cursor-cloud-specific-instructions). +Existing fixtures do not automatically cover a new model's operations or +demonstrate predictive accuracy. + +There is currently no public model-readiness call that evaluates an arbitrary +retained graph through the exact FPM `sol_total`/`fpm_sol` path without a pair. +`static_phase_diagnostics()` above uses a different diagnostic path and integer +workload coordinates. FPM transfer maps iteration totals to potentially +fractional per-request coordinates. Also, exact FPM hits and some linear +interpolations can succeed without ever evaluating `sol_ops`. + +Before declaring the new model ready for collection, add a model-specific CPU +regression that constructs the model through `best_available` with explicit +`fpm_interpolation`/`sol` and denied fallback, then queries deliberately synthetic +FPM cells at an off-site coordinate requiring SOL transfer. Cover both phases +and any relevant fractional coordinates. Use the synthetic pair layout in +[`test_fpm_forward.py`](../../tests/unit/sdk/test_fpm_forward.py), the canonical +constructor/query examples in [`test_fpm_profile.py`](../../tests/unit/sdk/test_fpm_profile.py), and the native +tests `prefill_off_site_kv_resolves_through_the_dsa_sol_transfer`, +`decode_pairless_batch_resolves_through_the_dsa_sol_transfer`, and +`unsupported_sol_family_is_lazy` in +[`fpm_forward.rs`](../../../../crates/core/src/perfmodel/operators/fpm_forward.rs) +as contract examples. Assert an independently justified result and a clear +failure for an unsupported shape or identity; an exact-hit-only fixture misses +this dependency. Keep synthetic timings in test fixtures, never publish them as +collected data. If an operation is unsupported, implement and numerically test +its Rust FPM SOL path before relying on sparse collection. + +## 5. Review, install, and hand off to collection + +Submit the model/parser/operation changes with the focused model tests, CPU +reports, resolved deployment inputs, and any explicit unsupported cases. Follow +the repository [review contract](../../../../REVIEW.md). After review, install +the approved AISimulate revision using the [development setup](../../../../DEVELOPMENT.md) +in the environment that will plan collection and run prediction. From that +checkout with its environment active: + +```bash +git rev-parse HEAD +python -c 'import importlib.metadata; print(importlib.metadata.version("aisimulate"))' +aisimulate onboard --help +aisimulate predict --help +aisimulate recommend --help +``` + +Record the commit as well as the package version; several source revisions may +share a version. Rerun the CPU checks in a fresh process against the installed +integration. Preserve the reports with the pinned config files and test output. + +The collection handoff consists of the canonical checkpoint identity/revision, +config files/hashes, reviewed AISimulate revision, exact runtime image/version, +hardware/topology/precision choices, runtime and collection limits, and the CPU evidence. +The metadata used for local construction must describe the same checkpoint +mounted in the collection runtime; a temporary local path is not a portable +replacement for the canonical identity in the published pair. + +For the config/profile route, continue the [self-service workflow](../../../../docs/fpm-self-service.md#plan-preview-and-explicitly-execute) +with the accepted profile. The [collection campaign example](self-benchmarking-and-onboarding.md#b2-freeze-and-inspect-the-plan-step-2) +illustrates guided collection and formal pair publication without requiring an analytical class. Only +whole-forward silicon timings are collected for FPM. Exercise ordinary `predict` +and `recommend` on covered candidates and retain the matching inputs, data +provenance, reports and reproduction commands. CPU integration checks, +collection completion, replay coverage, completed simulations and agreement +with independent silicon measurements are separate acceptance results. diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md index 4e27719a5..495ea283b 100644 --- a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -11,8 +11,12 @@ means running the engine's self-benchmark on your target configuration. **Onboarding** means validating and loading its measured profile through the canonical performance-model API, then using it for prediction. -The procedure has seven steps. The model, hardware, and engine settings are -inputs to that procedure. [Kimi K3 TP8+DCP8](#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile) +The procedure has seven technical steps. For an agent-led session, follow the +[six onboarding stages](../../../../docs/fpm-self-service.md#onboard-with-an-agent), +which also cover input discovery, review, acceptance, and checkpointing. The +`aisimulate onboard ...` commands below implement that guided path. This guide +then follows the data from collection through consumption and validation. +The model, hardware, and engine settings are inputs to that procedure. [Kimi K3 TP8+DCP8](#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile) and [MiniMax-M2.7 TP4](#example-b-collect-a-minimax-m27-tp4-profile) are worked examples at the end; their values are not defaults for other deployments. @@ -34,7 +38,9 @@ For example, collecting Kimi TP8+DCP8 measures that configuration's prefill and decode costs, including DCP communication inside the timing boundary. A TP8 profile without DCP is a different measurement identity and cannot substitute for those observations. If suitable measured data already exists, reuse it -starting at step 4 after checking its provenance and coverage. +starting at step 4 after checking its provenance, matching resource information, +and coverage. A measured timing table and an onboarding resource profile are +different artifacts; a timing table alone does not determine serving memory. ## What self-collection covers @@ -72,28 +78,32 @@ engine versions can require changes in: convolution state, per-request state allocation, and distributed cache layout. - **Scheduling and measurement hooks:** requested versus actually scheduled token counts, graph paths, completion checks, and per-rank timing aggregation. -- **AISimulate model construction:** architecture/configuration parsing, model - registration, resident weights, and memory/cache descriptions used by the - consumer. A collection plan's bootstrap template is not proof of consumer support. +- **AISimulate input and memory support:** architecture/configuration parsing, + precision identity, and cache geometry/capacity used by the consumer. Direct + FPM uses an explicit identity/resource profile without an op-level model class; + unknown runtime layouts still need an audited observation adapter or supported + resource declaration. A collection plan is not proof of consumer support. Use a small smoke run to qualify those paths before collecting the full surface. If it fails or silently schedules different work, fix and revalidate the -benchmark adapter before using the measurements. For an architecture absent -from AISimulate, follow [How to add a new model](../add_a_new_model.md) as well. -An already registered architecture needs no new class or parallelization mapping -merely to add another supported measured configuration. +benchmark adapter before using the measurements. For a new architecture, start from a +[local model config](../../../../docs/fpm-self-service.md#start-from-a-local-model-config) +and investigate unresolved precision/cache facts before asking the user. +[How to add a new model](../add_a_new_model.md) describes the separate registered +analytical-model route when op-level or SOL estimation is wanted. Neither a +dedicated class nor per-operation silicon data is required for direct FPM. ## Workflow at a glance | Step | Action | Required result | | --- | --- | --- | | [1](#1-check-support-and-prepare-the-environment) | Check backend/model support and prepare the environment | Compatible engine, benchmark, and AISimulate builds; known adaptation gaps | -| [2](#2-freeze-the-engine-configuration-and-measurement-plan) | Freeze configuration, timing semantics, and sampling coverage | Reproducible configuration and a plan covering the intended workloads | -| [3](#3-smoke-test-and-collect) | Smoke test, adapt if needed, then collect | Validated prefill/decode measurements plus raw evidence and provenance | -| [4](#4-validate-and-install-the-fpm-profile) | Validate the data contract and install the pair | A schema-v6 Parquet/metadata pair in a local systems root | -| [5](#5-construct-and-check-the-performance-model) | Construct the canonical model and check queries | Matching identity, ready estimator, and verified measured-point lookups | -| [6](#6-connect-the-model-to-replay) | Supply the same profile plus serving configuration | A reproducible prediction configuration and successful Replay smoke | -| [7](#7-validate-coverage-and-accuracy) | Check coverage and compare held-out measurements | Errors and limitations reported separately from calibration checks | +| [2](#2-freeze-the-engine-configuration-and-measurement-plan) | `onboard init`, review/accept, then `onboard plan` | Reviewed role/configuration and a saved plan; runtime readiness still unchecked | +| [3](#3-smoke-test-and-collect) | `onboard collect-fpm`: preview, smoke/readiness, then execute | Measurements for each selected role/phase plus raw evidence and provenance | +| [4](#4-validate-and-install-the-fpm-profile) | Inspect the pair and run `onboard validate-collection` | Matching Parquet/metadata pair and separate validity, repeatability, and holdout results | +| [5](#5-construct-and-check-the-performance-model) | `onboard finalize` when memory is pending; check canonical queries | Reviewed resolved resource profile, matching estimator, and checked lookups | +| [6](#6-connect-the-model-to-replay) | `onboard validate-fpm` and ordinary `predict`/`recommend` for aggregated profiles | Reproducible replay and coverage reports, separate from accuracy | +| [7](#7-validate-coverage-and-accuracy) | Assess matched serving with `onboard validate-serving` | Coverage, forward errors, and serving accuracy reported separately | With a validated profile from someone else, review steps 1–2 and start execution at step 4. New GPU collection follows all seven steps. A failure in runtime or @@ -104,52 +114,58 @@ resolve it. ## 1. Check support and prepare the environment -Use Python 3.11–3.13 and install AISimulate from the latest `main`. -For an existing checkout, run `git switch main` and start from the pull command: +Use Python 3.11–3.13 and a source revision containing the commands you need. +Keep your existing working checkout; switching branches can remove the guided +CLI. Follow the [installation guide](../../../../docs/installation.md#use-current-source) +when creating a new checkout. From its repository root: ```bash -git clone --branch main https://github.com/ai-dynamo/aisimulate.git -cd aisimulate -git pull --ff-only origin main uv sync --project python/aisimulate --extra dev source python/aisimulate/.venv/bin/activate -``` - -Source installation requires `uv`, Cargo/Rust, and a C/C++ compiler/linker. -See the [installation guide](../../../../docs/installation.md#use-current-source) -for platform requirements. No serving GPU or model weights are needed to query -an existing profile or run the engine-stack simulation. - -From the repository root, record the source and package in the activated session: - -```bash export AIS_REPO="$PWD" +export PYTHONPATH="$AIS_REPO/python/aisimulate${PYTHONPATH:+:$PYTHONPATH}" git rev-parse HEAD python -c 'from importlib.metadata import version; print(version("aisimulate"))' aisimulate predict --help ``` -Use the same activated Bash or Zsh session for the remaining commands. -For existing data, review its recorded identity and conditions in step 2, then -continue at step 4. For new measurements, continue through steps 2 and 3. - -For **new GPU collection**, additionally check: +For the guided workflow and Example B, additionally verify: ```bash -cd "$AIS_REPO/python/aisimulate" -python -m collector.fpm_forward --help +aisimulate onboard --help +aisimulate onboard init --help ``` -The collection host needs model configuration metadata and `kubectl` access. -The target runtime needs the model weights, allocated GPUs, and matching engine -and benchmark code. The Kubernetes Collector additionally needs an existing -namespace, checkpoint PVC/path, collection image, and Generator target release. -Record the image digest, model revision, GPU type, and resolved configuration. -Before launching, read [recovery and cleanup](#recovery-and-cleanup). +Source installation needs `uv`, Cargo/Rust, and a C/C++ compiler/linker. A missing +guided command is a build/version mismatch for that workflow. Existing-profile +consumers such as Example A can skip the `onboard` checks; they use the canonical +SDK and ordinary prediction in a compatible consumer build. Inspect each later +subcommand's help before using it. Keep the same activated Bash or Zsh session for the examples. No serving +GPU or model weights are needed to query existing timings or run simulation. +The Kimi example has [additional consumer-build prerequisites](#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile). + +For **new GPU collection**, also check `python -m collector.fpm_forward --help`. +The host needs local model configuration metadata; the target runtime needs the +pinned checkpoint, allocated GPUs, and matching engine/benchmark code. Kubernetes +needs `kubectl`, an existing namespace, model-cache PVC and deployment permissions. +Slurm needs a caller-owned `sbatch`/`salloc` allocation, Pyxis/Enroot, shared storage +and matching CPU/GPU resources. Neither path needs a prestarted HTTP server: +AISimulate launches benchmark workers and Dynamo self-benchmark generates their +measurements. Record the immutable image/model revisions and actual full runtime +version, including custom-build suffixes. See +[executor prerequisites](../../../../docs/fpm-self-service.md#choose-the-collection-executor) +and [recovery and cleanup](#recovery-and-cleanup). ## 2. Freeze the engine configuration and measurement plan -Define the target deployment before measuring it. Keep a record of: +Create one `onboard checkpoint` file outside fresh request/plan output roots, +then use `onboard init --model-config` (or an existing `--fpm-profile`) to prepare +the inputs. Resolve facts from the pinned checkpoint, hardware and runtime; +present supported precision options before requesting the user's choice. Review +the exact request, save edits and record acceptance as described in the +[checkpoint workflow](../../../../docs/fpm-self-service.md#preserve-partial-profile-review). +`onboard plan` saves the accepted request, collector command and simulation inputs; +it does not verify a target runtime. Keep a record of: - **Identity:** model/config/checkpoint revision, GPU and interconnect topology, engine version and image/source revision, TP/PP/attention-DP/MoE TP/EP, and DCP @@ -163,6 +179,20 @@ Define the target deployment before measuring it. Keep a record of: - **Coverage:** prefill request/new-token/prefix coordinates and decode batch/context coordinates needed by the intended Replay workload. +Choose `--worker-type aggregated` for one serving worker with both collection +phases. For P/D disaggregation, onboard `prefill` and `decode` independently; +each may have its own topology, graph policy, scheduler bounds and memory budget. +Collection initializes state locally and does not need an actual P/D transfer. +Explicit aggregated probes receive the same intended serving configuration; +vLLM may naturally resolve different per-batch prefill/decode graph paths. See +[serving roles](../../../../docs/fpm-self-service.md#serving-roles-and-cuda-graph-settings). + +AISimulate sets runtime limits and collection policies. Dynamo self-benchmark +uses the initialized engine, capture configuration and feasibility checks to +generate the exact grid. `--prefill-cudagraph-policy runtime` leaves graph +selection to the pinned runtime; it does not request a second trace-derived grid. +Validation trace lengths and SLAs do not determine collection coordinates. + The balanced profile uses iteration totals: prefill coordinates are `(batch_size, total_prefill_tokens, total_kv_read_tokens)`; decode coordinates are `(batch_size, total_kv_read_tokens)`. A rectangular min/max range does not @@ -184,8 +214,13 @@ assumptions. Use a separate data identity or campaign when changing those condit ## 3. Smoke test and collect -First execute a small prefill and decode measurement for each intended -configuration. Check that the engine initializes, the requested work is actually +Preview `onboard collect-fpm` without `--execute`, keeping the deployment +options identical to the intended launch. Inspect the generated read-only +collector plan too; the guided preview itself does not run that plan. Run +`onboard collect-fpm --check-readiness` against any existing native evidence. +Where evidence is missing, use `--execute --smoke` for every phase required by +each selected role; guided smoke checks all cells by default and publishes no +formal pair. Aggregated configurations require both prefill and decode. Check that the engine initializes, the requested work is actually scheduled, timings are finite and use the intended boundary, and all required ranks participate. Inspect whether attention KV and any recurrent state are warmed, synthesized, or skipped. A successful launch or prefill alone does not @@ -197,14 +232,17 @@ limit may not account for per-request linear/recurrent cache state. State initialization can change performance; document and validate synthetic-state measurements instead of treating them as equivalent to real request history. -After smoke passes, run the full planned surface. Save raw FPMs, actual scheduled +After matching readiness passes, use `onboard collect-fpm --execute` to run +the full planned surface, omitting `--smoke` and `--limit`. Save raw FPMs, actual scheduled shapes, effective engine configuration, versions/hashes, repeat counts, and failure/skip reasons. Keep incomplete or alternate-state results distinguishable from the primary profile. Validate that the intended phases and cells were published; exit code 0 by itself is not evidence that old cells were replaced. The [MiniMax collection example](#b3-run-smoke-then-formal-collection-step-3) -shows the Collector commands and publication checks. Use +shows the guided commands and publication checks. Track independent +configurations separately; sharing GPUs does not make one configuration's +publication a prerequisite for another. Use [recovery and cleanup](#recovery-and-cleanup) for retries and owned resources. **Result:** measurements and provenance that satisfy the frozen plan, with @@ -227,8 +265,11 @@ The runtime consumes a pair of files, not raw self-benchmark JSON: ``` Use the [whole-forward publication contract](../../collector/README.md#whole-forward-fpm-campaign) -to produce schema `aic_fpm_forward_perf`, version 6. Collector performs this -conversion for its campaigns. Importing raw output from another benchmark path +for schema `aic_fpm_forward_perf`. New guided collection publishes version 7, +including its execution identity. The published Kimi example retains its +original version-6 pair and requires a consumer that supports that identity. +Do not rewrite a schema label to migrate data. Collector performs conversion +for its own campaigns. Importing raw output from another benchmark path requires an explicit conversion preserving the recorded workload and identity. The `systems_paths` configuration points to ``, not its `data/` subdirectory. Whole-forward profiles use the layout above; operation tables may @@ -241,6 +282,16 @@ measurements, the sidecar policy and each row's policy/repeat count must agree. Keep the producer's aggregated latency; importing the pair does not train or recompute it. +Run `onboard validate-collection` for native guided campaigns: point validity, +execution evidence, matched-context full-grid repetitions, and withheld-coordinate +interpolation are separate checks. Without `--execute`, the command launches no +GPU repetitions; it can assess compatible existing evidence, while missing +repetitions leave repeatability incomplete. Review the editable policy before +collecting repetitions; bounded subsets remain diagnostics. Imported historical +tables without the required native evidence stay unqualified rather than +acquiring a passing status from successful loading. See the +[collection-quality procedure](../../../../docs/fpm-self-service.md#validate-collection-and-serving-accuracy). + Keep the hardware/query-policy files and source manifest with the pair. Preserve both files together when copying or renaming them to the canonical filenames. A dataset may contain comparator profiles or nonuniform workloads: choose the @@ -248,7 +299,7 @@ appropriate primary profile instead of merging every artifact into one table. [Example A1](#a1-import-the-measured-pair-step-4) downloads our published DCP8 profile at a pinned HF revision and checks its hashes. [Example B4](#b4-inspect-the-published-pair-step-4) -checks a newly collected pair. Publishing a pair does not automatically bundle +checks a newly collected pair and prepares its quality assessment. Publishing a pair does not automatically bundle it into an AISimulate release: use an external systems root or submit it through the repository's data-publication process. @@ -259,12 +310,26 @@ service or regression-training command is needed for measured FPM interpolation. ## 5. Construct and check the performance model +For config-derived profiles with pending memory, first run `onboard finalize` +against the original collection, writing to a fresh resolved directory. This +verifies runtime observations and the formal pair; it does not infer a required +activation/non-KV bound from guesswork. Review and accept the resolved profile +separately. Timing publication or a quality pass does not itself resolve memory. +An unsupported runtime needs an audited +[observation adapter](../../../../docs/fpm-self-service.md#resolve-cache-geometry-with-a-runtime-probe); +a preexisting complete resource profile can be used without this finalization. +Historical version-6 tables cannot be finalized by the current native campaign +finalizer; retain their resource provenance and use the compatible consumer route. + Use `RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig(...))`. Supply the exact model, system, backend/version, worker role, topology, precision, and backend identity of the profile, together with its `systems_paths` root. Use `estimation_mode="fpm_interpolation"` and `fallback_policy="deny"` to select measured FPM data explicitly. The general defaults are `auto` + `deny`; auto tries op-level, FPM interpolation, then regression at construction. +Guided config/profile onboarding instead supplies `fpm_profile` and explicitly +sets `estimator_config.fpm_interpolation.method="direct"`, avoiding analytical +model construction and SOL estimation. Keep this choice in generated configs. Check `diagnostics()` for readiness, the selected estimator, resolved identity, and systems root. Query a known measured prefill point and a known measured @@ -273,7 +338,9 @@ iteration totals; they are not request TTFT/ITL. With correction disabled, verif that measured-point queries reproduce their stored latency before testing interpolation. Save the resolved provenance with your results. -For profiles that require them: +The following additional options belong to the consumer build used by +[Example A](#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile); do not assume +a build with `onboard` already supports these independently added features: - `dcp` records decode context parallelism within the TP group; it does not multiply physical GPU count. Current DCP timing uses measured vLLM FPM; @@ -291,8 +358,8 @@ does not switch estimators on a query miss, and untrained regression cannot run offline prediction. [Example A2](#a2-query-the-canonical-model-step-5) checks both phases of the -collected Kimi profile. [Example B5](#b5-query-the-canonical-model-step-5) shows -the same API for MiniMax. See the [Core API](../../../../docs/core-api.md#choosing-a-forward-pass-api) +collected Kimi profile. [Example B5](#b5-finalize-memory-and-check-queries-step-5) checks +the resolved profile through the same API for MiniMax. See the [Core API](../../../../docs/core-api.md#choosing-a-forward-pass-api) for the complete contract. **Result:** a ready estimator with the intended data identity and checked queries. @@ -301,13 +368,23 @@ for the complete contract. ## 6. Connect the model to Replay +For an aggregated guided campaign, use the resolved directory's generated +`predict/pilot.yaml` and `recommend/pilot.yaml`; `onboard validate-fpm` generates +and runs a cold aggregated trace replay with direct-FPM coverage reporting. +It does not download traces or collect GPU data. An independent role exports +`worker.yaml`, not an executable full P/D prediction/recommendation; qualify the +counterpart independently and compose serving only where Replay supports it. +Current grouped-cache P/D handoff limits concern replay, not FPM collection. + Configure `aisimulate predict` with the same data root, estimator selection, identity, and engine build assumptions. Then supply the serving behavior that an FPM table does not define: worker layout, scheduler token/request limits, context limit, block geometry and capacity, prefix-cache policy, traffic, and any transfer/offload settings. -Relevant YAML fields include: +Relevant ordinary prediction YAML fields include the following. DCP and the +precision/backend override support require the consumer build in Example A; +they are not additional guided topology choices. | Purpose | Configuration | | --- | --- | @@ -328,7 +405,7 @@ For `decode_context`, use **`--stack engine`**. Dynamo Replay and Planner do not yet support this field; their DCP integration requires a separate downstream update. Prediction accepts the timing identity overrides above for regular language -workers with default timing. Recommendation currently rejects those overrides, +workers with default timing. In the Example A consumer build, recommendation rejects those overrides, and DCP is not a recommendation search dimension. [Example A3](#a3-run-and-check-replay-steps-67) and [Example B6](#b6-run-replay-step-6) provide complete prediction inputs. @@ -339,6 +416,14 @@ accuracy and unsupported regions are checked in step 7. ## 7. Validate coverage and accuracy +Use `onboard validate-collection` for native collection quality and +`onboard validate-fpm` for aggregated trace coverage. Follow +[`onboard validate-serving`](../../../../docs/fpm-self-service.md#stage-6-matched-serving) +to prepare a frozen matched workload, retain serving measurements and assess it +in a separate output directory. The collection/report, replay/report and +serving/report remain separate; ordinary exploratory prediction is available +when memory is resolved even if accuracy is unqualified. + First verify exact measured-point queries, then test the interpolation and Replay shapes required by your workload. Report unsupported queries as coverage gaps. Successfully loading a table and replaying its calibration points does not @@ -377,6 +462,15 @@ whether it fits their workload. ## Example A: onboard the collected Kimi K3 TP8+DCP8 profile +This example needs an AISimulate build containing the DCP profile-consumer +changes from [#284](https://github.com/ai-dynamo/aisimulate/pull/284) +(merged commit `982c18f4e8e33c124b5195945db269b83168ed3f`), including `dcp`, +`text_only`, `unrecorded_quant_modes`, and engine-stack `decode_context`. Verify +that build separately if your guided-onboarding checkout lacks these options. +`onboard init` and its collection presets do not enumerate DCP, so this published +profile uses the canonical SDK and ordinary prediction directly. Do not replace +DCP8 with TP8-only to make a guided request pass. + This worked example illustrates steps 4–7 using a profile already collected through Dynamo self-benchmarking. Complete step 1 first and review the retained collection configuration as described in step 2. It runs on CPU and downloads @@ -605,264 +699,272 @@ use a separate directory if running both examples. ## Example B: collect a MiniMax-M2.7 TP4 profile -This example illustrates steps 2–6 with a new vLLM collection on four H200 GPUs. -Complete step 1 first. It assumes a working GPU collection environment; replace -the namespace, checkpoint, image, and release placeholders with its values. -This example does not collect DCP. Use step 7 for accuracy validation afterward. +This example uses the guided config/profile route for an **aggregated** worker +on four H200 GPUs. It collects whole-forward timings without requiring a new +op-level class. Its topology and workload values illustrate one campaign, not +MiniMax defaults. Complete step 1 and the pinned-runtime investigation first; +replace every placeholder with inspected deployment facts. This example does +not collect DCP or compose independent P/D workers. ### B1. Prepare the campaign (step 2) -This example uses an external directory to keep its results separate -from bundled databases and previous campaigns. The directory name and layout -are choices for this walkthrough, not mandatory FPM setup steps. The later -commands use these paths consistently; replace the absolute path below with a -new persistent directory. +Keep the local `config.json` and any quantization sidecar from the immutable +checkpoint. Derive supported precision/cache geometry and prepare a reviewed +`resource-overrides.yaml` only for facts the config cannot establish; its +[fields and source requirements](../../../../docs/fpm-self-service.md#start-from-a-local-model-config) +are documented separately. Do not infer FMHA or KV precision solely from the +weight format, copy a historical BF16 cell's identity onto a native FP8 runtime, +or invent activation/non-KV bounds. Memory may remain pending until initialization. +For unsupported runtime geometry, first use the +[probe/import workflow](../../../../docs/fpm-self-service.md#resolve-cache-geometry-with-a-runtime-probe) +with an audited adapter for the exact framework build; keep the resulting +reviewed request instead of regenerating it from unobserved config estimates. + +Use a fresh absolute directory, on shared storage for a multi-node Slurm campaign: ```bash export FPM_RUN=/absolute/new/path/m27-h200-tp4 +export FPM_MODEL_CONFIG=/absolute/path/to/pinned-checkpoint/config.json +export FPM_MODEL_REVISION=REPLACE_WITH_IMMUTABLE_CHECKPOINT_REVISION +export FPM_VLLM_VERSION=REPLACE_WITH_FULL_PINNED_VLLM_VERSION +export FPM_RESOURCE_OVERRIDES=/absolute/path/to/reviewed-resource-overrides.yaml mkdir -p "$(dirname "$FPM_RUN")" mkdir "$FPM_RUN" -mkdir -p "$FPM_RUN/systems/data" -cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/h200_sxm.yaml" \ - "$FPM_RUN/systems/" -cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/query_versions.yaml" \ - "$FPM_RUN/systems/" -cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/attention_lane_defaults.yaml" \ - "$FPM_RUN/systems/" +aisimulate onboard checkpoint --file "$FPM_RUN/onboarding-checkpoint.json" + +init_args=( + --model MiniMaxAI/MiniMax-M2.7 + --model-config "$FPM_MODEL_CONFIG" + --model-revision "$FPM_MODEL_REVISION" + --framework-version "$FPM_VLLM_VERSION" + --resource-overrides "$FPM_RESOURCE_OVERRIDES" + --gpu h200_sxm --interconnect nvswitch + --worker-type aggregated + --tensor-parallel 4 --attention-data-parallel 1 + --moe-tensor-parallel 4 --moe-expert-parallel 1 + --context-length 8192 --max-num-tokens 8192 --max-batch-size 256 + --gpu-memory-utilization 0.90 --prefill-cudagraph-policy runtime + --input-tokens 1024 --output-tokens 32 --concurrency 4 --request-count 8 +) +aisimulate onboard init "${init_args[@]}" --output "$FPM_RUN/draft-request.yaml" ``` -`h200_sxm.yaml` provides hardware specifications for the later SDK/prediction -steps when they load this external systems root. `query_versions.yaml` preserves -the version-slot policy, and `attention_lane_defaults.yaml` preserves backend -lane defaults. These files describe hardware and query policy; they contain no -measured FPM rows. The collector's -`--fpm-database-root` override is optional; the later examples select it -explicitly so they publish and consume data from the same directory. +The last four flags describe an optional synthetic validation workload. They do +not define the collection grid. The context, scheduler, memory fraction and graph +policy are collection inputs and must reflect the intended serving configuration. +Keep the actual checkpoint mounted in the runtime consistent with its declared pin. +A config hash alone does not prove a checkpoint revision. -Save the following as `$FPM_RUN/collector-k8s.yaml`, replacing every -`REPLACE_...` value with the corresponding existing collection setting: +The command saves a draft without GPU work. If inputs remain unresolved, inspect +the diagnostics, investigate the missing facts and regenerate the draft with +corrected overrides. Save the complete draft in the session checkpoint: -```yaml -K8sConfig: - k8s_namespace: REPLACE_WITH_EXISTING_NAMESPACE - k8s_image: REPLACE_WITH_EXISTING_COLLECTION_IMAGE_DIGEST - k8s_pvc_name: REPLACE_WITH_MODEL_PVC - k8s_pvc_mount_path: /models - k8s_model_path_in_pvc: REPLACE_WITH_RELATIVE_CHECKPOINT_DIRECTORY - k8s_hf_home: /tmp/fpm-hf-cache +```bash +python - <<'PY' +import json +import os +from pathlib import Path +import yaml + +run = Path(os.environ["FPM_RUN"]) +request = yaml.safe_load((run / "draft-request.yaml").read_text()) +update = {"configurations": {"tp4": {"draft_request": request}}} +(run / "draft-update.json").write_text(json.dumps(update, indent=2) + "\n") +PY +FPM_CHECKPOINT_REVISION=$(aisimulate onboard checkpoint \ + --file "$FPM_RUN/onboarding-checkpoint.json" | \ + python -c 'import json, sys; print(json.load(sys.stdin)["state"]["revision"])') +aisimulate onboard checkpoint --file "$FPM_RUN/onboarding-checkpoint.json" \ + --expect-revision "$FPM_CHECKPOINT_REVISION" --update "$FPM_RUN/draft-update.json" +``` + +Present the exact profile, sources, precision, topology and collection settings +for review. Save edits and show the revised draft; only after explicit acceptance +of those values, record it and regenerate the same request at a fresh final path: + +```bash +FPM_CHECKPOINT_REVISION=$(aisimulate onboard checkpoint \ + --file "$FPM_RUN/onboarding-checkpoint.json" | \ + python -c 'import json, sys; print(json.load(sys.stdin)["state"]["revision"])') +aisimulate onboard checkpoint --file "$FPM_RUN/onboarding-checkpoint.json" \ + --expect-revision "$FPM_CHECKPOINT_REVISION" --accept-profile tp4 +aisimulate onboard init "${init_args[@]}" --output "$FPM_RUN/request.yaml" +cmp "$FPM_RUN/draft-request.yaml" "$FPM_RUN/request.yaml" ``` -The dedicated collector accepts deployment-only `K8sConfig` fields. In -particular, do not copy `Workers.agg.extra_cli_args` or `ServiceConfig` from a -standalone Generator example into this file. The accepted fields are defined in -[the collector input resolver](../../collector/fpm_forward/entry.py). +Keep `init_args` and input files synchronized with reviewed edits; a differing +request needs review again. Register the final request and later reports using +[checkpoint artifact references](../../../../docs/fpm-self-service.md#record-artifacts-and-resume). +Record findings, commands/exit statuses and blockers throughout, not only at the +end. The checkpoint command does not capture other commands automatically. ### B2. Freeze and inspect the plan (step 2) -Use the Generator target release from the existing collection configuration -and define the campaign arguments once: +Generate the plan from the accepted request and preview the intended deployment: ```bash -export AIS_DYNAMO_RELEASE=REPLACE_WITH_MATCHING_DYNAMO_RELEASE -cd "$AIS_REPO/python/aisimulate" -fpm_args=( - --backend vllm - --model-path MiniMaxAI/MiniMax-M2.7 - --gpu h200_sxm - --fpm-max-gpus 4 - --fpm-gpu-counts 4 - --fpm-parallel-presets pure_tp - --fpm-kv-cache-dtypes auto - --fpm-max-prefill-isl 8192 - --fpm-max-prefill-cudagraph-size 2048 - --dynamo-version "$AIS_DYNAMO_RELEASE" - --generator-config "$FPM_RUN/collector-k8s.yaml" - --checkpoint-dir "$FPM_RUN/checkpoints" - --fpm-artifact-root "$FPM_RUN/artifacts" - --fpm-database-root "$FPM_RUN/systems/data" +aisimulate onboard plan \ + --config "$FPM_RUN/request.yaml" --output-dir "$FPM_RUN/collection" + +deployment_args=( + --executor kubernetes + --namespace REPLACE_WITH_EXISTING_NAMESPACE + --image REPLACE_WITH_IMMUTABLE_COLLECTION_IMAGE + --model-cache REPLACE_WITH_MODEL_PVC:/models:REPLACE_WITH_RELATIVE_CHECKPOINT_DIRECTORY + --dynamo-version REPLACE_WITH_MATCHING_DYNAMO_RELEASE + --transport nvlink +) +collect_args=( + --config "$FPM_RUN/collection/request.yaml" + --output-dir "$FPM_RUN/collection" + "${deployment_args[@]}" ) -python3 -m collector.fpm_forward "${fpm_args[@]}" \ - --plan-only > "$FPM_RUN/plan.json" +aisimulate onboard collect-fpm "${collect_args[@]}" ``` -If the checkpoint exists only inside the target Pod and its metadata cannot be -resolved on the machine running Collector, add `--fpm-model-config /absolute/config.json` -to `fpm_args`. Keep the canonical model ID in `--model-path`. - -**Expected:** exit code 0 and JSON describing the frozen plan, without creating -Kubernetes workloads. Inspect the model/config hash, system, topology, -quantization, phase cells, sampling policy, and admission decisions before -continuing. For this single-topology, single-KV-setting example, expect one -prefill cell and one decode cell if both are admitted. The plan is authoritative; -do not proceed with an empty plan or unexplained omissions. -`counts.cells` should be 2 for that case, while `counts.points` is -`"runtime-determined"`: runtime initialization determines the actual point set. - -The sampling settings have distinct meanings: - -- `--fpm-max-prefill-isl 8192` bounds the scheduled prefill **new-token** axis. - It does not set runtime context length. -- `--fpm-max-prefill-cudagraph-size 2048` is the prefill capture-size policy. - Align it with the deployment being modeled before freezing a formal campaign. -- The collector currently sets runtime `max_model_len=-1` for vLLM auto-fit and - has no CLI override for that setting. Record the resolved value with the - collection results. -- The default is five global warm-up iterations and one measurement per point. - Published `latency_ms` uses the maximum rank wall time, converted to ms; - it is not a mean or median over repeated measurements. - -The complete option list is available from `python3 -m collector.fpm_forward ---help` and [its argument definitions](../../collector/fpm_forward/config.py). - -#### Include multiple GPU counts and parallel configurations - -Both GPU counts and parallel presets accept multiple values. To expand this -MiniMax campaign, replace these three entries in `fpm_args` before generating -the plan: +Set the final `--model-cache` component to the pinned checkpoint directory +relative to the PVC root; use `.` when its weights and config are at that root. +Keep the canonical Hugging Face model ID and declared revision in the request +and profile; this path locates the same checkpoint on the PVC. + +For Slurm, replace the deployment array with the actual `--executor slurm`, +`--image`, repeated `--container-mount`, `--cpus-per-task`, `--cpu-bind`, template +version and transport options. Use a caller-owned allocation; Kubernetes-only +namespace/PVC options do not apply. Inspect the +[executor requirements](../../../../docs/fpm-self-service.md#choose-the-collection-executor) +before launch. Both executors are supported by the same guided commands. + +**Expected:** `collection/` contains `request.yaml`, `fpm-model-profile.json`, +`support-plan.json`, `commands.json`, local `systems/`, and generated +`predict/pilot.yaml` and `recommend/pilot.yaml`. Planning copies the hardware and +query-policy files; do not create or relabel timing data. The preview prints the +collector invocation without running it or checking target readiness. + +Inspect the generated command and the read-only collector plan in +`commands.json` (`fpm_plan_local`). When executing that lower-level plan command, +apply the same deployment options in its collector spelling as shown by the +guided preview; it must retain `--plan-only`. Preserve all generated model, +profile, runtime and output arguments. The collector plan is authoritative for +admitted cells: for this one-topology aggregated request, expect prefill and +decode cells, or investigate omissions. The exact grid and point count remain +runtime-determined. AISimulate bounds context and scheduled new tokens +separately; Dynamo generates feasible points from those limits and the initialized +engine's capture configuration. Runtime graph policy does not force a prefill +capture limit of 2,048 or assume prefill and decode execute the same graph. + +#### Include multiple parallel configurations + +Use `onboard init --model-config ... --parallel-configs FILE --output-dir ROOT` +for separately reviewed TP/DEP/TEP requests. Do not combine `--parallel-configs` +with the single-topology flags above. For this MoE model, a list can select: -```text ---fpm-max-gpus 8 ---fpm-gpu-counts 2 4 8 ---fpm-parallel-presets pure_tp tep dep +```yaml +- worker_type: aggregated + tensor_parallel: 4 + attention_data_parallel: 1 + moe_tensor_parallel: 4 + moe_expert_parallel: 1 +- worker_type: aggregated + tensor_parallel: 1 + attention_data_parallel: 4 + moe_tensor_parallel: 1 + moe_expert_parallel: 4 ``` -For each selected GPU count `N`, the MoE presets request these configurations -(`PP=CP=1`): - -| Preset | Attention TP | Attention DP | MoE TP | MoE EP | -| --- | --- | --- | --- | --- | -| `pure_tp` | N | 1 | N | 1 | -| `tep` | N | 1 | 1 | N | -| `dep` | 1 | N | 1 | N | - -The planner combines the selected counts and presets, deduplicates identical -topologies, and applies model/backend and memory admission. It then creates -prefill and decode cells for admitted topology, KV dtype, and backend-policy -combinations. Read `plan.json` for the actual configurations and cell count; -do not assume every requested combination is admitted. Dense models use the -`tp` preset; MoE `pure_tp` requires support in the model's capability profile. -See [topology enumeration](../../collector/fpm_forward/topology.py). - -`--fpm-max-gpus 8` is the maximum GPU count for one cell. The runner executes -cells sequentially; multiple Pods belonging to one cell execute together. -Selecting multiple counts/presets does not launch all configurations -simultaneously. - -Regenerate and inspect `plan.json` after changing these options. If the original -campaign has already run, first prepare a fresh campaign directory as in B1 -and redefine `fpm_args` with the new paths. The B5–B6 TP4 SDK/predict examples -still select only their matching cell from the larger database; they do not -automatically compare all collected configurations. +These are TP4 and DEP4, each requiring four GPUs per worker. Review each tuple's +own resource profile and collection limits; do not reuse TP rank-local bounds +for DEP. Read the emitted `onboarding.json` and plan each request separately. +Available total GPU allocation is a later execution/deployment choice, not a +sum of selected worker widths. Track each campaign independently and sequence +shared-resource reuse without coupling unrelated publication success. See +[multiple configurations](../../../../docs/fpm-self-service.md#onboard-multiple-parallel-configurations). ### B3. Run smoke, then formal collection (step 3) -First read the number of cells from the inspected plan and execute the minimal -profile for every cell. This works for both the two-cell TP4 example and an -expanded campaign: +First inspect any existing native evidence without launching workers: ```bash -FPM_CELL_COUNT=$(python3 -c \ - 'import json, sys; n = len(json.load(open(sys.argv[1]))["cells"]); assert n > 0; print(n)' \ - "$FPM_RUN/plan.json") -python3 -m collector.fpm_forward "${fpm_args[@]}" --smoke --limit "$FPM_CELL_COUNT" +aisimulate onboard collect-fpm "${collect_args[@]}" --check-readiness ``` -**Expected:** exit code 0, all selected smoke cells pass, raw runtime evidence is saved, -and no formal Parquet pair is published. `--smoke` without `--limit` runs only -the first cell. `--limit 2` is sufficient for the original two-cell example, -but in an expanded plan the first two cells can both be prefill: the planner -lists all prefill cells before decode cells. Use the full plan count to check -both phases of every admitted configuration. -In `checkpoints/fpm_forward_smoke.json`, check `smoke.status == "passed"`, -`smoke.cell_count == FPM_CELL_COUNT`, and `smoke.formal_database_written == false`. - -Only after the smoke command succeeds and its evidence is checked, run: +A fresh campaign returns incomplete/nonzero because it has no measurements. +When GPU execution is authorized and the environment is prepared, collect the +bounded smoke and inspect readiness: ```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" +aisimulate onboard collect-fpm "${collect_args[@]}" --execute --smoke +aisimulate onboard collect-fpm "${collect_args[@]}" --check-readiness ``` -Formal collection must omit both `--smoke` and `--limit`. Smoke and formal -checkpoints/artifacts are isolated automatically, so the roots above can be -shared; smoke samples are not reused as formal measurements. - -**Expected:** all frozen cells pass and the formal pair is published. Check the -formal checkpoint's database status, that `missing_cells` is empty, and that -`published_cells` covers the planned cells. In this fresh database, -`skipped_first_publisher_wins` should be empty. Exit code 0 alone is insufficient -to prove a new attempt replaced existing data: publication preserves a cell -already owned by another attempt. `--fpm-publish-partial` is an explicit opt-in -for incomplete campaigns and still returns a nonzero result. - -For this fresh, two-cell campaign, the formal checkpoint should contain: - -```json -{ - "database": { - "status": "passed", - "published_cells": 2, - "plan_cells": 2, - "missing_cells": [], - "skipped_first_publisher_wins": [] - } -} +Guided smoke covers all cells by default: both phases for this aggregated +request. `--limit` can omit a required phase and leave readiness incomplete. +Smoke saves raw diagnostics and `fpm-readiness.json`, but no formal pair. Check +the effective precision, graph settings, runtime pin, KV/state initialization, +actual scheduled shapes and per-rank evidence. In Slurm also inspect the observed +CPU pool and worker/scheduler affinity. A completed memory probe or successful +prefill alone does not qualify decode timing. Preserve failed evidence and fix +deterministic capability gaps before retrying unchanged work. + +Only once matching readiness is established, run formal collection: + +```bash +aisimulate onboard collect-fpm "${collect_args[@]}" --execute ``` -This is an excerpt of the expected fields, not a complete checkpoint file. -For an expanded campaign, use its actual plan count in place of 2; a complete -fresh publication has `published_cells == plan_cells == FPM_CELL_COUNT`. +Omit both `--smoke` and `--limit`. Formal and smoke checkpoints/artifacts are +separate. Inspect `fpm-checkpoint/fpm_forward.json`: publication must pass with +no missing cells and the expected published count. A successful process exit +alone does not prove it replaced an existing cell: check +`skipped_first_publisher_wins` too. A fresh two-cell campaign should publish both +cells with no skips. Keep worker completion, publication, memory and quality +statuses separate in the session checkpoint. -The resulting directory has this structure; `` is the first -16 characters of the plan hash: +The plan's main paths are: ```text $FPM_RUN/ - plan.json - checkpoints/ - fpm_forward_smoke.json - fpm_forward.json - artifacts// - collection-plan.json - run-manifest.json - cells//... - smoke/... - systems/ - h200_sxm.yaml - query_versions.yaml - attention_lane_defaults.yaml - data/h200_sxm/vllm// + onboarding-checkpoint.json + request.yaml + collection/ + request.yaml + support-plan.json + commands.json + fpm-model-profile.json + fpm-readiness.json + fpm-checkpoint/fpm_forward.json + fpm-artifacts//... + systems/data/h200_sxm/vllm// fpm_forward_perf.parquet fpm_forward_perf.metadata.json ``` -Retain generated manifests/scripts, raw benchmark files, resolved configs, -provenance, logs, checkpoints, and the published pair. Inspect -`kv_seed_regime`; do not assume every row contains real-KV evidence. See the -[collector contract](../../collector/README.md#whole-forward-fpm-campaign) for -accepted regimes and publication behavior. +Retain the frozen plan, manifests/scripts, raw measurements, runtime observations, +logs and pair. Global warm-up, per-point timing and KV/state preparation serve +different purposes. Keep KV warm-up enabled for both phases and inspect native +eligibility/regimes; a zero-KV prefill point needs no cached prefix. Formal +`fake_fallback` rows remain ineligible for direct FPM. Do not infer repeatability +from adjacent decode steps or consolidated provenance duplicates. The quality +stage below checks independent measurements under matched execution conditions. ### B4. Inspect the published pair (step 4) -The version directory is the actual Pod's `importlib.metadata.version("vllm")`, -also recorded in the published rows/metadata. It is not the AIS version or the -`--dynamo-version` used to select templates. Read it from the artifacts before -setting this variable: +Verify the actual container's full vLLM version matches `FPM_VLLM_VERSION` before +using that directory; the Generator's `--dynamo-version` is not the runtime +version. Keep the Parquet and metadata together: ```bash -export FPM_VLLM_VERSION=REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION -export FPM_PARQUET="$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION/fpm_forward_perf.parquet" -python3 - "$FPM_PARQUET" > "$FPM_RUN/pair-inspection.json" <<'PY' +export FPM_PARQUET="$FPM_RUN/collection/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION/fpm_forward_perf.parquet" +python - "$FPM_PARQUET" > "$FPM_RUN/pair-inspection.json" <<'PY' import hashlib import json import sys from pathlib import Path - import pyarrow.parquet as pq parquet = Path(sys.argv[1]) @@ -870,7 +972,7 @@ metadata = json.loads(parquet.with_suffix(".metadata.json").read_text()) with parquet.open("rb") as stream: digest = hashlib.file_digest(stream, "sha256").hexdigest() assert metadata["schema_name"] == "aic_fpm_forward_perf" -assert metadata["schema_version"] == 6 +assert metadata["schema_version"] == 7 assert metadata["parquet_sha256"] == digest table = pq.read_table(parquet) assert table.num_rows == metadata["row_count"] and table.num_rows > 0 @@ -879,197 +981,203 @@ identity = [ "tp", "pp", "dp", "moe_tp", "moe_ep", "cp", "gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "comm_quant_mode", "kv_cache_dtype", "moe_backend", "attention_backend", "enable_wideep", "enable_eplb", + "model_config_sha256", "execution_profile", "engram_residency", "input_modality", ] -identity += ["dcp"] if "dcp" in table.column_names else [] frame = table.to_pandas() -report = { +print(json.dumps({ "hash_match": True, "row_count": table.num_rows, "phase_cells": frame[identity].drop_duplicates().to_dict(orient="records"), - "kv_seed_regimes": frame["kv_seed_regime"].value_counts().to_dict() if "kv_seed_regime" in frame else {}, -} -print(json.dumps(report, indent=2)) + "kv_seed_regimes": frame["kv_seed_regime"].value_counts().to_dict(), +}, indent=2)) PY ``` -**Expected:** exit code 0, `hash_match: true`, a positive row count, and phase -identities matching the campaign. This structural check does not replace the -native loader's validation or prove interpolation coverage. Inspect the -`batch_size`, `total_prefill_tokens`, and `total_kv_read_tokens` coordinates -before selecting queries; axis minima/maxima alone do not prove that all -interior points can be answered. +Expect matching identities, positive rows and a verified hash. This structural +check does not replace the native loader or quality gates. Historical version-6 +reuse in Example A follows its own schema check, not this new-publication check. +`systems_paths` points to `collection/systems`, while the collector database root +is `collection/systems/data`; there is no extra model-family directory. -Keep these paths distinct: +Review and save the editable +[validation policy](../../../../docs/fpm-self-service.md#validate-collection-and-serving-accuracy) +at `$FPM_RUN/validation-policy.json`, then prepare point-validity, execution and +holdout checks without GPU work: -| Setting | Value in this guide | -| --- | --- | -| Collector `--fpm-database-root` | `$FPM_RUN/systems/data` | -| Canonical SDK/Replay `systems_paths` entry | `$FPM_RUN/systems` | -| Pair directory | `$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION` | +```bash +aisimulate onboard validate-collection \ + --config "$FPM_RUN/collection/request.yaml" --output-dir "$FPM_RUN/collection" \ + --validation-output-dir "$FPM_RUN/quality" --policy "$FPM_RUN/validation-policy.json" +``` + +The initial report normally remains incomplete until matched full-grid repeats +exist. Numerical failures return nonzero even during preparation. When repetition +execution is authorized, use the frozen assessment and source deployment: + +```bash +aisimulate onboard validate-collection \ + --config "$FPM_RUN/collection/request.yaml" --output-dir "$FPM_RUN/collection" \ + --validation-output-dir "$FPM_RUN/quality" --resume --execute +``` + +Preserve valid slow samples, individual attempts and the independent-launch +median. Bounded subsets have different sweep history and cannot qualify the +full-grid source. Do not relax CV/error thresholds or replace a source merely +because a newer attempt is faster. Missing producer evidence keeps comparability +unestablished; stable numbers alone do not prove comparable execution. -Whole-forward FPM has no extra model-family directory between the system and -backend. Always transfer the metadata sidecar together with the Parquet file. +### B5. Finalize memory and check queries (step 5) -### B5. Query the canonical model (step 5) +When the resource profile is pending, finalize from verified native observations: -This query explicitly selects the external root. The example asks for a -four-request prefill with 1,024 new tokens per request and no cached prefix; -confirm this coordinate is covered by the new collection. +```bash +aisimulate onboard finalize \ + --config "$FPM_RUN/collection/request.yaml" --output-dir "$FPM_RUN/collection" \ + --collection-report "$FPM_RUN/quality/collection-validation.json" \ + --resolved-output-dir "$FPM_RUN/resolved" +``` + +The report preserves its passed/failed/incomplete quality independently from +memory readiness. Missing or incompatible runtime memory cannot be replaced by +a passing timing table. Review and accept the resolved profile as a new checkpoint +draft, retaining the original collection references. Use that fresh directory +for simulation only after review. An already complete reviewed profile can use +its existing plan directory instead; do not run finalization just to copy data. +Set `FPM_SIM` to the applicable complete, accepted plan: ```bash -AIC_ALLOW_UNLISTED_VERSIONS=1 \ -python3 - "$FPM_RUN/systems" "$FPM_VLLM_VERSION" <<'PY' +export FPM_SIM="$FPM_RUN/resolved" +``` + +Use the resolved identity/resource profile for a CPU query of one measured point +per phase. This selects direct interpolation without an analytical class: + +```bash +python - <<'PY' import json import math -import sys +import os from pathlib import Path - +import pyarrow.parquet as pq from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel +root = Path(os.environ["FPM_SIM"]) +profile = json.loads((root / "fpm-model-profile.json").read_text()) +deployment, = profile["deployments"] config = ForwardPassPerfModelConfig( - model="MiniMaxAI/MiniMax-M2.7", system="h200_sxm", backend="vllm", - worker_type="aggregated", backend_version=sys.argv[2], - systems_paths=(str(Path(sys.argv[1]).resolve(strict=True)),), + model=profile["model"], system=deployment["system"], backend=deployment["backend"], + backend_version=deployment["backend_version"], worker_type="aggregated", + tp=deployment["tp"], pp=deployment["pp"], attention_dp=deployment["dp"], + moe_tp_size=deployment["moe_tp"], moe_ep_size=deployment["moe_ep"], + fpm_profile=profile, systems_paths=(str(root / "systems"),), estimation_mode="fpm_interpolation", fallback_policy="deny", - tp=4, pp=1, attention_dp=1, moe_tp_size=4, moe_ep_size=1, - estimator_config={"correction": {"enabled": False}}, + estimator_config={"fpm_interpolation": {"method": "direct"}, "correction": {"enabled": False}}, ) +pair = root / "systems/data" / deployment["system"] / deployment["backend"] / deployment["backend_version"] / "fpm_forward_perf.parquet" +rows = pq.read_table(pair).to_pylist() model = RustForwardPassPerfModel.best_available(config) try: - latency_ms = model.estimate_forward_pass_time_ms({ - "version": 1, - "scheduled_requests": { - "num_prefill_requests": 4, - "sum_prefill_tokens": 4096, - "sum_prefill_kv_tokens": 0, - }, - }) - assert latency_ms is not None and math.isfinite(latency_ms) and latency_ms > 0 - print(json.dumps({"prefill_ms": latency_ms, "provenance": model.diagnostics()["provenance"]})) + for phase in ("prefill", "decode"): + row = next(row for row in rows if row["workload_kind"] == phase and row["kv_seed_regime"] != "fake_fallback") + scheduled = ( + {"num_prefill_requests": row["batch_size"], "sum_prefill_tokens": row["total_prefill_tokens"], + "sum_prefill_kv_tokens": row["total_kv_read_tokens"]} + if phase == "prefill" else + {"num_decode_requests": row["batch_size"], "sum_decode_kv_tokens": row["total_kv_read_tokens"]} + ) + latency = model.estimate_forward_pass_time_ms({"version": 1, "scheduled_requests": scheduled}) + assert latency is not None and math.isclose(latency, row["latency_ms"], rel_tol=1e-9) + print(json.dumps({"phase": phase, "predicted_ms": latency, "measured_ms": row["latency_ms"]})) + (root / "estimator-provenance.json").write_text(json.dumps(model.diagnostics(), indent=2) + "\n") finally: model.close() PY ``` -**Expected:** exit code 0 and a finite, positive `prefill_ms`. No serving GPU is -used for this query. A successful lookup confirms this query can consume the -pair, not that all replay steps will be in-domain or that the prediction is -accurate. - -The model, hardware, backend version, parallel shape, quantization, and backend -knobs must match the collected identity. This example uses checkpoint defaults. -For explicitly collected nondefault precision, use the SDK's -`gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `kvcache_quant_mode`, and -`comm_quant_mode` arguments with the corresponding enum names as strings. See -[the canonical API](../../../../docs/core-api.md#choosing-a-forward-pass-api). - -This guide copies the version policy file and uses literal backend versions. -If the collected version is already queryable, omit -`AIC_ALLOW_UNLISTED_VERSIONS=1`. The override is a transitional measure, not a -support guarantee. Do not rename version directories or edit identities to -bypass a mismatch. External trees without `query_versions.yaml` use a different -policy that disables the version-slot gate; omitting the file is not needed for -this workflow. +If a custom runtime version is outside the copied query-version policy, use +`AIC_ALLOW_UNLISTED_VERSIONS=1` for the query and later simulation only after +confirming the actual identity. This permits version selection, not mismatched +precision/topology or unsupported queries. A measured-point match validates the +lookup path; it is not held-out accuracy evidence. ### B6. Run Replay (step 6) -Save this as `$FPM_RUN/predict.yaml`. Replace the backend version with the -collected version and `systems_paths` with the absolute `$FPM_RUN/systems` -directory. Confirm block size, context length, scheduler limits, and -prefix-cache policy against the resolved runtime configuration and the target -workload. These settings are explicit simulation inputs, not reconstructed -automatically from the FPM table. +For a first trace check, select one complete local AgentX play and preserve its +request dependencies. Run the existing cold aggregated replay with coverage: -```yaml -engine: - mode: aggregated - model: MiniMaxAI/MiniMax-M2.7 - hardware: h200_sxm - backend: vllm - backend_version: "REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION" - context_length: 8192 - systems_paths: ["REPLACE_WITH_ABSOLUTE_SYSTEMS_DIRECTORY"] - estimation_mode: fpm_interpolation - fallback_policy: deny - estimator_config: - correction: {enabled: false} - workers: - aggregated: - parallelism: - replicas: 1 - tensor: 4 - pipeline: 1 - attention_data: 1 - moe_tensor: 4 - moe_expert: 1 - scheduler: - max_batched_tokens: 8192 - max_sequences: 256 - kv_cache: - block_size: 64 - prefix_caching: false - capacity: {type: default} - timing: {type: default} -traffic: - source: {type: synthetic, input_tokens: 1024, output_tokens: 32} - load: {type: concurrency, concurrency: 4} - stop: {requests: 8} +```bash +aisimulate onboard validate-fpm \ + --config "$FPM_SIM/request.yaml" --output-dir "$FPM_SIM" \ + --trace /absolute/path/to/selected-play.jsonl \ + --validation-output-dir "$FPM_RUN/replay-validation" ``` -`engine.systems_paths` selects the data root for this prediction. No global -Python root overrides are needed. Run the CLI directly: +The result records completed replay separately from exact/interpolated/unsupported +queries. A missing query stops execution; the partial report does not audit the +remaining trace. The supplied file is replayed completely. See +[AgentX scope and coverage results](../../../../docs/fpm-self-service.md#validate-fpm-query-coverage-with-agentx-replay). +This does not qualify unsupported grouped-cache P/D handoff or full multimodal +execution. Independent prefill/decode plans must be composed separately where +Replay supports their memory/transfer semantics. + +The generated synthetic examples preserve the same accepted profile, data root, +worker width and strict direct-FPM selection: ```bash -AIC_ALLOW_UNLISTED_VERSIONS=1 \ -aisimulate predict --stack engine --config "$FPM_RUN/predict.yaml" \ +aisimulate predict --config "$FPM_SIM/predict/pilot.yaml" \ --output-dir "$FPM_RUN/prediction-fpm" --capture-per-request +aisimulate recommend --config "$FPM_SIM/recommend/pilot.yaml" \ + --output-dir "$FPM_RUN/recommendation-fpm" ``` -**Expected:** exit code 0 and a new `prediction-fpm/` directory containing: - -- `prediction.json`: the durable prediction report; check the completion count - against the eight requested completions and inspect the reported metrics. -- `requests.jsonl`: per-request records enabled by `--capture-per-request`. - -Keep the input YAML, resolved estimator provenance, package/source revision, -and pair hash with the report so its data source can be reproduced. Use a new output -directory for another prediction; the CLI rejects nonempty output directories -unless overwrite is explicitly requested. - -Every simulated step needs a matching identity and supported coordinates. A -missing cell fails rather than silently switching to `op_level`. With default -capacity estimation, FPM also caps KV capacity at the collected decode-KV -ceiling. Inspect that cap when explaining concurrency and throughput. +Use fresh output paths and check exit status, completed requests/tokens and +estimator provenance. For this example the prediction asks for eight completions +with 32 output tokens each; recommendation evaluates one exact worker preset, +not every possible topology. Actual deployment replicas/GPU budgets belong in +ordinary prediction/recommendation configs. Neither successful synthetic output +nor covered trace replay establishes accuracy: finish the matched serving +validation in step 7 and retain the collection, memory, coverage and serving +results separately in the checkpoint. ## Recovery and cleanup -For a Collector campaign such as Example B, use the same frozen `fpm_args` -and output roots to resume: - -```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" --resume -``` - -Before retrying failed cells, copy their existing raw artifacts and logs to a -separate archive. The runner removes old raw/log directories for cells it -re-executes. After preserving that evidence: +Resume the agent session with `aisimulate onboard resume --checkpoint PATH` +and inspect integrity issues before acting on saved progress. This checkpoint +restores inputs/conversation state; it does not automatically restart collection. +For Example B, retain the original request, plan and `deployment_args`: ```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" --resume --resume-retry-failed +aisimulate onboard collect-fpm "${collect_args[@]}" --resume --execute ``` -A completed, validated schema-v6 publication is terminal on resume even if raw -artifacts were pruned. Changing the model, image, sampling policy, or source can -change the frozen plan. Use fresh campaign, checkpoint, artifact, and database -roots for independent recollection or A/B comparisons. Do not run concurrent -campaigns with colliding workload names. +Use the same executor, image, mounts, transport and CPU policy. For a failed cell +that needs another GPU attempt, preserve the old raw/log evidence first; the +collector removes those directories when it re-executes the cell. The guided CLI +does not expose `--resume-retry-failed`. Follow the emitted lower-level +`python -m collector.fpm_forward` command with its **complete frozen arguments**, +adding `--resume --resume-retry-failed`; then return through guided +`collect-fpm --resume --execute` for runtime compatibility checks. Do not rerun +GPUs for a post-processing failure when intact artifacts can be recovered. + +A completed validated publication is terminal on compatible collector resume. +New native campaigns use schema-11 plans and schema-7 pairs. Historical plans and +schema-6 pairs have separate compatibility limits; see +[finalization and migration limits](../../../../docs/fpm-self-service.md#finalize-runtime-memory). +Changing the model, image, sampling/CPU policy or source can change the frozen +plan. Use fresh request/plan, checkpoint, artifact and database roots for a +changed campaign or independent A/B recollection. Preserve the prior outputs. +Schedule independent configurations independently; use success dependencies only +for actual prerequisites such as validation after its own publication. See +[campaign orchestration](../../../../docs/fpm-self-service.md#orchestrate-independent-collection-campaigns). The collector normally cleans up each cell's workload in its finalization path; cleanup failure is an error. After an abnormal interruption, preserve available logs and inspect the manifests recorded under this campaign's cell artifacts. -For each owned manifest that still has resources, run: +For Kubernetes, run the following only for each owned manifest that still has +resources. Slurm cleanup concerns the collector's named job steps, not the +caller's allocation; preserve its shared-storage results. ```bash kubectl delete -f /absolute/path/to/this/cell/k8s_deploy.yaml \ @@ -1089,10 +1197,11 @@ files. | Symptom | Check and next action | | --- | --- | -| Plan fails to resolve a model | Supply a real local config with `--fpm-model-config`; preserve the canonical model identity | +| Setup or plan cannot resolve model facts | Use `onboard init --model-config` and reviewed overrides for unresolved supported fields; preserve the checkpoint identity and pin | | Smoke succeeds but no Parquet appears | Expected for smoke; inspect both phase results before formal collection | | Partial publication or skipped existing cells | Inspect `missing_cells` and `skipped_first_publisher_wins`; use a new database root for independent recollection | | Pair hash/schema validation fails | Preserve both files, recover a matching published pair, and do not edit the sidecar to conceal a mismatch | -| FPM lookup fails | Check root layout, actual backend version, full identity, and supported coordinates; some queries also require supported SOL operations | -| SDK query works but replay fails | Inspect the later simulated shapes and decode-KV ceiling; one successful prefill query does not establish replay coverage | +| FPM lookup fails | Check root layout, runtime version, exact identity and query coordinates; direct FPM does not use SOL to fill missing coverage | +| SDK query works but replay fails | Inspect missing coordinates, resolved memory and supported Replay behavior; one lookup does not establish trace coverage | +| Quality or memory remains incomplete | Preserve the reports and resolve the missing native/repeat/runtime evidence; publication alone is not qualification | | Prediction output directory is not empty | Select a new directory to retain the previous report | diff --git a/python/aisimulate/pyproject.toml b/python/aisimulate/pyproject.toml index 61bde3f17..9fdc2e472 100644 --- a/python/aisimulate/pyproject.toml +++ b/python/aisimulate/pyproject.toml @@ -130,6 +130,8 @@ include = [ { path = "collector/fpm_forward/**/*.py", format = "wheel" }, { path = "collector/fpm_forward/runtime/fpm_exec.sh", format = "wheel" }, { path = "collector/fpm_forward/runtime/fpm_text.txt", format = "wheel" }, + { path = "collector/fpm_forward/runtime/vllm-*.json", format = "wheel" }, + { path = "collector/fpm_forward/runtime/instrumentation/*.md", format = "wheel" }, { path = "collector/fpm_forward/runtime/dsv41/*.json", format = "wheel" }, { path = "collector/fpm_forward/runtime/dsv41/README.md", format = "wheel" }, { path = "collector/fpm_forward/runtime/dsv41/LICENSE", format = "wheel" }, diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 08b05d8b5..20bd0a874 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -41,7 +41,14 @@ def finish_lowering(value): if canonical_result is None: return value result = dict(canonical_result) - for name in ("timing_model", "num_gpu_blocks", "tensor_parallel_size", "dp_size"): + for name in ( + "timing_model", + "num_gpu_blocks", + "tensor_parallel_size", + "dp_size", + "kv_cache_groups", + "kv_cache_capacity_bytes", + ): if name in value: result[name] = value[name] for name in ( @@ -89,6 +96,10 @@ def finish_lowering(value): raise ValueError(f"{name} conflicts with canonical timing configuration") memory_fields[name] = lowered[name] request = {key: value for key, value in authored.items() if key not in memory_fields} + if request.get("fpm_profile") is not None: + from aisimulate_core.sdk.fpm_profile import _require_forward_pass_profile_memory + + _require_forward_pass_profile_memory(RustForwardPassPerfModel.normalize_config(request)) model = RustForwardPassPerfModel.best_available(request) try: diagnostics = model.diagnostics() @@ -98,6 +109,8 @@ def finish_lowering(value): raise ValueError("regression estimator is not ready; replay requires training observations") canonical_result = dict(raw) resolved = diagnostics["provenance"]["config"] + if "worker_type" in lowered and lowered["worker_type"] != resolved["worker_type"]: + raise ValueError("worker_type conflicts with canonical timing configuration") lowered["timing_model"] = {**timing, "config": {**resolved, **memory_fields}} for names, expected in ( (("tensor_parallel_size", "aic_tp_size"), resolved["tp"]), @@ -120,6 +133,7 @@ def finish_lowering(value): "aic_moe_tp_size": resolved["moe_tp_size"], "aic_moe_ep_size": resolved["moe_ep_size"], "aic_fpm_profile": resolved.get("fpm_profile"), + "worker_type": resolved["worker_type"], }.items() if value is not None } @@ -155,6 +169,8 @@ def finish_lowering(value): if attention_dp is not None and dp > 1: lowered["dp_size"] = dp + if _materialize_profile_cache_groups(lowered, memory_diagnostics): + return finish_lowering(lowered) if lowered.get("num_gpu_blocks") is not None: return finish_lowering(lowered) if canonical_result is not None and (resolved.get("dcp") or 1) > 1: @@ -221,12 +237,86 @@ def finish_lowering(value): }, systems_path=capacity_systems_path, cuda_graph_reserved_bytes=lowered.get("cuda_graph_reserved_bytes", 0), - **({"fpm_profile": lowered["aic_fpm_profile"]} if lowered.get("aic_fpm_profile") is not None else {}), + **( + {"fpm_profile": lowered["aic_fpm_profile"], "worker_type": lowered.get("worker_type", "aggregated")} + if lowered.get("aic_fpm_profile") is not None + else {} + ), + **({"context_length": lowered.get("max_model_len")} if lowered.get("aic_fpm_profile") is not None else {}), **({"diagnostics": memory_diagnostics} if memory_diagnostics is not None else {}), ) return finish_lowering(lowered) +def _materialize_profile_cache_groups(raw: dict[str, Any], diagnostics: dict[str, Any] | None) -> bool: + """Transport the canonical grouped byte budget without a scalar capacity.""" + profile = raw.get("aic_fpm_profile") + if profile is None: + return False + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + from aisimulate_core.sdk.memory import estimate_kv_cache + + identity = { + "model": raw.get("aic_model_path"), + "worker_type": raw.get("worker_type", "aggregated"), + "system": raw.get("aic_system"), + "backend": raw.get("aic_backend"), + "backend_version": raw.get("aic_backend_version", raw.get("backend_version")), + "tp_size": raw.get("aic_tp_size", 1), + "pp_size": raw.get("aic_pp_size", 1), + "attention_dp_size": raw.get("aic_attention_dp_size", 1), + "moe_tp_size": raw.get("aic_moe_tp_size"), + "moe_ep_size": raw.get("aic_moe_ep_size"), + } + deployment = load_fpm_profile(profile).select(**identity) + deployment.resources.require_memory() + if deployment.resources.runtime_memory is not None and raw.get("num_gpu_blocks") is not None: + raise ValueError("runtime FPM memory cannot use fixed num_gpu_blocks; use the recorded cache allocation") + if deployment.resources.cache_layout != "grouped": + return False + if raw.get("num_gpu_blocks") is not None: + raise ValueError("grouped FPM cache requires a byte budget; fixed num_gpu_blocks is unsupported") + estimate = estimate_kv_cache( + identity.pop("model"), + identity.pop("system"), + identity.pop("backend"), + **identity, + max_num_tokens=raw.get("max_num_batched_tokens", _DEFAULT_MAX_NUM_BATCHED_TOKENS), + max_batch_size=raw.get("max_num_seqs", _DEFAULT_MAX_NUM_SEQUENCES), + context_length=raw.get("max_model_len"), + memory_fraction_kind="of_total", + memory_fraction_value=raw.get("gpu_memory_utilization", DEFAULT_GPU_MEMORY_UTILIZATION), + cuda_graph_reserved_bytes=raw.get("cuda_graph_reserved_bytes", 0), + systems_path=raw.get("systems_path"), + fpm_profile=profile, + **{ + target: raw[source] + for target, source in ( + ("gemm_quant_mode", "aic_gemm_dtype"), + ("moe_quant_mode", "aic_moe_dtype"), + ("fmha_quant_mode", "aic_fmha_dtype"), + ("kvcache_quant_mode", "aic_kv_cache_dtype"), + ("comm_quant_mode", "aic_comm_dtype"), + ("moe_backend", "aic_moe_backend"), + ("attention_backend", "aic_attention_backend"), + ("enable_eplb", "aic_enable_eplb"), + ("wideep_num_slots", "aic_wideep_num_slots"), + ) + if raw.get(source) is not None + }, + ) + for name, value in ( + ("kv_cache_groups", estimate["cache_groups"]), + ("kv_cache_capacity_bytes", estimate["total_kv_size_bytes"]), + ): + if name in raw and raw[name] != value: + raise ValueError(f"{name} conflicts with the canonical FPM resource budget") + raw[name] = value + if diagnostics is not None: + diagnostics.update(estimate) + return True + + def estimate_num_gpu_blocks( *, backend_name: str, @@ -258,6 +348,8 @@ def estimate_num_gpu_blocks( cuda_graph_reserved_bytes: int = 0, diagnostics: dict[str, Any] | None = None, fpm_profile: dict[str, Any] | None = None, + worker_type: str = "aggregated", + context_length: int | None = None, ) -> int: """Estimate per-rank KV blocks using the replay-wide AIC contract. @@ -335,7 +427,8 @@ def estimate_num_gpu_blocks( }, systems_path=systems_path, cuda_graph_reserved_bytes=cuda_graph_reserved_bytes, - **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), + **({"fpm_profile": fpm_profile, "worker_type": worker_type} if fpm_profile is not None else {}), + **({"context_length": context_length} if context_length is not None else {}), **({"diagnostics": diagnostics} if diagnostics is not None else {}), ) ) @@ -350,6 +443,7 @@ def estimate_kv_bytes_per_token( moe_ep_size: int = 1, kvcache_quant_mode: str | None = None, fpm_profile: dict[str, Any] | None = None, + worker_type: str = "aggregated", system: str | None = None, backend: str = "vllm", backend_version: str | None = None, @@ -362,6 +456,7 @@ def estimate_kv_bytes_per_token( deployment = load_fpm_profile(fpm_profile).select( model=model_name, + worker_type=worker_type, system=system, backend=backend, backend_version=backend_version, @@ -372,6 +467,8 @@ def estimate_kv_bytes_per_token( moe_ep_size=moe_ep_size, ) deployment.validate_overrides(kvcache_quant_mode=kvcache_quant_mode) + if deployment.resources.cache_layout == "grouped": + raise ValueError("grouped FPM cache has no scalar bytes per token; use its cache groups and byte budget") return deployment.resources.kv_bytes_per_token from aisimulate_core.sdk.memory import NaiveKVCacheEstimator diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index fb590baf3..ea7b3de9d 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -265,6 +265,7 @@ def _deployment( engine, prefill, engine.kv_transfer.bytes_per_token, + role="prefill", ) parallel = { **_parallel_mapping(prefill, prefix="prefill_"), @@ -651,6 +652,7 @@ def _worker_engine_args( engine, worker, cache.bytes_per_token, + role=role, ) if transfer_bytes_per_token is not None: payload["kv_transfer_bytes_per_token"] = transfer_bytes_per_token @@ -670,6 +672,8 @@ def _resolve_kv_bytes_per_token( engine: EnginePredictionConfig, worker: WorkerPredictionConfig, configured: int | str, + *, + role: str, ) -> int: if configured != "auto": return configured @@ -685,6 +689,7 @@ def _resolve_kv_bytes_per_token( **( { "fpm_profile": engine.fpm_profile.model_dump(mode="json"), + "worker_type": role, "system": worker.hardware or engine.hardware, "backend": engine.backend, "backend_version": engine.backend_version, diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index f4d1b82b5..207ea7742 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -798,6 +798,7 @@ def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: system=worker.hardware or engine.hardware, backend=engine.backend, backend_version=backend_version, + worker_type=role, tp_size=parallel.tensor, pp_size=parallel.pipeline, attention_dp_size=parallel.attention_data, @@ -808,6 +809,20 @@ def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: max_num_tokens=worker.scheduler.max_batched_tokens, max_batch_size=worker.scheduler.max_sequences, ) + if deployment.resources.cache_layout == "grouped": + if role != "aggregated" or engine.nextn or engine.speculation is not None: + raise ValueError("grouped FPM cache supports only aggregated vLLM without speculative decoding") + cache = worker.kv_cache + if cache.prefix_caching: + raise ValueError("grouped FPM cache requires kv_cache.prefix_caching=false for cold replay") + if cache.host_offload is not None or cache.g3_offload is not None: + raise ValueError("grouped FPM cache supports only HBM; host and G3 offload are unsupported") + if cache.capacity.type == "fixed" or cache.bytes_per_token != "auto": + raise ValueError( + "grouped FPM cache uses profile groups and a byte budget, not fixed blocks or bytes_per_token" + ) + # Recommendation topology domains and the GPU budget are resolved by + # lowering to SearchSpace, which validates reachable grouped workers. def _validate_speculation(engine, *, modes: set[str], backends: set[str]) -> None: diff --git a/python/aisimulate/src/aisimulate/fpm_advisory.py b/python/aisimulate/src/aisimulate/fpm_advisory.py new file mode 100644 index 000000000..a0b517334 --- /dev/null +++ b/python/aisimulate/src/aisimulate/fpm_advisory.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FPM collection advisories shared without importing collector or estimator runtime.""" + + +def prefill_graph_advisory(policy: str, *, enforce_eager: bool) -> str | None: + if policy != "explicit" or enforce_eager: + return None + return ( + "Explicit CUDA-graph captures apply to prefill only; decode keeps runtime-selected captures. " + "A shared memory profile requires matching effective graph settings in both phases. " + "Compatibility is unresolved until observed: use prefill_cudagraph_policy=runtime, " + "choose explicit captures verified to match decode, or select enforce_eager=true for both " + "phases when that matches the serving target." + ) diff --git a/python/aisimulate/src/aisimulate/fpm_profile.py b/python/aisimulate/src/aisimulate/fpm_profile.py index c81b1ebaa..14861f0bc 100644 --- a/python/aisimulate/src/aisimulate/fpm_profile.py +++ b/python/aisimulate/src/aisimulate/fpm_profile.py @@ -3,7 +3,9 @@ """Compatibility exports for the lightweight core fpm_profile types.""" +from aisimulate_core.fpm_profile import FpmCacheGroup as FpmCacheGroup from aisimulate_core.fpm_profile import FpmDeploymentProfile as FpmDeploymentProfile from aisimulate_core.fpm_profile import FpmModelProfile as FpmModelProfile from aisimulate_core.fpm_profile import FpmResourceProfile as FpmResourceProfile +from aisimulate_core.fpm_profile import FpmRuntimeMemoryProfile as FpmRuntimeMemoryProfile from aisimulate_core.fpm_profile import load_fpm_profile as load_fpm_profile diff --git a/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py b/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py index 9f555739d..a550b6ef5 100644 --- a/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py +++ b/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py @@ -763,7 +763,10 @@ def _resource_documents( raise ValueError("FPM single-node workload must not require a ComputeDomain document") if compute_domain is None: raise ValueError("FPM multinode workload requires a ComputeDomain document") - return [compute_domain.to_dict(), workload] + document = compute_domain.to_dict() + metadata = document["metadata"] + metadata["labels"] = {**(metadata.get("labels") or {}), **workload["metadata"]["labels"]} + return [document, workload] def _resource_metadata(context: dict[str, Any]) -> dict[str, Any]: diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index 817679d15..56d60013a 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -33,6 +33,7 @@ format_prediction_stdout, format_recommendation_stdout, prepare_output_directory, + write_fpm_coverage, write_prediction_report, write_recommendation_csv, write_recommendation_result, @@ -131,9 +132,14 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: capture_performance_diagnostics=bool({"time", "source"}.intersection(args.detail)), ), ) - except (KeyboardInterrupt, ResourceLimitError): - raise - except Exception as exc: + except BaseException as exc: + coverage = getattr(exc, "fpm_query_coverage", None) + if isinstance(coverage, str): + coverage = json.loads(coverage) + if isinstance(coverage, dict): + write_fpm_coverage(root, {**coverage, "status": "incomplete", "error": str(exc)}) + if isinstance(exc, (KeyboardInterrupt, ResourceLimitError)) or not isinstance(exc, Exception): + raise raise _CliExecutionError(f"{type(exc).__name__}: {exc}") from exc finally: mark_shutdown() @@ -149,6 +155,8 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: # JSON stdout, like prediction.json, must identify the approximation. native["summary"]["metric_semantics"] = report.metadata["metric_semantics"] native["summary"]["total_gpus"] = report.metadata["total_gpus"] + if isinstance(native.get("fpm_query_coverage"), dict): + write_fpm_coverage(root, native["fpm_query_coverage"]) summary = prediction_summary(native) summary.update(normalize_power_summary(report.metrics)) if "summary" in native: diff --git a/python/aisimulate/src/aisimulate/output.py b/python/aisimulate/src/aisimulate/output.py index acacd8c6a..f2503a41e 100644 --- a/python/aisimulate/src/aisimulate/output.py +++ b/python/aisimulate/src/aisimulate/output.py @@ -34,6 +34,7 @@ def prepare_output_directory(path: str | Path, *, overwrite: bool) -> Path: if overwrite: for name in ( "prediction.json", + "fpm-coverage.json", "recommendation.json", "recommendation.csv", "requests.jsonl", @@ -60,6 +61,14 @@ def write_prediction_report(root: Path, report: dict[str, Any]) -> Path: return path +def write_fpm_coverage(root: Path, coverage: dict[str, Any]) -> Path: + """Persist native lookup evidence even when the replay fails.""" + + path = root / "fpm-coverage.json" + path.write_text(json.dumps(coverage, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return path + + def write_requests(root: Path, records: list[dict[str, Any]]) -> Path: path = root / "requests.jsonl" with path.open("w", encoding="utf-8") as output: diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index efd627aeb..ce28a217f 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -390,6 +390,7 @@ def capabilities(self) -> RunnerCapabilities: supports_mtp_expected_acceptance=True, supported_engine_model_controls=ENGINE_MODEL_CONTROL_FIELDS, supports_state_cache=True, + supports_grouped_kv_cache=True, supported_trace_formats=( "mooncake", "mooncake-delta", @@ -553,7 +554,10 @@ def run( from .resources import ResourceLimitError except ImportError: raise error from None - raise ResourceLimitError(str(error)) from error + failure = ResourceLimitError(str(error)) + if hasattr(error, "fpm_query_coverage"): + failure.fpm_query_coverage = error.fpm_query_coverage + raise failure from error if not isinstance(report_json, str): raise InvalidRunnerError("AISimulate engine replay runtime report must be a JSON string") try: @@ -1420,9 +1424,12 @@ def _materialize_engine_role( ) if role_memory is not None and "total_gpu_capacity_bytes" in role_memory: role_memory["status"] = "available" - role_memory["estimated_num_gpu_blocks"] = role_memory.pop("num_gpu_blocks") + if "num_gpu_blocks" in role_memory: + role_memory["estimated_num_gpu_blocks"] = role_memory.pop("num_gpu_blocks") role_memory.pop("unavailable_reason", None) - capacity_materialized = role_config.get("num_gpu_blocks") is not None + capacity_materialized = ( + role_config.get("num_gpu_blocks") is not None or role_config.get("kv_cache_capacity_bytes") is not None + ) for name in ("engine_type", "aic_backend"): configured = role_config.pop(name, None) if configured is not None and configured != deployment_backend: diff --git a/python/aisimulate/src/aisimulate/sdk/fpm_model_metadata.py b/python/aisimulate/src/aisimulate/sdk/fpm_model_metadata.py new file mode 100644 index 000000000..7700da558 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/fpm_model_metadata.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility alias for aisimulate_core.sdk.fpm_model_metadata.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.fpm_model_metadata") diff --git a/python/aisimulate/src/aisimulate/support/checkpoint.py b/python/aisimulate/src/aisimulate/support/checkpoint.py new file mode 100644 index 000000000..45599652a --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/checkpoint.py @@ -0,0 +1,512 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Durable agent-supplied onboarding context, not a workflow executor.""" + +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +import tempfile +from copy import deepcopy +from pathlib import Path +from typing import Annotated, Any, Literal + +from pydantic import Field, JsonValue + +from aisimulate.config.common import StrictModel, load_yaml + +from .schema import SupportRequest + +Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] + + +class Progress(StrictModel): + """Agent-reported progress; completion is not runtime verification.""" + + stage: int = Field(default=1, strict=True, ge=1, le=6) + status: Literal["in_progress", "blocked", "complete"] = "in_progress" + blockers: list[str] = Field(default_factory=list) + next_action: str | None = None + + +class Artifact(StrictModel): + path: str = Field(min_length=1) + kind: Literal["file", "request", "profile", "collector_checkpoint"] = "file" + scope: Literal["input", "collection", "validation"] = "input" + archived: bool = Field(default=False, strict=True) + sha256: Digest | None = None + context_sha256: Digest | None = None + + +class Acceptance(StrictModel): + sha256: Digest + revision: int = Field(strict=True, ge=1) + + +class Configuration(StrictModel): + inputs: dict[str, JsonValue] = Field(default_factory=dict) + validation_inputs: dict[str, JsonValue] = Field(default_factory=dict) + draft_request: dict[str, JsonValue] = Field(default_factory=dict) + progress: Progress = Field(default_factory=lambda: Progress(stage=3)) + artifacts: dict[str, Artifact] = Field(default_factory=dict) + acceptance: Acceptance | None = None + history: list[dict[str, JsonValue]] = Field(default_factory=list) + + +class Checkpoint(StrictModel): + schema_version: Literal["aisimulate-onboarding-checkpoint/v1"] = "aisimulate-onboarding-checkpoint/v1" + revision: int = Field(default=0, strict=True, ge=0) + inputs: dict[str, JsonValue] = Field(default_factory=dict) + validation_inputs: dict[str, JsonValue] = Field(default_factory=dict) + research: dict[str, JsonValue] = Field(default_factory=dict) + decisions: dict[str, JsonValue] = Field(default_factory=dict) + pending_questions: list[JsonValue] = Field(default_factory=list) + progress: Progress = Field(default_factory=Progress) + artifacts: dict[str, Artifact] = Field(default_factory=dict) + configurations: dict[str, Configuration] = Field(default_factory=dict) + + +def add_checkpoint_parsers(actions: Any) -> None: + checkpoint = actions.add_parser("checkpoint", help="Save or inspect one resumable onboarding session.") + checkpoint.add_argument("--file", required=True, help="Session JSON path; keep outside fresh init output roots.") + checkpoint.add_argument("--update", help="JSON merge patch file, or - for stdin; supplied values are data only.") + checkpoint.add_argument("--expect-revision", type=int, help="Required current revision for existing updates.") + checkpoint.add_argument( + "--accept-profile", + action="append", + default=[], + metavar="CONFIG_ID", + help="Record user approval of the exact complete draft; repeat to accept several configurations.", + ) + resume = actions.add_parser("resume", help="Verify and report saved onboarding context without executing work.") + resume.add_argument("--checkpoint", required=True) + + +def _json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _digest(value: Any) -> str: + return hashlib.sha256(_json(value).encode()).hexdigest() + + +def _read_json(text: str) -> dict[str, Any]: + def reject_constant(value: str) -> None: + raise ValueError(f"non-finite JSON number {value} is unsupported") + + def unique_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + values = {} + for key, value in pairs: + if key in values: + raise ValueError(f"duplicate JSON key {key!r}") + values[key] = value + return values + + value = json.loads(text, parse_constant=reject_constant, object_pairs_hook=unique_keys) + if not isinstance(value, dict): + raise ValueError("checkpoint and update must be JSON objects") + return value + + +def _load(path: Path) -> Checkpoint: + payload = _read_json(path.read_text(encoding="utf-8")) + if "schema_version" not in payload or "revision" not in payload: + raise ValueError("saved checkpoint requires schema_version and revision") + state = Checkpoint.model_validate(payload) + if state.revision < 1: + raise ValueError("saved checkpoint revision must be at least 1") + groups = [(state.artifacts, None)] + [(config.artifacts, config) for config in state.configurations.values()] + for refs, config in groups: + for ref in refs.values(): + if ref.kind == "collector_checkpoint": + if config is None or ref.scope != "collection" or ref.sha256 is not None: + raise ValueError( + "collector checkpoint references require configuration collection scope without SHA-256" + ) + elif ref.sha256 is None: + raise ValueError("saved immutable artifact is missing its SHA-256 snapshot") + if (ref.context_sha256 is None) != (ref.scope == "input"): + raise ValueError("saved artifact context SHA-256 must exist only for collection/validation scope") + for config in state.configurations.values(): + if config.acceptance is not None and config.acceptance.revision > state.revision: + raise ValueError("profile acceptance revision cannot exceed checkpoint revision") + for ref in state.artifacts.values(): + if ref.scope != "input": + raise ValueError( + "shared artifacts must use scope=input; put collection/validation files in a configuration" + ) + return state + + +def _merge(original: Any, patch: Any) -> Any: + if not isinstance(patch, dict): + return deepcopy(patch) + result = deepcopy(original) if isinstance(original, dict) else {} + for key, value in patch.items(): + if value is None: + result.pop(key, None) + else: + result[key] = _merge(result.get(key), value) + return result + + +def _check_patch(patch: dict[str, Any]) -> None: + if {"schema_version", "revision"} & patch.keys(): + raise ValueError("schema_version and revision are CLI-owned fields") + groups = [patch.get("artifacts", {})] + configurations = patch.get("configurations") or {} + if not isinstance(configurations, dict): + raise ValueError("configurations must be a JSON object keyed by configuration ID") + for config in configurations.values(): + if not isinstance(config, dict): + continue + if "acceptance" in config: + raise ValueError("acceptance is CLI-owned; use --accept-profile after explicit user approval") + groups.append(config.get("artifacts", {})) + for group in groups: + if isinstance(group, dict): + for ref in group.values(): + if isinstance(ref, dict) and {"sha256", "context_sha256"} & ref.keys(): + raise ValueError("artifact sha256 and context_sha256 are captured by the CLI") + + +def _canonical_request(config: Configuration, *, require_complete: bool = False) -> dict[str, Any]: + try: + request = SupportRequest.model_validate(config.draft_request) + if request.fpm_profile is None: + raise ValueError("acceptance requires a complete draft_request with fpm_profile") + request.scheduler_limits() + if require_complete: + from .runtime import verify_runtime_profile + + verify_runtime_profile(request) + return request.model_dump(mode="json", exclude_none=True) + except ValueError: + if require_complete: + raise + return config.draft_request + + +def _source_refs(artifacts: dict[str, Artifact]) -> dict[str, Any]: + return { + name: {"path": ref.path, "sha256": ref.sha256, "kind": ref.kind} + for name, ref in artifacts.items() + if ref.scope == "input" and not ref.archived + } + + +def _context(state: Checkpoint, config: Configuration, scope: str) -> str: + request = deepcopy(_canonical_request(config)) + if scope == "collection": + request.pop("workload", None) + if isinstance(request.get("search"), dict): + for key in ("objective", "seed"): + request["search"].pop(key, None) + value: dict[str, Any] = { + "inputs": state.inputs, + "configuration_inputs": config.inputs, + "source_artifacts": _source_refs(state.artifacts), + "configuration_source_artifacts": _source_refs(config.artifacts), + "request": request, + } + if scope == "validation": + value.update( + validation_inputs=state.validation_inputs, configuration_validation_inputs=config.validation_inputs + ) + return _digest(value) + + +def _resolve(base: Path, value: str) -> Path: + path = Path(value).expanduser() + if "\x00" in value: + raise ValueError("artifact paths must contain no NUL characters") + return (path if path.is_absolute() else base / path).resolve() + + +def _file_digest(path: Path) -> str: + with path.open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def _capture_refs(state: Checkpoint, old: Checkpoint, target: Path) -> None: + groups = [(state.artifacts, old.artifacts)] + [ + (config.artifacts, old.configurations[name].artifacts if name in old.configurations else {}) + for name, config in state.configurations.items() + ] + for refs, old_refs in groups: + for name, ref in refs.items(): + path = _resolve(target.parent, ref.path) + if path in {target, _lock_path(target)}: + raise ValueError("artifact path collides with the session checkpoint or its lock") + ref.path = os.path.relpath(path, target.parent) + previous = old_refs.get(name) + same_path = previous is not None and _resolve(target.parent, previous.path) == path + if same_path: + if (ref.kind, ref.scope) != (previous.kind, previous.scope): + raise ValueError("remove an artifact reference before changing its kind or scope") + ref.sha256 = previous.sha256 + ref.context_sha256 = previous.context_sha256 + continue + if not path.is_file(): + raise ValueError(f"new artifact must reference an existing file: {path}") + ref.sha256 = None if ref.kind == "collector_checkpoint" else _file_digest(path) + ref.context_sha256 = None + for ref in state.artifacts.values(): + if ref.scope != "input": + raise ValueError( + "shared artifacts must use scope=input; put collection/validation files in a configuration" + ) + if ref.kind == "collector_checkpoint": + raise ValueError("mutable collector checkpoints belong to a configuration with scope=collection") + for config in state.configurations.values(): + for ref in config.artifacts.values(): + if ref.kind == "collector_checkpoint" and ref.scope != "collection": + raise ValueError("mutable collector checkpoints require scope=collection") + if ref.scope != "input" and ref.context_sha256 is None: + ref.context_sha256 = _context(state, config, ref.scope) + + +def _artifact_issues(ref: Artifact, state: Checkpoint, config: Configuration | None, target: Path) -> list[str]: + if ref.archived: + return [] + path = _resolve(target.parent, ref.path) + issues = [] + if path in {target, _lock_path(target)}: + return ["reference collides with checkpoint or lock"] + if config is not None and ref.scope != "input" and ref.context_sha256 != _context(state, config, ref.scope): + issues.append("stale: inputs or draft changed; archive this superseded reference and register current output") + try: + if not path.is_file(): + issues.append(f"missing file: {path}") + elif ref.kind == "collector_checkpoint": + _read_json(path.read_text(encoding="utf-8")) + elif ref.sha256 is None or _file_digest(path) != ref.sha256: + issues.append("file content differs from the saved SHA-256 snapshot") + elif config is not None and ref.kind in {"request", "profile"}: + expected = _canonical_request(config, require_complete=True) + actual = load_yaml(path) + if ref.kind == "request": + actual = SupportRequest.model_validate(actual).model_dump(mode="json", exclude_none=True) + else: + from aisimulate.fpm_profile import FpmModelProfile + + actual = FpmModelProfile.model_validate(actual).model_dump(mode="json", exclude_none=True) + expected = expected["fpm_profile"] + if _digest(actual) != _digest(expected): + issues.append(f"saved {ref.kind} does not match the current reviewed draft") + except (OSError, ValueError) as exc: + issues.append(f"cannot verify artifact: {exc}") + return issues + + +def report(state: Checkpoint, target: Path) -> dict[str, Any]: + issues = [] + shared_invalid = False + for name, ref in state.artifacts.items(): + for detail in _artifact_issues(ref, state, None, target): + issues.append({"configuration": None, "artifact": name, "detail": detail}) + shared_invalid = True + configurations = {} + for name, config in state.configurations.items(): + local_issues = [] + input_invalid = shared_invalid + for artifact, ref in config.artifacts.items(): + for detail in _artifact_issues(ref, state, config, target): + local_issues.append({"configuration": name, "artifact": artifact, "detail": detail}) + input_invalid |= ref.scope == "input" + accepted = config.acceptance is not None + if accepted: + try: + _canonical_request(config, require_complete=True) + from .runtime import verify_checkpoint_runtime + + verify_checkpoint_runtime(state, name, target) + if config.acceptance.sha256 != _context(state, config, "acceptance"): + raise ValueError("accepted inputs or draft no longer match; review and accept again") + except (OSError, ValueError) as exc: + accepted = False + local_issues.append({"configuration": name, "artifact": None, "detail": str(exc)}) + accepted &= not input_invalid + needs_attention = input_invalid or bool(local_issues) + configurations[name] = { + "profile_accepted": accepted, + "effective_status": "needs_attention" if needs_attention else "accepted" if accepted else "draft", + "reported_progress": config.progress.model_dump(mode="json"), + "next_action": "Resolve the integrity issues and review affected inputs before continuing." + if needs_attention + else config.progress.next_action + or ( + "Plan collection from the accepted draft." if accepted else "Continue investigation and profile review." + ), + } + issues.extend(local_issues) + return { + "checkpoint": str(target), + "state": state.model_dump(mode="json"), + "configurations": configurations, + "integrity_issues": issues, + "archived_artifacts": [ + {"configuration": owner, "artifact": name, "path": ref.path} + for owner, refs in [(None, state.artifacts)] + + [(name, config.artifacts) for name, config in state.configurations.items()] + for name, ref in refs.items() + if ref.archived + ], + "verification": { + "scope": ( + "Canonical profile acceptance and current file integrity only; stage completion is agent-reported." + ), + "archived_artifacts": "Historical references only; excluded from current file and semantic verification.", + "collector_checkpoints": ( + "Existence and JSON readability only; the collector verifies frozen-plan identity on actual resume." + ), + "executed": False, + }, + } + + +def _lock_path(target: Path) -> Path: + return target.with_name(f".{target.name}.lock") + + +def _target(value: str) -> Path: + path = Path(value).expanduser().absolute() + if path.is_symlink(): + raise ValueError("checkpoint path must not be a symlink") + path = path.parent.resolve() / path.name + if path.exists() and not path.is_file(): + raise ValueError("checkpoint path must be a file") + return path + + +def _atomic_save(state: Checkpoint, target: Path) -> None: + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=target.parent, prefix=f".{target.name}.", delete=False + ) as stream: + temporary = Path(stream.name) + stream.write(json.dumps(state.model_dump(mode="json"), indent=2, sort_keys=True, allow_nan=False) + "\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, target) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def save_checkpoint( + target: Path, *, patch: dict[str, Any] | None, expected_revision: int | None, accept: list[str] +) -> tuple[Checkpoint, bool]: + target.parent.mkdir(parents=True, exist_ok=True) + with os.fdopen(os.open(_lock_path(target), os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600), "a+") as lock: + try: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + raise ValueError( + "another checkpoint operation is writing this session; read the latest revision and retry" + ) from exc + _target(str(target)) + exists = target.exists() + old = _load(target) if exists else Checkpoint() + if exists and patch is None and not accept: + return old, False + if (exists and expected_revision is None) or ( + expected_revision is not None and expected_revision != old.revision + ): + raise ValueError( + f"checkpoint revision is {old.revision}; retry after reading it with --expect-revision {old.revision}" + ) + patch = patch or {} + _check_patch(patch) + state = Checkpoint.model_validate(_merge(old.model_dump(mode="json"), patch)) + state.revision = old.revision + 1 + for name, config in state.configurations.items(): + previous = old.configurations.get(name) + if ( + previous is not None + and previous.draft_request != config.draft_request + and previous.draft_request.get("fpm_profile") is not None + ): + from .runtime import preserve_runtime_overrides + + config.draft_request = preserve_runtime_overrides(previous.draft_request, config.draft_request) + _capture_refs(state, old, target) + for name, config in state.configurations.items(): + previous = old.configurations.get(name) + if previous is None: + continue + changed = _context(state, config, "acceptance") != _context(old, previous, "acceptance") + collection_changed = _context(state, config, "collection") != _context(old, previous, "collection") + validation_changed = _context(state, config, "validation") != _context(old, previous, "validation") + draft_changed = previous.draft_request != config.draft_request + if draft_changed or (changed and previous.acceptance is not None): + config.history.append( + { + "event": "draft_replaced" if draft_changed else "acceptance_invalidated", + "revision": old.revision, + "draft_request": deepcopy(previous.draft_request), + "inputs": deepcopy(previous.inputs), + "shared_inputs": deepcopy(old.inputs), + "acceptance": previous.acceptance.model_dump(mode="json") if previous.acceptance else None, + "artifacts": {key: ref.model_dump(mode="json") for key, ref in previous.artifacts.items()}, + "shared_artifacts": {key: ref.model_dump(mode="json") for key, ref in old.artifacts.items()}, + } + ) + if changed: + config.acceptance = None + progress_patch = patch.get("configurations", {}).get(name, {}).get("progress") + explicit_blockers = isinstance(progress_patch, dict) and progress_patch.get("status") == "blocked" + if config.progress.stage >= 3 and not explicit_blockers: + config.progress = Progress(stage=3, next_action="Review and accept the changed draft.") + elif validation_changed and not collection_changed and config.progress.stage >= 6: + config.progress = Progress(stage=6, next_action="Repeat validation with the changed validation inputs.") + for name in accept: + if name not in state.configurations: + raise ValueError(f"unknown configuration {name!r}") + config = state.configurations[name] + _canonical_request(config, require_complete=True) + from .runtime import verify_checkpoint_runtime + + verify_checkpoint_runtime(state, name, target) + source_refs = [(ref, None) for ref in state.artifacts.values()] + [ + (ref, config) for ref in config.artifacts.values() if ref.scope == "input" + ] + if any(_artifact_issues(ref, state, owner, target) for ref, owner in source_refs): + raise ValueError(f"cannot accept {name!r}: input artifacts failed integrity checks") + config.acceptance = Acceptance(sha256=_context(state, config, "acceptance"), revision=state.revision) + _atomic_save(state, target) + return state, True + + +def run_checkpoint_command(args: argparse.Namespace) -> int: + if args.support_action == "resume": + target = _target(args.checkpoint) + output = report(_load(target), target) + output["saved"] = False + else: + import sys + + target = _target(args.file) + if ( + args.update + and args.update != "-" + and Path(args.update).expanduser().resolve() in {target, _lock_path(target)} + ): + raise ValueError("update file must be separate from the checkpoint and its lock") + patch = None + if args.update: + patch = _read_json( + sys.stdin.read() if args.update == "-" else Path(args.update).expanduser().read_text(encoding="utf-8") + ) + state, saved = save_checkpoint( + target, patch=patch, expected_revision=args.expect_revision, accept=args.accept_profile + ) + output = report(state, target) + output["saved"] = saved + print(json.dumps(output, sort_keys=True, allow_nan=False)) + return 2 if output["integrity_issues"] else 0 diff --git a/python/aisimulate/src/aisimulate/support/cli.py b/python/aisimulate/src/aisimulate/support/cli.py index 08dc3c2c6..ea3540ba6 100644 --- a/python/aisimulate/src/aisimulate/support/cli.py +++ b/python/aisimulate/src/aisimulate/support/cli.py @@ -8,9 +8,12 @@ import argparse import json import os +import re import shlex import sys import tempfile +from copy import deepcopy +from fractions import Fraction from pathlib import Path from typing import Any, get_args @@ -19,9 +22,31 @@ from aisimulate.config.common import load_yaml +from .config_profile import ( + FIELD_CHOICES, + INTEGER_FIELDS, + OVERRIDE_FIELDS, + ModelConfig, + ProfileDraft, + ProfileRequestError, + derive_profile, + load_model_config, + validate_overrides, +) from .fpm import run_fpm from .plan import create_plan, request_id -from .schema import FPMDeployment, SearchProfile, SloSpec, SupportIdentity, SupportRequest, WorkloadSpec +from .schema import ( + AGENTX_REFERENCE_CONTEXT, + CUDAGRAPH_MODES, + CollectionSpec, + FPMDeployment, + SearchProfile, + SloSpec, + SupportIdentity, + SupportRequest, + WorkloadSpec, +) +from .topology import TopologySuggestions, suggest_topologies def add_support_parser(subparsers: Any) -> None: @@ -34,16 +59,32 @@ def add_support_parser(subparsers: Any) -> None: ), ) actions = onboard.add_subparsers(dest="support_action", required=True) + from .checkpoint import add_checkpoint_parsers + + add_checkpoint_parsers(actions) init = actions.add_parser( "init", - help="Declare a model, target hardware, and FPM pilot workload.", + help="Declare a model, target hardware, worker topology, and FPM collection limits.", description=( - "Use --interactive for terminal prompts. Scripted setup requires --model, --model-revision, " - "--model-kind, --framework-version, --gpu, --gpu-count, and --interconnect. " - "One node and a small TP1 synthetic pilot are defaults; review the saved scope before collection." + "Use --interactive for terminal prompts. --model-config derives supported local model metadata; " + "supply unresolved profile fields with --resource-overrides. Otherwise scripted setup requires " + "--model, --model-revision, --model-kind, --framework-version, --gpu, and --interconnect. " + "Collection GPUs are derived from the selected topology. Model-config setup suggests model/hardware-aware " + "topologies; other routes default to TP1. AISimulate launches benchmark workers when collection is " + "executed; no separately launched HTTP server is required. An existing vLLM launch configuration is " + "optional evidence for reviewing settings, not a required input or automatically imported file. " + "AISimulate sets runtime limits and optional capture overrides; Dynamo self-benchmark combines them " + "with image sampling defaults and runtime " + "feasibility checks to generate the exact grid. A complete grid does not establish query coverage. " + "Synthetic workload and SLA options customize optional validation examples only. " + "Set deployment replicas and GPU budgets in ordinary predict/recommend configs." ), ) - init.add_argument("--interactive", action="store_true", help="Guide setup with terminal prompts.") + init.add_argument( + "--interactive", + action="store_true", + help="Guide setup with terminal prompts; review and accept config-based profiles before saving.", + ) init.add_argument("--profile", choices=("onboarding",), help="Optional alias; all setup uses onboarding.") init.add_argument("--model") init.add_argument("--model-revision", help="Pinned model revision; moving labels such as main are unsupported.") @@ -51,35 +92,116 @@ def add_support_parser(subparsers: Any) -> None: init.add_argument("--framework", choices=("vllm",), help="Collection runtime (default: vllm).") init.add_argument("--framework-version", help="Pinned vLLM version in the collection environment.") init.add_argument("--gpu", help="Target GPU system name, for example h200_sxm.") - init.add_argument("--gpu-count", type=int, help="Total GPUs available.") - init.add_argument("--node-count", type=int, help="Number of nodes (default: 1).") - init.add_argument("--gpus-per-node", type=int, help="GPUs available per node; defaults to gpu-count on one node.") init.add_argument("--interconnect", help="Interconnect, for example nvswitch, pcie, or none.") + init.add_argument( + "--worker-type", + choices=("aggregated", "prefill", "decode"), + help="Serving role: prefill/decode collect independently for P/D disaggregation; " + "aggregated probes both phases.", + ) init.add_argument("--sm", type=int, help="GPU SM version override.") init.add_argument("--tokenizer-revision") init.add_argument("--chat-template-revision") init.add_argument("--aisimulate-revision", help="Optional pinned AISimulate source revision.") - init.add_argument("--tensor-parallel", type=int, help="Attention tensor-parallel size (default: 1).") + init.add_argument( + "--tensor-parallel", type=int, help="Attention tensor-parallel size; bypasses model-config suggestions." + ) init.add_argument("--attention-data-parallel", type=int, help="Attention data-parallel size (default: 1).") init.add_argument("--moe-tensor-parallel", type=int, help="Expert tensor-parallel size (default: TP for MoE).") init.add_argument("--moe-expert-parallel", type=int, help="Expert-parallel size (default: 1).") - init.add_argument("--fpm-profile", help="JSON/YAML identity and resource profile for class-independent FPM.") - init.add_argument("--input-tokens", type=int, help="Input tokens per request (default: 1024).") - init.add_argument("--output-tokens", type=int, help="Output tokens per request (default: 128).") - init.add_argument("--concurrency", type=int, help="Concurrent requests (default: 1).") - init.add_argument("--request-count", type=int, help="Synthetic request count (default: 4).") - init.add_argument("--context-length", type=int, help="Pilot context limit in tokens (default: 16384).") - init.add_argument("--ttft-ms", type=float, help="Target time to first token in ms (default: 1000).") - init.add_argument("--tpot-ms", type=float, help="Target time per output token in ms (default: 100).") - init.add_argument("--max-candidates", type=int, help="One worker, or also the largest replica count (1 or 2).") + profile_source = init.add_mutually_exclusive_group() + profile_source.add_argument( + "--fpm-profile", help="JSON/YAML identity and resource profile for class-independent FPM." + ) + profile_source.add_argument("--model-config", metavar="PATH", help="Local Hugging Face-style config.json.") + init.add_argument( + "--resource-overrides", + metavar="PATH", + help="JSON/YAML profile field overrides, including cache_groups/cache_block_sizes; " + "requires --model-config. Resource bytes are per rank.", + ) + init.add_argument( + "--suggest-parallel", + action="store_true", + help="Preview model-config topology suggestions as JSON; no prompts or writes, including to --output.", + ) + init.add_argument( + "--parallel-configs", + metavar="PATH", + help="JSON/YAML list of parallel configurations with optional worker_type, collection and resource_overrides; " + "requires --model-config and --output-dir; conflicts with topology flags and --suggest-parallel.", + ) + init.add_argument("--input-tokens", type=int, help="Synthetic validation input tokens (default: 1024).") + init.add_argument("--output-tokens", type=int, help="Synthetic validation output tokens (default: 128).") + init.add_argument("--concurrency", type=int, help="Synthetic validation concurrency (default: 1).") + init.add_argument("--request-count", type=int, help="Synthetic validation request count (default: 4).") + init.add_argument( + "--context-length", + type=int, + help="Runtime per-request context limit; defaults to min(model/profile context, AgentX reference 256000).", + ) + init.add_argument( + "--max-num-tokens", type=int, help="Rank-local scheduled token budget; profile bound or initial policy 8192." + ) + init.add_argument( + "--max-batch-size", + type=int, + help="Rank-local scheduler sequence bound; profile bound or initial policy 256. " + "Does not request every prefill batch.", + ) + init.add_argument( + "--prefill-cudagraph-policy", + choices=("runtime", "explicit"), + help="Prefill CUDA graphs: runtime selection (new setup default), or explicit capture override.", + ) + init.add_argument( + "--max-prefill-cudagraph-size", + type=int, + help="Select explicit capture with this limit; explicit policy without a limit uses 2048. " + "Conflicts with runtime policy; match the serving target.", + ) + init.add_argument( + "--gpu-memory-utilization", + type=float, + help="Fraction of total GPU memory for collection and simulation, in (0, 1] (initial policy: 0.90).", + ) + init.add_argument( + "--cudagraph-mode", choices=CUDAGRAPH_MODES, help="vLLM graph policy for this explicit worker role." + ) + init.add_argument( + "--cudagraph-capture-sizes", + type=int, + nargs="+", + metavar="TOKENS", + help="Increasing, distinct vLLM capture sizes for this role; requires runtime prefill policy.", + ) + init.add_argument( + "--max-cudagraph-capture-size", + type=int, + help="vLLM capture limit for this role; aggregated applies it to both launches. " + "Requires runtime prefill policy.", + ) + init.add_argument("--ttft-ms", type=float, help="Synthetic validation TTFT target in ms (default: 1000).") + init.add_argument("--tpot-ms", type=float, help="Synthetic validation TPOT target in ms (default: 100).") init.add_argument( "--objective", choices=get_args(SearchProfile.model_fields["objective"].annotation), help="Recommendation objective (default: throughput).", ) init.add_argument("--seed", type=int, help="Recommendation search seed (default: 42).") - init.add_argument("--output", default="support-request.yaml") - init.add_argument("--overwrite", action="store_true") + output = init.add_mutually_exclusive_group() + output.add_argument( + "--output", default="support-request.yaml", help="Save one request (default: support-request.yaml)." + ) + output.add_argument( + "--output-dir", + metavar="ROOT", + help="Save a request and FPM profile for each selected configuration in a fresh directory; " + "requires --model-config. Interactive suggestions accept comma-separated choices.", + ) + init.add_argument( + "--overwrite", action="store_true", help="Allow replacement of --output; unsupported with --output-dir." + ) plan = actions.add_parser( "plan", help="Plan FPM collection and write predict/recommend configs for the target hardware." @@ -98,19 +220,65 @@ def add_support_parser(subparsers: Any) -> None: "--execute", action="store_true", help="Execute collection using a matching saved onboarding plan." ) collect.add_argument( - "--smoke", action="store_true", help="Diagnostic smoke collection; does not publish formal data." + "--check-readiness", + action="store_true", + help="Inspect saved native timing evidence as JSON without GPU work or writes.", + ) + collect.add_argument( + "--smoke", + action="store_true", + help="Bounded collection of every selected cell, including both phases; no formal publication.", ) collect.add_argument( "--limit", type=int, help="Limit diagnostic cases; requires --smoke, no formal data publication." ) collect.add_argument("--resume", action="store_true", help="Resume the existing collector checkpoint.") collect.add_argument("--checkpoint-dir", help="Checkpoint directory inside the plan's fpm-checkpoint directory.") - deployment = collect.add_argument_group("Collector deployment") + add_deployment_arguments(collect) + from .finalization import add_finalization_parser + from .runtime import add_runtime_parsers + from .validation import add_validation_parser + from .validation_workflow import add_quality_parsers + + add_validation_parser(actions) + add_quality_parsers(actions) + add_finalization_parser(actions) + add_runtime_parsers(actions) + + +def add_deployment_arguments(parser: Any, *, default_executor: str | None = "kubernetes") -> None: + deployment = parser.add_argument_group("Collector deployment") + deployment.add_argument( + "--executor", + choices=("kubernetes", "slurm"), + default=default_executor, + help="Collection executor (default: kubernetes); Slurm uses an existing sbatch/salloc allocation.", + ) deployment.add_argument("--dynamo-version", help="Target Dynamo release used to resolve collector templates.") - deployment.add_argument("--image", help="Collector container image; prefer an immutable digest.") + deployment.add_argument( + "--image", help="Collector container image; required for Slurm. Prefer an immutable digest." + ) + deployment.add_argument( + "--container-mount", + action="append", + metavar="SRC[:DST[:FLAGS]]", + help="Slurm/Pyxis container mount; repeat for several mounts. Kubernetes uses --model-cache.", + ) deployment.add_argument("--namespace", help="Kubernetes namespace for collector resources.") + deployment.add_argument( + "--cpus-per-task", + type=int, + help="Slurm CPUs for each node's shared worker/scheduler pool (new campaigns: 16).", + ) + deployment.add_argument( + "--cpu-bind", + choices=("cores", "none"), + help="Slurm node-task binding (new campaigns: cores); does not pin individual ranks.", + ) deployment.add_argument("--model-cache", metavar="NAME[:MOUNT[:SUBPATH]]", help="Model-cache PVC and mount.") - deployment.add_argument("--transport", choices=("nvlink", "ib", "efa")) + deployment.add_argument( + "--transport", choices=("nvlink", "ib", "efa"), help="GPU networking transport, independent of the executor." + ) deployment.add_argument("--image-pull-secret", help="Kubernetes secret for pulling the collector image.") @@ -119,17 +287,33 @@ def _values(args: argparse.Namespace, model: Any) -> dict[str, Any]: def _request_from_args(args: argparse.Namespace) -> SupportRequest: - identity = _values(args, SupportIdentity) - if args.gpus_per_node is None and identity.get("node_count", 1) == 1: - identity["gpus_per_node"] = args.gpu_count - return SupportRequest.model_validate( - { - "identity": identity, - "workload": {**_values(args, WorkloadSpec), "slo": _values(args, SloSpec)}, - "search": _values(args, SearchProfile), - **({"fpm_profile": load_yaml(args.fpm_profile)} if getattr(args, "fpm_profile", None) else {}), - } + collection = _values(args, CollectionSpec) + collection.setdefault( + "prefill_cudagraph_policy", "explicit" if "max_prefill_cudagraph_size" in collection else "runtime" ) + payload = { + "identity": _values(args, SupportIdentity), + "workload": {**_values(args, WorkloadSpec), "slo": _values(args, SloSpec)}, + "search": _values(args, SearchProfile), + "collection": collection, + **({"worker_type": args.worker_type} if getattr(args, "worker_type", None) is not None else {}), + **({"fpm_profile": load_yaml(args.fpm_profile)} if getattr(args, "fpm_profile", None) else {}), + } + request = SupportRequest.model_validate(payload) + if getattr(args, "context_length", None) is None and not getattr(args, "fpm_profile", None): + raise ValueError("--context-length is required when no model config or FPM profile provides a context limit") + deployment = request.profile_deployment() + runtime = deployment.resources.runtime_memory if deployment is not None else None + if runtime is not None: + payload["search"].setdefault("context_length", min(request.search.context_length, runtime.max_model_len)) + collection.setdefault("gpu_memory_utilization", runtime.gpu_memory_utilization) + if payload["search"]["context_length"] > runtime.max_model_len: + raise ValueError("context_length exceeds the observed runtime memory max_model_len") + if collection["gpu_memory_utilization"] != runtime.gpu_memory_utilization: + raise ValueError("gpu_memory_utilization must match the observed runtime memory") + else: + collection.setdefault("gpu_memory_utilization", 0.9) + return SupportRequest.model_validate(payload) def _request_target(path: str | Path, *, overwrite: bool) -> Path: @@ -147,6 +331,7 @@ def _request_target(path: str | Path, *, overwrite: bool) -> Path: def _write_request(request: SupportRequest, path: str | Path, *, overwrite: bool) -> Path: + request.scheduler_limits() target = _request_target(path, overwrite=overwrite) payload = yaml.safe_dump(request.model_dump(mode="json", exclude_none=True), sort_keys=False) target.parent.mkdir(parents=True, exist_ok=True) @@ -180,27 +365,35 @@ def _print(value: dict[str, Any], output_format: str = "table") -> None: "model_kind": ("Model kind (dense/moe)", str), "framework_version": ("Pinned vLLM version", str), "gpu": ("Target GPU system name (for example h200_sxm)", str), - "gpu_count": ("Total GPUs available", int), "interconnect": ("GPU interconnect (for example nvswitch, pcie, or none)", str), - "tensor_parallel": ("GPUs per pure-TP worker", int), + "tensor_parallel": ("Attention tensor-parallel size", int), +} +_CORRECTION_PROMPTS = { + **_PROMPTS, "input_tokens": ("Input tokens per request", int), "output_tokens": ("Output tokens per request", int), "concurrency": ("Concurrent requests", int), - "context_length": ("Pilot context limit in tokens", int), + "context_length": ("Runtime per-request context limit in tokens", int), "ttft_ms": ("Target time to first token (ms)", float), "tpot_ms": ("Target time per output token (ms)", float), -} -_CORRECTION_PROMPTS = { - **_PROMPTS, + "max_num_tokens": ("Rank-local scheduled token budget", int), + "max_batch_size": ("Rank-local scheduler sequence bound", int), + "max_prefill_cudagraph_size": ("Prefill CUDA graph capture limit", int), + "prefill_cudagraph_policy": ("Prefill CUDA graph policy (runtime/explicit)", str), + "gpu_memory_utilization": ("Fraction of total GPU memory (0 < value <= 1)", float), + "worker_type": ("Serving role (aggregated/prefill/decode)", str), + "cudagraph_mode": ("vLLM CUDA graph mode", str), + "cudagraph_capture_sizes": ( + "CUDA graph capture sizes (space-separated integers)", + lambda value: [int(x) for x in value.split()], + ), + "max_cudagraph_capture_size": ("CUDA graph capture limit for this role", int), "framework": ("Runtime (vllm)", str), "tokenizer_revision": ("Pinned tokenizer revision", str), "chat_template_revision": ("Pinned chat-template revision", str), "aisimulate_revision": ("Pinned AISimulate revision", str), - "node_count": ("Number of nodes", int), - "gpus_per_node": ("GPUs per node", int), "sm": ("GPU SM version", int), "request_count": ("Number of synthetic requests", int), - "max_candidates": ("Maximum candidates (1 or 2)", int), "objective": ("Recommendation objective", str), "seed": ("Recommendation search seed", int), "attention_data_parallel": ("Attention data-parallel size", int), @@ -212,13 +405,11 @@ def _print(value: dict[str, Any], output_format: str = "table") -> None: def _prompt(args: argparse.Namespace, name: str) -> None: label, convert = _CORRECTION_PROMPTS[name] default = getattr(args, name) - if default is None: - for model in (SupportIdentity, WorkloadSpec, SloSpec, SearchProfile): + if default is None and name != "context_length": + for model in (SupportIdentity, WorkloadSpec, SloSpec, SearchProfile, CollectionSpec): field = model.model_fields.get(name) if field is not None and not field.is_required(): default = field.default - if name == "gpu_count": - default = 1 while True: answer = input(f"{label}{f' [{default}]' if default is not None else ''}: ").strip() if not answer and default is None: @@ -230,60 +421,833 @@ def _prompt(args: argparse.Namespace, name: str) -> None: print(f"Enter a valid {'integer' if convert is int else 'number'}.") continue setattr(args, name, value) + if name == "prefill_cudagraph_policy" and value == "runtime": + args.max_prefill_cudagraph_size = None + elif name == "max_prefill_cudagraph_size": + args.prefill_cudagraph_policy = "explicit" return -def _guided_request(args: argparse.Namespace) -> SupportRequest: +def _require_terminal() -> None: if not sys.stdin.isatty(): raise ValueError( "--interactive requires a terminal; for automation supply options (aisimulate onboard init --help)" ) + + +def _guided_request(args: argparse.Namespace, *, skip_topology: bool = False) -> SupportRequest: + _require_terminal() print("Onboard a model for FPM simulation on a target hardware platform.") - print("Start with a pure-TP worker and a small synthetic workload on vLLM.") + selection = "one or more" if getattr(args, "output_dir", None) is not None else "one" + print(f"Select {selection} TP, DEP, or TEP worker configurations and review collection limits for vLLM.") + print("AISimulate launches benchmark workers during collection; you do not need to start an HTTP server.") + print("An existing vLLM launch configuration is optional evidence for reviewing these settings.") print("Enter accepts a displayed default. Ctrl-C cancels without saving. Supplied options skip their prompts.") for name in _PROMPTS: + if skip_topology and name == "tensor_parallel": + continue if getattr(args, name) is None: _prompt(args, name) + if args.context_length is None and not args.fpm_profile: + _prompt(args, "context_length") + return _validate_guided_request(args) + + +def _correct_option(args: argparse.Namespace, exc: ValueError, *, require_known_field: bool = False) -> None: + name, message = None, str(exc) + if isinstance(exc, ProfileRequestError): + name = exc.field + if isinstance(exc, ValidationError): + error = exc.errors(include_url=False)[0] + name = error["loc"][-1] if error["loc"] else None + message = error["msg"] + if name == "backend_version": + name = "framework_version" + if require_known_field and name not in _CORRECTION_PROMPTS: + raise exc + print(f"{name or 'Request'}: {message}") + while name not in _CORRECTION_PROMPTS: + name = input("Option to correct (for example --tensor-parallel): ").strip().removeprefix("--").replace("-", "_") + if name not in _CORRECTION_PROMPTS: + print("Choose one of: " + ", ".join(key.replace("_", "-") for key in _CORRECTION_PROMPTS)) + _prompt(args, name) + + +def _validate_guided_request(args: argparse.Namespace) -> SupportRequest: while True: try: return _request_from_args(args) except ValidationError as exc: - error = exc.errors(include_url=False)[0] - name = error["loc"][-1] if error["loc"] else None - print(f"{name or 'Request'}: {error['msg']}") - while name not in _CORRECTION_PROMPTS: - name = ( - input("Option to correct (for example --tensor-parallel): ") - .strip() - .removeprefix("--") - .replace("-", "_") + _correct_option(args, exc) + + +class _ResourceOverridesLoader(yaml.SafeLoader): + def construct_mapping(self, node, deep=False): + values = {} + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) + if not isinstance(key, str): + raise ValueError("resource override field names must be strings") + if key in values: + raise ValueError(f"duplicate resource override field: {key}") + values[key] = self.construct_object(value_node, deep=deep) + return values + + +def _load_resource_overrides(path: str | None) -> dict[str, Any]: + if path is None: + return {} + source = Path(path).expanduser() + try: + values = yaml.load(source.read_text(encoding="utf-8"), Loader=_ResourceOverridesLoader) + except yaml.YAMLError as exc: + raise ValueError(f"malformed resource overrides in {source}: {exc}") from exc + if not isinstance(values, dict): + raise ValueError(f"resource overrides in {source} must be a flat JSON/YAML mapping") + return values + + +def _memory_field(name: str) -> bool: + return name.endswith("_bytes") or name == "kv_bytes_per_token" + + +def _resource_answer(name: str, answer: str) -> Any: + if name in {"cache_groups", "cache_block_sizes"}: + try: + return yaml.load(answer, Loader=_ResourceOverridesLoader) + except yaml.YAMLError as exc: + raise ValueError(f"enter a valid JSON {'list' if name == 'cache_groups' else 'object'}") from exc + if name not in INTEGER_FIELDS: + return answer + if _memory_field(name): + units = {"B": 1, **{f"{prefix}B": 1000**power for power, prefix in enumerate("KMGT", 1)}} + units.update({f"{prefix}IB": 1024**power for power, prefix in enumerate("KMGT", 1)}) + match = re.fullmatch(r"([+-]?\d+(?:\.\d+)?)\s*([KMGT]i?B|B)", answer, re.IGNORECASE) + if match: + value = Fraction(match[1]) * units[match[2].upper()] + if value.denominator != 1: + raise ValueError("memory quantities must resolve to a whole number of bytes") + return value.numerator + try: + return int(answer) + except ValueError as exc: + raise ValueError("enter integer bytes or an explicit unit, for example 70 GiB or 512 MiB") from exc + try: + return int(answer) + except ValueError as exc: + raise ValueError("enter a valid integer") from exc + + +def _read_profile_value(name: str) -> dict[str, Any]: + choices = f" ({', '.join(FIELD_CHOICES[name])})" if name in FIELD_CHOICES else "" + if _memory_field(name): + choices += " (per rank; integer bytes or units such as GiB/MiB)" + elif name == "cache_groups": + choices += " (JSON list; page_size_bytes includes every group layer and runtime padding per rank)" + elif name == "cache_block_sizes": + choices += " (JSON object mapping displayed group names to runtime tokens per block)" + while True: + answer = input(f"{name}{choices}: ").strip() + if not answer: + print("A value is required.") + continue + try: + return validate_overrides({name: _resource_answer(name, answer)}) + except ValueError as exc: + print(f"{name}: {exc}") + + +def _prompt_profile_value( + config: ModelConfig, + request: SupportRequest | None, + overrides: dict[str, Any], + draft: ProfileDraft, + name: str, +) -> tuple[dict[str, Any], ProfileDraft]: + print(f"{name}: {draft.missing[name]}") + # Retain a valid answer if final profile validation asks for an identity + # correction. Other derivation failures cannot be repaired by repeating + # this answer, so let the caller report the actual error. + answer = _read_profile_value(name) + if name == "cache_block_sizes": + answer[name] = {**overrides.get(name, {}), **answer[name]} + overrides.update(answer) + updated = derive_profile(config, request, overrides) + if name in updated.resolved: + print(f" {name}: {updated.resolved[name]} (source: {updated.sources[name]})") + else: + print(f" {name}: {updated.missing[name]}") + return overrides, updated + + +def _review_config_profile( + config: ModelConfig, + request: SupportRequest, + overrides: dict[str, Any], + draft: ProfileDraft, + *, + directory_output: bool = False, +) -> SupportRequest: + while True: + identity = request.identity + print("Review FPM profile before saving:") + print(f" Modeling scope: {config.notes['modeling_scope']}") + print( + f" Target: {identity.model} ({identity.model_kind}) @ {identity.model_revision}; " + f"{identity.framework} {identity.framework_version}; " + f"{identity.gpu}; {identity.interconnect}" + ) + print(f" Parallelism: {request.parallelism()}") + if request.worker_type is not None: + print(f" Serving role: {request.worker_type}; collection phases: {', '.join(request.collection_phases)}.") + if request.worker_type != "aggregated": + print(" P/D disaggregation: this role has its own accepted settings, memory and collection evidence.") + print(" The other role is onboarded independently; collection does not require actual P/D transfer.") + print(f" Collection GPUs required: {request.worker_gpus} for one selected worker; availability is unchecked.") + settings = request.collection_settings() + print(f" Runtime context limit: {request.search.context_length} tokens per request.") + print(f" Collection: {json.dumps(settings, sort_keys=True)}") + print(" AISimulate launches and manages benchmark workers when collection is executed.") + print(" vLLM loads the model and allocates caches; Dynamo self-benchmark generates and measures the grid.") + print(" AISimulate supplies runtime limits and collection policies; no HTTP server is needed beforehand.") + print( + " Dynamo combines them with image sampling defaults and runtime feasibility checks to generate the grid." + ) + if any( + getattr(request.collection, name) is not None + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size") + ): + print(" Reviewed role CUDA graph settings configure the selected launch; verify the resolved runtime.") + elif request.collection.prefill_cudagraph_policy == "runtime": + print(" CUDA graph sizes resolve in the initialized runtime; no explicit capture override is proposed.") + elif request.worker_type == "aggregated": + print(" The reviewed capture override configures both aggregated probe launches.") + else: + print(" Prefill capture overrides configure the engine; match the serving target.") + print(" The sequence bound does not request every prefill batch; total KV capacity resolves at runtime.") + print(" A complete generated grid does not establish AgentX/direct-FPM query coverage.") + print(" Synthetic validation traffic and SLAs do not determine these collection limits.") + print(" Resource *_bytes values are bytes per rank; kv_bytes_per_token is bytes per cached token per rank.") + if draft.resolved.get("cache_layout") == "grouped": + print(" Grouped pages retain full history or the declared window, with separate convolution state.") + print(" Group page_size_bytes sums every group layer per rank; verify runtime block sizes and padding.") + print( + " Config-derived pages are minimum packed tensor estimates; edit cache_groups for runtime allocation." ) - if name not in _CORRECTION_PROMPTS: - print("Choose one of: " + ", ".join(key.replace("_", "-") for key in _CORRECTION_PROMPTS)) - _prompt(args, name) + print(" Logical context determines FPM timing queries; retained cache length is only memory accounting.") + print(" max_num_tokens/max_batch_size are per-rank scheduler limits. Estimates require runtime verification.") + memory = request.profile_deployment().resources + print(f" Memory source: {memory.memory_source}.") + if not memory.memory_ready: + print(" No activation or non-KV byte estimate is required.") + print(" Displayed byte estimates are planning evidence. Finalize the collected memory before simulation.") + if request.identity.framework_version == "0.27.0": + print(" Memory is observed during initialized vLLM 0.27.0 collection.") + print( + " Collection policy: synchronous scheduling in both phases, " + "matching supported prefill benchmarking." + ) + else: + print(" Memory observation currently supports vLLM 0.27.0; this version can collect timings only.") + for name, value in draft.resolved.items(): + print(f" {name}: {value} (source: {draft.sources[name]})") + if draft.resolved.get("comm_quant_mode") != "half": + print( + " Collection compatibility: this collector currently supports comm_quant_mode=half. " + "Your override is retained, but new collection will reject a mismatched identity." + ) + if "provenance" not in draft.resolved: + print( + " Provenance records the config SHA-256, deployment, values and sources; " + "edit provenance to add a note." + ) + if directory_output: + print("Nothing has been saved. All selected profiles must be accepted before any files are saved.") + print("Accept keeps this profile; edit changes a profile field. Ctrl-C cancels setup.") + else: + print( + "Nothing has been saved. Accept saves these values; edit changes a profile field. Ctrl-C cancels setup." + ) + while True: + action = input("Review action (accept/edit/cancel): ").strip().lower() + if action in ("accept", "edit", "cancel"): + break + print("Enter accept, edit, or cancel. Explicit acceptance is required to save.") + if action == "accept": + return request + if action == "cancel": + raise KeyboardInterrupt + collection_fields = { + "runtime_context_length", + "max_prefill_cudagraph_size", + "prefill_cudagraph_policy", + "gpu_memory_utilization", + "worker_type", + "cudagraph_mode", + "cudagraph_capture_sizes", + "max_cudagraph_capture_size", + } + print("Editable fields: " + ", ".join(sorted(OVERRIDE_FIELDS | collection_fields))) + while True: + name = input("Field to edit: ").strip() + if name in OVERRIDE_FIELDS | collection_fields: + break + print("Choose an editable field by its displayed name.") + # Derive from explicit inputs again so inferred dependents can change. + # Only publish the staged values after the entire request validates. + staged = deepcopy(overrides) + payload = request.model_dump(exclude={"fpm_profile"}) + if name in collection_fields: + if name in {"worker_type", "cudagraph_mode"}: + choices = ("aggregated", "prefill", "decode") if name == "worker_type" else (*CUDAGRAPH_MODES, "clear") + while (value := input(f"{name} ({'/'.join(choices)}): ").strip()) not in choices: + print("Enter one of: " + ", ".join(choices)) + if value == "clear": + value = None + elif name == "cudagraph_capture_sizes": + while True: + try: + answer = input(f"{name} (space-separated positive integers, or clear): ").strip() + if answer == "clear": + value = None + break + value = [int(item) for item in answer.split()] + if not value or min(value) <= 0 or value != sorted(set(value)): + raise ValueError + break + except ValueError: + print("Enter increasing, distinct positive integers.") + elif name == "prefill_cudagraph_policy": + while (value := input(f"{name} (runtime/explicit): ").strip()) not in {"runtime", "explicit"}: + print("Enter runtime or explicit.") + if value == "runtime": + payload["collection"]["max_prefill_cudagraph_size"] = None + elif name == "gpu_memory_utilization": + while True: + try: + value = float(input(f"{name} (0 < fraction <= 1): ").strip()) + if not 0 < value <= 1: + raise ValueError + break + except ValueError: + print("Enter a finite fraction greater than 0 and at most 1.") + else: + while True: + try: + clear = ", or clear" if name == "max_cudagraph_capture_size" else "" + answer = input(f"{name} (positive integer{clear}): ").strip() + if clear and answer == "clear": + value = None + break + value = int(answer) + if value <= 0: + raise ValueError + break + except ValueError: + print("Enter a positive integer.") + if name == "runtime_context_length": + payload["search"]["context_length"] = value + elif name == "worker_type": + payload["worker_type"] = value + else: + payload["collection"][name] = value + if name == "max_prefill_cudagraph_size": + payload["collection"]["prefill_cudagraph_policy"] = "explicit" + else: + staged.update(_read_profile_value(name)) + if name in {"max_num_tokens", "max_batch_size"}: + payload["collection"][name] = staged[name] + try: + edited_request = SupportRequest.model_validate(payload) + updated = derive_profile(config, edited_request, staged) + while updated.missing: + field = min(updated.missing, key=_memory_field) + staged, updated = _prompt_profile_value(config, edited_request, staged, updated, field) + if updated.profile is None: + raise ValueError("model config did not produce a complete FPM profile") + validated = SupportRequest.model_validate({**edited_request.model_dump(), "fpm_profile": updated.profile}) + except ValueError as exc: + print(f"Edit rejected: {exc}. Previous profile values retained.") + continue + overrides, draft, request = staged, updated, validated + + +_TOPOLOGY_OPTIONS = ("tensor_parallel", "attention_data_parallel", "moe_tensor_parallel", "moe_expert_parallel") + + +def _explicit_topology(args: argparse.Namespace) -> bool: + return any(getattr(args, name) is not None for name in _TOPOLOGY_OPTIONS) + + +def _config_inputs(args: argparse.Namespace) -> tuple[ModelConfig, dict[str, Any], SupportRequest]: + config = load_model_config(args.model_config) + overrides = _load_resource_overrides(args.resource_overrides) + preview = derive_profile(config, None, overrides) + if not args.suggest_parallel: + print(f"Model config source: {Path(args.model_config).expanduser()} (SHA-256 {config.sha256}).") + for name, note in config.notes.items(): + print(f"Config {name}: {note}") + for name in ("architecture", "context_length", "num_experts"): + if name in preview.resolved: + print(f"{name}: {preview.resolved[name]} (source: {preview.sources[name]})") + if args.interactive: + _require_terminal() + for name in ("architecture", "context_length", "num_experts"): + if name in preview.missing: + overrides, preview = _prompt_profile_value(config, None, overrides, preview, name) + if args.model is None: + args.model = config.suggestions.get("model") + if args.model_kind is None: + experts = preview.resolved.get("num_experts") + args.model_kind = ( + ("moe" if experts > 0 else "dense") if experts is not None else config.suggestions.get("model_kind") + ) + if args.context_length is None and "context_length" in preview.resolved: + args.context_length = min(preview.resolved["context_length"], AGENTX_REFERENCE_CONTEXT) + if args.interactive: + request = _guided_request(args, skip_topology=not _explicit_topology(args)) + else: + try: + request = _request_from_args(args) + except ValidationError as exc: + details = [] + for error in exc.errors(include_url=False): + name = error["loc"][-1] if error["loc"] else None + option = "--" + name.replace("_", "-") if name in _CORRECTION_PROMPTS else "request" + details.append(f" {option}: {error['msg']}") + raise ValueError( + "Complete target identity and collection options before dependent profile derivation:\n" + + "\n".join(details) + + "\nThen supply unresolved profile fields with --resource-overrides PATH (flat JSON/YAML), " + "or use --interactive." + ) from exc + return config, overrides, request + + +def _topology_summary(suggestions: TopologySuggestions) -> str: + lines = [ + "Suggested worker topologies (collection starting points, not a performance ranking):", + f" Hardware: {suggestions.hardware.source}; {suggestions.hardware.assumption}", + ] + for number, candidate in enumerate(suggestions.candidates, 1): + default = " (default)" if candidate == suggestions.default else "" + lines.append( + f" {number}. {candidate.family}, {candidate.required_gpus} GPUs, {candidate.status}{default}: " + + shlex.join(candidate.cli_flags) + ) + lines.extend(f" {reason}" for reason in candidate.reasons) + lines.extend(f" Missing {name}: {reason}" for name, reason in candidate.missing.items()) + if not suggestions.candidates: + lines.append(" No eligible topology remains in the automatic hardware domain.") + for candidate in suggestions.rejected_candidates: + lines.append( + f" Rejected {candidate.family}, {candidate.required_gpus} GPUs: " + "; ".join(candidate.reasons) + ) + lines.append("Runtime compatibility, actual GPU availability and measured performance remain unchecked.") + return "\n".join(lines) + + +def _resolve_shared_profile_inputs( + config: ModelConfig, shared: dict[str, Any], entry_overrides: list[dict[str, Any]] +) -> None: + while True: + for overrides in entry_overrides: + draft = derive_profile(config, None, {**shared, **overrides}) + pending = [ + name + for name in draft.missing + if not _memory_field(name) and name not in {"cache_groups", "cache_block_sizes"} + ] + if pending: + name = pending[0] + print(f"{name}: {draft.missing[name]}") + shared.update(_read_profile_value(name)) + break + else: + return + + +def _select_topologies( + args: argparse.Namespace, + config: ModelConfig, + overrides: dict[str, Any], + request: SupportRequest, + *, + multiple: bool = False, +) -> list[SupportRequest]: + while True: + try: + suggestions = suggest_topologies(config, request, overrides) + if args.interactive and suggestions.candidates: + _resolve_shared_profile_inputs(config, overrides, [{}]) + suggestions = suggest_topologies(config, request, overrides) + break + except (ProfileRequestError, ValidationError) as exc: + if not args.interactive: + raise + _correct_option(args, exc, require_known_field=True) + request = _validate_guided_request(args) + summary = _topology_summary(suggestions) + if not suggestions.candidates or (not args.interactive and suggestions.default is None): + raise ValueError( + "No fully assessed topology default is available.\n" + + summary + + "\nSupply shared precision/layout inputs with --resource-overrides or use --interactive. " + "When planning estimates are unavailable, select an exact topology with the displayed flags. " + "Explicit topology may be outside the automatic shortlist." + ) + print(summary) + selected = [suggestions.default] if suggestions.default is not None else [] + if args.interactive: + default = suggestions.candidates.index(selected[0]) + 1 if selected else None + if default is None: + print("Memory fit is unresolved; select a candidate explicitly. Memory will be observed during collection.") + while True: + answer = input( + f"Choose {'topologies' if multiple else 'topology'} 1-{len(suggestions.candidates)}" + + (f" [{default}]" if default is not None else "") + + (" (comma-separated, or cancel): " if multiple else " (or cancel): ") + ).strip() + if answer.lower() == "cancel": + raise KeyboardInterrupt + if not answer and default is not None: + break + choices = [value.strip() for value in answer.split(",")] if multiple else [answer] + if all(value.isdecimal() and 1 <= int(value) <= len(suggestions.candidates) for value in choices) and len( + {int(value) for value in choices} + ) == len(choices): + selected = [suggestions.candidates[int(value) - 1] for value in choices] + break + numbers = "unique comma-separated candidate numbers" if multiple else "a displayed candidate number" + print(f"Enter {numbers} or cancel. Enter accepts only a fully assessed default.") + for candidate in selected: + print( + f"Selected {candidate.family} with {candidate.required_gpus} collection GPUs: " + + shlex.join(candidate.cli_flags) + ) + return [candidate.apply(request) for candidate in selected] + + +def _topology_values(request: SupportRequest) -> dict[str, int]: + parallel = request.parallelism() + return dict( + zip( + _TOPOLOGY_OPTIONS, + (parallel[name] for name in ("tensor", "attention_data", "moe_tensor", "moe_expert")), + strict=True, + ) + ) + + +def _config_request(args: argparse.Namespace) -> SupportRequest: + automatic = not _explicit_topology(args) + config, overrides, request = _config_inputs(args) + if automatic: + request = _select_topologies(args, config, overrides, request)[0] + for name, value in _topology_values(request).items(): + setattr(args, name, value) + return _complete_config_request(args, config, overrides, request) + + +def _complete_config_request( + args: argparse.Namespace, config: ModelConfig, overrides: dict[str, Any], request: SupportRequest +) -> SupportRequest: + while True: + try: + draft = derive_profile(config, request, overrides) + except (ProfileRequestError, ValidationError) as exc: + if not args.interactive: + raise + _correct_option(args, exc, require_known_field=True) + request = _validate_guided_request(args) + continue + print("Profile inputs (resource bytes per rank; estimates require runtime verification):") + for name, value in draft.resolved.items(): + print(f" {name}: {value} (source: {draft.sources[name]})") + if draft.missing and not args.interactive: + raise ValueError( + "Unresolved profile fields for this deployment:\n" + + "\n".join(f" {name}: {reason}" for name, reason in draft.missing.items()) + + "\nInspect the pinned runtime and checkpoint metadata for unresolved precision/layout facts. " + "An existing launch configuration is optional evidence; values not established locally remain " + "explicit inputs. Supply them in --resource-overrides PATH (JSON/YAML; bytes per rank), " + "or use --interactive." + ) + try: + while draft.missing: + # Resolve dtype/layout/envelope inputs before asking for dependent memory amounts. + name = min(draft.missing, key=_memory_field) + overrides, draft = _prompt_profile_value(config, request, overrides, draft, name) + except (ProfileRequestError, ValidationError) as exc: + _correct_option(args, exc, require_known_field=True) + request = _validate_guided_request(args) + continue + if draft.profile is None: + raise ValueError("model config did not produce a complete FPM profile") + try: + completed = SupportRequest.model_validate({**request.model_dump(), "fpm_profile": draft.profile}) + except ValidationError as exc: + if not args.interactive: + raise + _correct_option(args, exc) + request = _validate_guided_request(args) + continue + return ( + _review_config_profile( + config, completed, overrides, draft, directory_output=getattr(args, "output_dir", None) is not None + ) + if args.interactive + else completed + ) + + +def _load_parallel_configs(path: str) -> list[dict[str, Any]]: + source = Path(path).expanduser() + try: + entries = yaml.load(source.read_text(encoding="utf-8"), Loader=_ResourceOverridesLoader) + except yaml.YAMLError as exc: + raise ValueError(f"malformed parallel configurations in {source}: {exc}") from exc + if not isinstance(entries, list) or not entries: + raise ValueError("--parallel-configs must contain a nonempty JSON/YAML list of objects") + for number, entry in enumerate(entries, 1): + if not isinstance(entry, dict): + raise ValueError(f"parallel configuration {number} must be an object") + unknown = entry.keys() - {*_TOPOLOGY_OPTIONS, "resource_overrides", "worker_type", "collection"} + if unknown: + raise ValueError(f"parallel configuration {number} has unknown fields: {', '.join(sorted(unknown))}") + if "tensor_parallel" not in entry: + raise ValueError(f"parallel configuration {number} requires tensor_parallel") + for name in entry.keys() & set(_TOPOLOGY_OPTIONS): + if type(entry[name]) is not int or entry[name] < 1: + raise ValueError(f"parallel configuration {number}: {name} must be a positive integer") + overrides = entry.get("resource_overrides", {}) + if not isinstance(overrides, dict): + raise ValueError(f"parallel configuration {number}: resource_overrides must be a mapping") + entry["resource_overrides"] = validate_overrides(overrides) + if "worker_type" in entry and entry["worker_type"] not in {"aggregated", "prefill", "decode"}: + raise ValueError(f"parallel configuration {number}: worker_type must be aggregated, prefill or decode") + if "collection" in entry and not isinstance(entry["collection"], dict): + raise ValueError(f"parallel configuration {number}: collection must be a mapping") + return entries + + +def _check_unique_topologies(requests: list[SupportRequest]) -> None: + seen = set() + for request in requests: + topology = _topology_values(request) + key = (request.worker_type, *topology.values()) + if key in seen: + raise ValueError(f"duplicate resolved parallel configuration: {topology}") + seen.add(key) + + +def _config_requests(args: argparse.Namespace) -> list[SupportRequest]: + entries = _load_parallel_configs(args.parallel_configs) if args.parallel_configs is not None else None + automatic = not _explicit_topology(args) and entries is None + implicit_context = args.context_length is None + shared_graph_settings = {} + if entries is not None: + # Validate generic graph settings with each entry's role and overrides. + args = deepcopy(args) + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"): + value = getattr(args, name) + if value is not None: + shared_graph_settings[name] = value + setattr(args, name, None) + config, shared, request = _config_inputs(args) + if entries is not None: + selected = [] + for entry in entries: + payload = request.model_dump() + payload["search"].update({name: entry[name] for name in _TOPOLOGY_OPTIONS if name in entry}) + if "worker_type" in entry: + payload["worker_type"] = entry["worker_type"] + payload["collection"].update(shared_graph_settings) + payload["collection"].update(entry.get("collection", {})) + selected.append(SupportRequest.model_validate(payload)) + else: + selected = _select_topologies(args, config, shared, request, multiple=True) if automatic else [request] + _check_unique_topologies(selected) + if len(selected) > 1 and ( + rank_fields := [name for name in shared if _memory_field(name) or name == "cache_groups"] + ): + raise ValueError( + "Multiple configurations cannot share rank-local resource overrides: " + + ", ".join(sorted(rank_fields)) + + ". Put them in each --parallel-configs entry's resource_overrides or edit each profile interactively." + ) + if args.interactive and not automatic: + _resolve_shared_profile_inputs( + config, shared, [entry["resource_overrides"] for entry in entries] if entries is not None else [{}] + ) + completed = [] + for number, candidate in enumerate(selected, 1): + # Corrections, nested resource inputs and review edits belong only to + # this tuple. Start every profile from the original shared intake. + candidate_args = deepcopy(args) + candidate = candidate.model_copy(deep=True) + candidate_args.worker_type = candidate.worker_type + for name in CollectionSpec.model_fields: + setattr(candidate_args, name, getattr(candidate.collection, name)) + for name, value in _topology_values(candidate).items(): + setattr(candidate_args, name, value) + overrides = deepcopy(shared) + if entries is not None: + overrides.update(deepcopy(entries[number - 1]["resource_overrides"])) + if implicit_context: + context = derive_profile(config, None, overrides).resolved.get("context_length") + if context is not None: + candidate_args.context_length = min(context, AGENTX_REFERENCE_CONTEXT) + candidate.search.context_length = candidate_args.context_length + print(f"Configuration {number}/{len(selected)}: {_topology_values(candidate)}") + completed.append(_complete_config_request(candidate_args, config, overrides, candidate)) + _check_unique_topologies(completed) + return completed + + +def _onboarding_target(path: str | Path) -> Path: + if not str(path): + raise ValueError("onboarding output directory must be a nonempty path") + target = Path(path).expanduser().absolute() + for parent in (target, *target.parents): + if parent.is_symlink(): + raise ValueError(f"onboarding output path must not contain symlinks: {parent}") + if target.exists() and (not target.is_dir() or any(target.iterdir())): + raise ValueError( + f"onboarding output directory {target} must be new or empty; collection results cannot be replaced" + ) + parent = target.parent + while not parent.exists(): + parent = parent.parent + if not parent.is_dir() or not os.access(parent, os.W_OK): + raise ValueError(f"onboarding output parent {parent} must be a writable directory") + return target + + +def _write_onboarding(requests: list[SupportRequest], root: Path) -> dict[str, Any]: + _check_unique_topologies(requests) + for request in requests: + request.scheduler_limits() + root = _onboarding_target(root) + root.parent.mkdir(parents=True, exist_ok=True) + index = {"schema_version": "aisimulate-onboarding/v1", "configurations": []} + with tempfile.TemporaryDirectory(prefix=f".{root.name}-", dir=root.parent) as temporary: + staged = Path(temporary) + for request in requests: + topology = _topology_values(request) + tp, dp, mtp, ep = topology.values() + name = f"tp{tp}-dp{dp}-moe-tp{mtp}-moe-ep{ep}" + if request.worker_type is not None: + name = f"{request.worker_type}-{name}" + final = root / name + _write_request(request, staged / name / "request.yaml", overwrite=False) + profile = request.fpm_profile + assert profile is not None + (staged / name / "fpm-profile.json").write_text(profile.model_dump_json(indent=2) + "\n", encoding="utf-8") + index["configurations"].append( + { + **({"worker_type": request.worker_type} if request.worker_type is not None else {}), + "topology": topology, + "collection_gpus_required": request.worker_gpus, + "request_id": request_id(request), + "request": str(final / "request.yaml"), + "fpm_profile": str(final / "fpm-profile.json"), + "collection_dir": str(final / "collection"), + "plan_command": shlex.join( + [ + "aisimulate", + "onboard", + "plan", + "--config", + str(final / "request.yaml"), + "--output-dir", + str(final / "collection"), + ] + ), + } + ) + (staged / "onboarding.json").write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8") + # Recheck after all reviews and serialization. Rename publishes only a + # complete directory and cannot replace a nonempty collection directory. + _onboarding_target(root) + staged.rename(root) + return index def _init(args: argparse.Namespace) -> int: + if args.resource_overrides and not args.model_config: + raise ValueError("--resource-overrides requires --model-config") + if args.parallel_configs is not None: + if not args.model_config or args.output_dir is None: + raise ValueError("--parallel-configs requires --model-config and --output-dir") + if _explicit_topology(args): + raise ValueError("--parallel-configs cannot be combined with explicit topology flags") + if args.output_dir is not None: + if not args.model_config: + raise ValueError("--output-dir requires --model-config") + if args.overwrite: + raise ValueError("--output-dir cannot be combined with --overwrite; use a fresh output directory") + if args.suggest_parallel: + if not args.model_config: + raise ValueError("--suggest-parallel requires --model-config") + if ( + args.interactive + or args.fpm_profile + or args.profile + or _explicit_topology(args) + or args.output_dir is not None + or args.parallel_configs is not None + ): + raise ValueError( + "--suggest-parallel cannot be combined with --interactive, profiles, explicit topology flags, " + "--output-dir, or --parallel-configs" + ) + config, overrides, request = _config_inputs(args) + _print(suggest_topologies(config, request, overrides).to_dict(), "json") + return 0 + if args.output_dir is not None: + root = _onboarding_target(args.output_dir) + try: + index = _write_onboarding(_config_requests(args), root) + except (EOFError, KeyboardInterrupt): + print("Setup cancelled.", file=sys.stderr) + return 130 + print( + "Requests and profiles saved. Runtime compatibility and FPM data are unchecked; " + "setup has not launched GPU work." + ) + print("Separate collection plan commands are ready; the same GPUs may be reused across separate runs.") + _print({"onboarding_index": str(root / "onboarding.json"), **index}) + return 0 _request_target(args.output, overwrite=args.overwrite) try: - request = _guided_request(args) if args.interactive else _request_from_args(args) + if args.model_config: + request = _config_request(args) + else: + request = _guided_request(args) if args.interactive else _request_from_args(args) path = _write_request(request, args.output, overwrite=args.overwrite) except (EOFError, KeyboardInterrupt): print("Setup cancelled.", file=sys.stderr) return 130 - workload = request.workload print( f"Scope: FPM simulation for {request.identity.model} ({request.identity.model_kind}), " f"vLLM {request.identity.framework_version}, " - f"{request.identity.gpu_count} {request.identity.gpu} GPU(s) on {request.identity.node_count} node(s); " - f"{request.parallel_preset} {request.parallelism()}, {workload.input_tokens}/{workload.output_tokens} tokens, " - f"concurrency {workload.concurrency}, {workload.request_count} requests, " - f"up to {request.search.max_candidates} candidate(s)." - ) - print( - "Request saved. Model metadata/resources, runtime compatibility and FPM data are unchecked; " - "accuracy is not assessed. Setup has not launched GPU work." + f"{request.identity.gpu}; {request.identity.interconnect}; " + f"{request.parallel_preset} {request.parallelism()}; runtime context {request.search.context_length} tokens." ) + print("Collection settings: " + json.dumps(request.collection_settings(), sort_keys=True)) + print(f"Collection GPUs required: {request.worker_gpus} for one selected worker; availability is unchecked.") + if args.model_config: + print( + f"Request saved with FPM metadata; memory source: {request.profile_deployment().resources.memory_source}. " + "Runtime compatibility and FPM data are unchecked; accuracy is not assessed. " + "Setup has not launched GPU work." + ) + else: + print( + "Request saved. Model metadata/resources, runtime compatibility and FPM data are unchecked; " + "accuracy is not assessed. Setup has not launched GPU work." + ) _print( { "request_id": request_id(request), @@ -302,10 +1266,15 @@ def _plan(args: argparse.Namespace) -> int: { "request_id": plan["request_id"], "candidate_count": plan["search"]["candidate_count"], + "collection_gpus_required": request.worker_gpus, "plan": str(root / "support-plan.json"), "commands": plan["outputs"]["commands"], "prerequisites": "runtime_and_fpm_data_unchecked" if request.fpm_profile is not None else "unchecked", - **({"resources": "estimated_from_declared_profile"} if request.fpm_profile is not None else {}), + **( + {"resources": request.profile_deployment().resources.memory_source} + if request.fpm_profile is not None + else {} + ), "accuracy": "not assessed", "next": shlex.join( [ @@ -325,6 +1294,26 @@ def _plan(args: argparse.Namespace) -> int: def run_support_command(args: argparse.Namespace) -> int: + if args.support_action in {"probe-runtime", "import-observations"}: + from .runtime import run_runtime_command + + return run_runtime_command(args) + if args.support_action in {"validate-collection", "validate-serving"}: + from .validation_workflow import validate_collection, validate_serving + + return validate_collection(args) if args.support_action == "validate-collection" else validate_serving(args) + if args.support_action == "finalize": + from .finalization import run_finalization + + return run_finalization(args) + if args.support_action in {"checkpoint", "resume"}: + from .checkpoint import run_checkpoint_command + + return run_checkpoint_command(args) + if args.support_action == "validate-fpm": + from .validation import run_validation + + return run_validation(args) if args.support_action == "init": return _init(args) if args.support_action == "plan": @@ -334,6 +1323,7 @@ def run_support_command(args: argparse.Namespace) -> int: SupportRequest.from_yaml(args.config), output_dir=args.output_dir, execute=args.execute, + check_readiness=args.check_readiness, smoke=args.smoke, limit=args.limit, resume=args.resume, diff --git a/python/aisimulate/src/aisimulate/support/collection_readiness.py b/python/aisimulate/src/aisimulate/support/collection_readiness.py new file mode 100644 index 000000000..8bd8f1ca4 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/collection_readiness.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read-only collection readiness from attempt-bound native evidence.""" + +from __future__ import annotations + +import json +import re +from collections import Counter +from pathlib import Path +from typing import Any + +from .schema import SupportRequest + +REPORT_FILENAME = "fpm-readiness.json" + + +def _cell_report(plan, cell, campaign: Path, entry: Any) -> dict[str, Any]: + from collector.fpm_forward.database import aggregate_cell + from collector.fpm_forward.execution_evidence import _effective_launch, file_evidence, inspect_execution_evidence + from collector.fpm_forward.native_artifact import _rank_artifacts, validate_native_collection + + directory = campaign / "cells" / cell.cell_id + raw = directory / "raw" + entry = entry if isinstance(entry, dict) else {} + result = { + "cell": cell.to_dict(), + "checkpoint_status": entry.get("status", "missing"), + "attempt_id": entry.get("attempt_id"), + "status": "incomplete", + "blockers": [], + "raw_artifact_root": str(raw), + "native_regime_counts": {}, + "formal_regime_counts": {}, + "classification_scope": ( + "formal_regime_counts and direct_eligible_points describe current reaggregation of verified native " + "artifacts; they do not assess or upgrade an existing published table. Existing publications keep " + "their saved classification. Publish revised classification to a fresh output tree and retain the " + "original pair." + ), + "direct_eligible_points": 0, + "coverage": "unverified", + "rank_diagnostics": [], + "launch_evidence": [], + "cache_observation_evidence": [], + } + try: + result["launch_evidence"] = [ + file_evidence(directory / name) + for name in ("generator-request.json", "run.sh", "fpm_env.sh", "collector-runtime-env.sh") + if (directory / name).is_file() + ] + result["cache_observation_evidence"] = [ + file_evidence(path) for path in sorted(raw.glob("**/fpm-memory-*.json")) + ] + # Preserve diagnostics even when protocol validation rejects the attempt. + # They remain explicitly unverified; never use these counts for admission. + for path, payload in _rank_artifacts(raw): + result["rank_diagnostics"].append( + { + "source": file_evidence(path), + "verified": False, + "kvwarm": payload.get("kvwarm"), + "regime_counts": dict( + Counter( + row["kv_seed_regime"] + if isinstance(row.get("kv_seed_regime"), str) + else "unreported" + if row.get("kv_seed_regime") is None + else "invalid" + for row in payload.get("results", []) + if isinstance(row, dict) + ) + ), + } + ) + attempt = entry.get("attempt_id") + if not isinstance(attempt, str) or not attempt: + raise ValueError("checkpoint has no current attempt identity; inspect existing artifacts before retrying") + collection = validate_native_collection( + cell, raw, expected_plan_sha256=plan.sha256, expected_attempt_id=attempt + ) + result.update( + native_regime_counts=dict(Counter(item.kv_seed_regime or "unreported" for item in collection.points)), + kvwarm=collection.kvwarm_meta, + observed_backend_version=collection.backend_version, + runtime_run_id=collection.runtime_run_id, + runtime_grid_digest=collection.runtime_grid_digest, + ) + for diagnostic in result["rank_diagnostics"]: + diagnostic["verified"] = True + rows = aggregate_cell(plan, cell, directory, expected_attempt_id=attempt) + formal_counts = Counter(row["kv_seed_regime"] for row in rows) + # This is the existing direct consumer's row eligibility rule, not an + # interpolation implementation or a claim of query-domain coverage. + eligible = len(rows) - formal_counts["fake_fallback"] + result.update( + formal_regime_counts=dict(formal_counts), + direct_eligible_points=eligible, + coverage="partial" if formal_counts["fake_fallback"] else "sampled_points_only", + execution=inspect_execution_evidence(cell, raw, collection, plan=plan), + ) + if cell.workload_kind == "decode" and collection.kvwarm_meta is None: + result["blockers"].append("decode KV warm-up metadata is unreported; inspect the runtime protocol") + if not eligible: + result["blockers"].append( + "no direct-eligible timing points; inspect KV warm-up eligibility/skip reason and fallback logs; " + "an unchanged full collection is not justified" + ) + captures = plan.options.prefill_sampling.cudagraph_capture_sizes + expected = dict(plan.options.compilation_config) + if captures is not None and (cell.workload_kind == "prefill" or plan.options.worker_type is not None): + expected.update( + cudagraph_capture_sizes=list(captures), + max_cudagraph_capture_size=plan.options.prefill_sampling.max_cudagraph_capture_size, + ) + if expected and not plan.options.enforce_eager: + flags = _effective_launch(plan, directory) + configured = json.loads(flags.get("--compilation-config", "{}")) + if not isinstance(configured, dict) or any(configured.get(key) != value for key, value in expected.items()): + result["blockers"].append( + "saved serving launch does not preserve the reviewed explicit graph configuration" + ) + # Effective modes may downgrade for the initialized backend. Full + # memory import verifies initial settings and those downgrades; + # the graph size fields must still match here when reported. + expected.pop("cudagraph_mode", None) + for worker in result["execution"]["observed_workers"]: + graph = worker.get("graph_config") + if isinstance(graph, dict) and any( + key in graph and graph[key] != value for key, value in expected.items() + ): + result["blockers"].append( + "observed prefill graph configuration differs from the reviewed explicit settings" + ) + # Dynamo b83b1d9304ebfc624709ac46db32b1b6f1ff1615 instrumented_scheduler.py:2302-2350,6223-6236: + # Full native captures reflect compilation_config; phase lists may be filtered. + native_expected = { + native: expected[configured] + for configured, native in ( + ("cudagraph_capture_sizes", "capture_sizes"), + ("max_cudagraph_capture_size", "max_capture_size"), + ) + if configured in expected + } + for native in result["execution"]["native_graph_config"]: + graph = native["config"] + if isinstance(graph, dict) and any( + key in graph and graph[key] != value for key, value in native_expected.items() + ): + result["blockers"].append( + "native prefill graph configuration differs from the reviewed explicit settings" + ) + if result["execution"]["failures"]: + result["blockers"].extend(result["execution"]["failures"]) + cpu = result["execution"]["cpu_affinity"] + if cpu["policy_required"]: + result["blockers"].extend(f"CPU affinity: {message}" for message in cpu["missing_evidence"]) + if entry.get("cleanup_error"): + result["blockers"].append( + "current attempt has unresolved resource cleanup; retain its cleanup error and inspect" + ) + if entry.get("status") != "passed": + result["blockers"].append( + "current attempt is not passed; inspect or resume saved evidence before collection" + ) + result["status"] = "blocked" if result["blockers"] else "ready" + except (OSError, TypeError, ValueError, KeyError) as error: + result["blockers"].append(str(error)) + result["status"] = "blocked" if result["rank_diagnostics"] else "incomplete" + result["next_action"] = ( + "full collection may proceed; repeatability, query coverage, memory and serving accuracy remain separate" + if result["status"] == "ready" + else "inspect preserved evidence and resolve the blocker before an unchanged full collection; " + "qualify a runtime/launch change with a new bounded campaign when required" + ) + return result + + +def _matches_request(request: SupportRequest, plan) -> bool: + if ( + plan.model_path != request.identity.model + or plan.system != request.identity.gpu + or plan.backend != request.identity.framework + or plan.capability.aic_database_version != request.identity.framework_version + or plan.fpm_profile != request.fpm_profile + ): + return False + scheduler = request.scheduler_limits() + parallel = request.parallelism() + return ( + plan.options.vllm_max_model_len == request.search.context_length + and plan.options.max_num_batched_tokens == scheduler["max_batched_tokens"] + and plan.options.max_num_seqs == scheduler["max_sequences"] + and plan.options.prefill_cudagraph_policy == request.collection.prefill_cudagraph_policy + and plan.options.max_prefill_cudagraph_size == request.collection.max_prefill_cudagraph_size + and plan.options.worker_type == request.worker_type + and plan.options.cudagraph_mode == request.collection.cudagraph_mode + and plan.options.cudagraph_capture_sizes + == ( + tuple(request.collection.cudagraph_capture_sizes) + if request.collection.cudagraph_capture_sizes is not None + else None + ) + and plan.options.max_cudagraph_capture_size == request.collection.max_cudagraph_capture_size + and (plan.options.gpu_memory_utilization or 0.9) == request.collection.memory_fraction + and all( + cell.topology.tp == parallel["tensor"] + and cell.topology.dp == parallel["attention_data"] + and cell.topology.moe_tp == parallel["moe_tensor"] + and cell.topology.moe_ep == parallel["moe_expert"] + for cell in plan.cells + ) + ) + + +def assess_readiness( + request: SupportRequest, + root: Path, + checkpoint_root: Path, + *, + expected_plan=None, +) -> dict[str, Any]: + """Inspect saved sources without model fetching, reruns or checkpoint writes. + + A formal attempt supersedes smoke evidence for the same cell, even when it + failed. Only the selected checkpoint's current attempt can qualify a cell. + """ + from collector.fpm_forward.execution_evidence import file_evidence + from collector.fpm_forward.repeatability import load_repeatability_deployment, load_repeatability_source + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + campaigns = [] + selected = {} + blockers = [] + for smoke in (True, False): + checkpoint_path = checkpoint_root / ("fpm_forward_smoke.json" if smoke else "fpm_forward.json") + if not checkpoint_path.exists(): + continue + report = {"kind": "smoke" if smoke else "formal", "cells": [], "errors": []} + campaigns.append(report) + try: + report["checkpoint"] = file_evidence(checkpoint_path) + checkpoint = json.loads(checkpoint_path.read_text()) + if not isinstance(checkpoint, dict): + raise ValueError("collector checkpoint must be an object") + sha = checkpoint.get("plan_sha256") + if ( + checkpoint.get("schema") != CHECKPOINT_SCHEMA + or not isinstance(sha, str) + or not re.fullmatch(r"[0-9a-f]{64}", sha) + or not isinstance(checkpoint.get("cells"), dict) + ): + raise ValueError("invalid collector checkpoint identity or cell records") + campaign = root / "fpm-artifacts" / sha[:16] + if smoke: + campaign /= "smoke" + plan = load_repeatability_source(campaign) + if plan.sha256 != sha or not _matches_request(request, plan): + raise ValueError("saved collection differs from the reviewed onboarding request") + if expected_plan is not None and plan.sha256 != expected_plan.sha256: + raise ValueError( + "saved collection differs from the current frozen runtime/launch plan; use a new campaign" + ) + deployment = load_repeatability_deployment(campaign) + report.update( + plan_sha256=plan.sha256, + source_plan=file_evidence(campaign / "collection-plan.json"), + runtime_identity={ + "model": plan.model_path, + "model_revision": request.identity.model_revision, + "system": plan.system, + "backend": plan.backend, + "backend_version": plan.capability.aic_database_version, + "collector_revision": plan.aic_revision, + "generator_config_sha256": plan.generator_config_sha256, + "deployment": deployment, + "executor": plan.options.executor, + "slurm_image": plan.options.slurm_container_image, + "runtime_observation": plan.to_dict().get("runtime_observation"), + "scope": "frozen requested identity; observed backend and launch evidence are reported per cell", + }, + ) + for cell in plan.cells: + entry = checkpoint["cells"].get(cell.cell_id) + cell_report = _cell_report(plan, cell, campaign, entry) + report["cells"].append(cell_report) + if cell.cell_id not in selected or cell.cell_id in checkpoint["cells"]: + selected[cell.cell_id] = {"campaign": report["kind"], "plan_sha256": plan.sha256, **cell_report} + except (OSError, TypeError, ValueError, KeyError) as error: + report["errors"].append(str(error)) + blockers.append(f"{report['kind']}: {error}") + if len({item["plan_sha256"] for item in selected.values()}) > 1: + blockers.append("selected phase evidence belongs to different frozen runtime/launch plans") + if expected_plan is not None and set(selected) != {cell.cell_id for cell in expected_plan.cells}: + blockers.append("not every current planned cell has saved readiness evidence") + if {item["cell"]["workload_kind"] for item in selected.values()} != set(request.collection_phases): + blockers.append("readiness requires every phase of the selected serving role") + for item in selected.values(): + blockers.extend( + f"{item['cell']['workload_kind']} {item['cell']['cell_id']}: {reason}" for reason in item["blockers"] + ) + ready = bool(selected) and not blockers and all(item["status"] == "ready" for item in selected.values()) + return { + "status": "ready" if ready else "blocked" if campaigns else "incomplete", + "ready_for_full_collection": ready, + "scope": "sampled native timing readiness only; not complete query coverage, memory or serving accuracy", + "request_identity": request.identity.model_dump(mode="json"), + "blockers": blockers, + "selected_cells": list(selected.values()), + "campaigns": campaigns, + } + + +def resume_without_workers(plan, checkpoint_path: Path) -> bool: + """Recognize only the runner's guaranteed no-worker resume path. + + Running/interrupted/missing cells can launch cleanup or workers. Do not + admit those states, even if their raw files look complete. Failed and + cleanup-failed cells are skipped without --resume-retry-failed, which the + guided CLI does not expose. The collector still revalidates recovery. + """ + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + try: + checkpoint = json.loads(checkpoint_path.read_text()) + return ( + isinstance(checkpoint, dict) + and checkpoint.get("schema") == CHECKPOINT_SCHEMA + and checkpoint.get("plan_sha256") == plan.sha256 + and isinstance(checkpoint.get("cells"), dict) + and bool(plan.cells) + and all( + isinstance(entry := checkpoint["cells"].get(cell.cell_id), dict) + and entry.get("status") in {"passed", "failed", "cleanup_failed"} + and isinstance(entry.get("attempt_id"), str) + and bool(entry["attempt_id"]) + for cell in plan.cells + ) + ) + except (OSError, TypeError, ValueError): + return False diff --git a/python/aisimulate/src/aisimulate/support/config_profile.py b/python/aisimulate/src/aisimulate/support/config_profile.py new file mode 100644 index 000000000..2b348e8c2 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/config_profile.py @@ -0,0 +1,1337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Adapted and modified from vLLM at commit 98dff2a81d747d1dba01a47f939f48c3526d4206: +# https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py +# https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/model.py +# https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/configs.py + +"""Local config facts and explicitly bounded resource estimates for onboarding. + +This module reads JSON and packaged hardware metadata. It never instantiates a +model, imports remote configuration code, or builds an operation graph. Tensor +storage estimates cover the declared vanilla decoder layouts, not arbitrary +implementations sharing their field names. Runtime reservations and activation +estimates are not measurements or a guarantee that a deployment fits. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate import quantization +from aisimulate.fpm_profile import FpmCacheGroup, FpmModelProfile + +from .schema import SupportRequest + +_RESOURCE_FIELDS = ( + "weights_bytes", + "activations_bytes", + "runtime_overhead_bytes", + "comm_overhead_bytes", + "kv_bytes_per_token", + "cache_layout", + "cache_groups", + "max_num_tokens", + "max_batch_size", +) +_MODE_ENUMS = { + "gemm_quant_mode": quantization.GEMMQuantMode, + "moe_quant_mode": quantization.MoEQuantMode, + "fmha_quant_mode": quantization.FMHAQuantMode, + "comm_quant_mode": quantization.CommQuantMode, + "kv_cache_dtype": quantization.KVCacheQuantMode, +} +_DEPLOYMENT_FIELDS = (*_MODE_ENUMS, "moe_backend", "attention_backend") +_MODEL_FIELDS = ("architecture", "context_length", "num_experts") +OVERRIDE_FIELDS = frozenset((*_MODEL_FIELDS, *_RESOURCE_FIELDS, *_DEPLOYMENT_FIELDS, "cache_block_sizes", "provenance")) +INTEGER_FIELDS = frozenset( + ( + "context_length", + "num_experts", + *(field for field in _RESOURCE_FIELDS if field not in {"cache_layout", "cache_groups"}), + ) +) +FIELD_CHOICES = {field: tuple(enum.__members__) for field, enum in _MODE_ENUMS.items()} +FIELD_CHOICES["cache_layout"] = ("linear", "grouped") +_NONNEGATIVE_FIELDS = frozenset( + ("num_experts", "weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes") +) +_BYTE_FIELDS = frozenset(field for field in INTEGER_FIELDS if "bytes" in field) +_NON_KV_FIELDS = _BYTE_FIELDS - {"kv_bytes_per_token"} +_CACHE_GROUP_NAMES = frozenset(("full_attention", "sliding_attention", "full_convolution", "sliding_convolution")) +_INKLING_ARCHITECTURE = "InklingForConditionalGeneration" +_INKLING_CACHE_SOURCE = ( + "https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/" + "vllm/models/inkling/nvidia/sconv_swa_attn.py" +) + +_ARCHITECTURES = { + "llama": "LlamaForCausalLM", + "mistral": "MistralForCausalLM", + "mixtral": "MixtralForCausalLM", + "qwen2": "Qwen2ForCausalLM", + "qwen3": "Qwen3ForCausalLM", + "qwen3_moe": "Qwen3MoeForCausalLM", + "minimax_m2": "MiniMaxM2ForCausalLM", + "glm_moe_dsa": "GlmMoeDsaForCausalLM", + "deepseek_v3": "DeepseekV3ForCausalLM", + "deepseek_v32": "DeepseekV32ForCausalLM", +} +_DENSE = frozenset(("LlamaForCausalLM", "MistralForCausalLM", "Qwen2ForCausalLM", "Qwen3ForCausalLM")) +_FULL_ATTENTION = _DENSE | {"MixtralForCausalLM", "Qwen3MoeForCausalLM", "MiniMaxM2ForCausalLM"} +_WEIGHT_LAYOUTS = _DENSE | {"MixtralForCausalLM"} +_EXPERT_ALIASES = ("num_local_experts", "n_routed_experts", "num_experts") +_INT_ALIASES = { + "hidden_size": ("hidden_size", "n_embd", "d_model"), + "intermediate_size": ("intermediate_size",), + "num_hidden_layers": ("num_hidden_layers", "n_layer", "num_layers"), + "num_attention_heads": ("num_attention_heads", "n_head"), + "num_key_value_heads": ("num_key_value_heads", "num_kv_heads"), + "head_dim": ("head_dim",), + "vocab_size": ("vocab_size",), + "context_length": ("max_position_embeddings", "n_positions", "max_seq_len", "seq_length", "model_max_length"), + "num_experts": _EXPERT_ALIASES, +} +_OPTIONAL_DIMENSIONS = frozenset(("num_key_value_heads", "head_dim", "intermediate_size")) +_EXTRA_DIMENSIONS = ( + "moe_intermediate_size", + "num_experts_per_tok", + "kv_lora_rank", + "q_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "index_head_dim", + "index_n_heads", +) +_AUXILIARY_DIMENSIONS = ( + "n_shared_experts", + "shared_intermediate_size", + "shared_expert_intermediate_size", + "num_nextn_predict_layers", + "num_mtp_modules", + "mtp_transformer_layers", + "first_k_dense_replace", +) + + +@dataclass(frozen=True) +class ModelConfig: + raw: dict[str, Any] + # Checkpoint/collector identity in suggestions may name a multimodal wrapper. + decoder_architecture: str | None + sha256: str + suggestions: dict[str, Any] + notes: dict[str, str] + source_files: dict[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ProfileDraft: + profile: FpmModelProfile | None + resolved: dict[str, Any] + missing: dict[str, str] + sources: dict[str, str] + + +class ProfileRequestError(ValueError): + """Config facts conflict with a correctable onboarding request field.""" + + def __init__(self, field: str, message: str): + super().__init__(message) + self.field = field + + +def _integer(value: Any, field: str, *, minimum: int = 1, maximum: int | None = None) -> int: + if type(value) is not int or value < minimum or (maximum is not None and value > maximum): + bound = f"{minimum}..{maximum}" if maximum is not None else f">={minimum}" + raise ValueError(f"{field} must be an integer {bound}, not a boolean, float or numeric string") + return value + + +def _text(value: Any, field: str) -> str: + if not isinstance(value, str) or not value.strip() or "\x00" in value: + raise ValueError(f"{field} must be a nonempty string without NUL characters") + return value.strip() + + +def _aliased_int(raw: Mapping[str, Any], name: str) -> int | None: + supplied = {} + for key in _INT_ALIASES[name]: + if key not in raw or (raw[key] is None and name in _OPTIONAL_DIMENSIONS): + continue + supplied[key] = _integer(raw[key], key, minimum=0 if name == "num_experts" else 1) + if len(set(supplied.values())) > 1: + raise ValueError(f"conflicting {name} fields: {supplied}") + return next(iter(supplied.values()), None) + + +def _dtype(raw: Mapping[str, Any]) -> str | None: + values = { + key: _text(raw[key], key).removeprefix("torch.") for key in ("dtype", "torch_dtype") if raw.get(key) is not None + } + if len(set(values.values())) > 1: + raise ValueError(f"conflicting dtype fields: {values}") + return next(iter(values.values()), None) + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate config JSON key: {key}") + result[key] = value + return result + + +def _finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError("config JSON numbers must be finite") + return result + + +def _architecture(raw: Mapping[str, Any]) -> str | None: + architectures = raw.get("architectures") + architecture = None + if architectures is not None: + if not isinstance(architectures, list) or len(architectures) != 1: + raise ValueError("architectures must identify exactly one decoder architecture") + architecture = _text(architectures[0], "architectures[0]") + model_type = _text(raw["model_type"], "model_type") if "model_type" in raw else None + expected = _ARCHITECTURES.get(model_type) + if expected and architecture and expected != architecture: + raise ValueError(f"model_type={model_type!r} conflicts with architectures={architecture!r}") + return architecture or expected + + +def _kv_cache_scheme(quant: Mapping[str, Any]) -> dict[str, Any]: + scheme = quant.get("kv_cache_scheme") + if scheme is None: + return {} + if isinstance(scheme, str): + if scheme.strip().upper() == "FP8": + return {"num_bits": 8, "type": "float"} + raise ValueError("unsupported quantization_config.kv_cache_scheme string; supported value is FP8") + if not isinstance(scheme, dict): + raise ValueError("quantization_config.kv_cache_scheme must be an object or FP8") + return scheme + + +def _validate_quantization(raw: Mapping[str, Any]) -> None: + quant = raw.get("quantization_config") + if quant is None: + return + if not isinstance(quant, dict): + raise ValueError("quantization_config must be an object") + for key in ("quant_method", "quant_algo", "activation_scheme"): + if key in quant: + _text(quant[key], f"quantization_config.{key}") + if quant.get("quant_method") == "fp8" and quant.get("quant_algo") not in (None, "FP8"): + raise ValueError("conflicting quantization_config quant_method and quant_algo") + blocks = quant.get("weight_block_size") + if blocks is not None: + if not isinstance(blocks, list) or len(blocks) != 2: + raise ValueError("quantization_config.weight_block_size must contain two positive integers") + for block in blocks: + _integer(block, "quantization_config.weight_block_size") + for key in ("ignore", "modules_to_not_convert"): + if key in quant: + if not isinstance(quant[key], list): + raise ValueError(f"quantization_config.{key} must be a list of module names") + for name in quant[key]: + _text(name, f"quantization_config.{key}") + if quant.get("config_groups") is not None and not isinstance(quant["config_groups"], dict): + raise ValueError("quantization_config.config_groups must be an object") + for group_name, group in (quant.get("config_groups") or {}).items(): + if not isinstance(group, dict): + raise ValueError(f"quantization_config.config_groups.{group_name} must be an object") + for key in ("input_activations", "weights"): + tensor = group.get(key) + if tensor is None: + continue + if not isinstance(tensor, dict): + raise ValueError(f"quantization_config.config_groups.{group_name}.{key} must be an object") + if "dynamic" in tensor and type(tensor["dynamic"]) is not bool: + raise ValueError(f"quantization_config.config_groups.{group_name}.{key}.dynamic must be a boolean") + if "num_bits" in tensor: + _integer(tensor["num_bits"], f"quantization_config.config_groups.{group_name}.{key}.num_bits") + scheme = _kv_cache_scheme(quant) + if scheme.get("type") is not None: + _text(scheme["type"], "quantization_config.kv_cache_scheme.type") + if "num_bits" in scheme: + _integer(scheme["num_bits"], "quantization_config.kv_cache_scheme.num_bits") + if "dynamic" in scheme and type(scheme["dynamic"]) is not bool: + raise ValueError("quantization_config.kv_cache_scheme.dynamic must be a boolean") + + +def _cache_window(raw: Mapping[str, Any]) -> int | None: + supplied = { + key: _integer(raw[key], key, minimum=0) + for key in ("sliding_window", "sliding_window_size") + if raw.get(key) is not None + } + if len(set(supplied.values())) > 1: + raise ValueError(f"conflicting sliding-window fields: {supplied}") + return next(iter(supplied.values()), None) + + +def _cache_layer_ids(raw: Mapping[str, Any], architecture: str | None, *, inkling: bool = False) -> tuple[int, ...]: + """Resolve explicit attention retention without assuming an operation graph.""" + window = _cache_window(raw) + if raw.get("use_sliding_window") is True and not window: + raise ValueError("use_sliding_window=true requires a positive sliding_window or sliding_window_size") + layers = _aliased_int(raw, "num_hidden_layers") + # The pinned Inkling config defaults absent/null local_layer_ids to []; a + # window size alone must not turn its full-attention layers into SWA. + # https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/configs.py + local = () if inkling else None + if "local_layer_ids" in raw and not (inkling and raw["local_layer_ids"] is None): + local = raw["local_layer_ids"] + if not isinstance(local, list) or any(type(value) is not int or value < 0 for value in local): + raise ValueError("local_layer_ids must be a list of nonnegative layer indexes") + if len(local) != len(set(local)) or (layers is not None and any(value >= layers for value in local)): + raise ValueError("local_layer_ids must be unique indexes below num_hidden_layers") + local = tuple(sorted(local)) + if "layer_types" in raw: + types = raw["layer_types"] + if not isinstance(types, list) or (layers is not None and len(types) != layers): + raise ValueError("layer_types must have one entry per num_hidden_layers") + allowed = {"full_attention", "sliding_attention"} + if architecture in {"GlmMoeDsaForCausalLM", "DeepseekV32ForCausalLM"}: + allowed.add("deepseek_sparse_attention") + if any(not isinstance(value, str) or value not in allowed for value in types): + raise ValueError("unsupported mixed or recurrent attention layout in layer_types") + indexed = tuple(index for index, value in enumerate(types) if value == "sliding_attention") + if local is not None and indexed != local: + raise ValueError("conflicting layer_types and local_layer_ids attention layouts") + local = indexed + if raw.get("use_sliding_window") is False: + if local: + raise ValueError("use_sliding_window=false conflicts with sliding attention layers") + return () + if local is None and window: + if any(raw.get(key) for key in ("max_window_layers", "sliding_window_pattern")): + raise ValueError("unsupported implicit sliding-window schedule; provide explicit layer_types") + local = tuple(range(layers or 0)) + if local and not window: + raise ValueError("sliding attention layers require a positive sliding_window or sliding_window_size") + return local or () + + +def _uses_grouped_cache(raw: Mapping[str, Any], architecture: str | None) -> bool: + return bool( + raw.get("use_sconv") + or _cache_layer_ids(raw, architecture) + or ( + _cache_window(raw) + and raw.get("use_sliding_window") is not False + and "layer_types" not in raw + and "local_layer_ids" not in raw + ) + ) + + +def _validate_layout(raw: Mapping[str, Any], architecture: str | None, *, inkling: bool = False) -> None: + for key in ( + "attention_bias", + "mlp_bias", + "tie_word_embeddings", + "is_encoder_decoder", + "use_sliding_window", + "use_sconv", + ): + if key in raw and type(raw[key]) is not bool: + raise ValueError(f"{key} must be a boolean") + if raw.get("is_encoder_decoder"): + raise ValueError("unsupported encoder-decoder text layout; supply a decoder-only config or FPM profile") + if inkling and raw.get("use_sconv") is not True: + raise ValueError( + "pinned Inkling runtime allocates convolution state in every decoder layer; config-based derivation " + "requires explicit use_sconv=true. Verify the runtime/config or provide a complete FPM profile. " + f"Source: {_INKLING_CACHE_SOURCE.rsplit('/', 1)[0]}/model.py" + ) + if any(raw.get(key) is not None for key in ("ssm_cfg", "mamba_d_state", "linear_conv_kernel_dim")) or any( + name in (architecture or "").lower() for name in ("mamba", "jamba", "rwkv") + ): + raise ValueError( + "unsupported recurrent/hybrid state; only explicit attention and windowed convolution caches are supported" + ) + sliding = _cache_layer_ids(raw, architecture, inkling=inkling) + layers = _aliased_int(raw, "num_hidden_layers") + if "attn_type_list" in raw: + values = raw["attn_type_list"] + if not isinstance(values, list) or (layers is not None and len(values) != layers): + raise ValueError("attn_type_list must have one entry per num_hidden_layers") + if any(type(value) is not int or value != 1 for value in values): + raise ValueError("unsupported mixed or non-full attention layout in attn_type_list") + if sliding: + raise ValueError("full-attention attn_type_list conflicts with sliding-window layer declarations") + if architecture in _FULL_ATTENTION and any( + raw.get(key) for key in ("kv_lora_rank", "q_lora_rank", "compress_ratios", "attention_k_eq_v") + ): + raise ValueError("unsupported cache/projection modifiers conflict with the declared full-attention layout") + for key in ( + "swa_num_attention_heads", + "swa_num_key_value_heads", + "swa_head_dim", + "swa_v_head_dim", + "sconv_kernel_size", + "conv_kernel_size", + ): + if raw.get(key) is not None: + _integer(raw[key], key) + if ( + raw.get("sconv_kernel_size") is not None + and raw.get("conv_kernel_size") is not None + and raw["sconv_kernel_size"] != raw["conv_kernel_size"] + ): + raise ValueError("conflicting sconv_kernel_size and conv_kernel_size") + if raw.get("use_sconv"): + if not inkling: + raise ValueError("unsupported use_sconv state outside the verified Inkling cache layout") + if not (raw.get("sconv_kernel_size") or raw.get("conv_kernel_size")): + raise ValueError("use_sconv requires an explicit sconv_kernel_size or conv_kernel_size") + if _dtype(raw) not in (None, "bfloat16"): + raise ValueError("Inkling convolution cache requires bfloat16 runtime storage") + + +def _validate_architecture( + raw: Mapping[str, Any], architecture: str | None, experts: int | None, *, inkling: bool = False +) -> None: + _validate_layout(raw, architecture, inkling=inkling) + experts_per_token = raw.get("num_experts_per_tok") + if experts and experts_per_token is not None and experts_per_token > experts: + raise ValueError("num_experts_per_tok cannot exceed the routed expert count") + if architecture in _DENSE and experts not in (None, 0): + raise ValueError("num_experts conflicts with the declared dense architecture") + if architecture in set(_ARCHITECTURES.values()) - _DENSE and experts == 0: + raise ValueError("num_experts must be positive for the declared MoE architecture") + + +def load_model_config(path: str | Path) -> ModelConfig: + """Load local config/ModelOpt metadata with separate byte-level provenance.""" + path = Path(path).expanduser().absolute() + payload = path.read_bytes() + raw = json.loads(payload, object_pairs_hook=_unique_object, parse_float=_finite_float, parse_constant=_finite_float) + if not isinstance(raw, dict): + raise ValueError("model config JSON must be an object") + document = raw + source_files = {} + notes = { + "modeling_scope": ( + "Text decoder only. FPM excludes multimodal encoders, projectors, preprocessing and other non-text " + "components from timing and config-derived planning estimates. Observed runtime cache capacity accounts " + "for all components actually loaded by the worker; full multimodal latency is not modeled." + ) + } + if "text_config" in document: + text_config = document["text_config"] + if not isinstance(text_config, dict) or not text_config: + raise ValueError( + "text_config must be a nonempty object describing one text decoder; " + "supply its decoder config or a complete --fpm-profile" + ) + if "text_config" in text_config: + raise ValueError("ambiguous nested text_config; supply one text decoder config or a complete --fpm-profile") + raw = dict(text_config) + notes["decoder_config"] = "config text_config; outer model and encoder geometry are excluded" + inherited = [] + for keys, inferred_keys in ( + (("dtype", "torch_dtype"), ()), + (("quantization_config", "hf_quant_config", "quant_algo"), ("quant_dynamic", "kv_cache_quant_algo")), + ): + if any(raw.get(key) is not None for key in keys): + continue + for key in keys: + if document.get(key) is not None: + raw[key] = document[key] + inherited.append(key) + for key in inferred_keys: + if raw.get(key) is None and document.get(key) is not None: + raw[key] = document[key] + inherited.append(key) + if inherited: + notes["shared_metadata"] = ( + f"config-level {', '.join(inherited)} inherited because text_config does not declare that metadata" + ) + sidecar_path = path.parent / "hf_quant_config.json" + if sidecar_path.is_file(): + from aisimulate_core.sdk.fpm_model_metadata import attach_quantization_sidecar + + sidecar_payload = sidecar_path.read_bytes() + sidecar = json.loads( + sidecar_payload, object_pairs_hook=_unique_object, parse_float=_finite_float, parse_constant=_finite_float + ) + if not isinstance(sidecar, dict): + raise ValueError("hf_quant_config.json must be an object") + if document.get("hf_quant_config") is not None and document["hf_quant_config"] != sidecar: + raise ValueError("conflicting inline hf_quant_config and adjacent hf_quant_config.json") + source_files[sidecar_path.name] = hashlib.sha256(sidecar_payload).hexdigest() + notes["quantization_sidecar"] = ( + f"local ModelOpt sidecar {sidecar_path} sha256={source_files[sidecar_path.name]}; " + "checkpoint quantization metadata; runtime attention precision remains independent" + ) + _validate_quantization(raw) + merged = attach_quantization_sidecar(raw, sidecar) + if raw is document: + document = merged + raw = merged + geometry = {key: _aliased_int(raw, key) for key in _INT_ALIASES} + for key in _EXTRA_DIMENSIONS: + if raw.get(key) is not None: + _integer(raw[key], key) + for key in _AUXILIARY_DIMENSIONS: + if raw.get(key) is not None: + _integer(raw[key], key, minimum=0) + decoder_architecture = _architecture(raw) + architecture = decoder_architecture + if raw is not document: + if raw.get("architectures") is None and (wrapper_architecture := _architecture(document)): + architecture = wrapper_architecture + notes["architecture"] = ( + f"config wrapper architecture={architecture} retained for checkpoint/collector identity" + ) + notes["decoder_architecture"] = ( + f"text_config decoder architecture={decoder_architecture}; used for decoder validation/resource estimates" + if decoder_architecture + else "text_config does not identify a decoder architecture; decoder resource bounds require explicit input" + ) + experts = geometry["num_experts"] + _validate_architecture(raw, decoder_architecture, experts, inkling=architecture == _INKLING_ARCHITECTURE) + _validate_quantization(raw) + dtype = _dtype(raw) + heads, kv_heads = geometry["num_attention_heads"], geometry["num_key_value_heads"] + if heads and kv_heads and (kv_heads > heads or heads % kv_heads): + raise ValueError("num_key_value_heads must divide num_attention_heads and cannot exceed it") + dense_layers = raw.get("first_k_dense_replace") + if geometry["num_hidden_layers"] and dense_layers is not None and dense_layers > geometry["num_hidden_layers"]: + raise ValueError("first_k_dense_replace cannot exceed num_hidden_layers") + if decoder_architecture in _DENSE: + experts = 0 + suggestions: dict[str, Any] = {} + notes.update({key: f"config {key}={value}" for key, value in geometry.items() if value is not None}) + for key in (*_EXTRA_DIMENSIONS, *_AUXILIARY_DIMENSIONS): + if raw.get(key) is not None: + notes[key] = f"config {key}={raw[key]}" + if _uses_grouped_cache(raw, decoder_architecture): + notes["cache_layout"] = ( + "grouped cache: full-history attention, explicit sliding attention windows, and declared convolution " + "state are accounted separately; runtime KV dtype, block sizes and page padding are not model-config facts" + ) + if raw.get("use_sconv"): + notes["convolution_cache"] = ( + "Inkling convolution state uses BF16 packed K/V and output streams; " + f"pinned runtime layout: {_INKLING_CACHE_SOURCE}" + ) + names = {_text(document[key], key) for key in ("_name_or_path", "name_or_path") if document.get(key)} + if len(names) > 1: + raise ValueError("conflicting _name_or_path and name_or_path identity hints") + if names: + suggestions["model"] = names.pop() + notes["model"] = "config _name_or_path/name_or_path hint; explicit deployment identity takes precedence" + for key, value in ( + ("architecture", architecture), + ("context_length", geometry["context_length"]), + ("num_experts", experts), + ): + if value is not None: + suggestions[key] = value + notes.setdefault( + key, f"config declared {key}={value}" if key != "num_experts" else "known dense architecture" + ) + if experts is not None: + suggestions["model_kind"] = "moe" if experts else "dense" + if dtype: + notes["dtype"] = ( + f"config tensor dtype={dtype}; checkpoint storage still needs verification; " + "does not declare runtime attention or KV-cache precision" + ) + if raw.get("auto_map") or document.get("auto_map"): + notes["auto_map"] = "remote-code references present; read as metadata only and never executed" + if raw.get("quantization_config") is not None: + quant = raw["quantization_config"] + summary = { + key: quant[key] + for key in ("quant_method", "quant_algo", "weight_block_size", "kv_cache_scheme") + if key in quant + } + notes["quantization"] = ( + f"config quantization={json.dumps(summary, sort_keys=True)}; " + "tensor scales/exclusions require explicit weight bytes" + ) + for key in ("num_nextn_predict_layers", "num_mtp_modules", "mtp_transformer_layers", "use_mtp"): + if key in raw: + notes[key] = f"config {key}={raw[key]}; checkpoint capability does not enable runtime speculative decoding" + return ModelConfig(raw, decoder_architecture, hashlib.sha256(payload).hexdigest(), suggestions, notes, source_files) + + +def validate_overrides(overrides: Mapping[str, Any] | None) -> dict[str, Any]: + """Validate supplied values independently of a complete derived profile.""" + if overrides is None: + return {} + if not isinstance(overrides, Mapping): + raise ValueError("resource overrides must be a flat object") + unknown = set(overrides) - OVERRIDE_FIELDS + if unknown: + raise ValueError(f"unknown resource override fields: {', '.join(sorted(map(str, unknown)))}") + result = {} + for key, value in overrides.items(): + if key == "cache_groups": + if not isinstance(value, list) or not value: + raise ValueError("cache_groups must be a nonempty list of runtime cache group objects") + value = [FpmCacheGroup.model_validate(group).model_dump(mode="json") for group in value] + if len({group["name"] for group in value}) != len(value): + raise ValueError("cache_groups names must be unique") + elif key == "cache_block_sizes": + if not isinstance(value, dict) or not value or set(value) - _CACHE_GROUP_NAMES: + raise ValueError("cache_block_sizes must map generated cache group names to positive block sizes") + value = {name: _integer(size, f"cache_block_sizes.{name}") for name, size in value.items()} + elif key in INTEGER_FIELDS: + value = _integer( + value, + key, + minimum=0 if key in _NONNEGATIVE_FIELDS else 1, + maximum=2**53 if key in _BYTE_FIELDS else None, + ) + else: + value = _text(value, key) + if key in FIELD_CHOICES and value not in FIELD_CHOICES[key]: + raise ValueError(f"unknown {key}={value!r}; choose {', '.join(FIELD_CHOICES[key])}") + result[key] = value + if "cache_groups" in result and ("kv_bytes_per_token" in result or result.get("cache_layout") == "linear"): + raise ValueError("cache_groups cannot be combined with linear cache_layout or kv_bytes_per_token") + if sum(result.get(key, 0) for key in _BYTE_FIELDS - {"kv_bytes_per_token"}) > 2**53: + raise ValueError("total non-KV resource bytes must not exceed 2**53") + return result + + +def _precision_facts(config: ModelConfig) -> dict[str, tuple[str, str]]: + raw = config.raw + if config.source_files: + from aisimulate_core.sdk.models.helpers import _infer_quant_modes_from_raw_config + + inferred = { + key: raw[key] for key in ("quant_algo", "quant_dynamic", "kv_cache_quant_algo") if raw.get(key) is not None + } + result = {} + if inferred.get("quant_algo") in { + "fp8", + "fp8_block", + "nvfp4", + "mxfp4", + "mixed_precision", + "compressed-tensors", + }: + # Reuse collector weight classification, discarding its legacy + # attention/KV defaults: weight precision never establishes them. + modes = _infer_quant_modes_from_raw_config({**raw, "kv_cache_quant_algo": None}) + result.update( + ( + key, + ( + modes[key].name, + "checkpoint/FPM weight identity from config and ModelOpt sidecar; " + "tensor scales/exclusions require explicit weight bytes", + ), + ) + for key in ("gemm_quant_mode", "moe_quant_mode") + if key in modes + ) + kv = inferred.get("kv_cache_quant_algo") + if kv in {"fp8", "bfloat16"}: + result["kv_cache_dtype"] = ( + kv, + "explicit config/ModelOpt KV quantization metadata (independent of weight quantization)", + ) + else: + scheme = _kv_cache_scheme(raw.get("quantization_config") or {}) + if scheme.get("num_bits") == 8 and scheme.get("type") in {"float", "int"}: + result["kv_cache_dtype"] = ( + "fp8" if scheme["type"] == "float" else "int8", + "explicit config quantization_config.kv_cache_scheme (independent of weight quantization)", + ) + return result + quant = raw.get("quantization_config") + if quant is None and (raw.get("hf_quant_config") is not None or raw.get("quant_algo") is not None): + return {} # Unrecognized sidecar/processed metadata must not become unquantized defaults. + if quant is None: + if _dtype(raw) == "bfloat16": + return dict.fromkeys( + ("gemm_quant_mode", "moe_quant_mode"), + ( + "bfloat16", + "assumption: bfloat16 tensor dtype with no inline or adjacent ModelOpt quantization metadata; " + "verify deployed weight precision (checkpoint tensors are not inspected)", + ), + ) + return {} + result = {} + method = str(quant.get("quant_method", "")).lower() + algorithm = str(quant.get("quant_algo", method)).lower() + source = ( + "config checkpoint/FPM quantization label, matching sdk/models/helpers.py; " + "exclusions and tensor scales prevent inferring all tensors' storage from this label" + ) + if algorithm == "fp8" and method in {"fp8", "modelopt", ""}: + if quant.get("weight_block_size") == [128, 128]: + result = dict.fromkeys(("gemm_quant_mode", "moe_quant_mode"), ("fp8_block", source)) + elif not quant.get("weight_block_size"): + scheme = quant.get("activation_scheme") + dynamic = scheme == "dynamic" or ( + scheme is None + and any( + (group.get(key) or {}).get("dynamic") is True + for group in (quant.get("config_groups") or {}).values() + for key in ("input_activations", "weights") + ) + ) + if scheme in (None, "static", "dynamic"): + result["gemm_quant_mode"] = "fp8" if dynamic else "fp8_static", source + result["moe_quant_mode"] = "fp8", source + elif algorithm == "nvfp4" and method in {"nvfp4", "modelopt", ""}: + result = dict.fromkeys(("gemm_quant_mode", "moe_quant_mode"), ("nvfp4", source)) + scheme = _kv_cache_scheme(quant) + if scheme.get("num_bits") == 8 and scheme.get("type") in {"float", "int"}: + result["kv_cache_dtype"] = ( + "fp8" if scheme["type"] == "float" else "int8", + "explicit config quantization_config.kv_cache_scheme (independent of weight quantization)", + ) + return result + + +def _geometry(config: ModelConfig, architecture: str | None, request: SupportRequest | None) -> dict[str, int]: + dimensions = {key: value for key in _INT_ALIASES if (value := _aliased_int(config.raw, key)) is not None} + if architecture not in _FULL_ATTENTION and not _uses_grouped_cache(config.raw, architecture): + return dimensions + heads = dimensions.get("num_attention_heads") + hidden = dimensions.get("hidden_size") + if "head_dim" not in dimensions and heads and hidden: + if hidden % heads: + raise ValueError("hidden_size must divide into attention heads when head_dim is absent") + dimensions["head_dim"] = hidden // heads + if request is not None and heads: + tp = request.search.tensor_parallel + if heads % tp: + raise ProfileRequestError( + "tensor_parallel", "num_attention_heads must divide into the selected TP attention heads" + ) + kv = dimensions.get("num_key_value_heads") + if kv and ((kv >= tp and kv % tp) or (kv < tp and tp % kv)): + raise ProfileRequestError( + "tensor_parallel", "num_key_value_heads cannot be evenly sharded or replicated across the selected TP" + ) + if kv: + dimensions["local_kv_heads"] = max(1, kv // tp) + intermediate = dimensions.get("intermediate_size") + mlp_tp = request.parallelism()["moe_tensor"] if architecture == "MixtralForCausalLM" else tp + if architecture in _WEIGHT_LAYOUTS and intermediate and intermediate % mlp_tp: + raise ProfileRequestError( + "moe_tensor_parallel" if architecture == "MixtralForCausalLM" else "tensor_parallel", + "intermediate_size must be divisible by the selected tensor parallel dimension", + ) + return dimensions + + +def _derive_cache_groups( + config: ModelConfig, + architecture: str | None, + request: SupportRequest | None, + geometry: dict[str, int], + values: dict[str, Any], + sources: dict[str, str], + missing: dict[str, str], +) -> None: + raw = config.raw + if values.get("cache_layout") == "linear" or "kv_bytes_per_token" in values: + raise ValueError( + "unsupported linear cache override for declared sliding-window/convolution state; use cache_groups" + ) + if "cache_layout" not in values: + values["cache_layout"] = "grouped" + sources["cache_layout"] = "config retention metadata or explicit grouped cache resource override" + missing.pop("kv_bytes_per_token", None) + layers = geometry.get("num_hidden_layers") + sliding = _cache_layer_ids(raw, architecture, inkling=values.get("architecture") == _INKLING_ARCHITECTURE) + window = _cache_window(raw) + if not layers: + missing["cache_groups"] = "provide explicit cache_groups; config does not establish the layer count" + return + tp = request.search.tensor_parallel if request else None + scheme = _kv_cache_scheme(raw.get("quantization_config") or {}) + ordinary_kv = ( + scheme.get("dynamic") is not True + and scheme.get("strategy") in (None, "tensor") + and not any(raw.get(key) for key in ("kv_lora_rank", "q_lora_rank", "compress_ratios", "attention_k_eq_v")) + ) + dtype = values.get("kv_cache_dtype") + groups = [] + geometry_notes = [] + for name, indexes, retention in ( + ("full", tuple(index for index in range(layers) if index not in sliding), None), + ("sliding", sliding, window), + ): + if not indexes: + continue + prefix = "swa_" if name == "sliding" else "" + heads = raw.get(prefix + "num_attention_heads") or geometry.get("num_attention_heads") + kv_heads = raw.get(prefix + "num_key_value_heads") or geometry.get("num_key_value_heads") + head_dim = raw.get(prefix + "head_dim") or geometry.get("head_dim") + value_dim = raw.get(prefix + "v_head_dim") or head_dim + if values.get("architecture") == _INKLING_ARCHITECTURE and value_dim != head_dim: + raise ValueError("pinned Inkling cache layout requires equal key and value head dimensions") + local_heads = None + if heads and kv_heads and (kv_heads > heads or heads % kv_heads): + raise ValueError(f"{prefix}num_key_value_heads must divide {prefix}num_attention_heads") + if tp and heads: + if heads % tp or (kv_heads and ((kv_heads >= tp and kv_heads % tp) or (kv_heads < tp and tp % kv_heads))): + raise ProfileRequestError( + "tensor_parallel", f"{name} attention head geometry is incompatible with TP={tp}" + ) + if kv_heads: + local_heads = max(1, kv_heads // tp) + bytes_per_token = None + if local_heads and head_dim and value_dim and dtype and ordinary_kv: + element_bytes = quantization.KVCacheQuantMode[dtype].value.memory + bytes_per_token = int(len(indexes) * local_heads * (head_dim + value_dim) * element_bytes) + group = { + "name": name + "_attention", + "kind": "attention", + "num_layers": len(indexes), + "sliding_window": retention, + } + groups.append((group, bytes_per_token)) + geometry_notes.append( + { + **group, + "layer_ids": indexes, + "rank_local_kv_heads": local_heads, + "head_dim": head_dim, + "value_head_dim": value_dim, + "dtype": dtype, + } + ) + if raw.get("use_sconv"): + hidden = geometry.get("hidden_size") + if tp and kv_heads and tp > kv_heads: + raise ProfileRequestError( + "tensor_parallel", + f"Inkling convolution cache requires attention TP <= {kv_heads} for {name} layers; " + f"attention DP and expert parallelism are separate; source: {_INKLING_CACHE_SOURCE}", + ) + if hidden and kv_heads and hidden % kv_heads: + raise ValueError("Inkling convolution hidden_size must divide evenly into KV-head streams") + # Adapted from vLLM contributors' Apache-2.0 Inkling cache layout: + # https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py + # Modified: metadata-only aggregate sizing, without tensors or model construction. + packed_width = None + if hidden and kv_heads and head_dim: + packed_elements = 2 * head_dim + 2 * (hidden // kv_heads) + packed_width = 2 ** ((packed_elements - 1).bit_length()) + conv_bytes = len(indexes) * local_heads * packed_width * 2 if local_heads and packed_width else None + kernel = raw.get("sconv_kernel_size") or raw["conv_kernel_size"] + conv = { + "name": name + "_convolution", + "kind": "convolution", + "num_layers": len(indexes), + "sliding_window": kernel, + } + groups.append((conv, conv_bytes)) + geometry_notes.append( + { + **conv, + "layer_ids": indexes, + "rank_local_heads": local_heads, + "packed_head_width": packed_width, + "dtype": "bfloat16", + "source": _INKLING_CACHE_SOURCE, + } + ) + expected = {} + for group, _ in groups: + key = group["kind"], group["sliding_window"] + expected[key] = expected.get(key, 0) + group["num_layers"] + if "cache_groups" in values: + observed = {} + for group in values["cache_groups"]: + key = group["kind"], group.get("sliding_window") + observed[key] = observed.get(key, 0) + group["num_layers"] + if observed != expected: + raise ValueError( + "cache_groups must preserve the config's attention/convolution layer counts and retention windows" + ) + sources["cache_groups"] += f"; config geometry reference: {json.dumps(geometry_notes, sort_keys=True)}" + return + block_sizes = values.get("cache_block_sizes", {}) + if set(block_sizes) - {group["name"] for group, _ in groups}: + raise ValueError("cache_block_sizes contains groups not present in this model's cache layout") + unresolved = [ + group["name"] for group, _ in groups if group["kind"] == "attention" and group["name"] not in block_sizes + ] + if unresolved: + missing["cache_block_sizes"] = ( + "provide a JSON mapping of runtime block sizes after page unification for " + + ", ".join(unresolved) + + "; config cannot determine them. A complete cache_groups override can instead declare runtime pages." + ) + # A supplied partial mapping is not a complete resolution of this field. + values.pop("cache_block_sizes", None) + return + if any(size is None for _, size in groups): + missing["cache_groups"] = ( + "grouped cache sizing needs layer/head geometry, runtime KV dtype and completed topology; " + "nonstandard KV scales/packing require explicit cache_groups including their storage" + ) + return + values["cache_groups"] = [ + FpmCacheGroup.model_validate( + { + **group, + "block_size_tokens": block_sizes.get(group["name"], group["sliding_window"]), + "page_size_bytes": size * block_sizes.get(group["name"], group["sliding_window"]), + } + ).model_dump(mode="json") + for group, size in groups + ] + sources["cache_groups"] = ( + "config tensor geometry and selected attention TP; minimum packed rank-local page bytes summed over each " + "group's layers. Runtime allocation padding is not known from config; review/edit cache_groups against " + "runtime metadata. Attention DP does not divide a rank's cache. Convolution block sizes default to the " + f"pinned kernel width unless overridden. Geometry: {json.dumps(geometry_notes, sort_keys=True)}" + ) + + +def _weight_estimate( + config: ModelConfig, + request: SupportRequest | None, + values: dict[str, Any], + geometry: dict[str, int], + architecture: str | None, +) -> tuple[int | None, str]: + raw = config.raw + if any(raw.get(key) is not None for key in ("quantization_config", "hf_quant_config", "quant_algo")): + return None, "quantized/mixed tensor storage, packing and scales require explicit rank-local weights_bytes" + if architecture not in _WEIGHT_LAYOUTS: + return ( + None, + "tensor layout has no supported weight estimator; " + "provide rank-local weights_bytes including auxiliary tensors", + ) + if request is None: + return None, "complete deployment identity and topology before estimating rank-local weights_bytes" + if values.get("gemm_quant_mode") != "bfloat16" or ( + architecture == "MixtralForCausalLM" and values.get("moe_quant_mode") != "bfloat16" + ): + return None, "weight estimator supports unquantized bfloat16 deployment tensors only" + needed = ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "head_dim", + "local_kv_heads", + "vocab_size", + ) + absent = [key for key in needed if key not in geometry] + if absent: + return None, f"weight geometry missing {', '.join(absent)}; provide rank-local weights_bytes" + if any(raw.get(key) for key in (*_AUXILIARY_DIMENSIONS, "lm_head_bias")): + return None, "additional prediction/shared-expert tensors require explicit rank-local weights_bytes" + if raw.get("use_qk_norm") and architecture != "Qwen3ForCausalLM": + return None, "nonstandard Q/K normalization tensors require explicit weights_bytes" + if architecture == "Qwen2ForCausalLM" and "attention_bias" in raw: + return None, "Qwen2 attention-bias override is outside the vanilla tensor layout; provide weights_bytes" + if architecture in {"MistralForCausalLM", "MixtralForCausalLM"} and raw.get("attention_bias"): + return None, "nonstandard attention biases require explicit weights_bytes" + if architecture != "LlamaForCausalLM" and raw.get("mlp_bias"): + return None, "nonstandard MLP biases require explicit weights_bytes" + h, intermediate, layers, heads, head_dim, local_kv, vocab = (geometry[key] for key in needed) + tp = request.search.tensor_parallel + if vocab % 64 or vocab % tp: + return None, "runtime vocabulary padding is not established for this vocabulary/TP; provide weights_bytes" + q_width, kv_width = heads * head_dim // tp, local_kv * head_dim + attention = 2 * h * (q_width + kv_width) + qkv_bias = raw.get("attention_bias", architecture == "Qwen2ForCausalLM") + output_bias = raw.get("attention_bias", False) if architecture == "LlamaForCausalLM" else False + attention += (q_width + 2 * kv_width) * qkv_bias + h * output_bias + norms = 2 * h + (2 * head_dim if architecture == "Qwen3ForCausalLM" else 0) + if architecture == "MixtralForCausalLM": + experts = values.get("num_experts") + if experts is None: + return None, "num_experts is required for the MoE weight estimate" + if raw.get("mlp_bias"): + return None, "biased MoE tensor layout requires explicit weights_bytes" + parallel = request.parallelism() + mlp = 3 * h * (intermediate // parallel["moe_tensor"]) * (experts // parallel["moe_expert"]) + router = h * experts # replicated on every rank, including expert-parallel ranks + if raw.get("moe_router_dtype") not in (None, "bfloat16"): + return None, "mixed router dtype requires explicit weights_bytes" + else: + mlp = 3 * h * (intermediate // tp) + mlp += (2 * intermediate // tp + h) * raw.get("mlp_bias", False) + router = 0 + embeddings = (vocab // tp) * h * (1 if raw.get("tie_word_embeddings", False) else 2) + result = 2 * (embeddings + layers * (attention + mlp + router + norms) + h) + return result, ( + f"estimate: vanilla {architecture} bfloat16 tensors, TP={tp}, MoE TP={request.parallelism()['moe_tensor']}, " + f"EP={request.parallelism()['moe_expert']}; " + f"sharded embedding/head (tied={raw.get('tie_word_embeddings', False)}), " + "Q/K/V/output and gated MLP, configured biases, replicated RMS norms/router; " + "no quantization/padding/allocator extras" + ) + + +def _activation_estimate( + values: dict[str, Any], geometry: dict[str, int], request: SupportRequest | None, architecture: str | None +) -> tuple[int | None, str]: + if architecture not in _FULL_ATTENTION or request is None: + return ( + None, + "activation/workspace estimate unavailable for this layout or topology; " + "memory remains pending runtime observation", + ) + if values.get("fmha_quant_mode") not in {"bfloat16", "float16"}: + return ( + None, + "activation estimate requires an explicit 16-bit FMHA mode; memory remains pending runtime observation", + ) + if "num_attention_heads" not in geometry or "head_dim" not in geometry: + return None, "attention width missing; memory remains pending runtime observation" + tp = request.search.tensor_parallel + # AISimulate's local vLLM backend shares the TRT-LLM activation estimate. + # Keep this narrow subset independent of backend/model imports and do not + # extrapolate the coefficient table to other TP sizes or DSA workspaces. + coefficients = {1: 22, 2: 13, 4: 10, 8: 10} if architecture not in _DENSE else {1: 11, 2: 6.5, 4: 5, 8: 5} + if tp not in coefficients: + return None, f"no activation coefficient for TP={tp}; memory remains pending runtime observation" + tokens = values["max_num_tokens"] + width = geometry["num_attention_heads"] * geometry["head_dim"] + result = max(70 * 1024 * 1024, tokens * width * int(2 * coefficients[tp])) + return result, ( + f"estimate: sdk/backends/base_backend.py + trtllm_backend.py vLLM policy; " + f"max(70 MiB, 2 * max_num_tokens={tokens} * attention_width={width} * coefficient={coefficients[tp]}); " + f"rank-local max_batch_size={values['max_batch_size']}; no speculative decoding; not memory qualification" + ) + + +def packaged_hardware_path(system: str) -> Path: + """Locate wheel-owned metadata without importing the estimator SDK or native runtime.""" + if not isinstance(system, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", system): + raise ValueError("system must be a packaged system name, not a path") + return Path(__file__).resolve().parents[2] / "aisimulate_core" / "systems" / f"{system}.yaml" + + +def _hardware_estimates(request: SupportRequest | None) -> dict[str, tuple[int, str]]: + if request is None: + return {} + resource = packaged_hardware_path(request.identity.gpu) + if not resource.is_file(): + return {} + payload = resource.read_bytes() + spec = yaml.safe_load(payload) + if not isinstance(spec, dict) or not isinstance(spec.get("misc", {}), dict): + raise ValueError(f"invalid packaged hardware metadata for {request.identity.gpu}") + misc = spec.get("misc", {}) + source = f"estimate: packaged systems/{request.identity.gpu}.yaml sha256={hashlib.sha256(payload).hexdigest()}" + result = {} + if "other_mem" in misc: + value = _integer(misc["other_mem"], "system misc.other_mem", minimum=0, maximum=2**53) + result["runtime_overhead_bytes"] = ( + value, + source + " misc.other_mem per rank; excludes CUDA graph reservation; not a measurement", + ) + parallel = request.parallelism() + nccl = misc.get("nccl_mem", {}) + if not isinstance(nccl, dict): + raise ValueError("system misc.nccl_mem must be a TP-indexed object") + if parallel["attention_data"] == parallel["moe_expert"] == 1 and parallel["tensor"] in nccl: + value = _integer(nccl[parallel["tensor"]], "system misc.nccl_mem", minimum=0, maximum=2**53) + result["comm_overhead_bytes"] = ( + value, + source + f" misc.nccl_mem[{parallel['tensor']}] for pure TP; not a measurement", + ) + return result + + +def derive_profile( + config: ModelConfig, request: SupportRequest | None = None, overrides: Mapping[str, Any] | None = None +) -> ProfileDraft: + """Return every resolved value and actionable missing field, or a full profile. + + Invalid structural facts and overrides raise immediately. ``request=None`` + supports validating overrides and previewing config facts before collecting + deployment identity; it can never produce a complete deployment profile. + """ + supplied = validate_overrides(overrides) + values = {key: config.suggestions[key] for key in _MODEL_FIELDS if key in config.suggestions} + sources = {key: config.notes[key] for key in values} + for key, (value, source) in _precision_facts(config).items(): + values[key], sources[key] = value, source + values["comm_quant_mode"] = "half" + sources["comm_quant_mode"] = ( + "collector FPM identity default: capabilities.py resolves comm_quant_mode=half; " + "not a claim about every runtime communication tensor" + ) + for key in ("moe_backend", "attention_backend"): + values[key], sources[key] = "auto", "existing FPM deployment schema default: automatic backend selection" + for key, default in (("max_num_tokens", 8192), ("max_batch_size", 256)): + declared = getattr(request.collection, key) if request is not None else None + if declared is not None and key in supplied and declared != supplied[key]: + raise ValueError(f"collection.{key} conflicts with the resource override; provide one consistent bound") + values[key] = declared if declared is not None else default + sources[key] = ( + "user collection override; per-rank runtime bound" + if declared is not None + else f"initial vLLM collection policy: rank-local {key}={default}; independent of validation traffic" + ) + for key, value in supplied.items(): + if key in ("architecture", "num_experts") and key in values and values[key] != value: + raise ValueError(f"{key} override conflicts with the source config: {value!r} != {values[key]!r}") + if key == "context_length" and key in values and value > values[key]: + raise ValueError("context_length override exceeds the source config context limit") + values[key], sources[key] = value, "user override; caller-declared value, not an inferred measurement" + if key == "comm_quant_mode" and value != "half": + sources[key] += "; current collector supports half and will reject this identity for new collection" + architecture = config.decoder_architecture + if "architecture" not in config.suggestions: + architecture = values.get("architecture") + _validate_architecture( + config.raw, architecture, values.get("num_experts"), inkling=values.get("architecture") == _INKLING_ARCHITECTURE + ) + geometry = _geometry(config, architecture, request) + if request: + if "num_experts" in values: + experts = values["num_experts"] + if (experts > 0) != (request.identity.model_kind == "moe"): + raise ProfileRequestError("model_kind", "num_experts and request identity.model_kind disagree") + if experts and experts % request.parallelism()["moe_expert"]: + raise ProfileRequestError( + "moe_expert_parallel", "num_experts must be divisible by the selected moe_expert_parallel" + ) + if "context_length" in values and request.search.context_length > values["context_length"]: + raise ProfileRequestError( + "context_length", "request search.context_length exceeds the configured profile context_length" + ) + missing = { + "architecture": "provide the decoder architecture; config architectures/model_type did not identify it", + "context_length": "provide the supported model context length; not declared in config", + "num_experts": "provide the routed expert count (0 for dense); model kind is unknown from this config", + "gemm_quant_mode": ( + "declare the deployed GEMM quantization mode; config does not identify a supported unambiguous mode" + ), + "moe_quant_mode": "declare the deployed MoE quantization identity, also required for dense profiles", + "fmha_quant_mode": ( + "resolve attention precision from the pinned runtime or verified launch settings; " + "checkpoint weight dtype does not establish it" + ), + "kv_cache_dtype": ( + "resolve KV-cache dtype from explicit checkpoint metadata, pinned runtime or verified launch settings; " + "weight precision does not establish it" + ), + "cache_layout": ( + "confirm linear cache layout; unknown cache semantics cannot be inferred from ordinary head counts" + ), + "kv_bytes_per_token": ( + "provide rank-local bytes per cached token for all persistent cache tensors; " + "MLA/DSA and unknown layouts need explicit accounting" + ), + } + grouped_cache = ( + _uses_grouped_cache(config.raw, architecture) + or "cache_groups" in values + or values.get("cache_layout") == "grouped" + ) + if grouped_cache: + _derive_cache_groups(config, architecture, request, geometry, values, sources, missing) + elif "cache_block_sizes" in values: + raise ValueError("cache_block_sizes requires a grouped cache layout") + elif architecture in _FULL_ATTENTION: + if "cache_layout" not in values: + values["cache_layout"], sources["cache_layout"] = ( + "linear", + "known full-attention GQA/MHA decoder layout; no sliding/recurrent state", + ) + needed = ("num_hidden_layers", "local_kv_heads", "head_dim") + kv_scheme = _kv_cache_scheme(config.raw.get("quantization_config") or {}) + ordinary_kv = kv_scheme.get("dynamic") is not True and kv_scheme.get("strategy") in (None, "tensor") + if ( + "kv_bytes_per_token" not in values + and all(key in geometry for key in needed) + and "kv_cache_dtype" in values + and ordinary_kv + ): + layers, local_kv, head_dim = (geometry[key] for key in needed) + element_bytes = quantization.KVCacheQuantMode[values["kv_cache_dtype"]].value.memory + values["kv_bytes_per_token"] = int(2 * layers * local_kv * head_dim * element_bytes) + sources["kv_bytes_per_token"] = ( + f"exact linear tensor geometry: 2 K/V * {layers} layers * {local_kv} rank-local KV heads " + f"* {head_dim} head_dim * {element_bytes} bytes; KV heads replicate when TP exceeds head count; " + "attention DP does not divide a rank's cache" + ) + missing["kv_bytes_per_token"] = ( + "linear KV calculation needs layer/head geometry, explicit runtime kv_cache_dtype and completed topology; " + "otherwise provide bytes per token" + ) + if not ordinary_kv: + missing["kv_bytes_per_token"] = ( + "nonstandard KV scaling/packing needs explicit bytes per token including scales" + ) + for key, (value, source) in _hardware_estimates(request).items(): + if key not in values: + values[key], sources[key] = value, source + for key, estimate in (("weights_bytes", _weight_estimate), ("activations_bytes", _activation_estimate)): + if key in values: + continue + value, source = ( + estimate(config, request, values, geometry, architecture) + if key == "weights_bytes" + else estimate(values, geometry, request, architecture) + ) + if value is not None: + values[key], sources[key] = value, source + missing = {key: value for key, value in missing.items() if key not in values} + for key in _BYTE_FIELDS & values.keys(): + _integer(values[key], key, minimum=1 if key == "kv_bytes_per_token" else 0, maximum=2**53) + if sum(values.get(key, 0) for key in _BYTE_FIELDS - {"kv_bytes_per_token"}) > 2**53: + raise ValueError("total non-KV resource bytes must not exceed 2**53") + if missing or request is None: + return ProfileDraft(None, values, missing, sources) + declared_memory = supplied.keys() >= _NON_KV_FIELDS + provenance = json.dumps( + { + "config_sha256": config.sha256, + **({"config_source_files": config.source_files} if config.source_files else {}), + "method": ("local config facts; byte estimates are planning evidence, not runtime limits or measurements"), + "memory_source": "declared" if declared_memory else "pending", + "planning_estimates": { + key: {"value": values[key], "source": sources[key]} + for key in sorted(_NON_KV_FIELDS & values.keys()) + if key not in supplied + }, + "config_notes": config.notes, + "fields": {key: {"value": values[key], "source": sources[key]} for key in sorted(values)}, + "deployment_identity": request.identity.model_dump(mode="json"), + "parallelism": request.parallelism(), + **({"worker_type": request.worker_type} if request.worker_type is not None else {}), + }, + sort_keys=True, + ) + parallel = request.parallelism() + profile = FpmModelProfile.model_validate( + { + "schema_version": 1, + "model": request.identity.model, + "model_revision": request.identity.model_revision, + **{key: values[key] for key in _MODEL_FIELDS}, + "provenance": provenance, + "deployments": [ + { + **({"worker_type": request.worker_type} if request.worker_type is not None else {}), + "system": request.identity.gpu, + "backend": request.identity.framework, + "backend_version": request.identity.framework_version, + "tp": parallel["tensor"], + "dp": parallel["attention_data"], + "moe_tp": parallel["moe_tensor"], + "moe_ep": parallel["moe_expert"], + **{key: values[key] for key in _DEPLOYMENT_FIELDS}, + "resources": { + **{ + key: values[key] + for key in _RESOURCE_FIELDS + if key in values and (key not in _NON_KV_FIELDS or declared_memory) + }, + "provenance": provenance, + }, + } + ], + } + ) + return ProfileDraft(profile, values, {}, sources) + + +def profile_from_observations( + config: ModelConfig, + request: SupportRequest, + precision: Mapping[str, Any], + resources: Mapping[str, Any], + overrides: Mapping[str, Any] | None = None, +) -> FpmModelProfile: + """Combine validated model metadata with independently observed cache groups. + + Runtime grouping need not match config-derived planning groups. No estimated + geometry or non-KV byte declaration is substituted for the observation. + """ + supplied = validate_overrides(overrides) + metadata = {key: config.suggestions[key] for key in _MODEL_FIELDS if key in config.suggestions} + for key in _MODEL_FIELDS: + if key in supplied: + if key in metadata and ( + supplied[key] > metadata[key] if key == "context_length" else supplied[key] != metadata[key] + ): + raise ValueError(f"{key} override conflicts with the source config") + metadata[key] = supplied[key] + missing = set(_MODEL_FIELDS) - metadata.keys() + if missing: + raise ValueError("model metadata remains unresolved: " + ", ".join(sorted(missing))) + _validate_architecture( + config.raw, + config.decoder_architecture or metadata["architecture"], + metadata["num_experts"], + inkling=metadata["architecture"] == _INKLING_ARCHITECTURE, + ) + _geometry(config, config.decoder_architecture or metadata["architecture"], request) + for key in ("cache_layout", "cache_groups", "kv_bytes_per_token", "max_num_tokens", "max_batch_size"): + if key in supplied and supplied[key] != resources.get(key): + raise ValueError( + f"resource_overrides.{key} differs from observed runtime; a compatible fresh probe is required" + ) + for key in _DEPLOYMENT_FIELDS: + if key in supplied and supplied[key] != precision.get(key): + raise ValueError(f"resource_overrides.{key} differs from observed precision") + parallel = request.parallelism() + return FpmModelProfile.model_validate( + { + "schema_version": 1, + "model": request.identity.model, + "model_revision": request.identity.model_revision, + **metadata, + "provenance": json.dumps( + { + "method": "validated runtime resources with local model configuration metadata", + "config_sha256": config.sha256, + **({"config_source_files": config.source_files} if config.source_files else {}), + "config_notes": config.notes, + "user_overrides": { + key: {"value": value, "source": "user assumption"} for key, value in supplied.items() + }, + "source_profile_provenance": request.fpm_profile.provenance if request.fpm_profile else None, + }, + sort_keys=True, + ), + "deployments": [ + { + **({"worker_type": request.worker_type} if request.worker_type is not None else {}), + "system": request.identity.gpu, + "backend": request.identity.framework, + "backend_version": request.identity.framework_version, + "tp": parallel["tensor"], + "dp": parallel["attention_data"], + "moe_tp": parallel["moe_tensor"], + "moe_ep": parallel["moe_expert"], + **precision, + "resources": dict(resources), + } + ], + } + ) diff --git a/python/aisimulate/src/aisimulate/support/finalization.py b/python/aisimulate/src/aisimulate/support/finalization.py new file mode 100644 index 000000000..c1a707b61 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/finalization.py @@ -0,0 +1,600 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Finalize memory from an immutable collection into a fresh simulation plan.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any + +from aisimulate.fpm_profile import FpmResourceProfile + +from .plan import _plan_documents, check_plan, plan_lock, request_id +from .schema import FPMDeployment, SupportRequest + +_SCHEMA = "aisimulate-onboarding-finalization/v1" + + +def add_finalization_parser(actions: Any) -> None: + parser = actions.add_parser( + "finalize", + help="Resolve memory from collected runtime evidence and create fresh simulation inputs.", + description=( + "Verify a completed formal FPM collection, resolve its runtime cache allocation, and write a fresh " + "request/profile, predict/recommend configs and a verified copy of the timing data. " + "The original collection is unchanged. Review the resolved profile before accepting it in your " + "onboarding checkpoint. No GPU work is launched." + ), + ) + parser.add_argument("-c", "--config", required=True, help="Original reviewed onboarding request.") + parser.add_argument("--output-dir", required=True, help="Original completed collection plan directory.") + parser.add_argument("--resolved-output-dir", required=True, help="Fresh directory for resolved simulation inputs.") + parser.add_argument( + "--memory-config", + help="Accepted capacity-only request revision imported from this collection's formal observations.", + ) + parser.add_argument( + "--collection-report", + help=( + "Collection-validation.json to independently recheck and bind to the exported profile. " + "Omitting it records collection quality as not_assessed; finalization does not assess serving accuracy." + ), + ) + + +def _json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object in {path}") + return value + + +def _digest(content: bytes) -> str: + return hashlib.sha256(content).hexdigest() + + +def _canonical(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _inside(path: Path, parent: Path) -> Path: + if path.is_symlink() or not path.is_file() or not path.resolve().is_relative_to(parent): + raise ValueError(f"expected a regular collection artifact inside {parent}: {path}") + for directory in path.parents: + if directory == parent: + break + if directory.is_symlink(): + raise ValueError(f"refusing symlinked collection artifact: {path}") + return path + + +def finalization_manifest(request: SupportRequest) -> dict[str, Any] | None: + """Return an onboarding-owned manifest, without interpreting other provenance.""" + deployment = request.profile_deployment() + runtime = deployment.resources.runtime_memory if deployment is not None else None + if runtime is None: + return None + try: + value = json.loads(runtime.provenance) + except ValueError: + return None + if not isinstance(value, dict) or value.get("schema_version") != _SCHEMA: + return None + if "observation_provenance" in value and value["observation_provenance"] != "source_references": + raise ValueError("unsupported finalized observation provenance format") + try: + profile_provenance = json.loads(request.fpm_profile.provenance) + except ValueError: + profile_provenance = {} + profile_quality = profile_provenance.get("collection_quality") if isinstance(profile_provenance, dict) else None + quality = value.get("collection_quality") + if quality != profile_quality and (quality is not None or profile_quality != _unassessed_quality()): + raise ValueError("finalized profile collection quality differs from its memory manifest") + return value + + +def _unassessed_quality() -> dict[str, Any]: + return {"status": "not_assessed", "accuracy": "not_assessed"} + + +def _collection_quality( + request: SupportRequest, root: Path, manifest: dict[str, Any], report_path: Path | None +) -> dict[str, Any]: + """Recheck native evidence, preserving failures without making an accuracy claim.""" + if report_path is None: + return _unassessed_quality() + from .validation_workflow import _checked, _identity, check_collection_report + + reference = _identity(report_path) + saved = _json(report_path) + inputs = saved["inputs"] + # Establish the source before following a report into its own saved plan. + # A capacity-only revision does not replace this original timing request. + original = SupportRequest.from_yaml(_checked(inputs["request"])) + if original != request or Path(inputs["collection_directory"]).resolve() != root: + raise ValueError("collection quality report differs from the original request or collection directory") + report, original, plan, _ = check_collection_report(report_path) + if ( + original != request + or request_id(original) != manifest["source_request_id"] + or plan.sha256 != manifest["source_collection_plan_sha256"] + ): + raise ValueError("collection quality report differs from the finalized collection identity") + _checked(reference) + return { + "status": report["status"], + "accuracy": "not_assessed", + "report": reference, + "policy": report["policy"], + "gates": {name: gate["status"] for name, gate in report["gates"].items()}, + } + + +def verify_finalized_quality(manifest: dict[str, Any]) -> dict[str, Any] | None: + """Revalidate an exported assessment; absent historical markers stay absent.""" + if "collection_quality" not in manifest: + return None + quality = manifest["collection_quality"] + if not isinstance(quality, dict): + raise ValueError("invalid finalized collection quality") + reference = quality.get("report") + if reference is None: + expected = _unassessed_quality() + else: + from .validation_workflow import _checked + + root = Path(manifest["source_directory"]).resolve() + request = SupportRequest.from_yaml(root / "request.yaml") + if request_id(request) != manifest["source_request_id"]: + raise ValueError("finalized collection quality source request changed") + expected = _collection_quality(request, root, manifest, _checked(reference)) + if quality != expected: + raise ValueError("finalized collection quality differs from verified assessment") + return expected + + +def verify_finalized_data(request: SupportRequest, root: Path) -> None: + manifest = finalization_manifest(request) + if manifest is None: + return + artifacts = manifest.get("formal_data") + if not isinstance(artifacts, list) or len(artifacts) != 2: + raise ValueError("finalized data manifest requires the verified formal data pair") + for artifact in artifacts: + if not isinstance(artifact, dict) or not isinstance(artifact.get("relative_path"), str): + raise ValueError("finalized data manifest contains an invalid artifact") + relative = Path(artifact["relative_path"]) + if relative.is_absolute() or ".." in relative.parts or relative.parts[:2] != ("systems", "data"): + raise ValueError("finalized data manifest contains an invalid relative path") + path = _inside(root / relative, root / "systems/data") + if _digest(path.read_bytes()) != artifact.get("sha256"): + raise ValueError(f"finalized FPM data changed: {path}") + verify_finalized_quality(manifest) + + +def _merge_resources( + observations: list[dict[str, Any]], manifest: dict[str, Any], *, source_references: bool = False +) -> dict[str, Any]: + if not observations: + raise ValueError("collection has no runtime memory observations") + canonical = None + runtime_settings = None + runtime_geometry = None + minimum = None + summaries = [] + for observation in observations: + resources = FpmResourceProfile.model_validate(observation) + if resources.memory_source != "runtime" or resources.cache_layout != "grouped": + raise ValueError("finalization requires observed grouped runtime cache resources") + evidence = json.loads(resources.runtime_memory.provenance) + settings = evidence.get("runtime_settings") if isinstance(evidence, dict) else None + if not isinstance(settings, dict) or not settings: + raise ValueError("runtime memory evidence lacks resolved launch settings") + if runtime_settings is None: + runtime_settings = settings + elif settings != runtime_settings: + raise ValueError( + "collection cells have incompatible resolved launch settings, including graph configuration; " + "collect memory with the intended serving configuration" + ) + geometry = next( + ( + item["evidence"]["cache"]["groups"] + for item in evidence["artifacts"] + if item["evidence"]["kind"] == "worker" + ), + None, + ) + if not geometry: + raise ValueError("runtime memory evidence lacks physical cache group geometry") + if runtime_geometry is None: + runtime_geometry = geometry + elif geometry != runtime_geometry: + raise ValueError("collection cells have incompatible runtime cache layouts or layer groups") + comparable = resources.model_dump(mode="json") + comparable.pop("provenance") + runtime = comparable["runtime_memory"] + runtime.pop("provenance") + capacity = runtime.pop("kv_cache_bytes") + if canonical is None: + canonical = comparable + elif comparable != canonical: + raise ValueError("collection cells have incompatible runtime cache layouts or settings") + minimum = capacity if minimum is None else min(minimum, capacity) + if source_references: + # Raw records remain immutable collection artifacts. Reference their + # exact bytes instead of copying every rank/layer into each candidate. + evidence["artifacts"] = [ + {key: value for key, value in artifact.items() if key != "evidence"} + for artifact in evidence["artifacts"] + ] + summaries.append( + { + **observation, + "runtime_memory": {**observation["runtime_memory"], "provenance": _canonical(evidence)}, + } + ) + assert canonical is not None and minimum is not None + if source_references: + manifest["observation_provenance"] = "source_references" + manifest["cell_observations"] = summaries if source_references else observations + manifest["capacity_policy"] = "minimum observed request-usable rank-local capacity across compatible cells" + provenance = _canonical(manifest) + return { + **canonical, + "runtime_memory": {**canonical["runtime_memory"], "kv_cache_bytes": minimum, "provenance": provenance}, + "provenance": ( + "Runtime cache resources resolved from verified collection evidence. Raw records are retained at " + "the SHA-256-bound sources in runtime_memory.provenance." + if source_references + else provenance + ), + } + + +def _verify_collection( + request: SupportRequest, root: Path, *, memory_revision: dict[str, Any] | None = None +) -> tuple[list[dict[str, Any]], dict[str, Any], dict[Path, bytes], dict[Path, str]]: + from collector.fpm_forward.cli import _parser + from collector.fpm_forward.config import FPMCollectionOptions, collection_phases + from collector.fpm_forward.database import ( + aggregate_cell, + published_dense_synthetic_cells, + validate_formal_database_commit, + ) + from collector.fpm_forward.planner import backend_identity_columns + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + from collector.fpm_forward.runtime_memory import ( + cell_from_dict, + resolve_runtime_resources, + saved_plan_identity, + ) + + deployment = request.profile_deployment() + if deployment is None: + raise ValueError("finalization requires the original FPM profile used for collection") + paths = list((root / "fpm-checkpoint").rglob("fpm_forward.json")) + if len(paths) != 1: + raise ValueError("finalization requires exactly one formal collector checkpoint under fpm-checkpoint") + checkpoint_path = _inside(paths[0], root) + checkpoint = _json(checkpoint_path) + sha = checkpoint.get("plan_sha256") + if checkpoint.get("schema") != CHECKPOINT_SCHEMA or not isinstance(sha, str) or len(sha) != 64: + raise ValueError("invalid formal collector checkpoint identity") + collection_root = root / "fpm-artifacts" / sha[:16] + collection_path = _inside(collection_root / "collection-plan.json", root) + payload = _json(collection_path) + frozen = saved_plan_identity(payload) + if frozen.sha256 != sha: + raise ValueError("formal collector checkpoint differs from the saved collection plan") + expected = { + "model_path": request.identity.model, + "system": request.identity.gpu, + "backend": request.identity.framework, + "fpm_profile": request.fpm_profile.model_dump(mode="json"), + } + if any(payload.get(key) != value for key, value in expected.items()): + raise ValueError("saved collection plan differs from the original onboarding profile or identity") + if frozen.capability.aic_database_version != request.identity.framework_version: + raise ValueError("saved collection runtime version differs from onboarding") + scheduler = request.scheduler_limits() + from .fpm import fpm_cli_args + from .runtime import runtime_probe_manifest, verify_collection_runtime + + probe_manifest = runtime_probe_manifest(request) + memory_request = _memory_revision_request(memory_revision) if memory_revision is not None else None + if memory_request is not None and probe_manifest is None: + raise ValueError("memory revision requires the original imported runtime profile") + + # Normalize the reviewed inputs through the same local collector options + # parser. This does not resolve a model or regenerate Dynamo's runtime grid. + saved_options = dict(payload["options"]) + # Deployment is selected at collection, after the request is saved. Retain + # its verified frozen values while deriving runtime/grid limits from the + # reviewed request. Parsing these values never invokes either executor. + collection_deployment = FPMDeployment( + executor=saved_options.get("executor", "kubernetes"), + image=saved_options.get("slurm_container_image") or None, + container_mount=saved_options.get("slurm_container_mounts", []), + cpus_per_task=saved_options.get("slurm_cpus_per_task"), + cpu_bind=saved_options.get("slurm_cpu_bind"), + ) + if probe_manifest is not None: + source_index = _json(Path(probe_manifest["observations_index"])) + collection_deployment = FPMDeployment.model_validate( + source_index["configurations"][probe_manifest["configuration"]]["launch"]["deployment"] + ) + reviewed_collection = FPMCollectionOptions.from_args( + _parser().parse_args( + fpm_cli_args( + request, + output_dir=root, + checkpoint_dir=checkpoint_path.parent, + plan_only=True, + resume=True, + deployment=collection_deployment, + )[3:] + ) + ) + reviewed_options = reviewed_collection.to_dict() + for options in (reviewed_options, saved_options): + options.setdefault("gpu_memory_utilization", request.collection.memory_fraction) + if saved_options != reviewed_options: + raise ValueError("saved collection options differ from the reviewed runtime and collection settings") + expected_graph = reviewed_collection.compilation_config + if request.worker_type is not None and reviewed_collection.prefill_cudagraph_policy == "explicit": + sampling = reviewed_collection.prefill_sampling + expected_graph.update( + cudagraph_capture_sizes=list(sampling.cudagraph_capture_sizes), + max_cudagraph_capture_size=sampling.max_cudagraph_capture_size, + ) + cells = [cell_from_dict(value) for value in payload["cells"]] + if not cells or {cell.workload_kind for cell in cells} != set(collection_phases(request.worker_type)): + raise ValueError("finalization requires complete collection for the selected serving role") + if len({cell.cell_id for cell in cells}) != len(cells): + raise ValueError("saved collection has duplicate cell identities") + database = checkpoint.get("database", {}) + if ( + not isinstance(database, dict) + or database.get("status") != "passed" + or database.get("missing_cells") != [] + or database.get("skipped_first_publisher_wins") != [] + or database.get("plan_cells") != len(cells) + or database.get("published_cells") != len(cells) + ): + raise ValueError("finalization requires complete formal publication without missing or reused cells") + if any(not isinstance(database.get(name), str) for name in ("parquet", "metadata")): + raise ValueError("formal collector checkpoint is missing its published data paths") + parquet = _inside(Path(database["parquet"]), root / "systems/data") + metadata = _inside(Path(database["metadata"]), root / "systems/data") + if parquet.name != "fpm_forward_perf.parquet" or metadata != parquet.with_name("fpm_forward_perf.metadata.json"): + raise ValueError("collector checkpoint does not reference a canonical formal FPM pair") + published = validate_formal_database_commit(parquet, metadata, frozen) + dense_cells = published_dense_synthetic_cells(published) + snapshots = {path: _digest(path.read_bytes()) for path in (checkpoint_path, collection_path, parquet, metadata)} + rows, observations = [], [] + expected_topology = { + "tp": deployment.tp, + "pp": deployment.pp, + "dp": deployment.dp, + "moe_tp": deployment.moe_tp, + "moe_ep": deployment.moe_ep, + "cp": deployment.cp, + } + if not isinstance(checkpoint.get("cells"), dict): + raise ValueError("formal collector checkpoint is missing successful cell attempts") + for cell in cells: + if any(getattr(cell.topology, key) != value for key, value in expected_topology.items()): + raise ValueError("saved collection topology differs from the reviewed deployment") + for name in ("gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "comm_quant_mode", "kv_cache_dtype"): + if getattr(cell, name) != getattr(deployment, name): + raise ValueError(f"saved collection {name} differs from the reviewed deployment") + for name, value in backend_identity_columns(cell.backend_policy).items(): + if value != getattr(deployment, name): + raise ValueError(f"saved collection {name} differs from the reviewed deployment") + entry = checkpoint.get("cells", {}).get(cell.cell_id, {}) + attempt = entry.get("attempt_id") if isinstance(entry, dict) else None + if ( + not isinstance(entry, dict) + or entry.get("status") != "passed" + or not isinstance(attempt, str) + or not attempt + ): + raise ValueError(f"collection cell {cell.cell_id} has no successful attempt") + cell_dir = collection_root / "cells" / cell.cell_id + if _json(_inside(cell_dir / "cell.json", root)) != cell.to_dict(): + raise ValueError(f"saved cell differs from frozen plan: {cell.cell_id}") + for path in cell_dir.rglob("*"): + if path.is_symlink(): + raise ValueError(f"refusing symlinked collection artifact: {path}") + if path.is_file(): + snapshots[path] = _digest(path.read_bytes()) + rows.extend( + aggregate_cell( + frozen, cell, cell_dir, expected_attempt_id=attempt, dense_synthetic_kv=cell.cell_id in dense_cells + ) + ) + if probe_manifest is None: + observations.append( + resolve_runtime_resources( + cell, + cell_dir / "raw", + expected_plan_sha256=sha, + expected_attempt_id=attempt, + expected_backend_version=request.identity.framework_version, + expected_context_length=request.search.context_length, + expected_max_num_tokens=scheduler["max_batched_tokens"], + expected_max_batch_size=scheduler["max_sequences"], + expected_gpu_memory_utilization=request.collection.memory_fraction, + expected_model_revision=request.identity.model_revision, + expected_compilation_config=expected_graph, + ) + ) + runtime_compatibility = None + if probe_manifest is not None: + if not isinstance(checkpoint.get("runtime_observations"), str): + raise ValueError("formal collection is missing its runtime observation index") + index = _inside(Path(checkpoint["runtime_observations"]), root) + observed = verify_collection_runtime( + request, index, collection_checkpoint=checkpoint, memory_request=memory_request + ) + observations.append(observed["resources"]) + runtime_compatibility = observed["compatibility"] + snapshots[index] = _digest(index.read_bytes()) + evidence = observed["provenance"] + for reference in [*evidence["launch_artifacts"], *evidence["artifacts"], *evidence["runtime_artifacts"]]: + path = _inside(index.parent / reference["path"], root) + snapshots[path] = _digest(path.read_bytes()) + from collector.fpm_forward.runtime_instrumentation import load_instrumentation + + bundle_path = _inside(index.parent / evidence["instrumentation"]["manifest"], root) + bundle = load_instrumentation(bundle_path) + for path in [bundle_path, *(bundle.root / relative for relative in bundle.files)]: + snapshots[_inside(path, root)] = _digest(path.read_bytes()) + import pyarrow.parquet as pq + + actual = pq.read_table(parquet).to_pylist() + if sorted(map(_canonical, actual)) != sorted(map(_canonical, rows)): + raise ValueError("formal FPM rows differ from the verified native collection or include unrelated cells") + files = {path.relative_to(root): path.read_bytes() for path in (parquet, metadata)} + manifest = { + "schema_version": _SCHEMA, + "source_request_id": request_id(request), + "source_collection_plan_sha256": sha, + "source_directory": str(root), + "review_status": "requires_review", + "modeling_scope": ( + "Text-decoder timing; observed cache capacity accounts for every component loaded by the benchmark " + "worker. No guessed multimodal encoder allocation is subtracted." + ), + "source_artifacts": [{"path": str(path), "sha256": digest} for path, digest in sorted(snapshots.items())], + "formal_data": [ + {"relative_path": str(path), "sha256": _digest(content)} for path, content in sorted(files.items()) + ], + } + if runtime_compatibility is not None: + manifest["runtime_probe"] = probe_manifest + manifest["runtime_compatibility"] = runtime_compatibility + if memory_revision is not None: + manifest["memory_revision"] = memory_revision + return observations, manifest, files, snapshots + + +def _memory_revision_request(reference: dict[str, Any]) -> SupportRequest: + if not isinstance(reference, dict) or not isinstance(reference.get("path"), str): + raise ValueError("invalid memory revision reference") + path = Path(reference["path"]) + if path.is_symlink() or not path.is_file() or _digest(path.read_bytes()) != reference.get("sha256"): + raise ValueError(f"memory revision request changed: {path}") + request = SupportRequest.from_yaml(path) + if request_id(request) != reference.get("request_id"): + raise ValueError("memory revision request identity changed") + return request + + +def finalize( + request: SupportRequest, + output_dir: str | Path, + resolved_output_dir: str | Path, + *, + memory_config: str | Path | None = None, + collection_report: str | Path | None = None, +) -> dict[str, Any]: + root = Path(output_dir).expanduser().resolve() + raw_target = Path(resolved_output_dir).expanduser() + if raw_target.is_symlink(): + raise ValueError("resolved output must not be a symlink") + target = raw_target.resolve() + if target == root or target.is_relative_to(root) or root.is_relative_to(target): + raise ValueError("resolved output must be separate from the original collection directory") + if target.exists(): + raise ValueError("resolved output directory must be new; existing artifacts cannot be replaced") + with plan_lock(root): + check_plan(request, root) + memory_revision = None + if memory_config is not None: + from .runtime import verify_runtime_acceptance + + path = Path(memory_config).expanduser().absolute() + revised = SupportRequest.from_yaml(path) + memory_revision = { + "path": str(path), + "sha256": _digest(path.read_bytes()), + "request_id": request_id(revised), + } + _memory_revision_request(memory_revision) + verify_runtime_acceptance(revised) + observations, manifest, formal_files, snapshots = _verify_collection( + request, root, memory_revision=memory_revision + ) + manifest["collection_quality"] = _collection_quality( + request, + root, + manifest, + Path(collection_report).expanduser().resolve() if collection_report is not None else None, + ) + resources = _merge_resources(observations, manifest, source_references=True) + payload = request.model_dump(mode="json") + payload["fpm_profile"]["provenance"] = _canonical( + { + "memory_source": "runtime", + "method": "runtime memory resolved from verified collector artifacts", + "source_collection_plan_sha256": manifest["source_collection_plan_sha256"], + "source_profile_provenance": request.fpm_profile.provenance, + "collection_quality": manifest["collection_quality"], + } + ) + selected = request.profile_deployment() + for index, deployment in enumerate(request.fpm_profile.deployments): + if deployment == selected: + payload["fpm_profile"]["deployments"][index]["resources"] = resources + resolved = SupportRequest.model_validate(payload) + plan, documents = _plan_documents(resolved, target) + documents.update(formal_files) + target.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{target.name}-", dir=target.parent)) + try: + for relative, content in documents.items(): + destination = staging / relative + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(content) + check_plan(request, root) + for path, digest in snapshots.items(): + if _digest(_inside(path, root).read_bytes()) != digest: + raise ValueError(f"collection artifact changed during finalization: {path}") + if memory_revision is not None: + _memory_revision_request(memory_revision) + verify_finalized_quality(manifest) + if target.exists(): + raise ValueError("resolved output appeared during finalization; refusing to replace it") + os.rename(staging, target) + finally: + if staging.exists(): + shutil.rmtree(staging) + return { + "resolved_directory": str(target), + "request": plan["outputs"]["request"], + "review_status": "requires_review", + "collection_quality": manifest["collection_quality"]["status"], + "accuracy": "not_assessed", + } + + +def run_finalization(args: argparse.Namespace) -> int: + result = finalize( + SupportRequest.from_yaml(args.config), + args.output_dir, + args.resolved_output_dir, + memory_config=args.memory_config, + collection_report=args.collection_report, + ) + print(json.dumps(result, indent=2, sort_keys=True)) + print("Runtime memory resolved. Review the new profile before accepting it in the onboarding checkpoint.") + return 0 diff --git a/python/aisimulate/src/aisimulate/support/fpm.py b/python/aisimulate/src/aisimulate/support/fpm.py index b5f642a2d..62a58c531 100644 --- a/python/aisimulate/src/aisimulate/support/fpm.py +++ b/python/aisimulate/src/aisimulate/support/fpm.py @@ -10,6 +10,8 @@ import sys from pathlib import Path +from aisimulate.fpm_advisory import prefill_graph_advisory + from .schema import FPMDeployment, SupportRequest @@ -30,17 +32,27 @@ def fpm_cli_args( checkpoint = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" if checkpoint != root / "fpm-checkpoint" and root / "fpm-checkpoint" not in checkpoint.parents: raise ValueError("checkpoint_dir must stay within the plan's fpm-checkpoint directory") - profile = request.profile_deployment() - max_prefill_tokens = max(2, request.workload.input_tokens * request.workload.concurrency) - max_prefill_batch = request.workload.concurrency - if profile is not None: - scheduler = request.scheduler_limits() - max_prefill_batch = scheduler["max_sequences"] - max_prefill_tokens = min( - max(2, request.workload.input_tokens * max_prefill_batch), scheduler["max_batched_tokens"] + if deployment is not None and deployment.executor == "slurm": + from collector.fpm_forward.config import resolve_slurm_cpu_policy + + cpus, binding = resolve_slurm_cpu_policy( + deployment.cpus_per_task, + deployment.cpu_bind, + resume=resume, + checkpoint_dir=checkpoint, + artifact_root=root / "fpm-artifacts", + smoke=smoke, ) - if max_prefill_tokens < 2: - raise ValueError("FPM collection requires a rank-local token limit of at least 2 in the resource profile") + deployment = deployment.model_copy(update={"cpus_per_task": cpus, "cpu_bind": binding}) + profile = request.profile_deployment() + from .runtime import runtime_collection_inputs + + runtime_arguments, deployment = runtime_collection_inputs(request, deployment) + scheduler = request.scheduler_limits() + max_prefill_tokens = scheduler["max_batched_tokens"] + max_prefill_batch = scheduler["max_sequences"] + if max_prefill_tokens < 2: + raise ValueError("FPM collection requires a rank-local token limit of at least 2 in the resource profile") command = [ "python3", "-m", @@ -57,6 +69,12 @@ def fpm_cli_args( str(request.worker_gpus), "--fpm-parallel-presets", request.parallel_preset, + "--fpm-max-model-len", + str(request.search.context_length), + "--fpm-max-num-batched-tokens", + str(max_prefill_tokens), + "--fpm-max-num-seqs", + str(max_prefill_batch), "--fpm-max-prefill-isl", str(max_prefill_tokens), "--fpm-max-prefill-batch-size", @@ -68,6 +86,21 @@ def fpm_cli_args( "--fpm-database-root", str(root / "systems/data"), ] + command.extend(runtime_arguments) + if request.worker_type is not None: + command.extend(("--fpm-worker-type", request.worker_type)) + if request.collection.cudagraph_mode is not None: + command.extend(("--fpm-cudagraph-mode", request.collection.cudagraph_mode)) + if request.collection.cudagraph_capture_sizes is not None: + command.extend(("--fpm-cudagraph-capture-sizes", *map(str, request.collection.cudagraph_capture_sizes))) + if request.collection.max_cudagraph_capture_size is not None: + command.extend(("--fpm-max-cudagraph-capture-size", str(request.collection.max_cudagraph_capture_size))) + if "prefill_cudagraph_policy" in request.collection.model_fields_set: + command.extend(("--fpm-prefill-cudagraph-policy", request.collection.prefill_cudagraph_policy)) + if request.collection.max_prefill_cudagraph_size is not None: + command.extend(("--fpm-max-prefill-cudagraph-size", str(request.collection.max_prefill_cudagraph_size))) + if request.collection.gpu_memory_utilization is not None: + command.extend(("--fpm-gpu-memory-utilization", str(request.collection.gpu_memory_utilization))) if profile is not None: command.extend( ( @@ -88,11 +121,19 @@ def fpm_cli_args( if request.identity.sm is not None: command.extend(("--sm", str(request.identity.sm))) if deployment is not None: - for name, value in deployment.model_dump(exclude_none=True).items(): + if deployment.executor == "slurm": + command.extend(("--fpm-executor", "slurm")) + for name, value in deployment.model_dump(exclude_none=True, exclude={"executor", "container_mount"}).items(): if name == "image": - command.extend(("--generator-set", f"K8sConfig.k8s_image={json.dumps(value)}")) + if deployment.executor == "slurm": + command.extend(("--fpm-slurm-container-image", value)) + else: + command.extend(("--generator-set", f"K8sConfig.k8s_image={json.dumps(value)}")) else: - command.extend(("--" + name.replace("_", "-"), value)) + prefix = "--fpm-slurm-" if name in {"cpus_per_task", "cpu_bind"} else "--" + command.extend((prefix + name.replace("_", "-"), str(value))) + for mount in deployment.container_mount: + command.extend(("--fpm-slurm-container-mount", mount)) if plan_only: command.append("--plan-only") if smoke: @@ -143,11 +184,46 @@ def _check_campaign_outputs(root: Path, *, smoke: bool, resume: bool, checkpoint # The collector remains responsible for the checkpoint schema and frozen-plan identity. +def _resolve_execution(command: list[str]): + """Use the collector's public entry boundary, resolving the plan once.""" + from collector.fpm_forward.cli import _INPUT_ERRORS, _parser + from collector.fpm_forward.entry import resolve_run_inputs + from collector.model_cases import build_collection_case_plan + + parser = _parser() + args = parser.parse_args(command[3:]) + try: + case_plan = build_collection_case_plan( + backend=args.backend, + model_path=args.model_path, + model_architecture=args.model_architecture, + gpu_type=args.gpu, + sm_version=args.sm, + model_cases_path=args.model_cases, + ) + resolved = resolve_run_inputs(args, case_plan) + except _INPUT_ERRORS as error: + parser.error(str(error)) + if args.smoke and args.limit is None: + # Native smoke defaults to one cell. Onboarding must exercise decode + # and every selected backend/precision cell, regardless of ordering. + args.limit = len(resolved[0].cells) + return args, resolved + + +def _checkpoint_root(root: Path, checkpoint_dir: str | Path | None) -> Path: + selected = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" + if selected != root / "fpm-checkpoint" and root / "fpm-checkpoint" not in selected.parents: + raise ValueError("checkpoint_dir must stay within the plan's fpm-checkpoint directory") + return selected + + def run_fpm( request: SupportRequest, *, output_dir: str | Path, execute: bool = False, + check_readiness: bool = False, smoke: bool = False, limit: int | None = None, resume: bool = False, @@ -159,6 +235,16 @@ def run_fpm( from .plan import check_plan, plan_lock root = Path(output_dir).expanduser().resolve() + selected_checkpoint = _checkpoint_root(root, checkpoint_dir) + if check_readiness: + if execute or smoke or limit is not None or resume: + raise ValueError("--check-readiness is read-only and cannot use --execute, --smoke, --limit or --resume") + check_plan(request, root) + from .collection_readiness import assess_readiness + + report = assess_readiness(request, root, selected_checkpoint) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["ready_for_full_collection"] else 1 command = fpm_cli_args( request, output_dir=root, @@ -169,24 +255,62 @@ def run_fpm( checkpoint_dir=checkpoint_dir, deployment=deployment, ) + advisory = ( + prefill_graph_advisory( + request.collection.prefill_cudagraph_policy, enforce_eager="--fpm-enforce-eager" in command + ) + if request.worker_type is None + else None + ) + if advisory: + print(advisory, file=sys.stderr) if not execute: if resume or (root.exists() and any(root.iterdir())): check_plan(request, root) print(shlex.join(command)) return 0 check_plan(request, root) + from .runtime import runtime_probe_manifest, verify_collection_runtime, verify_runtime_acceptance + + verify_runtime_acceptance(request) with plan_lock(root): check_plan(request, root) _check_campaign_outputs(root, smoke=smoke, resume=resume, checkpoint_dir=checkpoint_dir) - # The collector reports input/plan failures through argparse before - # entering run_resolved; imports and execution failures return exit 1. + from .collection_readiness import REPORT_FILENAME, assess_readiness, resume_without_workers + + frozen = None + collector_status = None + execution_error = None + recovery_only = False + collector_available = False try: - from collector.fpm_forward.cli import main as fpm_main + from collector.fpm_forward.entry import run_resolved + from collector.fpm_forward.runner import _atomic_json - result = fpm_main(command[3:]) - if result == 0 and not smoke: - selected = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" - checkpoint = json.loads((selected / "fpm_forward.json").read_text(encoding="utf-8")) + collector_available = True + args, resolved = _resolve_execution(command) + frozen = resolved[0] + if not smoke: + before = assess_readiness(request, root, selected_checkpoint, expected_plan=frozen) + recovery_only = resume and resume_without_workers(frozen, selected_checkpoint / "fpm_forward.json") + if not before["ready_for_full_collection"] and not recovery_only: + raise ValueError( + "full collection requires usable prefill and decode readiness for this exact runtime/launch; " + "inspect --check-readiness and run --smoke --execute first" + ) + errors = run_resolved(args, resolved) + status = collector_status = 1 if errors else 0 + if errors: + print(json.dumps(errors, indent=2, sort_keys=True), file=sys.stderr) + if status == 0 and not smoke and runtime_probe_manifest(request) is not None: + payload = json.loads((selected_checkpoint / "fpm_forward.json").read_text(encoding="utf-8")) + index = Path(payload["runtime_observations"]) + if not index.resolve().is_relative_to(root): + raise ValueError("formal runtime observation index must stay inside the collection directory") + observed = verify_collection_runtime(request, index, collection_checkpoint=payload) + (root / "runtime-compatibility.json").write_text(json.dumps(observed, indent=2, sort_keys=True) + "\n") + if status == 0 and not smoke: + checkpoint = json.loads((selected_checkpoint / "fpm_forward.json").read_text(encoding="utf-8")) publication = checkpoint.get("database") if not isinstance(publication, dict) or publication.get("status") != "passed": raise RuntimeError("collector did not record a completed FPM database publication") @@ -203,7 +327,18 @@ def run_fpm( "Collected artifacts are preserved. Create a new plan in a new output directory " "using a matching runtime or the observed framework_version." ) - return result except Exception as exc: print(f"aisimulate onboard collect-fpm failed: {exc}", file=sys.stderr) - return 1 + execution_error = str(exc) + status = 1 + if not collector_available: + # Readiness reporting also imports the collector; do not retry an + # unavailable dependency while handling its initialization failure. + return status + report = assess_readiness(request, root, selected_checkpoint, expected_plan=frozen) + report.update(collector_exit_status=collector_status, recovery_only=recovery_only) + if execution_error is not None: + report["execution_error"] = execution_error + _atomic_json(root / REPORT_FILENAME, report) + print(json.dumps(report, indent=2, sort_keys=True)) + return status or (0 if report["ready_for_full_collection"] else 1) diff --git a/python/aisimulate/src/aisimulate/support/interpolation_validation.py b/python/aisimulate/src/aisimulate/support/interpolation_validation.py new file mode 100644 index 000000000..ce29e248e --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/interpolation_validation.py @@ -0,0 +1,577 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Evaluate direct FPM interpolation against withheld measured coordinates.""" + +from __future__ import annotations + +import hashlib +import itertools +import json +import math +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel +from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY + +from .schema import SupportRequest + +_PHASES = ("prefill", "decode") +_COORDINATES = ("batch_size", "total_prefill_tokens", "total_kv_read_tokens") + + +def _json_bytes(value: Any) -> bytes: + return (json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n").encode() + + +def _identity(path: Path) -> dict[str, Any]: + content = path.read_bytes() + return {"path": str(path), "sha256": hashlib.sha256(content).hexdigest(), "size_bytes": len(content)} + + +def _config(request: SupportRequest, root: Path) -> ForwardPassPerfModelConfig: + deployment = request.profile_deployment() + if deployment is None: + raise ValueError("interpolation validation requires an explicit FPM model profile") + return ForwardPassPerfModelConfig( + model=request.identity.model, + system=request.identity.gpu, + backend=request.identity.framework, + backend_version=request.identity.framework_version, + worker_type=request.worker_type or "aggregated", + tp=deployment.tp, + pp=deployment.pp, + attention_dp=deployment.dp, + moe_tp_size=deployment.moe_tp, + moe_ep_size=deployment.moe_ep, + fpm_profile=request.fpm_profile.model_dump(mode="json"), + systems_paths=(str(root),), + estimation_mode="fpm_interpolation", + fallback_policy="deny", + estimator_config={"fpm_interpolation": {"method": "direct", "collect_coverage": True}}, + ) + + +def _matches(row: dict[str, Any], request: SupportRequest) -> bool: + deployment = request.profile_deployment().model_dump(mode="json", exclude={"resources", "worker_type"}) + return ( + row["model_path"] == request.identity.model + and row["workload_kind"] in request.collection_phases + and all(row.get(name) == value for name, value in deployment.items()) + and all( + row.get(name, default) == default + for name, default in zip(EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, strict=True) + ) + ) + + +def _coordinate(row: dict[str, Any]) -> tuple[str, int, int, int]: + return (row["workload_kind"], *(row[name] for name in _COORDINATES)) + + +def _point(coordinate: tuple[str, int, int, int]) -> dict[str, Any]: + return {"phase": coordinate[0], **dict(zip(_COORDINATES, coordinate[1:], strict=True))} + + +def _execution_boundaries( + rows: list[dict[str, Any]], evidence: list[dict[str, Any]] | None +) -> tuple[set[tuple[str, int, int, int]], dict[str, Any]]: + """Retain adjacent measured mode transitions, without predicting dispatch or timing.""" + report: dict[str, Any] = { + "status": "not_assessed", + "basis": "native_expected_point_modes_and_resolved_graph_configuration", + "per_point_observed_dispatch": "unreported", + "cells": [], + "transitions": [], + "retained_anchors": [], + "additional_to_envelope": [], + "issues": [], + } + if evidence is None: + report["issues"].append("formal table coordinates alone do not establish CUDA graph boundaries") + return set(), report + if {cell["cell_id"] for cell in evidence} != {row["cell_id"] for row in rows} or len(evidence) != len( + {cell["cell_id"] for cell in evidence} + ): + raise ValueError("holdout execution evidence must match the formal deployment cells exactly") + modes = {} + for cell in sorted(evidence, key=lambda item: item["cell_id"]): + cell_id = cell["cell_id"] + cell_rows = [row for row in rows if row["cell_id"] == cell_id] + identity = cell["identity"] + if set(identity) != { + "source_plan_sha256", + "collector_attempt_id", + "runtime_run_id", + "runtime_grid_digest", + } or any(any(not value or row.get(name) != value for name, value in identity.items()) for row in cell_rows): + raise ValueError(f"holdout execution evidence has a different source identity: {cell_id}") + points = {_coordinate(point): point for point in cell["points"]} + if len(points) != len(cell["points"]) or set(points) != {_coordinate(row) for row in cell_rows}: + raise ValueError(f"holdout execution evidence coordinates differ from the formal cell: {cell_id}") + graphs = sorted(cell["native_graph_config"], key=lambda graph: graph["dp_rank"]) + ranks = [graph["dp_rank"] for graph in graphs] + if len(ranks) != len(set(ranks)) or set(ranks) != set(range(cell_rows[0]["dp"])): + raise ValueError(f"holdout graph evidence has inconsistent DP ranks: {cell_id}") + for graph in graphs: + if _identity(Path(graph["source"]["path"])) != graph["source"]: + raise ValueError(f"holdout native graph source changed: {cell_id}") + configs = [graph["config"] for graph in graphs] + present = [config for config in configs if isinstance(config, dict)] + if any(config != present[0] for config in present): + raise ValueError(f"holdout native graph configurations disagree across ranks: {cell_id}") + config = present[0] if len(present) == len(configs) else {} + phase = cell_rows[0]["workload_kind"] + captures = config.get(f"{phase}_capture_sizes", config.get("capture_sizes")) + configured_mode = config.get(f"{phase}_mode", config.get("mode")) + configured_modes = { + "NONE": {"NONE"}, + "FULL": {"NONE", "FULL"}, + "PIECEWISE": {"NONE", "PIECEWISE"}, + "FULL_AND_PIECEWISE": {"NONE", "PIECEWISE" if phase == "prefill" else "FULL"}, + "FULL_DECODE_ONLY": {"NONE"} if phase == "prefill" else {"NONE", "FULL"}, + }.get(configured_mode if isinstance(configured_mode, str) else None) + if ( + configured_modes is None + or not isinstance(captures, list) + or any(type(size) is not int or size < 1 for size in captures) + ): + report["issues"].append(f"{cell_id}: resolved native graph configuration is missing or unknown") + configured_modes = None + maximum = config.get("max_capture_size") + if ( + configured_modes is not None + and maximum is not None + and (type(maximum) is not int or maximum < max(captures, default=0)) + ): + raise ValueError(f"holdout native capture sizes exceed the reported maximum: {cell_id}") + known = 0 + for row in cell_rows: + if row.get("kv_seed_regime") == "fake_fallback": + continue + coordinate = _coordinate(row) + point = points[coordinate] + mode, capture = point.get("expected_cudagraph_mode"), point.get("expected_capture_size") + if configured_modes is None or not isinstance(mode, str) or mode not in {"NONE", "FULL", "PIECEWISE"}: + continue + if mode != "NONE" and capture is None: + continue + scheduled = row["total_prefill_tokens"] if phase == "prefill" else row["batch_size"] + if mode not in configured_modes or ( + capture is not None + if mode == "NONE" + else type(capture) is not int or capture < scheduled or capture not in captures + ): + raise ValueError(f"holdout point mode/capture contradicts native graph configuration: {cell_id}") + modes[coordinate] = mode + known += 1 + eligible = sum(row.get("kv_seed_regime") != "fake_fallback" for row in cell_rows) + if known < eligible: + report["issues"].append(f"{cell_id}: {eligible - known} eligible points lack usable expected-mode evidence") + report["cells"].append( + { + "cell_id": cell_id, + "identity": identity, + "native_graph_config": graphs, + "eligible_point_count": eligible, + "known_mode_point_count": known, + } + ) + anchors = set() + eligible_rows = [row for row in rows if row.get("kv_seed_regime") != "fake_fallback"] + for axis in _COORDINATES: + fixed = [name for name in _COORDINATES if name != axis] + curves: dict[tuple, list[dict[str, Any]]] = {} + for row in eligible_rows: + curves.setdefault((row["workload_kind"], *(row[name] for name in fixed)), []).append(row) + for _key, curve in sorted(curves.items()): + ordered = sorted(curve, key=lambda row: row[axis]) + for lower, upper in itertools.pairwise(ordered): + lo, hi = _coordinate(lower), _coordinate(upper) + if lo in modes and hi in modes and modes[lo] != modes[hi]: + anchors.update((lo, hi)) + report["transitions"].append( + { + "axis": axis, + "lower": {**_point(lo), "mode": modes[lo]}, + "upper": {**_point(hi), "mode": modes[hi]}, + } + ) + envelope = set().union( + *(_anchors([row for row in eligible_rows if row["workload_kind"] == phase]) for phase in _PHASES) + ) + report.update( + status="incomplete" if report["issues"] else "assessed", + retained_anchors=[_point(coordinate) for coordinate in sorted(anchors)], + additional_to_envelope=[_point(coordinate) for coordinate in sorted(anchors - envelope)], + ) + return anchors, report + + +def _anchors(rows: list[dict[str, Any]]) -> set[tuple[str, int, int, int]]: + """Retain the envelope at every batch, including both outer KV curves.""" + anchors = set() + for batch in sorted({row["batch_size"] for row in rows}): + group = [row for row in rows if row["batch_size"] == batch] + if group[0]["workload_kind"] == "decode": + anchors.update(_coordinate(fn(group, key=_coordinate)) for fn in (min, max)) + continue + kv_values = {row["total_kv_read_tokens"] for row in group} + for kv in (min(kv_values), max(kv_values)): + curve = [row for row in group if row["total_kv_read_tokens"] == kv] + anchors.update(_coordinate(fn(curve, key=_coordinate)) for fn in (min, max)) + # A ragged interior KV curve may extend beyond both outer curves. + # Preserve those token-envelope extrema as well. + anchors.update(_coordinate(fn(group, key=_coordinate)) for fn in (min, max)) + return anchors + + +def _select( + rows: list[dict[str, Any]], limit: int, seed: int, execution_anchors: set[tuple[str, int, int, int]] +) -> list[dict[str, Any]]: + """Target 20% up to the cap, retaining anchors and spreading over log-scaled shapes.""" + if len(rows) < 2: + return [] + anchors = _anchors(rows) | execution_anchors + ordered = sorted((row for row in rows if _coordinate(row) not in anchors), key=_coordinate) + if not ordered: + return [] + count = min(limit, max(1, len(rows) // 5), len(ordered)) + values = [[math.log2(1 + row[name]) for name in _COORDINATES] for row in ordered] + bounds = [(min(axis), max(axis)) for axis in zip(*values, strict=True)] + positions = [ + [(value - lo) / (hi - lo) if hi > lo else 0.0 for value, (lo, hi) in zip(vector, bounds, strict=True)] + for vector in values + ] + + def tie_break(index: int) -> str: + return hashlib.sha256(_json_bytes([seed, _coordinate(ordered[index])])).hexdigest() + + selected = [min(range(len(ordered)), key=tie_break)] + remaining = set(range(len(ordered))) - set(selected) + while len(selected) < count: + index = max( + remaining, + key=lambda candidate: ( + min( + sum((a - b) ** 2 for a, b in zip(positions[candidate], positions[chosen], strict=True)) + for chosen in selected + ), + tie_break(candidate), + ), + ) + selected.append(index) + remaining.remove(index) + return [ordered[index] for index in selected] + + +def _boundary_axes(row: dict[str, Any], rows: list[dict[str, Any]]) -> list[str]: + axes = [] + for name in _COORDINATES: + values = [point[name] for point in rows] + lo, hi = min(values), max(values) + if lo != hi: + if row[name] == lo: + axes.append(f"min_{name}") + if row[name] == hi: + axes.append(f"max_{name}") + return axes + + +def _metrics(row: dict[str, Any]) -> dict[str, Any]: + if row["workload_kind"] == "prefill": + scheduled = { + "num_prefill_requests": row["batch_size"], + "sum_prefill_tokens": row["total_prefill_tokens"], + "sum_prefill_kv_tokens": row["total_kv_read_tokens"], + } + else: + scheduled = {"num_decode_requests": row["batch_size"], "sum_decode_kv_tokens": row["total_kv_read_tokens"]} + return {"scheduled_requests": scheduled} + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + index = (len(ordered) - 1) * percentile + lower = math.floor(index) + upper = math.ceil(index) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (index - lower) + + +def _phase_report( + rows: list[dict[str, Any]], *, max_p95_relative_error: float, max_unsupported_fraction: float +) -> dict[str, Any]: + errors = [row["relative_error"] for row in rows if row["status"] == "predicted"] + unsupported = sum(row["status"] == "unsupported" for row in rows) + fraction = unsupported / len(rows) if rows else None + absolute = [abs(error) for error in errors] + distribution = ( + { + "p50": _percentile(absolute, 0.5), + "p95": _percentile(absolute, 0.95), + "max": max(absolute), + "bias": sum(errors) / len(errors), + } + if errors + else None + ) + issues = [] + if not rows: + issues.append("no eligible coordinates remain after retaining measured envelope and execution-boundary anchors") + elif not errors: + issues.append("no withheld coordinate has a supported interpolation prediction") + if fraction is not None and fraction > max_unsupported_fraction: + issues.append("unsupported query fraction exceeds the configured threshold") + if distribution is not None and distribution["p95"] > max_p95_relative_error: + issues.append("p95 absolute relative error exceeds the configured threshold") + return { + "status": "passed" if not issues else ("failed" if rows else "not_assessed"), + "selected_count": len(rows), + "predicted_count": len(errors), + "unsupported_count": unsupported, + "unsupported_fraction": fraction, + "boundary_unsupported_count": sum( + row["status"] == "unsupported" and bool(row["boundary_axes"]) for row in rows + ), + "absolute_relative_error": distribution, + "bias_definition": "mean((predicted_ms - measured_ms) / measured_ms)", + "issues": issues, + } + + +def evaluate_interpolation_holdout( + request: SupportRequest, + *, + systems_root: str | Path, + output_dir: str | Path, + max_points_per_phase: int = 16, + seed: int = 42, + max_p95_relative_error: float = 0.20, + max_unsupported_fraction: float = 0.0, + execution_evidence: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Write a reproducible native direct-interpolation assessment to a fresh directory. + + Forward timing permits pending memory. This tests interpolation, not request + admission, measurement repeatability or serving accuracy. The measured + envelope and evidenced execution-mode boundaries are retained; any missing + interior brackets remain unsupported. Changing policy or inputs requires + another directory; source measurements are never edited. + """ + import pyarrow as pa + import pyarrow.parquet as pq + + if type(max_points_per_phase) is not int or max_points_per_phase < 1: + raise ValueError("max_points_per_phase must be a positive integer") + if type(seed) is not int or seed < 0: + raise ValueError("seed must be a non-negative integer") + for name, value in ( + ("max_p95_relative_error", max_p95_relative_error), + ("max_unsupported_fraction", max_unsupported_fraction), + ): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be a finite non-negative number") + if max_unsupported_fraction > 1: + raise ValueError("max_unsupported_fraction must be at most 1") + source = Path(systems_root).expanduser().resolve() + output = Path(output_dir).expanduser().absolute() + if output.is_symlink(): + raise ValueError("holdout output must not be a symbolic link") + output = output.resolve() + if output == source or output in source.parents or source in output.parents: + raise ValueError("holdout output must be separate from the source systems directory") + if output.exists() and (not output.is_dir() or any(output.iterdir())): + raise ValueError("holdout output must be a fresh or empty directory") + deployment = request.profile_deployment() + if deployment is None: + raise ValueError("interpolation validation requires an explicit FPM model profile") + relative = Path("data") / deployment.system / deployment.backend / deployment.backend_version + parquet_relative = relative / "fpm_forward_perf.parquet" + metadata_relative = parquet_relative.with_suffix(".metadata.json") + system_relative = Path(f"{deployment.system}.yaml") + source_paths = [source / relative for relative in (system_relative, parquet_relative, metadata_relative)] + if any(not path.is_file() or not path.resolve().is_relative_to(source) for path in source_paths): + raise ValueError("source systems directory must contain a complete local system YAML and FPM data pair") + source_bytes = [path.read_bytes() for path in source_paths] + inputs = { + name: {"path": str(path), "sha256": hashlib.sha256(content).hexdigest(), "size_bytes": len(content)} + for name, path, content in zip(("system", "parquet", "metadata"), source_paths, source_bytes, strict=True) + } + system = yaml.safe_load(source_bytes[0]) + if not isinstance(system, dict) or system.get("data_dir") != f"data/{deployment.system}": + raise ValueError("holdout validation requires a system YAML with local data_dir=data/") + table = pq.read_table(pa.BufferReader(source_bytes[1])) + metadata = json.loads(source_bytes[2]) + output.mkdir(parents=True, exist_ok=True) + snapshot = output / "source-systems" + for relative_path, content in zip( + (system_relative, parquet_relative, metadata_relative), source_bytes, strict=True + ): + destination = snapshot / relative_path + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(content) + # Native construction checks the full source pair, exact identity, schema, + # digest, duplicate physical keys and coordinate collisions before selection. + baseline = RustForwardPassPerfModel.best_available(_config(request, snapshot)) + baseline.close() + all_rows = table.to_pylist() + matching = [row for row in all_rows if _matches(row, request)] + if not matching: + raise ValueError("FPM data contains no exact profile deployment cell") + eligible = [row for row in matching if row.get("kv_seed_regime") != "fake_fallback"] + execution_anchors, boundaries = _execution_boundaries(matching, execution_evidence) + selected = [] + for phase in request.collection_phases: + phase_rows = [row for row in eligible if row["workload_kind"] == phase] + selected.extend( + {**row, "boundary_axes": _boundary_axes(row, phase_rows)} + for row in _select(phase_rows, max_points_per_phase, seed, execution_anchors) + ) + coordinates = {_coordinate(row) for row in selected} + # Exclude every physical row at each selected cell-coordinate, independently + # of collector cell ID, sample ID or DP-rank/provenance fields. + retained_indices = [ + index for index, row in enumerate(all_rows) if not (_matches(row, request) and _coordinate(row) in coordinates) + ] + plan = { + "schema_version": 1, + "scope": "withheld_coordinate_direct_interpolation", + "inputs": inputs, + "request_sha256": hashlib.sha256(_json_bytes(request.model_dump(mode="json"))).hexdigest(), + "policy": { + "seed": seed, + "max_points_per_phase": max_points_per_phase, + "target_holdout_fraction": 0.20, + "small_table_minimum_points": 1, + "max_p95_relative_error": max_p95_relative_error, + "max_unsupported_fraction": max_unsupported_fraction, + "selection": "retained_envelope_execution_boundaries_seeded_farthest_shape_v2", + "fold": "simultaneously_remove_all_selected_coordinates", + }, + "capture_boundaries": boundaries, + "source_row_count": len(all_rows), + "matching_row_count": len(matching), + "eligible_row_count": len(eligible), + "excluded_fake_fallback_count": len(matching) - len(eligible), + "removed_row_count": len(all_rows) - len(retained_indices), + "training_row_count": len(retained_indices), + "retained_envelope_anchors": [ + _point(coordinate) + for phase in request.collection_phases + for coordinate in sorted(_anchors([row for row in eligible if row["workload_kind"] == phase])) + ], + "selected_points": selected, + } + if request.worker_type is not None: + plan["worker_type"] = request.worker_type + plan["collection_phases"] = list(request.collection_phases) + plan_path = output / "holdout-plan.json" + plan_path.write_bytes(_json_bytes(plan)) + training = output / "systems" + (training / system_relative).parent.mkdir(parents=True, exist_ok=True) + (training / system_relative).write_bytes(source_bytes[0]) + parquet = training / parquet_relative + parquet.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table.take(pa.array(retained_indices, type=pa.int64())), parquet, compression="zstd") + metadata.update( + row_count=len(retained_indices), + parquet_sha256=_identity(parquet)["sha256"], + validation_derivation={"kind": "coordinate_holdout", "plan": _identity(plan_path)}, + ) + (training / metadata_relative).write_bytes(_json_bytes(metadata)) + config = _config(request, training) + config_path = output / "perfmodel-config.json" + config_path.write_bytes(_json_bytes(config.to_dict())) + model = RustForwardPassPerfModel.best_available(config) + predictions = [] + try: + diagnostics = model.diagnostics() + for row in selected: + point = { + "phase": row["workload_kind"], + "coordinates": {name: row[name] for name in _COORDINATES}, + "cell_id": row["cell_id"], + "kv_seed_regime": row.get("kv_seed_regime"), + "boundary_axes": row["boundary_axes"], + "measured_ms": row["latency_ms"], + } + try: + predicted = model.estimate_forward_pass_time_ms(_metrics(row)) + except PerfDataNotAvailableError as error: + point.update(status="unsupported", reason=str(error)) + else: + if predicted is None or not math.isfinite(predicted) or predicted <= 0: + raise ValueError("native interpolation returned a non-positive or non-finite timing") + point.update( + status="predicted", + predicted_ms=predicted, + relative_error=(predicted - row["latency_ms"]) / row["latency_ms"], + ) + predictions.append(point) + coverage = model.fpm_query_coverage() + finally: + model.close() + if not coverage or coverage["queries"]["measured"]: + raise ValueError("holdout evaluation encountered a measured lookup; withheld data leaked into training") + phases = { + phase: _phase_report( + [row for row in predictions if row["phase"] == phase], + max_p95_relative_error=max_p95_relative_error, + max_unsupported_fraction=max_unsupported_fraction, + ) + for phase in request.collection_phases + } + unchanged = all(_identity(Path(identity["path"])) == identity for identity in inputs.values()) + unchanged &= all( + _identity(Path(graph["source"]["path"])) == graph["source"] + for cell in boundaries["cells"] + for graph in cell["native_graph_config"] + ) + status = "passed" + if not unchanged or any(item["status"] == "failed" for item in phases.values()): + status = "failed" + elif boundaries["status"] == "incomplete" or any(item["status"] == "not_assessed" for item in phases.values()): + status = "incomplete" + report = { + "schema_version": 1, + "status": status, + "scope": "forward_timing_interpolation_only", + "serving_accuracy": "not_assessed", + "measurement_noise": "not_separated_from_interpolation_error", + "reference_measurement_policy": metadata["measurement_policy"], + "limitations": [ + "Each reference is the original single-sample formal-table measurement; separate subset " + "repeatability evidence does not remove measurement noise from these errors.", + "The bounded representative holdout is not a statistical confidence bound or full-domain accuracy claim.", + "When available, boundary retention uses native expected point modes and resolved graph configuration, " + "not observed per-call dispatch.", + "This onboarding selection retains evidenced mode boundaries; it does not test missing-boundary " + "interpolation or change predict/recommend.", + "Capture-bucket and kernel changes within one graph mode are not classified by this assessment.", + ], + "source_artifacts_unchanged": unchanged, + "policy": plan["policy"], + "capture_boundaries": boundaries, + "phases": phases, + "predictions": predictions, + "native_query_coverage": coverage, + "native_diagnostics": diagnostics, + "artifacts": { + "holdout_plan": _identity(plan_path), + "perfmodel_config": _identity(config_path), + "training_parquet": _identity(parquet), + "training_metadata": _identity(training / metadata_relative), + "training_system": _identity(training / system_relative), + "source": inputs, + }, + "issues": [] if unchanged else ["source FPM artifacts changed during validation"], + } + if request.worker_type is not None: + report["worker_type"] = request.worker_type + report["collection_phases"] = list(request.collection_phases) + (output / "interpolation-validation.json").write_bytes(_json_bytes(report)) + return report diff --git a/python/aisimulate/src/aisimulate/support/plan.py b/python/aisimulate/src/aisimulate/support/plan.py index 4a62c73a5..7b43aac40 100644 --- a/python/aisimulate/src/aisimulate/support/plan.py +++ b/python/aisimulate/src/aisimulate/support/plan.py @@ -10,6 +10,7 @@ import json import os import stat +import tempfile from collections.abc import Iterator from contextlib import contextmanager from importlib.resources import files @@ -28,22 +29,26 @@ def request_id(request: SupportRequest) -> str: - """Bind the entire declared request, including workload and search settings.""" + """Identify reusable collection inputs, independent of evaluation traffic.""" - payload = json.dumps(request.model_dump(mode="json", exclude_none=True), sort_keys=True, separators=(",", ":")) + collection = request.model_dump(mode="json", exclude_none=True) + collection.pop("workload") + for key in ("objective", "seed"): + collection["search"].pop(key) + payload = json.dumps(collection, sort_keys=True, separators=(",", ":")) return "onboarding-" + hashlib.sha256(payload.encode()).hexdigest() -def _parallelism(request: SupportRequest, replicas: int) -> dict[str, int]: - return request.parallelism(replicas) +def validation_id(request: SupportRequest) -> str: + """Bind generated validation examples to the complete saved request.""" + payload = json.dumps(request.model_dump(mode="json", exclude_none=True), sort_keys=True, separators=(",", ":")) + return "validation-" + hashlib.sha256(payload.encode()).hexdigest() def _configs( request: SupportRequest, root: Path ) -> tuple[dict[str, Any], dict[str, dict[str, Any]], list[dict[str, int]]]: - max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.worker_gpus) - replicas = list(dict.fromkeys((1, max_replicas)))[: request.search.max_candidates] - presets = [_parallelism(request, count) for count in replicas] + preset = request.parallelism() workload = request.workload common = { "traffic": { @@ -71,33 +76,54 @@ def _configs( "scheduler": request.scheduler_limits(), "timing": {"type": "default", "estimation_mode": "fpm_interpolation", "fallback_policy": "deny"}, } + if request.collection.gpu_memory_utilization is not None: + worker["kv_cache"] = {"capacity": {"memory_fraction": request.collection.gpu_memory_utilization}} if request.fpm_profile is not None: engine["fpm_profile"] = request.fpm_profile.model_dump(mode="json") worker["timing"]["estimator_config"] = {"fpm_interpolation": {"method": "direct"}} + if request.profile_deployment().resources.cache_layout == "grouped": + worker.setdefault("kv_cache", {})["prefix_caching"] = False + if request.worker_type in {"prefill", "decode"}: + # A role is independently collectable. A complete P/D serving config + # also needs the independently accepted counterpart and runtime support. + worker["timing"]["systems_paths"] = [str(root / "systems")] + return ( + { + "schema_version": "aisimulate-onboarding-worker/v1", + "worker_type": request.worker_type, + "serving_mode": "disaggregated", + "engine": { + **engine, + "mode": "disaggregated", + "workers": {request.worker_type: {**worker, "parallelism": preset}}, + }, + "collection": request.collection_settings(), + }, + {}, + [preset], + ) prediction = { **common, - "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": presets[0]}}}, + "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": preset}}}, } CorePredictionConfig.model_validate(prediction) - recommendations = {} - for preset in presets: - recommendation = { - **common, - "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": {"preset": [preset]}}}}, - "optimization": { - "target": request.search.objective, - "constraints": {"max_candidate_gpus": request.identity.gpu_count}, - }, - "optimizer": {"algorithm": "random", "max_trials": 1, "parallelism": 1, "seed": request.search.seed}, - } - CoreRecommendationConfig.model_validate(recommendation) - name = "pilot" if preset["replicas"] == 1 else f"replicas-{preset['replicas']}" - recommendations[name] = recommendation - return prediction, recommendations, presets + recommendation = { + **common, + "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": {"preset": [preset]}}}}, + "optimization": { + "target": request.search.objective, + # Match this validation worker even when it exceeds the runtime's default cap. + # This is derived scope, not a declaration of available deployment GPUs. + "constraints": {"max_candidate_gpus": request.worker_gpus}, + }, + "optimizer": {"algorithm": "random", "max_trials": 1, "parallelism": 1, "seed": request.search.seed}, + } + CoreRecommendationConfig.model_validate(recommendation) + return prediction, {"pilot": recommendation}, [preset] def _commands(request: SupportRequest, root: Path, recommendation_names: list[str]) -> dict[str, Any]: - return { + commands = { "fpm_plan_local": fpm_cli_args(request, output_dir=root, plan_only=True), "fpm_run_local": [ "aisimulate", @@ -135,6 +161,31 @@ def _commands(request: SupportRequest, root: Path, recommendation_names: list[st for name in recommendation_names ], } + if request.worker_type in {"prefill", "decode"}: + commands["predict"] = [] + commands["recommend"] = [] + from .runtime import runtime_collection_inputs, runtime_probe_manifest + + if runtime_probe_manifest(request) is not None: + _, deployment = runtime_collection_inputs(request, None) + for name, value in deployment.model_dump(exclude_none=True).items(): + for item in value if name == "container_mount" else [value]: + commands["fpm_run_local"].extend(("--" + name.replace("_", "-"), str(item))) + if request.fpm_profile is not None and ( + not request.profile_deployment().resources.memory_ready or runtime_probe_manifest(request) is not None + ): + commands["finalize"] = [ + "aisimulate", + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(root.with_name(root.name + "-resolved")), + ] + return commands def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any], dict[Path, bytes]]: @@ -150,32 +201,42 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] ) prediction, recommendations, presets = _configs(request, root) commands = _commands(request, root, list(recommendations)) + runtime_prerequisite = ( + "Prepare and verify the declared model/tokenizer snapshots, visible GPUs, model access and compatible " + "Dynamo self-benchmark runtime. At collection, choose Kubernetes with the packaged Generator and " + "required deployment permissions, or a caller-owned sbatch/salloc Slurm allocation with Pyxis/Enroot, " + "an explicit --image and a campaign directory shared at the same absolute path across allocated nodes. " + "The Slurm allocation must match the collector plan's node count and provide the required GPUs per node. " + "This plan does not download a pinned checkpoint or inspect a running worker. " + ) plan = { - "schema_version": "aisimulate-support-plan/v1", + "schema_version": "aisimulate-support-plan/v2", "request_id": request_id(request), + "validation_id": validation_id(request), "search": { "candidate_count": len(presets), - "max_candidates": request.search.max_candidates, - "candidates": [ - {"parallelism": preset, "total_gpus": preset["replicas"] * request.worker_gpus} for preset in presets - ], + "candidates": [{"parallelism": preset, "required_gpus": request.worker_gpus} for preset in presets], "detail": ( - "A single selected worker is planned; recommendation does not compare alternative configurations." - if len(presets) == 1 - else "Run one independent one-trial recommendation for the selected worker and one for the maximum " - "same-worker replicas that fit within each node. Compare their separate results; no combined ranking " - "or additional topology search is generated." + "Prediction and recommendation are single-worker validation examples for the selected topology. " + "The recommendation GPU cap equals that worker's required GPUs; no available allocation is declared. " + "Set deployment replicas and optimization budgets in ordinary predict/recommend configs." ), "baseline_rule": "Selected single worker; model parallelism legality and memory fit remain unchecked.", }, "fpm": { "status": "planned", - "worker_gpus": request.worker_gpus, + "collection_gpus_required": request.worker_gpus, + "resource_requirement": ( + "At least this many GPUs are required for one selected worker (attention TP times attention DP). " + "Actual available GPUs, placement and collector runtime resources remain unchecked." + ), "plan_command": commands["fpm_plan_local"], + "runtime_limits": request.collection_settings(), "sampling": ( - "Prefill is bounded by max(2, input_tokens * concurrency) and concurrency. Decode sampling and " - "runtime context auto-fitting use the collector's existing profile. The synthetic request count " - "does not bound timing samples. Smoke and limited runs do not publish formal FPM data." + "Dynamo self-benchmark generates the prefill/decode grid from CUDA graph sizes and the reviewed " + "runtime limits. The new-token budget, per-request context, scheduler sequence bound and total KV " + "capacity are separate quantities. Exact points resolve at runtime; validation traffic and SLAs " + "do not change these collection settings. Smoke and limited runs do not publish formal FPM data." ), }, "prerequisites": [ @@ -192,11 +253,10 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] { "id": "runtime", "status": "not_checked", - "detail": ( - "Prepare and verify the declared model/tokenizer revisions, vLLM version, visible GPUs, " - "model access, and the packaged Dynamo/Kubernetes/Generator collector runtime. Local invocation " - "uses that existing runtime. The collector does not apply revision/version declarations; " - "use a pinned local model snapshot where needed." + "detail": runtime_prerequisite + + ( + "Without an FPM profile, the collector does not enforce model/tokenizer revision or " + "vLLM version declarations; use and verify pinned local snapshots and a pinned runtime." ), }, { @@ -219,29 +279,53 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "systems_root": str(root / "systems"), }, } + if "prefill_cudagraph_policy" in request.collection.model_fields_set: + plan["fpm"]["launch"] = ( + "AISimulate launches and manages benchmark workers when collection is executed; no separately " + "launched HTTP server is required. vLLM initializes the model, caches and CUDA graphs; Dynamo " + "self-benchmark generates and measures the feasible grid. Initialization and planning do not launch GPUs." + ) if request.fpm_profile is not None: from aisimulate_core.sdk.memory import estimate_kv_cache deployment = request.profile_deployment() - scheduler = prediction["engine"]["workers"]["aggregated"]["scheduler"] - estimate = estimate_kv_cache( - request.identity.model, - request.identity.gpu, - request.identity.framework, - backend_version=request.identity.framework_version, - max_num_tokens=scheduler["max_batched_tokens"], - max_batch_size=scheduler["max_sequences"], - memory_fraction_kind="of_total", - memory_fraction_value=0.9, - tp_size=deployment.tp, - pp_size=deployment.pp, - attention_dp_size=deployment.dp, - moe_tp_size=deployment.moe_tp, - moe_ep_size=deployment.moe_ep, - fpm_profile=request.fpm_profile, + scheduler = request.scheduler_limits() + estimate = ( + estimate_kv_cache( + request.identity.model, + request.identity.gpu, + request.identity.framework, + backend_version=request.identity.framework_version, + max_num_tokens=scheduler["max_batched_tokens"], + max_batch_size=scheduler["max_sequences"], + memory_fraction_kind="of_total", + memory_fraction_value=request.collection.memory_fraction, + tp_size=deployment.tp, + pp_size=deployment.pp, + attention_dp_size=deployment.dp, + moe_tp_size=deployment.moe_tp, + moe_ep_size=deployment.moe_ep, + fpm_profile=request.fpm_profile, + context_length=request.search.context_length, + **({"worker_type": request.worker_type} if request.worker_type is not None else {}), + ) + if deployment.resources.memory_ready + else {"memory_source": "pending", "simulation_ready": False} ) - if estimate["total_kv_size_tokens"] <= request.search.context_length: - raise ValueError("declared FPM resources leave insufficient rank-local KV capacity for the pilot context") + if not deployment.resources.memory_ready: + plan["fpm"]["scheduling_policy"] = ( + "Synchronous scheduling in both phases; current collector prefill benchmarking requires it. " + "This is a collection policy, not the default for arbitrary vLLM serving." + if request.identity.framework_version == "0.27.0" + else "Native timing collection policy; runtime memory observation is unavailable for this vLLM version." + ) + elif deployment.resources.cache_layout == "grouped": + if estimate["request_peak_cache_bytes"] > estimate["total_kv_size_bytes"]: + raise ValueError( + "declared FPM resources leave insufficient rank-local bytes for grouped cache peak allocation" + ) + elif estimate["total_kv_size_tokens"] <= request.search.context_length: + raise ValueError("declared FPM resources leave insufficient rank-local KV capacity for the runtime context") plan["resources"] = estimate plan["search"]["baseline_rule"] = ( "Selected profile deployment; declared topology and rank-local resource bounds pass CPU admission. " @@ -255,27 +339,74 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "Generated configs use direct interpolation; matching timing coverage is still required." ), ) - plan["prerequisites"][1]["detail"] = ( - "Prepare and verify the declared model/tokenizer revisions, visible GPUs, model access, and the " - "packaged Dynamo/Kubernetes/Generator collector runtime. This plan does not download a pinned " - "checkpoint or inspect the running runtime. During collection execution, the collector checks " - "the observed Pod vLLM version against the profile's literal backend version before benchmarking." - ) - plan["fpm"]["sampling"] = ( - "Prefill sampling is capped by the pilot input length and the generated rank-local scheduler limits. " - "These bounds do not establish complete timing coverage of worker-wide FPM queries, including DEP. " - "Decode uses the collector's existing sampling profile within the declared resource envelope. " - "The synthetic request count does not bound timing samples. Smoke and limited runs do not publish " - "formal FPM data." + plan["prerequisites"][1]["detail"] = runtime_prerequisite + ( + "During collection execution, the collector checks the observed worker vLLM version against the " + "profile's literal backend version before benchmarking on either executor. Model/tokenizer " + "revision declarations and this version check do not verify the loaded checkpoint weights." ) plan["prerequisites"][2]["detail"] = ( "Collect matching prefill/decode timings into the local systems tree before prediction or recommendation. " "The profile resource estimate does not establish timing coverage or measured accuracy." ) + if not deployment.resources.memory_ready: + plan["search"]["baseline_rule"] = ( + "Selected profile deployment; memory remains pending runtime initialization. " + "Collection is permitted; prediction, recommendation and replay require finalized memory." + ) + plan["prerequisites"][0].update( + status="runtime_memory_pending", + detail=( + "Identity, precision, topology and cache geometry are declared. Runtime memory is collected " + "during worker initialization for supported vLLM 0.27.0 layouts; other versions can collect " + "timings but leave memory unresolved. No activation or non-KV byte declaration is required. " + "Run onboard finalize after successful collection, then review the resolved profile." + ), + ) + elif deployment.resources.memory_source == "runtime": + plan["search"]["baseline_rule"] = ( + "Selected deployment uses observed cache capacity for the exact recorded runtime settings. " + "Timing coverage and silicon accuracy require separate validation." + ) + plan["prerequisites"][0].update( + status="runtime_memory_resolved", + detail=( + "Observed cache capacity and grouped geometry are bound to this deployment and runtime envelope." + ), + ) plan["outputs"]["fpm_model_profile"] = str(root / "fpm-model-profile.json") + role_only = request.worker_type in {"prefill", "decode"} + if request.worker_type is not None: + plan["fpm"]["worker_type"] = request.worker_type + plan["fpm"]["collection_phases"] = list(request.collection_phases) + plan["fpm"]["serving_mode"] = "disaggregated" if role_only else "aggregated" + if role_only: + plan["outputs"]["prediction_configs"] = [] + plan["outputs"]["worker_config"] = str(root / "worker.yaml") + plan["search"]["detail"] = ( + "This plan collects one independently accepted P/D serving role. worker.yaml is a role configuration " + "artifact, not an executable prediction or recommendation config. Combine compatible, independently " + "accepted prefill and decode profiles in an ordinary serving configuration before simulation." + ) + plan["fpm"]["sampling"] = ( + f"Dynamo self-benchmark generates only the {request.worker_type} grid for this accepted role. " + "The other role is collected independently. Decode initializes representative state locally; actual " + "P/D transfer is not required. Timing collection and qualification do not establish replay readiness." + ) + plan["prerequisites"][2]["detail"] = ( + f"Collect and qualify matching {request.worker_type} timings in this role's systems tree. " + "Do not merge or replace another role's timing or memory evidence." + ) + plan["replay"] = { + "status": "not_assessed", + "detail": ( + "An independently accepted counterpart and supported P/D replay configuration are required. " + "The guided validate-fpm path currently handles aggregated serving only; grouped-cache P/D " + "handoff support is outside collection qualification." + ), + } yaml_documents = { Path("request.yaml"): request.model_dump(mode="json", exclude_none=True), - Path("predict/pilot.yaml"): prediction, + Path("worker.yaml" if role_only else "predict/pilot.yaml"): prediction, **{Path(f"recommend/{name}.yaml"): config for name, config in recommendations.items()}, } documents = {path: yaml.safe_dump(data, sort_keys=False).encode() for path, data in yaml_documents.items()} @@ -283,8 +414,16 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] documents[Path("fpm-model-profile.json")] = ( json.dumps(request.fpm_profile.model_dump(mode="json"), indent=2, sort_keys=True) + "\n" ).encode() + from .finalization import finalization_manifest + + manifest = finalization_manifest(request) + if manifest is not None: + documents[Path("finalization.json")] = (json.dumps(manifest, indent=2, sort_keys=True) + "\n").encode() documents[Path("systems") / source_spec.name] = source_spec.read_bytes() documents[Path("commands.json")] = (json.dumps(commands, indent=2, sort_keys=True) + "\n").encode() + plan["generated_files_sha256"] = { + path.as_posix(): hashlib.sha256(content).hexdigest() for path, content in documents.items() + } documents[Path("support-plan.json")] = (json.dumps(plan, indent=2, sort_keys=True) + "\n").encode() return plan, documents @@ -313,6 +452,7 @@ def _check_paths(root: Path) -> None: "support-plan.json", "fpm-model-profile.json", "commands.json", + "worker.yaml", "predict", "predict/pilot.yaml", "recommend", @@ -333,41 +473,72 @@ def _check_paths(root: Path) -> None: pending.extend(path.iterdir()) -def check_plan(request: SupportRequest, root: Path) -> None: +def check_plan(request: SupportRequest, root: Path, *, allow_missing: bool = False) -> None: """Verify saved identity before collecting or reusing any local output.""" _check_paths(root) + if allow_missing and not (root / "support-plan.json").exists(): + _, documents = _plan_documents(request, root) + _check_interrupted_plan(request, root, documents) + return try: prior = json.loads((root / "support-plan.json").read_text(encoding="utf-8")) except (OSError, ValueError) as exc: raise ValueError( f"a readable onboarding plan is required in {root}; run aisimulate onboard plan first" ) from exc - if not isinstance(prior, dict) or prior.get("request_id") != request_id(request): + if not isinstance(prior, dict) or prior.get("schema_version") != "aisimulate-support-plan/v2": + raise ValueError( + f"legacy or invalid onboarding plan in {root}; regenerate in a new output directory. " + "Legacy plans used validation traffic to choose collection bounds and cannot be reused implicitly." + ) + if prior.get("request_id") != request_id(request): raise ValueError(f"refusing to mix a different request identity in {root}; choose a new output directory") - _check_saved_request(request, root) - - -def _check_saved_request(request: SupportRequest, root: Path) -> None: try: saved = SupportRequest.from_yaml(root / "request.yaml") except ValueError as exc: raise ValueError(f"saved request in {root} was modified or cannot be read") from exc + if request_id(saved) != prior["request_id"] or validation_id(saved) != prior.get("validation_id"): + raise ValueError(f"saved request identity in {root} was modified or differs from the requested plan") + expected, documents = _plan_documents(saved, root) + if prior != expected: + raise ValueError( + f"generated plan input {root / 'support-plan.json'} was modified; choose a new output directory" + ) + # Recompute the manifest rather than trusting hashes in an edited plan. An + # evaluation-only refresh must never bless modified collector inputs. + for relative, content in documents.items(): + destination = root / relative + if destination.is_symlink(): + raise ValueError(f"refusing symlinked plan output {destination}") + if not destination.exists() and allow_missing: + continue + if not destination.is_file() or destination.read_bytes() != content: + raise ValueError( + f"generated plan input {destination} was modified or missing; choose a new output directory" + ) + from .finalization import verify_finalized_data + + verify_finalized_data(saved, root) + + +def _check_interrupted_plan(request: SupportRequest, root: Path, documents: dict[Path, bytes]) -> None: + """Require intact inputs before completing a plan whose manifest was never written.""" + + try: + saved = SupportRequest.from_yaml(root / "request.yaml") + except (OSError, ValueError) as exc: + raise ValueError(f"saved request in {root} was modified or cannot be read") from exc if saved != request: raise ValueError(f"saved request identity in {root} differs from the requested plan") - if request.fpm_profile is not None and (root / "fpm-model-profile.json").exists(): - from aisimulate_core.sdk.fpm_profile import load_fpm_profile - - try: - profile = load_fpm_profile((root / "fpm-model-profile.json").read_text(encoding="utf-8")) - except (OSError, ValueError) as exc: - raise ValueError(f"saved FPM model profile in {root} is missing or invalid") from exc - if profile != request.fpm_profile: - raise ValueError(f"saved FPM model profile in {root} differs from the requested identity") + for relative, content in documents.items(): + destination = root / relative + if destination.exists() and (not destination.is_file() or destination.read_bytes() != content): + raise ValueError(f"generated plan input {destination} was modified; choose a new output directory") def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: bool = False) -> dict[str, Any]: - """Write a new plan, or repair missing files for an identical saved request.""" + """Write a plan or safely refresh validation inputs for the same collection.""" root = Path(output_dir).expanduser().resolve() plan, documents = _plan_documents(request, root) @@ -377,28 +548,33 @@ def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: b nonempty = any(path.name != _LOCK_NAME for path in root.iterdir()) if nonempty: if not overwrite: - raise ValueError(f"output directory {root} is nonempty; use overwrite only for the same request") - if (root / "support-plan.json").exists(): - check_plan(request, root) - else: - # request.yaml is written first, so it can establish identity - # after an interruption before the completed plan was saved. - _check_saved_request(request, root) + raise ValueError(f"output directory {root} is nonempty; use overwrite only for the same collection") + check_plan(request, root, allow_missing=True) # Inspect every generated file before writing any file. Existing data, - # checkpoints, results, and edited inputs are never replaced or deleted. + # checkpoints and results are never replaced or deleted. Only verified + # generated inputs may change after an evaluation-only request edit. pending: dict[Path, bytes] = {} for relative, content in documents.items(): destination = root / relative if destination.is_symlink(): raise ValueError(f"refusing symlinked plan output {destination}") - if not destination.exists(): + if not destination.exists() or destination.read_bytes() != content: pending[relative] = content - elif destination.read_bytes() != content: - raise ValueError(f"generated plan input {destination} was modified; choose a new output directory") for relative, content in pending.items(): destination = root / relative destination.parent.mkdir(parents=True, exist_ok=True) - with destination.open("xb") as handle: - handle.write(content) + if not destination.exists(): + with destination.open("xb") as handle: + handle.write(content) + else: + temporary = None + try: + with tempfile.NamedTemporaryFile(dir=destination.parent, delete=False) as handle: + temporary = Path(handle.name) + handle.write(content) + os.replace(temporary, destination) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) (root / "systems/data").mkdir(parents=True, exist_ok=True) return plan diff --git a/python/aisimulate/src/aisimulate/support/runtime.py b/python/aisimulate/src/aisimulate/support/runtime.py new file mode 100644 index 000000000..2c6602b70 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/runtime.py @@ -0,0 +1,998 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint-backed runtime probing and CPU-only observation import.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import uuid +from copy import deepcopy +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate.fpm_profile import FpmModelProfile + +from .config_profile import load_model_config, profile_from_observations, validate_overrides +from .schema import ( + AGENTX_REFERENCE_CONTEXT, + CollectionSpec, + FPMDeployment, + SearchProfile, + SupportIdentity, + SupportRequest, +) + +PROVENANCE_SCHEMA = "aisimulate-onboarding-runtime-probe/v1" +_PRECISION = ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kv_cache_dtype", + "comm_quant_mode", + "moe_backend", + "attention_backend", + "enable_wideep", + "enable_eplb", +) +_TOPOLOGY = ("tensor_parallel", "attention_data_parallel", "moe_tensor_parallel", "moe_expert_parallel") + + +def add_runtime_parsers(actions: Any) -> None: + from .cli import add_deployment_arguments + + probe = actions.add_parser( + "probe-runtime", help="Preview or execute cache observations before profile geometry is known." + ) + imported = actions.add_parser( + "import-observations", help="Validate saved runtime evidence into fresh unaccepted profile drafts." + ) + for parser in (probe, imported): + parser.add_argument( + "--checkpoint", required=True, help="Existing single onboarding checkpoint, outside output directories." + ) + parser.add_argument( + "--configuration", action="append", help="Existing checkpoint key; repeat or omit for every configuration." + ) + parser.add_argument("--output-dir", required=True) + probe.add_argument( + "--instrumentation", help="Campaign-local YAML/JSON manifest; otherwise use a compatible bundled observer." + ) + probe.add_argument( + "--execute", action="store_true", help="Execute the serving role's phases for every selected configuration." + ) + probe.add_argument("--resume", action="store_true") + add_deployment_arguments(probe, default_executor=None) + imported.add_argument( + "--observations", required=True, help="Runner-owned observations.json index; observer code is never executed." + ) + + +def _json(value: Any) -> str: + return json.dumps(value, sort_keys=True, indent=2, allow_nan=False) + "\n" + + +def _hash(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _object(value: Any, label: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError(f"{label} must be an object") + return value + + +def _inputs(state: Any, config: Any) -> dict[str, Any]: + from .checkpoint import _merge + + return _merge(state.inputs, config.inputs) + + +def checkpoint_launch( + state: Any, + name: str, + checkpoint: Path, + deployment_overrides: dict[str, Any] | None = None, + *, + resolve_cpu_defaults: bool = False, + resume_output_dir: Path | None = None, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + """Resolve stage-2 facts; configuration names never determine topology.""" + from collector.fpm_forward.runtime_probe import normalize_probe_launch + + config = state.configurations[name] + inputs = _inputs(state, config) + payload = deepcopy(config.draft_request) + if "worker_type" in inputs: + if "worker_type" in payload and payload["worker_type"] != inputs["worker_type"]: + raise ValueError("inputs.worker_type conflicts with draft_request.worker_type") + payload["worker_type"] = inputs["worker_type"] + for section, model in (("identity", SupportIdentity), ("search", SearchProfile), ("collection", CollectionSpec)): + values = dict(_object(payload.get(section, {}), f"draft_request.{section}")) + for field in model.model_fields: + if field in inputs: + if field in values and values[field] != inputs[field]: + raise ValueError(f"inputs.{field} conflicts with draft_request.{section}.{field}") + values[field] = inputs[field] + payload[section] = values + if not config.draft_request: + payload["collection"].setdefault("prefill_cudagraph_policy", "runtime") + if not payload.get("fpm_profile"): + missing = [key for key in (*_TOPOLOGY, "context_length") if payload["search"].get(key) is None] + if missing: + raise ValueError("explicit search launch facts required: " + ", ".join(missing)) + # A probe validates launch selections, independently of a simulation + # workload or the previous profile's observed resource envelope. + identity = SupportIdentity.model_validate(payload["identity"]) + profile = FpmModelProfile.model_validate(payload["fpm_profile"]) if payload.get("fpm_profile") else None + search_values = dict(payload["search"]) + if profile is not None: + search_values.setdefault("context_length", min(profile.context_length, AGENTX_REFERENCE_CONTEXT)) + search = SearchProfile.model_validate(search_values) + collection = CollectionSpec.model_validate(payload["collection"]) + topology = { + "tp": search.tensor_parallel, + "pp": 1, + "dp": search.attention_data_parallel or 1, + "moe_tp": search.moe_tensor_parallel or (search.tensor_parallel if identity.model_kind == "moe" else 1), + "moe_ep": search.moe_expert_parallel or 1, + "cp": 1, + } + precision = {} + selected = ( + profile.select( + model=identity.model, + system=identity.gpu, + backend=identity.framework, + backend_version=identity.framework_version, + tp_size=topology["tp"], + attention_dp_size=topology["dp"], + moe_tp_size=topology["moe_tp"], + moe_ep_size=topology["moe_ep"], + worker_type=payload.get("worker_type") or "aggregated", + ) + if profile is not None + else None + ) + if selected is not None: + precision.update({name: getattr(selected, name) for name in _PRECISION}) + resource_overrides = validate_overrides(inputs.get("resource_overrides")) + provided = {key: value for key, value in resource_overrides.items() if key in _PRECISION} + explicit = _object(inputs.get("precision", {}), "inputs.precision") + if set(explicit) - set(_PRECISION): + raise ValueError("unknown inputs.precision fields: " + ", ".join(sorted(set(explicit) - set(_PRECISION)))) + for key, value in explicit.items(): + if key in provided and provided[key] != value: + raise ValueError(f"inputs.precision.{key} conflicts with resource_overrides") + provided[key] = value + for key, value in provided.items(): + if key in precision and precision[key] != value: + raise ValueError(f"inputs.precision.{key} conflicts with draft profile; edit the draft before probing") + precision[key] = value + precision.setdefault("comm_quant_mode", "half") + if "kv_cache_dtype" in precision: + precision["kvcache_quant_mode"] = precision.pop("kv_cache_dtype") + model = inputs.get("model_config") + if isinstance(model, str): + model = {"path": model} + model = _object(model, "inputs.model_config (path string or path/sha256 object)") + path = Path(model.get("path", "")).expanduser() + # HF cache snapshots use symlinks into a blob directory. Keep the logical + # parent so adjacent configuration sources are read from that snapshot. + path = (path if path.is_absolute() else checkpoint.parent / path).absolute() + model_config = load_model_config(path) + if model.get("sha256") is not None and model["sha256"] != model_config.sha256: + raise ValueError("model config content differs from inputs.model_config.sha256") + # The collector's explicit local config loader reads this adjacent ModelOpt + # sidecar. Preserve its bytes separately from the unmodified config.json. + sidecar = path.parent / "hf_quant_config.json" + source_files = {sidecar.name: _hash(sidecar)} if sidecar.is_file() else {} + if model.get("source_files", source_files) != source_files: + raise ValueError("model config sidecar files differ from inputs.model_config.source_files") + deployment = { + **_object(inputs.get("collection_deployment", {}), "inputs.collection_deployment"), + **(deployment_overrides or {}), + } + deployment = FPMDeployment.model_validate(deployment).model_dump(mode="json", exclude_none=True) + if resume_output_dir is not None: + from collector.fpm_forward.runtime_probe import resolve_probe_cpu_policy + + deployment = resolve_probe_cpu_policy(name, deployment, resume_output_dir, resume=True) + elif resolve_cpu_defaults and deployment["executor"] == "slurm": + from collector.fpm_forward.config import resolve_slurm_cpu_policy + + cpus, binding = resolve_slurm_cpu_policy(deployment.get("cpus_per_task"), deployment.get("cpu_bind")) + deployment.update(cpus_per_task=cpus, cpu_bind=binding) + max_tokens = collection.max_num_tokens or (selected.resources.max_num_tokens if selected else 8192) + max_sequences = collection.max_batch_size or (selected.resources.max_batch_size if selected else 256) + if max_tokens < max_sequences: + raise ValueError("resolved collection max_num_tokens must be at least max_batch_size") + launch = normalize_probe_launch( + { + **({"worker_type": payload["worker_type"]} if payload.get("worker_type") is not None else {}), + "identity": identity.model_dump(mode="json", exclude_none=True), + "topology": topology, + "precision": precision, + "collection": { + "max_model_len": search.context_length, + "max_num_batched_tokens": max_tokens, + "max_num_seqs": max_sequences, + "gpu_memory_utilization": collection.memory_fraction, + "prefill_cudagraph_policy": collection.prefill_cudagraph_policy, + "max_prefill_cudagraph_size": collection.max_prefill_cudagraph_size, + **{ + field: getattr(collection, field) + for field in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size") + if getattr(collection, field) is not None + }, + }, + "model_config": { + "path": str(path), + "sha256": model_config.sha256, + **({"source_files": source_files} if source_files else {}), + }, + "deployment": deployment, + } + ) + return payload, launch, resource_overrides + + +def _snapshot(index: Path, name: str, destination: Path) -> tuple[Path, list[str]]: + """Copy one active attempt, verifying bytes without executing observer code.""" + from collector.fpm_forward.runtime_instrumentation import contained_file, load_instrumentation, read_json + + content = index.read_bytes() + document = read_json(index) + saved = deepcopy(_object(document["configurations"].get(name, {}), "observation configuration")) + attempts = saved.get("attempts", []) + if not isinstance(attempts, list) or any(not isinstance(attempt, dict) for attempt in attempts): + raise ValueError("observation attempts must be a list of objects") + active = [attempt for attempt in attempts if attempt.get("attempt_id") == saved.get("active_attempt_id")] + if len(active) == 1: + saved["attempts"] = active + destination.mkdir(parents=True) + diagnostics = [] + + def copy_file(relative: str, digest: str | None = None) -> None: + try: + source = contained_file(index.parent, relative) + data = source.read_bytes() + if digest is not None and hashlib.sha256(data).hexdigest() != digest: + raise ValueError(f"artifact hash mismatch: {relative}") + target = destination / "raw" / relative + target.parent.mkdir(parents=True, exist_ok=True) + if target.exists() and target.read_bytes() != data: + raise ValueError(f"conflicting snapshot artifact: {relative}") + target.write_bytes(data) + except (OSError, ValueError) as exc: + diagnostics.append(str(exc)) + + for attempt in saved.get("attempts", []): + try: + manifest_path = contained_file(index.parent, attempt["bundle"]["manifest"]) + bundle = load_instrumentation(manifest_path) + if bundle.sha256 != attempt["bundle"]["sha256"]: + raise ValueError("instrumentation bundle hash mismatch") + copy_file(str(manifest_path.relative_to(index.parent))) + for relative, digest in bundle.files.items(): + copy_file(str((bundle.root / relative).relative_to(index.parent)), digest) + attempt["bundle"]["manifest"] = "raw/" + attempt["bundle"]["manifest"] + except (KeyError, TypeError, OSError, ValueError) as exc: + diagnostics.append(str(exc)) + for phase in _object(attempt.get("phases", {}), "observation phases").values(): + phase = _object(phase, "observation phase") + for ref in [phase.get("launch_manifest", {}), *phase.get("artifacts", [])]: + try: + copy_file(ref["path"], ref["sha256"]) + ref["path"] = "raw/" + ref["path"] + except (KeyError, TypeError) as exc: + diagnostics.append(f"invalid artifact reference: {exc}") + (destination / "source-index.json").write_bytes(content) + target = destination / "observations.json" + target.write_text(_json({"schema_version": document.get("schema_version"), "configurations": {name: saved}})) + return target, diagnostics + + +def runtime_probe_manifest(request: SupportRequest) -> dict[str, Any] | None: + selected = request.profile_deployment() + runtime = selected.resources.runtime_memory if selected else None + if runtime is None: + return None + return _probe_provenance(runtime.provenance) + + +def _probe_provenance(provenance: Any) -> dict[str, Any] | None: + try: + value = json.loads(provenance) + except (TypeError, ValueError): + return None + if not isinstance(value, dict) or value.get("schema_version") != PROVENANCE_SCHEMA: + return None + if ( + any( + not isinstance(value.get(key), str) or not value[key] + for key in ("checkpoint", "configuration", "observations_index") + ) + or any( + not isinstance(value.get(key), dict) for key in ("observed_resources", "model_metadata", "user_overrides") + ) + or not isinstance(value.get("source_artifacts"), list) + or any( + not isinstance(ref, dict) or not isinstance(ref.get("path"), str) or not isinstance(ref.get("sha256"), str) + for ref in value.get("source_artifacts", []) + ) + ): + raise ValueError("invalid onboarding runtime probe provenance") + return value + + +def _resource_values(resources: dict[str, Any]) -> dict[str, Any]: + values = deepcopy(resources) + values.pop("provenance", None) + values.get("runtime_memory", {}).pop("provenance", None) + return values + + +def _request_launch(request: SupportRequest, launch: dict[str, Any]) -> dict[str, Any]: + result = deepcopy(launch) + if request.worker_type is None: + result.pop("worker_type", None) + else: + result["worker_type"] = request.worker_type + result["identity"] = request.identity.model_dump(mode="json", exclude_none=True) + parallel = request.parallelism() + result["topology"].update( + tp=parallel["tensor"], + dp=parallel["attention_data"], + moe_tp=parallel["moe_tensor"], + moe_ep=parallel["moe_expert"], + ) + selected = request.profile_deployment() + result["precision"] = { + ("kvcache_quant_mode" if key == "kv_cache_dtype" else key): getattr(selected, key) for key in _PRECISION + } + scheduler = request.scheduler_limits() + result["collection"].update( + max_model_len=request.search.context_length, + max_num_batched_tokens=scheduler["max_batched_tokens"], + max_num_seqs=scheduler["max_sequences"], + gpu_memory_utilization=request.collection.memory_fraction, + prefill_cudagraph_policy=request.collection.prefill_cudagraph_policy, + max_prefill_cudagraph_size=request.collection.max_prefill_cudagraph_size + or (2048 if request.collection.prefill_cudagraph_policy == "explicit" else None), + ) + for field in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"): + value = getattr(request.collection, field) + if value is None: + result["collection"].pop(field, None) + else: + result["collection"][field] = value + return result + + +def verify_runtime_profile(request: SupportRequest) -> dict[str, Any] | None: + """Revalidate source files and the exact runtime binding of an imported draft.""" + manifest = runtime_probe_manifest(request) + if manifest is None: + from .finalization import ( + _merge_resources, + _verify_collection, + finalization_manifest, + verify_finalized_quality, + ) + from .plan import request_id + + finalized = finalization_manifest(request) + quality = verify_finalized_quality(finalized) if finalized is not None else None + if finalized is not None and "runtime_probe" in finalized: + root = Path(finalized["source_directory"]) + original = SupportRequest.from_yaml(root / "request.yaml") + if request_id(original) != finalized.get("source_request_id"): + raise ValueError("finalized runtime source request changed") + source = verify_runtime_profile(original) + if source is None or finalized["runtime_probe"] != runtime_probe_manifest(original): + raise ValueError("finalized runtime probe provenance changed") + observations, evidence, _, _ = _verify_collection( + original, root, memory_revision=finalized.get("memory_revision") + ) + if quality is not None: + evidence["collection_quality"] = quality + expected = _merge_resources( + observations, + evidence, + source_references=finalized.get("observation_provenance") == "source_references", + ) + if request.profile_deployment().resources.model_dump(mode="json") != expected: + raise ValueError("finalized runtime resources differ from verified formal observations") + launch = source["provenance"]["launch"] + if _request_launch(request, launch) != launch: + raise ValueError("finalized runtime settings changed; a compatible fresh probe is required") + metadata = { + key: getattr(request.fpm_profile, key) for key in ("architecture", "context_length", "num_experts") + } + if metadata != finalized["runtime_probe"]["model_metadata"]: + raise ValueError("finalized model metadata changed; a compatible fresh probe is required") + return None + from collector.fpm_forward.runtime_instrumentation import read_json + from collector.fpm_forward.runtime_observations import validate_observations + + covered = set() + for ref in manifest["source_artifacts"]: + path = Path(ref["path"]) + if path.is_symlink() or not path.is_file() or _hash(path) != ref["sha256"]: + raise ValueError(f"runtime probe source artifact changed: {path}") + covered.add(str(path)) + index = Path(manifest["observations_index"]) + if {str(index), str(index.parent / "launch.json"), str(index.parent / "model-config.json")} - covered: + raise ValueError("runtime probe provenance is missing immutable source artifacts") + name = manifest["configuration"] + launch = read_json(index)["configurations"][name]["launch"] + result = validate_observations(index, {name: launch})[name] + if result["status"] != "complete": + raise ValueError("runtime probe evidence is incomplete: " + "; ".join(result["diagnostics"])) + if _request_launch(request, launch) != launch: + raise ValueError("runtime-affecting draft edits require a compatible fresh probe") + expected = _resource_values(result["resources"]) + if manifest["observed_resources"] != expected: + raise ValueError("saved observed resources differ from immutable runtime evidence") + selected = request.profile_deployment() + actual = _resource_values(selected.resources.model_dump(mode="json")) + observed_capacity = expected["runtime_memory"].pop("kv_cache_bytes") + actual_capacity = actual["runtime_memory"].pop("kv_cache_bytes") + if actual != expected: + raise ValueError("runtime cache geometry or settings changed; a compatible fresh probe is required") + overrides = ( + {} + if actual_capacity == observed_capacity + else { + "runtime_memory.kv_cache_bytes": { + "observed": observed_capacity, + "value": actual_capacity, + "source": "user override; assumption, not a runtime measurement", + } + } + ) + if manifest.get("user_overrides", {}) != overrides: + raise ValueError( + "resource edits require checkpoint registration to preserve observed values and label user overrides" + ) + metadata = {key: getattr(request.fpm_profile, key) for key in ("architecture", "context_length", "num_experts")} + if metadata != manifest["model_metadata"]: + raise ValueError("model metadata edits require a compatible fresh probe") + return result + + +def preserve_runtime_overrides(previous: dict[str, Any], current: dict[str, Any]) -> dict[str, Any]: + """Keep observed source values when checkpoint edits change a resource assumption.""" + profile = previous.get("fpm_profile") + deployments = profile.get("deployments", []) if isinstance(profile, dict) else [] + if not isinstance(deployments, list): + return current + payload = deepcopy(current) + for position, deployment in enumerate(deployments): + if not isinstance(deployment, dict): + continue + resources = deployment.get("resources") + runtime = resources.get("runtime_memory") if isinstance(resources, dict) else None + manifest = _probe_provenance(runtime.get("provenance")) if isinstance(runtime, dict) else None + if manifest is None or not payload.get("fpm_profile"): + continue + # Removing the whole profile leaves an incomplete draft. Its prior + # request/evidence and acceptance are retained by checkpoint history. + updated_deployments = payload["fpm_profile"].get("deployments", []) + if position >= len(updated_deployments): + continue + updated_resources = updated_deployments[position].get("resources", {}) + updated_runtime = updated_resources.get("runtime_memory", {}) + updated = _probe_provenance(updated_runtime.get("provenance")) if isinstance(updated_runtime, dict) else None + if updated is None: + from .finalization import finalization_manifest + from .plan import request_id + + request = SupportRequest.model_validate(current) + finalized = finalization_manifest(request) + if finalized is not None and ( + ( + finalized.get("runtime_probe") == manifest + and finalized.get("source_request_id") == request_id(SupportRequest.model_validate(previous)) + ) + or ( + isinstance(finalized.get("memory_revision"), dict) + and finalized["memory_revision"].get("request_id") + == request_id(SupportRequest.model_validate(previous)) + ) + ): + verify_runtime_profile(request) + return current + raise ValueError( + "runtime evidence cannot be removed from an imported profile; import a fresh compatible probe" + ) + if updated["observations_index"] != manifest["observations_index"]: + verify_runtime_profile(SupportRequest.model_validate(current)) + return current + # Do not validate a complete simulation request here: runtime settings + # may be edited before a compatible replacement probe is available. + capacity = updated_runtime.get("kv_cache_bytes") + observed = manifest["observed_resources"]["runtime_memory"]["kv_cache_bytes"] + manifest["user_overrides"] = ( + {} + if capacity == observed + else { + "runtime_memory.kv_cache_bytes": { + "observed": observed, + "value": capacity, + "source": "user override; assumption, not a runtime measurement", + } + } + ) + updated_runtime["provenance"] = _json(manifest) + return payload + + +def verify_checkpoint_runtime(state: Any, name: str, checkpoint: Path) -> None: + request = SupportRequest.model_validate(state.configurations[name].draft_request) + manifest = runtime_probe_manifest(request) + if manifest is None: + return + from collector.fpm_forward.runtime_instrumentation import read_json + + _, launch, _ = checkpoint_launch(state, name, checkpoint) + saved = read_json(Path(manifest["observations_index"]))["configurations"][name]["launch"] + if launch != saved: + raise ValueError("checkpoint runtime inputs changed; a compatible fresh probe is required") + + +def runtime_collection_inputs( + request: SupportRequest, deployment: FPMDeployment | None +) -> tuple[list[str], FPMDeployment | None]: + """Supply the collector's explicit opt-in flags from verified source evidence.""" + manifest = runtime_probe_manifest(request) + if manifest is None: + return [], deployment + result = verify_runtime_profile(request) + if result is None: + return [], deployment + root = Path(manifest["observations_index"]).parent + observed_deployment = FPMDeployment.model_validate(result["provenance"]["launch"]["deployment"]) + if deployment is not None and deployment != observed_deployment: + raise ValueError( + "collection deployment differs from the runtime probe; use matching options or probe the changed deployment" + ) + return [ + "--fpm-runtime-instrumentation", + str(root / result["provenance"]["instrumentation"]["manifest"]), + "--fpm-runtime-launch", + str(root / "launch.json"), + "--fpm-runtime-configuration", + manifest["configuration"], + ], observed_deployment + + +def verify_runtime_acceptance(request: SupportRequest) -> None: + manifest = runtime_probe_manifest(request) + if manifest is None: + return + from .checkpoint import _canonical_request, _load, report + + checkpoint = Path(manifest["checkpoint"]) + state = _load(checkpoint) + name = manifest["configuration"] + if name not in state.configurations or not report(state, checkpoint)["configurations"][name]["profile_accepted"]: + raise ValueError( + "formal collection requires explicit acceptance of the imported runtime profile in its checkpoint" + ) + if _canonical_request(state.configurations[name], require_complete=True) != request.model_dump( + mode="json", exclude_none=True + ): + raise ValueError("formal collection request differs from the accepted checkpoint draft") + + +def verify_collection_runtime( + request: SupportRequest, + observations: Path, + *, + collection_checkpoint: dict[str, Any] | None = None, + memory_request: SupportRequest | None = None, +) -> dict[str, Any]: + """Validate new raw observations before asserting conservative compatibility.""" + from collector.fpm_forward.runtime_observations import validate_observations + + source = verify_runtime_profile(request) + if source is None: + raise ValueError("formal runtime comparison requires imported probe evidence") + manifest = runtime_probe_manifest(request) + name = manifest["configuration"] + result = validate_observations(observations, {name: source["provenance"]["launch"]})[name] + if result["status"] != "complete": + raise ValueError("formal runtime observations are incomplete: " + "; ".join(result["diagnostics"])) + if collection_checkpoint is not None: + from collector.fpm_forward.runtime_instrumentation import contained_file, read_json + + entry = read_json(observations)["configurations"][name] + active = entry["active_attempt_id"] + if collection_checkpoint.get("runtime_observation_attempt_id") != active: + raise ValueError("formal runtime observations belong to a different collection attempt") + attempt = next(value for value in entry["attempts"] if value["attempt_id"] == active) + seen = set() + for phase in attempt["phases"].values(): + context = read_json(contained_file(observations.parent, phase["launch_manifest"]["path"])) + cell = context.get("cell_id") + native_attempt = context.get("collector_attempt_id") + saved = collection_checkpoint.get("cells", {}).get(cell, {}) + if ( + not isinstance(cell, str) + or cell in seen + or not isinstance(native_attempt, str) + or phase.get("cell_id") != cell + or phase.get("collector_attempt_id") != native_attempt + or saved.get("status") != "passed" + or saved.get("attempt_id") != native_attempt + ): + raise ValueError("formal runtime phase does not match the successful native cell attempt") + seen.add(cell) + if seen != set(collection_checkpoint.get("cells", {})): + raise ValueError("formal runtime observations do not cover every collected cell") + for field in ("bundle_sha256", "runtime", "runtime_settings"): + if source["provenance"][field] != result["provenance"][field]: + raise ValueError(f"formal runtime {field} differs from the accepted probe; review a fresh compatible probe") + + def geometry(evidence): + groups = next( + item["evidence"]["cache"]["groups"] + for item in evidence["artifacts"] + if item["evidence"]["kind"] == "worker" + ) + return [{key: value for key, value in group.items() if key != "pool_id"} for group in groups] + + if geometry(source["provenance"]) != geometry(result["provenance"]): + raise ValueError("formal runtime cache group geometry differs from the accepted probe") + expected = _resource_values(source["resources"]) + actual = _resource_values(result["resources"]) + expected["runtime_memory"].pop("kv_cache_bytes") + capacity = actual["runtime_memory"].pop("kv_cache_bytes") + if actual != expected: + raise ValueError("formal runtime cache layout or settings differ from the accepted probe") + bound = request.profile_deployment().resources.runtime_memory.kv_cache_bytes + if memory_request is not None: + if collection_checkpoint is None: + raise ValueError("memory revision requires complete formal collection attempt bindings") + _verify_memory_revision(request, memory_request, result) + bound = memory_request.profile_deployment().resources.runtime_memory.kv_cache_bytes + if capacity < bound: + raise ValueError( + f"formal runtime usable capacity {capacity} is smaller than the accepted bound {bound}; review new evidence" + ) + result["compatibility"] = { + "status": "compatible", + "accepted_capacity_bytes": bound, + "observed_capacity_bytes": capacity, + "policy": "unchanged runtime/settings/geometry and observed usable capacity at least the accepted bound", + "source_observations_index": manifest["observations_index"], + "observations_index": str(observations), + "observations_sha256": _hash(observations), + } + return result + + +def _verify_memory_revision(original: SupportRequest, revised: SupportRequest, observed: dict[str, Any]) -> None: + """Allow only a smaller capacity measured by these exact formal attempts.""" + selected = original.profile_deployment() + position = original.fpm_profile.deployments.index(selected) + payload = revised.model_dump(mode="json") + try: + resources = payload["fpm_profile"]["deployments"][position]["resources"] + capacity = resources["runtime_memory"]["kv_cache_bytes"] + resources["runtime_memory"]["kv_cache_bytes"] = selected.resources.runtime_memory.kv_cache_bytes + resources["runtime_memory"]["provenance"] = selected.resources.runtime_memory.provenance + resources["provenance"] = selected.resources.provenance + payload["fpm_profile"]["provenance"] = original.fpm_profile.provenance + except (KeyError, TypeError, IndexError) as exc: + raise ValueError("memory revision requires the original runtime profile and deployment") from exc + if payload != original.model_dump(mode="json"): + raise ValueError("memory revision changed non-capacity onboarding inputs") + if capacity >= selected.resources.runtime_memory.kv_cache_bytes: + raise ValueError("memory revision must lower the original accepted capacity") + source = runtime_probe_manifest(original) + revision = runtime_probe_manifest(revised) + if ( + source is None + or revision is None + or any(revision[key] != source[key] for key in ("checkpoint", "configuration", "model_metadata")) + or revision["user_overrides"] + ): + raise ValueError("memory revision requires imported formal observations without capacity overrides") + verified = verify_runtime_profile(revised) + if verified is None or capacity != observed["resources"]["runtime_memory"]["kv_cache_bytes"]: + raise ValueError("memory revision capacity must equal the verified formal observed minimum") + + def identity(evidence): + value = deepcopy(evidence) + value.pop("observations_index", None) + # Observation imports relocate immutable bytes into a snapshot. Paths + # can change; attempts, record contents and every artifact digest cannot. + for field in ("artifacts", "launch_artifacts", "runtime_artifacts"): + value[field] = sorted( + [{key: item for key, item in ref.items() if key != "path"} for ref in value[field]], + key=_json, + ) + value["instrumentation"].pop("manifest") + return value + + if identity(verified["provenance"]) != identity(observed["provenance"]): + raise ValueError("memory revision is not derived from the same complete formal collection observations") + + +def _new_root(value: str, checkpoint: Path, *, fresh: bool) -> Path: + raw = Path(value).expanduser() + if raw.is_symlink(): + raise ValueError("runtime output must not be a symlink") + root = raw.resolve() + if checkpoint == root or checkpoint.is_relative_to(root): + raise ValueError("keep the session checkpoint outside runtime output directories") + if fresh and root.exists(): + raise ValueError("runtime import output must be a fresh directory") + return root + + +def _selected(state: Any, requested: list[str] | None) -> list[str]: + names = requested or list(state.configurations) + if not names or len(names) != len(set(names)): + raise ValueError("select at least one unique checkpoint configuration") + unknown = set(names) - state.configurations.keys() + if unknown: + raise ValueError("unknown checkpoint configurations: " + ", ".join(sorted(unknown))) + return names + + +def _draft( + payload: dict[str, Any], + launch: dict[str, Any], + overrides: dict[str, Any], + result: dict[str, Any], + snapshot: Path, + checkpoint: Path, + name: str, +) -> SupportRequest: + # Full workload validation belongs to simulation draft construction. A + # prior profile may describe different settings; replace it only after + # independently validating the newly observed resources below. + request = SupportRequest.model_validate({**payload, "fpm_profile": None}) + precision = { + ("kv_cache_dtype" if key == "kvcache_quant_mode" else key): value for key, value in launch["precision"].items() + } + model = load_model_config(launch["model_config"]["path"]) + if model.sha256 != launch["model_config"]["sha256"]: + raise ValueError("model config changed during runtime import") + profile = profile_from_observations(model, request, precision, result["resources"], overrides) + if payload.get("fpm_profile"): + provenance = json.loads(profile.provenance) + provenance["source_profile_provenance"] = payload["fpm_profile"].get("provenance") + profile.provenance = _json(provenance) + manifest = { + "schema_version": PROVENANCE_SCHEMA, + "checkpoint": str(checkpoint), + "configuration": name, + "observations_index": str(snapshot), + "source_artifacts": [ + {"path": str(path), "sha256": _hash(path)} for path in sorted(snapshot.parent.rglob("*")) if path.is_file() + ], + "observed_resources": _resource_values(result["resources"]), + "model_metadata": {key: getattr(profile, key) for key in ("architecture", "context_length", "num_experts")}, + "user_overrides": {}, + "review_status": "requires_review", + "timing_status": "not_established_by_runtime_probe", + } + profile.deployments[0].resources.runtime_memory.provenance = _json(manifest) + payload = deepcopy(payload) + payload["fpm_profile"] = profile.model_dump(mode="json", exclude_none=True) + return SupportRequest.model_validate(payload) + + +def import_observations( + checkpoint: Path, observations: Path, output: Path, configurations: list[str] | None = None +) -> dict[str, Any]: + from collector.fpm_forward.runtime_instrumentation import read_json + from collector.fpm_forward.runtime_observations import validate_observations + + from .checkpoint import _load, save_checkpoint + + state = _load(checkpoint) + names = _selected(state, configurations) + output = _new_root(str(output), checkpoint, fresh=True) + if observations.is_symlink(): + raise ValueError("runtime observation index cannot be a symlink") + document = read_json(observations) + if document.get("schema_version") != "aisimulate-runtime-observations/v1" or not isinstance( + document.get("configurations"), dict + ): + raise ValueError("unsupported runtime observation index schema") + output.mkdir(parents=True) + results, patches = {}, {} + for ordinal, name in enumerate(names, 1): + directory = output / f"configuration-{ordinal:03d}" + directory.mkdir() + result = { + "status": "incomplete", + "resources": None, + "diagnostics": [], + "provenance": {"source_index": str(observations)}, + } + request = None + try: + snapshot, diagnostics = _snapshot(observations, name, directory / "evidence") + original, launch, overrides = checkpoint_launch(state, name, checkpoint) + validated = validate_observations(snapshot, {name: launch})[name] + result.update(validated) + result["diagnostics"] += diagnostics + if diagnostics: + result["status"] = "incomplete" + if result["status"] == "complete": + # Frozen formal inputs belong to this configuration's immutable + # snapshot, alongside its model configuration and raw records. + (snapshot.parent / "launch.json").write_text(_json(launch)) + (snapshot.parent / "model-config.json").write_bytes(Path(launch["model_config"]["path"]).read_bytes()) + for relative, digest in launch["model_config"].get("source_files", {}).items(): + source = Path(launch["model_config"]["path"]).parent / relative + content = source.read_bytes() + if hashlib.sha256(content).hexdigest() != digest: + raise ValueError(f"model configuration sidecar changed during import: {relative}") + (snapshot.parent / relative).write_bytes(content) + request = _draft(original, launch, overrides, result, snapshot, checkpoint, name) + (directory / "request.yaml").write_text( + yaml.safe_dump(request.model_dump(mode="json", exclude_none=True), sort_keys=False) + ) + (directory / "fpm-model-profile.json").write_text( + _json(request.fpm_profile.model_dump(mode="json", exclude_none=True)) + ) + except (ValueError, OSError, KeyError, TypeError) as exc: + result["status"] = "incomplete" + result["diagnostics"].append(str(exc)) + (directory / "source-index.json").write_bytes(observations.read_bytes()) + validation = directory / "validation.json" + validation.write_text(_json(result)) + old = state.configurations[name] + artifacts = {} + if request is not None and result["status"] == "complete": + artifacts.update( + { + key: {"archived": True} + for key, ref in old.artifacts.items() + if not ref.archived + and (ref.kind in {"request", "profile"} or ref.scope != "input" or key.startswith("runtime-")) + } + ) + token = uuid.uuid4().hex[:12] + for number, path in enumerate(sorted(directory.rglob("*"))): + if path.is_file(): + kind = ( + "request" + if path.name == "request.yaml" + else "profile" + if path.name == "fpm-model-profile.json" + else "file" + ) + artifacts[f"runtime-{token}-{number}"] = { + "path": str(path), + "kind": kind, + "scope": "input", + "archived": result["status"] != "complete", + } + patch = { + "artifacts": artifacts, + "history": [ + *old.history, + { + "event": "runtime_observations_import", + "validation": str(validation), + "sha256": _hash(validation), + "status": result["status"], + "diagnostics": result["diagnostics"], + "archived_artifacts": [key for key, ref in artifacts.items() if ref.get("archived")], + "reason": "Fresh evidence supersedes prior drafts and dependent collection/validation outputs." + if result["status"] == "complete" + else "Failed import retained as historical evidence; current profile inputs are unchanged.", + }, + ], + "progress": { + "stage": 3, + "status": "in_progress" if result["status"] == "complete" else "blocked", + "blockers": result["diagnostics"], + "next_action": "Review and accept this exact draft." + if result["status"] == "complete" + else "Resolve runtime evidence diagnostics and import a fresh attempt.", + }, + } + if request is not None and result["status"] == "complete": + patch["draft_request"] = request.model_dump(mode="json", exclude_none=True) + patches[name] = patch + results[name] = { + "status": result["status"], + "diagnostics": result["diagnostics"], + "validation": str(validation), + } + if request is not None and result["status"] == "complete": + results[name].update( + request=str(directory / "request.yaml"), profile=str(directory / "fpm-model-profile.json") + ) + state, _ = save_checkpoint( + checkpoint, patch={"configurations": patches}, expected_revision=state.revision, accept=[] + ) + result = { + "checkpoint": str(checkpoint), + "revision": state.revision, + "status": "complete" if all(value["status"] == "complete" for value in results.values()) else "incomplete", + "configurations": results, + } + (output / "import.json").write_text(_json(result)) + return result + + +def run_runtime_command(args: argparse.Namespace) -> int: + from collector.fpm_forward.runtime_probe import probe_runtime + + from .checkpoint import _load, _target, save_checkpoint + + checkpoint = _target(args.checkpoint) + if args.support_action == "import-observations": + result = import_observations( + checkpoint, Path(args.observations).expanduser().absolute(), Path(args.output_dir), args.configuration + ) + print(_json(result), end="") + return 0 if result["status"] == "complete" else 1 + state = _load(checkpoint) + names = _selected(state, args.configuration) + output = _new_root(args.output_dir, checkpoint, fresh=False) + overrides = {key: getattr(args, key) for key in FPMDeployment.model_fields if getattr(args, key, None) is not None} + launches, failed = {}, {} + for name in names: + try: + _, launches[name], _ = checkpoint_launch( + state, + name, + checkpoint, + overrides, + resolve_cpu_defaults=not args.resume, + resume_output_dir=output if args.resume else None, + ) + except (ValueError, OSError) as exc: + failed[name] = {"status": "incomplete", "diagnostics": [str(exc)]} + result = ( + probe_runtime( + launches, instrumentation=args.instrumentation, output_dir=output, execute=args.execute, resume=args.resume + ) + if launches + else {"status": "failed", "configurations": {}} + ) + result["configurations"].update(failed) + if failed: + result["status"] = "partial" if launches else "failed" + output.mkdir(parents=True, exist_ok=True) + path = output / f"onboarding-probe-{uuid.uuid4().hex}.json" + path.write_text(_json(result)) + patches = {} + for name in names: + config = state.configurations[name] + patch = { + "history": [ + *config.history, + { + "event": "runtime_probe", + "report": str(path), + "sha256": _hash(path), + "result": result["configurations"][name], + }, + ] + } + if name in launches: + patch["inputs"] = {"collection_deployment": launches[name]["deployment"]} + patches[name] = patch + state, _ = save_checkpoint( + checkpoint, patch={"configurations": patches}, expected_revision=state.revision, accept=[] + ) + result.update(checkpoint=str(checkpoint), revision=state.revision) + print(_json(result), end="") + return 0 if result["status"] in {"preview", "completed"} else 1 diff --git a/python/aisimulate/src/aisimulate/support/schema.py b/python/aisimulate/src/aisimulate/support/schema.py index 532a6af58..c07a4bf0e 100644 --- a/python/aisimulate/src/aisimulate/support/schema.py +++ b/python/aisimulate/src/aisimulate/support/schema.py @@ -6,15 +6,19 @@ from __future__ import annotations import re +from collections.abc import Mapping from pathlib import Path, PurePosixPath -from typing import Literal +from typing import Any, Literal from packaging.version import Version -from pydantic import Field, field_validator, model_validator +from pydantic import Field, SerializerFunctionWrapHandler, field_validator, model_serializer, model_validator from aisimulate.config.common import PositiveFiniteFloat, PositiveStrictInt, StrictModel, load_yaml from aisimulate.fpm_profile import FpmModelProfile +AGENTX_REFERENCE_CONTEXT = 256000 +CUDAGRAPH_MODES = ("NONE", "PIECEWISE", "FULL", "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE") + class SupportIdentity(StrictModel): model: str @@ -23,9 +27,6 @@ class SupportIdentity(StrictModel): framework: Literal["vllm"] = "vllm" framework_version: str gpu: str - gpu_count: PositiveStrictInt - node_count: PositiveStrictInt = 1 - gpus_per_node: PositiveStrictInt interconnect: str sm: PositiveStrictInt | None = None tokenizer_revision: str | None = None @@ -67,12 +68,6 @@ def _system_name(cls, value: str) -> str: raise ValueError("gpu must be a packaged system name, not a path") return value - @model_validator(mode="after") - def _allocation(self) -> SupportIdentity: - if self.node_count * self.gpus_per_node != self.gpu_count: - raise ValueError("node_count * gpus_per_node must equal gpu_count") - return self - class SloSpec(StrictModel): ttft_ms: PositiveFiniteFloat = 1000.0 @@ -100,8 +95,9 @@ class SearchProfile(StrictModel): attention_data_parallel: PositiveStrictInt | None = None moe_tensor_parallel: PositiveStrictInt | None = None moe_expert_parallel: PositiveStrictInt | None = None - context_length: PositiveStrictInt = 16384 - max_candidates: int = Field(default=1, strict=True, ge=1, le=2) + # Historical field location retained for saved requests. This is a runtime + # limit, independent of the optional synthetic validation workload. + context_length: PositiveStrictInt = AGENTX_REFERENCE_CONTEXT objective: Literal[ "throughput", "throughput_per_gpu", @@ -115,6 +111,67 @@ class SearchProfile(StrictModel): seed: int = Field(default=42, strict=True, ge=0) +class CollectionSpec(StrictModel): + """AISimulate runtime and capture limits for Dynamo's native grid generation.""" + + max_num_tokens: PositiveStrictInt | None = None + max_batch_size: PositiveStrictInt | None = None + # Missing policy preserves saved requests' explicit-2048 collector behavior. + # Fresh onboarding sets runtime explicitly at the CLI input boundary. + prefill_cudagraph_policy: Literal["runtime", "explicit"] = "explicit" + max_prefill_cudagraph_size: PositiveStrictInt | None = None + gpu_memory_utilization: float | None = Field(default=None, strict=True, gt=0, le=1, allow_inf_nan=False) + cudagraph_mode: Literal["NONE", "PIECEWISE", "FULL", "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"] | None = None + cudagraph_capture_sizes: list[PositiveStrictInt] | None = None + max_cudagraph_capture_size: PositiveStrictInt | None = None + + @model_serializer(mode="wrap") + def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + values = handler(self) + # Preserve old request hashes and generated-plan verification. An + # explicitly reviewed policy is always serialized, including explicit. + if "prefill_cudagraph_policy" not in self.model_fields_set: + values.pop("prefill_cudagraph_policy", None) + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"): + if getattr(self, name) is None: + values.pop(name, None) + return values + + @model_validator(mode="after") + def _capture_policy(self) -> CollectionSpec: + if self.prefill_cudagraph_policy == "runtime" and self.max_prefill_cudagraph_size is not None: + raise ValueError("runtime prefill_cudagraph_policy rejects max_prefill_cudagraph_size; use explicit") + sizes = self.cudagraph_capture_sizes + if sizes is not None: + if not sizes or sizes != sorted(set(sizes)): + raise ValueError("cudagraph_capture_sizes must be a nonempty, increasing list of distinct sizes") + if self.max_cudagraph_capture_size is not None and sizes[-1] != self.max_cudagraph_capture_size: + raise ValueError("max_cudagraph_capture_size must equal the largest cudagraph_capture_sizes value") + if self.cudagraph_mode == "NONE" and ( + sizes is not None + or self.max_cudagraph_capture_size is not None + or self.prefill_cudagraph_policy == "explicit" + ): + raise ValueError("cudagraph_mode=NONE rejects CUDA graph capture settings") + if (sizes is not None or self.max_cudagraph_capture_size is not None) and ( + self.prefill_cudagraph_policy == "explicit" or self.max_prefill_cudagraph_size is not None + ): + raise ValueError("generic CUDA graph capture settings require prefill_cudagraph_policy=runtime") + return self + + @property + def memory_fraction(self) -> float: + """vLLM's fraction of total GPU memory, shared with simulation admission.""" + return self.gpu_memory_utilization if self.gpu_memory_utilization is not None else 0.9 + + @field_validator("max_num_tokens") + @classmethod + def _collector_token_minimum(cls, value: int | None) -> int | None: + if value is not None and value < 2: + raise ValueError("FPM collection requires max_num_tokens >= 2") + return value + + _DNS_LABEL = r"[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?" _DNS_SUBDOMAIN = rf"{_DNS_LABEL}(?:\.{_DNS_LABEL})*" @@ -122,13 +179,45 @@ class SearchProfile(StrictModel): class FPMDeployment(StrictModel): """Deployment-only inputs included in the collector's frozen-plan identity.""" + executor: Literal["kubernetes", "slurm"] = "kubernetes" dynamo_version: str | None = None image: str | None = Field(default=None, pattern=r"^[^\s\x00]+$") + container_mount: list[str] = Field(default_factory=list) + cpus_per_task: int | None = Field(default=None, strict=True, gt=0) + cpu_bind: Literal["cores", "none"] | None = None namespace: str | None = Field(default=None, pattern=rf"^{_DNS_LABEL}$") model_cache: str | None = None transport: Literal["nvlink", "ib", "efa"] | None = None image_pull_secret: str | None = Field(default=None, pattern=rf"^{_DNS_SUBDOMAIN}$", max_length=253) + @model_validator(mode="after") + def _executor_options(self) -> FPMDeployment: + if self.executor == "slurm": + incompatible = [ + "--" + name.replace("_", "-") + for name in ("namespace", "model_cache", "image_pull_secret") + if getattr(self, name) is not None + ] + if incompatible: + raise ValueError("--executor slurm rejects Kubernetes options: " + ", ".join(incompatible)) + if self.image is None: + raise ValueError("--executor slurm requires --image for the Pyxis container") + elif self.container_mount: + raise ValueError("--container-mount requires --executor slurm; use --model-cache for Kubernetes") + elif self.cpus_per_task is not None or self.cpu_bind is not None: + raise ValueError("--cpus-per-task and --cpu-bind require --executor slurm") + return self + + @field_validator("container_mount") + @classmethod + def _container_mounts(cls, values: list[str]) -> list[str]: + if any( + not value.strip() or any(ord(char) < 32 or ord(char) == 127 or char == "," for char in value) + for value in values + ): + raise ValueError("container mounts must be nonempty and contain no control characters or commas") + return values + @field_validator("dynamo_version") @classmethod def _release_version(cls, value: str | None) -> str | None: @@ -152,16 +241,68 @@ def _model_cache(cls, value: str | None) -> str | None: class SupportRequest(StrictModel): - schema_version: Literal["aisimulate-support-request/v1"] = "aisimulate-support-request/v1" + schema_version: Literal["aisimulate-support-request/v2"] = "aisimulate-support-request/v2" identity: SupportIdentity workload: WorkloadSpec = Field(default_factory=WorkloadSpec) search: SearchProfile = Field(default_factory=SearchProfile) + collection: CollectionSpec = Field(default_factory=CollectionSpec) fpm_profile: FpmModelProfile | None = None + worker_type: Literal["aggregated", "prefill", "decode"] | None = None + + @model_serializer(mode="wrap") + def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + values = handler(self) + if self.worker_type is None: + values.pop("worker_type", None) + return values - def parallelism(self, replicas: int = 1) -> dict[str, int]: + @property + def collection_phases(self) -> tuple[str, ...]: + return (self.worker_type,) if self.worker_type in {"prefill", "decode"} else ("prefill", "decode") + + @model_validator(mode="before") + @classmethod + def _legacy_allocation(cls, values: Any) -> Any: + if isinstance(values, Mapping): + obsolete = [] + for section, fields in ( + ("identity", ("gpu_count", "node_count", "gpus_per_node")), + ("search", ("max_candidates",)), + ): + fields_in_request = values.get(section) + if isinstance(fields_in_request, Mapping): + obsolete.extend(f"{section}.{name}" for name in fields if name in fields_in_request) + if obsolete: + raise ValueError( + "Legacy onboarding fields are no longer supported: " + ", ".join(obsolete) + ". " + "Copy the request, remove these fields, and regenerate the plan in a new output directory. " + "Set deployment replicas and GPU budgets in ordinary predict/recommend configs." + ) + if values.get("schema_version") == "aisimulate-support-request/v1": + search = values.get("search", {}) + values = {**values, "schema_version": "aisimulate-support-request/v2"} + if isinstance(search, Mapping): + search = dict(search) + search.setdefault("context_length", 16384) + values["search"] = search + # Preserve explicit legacy limits and profile resource bounds. Do + # not infer collection limits from a saved validation workload. + profile = values.get("fpm_profile") + search = values.get("search", {}) + if profile is not None and isinstance(search, Mapping) and "context_length" not in search: + context = ( + profile.get("context_length") + if isinstance(profile, Mapping) + else getattr(profile, "context_length", None) + ) + if type(context) is int and context > 0: + values = {**values, "search": {**search, "context_length": min(context, AGENTX_REFERENCE_CONTEXT)}} + return values + + def parallelism(self) -> dict[str, int]: search = self.search return { - "replicas": replicas, + "replicas": 1, "tensor": search.tensor_parallel, "pipeline": 1, "attention_data": search.attention_data_parallel or 1, @@ -172,6 +313,7 @@ def parallelism(self, replicas: int = 1) -> dict[str, int]: @property def worker_gpus(self) -> int: + """Minimum GPUs required for the selected collection worker, not available capacity.""" parallel = self.parallelism() return parallel["tensor"] * parallel["attention_data"] @@ -198,20 +340,94 @@ def profile_deployment(self): attention_dp_size=parallel["attention_data"], moe_tp_size=parallel["moe_tensor"], moe_ep_size=parallel["moe_expert"], + worker_type=self.worker_type or "aggregated", ) def scheduler_limits(self) -> dict[str, int]: """Rank-local limits shared by generated configs and collection bounds.""" deployment = self.profile_deployment() - if deployment is None: - return {"max_batched_tokens": 8192, "max_sequences": 256} - return { - "max_batched_tokens": min(8192, deployment.resources.max_num_tokens), - "max_sequences": min(self.workload.concurrency, deployment.resources.max_batch_size), + limits = { + "max_batched_tokens": self.collection.max_num_tokens + or (deployment.resources.max_num_tokens if deployment is not None else 8192), + "max_sequences": self.collection.max_batch_size + or (deployment.resources.max_batch_size if deployment is not None else 256), } + if limits["max_batched_tokens"] < limits["max_sequences"]: + raise ValueError( + f"resolved collection max_num_tokens ({limits['max_batched_tokens']}) must be at least " + f"max_batch_size ({limits['max_sequences']}); edit the collection or profile bounds" + ) + return limits + + def collection_settings(self) -> dict[str, Any]: + """Resolved runtime inputs and the source of each proposed bound.""" + scheduler = self.scheduler_limits() + profile = self.fpm_profile is not None + settings: dict[str, Any] = { + "context_length": self.search.context_length, + **scheduler, + "max_prefill_cudagraph_size": self.collection.max_prefill_cudagraph_size or 2048, + "sources": { + "context_length": ( + "reviewed runtime limit; model/profile context capped at the 256000-token AgentX reference " + "unless explicitly overridden" + ), + "max_batched_tokens": "user collection override" + if self.collection.max_num_tokens is not None + else "reviewed profile resource bound" + if profile + else "initial vLLM collection policy (8192); review for the target runtime", + "max_sequences": "user collection override" + if self.collection.max_batch_size is not None + else "reviewed profile resource bound" + if profile + else "initial vLLM collection policy (256); independent of validation concurrency", + "max_prefill_cudagraph_size": "user collection override" + if self.collection.max_prefill_cudagraph_size is not None + else "collector default (2048); review against the target runtime CUDA graph configuration", + }, + } + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"): + value = getattr(self.collection, name) + if value is not None: + settings[name] = value + settings["sources"][name] = "reviewed serving configuration for this worker role" + if "prefill_cudagraph_policy" in self.collection.model_fields_set: + policy = self.collection.prefill_cudagraph_policy + settings["prefill_cudagraph_policy"] = policy + settings["sources"]["prefill_cudagraph_policy"] = ( + "reviewed runtime selection; CUDA graph sizes resolve in the initialized vLLM engine" + if policy == "runtime" + else "reviewed explicit capture override; match the intended serving configuration" + ) + if policy == "runtime": + settings["max_prefill_cudagraph_size"] = None + settings["sources"]["max_prefill_cudagraph_size"] = ( + "deferred to the pinned runtime; onboarding does not prescribe CUDA graph sizes" + ) + if self.collection.gpu_memory_utilization is not None: + settings["gpu_memory_utilization"] = self.collection.gpu_memory_utilization + settings["sources"]["gpu_memory_utilization"] = ( + "reviewed fraction of total GPU memory; initial policy is 0.90; runtime fit remains unverified" + ) + deployment = self.profile_deployment() + runtime = deployment.resources.runtime_memory if deployment is not None else None + if runtime is not None: + settings["sources"]["context_length"] = ( + f"reviewed context within the observed runtime memory limit of {runtime.max_model_len} tokens" + ) + settings["sources"]["gpu_memory_utilization"] = ( + f"recorded runtime memory fraction {runtime.gpu_memory_utilization}; reused without scaling capacity" + ) + return settings @model_validator(mode="after") def _shape(self) -> SupportRequest: + if self.worker_type is None and any( + getattr(self.collection, name) is not None + for name in ("cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size") + ): + raise ValueError("generic CUDA graph settings require an explicit worker_type") parallel = self.parallelism() tp, dp, mtp, ep = (parallel[name] for name in ("tensor", "attention_data", "moe_tensor", "moe_expert")) if self.identity.model_kind == "dense": @@ -222,12 +438,19 @@ def _shape(self) -> SupportRequest: ) if not valid: raise ValueError("onboarding requires a complete TP, DEP, or TEP topology; set the attention and MoE sizes") - if self.worker_gpus > self.identity.gpus_per_node: - raise ValueError("search.tensor_parallel * attention_data_parallel must fit within gpus_per_node") if self.workload.input_tokens + self.workload.output_tokens > self.search.context_length: raise ValueError("search.context_length must cover the input and output tokens") if self.fpm_profile is not None: - self.profile_deployment() + deployment = self.profile_deployment() + if self.worker_type is not None and deployment.worker_type != self.worker_type: + raise ValueError( + "an explicit worker_type requires a matching role-tagged FPM deployment; " + "do not relabel a legacy shared profile as independent role evidence" + ) + scheduler = self.scheduler_limits() + deployment.resources.validate_envelope( + max_num_tokens=scheduler["max_batched_tokens"], max_batch_size=scheduler["max_sequences"] + ) if self.identity.model_revision != self.fpm_profile.model_revision: raise ValueError("identity.model_revision must match fpm_profile.model_revision") if (self.fpm_profile.num_experts > 0) != (self.identity.model_kind == "moe"): diff --git a/python/aisimulate/src/aisimulate/support/serving_validation.py b/python/aisimulate/src/aisimulate/support/serving_validation.py new file mode 100644 index 000000000..2efa2f606 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/serving_validation.py @@ -0,0 +1,1233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prepare and assess a small, independently measured serving accuracy check. + +The producer is upstream AIPerf, not another implementation of agentic scheduling. +Initially only one complete single-stream Weka play is qualified. Other plays +remain explicit incomplete checks; they must not be flattened or partially run. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +import subprocess +import uuid +from importlib.resources import files +from itertools import pairwise +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit +from urllib.request import Request, urlopen + +import yaml + +from aisimulate.config import CorePredictionConfig + +AIPERF_REVISION = "7db2ba37a62aa80c882bc90eaf61cc8073e2387b" +AIPERF_SOURCE = "https://github.com/ai-dynamo/aiperf" +SCHEMA = "aisimulate-serving-validation/v1" +EXECUTION_FIELDS = ( + "model", + "model_revision", + "backend", + "backend_version", + "image_digest", + "hardware", + "parallelism", + "precision", + "attention_groups", + "graph_config", + "scheduler", + "context_length", + "gpu_memory_utilization", + "runtime_config", +) + + +def normalize_runtime_config(config: dict[str, Any]) -> dict[str, Any]: + """Compare native runtime behavior while allowing checkpoint mount paths to differ. + + Model/checkpoint identity is verified separately. Revision, dtype, quantization, + cache, scheduler, topology and graph settings remain in this comparison. + """ + normalized = json.loads(json.dumps(config)) + model = normalized.get("model_config", {}) + for field in ("model", "tokenizer"): + model.pop(field, None) + # Collection uses phase-specific KV seeding. Serving's prefix cache flag is + # checked against the replay scope separately, using the native observation. + normalized.get("cache_config", {}).pop("enable_prefix_caching", None) + return normalized + + +def _identity(path: Path) -> dict[str, Any]: + path = path.expanduser().resolve() + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + return {"path": str(path), "sha256": digest, "size_bytes": path.stat().st_size} + + +def _write(path: Path, value: Any) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8") + + +def _number(value: Any, label: str, *, positive: bool = False) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + raise ValueError(f"{label} must be a finite number") + if value < 0 or (positive and value == 0): + raise ValueError(f"{label} must be {'positive' if positive else 'nonnegative'}") + return float(value) + + +def _threshold(value: float, label: str) -> float: + return _number(value, label) + + +def _read_play(trace: Path) -> tuple[dict[str, Any] | None, list[str]]: + text = trace.read_text(encoding="utf-8") + try: + document = json.loads(text) + except json.JSONDecodeError: + documents = [json.loads(line) for line in text.splitlines() if line.strip()] + if len(documents) != 1: + return None, ["choose one complete representative Weka play; serving qualification accepts one play"] + document = documents[0] + if not isinstance(document, dict) or not isinstance(document.get("requests"), list): + return None, ["the selected file must contain one complete Weka play object"] + issues = [] + if not isinstance(document.get("id"), str) or not document["id"]: + issues.append("the selected play needs its original nonempty Weka id") + requests = document["requests"] + if not requests: + issues.append("the selected play is empty") + previous = None + for index, row in enumerate(requests): + if not isinstance(row, dict) or row.get("type") not in {"n", "s"}: + issues.append("branching plays are not yet qualified; select a complete single-stream play") + break + for key in ("in", "out"): + if type(row.get(key)) is not int or row[key] <= 0: + issues.append(f"request {index} needs positive {key} tokens for HTTP serving validation") + try: + timestamp = _number(row.get("t"), f"request {index} timestamp") + _number(row.get("api_time", 0), f"request {index} api_time") + except ValueError as error: + issues.append(str(error)) + break + hashes = row.get("hash_ids") + if not isinstance(hashes, list) or not hashes: + issues.append("single-stream qualification requires the original nonempty prefix hash lists") + break + if previous is not None and ( + timestamp < previous["t"] + previous.get("api_time", 0) + or row.get("model") != previous.get("model") + or hashes[: len(previous["hash_ids"])] != previous["hash_ids"] + ): + issues.append( + "overlap, model changes or non-extending prefixes can create flattened agents; " + "select a complete single-stream play without these ambiguities" + ) + break + previous = row + if requests and not any(type(row) is dict and type(row.get("out")) is int and row["out"] > 1 for row in requests): + issues.append("at least one multi-token response is required to assess TPOT") + return document, issues + + +def _timing_subset(play: dict[str, Any]) -> str | None: + """Identify source-invariant equivalence to both native timing bounds. + + AIS enforces max(not-before timestamp, predecessor completion + idle gap). + AIPerf concurrency uses gaps, and fixed schedule uses timestamps. Either is + equivalent for the subsets below, regardless of the measured response time. + """ + pairs = list(pairwise(play["requests"])) + if all(previous.get("api_time", 0) == 0 for previous, _ in pairs): + return "zero_recorded_api_time" + if all(abs(current["t"] - previous["t"] - previous.get("api_time", 0)) < 1e-9 for previous, current in pairs): + return "zero_idle_gaps" + return None + + +def _scope(prediction: dict[str, Any], trace: Path) -> dict[str, Any]: + config = CorePredictionConfig.model_validate(prediction).model_dump(mode="json", exclude_none=True) + traffic, engine = config["traffic"], config["engine"] + source, load = traffic["source"], traffic["load"] + if ( + source.get("type") != "trace" + or source.get("format") != "weka" + or [str(Path(path).expanduser().resolve()) for path in source.get("paths", [])] != [str(trace)] + or source.get("nested_timestamp_basis") not in {None, "absolute"} + or load.get("type") != "trace_timestamps" + or load.get("agentic_lanes") != 1 + or load.get("agentic_snapshot") is not None + or load.get("agentic_warmup", False) + or load.get("speedup", 1) != 1 + or traffic.get("stop") + or engine["mode"] != "aggregated" + or engine.get("speculation") is not None + ): + raise ValueError("serving validation requires complete cold aggregated Weka replay at original speed, one lane") + worker = engine["workers"]["aggregated"] + if worker["parallelism"].get("replicas", 1) != 1: + raise ValueError("serving validation initially qualifies exactly one aggregated worker") + if worker["kv_cache"].get("host_offload") or worker["kv_cache"].get("g3_offload"): + raise ValueError("serving validation initially qualifies HBM-only cache") + timing = worker["timing"] + if ( + timing.get("type") != "default" + or timing.get("estimation_mode") != "fpm_interpolation" + or timing.get("fallback_policy") != "deny" + or timing.get("estimator_config", {}).get("fpm_interpolation", {}).get("method") != "direct" + ): + raise ValueError("serving validation requires ordinary strict direct-FPM prediction") + return { + "qualification": "one_complete_single_stream_weka_play", + "replay": "cold_aggregated", + "load_policy": "single_session_absolute_and_completion_bounds", + "agentic_lanes": 1, + "prefix_caching": worker["kv_cache"]["prefix_caching"], + "cache_state_at_start": "cold", + "kv_initialization": "real_request_history", + "cache_storage": "hbm_only", + "speculation": "disabled", + "model_projection": engine["model"], + "token_count_source": "server_usage", + } + + +def _benchmark_arguments(recipe: dict[str, Any]) -> list[str]: + return [ + "-m", + "aiperf", + "profile", + "--url", + recipe["endpoint"], + "--model", + recipe["scope"]["model_projection"], + "--tokenizer", + recipe["tokenizer"], + "--endpoint-type", + "chat", + "--streaming", + "--use-server-token-count", + "--input-file", + recipe.get("matched_workload", {}).get("payloads", recipe["benchmark_trace"])["path"], + "--custom-dataset-type", + "mooncake_trace" if recipe.get("matched_workload") else "weka_trace", + "--fixed-schedule" if recipe["timing_subset"] == "zero_idle_gaps" else "--no-fixed-schedule", + "--concurrency", + "1", + "--num-sessions", + "1", + "--request-count", + str(recipe["request_count"]), + "--extra-inputs", + "ignore_eos:true", + "--random-seed", + "0", + "--export-level", + "records", + "--artifact-dir", + recipe["artifacts_dir"], + ] + + +def _aiperf_installation(aiperf_python: Path) -> dict[str, Any]: + probe = ( + "import importlib.metadata,json;d=importlib.metadata.distribution('aiperf');" + "print(json.dumps({'version':d.version,'source':json.loads(d.read_text('direct_url.json') or '{}')}))" + ) + checked = subprocess.run([str(aiperf_python), "-c", probe], check=True, text=True, capture_output=True) + installation = json.loads(checked.stdout) + source = installation.get("source", {}) + if ( + source.get("vcs_info", {}).get("commit_id") != AIPERF_REVISION + or source.get("url", "").removesuffix(".git").rstrip("/") != AIPERF_SOURCE + ): + raise ValueError(f"install AIPerf from {AIPERF_SOURCE}.git@{AIPERF_REVISION} in the benchmark environment") + return installation + + +def _materialize_payloads(recipe: dict[str, Any], output: Path, aiperf_python: Path) -> dict[str, Any]: + installation = _aiperf_installation(aiperf_python) + helper = output / "serving_workload.py" + helper.write_bytes(files("aisimulate.support").joinpath("serving_workload.py").read_bytes()) + inputs = output / "materialization-input.json" + _write( + inputs, + { + "trace": recipe["benchmark_trace"]["path"], + "output": str(output), + "tokenizer": recipe["tokenizer"], + "model": recipe["scope"]["model_projection"], + "arguments": _benchmark_arguments(recipe)[2:], + }, + ) + environment = {key: value for key, value in os.environ.items() if not key.startswith("AIPERF_")} + environment.update(HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1") + with ( + (output / "materialization.stdout.log").open("w") as stdout, + (output / "materialization.stderr.log").open("w") as stderr, + ): + completed = subprocess.run( + [str(aiperf_python), str(helper), str(inputs)], env=environment, stdout=stdout, stderr=stderr, check=False + ) + if completed.returncode: + raise ValueError( + f"native AIPerf request materialization failed; inspect {output / 'materialization.stderr.log'}" + ) + return {"installation": installation, "helper": _identity(helper), "input": _identity(inputs)} + + +def _prepare_matched_workload( + recipe: dict[str, Any], prediction: dict[str, Any], output: Path, python: Path +) -> dict[str, Any]: + from aisimulate.supervision import main as predict + + from .validation import _completion_issues + + output.mkdir() + materialization = _materialize_payloads(recipe, output, python) + tokenization_path = output / "tokenization.json" + tokenization = json.loads(tokenization_path.read_text()) + play = json.loads(Path(recipe["benchmark_trace"]["path"]).read_text()) + original = json.loads(Path(recipe["prediction_report"]["path"]).read_text()) + identities = {int(row["request_id"].rsplit(":", 1)[-1]): row for row in original["per_request"]} + requests = tokenization["requests"] + if len(requests) != len(play["requests"]) or sorted(identities) != list(range(len(requests))): + raise ValueError("materialization or original replay does not cover the complete selected play") + # Explicit native graph input retains the source serial dependencies. Feeding + # newly tokenized hashes through Weka's fork detector would change causality. + header = { + "schema": "dynamo.agentic_mooncake", + "version": 2, + # Preserve equality at every engine cache-block boundary, including + # prefixes ending inside a larger source Weka block. + "block_size": 1, + "hash_id_scope": "local", + "source": {"format": "target_tokenized_weka", "digest": _identity(tokenization_path)["sha256"]}, + } + rows = [] + for index, (source, target) in enumerate(zip(play["requests"], requests, strict=True)): + tokens = target["input_token_ids"] + if target["turn_index"] != index or target["input_length"] != len(tokens): + raise ValueError("native materialization has inconsistent input token identities") + old = identities[index] + dependencies = [] + if index: + previous = play["requests"][index - 1] + dependencies.append( + { + "request_id": identities[index - 1]["request_id"], + "trigger": "completion", + "relation": "sequence", + "delay_ms": (source["t"] - previous["t"] - previous.get("api_time", 0)) * 1000, + } + ) + rows.append( + { + "request_id": old["request_id"], + "play_id": old["play_id"], + "session_id": old["session_id"], + "model": source["model"], + "input_length": len(tokens), + "output_length": source["out"], + "hash_ids": tokens, + "not_before_ms": (source["t"] - play["requests"][0]["t"]) * 1000, + "recorded_api_time_ms": source.get("api_time", 0) * 1000, + "dependencies": dependencies, + } + ) + trace_path = output / "target-tokenized.agentic.jsonl" + trace_path.write_text("".join(json.dumps(row) + "\n" for row in [header, *rows])) + derived = json.loads(json.dumps(prediction)) + derived["traffic"]["source"] = {"type": "trace", "paths": [str(trace_path)], "format": "agentic_mooncake"} + config_path = output / "predict.yaml" + config_path.write_text(yaml.safe_dump(derived, sort_keys=False)) + status = predict( + [ + "predict", + "--config", + str(config_path), + "--output-dir", + str(output / "prediction"), + "--capture-per-request", + "--format", + "json", + ] + ) + prediction_path = output / "prediction/prediction.json" + coverage_path = output / "prediction/fpm-coverage.json" + if status or not prediction_path.is_file() or not coverage_path.is_file(): + raise ValueError("ordinary prediction on frozen target-tokenized requests did not finish") + replay = json.loads(prediction_path.read_text()) + issues = _completion_issues(replay) + if issues or json.loads(coverage_path.read_text()).get("status") != "covered": + raise ValueError("frozen target-tokenized requests lack complete direct-FPM coverage: " + "; ".join(issues)) + return { + **materialization, + "payloads": _identity(output / "requests.jsonl"), + "tokenization": _identity(tokenization_path), + "trace": _identity(trace_path), + "prediction_config": _identity(config_path), + "prediction_report": _identity(prediction_path), + "coverage_report": _identity(coverage_path), + } + + +def prepare_serving_validation( + *, + trace: Path, + prediction_config: Path, + output: Path, + endpoint: str, + tokenizer: Path, + expected_execution: dict[str, Any], + latency_p95_error: float = 0.20, + throughput_error: float = 0.20, + prediction_report: Path | None = None, + aiperf_python: Path | None = None, + max_dispatch_delay_ms: float = 5.0, + max_dispatch_delay_fraction: float = 0.01, +) -> Path: + """Save a reviewed, reproducible producer recipe; never start serving or a benchmark.""" + trace, prediction_config, output = (path.expanduser().resolve() for path in (trace, prediction_config, output)) + tokenizer = tokenizer.expanduser().resolve() + parsed_url = urlsplit(endpoint) + if ( + parsed_url.scheme not in {"http", "https"} + or not parsed_url.hostname + or parsed_url.username + or parsed_url.password + ): + raise ValueError("endpoint must be an HTTP(S) serving URL without embedded credentials") + if not tokenizer.is_dir(): + raise ValueError("tokenizer must be an existing local directory for the pinned checkpoint") + for source in (trace, prediction_config, tokenizer): + if source == output or output in source.parents: + raise ValueError("serving output must not contain an input") + prediction = yaml.safe_load(prediction_config.read_text(encoding="utf-8")) + scope = _scope(prediction, trace) + play, issues = _read_play(trace) + timing_subset = _timing_subset(play) if play is not None and not issues else None + if play is not None and not issues and timing_subset is None: + issues.append( + "this play has both positive recorded API times and positive idle gaps: AIS replay enforces " + "both absolute and completion bounds, but pinned AIPerf cannot combine them. Choose a complete " + "representative play with zero/absent predecessor API times or all zero idle gaps; do not edit " + "its timestamps to force qualification. Broader trace timing alignment remains required." + ) + prediction_report = (prediction_report or prediction_config.parent / "prediction/prediction.json").resolve() + if not prediction_report.is_file(): + issues.append("run ordinary onboarding replay with per-request capture before preparing serving validation") + if aiperf_python is None: + issues.append("provide the pinned AIPerf Python for request materialization before matched prediction") + engine = prediction["engine"] + worker = engine["workers"]["aggregated"] + configured = {key: engine.get(key) for key in ("model", "backend", "backend_version", "hardware", "context_length")} + configured.update( + model_revision=engine.get("fpm_profile", {}).get("model_revision"), + parallelism=worker.get("parallelism"), + scheduler=worker.get("scheduler"), + gpu_memory_utilization=worker.get("kv_cache", {}).get("capacity", {}).get("memory_fraction"), + ) + for field, value in configured.items(): + if value is not None and expected_execution.get(field) != value: + issues.append(f"expected execution {field} differs from ordinary prediction configuration") + for field in EXECUTION_FIELDS: + if expected_execution.get(field) in (None, "", {}, []): + issues.append(f"expected collection execution has no resolved {field}; inspect all phases and ranks") + tokenizer_files = [ + _identity(path) + for path in sorted(tokenizer.rglob("*")) + if path.is_file() + and path.suffix in {".json", ".model", ".txt", ".tiktoken", ".jinja", ".jinja2", ".vocab", ".bpe"} + ] + if not tokenizer_files: + issues.append("the pinned tokenizer directory has no tokenizer/config artifacts") + timing_paths = { + path.resolve() + for root in engine.get("systems_paths", []) + for path in Path(root).expanduser().rglob("fpm_forward_perf.*") + if path.is_file() + } + external = worker["timing"].get("estimator_config", {}).get("fpm_interpolation", {}).get("fpm_parquet_path") + if external: + timing_paths.update((Path(external).resolve(), Path(external).with_suffix(".metadata.json").resolve())) + if not timing_paths: + issues.append("the prediction's collected FPM timing artifacts could not be located") + output.mkdir(parents=True, exist_ok=False) + observer_dir = output / "observer" + observer_dir.mkdir() + observer_files = [] + runtime = files("collector.fpm_forward.runtime") + for name in ("fpm_memory_observer.py", "fpm_memory_worker.py"): + target = observer_dir / name + target.write_bytes(runtime.joinpath(name).read_bytes()) + observer_files.append(_identity(target)) + benchmark_trace = output / "play.json" + if play is not None: + _write(benchmark_trace, play) + recipe = { + "schema_version": SCHEMA, + "status": "incomplete" if issues else "ready", + "issues": issues, + "trace": _identity(trace), + "prediction_config": _identity(prediction_config), + "prediction_report": _identity(prediction_report) if prediction_report.is_file() else None, + "benchmark_trace": _identity(benchmark_trace) if play is not None else None, + "scope": scope, + "play_id": play.get("id") if play else None, + "timing_subset": timing_subset, + "request_count": len(play["requests"]) if play else 0, + "aiperf": {"repository": AIPERF_SOURCE, "revision": AIPERF_REVISION}, + "endpoint": endpoint, + "tokenizer": str(tokenizer), + "tokenizer_files": tokenizer_files, + "timing_artifacts": [_identity(path) for path in sorted(timing_paths)], + "expected_execution": expected_execution, + "policy": { + "latency_p95_relative_error": _threshold(latency_p95_error, "latency_p95_error"), + "throughput_relative_error": _threshold(throughput_error, "throughput_error"), + "max_dispatch_delay_ms": _number(max_dispatch_delay_ms, "max_dispatch_delay_ms"), + "max_dispatch_delay_fraction": _number(max_dispatch_delay_fraction, "max_dispatch_delay_fraction"), + }, + "artifacts_dir": str(output / "aiperf"), + "serving_observer": { + "module_files": observer_files, + "worker_class": "fpm_memory_worker.FpmExecutionWorker", + "provenance_path": str(output / "serving-provenance.json"), + "output_dir": str(output / "runtime-observations"), + "launch_arguments": ["--worker-cls", "fpm_memory_worker.FpmExecutionWorker"], + "environment": { + "PYTHONPATH": str(observer_dir), + "FPM_EXECUTION_OUTPUT_DIR": str(output / "runtime-observations"), + "FPM_EXECUTION_PROVENANCE_FILE": str(output / "serving-provenance.json"), + }, + "instructions": "Add these arguments/environment to the reviewed ordinary vLLM serving launch. " + "Make the observer files and provenance available at these absolute paths inside every worker container. " + "Prepend the observer directory to any existing PYTHONPATH. Start a fresh server or reset its prefix " + "cache after initialization. A server started without the hook must be relaunched or remain incomplete. " + "The hook records selected attention groups, effective graphs and resolved runtime configuration. " + "Retain the scheduler/container launch record, mounted checkpoint revision and nvidia-smi output " + "as separate identity evidence; the observer does not cryptographically verify model weights or images.", + }, + } + if not issues: + try: + recipe["matched_workload"] = _prepare_matched_workload( + recipe, prediction, output / "matched", aiperf_python + ) + except (OSError, ValueError, RuntimeError) as error: + issues.append(str(error)) + recipe["status"] = "incomplete" if issues else "ready" + recipe["command_arguments"] = _benchmark_arguments(recipe) if not issues else None + path = output / "recipe.json" + _write(path, recipe) + _write( + output / "serving-provenance.json", + { + "schema_version": SCHEMA, + "run_id": str(uuid.uuid4()), + "purpose": "matched_serving_accuracy", + "recipe": _identity(path), + "trace": recipe["trace"], + "observer_modules": observer_files, + }, + ) + _write( + output / "observed-execution.template.json", + { + "schema_version": SCHEMA, + "recipe": _identity(path), + "observed_execution": dict.fromkeys(EXECUTION_FIELDS), + "scope": dict.fromkeys(scope), + "sources": [], + "note": "Populate from serving initialization and benchmark evidence; missing values never pass. " + "Do not copy expected values as observations. Record native attention groups and effective graphs. " + "Benchmark real_prefix/real_kv preparation is separate from serving's actual request history.", + }, + ) + return path + + +def _checked_recipe(path: Path) -> dict[str, Any]: + recipe = json.loads(path.read_text(encoding="utf-8")) + if recipe.get("schema_version") != SCHEMA: + raise ValueError("unsupported serving validation recipe schema") + if recipe.get("status") != "ready": + raise ValueError("serving validation preparation is incomplete: " + "; ".join(recipe.get("issues", []))) + if recipe.get("aiperf") != {"repository": AIPERF_SOURCE, "revision": AIPERF_REVISION}: + raise ValueError("serving recipe must use the qualified AIPerf source pin") + if recipe.get("command_arguments") != _benchmark_arguments(recipe): + raise ValueError("serving command differs from the fixed producer recipe") + for key in ("trace", "prediction_config", "prediction_report", "benchmark_trace"): + if _identity(Path(recipe[key]["path"])) != recipe[key]: + raise ValueError(f"serving recipe {key} changed after preparation") + for source in [*recipe["tokenizer_files"], *recipe["timing_artifacts"]]: + if _identity(Path(source["path"])) != source: + raise ValueError("pinned tokenizer or FPM timing files changed after serving preparation") + matched = recipe["matched_workload"] + for key in ( + "helper", + "input", + "payloads", + "tokenization", + "trace", + "prediction_config", + "prediction_report", + "coverage_report", + ): + if _identity(Path(matched[key]["path"])) != matched[key]: + raise ValueError(f"matched serving {key} changed after preparation") + for source in recipe["serving_observer"]["module_files"]: + if _identity(Path(source["path"])) != source: + raise ValueError("serving observer module changed after preparation") + provenance = json.loads(Path(recipe["serving_observer"]["provenance_path"]).read_text(encoding="utf-8")) + if ( + provenance.get("recipe") != _identity(path) + or provenance.get("observer_modules") != recipe["serving_observer"]["module_files"] + ): + raise ValueError("serving observer provenance does not match this recipe and its modules") + return recipe + + +def _tokenize_serving(recipe_path: Path, recipe: dict[str, Any]) -> Path: + """Check actual server rendering, without running a model or warming KV.""" + frozen = json.loads(Path(recipe["matched_workload"]["tokenization"]["path"]).read_text()) + payloads = [ + json.loads(row) for row in Path(recipe["matched_workload"]["payloads"]["path"]).read_text().splitlines() + ] + endpoint = recipe["endpoint"].rstrip("/").removesuffix("/v1") + "/tokenize" + rows = [] + for index, (request, tokens) in enumerate(zip(payloads, frozen["requests"], strict=True)): + payload = request["payload"] + body = {key: payload[key] for key in ("model", "messages", "add_generation_prompt", "continue_final_message")} + body["add_special_tokens"] = False + http = Request( + endpoint, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}, method="POST" + ) + with urlopen(http, timeout=60) as response: + actual = json.loads(response.read()) + rows.append({"turn_index": index, "request": body, "response": actual}) + if actual.get("tokens") != tokens["input_token_ids"] or actual.get("count") != tokens["input_length"]: + path = recipe_path.parent / "server-tokenization.json" + _write(path, {"recipe": _identity(recipe_path), "endpoint": endpoint, "requests": rows, "status": "failed"}) + raise ValueError(f"serving tokenizer/chat template differs at turn {index}; inspect {path}") + path = recipe_path.parent / "server-tokenization.json" + _write(path, {"recipe": _identity(recipe_path), "endpoint": endpoint, "requests": rows, "status": "matched"}) + return path + + +def run_serving_benchmark(recipe_path: Path, *, aiperf_python: Path) -> Path: + """Run only the fixed AIPerf producer against a caller-owned ready endpoint. + + Install AIPerf in a separate environment from the pinned Git URL. The PEP 610 + check rejects a mutable or unrelated installation before sending traffic. + """ + recipe_path = recipe_path.expanduser().resolve() + recipe = _checked_recipe(recipe_path) + output = recipe_path.parent + receipt_path = output / "benchmark-execution.json" + if receipt_path.exists() or Path(recipe["artifacts_dir"]).exists(): + raise ValueError("use a fresh serving validation directory for each measurement") + runtime_sources = [ + _identity(path) + for path in sorted(Path(recipe["serving_observer"]["output_dir"]).glob("fpm-execution-worker-*.json")) + ] + if not runtime_sources: + raise ValueError("start the ordinary serving worker with the prepared observer hook before benchmarking") + installation = _aiperf_installation(aiperf_python) + command = [str(aiperf_python), *_benchmark_arguments(recipe)] + receipt = { + "schema_version": SCHEMA, + "recipe": _identity(recipe_path), + "installation": installation, + "command": command, + "status": "incomplete", + "exit_code": None, + "serving_provenance": _identity(Path(recipe["serving_observer"]["provenance_path"])), + "observer_modules": recipe["serving_observer"]["module_files"], + "runtime_observations": runtime_sources, + } + missing, failures = [], [] + _runtime_observations(recipe, receipt, {"observed_execution": recipe["expected_execution"]}, missing, failures) + if missing or failures: + raise ValueError("serving worker does not match the prepared execution: " + "; ".join(missing + failures)) + receipt["server_tokenization"] = _identity(_tokenize_serving(recipe_path, recipe)) + _write(receipt_path, receipt) + # Preserve normal credentials/proxy access; prevent hidden AIPerf settings + # from changing the reviewed traffic or prefix reconstruction policy. + environment = {key: value for key, value in os.environ.items() if not key.startswith("AIPERF_")} + with (output / "aiperf.stdout.log").open("w") as stdout, (output / "aiperf.stderr.log").open("w") as stderr: + completed = subprocess.run(command, env=environment, stdout=stdout, stderr=stderr, check=False) + receipt.update(exit_code=completed.returncode, status="completed" if completed.returncode == 0 else "failed") + receipt["artifacts"] = [ + _identity(path) for path in sorted(Path(recipe["artifacts_dir"]).rglob("profile_export*")) if path.is_file() + ] + _write(receipt_path, receipt) + _checked_recipe(recipe_path) + return receipt_path + + +def _metric(row: dict[str, Any], name: str, unit: str) -> float: + metric = row.get("metrics", {}).get(name) + if not isinstance(metric, dict) or metric.get("unit") != unit: + raise ValueError(f"AIPerf {name} is missing or has a different unit from {unit}") + return _number(metric.get("value"), f"AIPerf {name}") + + +def _percentile(values: list[float], quantile: float) -> float: + ordered = sorted(values) + index = (len(ordered) - 1) * quantile + lower = math.floor(index) + return ordered[lower] + (ordered[math.ceil(index)] - ordered[lower]) * (index - lower) + + +def _error(predicted: float, measured: float) -> float: + if measured <= 0: + raise ValueError("relative timing error requires a positive measured value") + return abs(predicted - measured) / measured + + +def _compare_samples(samples: list[dict[str, Any]], metric: str, threshold: float) -> dict[str, Any]: + errors = [row[f"{metric}_relative_error"] for row in samples if row.get(f"{metric}_relative_error") is not None] + if not errors: + return {"status": "incomplete", "sample_count": 0, "issue": f"no valid {metric} samples"} + p95 = _percentile(errors, 0.95) + return { + "status": "passed" if p95 <= threshold else "failed", + "sample_count": len(errors), + "p95_relative_error": p95, + "max_relative_error": max(errors), + "threshold": threshold, + } + + +def _verify_sources(sources: Any, issues: list[str], failures: list[str]) -> None: + if not isinstance(sources, list) or not sources: + issues.append("observed execution needs hashed serving initialization evidence") + return + for source in sources: + try: + if _identity(Path(source["path"])) != source: + failures.append(f"execution source changed: {source.get('path')}") + except (OSError, KeyError, TypeError): + issues.append("an execution evidence source is missing or malformed") + + +def _runtime_observations( + recipe: dict[str, Any], receipt: dict[str, Any], observed: dict[str, Any], missing: list[str], failures: list[str] +) -> list[dict[str, Any]]: + sources = receipt.get("runtime_observations", []) + if not sources: + missing.append("native serving worker observations are missing; relaunch with the prepared observer hook") + return [] + _verify_sources(sources, missing, failures) + provenance_path = Path(recipe["serving_observer"]["provenance_path"]) + provenance = json.loads(provenance_path.read_text()) + if receipt.get("serving_provenance") != _identity(provenance_path): + failures.append("benchmark receipt is not bound to the prepared serving provenance") + if receipt.get("observer_modules") != recipe["serving_observer"]["module_files"]: + failures.append("benchmark receipt is not bound to the prepared observer modules") + ranks = set() + for source in sources: + try: + payload = json.loads(Path(source["path"]).read_text()) + if payload.get("schema_name") != "aisimulate_fpm_runtime_execution" or payload.get("schema_version") != 1: + raise ValueError("unknown native serving observation schema") + if payload.get("status") != "observed": + raise ValueError("serving runtime did not resolve selected attention groups and graph configuration") + if payload.get("execution_provenance") != provenance or payload.get("provenance_source") != _identity( + provenance_path + ): + failures.append("native serving observations belong to a different prepared server launch") + rank = tuple(payload.get(key) for key in ("dp_rank", "tp_rank", "pp_rank")) + if any(type(value) is not int or value < 0 for value in rank) or rank in ranks: + failures.append("native serving observation has a duplicate or invalid worker rank") + ranks.add(rank) + for key in ("attention_groups", "graph_config", "backend_version"): + actual = observed.get("observed_execution", {}).get(key) + if actual is not None and payload.get(key) != actual: + failures.append(f"reported serving {key} differs from the native initialized worker") + native_config = payload.get("resolved_config") + if not isinstance(native_config, dict) or not native_config: + missing.append("native serving worker resolved runtime configuration is missing") + else: + prefix = native_config.get("cache_config", {}).get("enable_prefix_caching") + if not isinstance(prefix, bool): + missing.append("native serving prefix cache setting is unresolved") + elif prefix != recipe["scope"]["prefix_caching"]: + failures.append("native serving prefix cache setting differs from the replay scope") + actual = observed.get("observed_execution", {}).get("runtime_config") + if actual is not None and normalize_runtime_config(native_config) != actual: + failures.append("reported serving runtime_config differs from the native initialized worker") + except (OSError, ValueError, TypeError) as error: + missing.append(f"serving worker observation is unavailable: {error}") + parallel = recipe["expected_execution"].get("parallelism", {}) + expected = { + (dp, tp, pp) + for dp in range(parallel.get("attention_data", 1)) + for tp in range(parallel.get("tensor", 1)) + for pp in range(parallel.get("pipeline", 1)) + } + if ranks != expected: + missing.append("native serving observations do not cover the selected worker ranks exactly") + return sources + + +def _check_server_tokenization( + recipe_path: Path, recipe: dict[str, Any], receipt: dict[str, Any], missing: list[str], failures: list[str] +) -> None: + source = receipt.get("server_tokenization") + if source is None: + missing.append("actual serving /tokenize evidence is missing") + return + try: + if _identity(Path(source["path"])) != source: + failures.append("actual serving tokenization evidence changed") + evidence = json.loads(Path(source["path"]).read_text()) + if evidence.get("recipe") != _identity(recipe_path) or evidence.get("status") != "matched": + failures.append("serving tokenization does not match the prepared workload") + frozen = json.loads(Path(recipe["matched_workload"]["tokenization"]["path"]).read_text()) + payloads = [ + json.loads(row) for row in Path(recipe["matched_workload"]["payloads"]["path"]).read_text().splitlines() + ] + if len(evidence.get("requests", [])) != len(frozen["requests"]): + missing.append("serving tokenization did not cover every prepared request") + return + for index, (actual, expected, payload) in enumerate( + zip(evidence["requests"], frozen["requests"], payloads, strict=True) + ): + request = { + key: payload["payload"][key] + for key in ("model", "messages", "add_generation_prompt", "continue_final_message") + } + request["add_special_tokens"] = False + if ( + actual.get("turn_index") != index + or actual.get("request") != request + or actual.get("response", {}).get("tokens") != expected["input_token_ids"] + or actual.get("response", {}).get("count") != expected["input_length"] + ): + failures.append(f"serving tokenization differs from frozen payload at turn {index}") + except (OSError, ValueError, KeyError, TypeError) as error: + missing.append(f"serving tokenization evidence is unavailable: {error}") + + +def _forward_validation(path: Path | None, recipe_path: Path, threshold: float) -> dict[str, Any]: + if path is None: + return {"status": "unavailable", "issue": "no independently instrumented serving forward timings supplied"} + try: + evidence = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + return {"status": "incomplete", "issue": f"forward evidence is unavailable: {error}"} + missing, failures = [], [] + _verify_sources(evidence.get("sources"), missing, failures) + if evidence.get("recipe") != _identity(recipe_path): + failures.append("forward evidence belongs to a different serving recipe") + from aisimulate_core.sdk import RustForwardPassPerfModel + + recipe = _checked_recipe(recipe_path) + prediction = yaml.safe_load(Path(recipe["prediction_config"]["path"]).read_text()) + engine = prediction["engine"] + worker = engine["workers"]["aggregated"] + parallel = worker["parallelism"] + config = { + "model": engine["model"], + "system": engine["hardware"], + "backend": engine["backend"], + "backend_version": engine["backend_version"], + "worker_type": "aggregated", + "tp": parallel.get("tensor", 1), + "pp": parallel.get("pipeline", 1), + "attention_dp": parallel.get("attention_data", 1), + "moe_tp_size": parallel.get("moe_tensor", 1), + "moe_ep_size": parallel.get("moe_expert", 1), + "fpm_profile": engine["fpm_profile"], + "systems_paths": engine["systems_paths"], + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": worker["timing"]["estimator_config"], + } + try: + model = RustForwardPassPerfModel.best_available(config) + except (ValueError, RuntimeError) as error: + return { + "status": "incomplete", + "issue": f"forward timing construction failed: {error}", + "evidence": _identity(path), + } + phases: dict[str, Any] = {} + try: + for phase in ("prefill", "decode"): + rows = evidence.get("samples", {}).get(phase, []) + if not rows: + phases[phase] = { + "status": "unavailable", + "sample_count": 0, + "issue": "no serving instrumentation for this phase", + } + continue + samples = [] + unsupported = [] + for row in rows: + try: + if row.get("unit") != "ms" or not isinstance(row.get("scheduled_requests"), dict): + raise ValueError("forward samples require native scheduled_requests coordinates and ms units") + measured = _number(row.get("measured_ms"), "measured forward latency", positive=True) + scheduled = row["scheduled_requests"] + if ( + phase == "prefill" + and ( + scheduled.get("num_prefill_requests", 0) < 1 or scheduled.get("num_decode_requests", 0) != 0 + ) + ) or ( + phase == "decode" + and ( + scheduled.get("num_decode_requests", 0) < 1 or scheduled.get("num_prefill_requests", 0) != 0 + ) + ): + raise ValueError("forward phase does not match scheduled request coordinates") + predicted = model.estimate_forward_pass_time_ms({"scheduled_requests": scheduled}) + predicted = _number(predicted, "native direct-FPM forward estimate", positive=True) + samples.append( + {**row, "predicted_ms": predicted, "forward_relative_error": _error(predicted, measured)} + ) + except (ValueError, RuntimeError) as error: + unsupported.append({"scheduled_requests": row.get("scheduled_requests"), "reason": str(error)}) + phases[phase] = { + **_compare_samples(samples, "forward", threshold), + "samples": samples, + "unsupported": unsupported, + } + if unsupported: + missing.append(f"{phase} forward timing queries were unsupported or invalid") + finally: + model.close() + statuses = [phase["status"] for phase in phases.values()] + if all(status == "unavailable" for status in statuses): + missing.append("supplied forward evidence contains no measured phase samples") + status = ( + "failed" + if failures or "failed" in statuses + else "incomplete" + if missing or "incomplete" in statuses + else "passed" + ) + return { + "status": status, + "phases": phases, + "missing_evidence": missing, + "failures": failures, + "evidence": _identity(path), + } + + +def validate_serving_measurements( + recipe_path: Path, + *, + prediction_report: Path, + execution_evidence_path: Path, + output_report: Path | None, + forward_evidence_path: Path | None = None, + latency_p95_error: float | None = None, + throughput_error: float | None = None, + max_dispatch_delay_ms: float | None = None, + max_dispatch_delay_fraction: float | None = None, +) -> dict[str, Any]: + """Compare real AIPerf records with ordinary prediction; retain all request evidence.""" + recipe_path = recipe_path.expanduser().resolve() + recipe = _checked_recipe(recipe_path) + policy = dict(recipe["policy"]) + if latency_p95_error is not None: + policy["latency_p95_relative_error"] = _threshold(latency_p95_error, "latency_p95_error") + if throughput_error is not None: + policy["throughput_relative_error"] = _threshold(throughput_error, "throughput_error") + for key, value in ( + ("max_dispatch_delay_ms", max_dispatch_delay_ms), + ("max_dispatch_delay_fraction", max_dispatch_delay_fraction), + ): + if value is not None: + policy[key] = _number(value, key) + if output_report is not None and output_report.exists(): + raise ValueError("save each serving assessment to a fresh report path; raw measurements can be reused") + missing, failures = [], [] + report: dict[str, Any] = { + "schema_version": SCHEMA, + "status": "incomplete", + "recipe": _identity(recipe_path), + "scope": recipe["scope"], + "timing_subset": recipe["timing_subset"], + "matched_workload": recipe["matched_workload"], + "policy": policy, + "missing_evidence": missing, + "failures": failures, + "definitions": { + "ttft_ms": "first token minus request dispatch; serving is client observed, simulation is engine observed", + "tpot_ms": "per-request mean inter-token latency; only responses with at least two tokens", + "output_throughput_tokens_per_second": ( + "total output tokens / (last response - first dispatch), including inter-turn delays" + ), + "relative_error": "abs(simulated - measured) / measured; p95 uses linear empirical quantile", + "kv_preparation": "benchmark seeds KV history; serving populates actual request history", + }, + "samples": [], + } + try: + prediction = json.loads( + Path(recipe["matched_workload"]["prediction_report"]["path"]).read_text(encoding="utf-8") + ) + observed = json.loads(execution_evidence_path.read_text(encoding="utf-8")) + receipt_path = recipe_path.parent / "benchmark-execution.json" + receipt = json.loads(receipt_path.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + missing.append(f"required prediction, execution receipt or serving observation is unavailable: {error}") + if output_report is not None: + _write(output_report, report) + return report + report["inputs"] = [_identity(path) for path in (prediction_report, execution_evidence_path, receipt_path)] + if _identity(prediction_report) != recipe["prediction_report"]: + failures.append("prediction report is not the ordinary replay bound to the prepared recipe") + if observed.get("recipe") != report["recipe"] or receipt.get("recipe") != report["recipe"]: + failures.append("serving observations or benchmark receipt belong to a different prepared recipe") + if receipt.get("status") != "completed" or receipt.get("exit_code") != 0: + missing.append("AIPerf benchmark did not complete successfully") + if receipt.get("command", [None])[1:] != _benchmark_arguments(recipe): + failures.append("executed benchmark command differs from the prepared workload") + if receipt.get("installation", {}).get("source", {}).get("vcs_info", {}).get("commit_id") != AIPERF_REVISION: + missing.append("executed AIPerf installation is not bound to the qualified source revision") + _verify_sources(observed.get("sources"), missing, failures) + report["runtime_observations"] = _runtime_observations(recipe, receipt, observed, missing, failures) + _check_server_tokenization(recipe_path, recipe, receipt, missing, failures) + for field in EXECUTION_FIELDS: + expected = recipe["expected_execution"].get(field) + actual = observed.get("observed_execution", {}).get(field) + if expected in (None, "", {}, []) or actual in (None, "", {}, []): + missing.append(f"missing expected or observed effective execution: {field}") + elif expected != actual: + failures.append(f"effective serving execution differs from collection: {field}") + for field, expected in recipe["scope"].items(): + actual = observed.get("scope", {}).get(field) + if actual is None: + missing.append(f"missing observed serving scope: {field}") + elif actual != expected: + failures.append(f"serving and simulated workload scope differ: {field}") + artifacts = receipt.get("artifacts", []) + summaries = [item for item in artifacts if Path(item.get("path", "")).name == "profile_export_aiperf.json"] + if len(summaries) != 1: + missing.append("one native AIPerf profile_export_aiperf.json summary is required") + else: + try: + source = summaries[0] + if _identity(Path(source["path"])) != source: + failures.append("AIPerf summary changed after the producer run") + summary = json.loads(Path(source["path"]).read_text(encoding="utf-8")) + if summary.get("was_cancelled") is not False or summary.get("is_complete") is False: + missing.append("native AIPerf summary did not report an uncancelled, complete run") + if summary.get("request_count", {}).get("avg") != recipe["request_count"]: + failures.append("native AIPerf summary request count differs from the selected play") + if summary.get("error_request_count", {}).get("avg", 0) != 0: + failures.append("native AIPerf summary includes failed requests") + except (OSError, ValueError, KeyError, TypeError) as error: + missing.append(f"AIPerf summary is unavailable: {error}") + record_files = [item for item in artifacts if Path(item.get("path", "")).name == "profile_export.jsonl"] + if len(record_files) != 1: + missing.append("one native AIPerf profile_export.jsonl file is required") + measured_rows = [] + else: + record_file = record_files[0] + try: + if _identity(Path(record_file["path"])) != record_file: + failures.append("AIPerf per-request records changed after the producer run") + measured_rows = [ + json.loads(line) for line in Path(record_file["path"]).read_text().splitlines() if line.strip() + ] + except (OSError, ValueError) as error: + missing.append(f"AIPerf per-request records are unavailable: {error}") + measured_rows = [] + predicted_rows = prediction.get("per_request", []) + count = recipe["request_count"] + if prediction.get("completed_requests") != count or len(predicted_rows) != count: + missing.append("ordinary prediction did not complete every selected request") + if len(measured_rows) != count: + missing.append("serving did not produce exactly one record for every selected request") + predicted_by_turn = {} + for row in predicted_rows: + match = re.search(r":request:outer:(\d+)$", str(row.get("request_id", ""))) + if match is None or int(match[1]) in predicted_by_turn: + failures.append("prediction request identity is missing, repeated or outside single-stream Weka scope") + else: + predicted_by_turn[int(match[1])] = row + play = json.loads(Path(recipe["benchmark_trace"]["path"]).read_text(encoding="utf-8")) + tokenization = json.loads(Path(recipe["matched_workload"]["tokenization"]["path"]).read_text()) + seen = set() + starts_ns, ends_ns, starts_ms, ends_ms = [], [], [], [] + for row in measured_rows: + metadata = row.get("metadata", {}) + turn = metadata.get("turn_index") + if ( + metadata.get("conversation_id") != recipe["play_id"] + or type(turn) is not int + or turn < 0 + or turn >= count + or turn in seen + or turn not in predicted_by_turn + ): + failures.append("AIPerf request identity is missing, repeated or does not match the selected play") + continue + seen.add(turn) + simulated = predicted_by_turn[turn] + if ( + row.get("error") is not None + or metadata.get("was_cancelled") is not False + or metadata.get("benchmark_phase") != "profiling" + or simulated.get("terminal_status") != "completed" + ): + failures.append(f"request {turn} failed, was cancelled or includes non-profile traffic") + continue + try: + isl, osl = _metric(row, "input_sequence_length", "tokens"), _metric(row, "output_sequence_length", "tokens") + if ( + isl != tokenization["requests"][turn]["input_length"] + or osl != play["requests"][turn]["out"] + or simulated.get("input_length") != isl + or simulated.get("output_length") != osl + or simulated.get("requested_output_length") != osl + ): + failures.append(f"request {turn} input/output token counts do not match trace and prediction") + continue + measured_ttft = _metric(row, "time_to_first_token", "ms") + predicted_ttft = _number(simulated.get("ttft_ms"), "simulated TTFT", positive=True) + start_ns = metadata.get("request_start_ns") + end_ns = metadata.get("request_end_ns") + if type(start_ns) is not int or type(end_ns) is not int or start_ns < 0 or end_ns <= start_ns: + raise ValueError("AIPerf request timestamps must be ordered integer nanoseconds") + start_ms = _number(simulated.get("arrival_time_ms"), "simulated arrival") + end_ms = _number(simulated.get("last_token_ms"), "simulated final token", positive=True) + if end_ms <= start_ms: + raise ValueError("simulated request timing interval is invalid") + sample = { + "turn_index": turn, + "prediction_request_id": simulated["request_id"], + "serving_request_id": metadata.get("x_request_id"), + "input_tokens": int(isl), + "source_trace_input_tokens": play["requests"][turn]["in"], + "output_tokens": int(osl), + "measured_ttft_ms": measured_ttft, + "predicted_ttft_ms": predicted_ttft, + "ttft_relative_error": _error(predicted_ttft, measured_ttft), + "measured_tpot_ms": None, + "predicted_tpot_ms": None, + "tpot_relative_error": None, + "request_start_ns": start_ns, + "request_end_ns": end_ns, + "simulated_arrival_ms": start_ms, + "simulated_last_token_ms": end_ms, + } + if osl > 1: + measured_tpot = _metric(row, "inter_token_latency", "ms") + predicted_tpot = _number(simulated.get("itl_ms"), "simulated TPOT", positive=True) + sample.update( + measured_tpot_ms=measured_tpot, + predicted_tpot_ms=predicted_tpot, + tpot_relative_error=_error(predicted_tpot, measured_tpot), + ) + report["samples"].append(sample) + starts_ns.append(start_ns) + ends_ns.append(end_ns) + starts_ms.append(start_ms) + ends_ms.append(end_ms) + except ValueError as error: + missing.append(f"request {turn}: {error}") + metrics = { + name: _compare_samples(report["samples"], name, policy["latency_p95_relative_error"]) + for name in ("ttft", "tpot") + } + if len(report["samples"]) == count and count > 0: + ordered_samples = sorted(report["samples"], key=lambda item: item["turn_index"]) + first_start_ns = ordered_samples[0]["request_start_ns"] + source_start = play["requests"][0]["t"] + for previous, current in pairwise(ordered_samples): + source_previous = play["requests"][previous["turn_index"]] + source_current = play["requests"][current["turn_index"]] + expected_gap_ms = (source_current["t"] - source_previous["t"] - source_previous.get("api_time", 0)) * 1000 + not_before_ns = first_start_ns + round((source_current["t"] - source_start) * 1e9) + ready_ns = max(not_before_ns, previous["request_end_ns"] + round(expected_gap_ms * 1e6)) + schedule_error_ms = (current["request_start_ns"] - ready_ns) / 1e6 + current["load_timing"] = { + "not_before_ns": not_before_ns, + "completion_gap_ms": expected_gap_ms, + "expected_ready_ns": ready_ns, + "dispatch_error_ms": schedule_error_ms, + } + if abs(schedule_error_ms) > max( + policy["max_dispatch_delay_ms"], expected_gap_ms * policy["max_dispatch_delay_fraction"] + ): + failures.append( + f"serving dispatch at turn {current['turn_index']} differs from trace bounds; " + "check producer load policy or client scheduling lag" + ) + tokens = sum(row["output_tokens"] for row in report["samples"]) + measured_rate = tokens / ((max(ends_ns) - min(starts_ns)) / 1e9) + predicted_rate = tokens / ((max(ends_ms) - min(starts_ms)) / 1000) + relative_error = _error(predicted_rate, measured_rate) + limit = policy["throughput_relative_error"] + metrics["output_throughput"] = { + "status": "passed" if relative_error <= limit else "failed", + "unit": "tokens/second", + "measured": measured_rate, + "predicted": predicted_rate, + "relative_error": relative_error, + "threshold": limit, + } + else: + metrics["output_throughput"] = {"status": "incomplete", "issue": "complete matched requests are required"} + report["metrics"] = metrics + report["forward"] = _forward_validation(forward_evidence_path, recipe_path, policy["latency_p95_relative_error"]) + statuses = [metric["status"] for metric in metrics.values()] + if failures or "failed" in statuses or report["forward"]["status"] == "failed": + report["status"] = "failed" + elif not missing and all(status == "passed" for status in statuses) and report["forward"]["status"] != "incomplete": + report["status"] = "passed" + _checked_recipe(recipe_path) + if output_report is not None: + _write(output_report, report) + return report + + +def check_serving_validation_report(path: Path) -> dict[str, Any]: + """Reassess raw evidence, including source hashes, before accepting a saved result.""" + report = json.loads(path.read_text(encoding="utf-8")) + try: + recipe = report["recipe"] + if _identity(Path(recipe["path"])) != recipe: + raise ValueError("serving recipe changed since assessment") + for source in report["inputs"]: + if _identity(Path(source["path"])) != source: + raise ValueError("serving validation input changed since assessment") + evidence = report.get("forward", {}).get("evidence") + if evidence is not None and _identity(Path(evidence["path"])) != evidence: + raise ValueError("serving forward evidence changed since assessment") + reassessed = validate_serving_measurements( + Path(recipe["path"]), + prediction_report=Path(report["inputs"][0]["path"]), + execution_evidence_path=Path(report["inputs"][1]["path"]), + output_report=None, + forward_evidence_path=Path(evidence["path"]) if evidence else None, + latency_p95_error=report["policy"]["latency_p95_relative_error"], + throughput_error=report["policy"]["throughput_relative_error"], + max_dispatch_delay_ms=report["policy"]["max_dispatch_delay_ms"], + max_dispatch_delay_fraction=report["policy"]["max_dispatch_delay_fraction"], + ) + except (KeyError, IndexError, TypeError) as error: + raise ValueError("serving validation report lacks required evidence references") from error + if reassessed != report: + raise ValueError("saved serving validation result differs from independent reassessment of its raw evidence") + return report diff --git a/python/aisimulate/src/aisimulate/support/serving_workload.py b/python/aisimulate/src/aisimulate/support/serving_workload.py new file mode 100644 index 000000000..4248c187e --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/serving_workload.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Freeze native AIPerf Weka chat requests in its isolated, pinned environment. + +Executed as a fixed package resource script with a structured JSON input. No +AISimulate imports or GPU model are needed in the AIPerf environment. Upstream +Weka reconstruction supplies content/history, and upstream Session/ChatEndpoint +supplies the complete request payload. Native Mooncake replay then sends that +payload verbatim with the original timestamp and end-to-start delay. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + + +def materialize(input_path: Path) -> None: + from aiperf.cli import app + from aiperf.common import random_generator + from aiperf.common.tokenizer import Tokenizer + from aiperf.config.flags.converter import convert_cli_to_aiperf + from aiperf.dataset.generator.prompt import PromptGenerator + from aiperf.dataset.loader.weka_trace import WekaTraceLoader + from aiperf.endpoints.openai_chat import ChatEndpoint + from aiperf.workers.session_manager import UserSession + from transformers import AutoTokenizer + + inputs = json.loads(input_path.read_text()) + trace = json.loads(Path(inputs["trace"]).read_text()) + output = Path(inputs["output"]) + # Parsing the real producer options preserves the native loader's options. + _, bound, _ = app.parse_args(inputs["arguments"], exit_on_error=False, print_error=False) + config = convert_cli_to_aiperf(bound.arguments["cli_config"]).benchmark + random_generator.reset() + random_generator.init(0) + # The pinned AIPerf offline Hub resolver treats a local directory as a Hub + # ID. Load the supplied local files directly, then use its native wrapper. + inner = AutoTokenizer.from_pretrained(inputs["tokenizer"], local_files_only=True, trust_remote_code=False) + tokenizer = Tokenizer._build_with_kwargs(inner, resolved_name=inputs["tokenizer"]) + template = inner.get_chat_template() + if not isinstance(template, str) or not template: + raise ValueError("the target tokenizer must define the serving chat template") + generator = PromptGenerator(prompts=None, prefix_prompts=None, tokenizer=tokenizer) + loader = WekaTraceLoader( + filename=inputs["trace"], + run=SimpleNamespace(cfg=config, random_seed=0), + prompt_generator=generator, + ) + conversations = loader.convert_to_conversations(loader.load_dataset()) + if len(conversations) != 1 or len(conversations[0].turns) != len(trace["requests"]): + raise ValueError("native AIPerf reconstruction changed the selected play's request graph") + conversation = conversations[0] + if conversation.branches: + raise ValueError("native AIPerf reconstructed branches; this complete play is not yet qualified") + session = UserSession( + x_correlation_id="preparation", + num_turns=len(conversation.turns), + conversation=conversation, + context_mode=conversation.context_mode or loader.get_default_context_mode(), + ) + if session.should_store_response(): + raise ValueError("request materialization requires prerecorded history, independent of generated responses") + endpoint_info = SimpleNamespace(primary_model_name=inputs["model"], endpoint=config.endpoint) + endpoint = ChatEndpoint(endpoint_info) + tokenization = [] + rows = [] + for index, turn in enumerate(conversation.turns): + if turn.source_outer_idx != index: + raise ValueError("native AIPerf reordered the selected play") + session.advance_turn(index) + payload = endpoint.format_payload( + SimpleNamespace( + turns=session.turn_list, + model_endpoint=endpoint_info, + system_message=conversation.system_message, + user_context_message=None, + ) + ) + # Explicit defaults make target chat tokenization reviewable. They do + # not override the server's template or enable remote template code. + payload.update(add_generation_prompt=True, continue_final_message=False) + if payload.get("ignore_eos") is not True: + raise ValueError("prepared request must force its requested output length") + tokens = inner.apply_chat_template( + payload["messages"], tokenize=True, add_generation_prompt=True, chat_template=template, return_dict=False + ) + if not isinstance(tokens, list) or not tokens or any(type(token) is not int for token in tokens): + raise ValueError("target chat tokenizer did not return an explicit nonempty token sequence") + rows.append( + { + "session_id": trace["id"], + "timestamp": turn.timestamp, + "delay": turn.delay, + "output_length": trace["requests"][index]["out"], + "payload": payload, + } + ) + tokenization.append({"turn_index": index, "input_token_ids": tokens, "input_length": len(tokens)}) + (output / "requests.jsonl").write_text("".join(json.dumps(row) + "\n" for row in rows)) + (output / "tokenization.json").write_text( + json.dumps( + { + "chat_template": template, + "chat_template_sha256": hashlib.sha256(template.encode()).hexdigest(), + "chat_template_kwargs": {"add_generation_prompt": True, "continue_final_message": False}, + "content": "native AIPerf Weka reconstruction from its packaged sonnet corpus, seed 0", + "history": "prerecorded synthetic history; generated responses do not alter future inputs", + "requests": tokenization, + }, + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + materialize(Path(sys.argv[1])) diff --git a/python/aisimulate/src/aisimulate/support/topology.py b/python/aisimulate/src/aisimulate/support/topology.py new file mode 100644 index 000000000..65db8eaa5 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/topology.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model- and hardware-aware starting points for one onboarding worker. + +These byte-budget prechecks do not predict performance or establish runtime +compatibility. The generated plan and serving runtime retain admission authority. +""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Mapping +from dataclasses import asdict, dataclass +from typing import Any, Literal + +import yaml + +from .config_profile import ( + ModelConfig, + ProfileDraft, + ProfileRequestError, + derive_profile, + packaged_hardware_path, + validate_overrides, +) +from .schema import SupportRequest + +_NON_KV_BYTES = ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes") +_RANK_BYTES = frozenset((*_NON_KV_BYTES, "kv_bytes_per_token")) +_TOPOLOGY_FIELDS = ("tensor_parallel", "attention_data_parallel", "moe_tensor_parallel", "moe_expert_parallel") +_FAMILIES = ("pure_tp", "dep", "tep") + + +@dataclass(frozen=True) +class HardwareEnvelope: + system: str + declared_interconnect: str + source: str + sha256: str + per_gpu_bytes: int + memory_budget_bytes: int + gpus_per_node: int + gpus_per_rack: int | None + fast_domain: Literal["single_gpu", "node", "rack"] + fast_domain_gpus: int + assumption: str + + +@dataclass(frozen=True) +class TopologyCandidate: + family: Literal["tp", "pure_tp", "dep", "tep"] + sizes: tuple[int, int, int, int] + status: Literal["estimated_fit", "needs_inputs", "rejected"] + reasons: tuple[str, ...] + draft: ProfileDraft | None + known_required_bytes: int | None + estimated_required_bytes: int | None + memory_budget_bytes: int + resident_sequences_per_rank: int | None + + @property + def topology(self) -> dict[str, int]: + return dict(zip(_TOPOLOGY_FIELDS, self.sizes, strict=True)) + + @property + def required_gpus(self) -> int: + return self.sizes[0] * self.sizes[1] + + @property + def cli_flags(self) -> tuple[str, ...]: + return tuple( + item for key, value in self.topology.items() for item in ("--" + key.replace("_", "-"), str(value)) + ) + + @property + def missing(self) -> dict[str, str]: + return dict(self.draft.missing) if self.draft is not None else {} + + def apply(self, request: SupportRequest) -> SupportRequest: + """Return the selected topology with the caller's other request inputs intact.""" + payload = request.model_dump() + payload["search"].update(self.topology) + return SupportRequest.model_validate(payload) + + def to_dict(self) -> dict[str, Any]: + return { + "family": self.family, + "topology": self.topology, + "required_gpus": self.required_gpus, + "cli_flags": list(self.cli_flags), + "status": self.status, + "reasons": list(self.reasons), + "missing": self.missing, + "known_required_bytes": self.known_required_bytes, + "estimated_required_bytes": self.estimated_required_bytes, + "memory_budget_bytes": self.memory_budget_bytes, + "resident_sequences_per_rank": self.resident_sequences_per_rank, + "resolved_profile_fields": dict(self.draft.resolved) if self.draft is not None else {}, + "field_sources": dict(self.draft.sources) if self.draft is not None else {}, + } + + +@dataclass(frozen=True) +class TopologySuggestions: + config_sha256: str + identity: dict[str, Any] + collection: dict[str, Any] + context_length: int + overrides: dict[str, Any] + hardware: HardwareEnvelope + candidates: tuple[TopologyCandidate, ...] + rejected_candidates: tuple[TopologyCandidate, ...] + default: TopologyCandidate | None + assumptions: tuple[str, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "config_sha256": self.config_sha256, + "identity": dict(self.identity), + "collection": dict(self.collection), + "context_length": self.context_length, + "overrides": dict(self.overrides), + "hardware": asdict(self.hardware), + "candidates": [candidate.to_dict() for candidate in self.candidates], + "rejected_candidates": [candidate.to_dict() for candidate in self.rejected_candidates], + "default": self.default.to_dict() if self.default is not None else None, + "assumptions": list(self.assumptions), + } + + +def _integer(value: Any, field: str, *, minimum: int = 1) -> int: + if type(value) is not int or not minimum <= value <= 2**53: + raise ValueError(f"packaged hardware {field} must be an integer in {minimum}..2**53") + return value + + +def _bandwidth(value: Any, field: str) -> int | float: + if type(value) not in (int, float) or not math.isfinite(value) or value <= 0: + raise ValueError(f"packaged hardware {field} must be a positive finite bandwidth in bytes/second") + return value + + +def _hardware_envelope(system: str, interconnect: str, memory_fraction: float) -> HardwareEnvelope: + path = packaged_hardware_path(system) + if not path.is_file(): + raise ValueError(f"no packaged hardware metadata for {system}; select an explicit topology") + payload = path.read_bytes() + spec = yaml.safe_load(payload) + if not isinstance(spec, dict): + raise ValueError(f"packaged hardware metadata for {system} must be an object") + for section in ("gpu", "node"): + if not isinstance(spec.get(section), dict): + raise ValueError(f"packaged hardware {section} must be an object") + capacity = _integer(spec["gpu"].get("mem_capacity"), "gpu.mem_capacity") + node = spec["node"] + node_width = _integer(node.get("num_gpus_per_node"), "node.num_gpus_per_node") + intra_bw = _bandwidth(node.get("intra_node_bw"), "node.intra_node_bw") + inter_bw = _bandwidth(node.get("inter_node_bw"), "node.inter_node_bw") + rack_width = None + if "num_gpus_per_rack" in node: + rack_width = _integer(node["num_gpus_per_rack"], "node.num_gpus_per_rack") + if rack_width < node_width or rack_width % node_width: + raise ValueError("packaged hardware node.num_gpus_per_rack must be a multiple of num_gpus_per_node") + # Missing reservation estimates remain explicit profile inputs. Malformed + # entries, including ones outside the candidate widths, are not absence. + misc = spec.get("misc", {}) + if not isinstance(misc, dict): + raise ValueError("packaged hardware misc must be an object") + if "other_mem" in misc: + _integer(misc["other_mem"], "misc.other_mem", minimum=0) + nccl = misc.get("nccl_mem", {}) + if not isinstance(nccl, dict): + raise ValueError("packaged hardware misc.nccl_mem must be a TP-indexed object") + for tp, value in nccl.items(): + _integer(tp, "misc.nccl_mem TP key") + _integer(value, f"misc.nccl_mem[{tp}]", minimum=0) + fabric = interconnect.lower().replace(" ", "").replace("-", "") + nvlink = fabric in {"nvlink", "nvswitch", "nvlink/nvswitch", "nvlink+nvswitch", "nvlinknvswitch"} + use_rack = rack_width is not None and inter_bw >= intra_bw and nvlink + domain, domain_width = "node", node_width + assumption = ( + "Use the packaged node width as the fast domain; this is hardware architecture, not available GPU capacity." + ) + if fabric == "none": + domain, domain_width = "single_gpu", 1 + assumption = "The declared interconnect is 'none'; automatic choices are limited to one GPU." + elif use_rack: + domain, domain_width = "rack", rack_width + assumption = ( + "Declared rack width is a fast domain because inter_node_bw >= intra_node_bw " + "and the request declares NVLink/NVSwitch; assumes the packaged within-rack fabric. " + "This is hardware architecture, not available GPU capacity." + ) + elif not nvlink: + assumption += ( + f" Declared fabric {interconnect!r} does not establish the packaged NVLink/NVSwitch rack; " + "fabric compatibility is unverified. Explicit wider topology choices remain available." + ) + return HardwareEnvelope( + system=system, + declared_interconnect=interconnect, + source=f"packaged systems/{system}.yaml", + sha256=hashlib.sha256(payload).hexdigest(), + per_gpu_bytes=capacity, + memory_budget_bytes=int(capacity * memory_fraction), + gpus_per_node=node_width, + gpus_per_rack=rack_width, + fast_domain=domain, + fast_domain_gpus=domain_width, + assumption=assumption, + ) + + +def _assess( + config: ModelConfig, + request: SupportRequest, + overrides: dict[str, Any], + hardware: HardwareEnvelope, + family: str, + width: int, +) -> TopologyCandidate: + sizes = { + "tp": (width, 1, 1, 1), + "pure_tp": (width, 1, width, 1), + "dep": (1, width, 1, width), + "tep": (width, 1, 1, width), + }[family] + payload = request.model_dump() + payload["search"].update(zip(_TOPOLOGY_FIELDS, sizes, strict=True)) + selected = SupportRequest.model_validate(payload) + try: + draft = derive_profile(config, selected, overrides) + except ProfileRequestError as error: + if error.field not in _TOPOLOGY_FIELDS: + raise + return TopologyCandidate( + family, sizes, "rejected", (str(error),), None, None, None, hardware.memory_budget_bytes, None + ) + resolved = draft.resolved + # Admission establishes room for one declared full-context request. The + # scheduler's sequence limit is not an allocation of full contexts; runtime + # total KV capacity is the remaining memory divided by bytes per token. + resident = 1 + known = sum(resolved[field] for field in _NON_KV_BYTES if field in resolved) + grouped = resolved.get("cache_layout") == "grouped" + no_grouped_cache_budget = grouped and known >= hardware.memory_budget_bytes + complete_memory_estimate = resolved.keys() >= set(_NON_KV_BYTES) + if grouped and draft.profile is not None and complete_memory_estimate and known < hardware.memory_budget_bytes: + try: + from aisimulate_core.sdk.rust_engine_step import RustForwardPassPerfModel + + # This temporary complete estimate is used only for the shortlist. + # The saved profile keeps memory pending until runtime observation. + estimate_profile = draft.profile.model_dump(mode="json") + estimate_profile["deployments"][0]["resources"].update({name: resolved[name] for name in _NON_KV_BYTES}) + estimate = RustForwardPassPerfModel.estimate_cache_budget( + { + "model": selected.identity.model, + "system": selected.identity.gpu, + "backend": selected.identity.framework, + "backend_version": selected.identity.framework_version, + "worker_type": selected.worker_type or "aggregated", + "tp": sizes[0], + "pp": 1, + "attention_dp": sizes[1], + "moe_tp_size": sizes[2], + "moe_ep_size": sizes[3], + "fpm_profile": estimate_profile, + }, + { + "total_gpu_capacity_bytes": hardware.per_gpu_bytes, + "memory_fraction_kind": "of_total", + "memory_fraction_value": request.collection.memory_fraction, + "max_num_tokens": resolved["max_num_tokens"], + "max_batch_size": resolved["max_batch_size"], + "context_length": request.search.context_length, + }, + ) + except (ImportError, AttributeError) as exc: + raise ValueError( + "grouped cache planning requires this checkout's compiled AISimulate extension; install or rebuild it" + ) from exc + known += estimate["request_peak_cache_bytes"] + elif "kv_bytes_per_token" in resolved: + known += resolved["kv_bytes_per_token"] * (request.search.context_length + 1) + complete_resources = ( + draft.profile is not None and complete_memory_estimate if grouped else resolved.keys() >= _RANK_BYTES + ) + estimated = known if complete_resources else None + if known > hardware.memory_budget_bytes or no_grouped_cache_budget: + status = "rejected" + bound = "Conservative grouped peak" if grouped and complete_resources else "Known per-rank resource lower bound" + reasons = ( + f"Non-KV resources {known} bytes leave no grouped cache budget ({hardware.memory_budget_bytes}-byte limit)." + if no_grouped_cache_budget + else f"{bound} {known} bytes exceeds the {request.collection.memory_fraction * 100:g}% per-GPU budget " + f"of {hardware.memory_budget_bytes} bytes.", + ) + elif draft.missing or not complete_resources: + status = "needs_inputs" + reasons = ( + "Known geometry checks passed; precision or planning estimates are incomplete. " + "Memory will be observed during collection; users need not supply non-KV byte estimates. " + "A known byte lower bound below the budget does not establish fit.", + ) + else: + status = "estimated_fit" + reasons = ( + f"Complete declared/estimated per-rank resources including one full-context cache need {known} bytes, " + f"within the {request.collection.memory_fraction * 100:g}% per-GPU budget " + f"of {hardware.memory_budget_bytes} bytes; " + "this is a conservative admission precheck, not runtime qualification.", + ) + return TopologyCandidate( + family, sizes, status, reasons, draft, known, estimated, hardware.memory_budget_bytes, resident + ) + + +def suggest_topologies( + config: ModelConfig, request: SupportRequest, overrides: Mapping[str, Any] | None = None +) -> TopologySuggestions: + """Shortlist one worker's starting points, without user allocation or timing data. + + The request supplies actual deployment identity, collection bounds and context. Its + topology is replaced during enumeration. Exact per-rank bounds and attached + deployment profiles must instead use the existing explicit-topology route. + """ + supplied = validate_overrides(overrides) + rank_overrides = sorted((_RANK_BYTES | {"cache_groups"}) & supplied.keys()) + if rank_overrides: + raise ValueError( + "automatic topology suggestions cannot transfer rank-local byte overrides across topologies: " + + ", ".join(rank_overrides) + + "; select an exact topology and use explicit onboarding with these bounds" + ) + if request.fpm_profile is not None: + raise ValueError("an attached FPM profile belongs to an exact topology; use explicit onboarding") + hardware = _hardware_envelope( + request.identity.gpu, request.identity.interconnect, request.collection.memory_fraction + ) + candidates = [] + rejected = [] + families = _FAMILIES if request.identity.model_kind == "moe" else ("tp",) + for family in families: + passing, pending = [], [] + width = 1 + while width <= hardware.fast_domain_gpus: + if width > 1 or family == families[0]: + candidate = _assess(config, request, supplied, hardware, family, width) + if candidate.status == "estimated_fit": + passing.append(candidate) + elif candidate.status == "needs_inputs": + pending.append(candidate) + else: + rejected.append(candidate) + width *= 2 + candidates.extend((passing or pending)[:2]) + passing = [candidate for candidate in candidates if candidate.status == "estimated_fit"] + default = min(passing, key=lambda candidate: candidate.required_gpus, default=None) + payload = request.model_dump() + payload["collection"].update( + {field: supplied[field] for field in ("max_num_tokens", "max_batch_size") if field in supplied} + ) + collection = SupportRequest.model_validate(payload).collection_settings() + for field, scheduler_field in (("max_num_tokens", "max_batched_tokens"), ("max_batch_size", "max_sequences")): + if field in supplied: + collection["sources"][scheduler_field] = "user resource override; shared rank-local scheduler bound" + grouped = any( + candidate.draft and candidate.draft.resolved.get("cache_layout") == "grouped" for candidate in candidates + ) + return TopologySuggestions( + config_sha256=config.sha256, + identity=request.identity.model_dump(mode="json"), + collection=collection, + context_length=request.search.context_length, + overrides=supplied, + hardware=hardware, + candidates=tuple(candidates), + rejected_candidates=tuple(rejected), + default=default, + assumptions=( + hardware.assumption, + "Enumerate powers-of-two widths within the fast domain; explicit topology choices may exceed it.", + ( + "Check one request's conservative peak grouped allocation using Rust's window, block-rounding " + "and transient prefill bound at the selected context and scheduler limits. No aggregate token " + "capacity is assigned to mixed cache groups. Concurrent workload fit remains unchecked." + if grouped + else "Check room for one full search.context_length request plus one cached token of strict capacity " + "headroom per rank. max_batch_size is a scheduler bound, not that many full-context cache allocations. " + "Total KV capacity is estimated from the remaining per-rank memory; this check does not guarantee " + "a particular concurrent workload fits or assume balanced attention-DP routing." + ), + f"The {request.collection.memory_fraction * 100:g}% byte budget uses the reviewed GPU memory fraction. " + "CUDA graph reservations are excluded; " + "generated plan and serving-runtime admission remain authoritative.", + "Only known config geometry constraints are checked. Unknown architecture constraints, " + "runtime compatibility, timing coverage and measured performance remain unchecked.", + "Shortlist the first two complete fits per family, or pending-input choices when none fit; " + "default to the smallest complete fit, breaking width ties by TP, DEP, then TEP. " + "This ordering is not a performance ranking.", + *([config.notes["modeling_scope"]] if "modeling_scope" in config.notes else []), + ), + ) diff --git a/python/aisimulate/src/aisimulate/support/validation.py b/python/aisimulate/src/aisimulate/support/validation.py new file mode 100644 index 000000000..22a2ea4cc --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/validation.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate collected FPM timings through ordinary cold aggregated replay.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate.config import CorePredictionConfig +from aisimulate.config.common import load_yaml +from aisimulate.output import prepare_output_directory + +from .plan import check_plan, plan_lock, request_id +from .schema import SupportRequest + + +def add_validation_parser(actions: Any) -> None: + parser = actions.add_parser( + "validate-fpm", + help="Check collected direct-FPM query coverage with a local Weka/AgentX trace.", + description=( + "Run ordinary predict with cold aggregated replay, one client lane, HBM-only cache and no speculation. " + "The trace may contain one play or a selected corpus. All selected requests must finish and every " + "native timing lookup must resolve for coverage to pass. Accuracy is not assessed. " + "This command does not download traces or collect GPU measurements." + ), + ) + parser.add_argument("-c", "--config", required=True, help="Reviewed onboarding request.") + parser.add_argument("--output-dir", required=True, help="Existing matching collection plan directory.") + parser.add_argument("--trace", required=True, help="Local Weka JSON/JSONL file; replay its full contents.") + parser.add_argument( + "--validation-output-dir", required=True, help="Separate directory for config, replay and coverage evidence." + ) + parser.add_argument("--overwrite", action="store_true", help="Replace previous validation outputs.") + + +def _file_identity(path: Path) -> dict[str, Any]: + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + return {"path": str(path), "sha256": digest, "size_bytes": path.stat().st_size} + + +def _fpm_artifacts(plan: Path) -> list[dict[str, Any]]: + return [ + _file_identity(path) for path in sorted((plan / "systems/data").rglob("fpm_forward_perf.*")) if path.is_file() + ] + + +def _prediction_config(request: SupportRequest, plan: Path, trace: Path) -> dict[str, Any]: + """Retain the reviewed deployment and replace only evaluation inputs.""" + + if request.worker_type in {"prefill", "decode"}: + raise ValueError( + "validate-fpm currently supports aggregated replay only. This independent P/D role can run " + "validate-collection; full replay requires a separately accepted counterpart and a supported " + "P/D serving configuration. Collection does not require actual P/D handoff." + ) + prediction = load_yaml(plan / "predict/pilot.yaml") + prediction["traffic"] = { + "source": { + "type": "trace", + "paths": [str(trace)], + "format": "weka", + "nested_timestamp_basis": "absolute", + }, + "load": {"type": "trace_timestamps", "agentic_lanes": 1}, + } + prediction["evaluation"] = {} + engine = prediction["engine"] + if engine["mode"] != "aggregated" or engine.get("speculation") is not None: + raise ValueError("FPM replay validation requires aggregated serving with speculation disabled") + worker = engine["workers"]["aggregated"] + timing = worker["timing"] + timing["estimation_mode"] = "fpm_interpolation" + timing["fallback_policy"] = "deny" + interpolation = timing.setdefault("estimator_config", {}).setdefault("fpm_interpolation", {}) + interpolation.update(method="direct", collect_coverage=True) + if engine["model"] != request.identity.model: + raise ValueError("saved prediction model differs from the reviewed collection identity") + CorePredictionConfig.model_validate(prediction) + return prediction + + +def _completion_issues(report: dict[str, Any]) -> list[str]: + """Check terminal evidence, including requests that a failed play never dispatched.""" + + issues = [] + graph = report.get("agentic_graph", {}) + count = graph.get("node_count", 0) + if not isinstance(count, int) or count <= 0: + return ["replay did not report a nonempty AgentX request graph"] + if report.get("num_requests") != count or report.get("completed_requests") != count: + issues.append("not every selected request completed") + records = report.get("per_request", []) + if len(records) != count or any( + row.get("terminal_status") != "completed" or row.get("output_length") != row.get("requested_output_length") + for row in records + ): + issues.append("request evidence is missing, failed, or has incomplete output") + outcomes = report.get("agentic_play_outcomes", []) + if ( + not outcomes + or len(outcomes) != graph.get("play_count") + or any(row.get("status") != "completed" or row.get("settled_at_ms") is None for row in outcomes) + ): + issues.append("not every selected play reached completed settlement") + return issues + + +def _write_report(path: Path, report: dict[str, Any]) -> None: + path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def run_validation(args: argparse.Namespace) -> int: + request = SupportRequest.from_yaml(args.config) + plan = Path(args.output_dir).expanduser().resolve() + trace = Path(args.trace).expanduser().resolve() + output = Path(args.validation_output_dir).expanduser().resolve() + if not trace.is_file() or trace.suffix.lower() not in {".json", ".jsonl"}: + raise ValueError("--trace must be a local Weka JSON or JSONL file") + if output == plan or plan in output.parents or output in plan.parents: + raise ValueError("validation output must be separate from the collection plan directory") + for source in (Path(args.config).expanduser().resolve(), trace): + if source == output or output in source.parents: + raise ValueError(f"validation output must not contain the input file {source}") + with plan_lock(plan): + check_plan(request, plan) + prediction = _prediction_config(request, plan, trace) + if request.fpm_profile is not None: + request.profile_deployment().resources.require_memory() + collection = json.loads((plan / "support-plan.json").read_text(encoding="utf-8")) + for relative in ("predict.yaml", "validation.json", "prediction"): + if (output / relative).is_symlink(): + raise ValueError(f"refusing symlinked validation output {output / relative}") + prepare_output_directory(output, overwrite=args.overwrite) + prepare_output_directory(output / "prediction", overwrite=args.overwrite) + prediction_path = output / "predict.yaml" + prediction_path.write_text(yaml.safe_dump(prediction, sort_keys=False), encoding="utf-8") + report_path = output / "validation.json" + report: dict[str, Any] = { + "schema_version": 1, + "status": "incomplete", + "accuracy": "not_assessed", + "collection_id": request_id(request), + "collection_plan": _file_identity(plan / "support-plan.json"), + "trace": _file_identity(trace), + "fpm_artifacts": _fpm_artifacts(plan), + "scope": { + "replay": "cold_aggregated", + "agentic_lanes": 1, + "nested_timestamp_basis": "absolute", + "cache": "hbm_only", + "speculation": "disabled", + "model_projection": request.identity.model, + "selected_corpus": "complete_local_file", + }, + "prediction_config": _file_identity(prediction_path), + "prediction_report": str(output / "prediction/prediction.json"), + "coverage_report": str(output / "prediction/fpm-coverage.json"), + "collection_schema_version": collection.get("schema_version"), + "issues": ["replay has not completed"], + } + _write_report(report_path, report) + # Use the existing supervised CLI: the same compiler, capacity checks, + # timing provider, source importer and host limits apply to ordinary predict. + from aisimulate.supervision import main as predict + + command = [ + "predict", + "--config", + str(prediction_path), + "--output-dir", + str(output / "prediction"), + "--capture-per-request", + "--format", + "json", + ] + if args.overwrite: + command.append("--overwrite") + try: + exit_code = predict(command) + except SystemExit as exc: + exit_code = int(exc.code) if isinstance(exc.code, int) else 1 + except KeyboardInterrupt: + report["issues"] = ["replay interrupted; coverage remains incomplete"] + _write_report(report_path, report) + raise + report["prediction_exit_code"] = exit_code + coverage_path = Path(report["coverage_report"]) + coverage = json.loads(coverage_path.read_text(encoding="utf-8")) if coverage_path.is_file() else None + report["fpm_query_coverage"] = coverage + if coverage_path.is_file(): + report["coverage_evidence"] = _file_identity(coverage_path) + issues = [] + if exit_code != 0: + issues.append(f"predict exited with status {exit_code}; replay coverage is incomplete") + prediction_report = Path(report["prediction_report"]) + if prediction_report.is_file(): + report["prediction_evidence"] = _file_identity(prediction_report) + replay = json.loads(prediction_report.read_text(encoding="utf-8")) + issues.extend(_completion_issues(replay)) + report["agentic_graph"] = replay.get("agentic_graph") + report["agentic_model_projection"] = replay.get("agentic_model_projection") + else: + issues.append("no completed prediction report was produced") + if not coverage or coverage.get("status") != "covered": + issues.append( + "native direct-FPM query coverage did not pass; inspect the coverage report for missing coordinates" + ) + if _file_identity(trace) != report["trace"]: + issues.append("trace file changed during validation") + if _fpm_artifacts(plan) != report["fpm_artifacts"]: + issues.append("FPM library files changed during validation") + with plan_lock(plan): + check_plan(request, plan) + report["issues"] = issues + report["status"] = "covered" if not issues else "incomplete" + _write_report(report_path, report) + print(f"FPM replay coverage: {report['status']}; accuracy: not assessed. Saved {report_path}") + return 0 if not issues else (exit_code or 1) diff --git a/python/aisimulate/src/aisimulate/support/validation_workflow.py b/python/aisimulate/src/aisimulate/support/validation_workflow.py new file mode 100644 index 000000000..e70695cb8 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/validation_workflow.py @@ -0,0 +1,934 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Source-bound, editable collection and serving qualification for onboarding. + +Validation files are separate from the collection plan. This module composes the +collector and native evaluators; it does not select another timing estimator or +change the runtime's collection grid. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import tempfile +from pathlib import Path +from typing import Annotated, Any, Literal + +from pydantic import Field, model_validator + +from aisimulate.config.common import StrictModel, load_yaml + +from .plan import check_plan, request_id +from .schema import SupportRequest + +Ratio = Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)] + + +class RepeatabilityPolicy(StrictModel): + comparison_mode: Literal["full_grid", "bounded"] = "full_grid" + samples: int = Field(default=5, strict=True, ge=2) + max_points_per_cell: int = Field(default=12, strict=True, ge=1) + max_attempts_per_sample: int = Field(default=2, strict=True, ge=1) + max_cv: Ratio = 0.05 + max_source_relative_difference: Ratio = 0.05 + + +class InterpolationPolicy(StrictModel): + max_points_per_phase: int = Field(default=16, strict=True, ge=1) + seed: int = Field(default=42, strict=True, ge=0) + max_p95_relative_error: Ratio = 0.20 + max_unsupported_fraction: float = Field(default=0.0, strict=True, ge=0, le=1, allow_inf_nan=False) + + +class ServingPolicy(StrictModel): + max_latency_p95_relative_error: Ratio = 0.20 + max_throughput_relative_error: Ratio = 0.20 + max_dispatch_delay_ms: Ratio = 5.0 + max_dispatch_delay_fraction: Ratio = 0.01 + + +class ValidationPolicy(StrictModel): + schema_version: Literal[ + "aisimulate-onboarding-validation-policy/v1", "aisimulate-onboarding-validation-policy/v2" + ] = "aisimulate-onboarding-validation-policy/v2" + repeatability: RepeatabilityPolicy = Field(default_factory=RepeatabilityPolicy) + interpolation: InterpolationPolicy = Field(default_factory=InterpolationPolicy) + serving: ServingPolicy = Field(default_factory=ServingPolicy) + + @model_validator(mode="before") + @classmethod + def preserve_legacy_subset_policy(cls, values: Any) -> Any: + """Reading a saved v1 campaign must not turn its subset into a full sweep.""" + if isinstance(values, dict) and values.get("schema_version") == "aisimulate-onboarding-validation-policy/v1": + repeatability = values.get("repeatability", {}) + if isinstance(repeatability, dict): + if repeatability.get("comparison_mode", "bounded") != "bounded": + raise ValueError("legacy v1 policy uses bounded repeats; choose v2 in a fresh campaign") + values = {**values, "repeatability": {**repeatability, "comparison_mode": "bounded"}} + return values + + +def _json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text()) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value + + +def _write(path: Path, value: Any) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") + + +def _identity(path: Path) -> dict[str, Any]: + path = path.expanduser().resolve() + raw = path.read_bytes() + return {"path": str(path), "sha256": hashlib.sha256(raw).hexdigest(), "size_bytes": len(raw)} + + +def _checked(ref: dict[str, Any]) -> Path: + path = Path(ref["path"]) + if _identity(path) != ref: + raise ValueError(f"validation evidence changed: {path}") + return path + + +def _policy(path: str | Path | None) -> ValidationPolicy: + if not path: + return ValidationPolicy() + try: + # YAML 1.1 treats JSON exponent notation such as 1e-05 as a string. + values = _json(Path(path)) + except (OSError, json.JSONDecodeError): + values = load_yaml(path) + return ValidationPolicy.model_validate(values) + + +def _fresh(output: Path, *inputs: Path) -> None: + if output.is_symlink(): + raise ValueError("validation output must not be a symbolic link") + output = output.resolve() + for source in inputs: + source = source.resolve() + if output == source or output in source.parents or (source.is_dir() and source in output.parents): + raise ValueError("validation output must be separate from all source artifacts") + if output.exists() and (not output.is_dir() or any(output.iterdir())): + raise ValueError("validation output must be fresh; resume the same policy or choose a new directory") + output.mkdir(parents=True, exist_ok=True) + + +def _referenced_paths(value: Any) -> list[Path]: + """Collect artifact identities and preserved directories before creating outputs.""" + paths = [] + if isinstance(value, dict): + for key, item in value.items(): + if key in { + "path", + "collection_directory", + "directory", + "tokenizer", + "artifacts_dir", + "output_dir", + } and isinstance(item, str): + paths.append(Path(item)) + else: + paths.extend(_referenced_paths(item)) + elif isinstance(value, list): + for item in value: + paths.extend(_referenced_paths(item)) + return paths + + +def _status(gates: dict[str, dict[str, Any]]) -> str: + statuses = [item["status"] for item in gates.values()] + return "failed" if "failed" in statuses else "passed" if all(x == "passed" for x in statuses) else "incomplete" + + +def _execution_status(frozen: dict[str, Any]) -> str: + return _status( + { + cell["cell_id"]: { + "status": "passed" if cell["execution"]["status"] == "qualified" else cell["execution"]["status"] + } + for cell in frozen["cells"] + } + ) + + +def _collection_source(request: SupportRequest, root: Path) -> tuple[Any, Path, Path, list[dict[str, Any]]]: + """Verify formal publication against every native row before qualification.""" + import pyarrow.parquet as pq + from collector.fpm_forward.config import PrefillSamplingProfile + from collector.fpm_forward.database import ( + aggregate_cell, + published_dense_synthetic_cells, + validate_formal_database_commit, + ) + from collector.fpm_forward.planner import backend_identity_columns + from collector.fpm_forward.repeatability import load_repeatability_source + + check_plan(request, root) + checkpoints = list((root / "fpm-checkpoint").rglob("fpm_forward.json")) + if len(checkpoints) != 1: + raise ValueError("collection validation requires exactly one formal collector checkpoint") + checkpoint_path = checkpoints[0] + checkpoint = _json(checkpoint_path) + sha = checkpoint.get("plan_sha256", "") + if not re.fullmatch(r"[0-9a-f]{64}", sha): + raise ValueError("invalid source collection plan identity") + campaign = root / "fpm-artifacts" / sha[:16] + plan = load_repeatability_source(campaign) + payload = plan.to_dict() + if ( + plan.sha256 != sha + or plan.model_path != request.identity.model + or plan.system != request.identity.gpu + or plan.backend != request.identity.framework + or plan.capability.aic_database_version != request.identity.framework_version + or request.fpm_profile is None + or payload.get("fpm_profile") != request.fpm_profile.model_dump(mode="json") + ): + raise ValueError("source collection differs from the original onboarding identity/profile") + scheduler = request.scheduler_limits() + sampling = PrefillSamplingProfile.build( + max_isl=scheduler["max_batched_tokens"], + max_batch_size=scheduler["max_sequences"], + max_cudagraph_capture_size=request.collection.max_prefill_cudagraph_size, + cudagraph_policy=request.collection.prefill_cudagraph_policy, + ) + if ( + plan.options.vllm_max_model_len != request.search.context_length + or plan.options.max_num_batched_tokens != scheduler["max_batched_tokens"] + or plan.options.max_num_seqs != scheduler["max_sequences"] + or (plan.options.gpu_memory_utilization or request.collection.memory_fraction) + != request.collection.memory_fraction + or plan.options.prefill_cudagraph_policy != request.collection.prefill_cudagraph_policy + or plan.options.prefill_sampling != sampling + or plan.options.worker_type != request.worker_type + or plan.options.cudagraph_mode != request.collection.cudagraph_mode + or plan.options.cudagraph_capture_sizes + != ( + tuple(request.collection.cudagraph_capture_sizes) + if request.collection.cudagraph_capture_sizes is not None + else None + ) + or plan.options.max_cudagraph_capture_size != request.collection.max_cudagraph_capture_size + ): + raise ValueError("source collection runtime limits differ from the reviewed onboarding settings") + deployment = request.profile_deployment() + if {cell.workload_kind for cell in plan.cells} != set(request.collection_phases): + raise ValueError( + "collection validation requires exactly the accepted role's phases: " + ", ".join(request.collection_phases) + ) + database = checkpoint.get("database", {}) + if ( + database.get("status") != "passed" + or database.get("missing_cells") != [] + or database.get("skipped_first_publisher_wins") != [] + or database.get("plan_cells") != len(plan.cells) + or database.get("published_cells") != len(plan.cells) + ): + raise ValueError("collection validation requires complete formal publication without reused cells") + parquet, metadata = (Path(database[name]).resolve() for name in ("parquet", "metadata")) + if any(not path.is_relative_to(root / "systems/data") for path in (parquet, metadata)): + raise ValueError("formal data must belong to this source campaign") + published = validate_formal_database_commit(parquet, metadata, plan) + dense_cells = published_dense_synthetic_cells(published) + rows = [] + for cell in plan.cells: + for field in ("tp", "pp", "dp", "cp", "moe_tp", "moe_ep"): + if getattr(cell.topology, field) != getattr(deployment, field): + raise ValueError("collection topology differs from onboarding") + for field in ("gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "comm_quant_mode", "kv_cache_dtype"): + if getattr(cell, field) != getattr(deployment, field): + raise ValueError("collection precision differs from onboarding") + if any( + getattr(deployment, name) != value for name, value in backend_identity_columns(cell.backend_policy).items() + ): + raise ValueError("collection backend identity differs from onboarding") + entry = checkpoint.get("cells", {}).get(cell.cell_id, {}) + if entry.get("status") != "passed" or not entry.get("attempt_id"): + raise ValueError(f"source cell is not complete: {cell.cell_id}") + rows.extend( + aggregate_cell( + plan, + cell, + campaign / "cells" / cell.cell_id, + expected_attempt_id=entry["attempt_id"], + dense_synthetic_kv=cell.cell_id in dense_cells, + ) + ) + canonical = lambda row: json.dumps(row, sort_keys=True, allow_nan=False) + if sorted(map(canonical, rows)) != sorted(map(canonical, pq.read_table(parquet).to_pylist())): + raise ValueError("published formal rows differ from verified native collection") + return plan, campaign, checkpoint_path, [_identity(parquet), _identity(metadata)] + + +def _holdout_evidence(plan: Any, campaign: Path, frozen: dict[str, Any]) -> list[dict[str, Any]]: + """Reuse full canonical native records, not the bounded repeatability subset.""" + from collector.fpm_forward.native_artifact import select_native_measurements, validate_native_collection + + evidence = [] + cells = {cell["cell_id"]: cell for cell in frozen["cells"]} + for cell in plan.cells: + source = cells[cell.cell_id] + collection = validate_native_collection( + cell, + campaign / "cells" / cell.cell_id / "raw", + expected_plan_sha256=plan.sha256, + expected_attempt_id=source["source_attempt_id"], + ) + evidence.append( + { + "cell_id": cell.cell_id, + "identity": { + "source_plan_sha256": plan.sha256, + "collector_attempt_id": collection.collector_attempt_id, + "runtime_run_id": collection.runtime_run_id, + "runtime_grid_digest": collection.runtime_grid_digest, + }, + "native_graph_config": source["execution"]["native_graph_config"], + "points": [ + { + "workload_kind": item.point["point_type"], + **{ + name: item.point.get(name) + for name in ( + "batch_size", + "total_prefill_tokens", + "total_kv_read_tokens", + "expected_cudagraph_mode", + "expected_capture_size", + ) + }, + } + for item in select_native_measurements(collection, cell_id=cell.cell_id) + ], + } + ) + return evidence + + +def _holdout( + request: SupportRequest, root: Path, output: Path, policy: ValidationPolicy, evidence: list[dict[str, Any]] +) -> dict[str, Any]: + from .interpolation_validation import evaluate_interpolation_holdout + + return evaluate_interpolation_holdout( + request, + systems_root=root / "systems", + output_dir=output, + execution_evidence=evidence, + **policy.interpolation.model_dump(), + ) + + +def _repeat_plan( + plan: Any, + campaign: Path, + checkpoint: Path, + policy: ValidationPolicy, + *, + observation_evidence_version: int = 2, +) -> dict[str, Any]: + from collector.fpm_forward.repeatability import freeze_repeatability_plan + + return freeze_repeatability_plan( + plan, + campaign, + checkpoint, + samples=policy.repeatability.samples, + max_points_per_cell=policy.repeatability.max_points_per_cell, + cv_threshold=policy.repeatability.max_cv, + comparison_mode=policy.repeatability.comparison_mode, + max_attempts_per_sample=policy.repeatability.max_attempts_per_sample, + source_agreement_threshold=policy.repeatability.max_source_relative_difference, + observation_evidence_version=observation_evidence_version, + ) + + +def _assess_repeats(directory: Path, policy: ValidationPolicy, expected_plan: dict[str, Any]) -> dict[str, Any]: + from collector.fpm_forward.repeatability import PLAN_FILENAME, REPORT_FILENAME, assess_repeatability + + frozen = _json(directory / PLAN_FILENAME) + if ( + {key: value for key, value in frozen.items() if key not in {"policy", "sha256"}} + != {key: value for key, value in expected_plan.items() if key not in {"policy", "sha256"}} + or frozen["policy"]["samples"] != policy.repeatability.samples + or frozen["policy"]["max_points_per_cell"] != policy.repeatability.max_points_per_cell + or frozen["policy"].get("comparison_mode", "bounded") != policy.repeatability.comparison_mode + or frozen["policy"].get("max_attempts_per_sample") != policy.repeatability.max_attempts_per_sample + ): + raise ValueError("repeatability measurements do not match the source or selected sample policy") + return assess_repeatability( + frozen, + _json(directory / REPORT_FILENAME), + cv_threshold=policy.repeatability.max_cv, + source_agreement_threshold=policy.repeatability.max_source_relative_difference, + ) + + +def _repeatability_gate_status(assessment: dict[str, Any], policy: ValidationPolicy) -> str: + """Fresh stable runs alone cannot qualify the original published source.""" + if ( + assessment.get("status") == "failed" + or assessment.get("source_qualification", {}).get("status") == "failed" + or assessment.get("new_population", {}).get("status") in {"failed", "unstable"} + ): + return "failed" + if ( + policy.schema_version == "aisimulate-onboarding-validation-policy/v2" + and policy.repeatability.comparison_mode == "full_grid" + and assessment.get("schema_version") == 2 + and assessment.get("policy", {}).get("comparison_mode") == "full_grid" + and assessment.get("source_qualification", {}).get("status") == "qualified" + and assessment.get("new_population", {}).get("status") == "qualified" + and assessment.get("status") == "passed" + ): + return "passed" + return "incomplete" + + +def validate_collection(args: argparse.Namespace) -> int: + from collector.fpm_forward.repeatability import PLAN_FILENAME, load_repeatability_deployment, run_repeatability + + request = SupportRequest.from_yaml(args.config) + root = Path(args.output_dir).expanduser().resolve() + output = Path(args.validation_output_dir).expanduser().resolve() + policy = _policy(args.policy or (output / "policy.json" if args.resume else None)) + if args.retry_failed and not (args.execute and args.resume): + raise ValueError("--retry-failed requires --execute --resume") + if args.repeatability_dir and (args.execute or args.resume): + raise ValueError("reuse --repeatability-dir only with a fresh offline assessment") + plan, campaign, checkpoint, formal = _collection_source(request, root) + previous_selection = ( + output / "repeatability-selection.json" + if args.resume + else Path(args.repeatability_dir).resolve() / PLAN_FILENAME + if args.repeatability_dir + else None + ) + evidence_version = ( + _json(previous_selection).get("observation_evidence_version", 1) if previous_selection is not None else 2 + ) + frozen = _repeat_plan(plan, campaign, checkpoint, policy, observation_evidence_version=evidence_version) + holdout_evidence = _holdout_evidence(plan, campaign, frozen) + inputs = { + "request": _identity(Path(args.config)), + "collection_directory": str(root), + "support_plan": _identity(root / "support-plan.json"), + "collection_plan": _identity(campaign / "collection-plan.json"), + "checkpoint": _identity(checkpoint), + "formal_data": formal, + } + if args.resume: + if _json(output / "campaign.json") != inputs or _json(output / "repeatability-selection.json") != frozen: + raise ValueError("validation campaign inputs changed; preserve it and use a fresh output directory") + if ValidationPolicy.model_validate(_json(output / "policy.json")) != policy: + raise ValueError("validation policy changed; reassess in a fresh output directory") + else: + _fresh( + output, + root, + Path(args.config), + *([Path(args.policy)] if args.policy else []), + *([Path(args.repeatability_dir)] if args.repeatability_dir else []), + ) + _write(output / "policy.json", policy.model_dump(mode="json")) + _write(output / "campaign.json", inputs) + _write(output / "repeatability-selection.json", frozen) + report: dict[str, Any] = { + "schema_version": "aisimulate-collection-validation/v1", + "status": "incomplete", + "accuracy": "not_assessed", + "policy": _identity(output / "policy.json"), + "inputs": inputs, + "selection": _identity(output / "repeatability-selection.json"), + "gates": { + "point_validity": {"status": "passed"}, + "execution": { + "status": _execution_status(frozen), + "cells": [{"cell_id": c["cell_id"], **c["execution"]} for c in frozen["cells"]], + }, + "repeatability": {"status": "incomplete", "issues": ["comparable GPU repeats have not completed"]}, + "interpolation": {"status": "incomplete"}, + }, + } + report_path = output / "collection-validation.json" + if not args.resume: + _write(report_path, report) + holdout_dir = output / "holdout" + if args.resume: + with tempfile.TemporaryDirectory(prefix="aisimulate-holdout-check-") as temporary: + holdout = _holdout(request, root, Path(temporary) / "holdout", policy, holdout_evidence) + old = _json(holdout_dir / "interpolation-validation.json") + _compare_holdout(old, holdout) + else: + holdout = _holdout(request, root, holdout_dir, policy, holdout_evidence) + report["gates"]["interpolation"] = { + "status": holdout["status"], + "report": _identity(holdout_dir / "interpolation-validation.json"), + } + _write(report_path, report) + repeats = Path(args.repeatability_dir).resolve() if args.repeatability_dir else output / "repeatability" + if args.execute: + run_repeatability( + plan, + generator_overrides=load_repeatability_deployment(campaign), + source_campaign_dir=campaign, + source_checkpoint_path=checkpoint, + output_dir=repeats, + resume=args.resume and repeats.exists(), + retry_failed=args.retry_failed, + samples=policy.repeatability.samples, + max_points_per_cell=policy.repeatability.max_points_per_cell, + cv_threshold=policy.repeatability.max_cv, + comparison_mode=policy.repeatability.comparison_mode, + max_attempts_per_sample=policy.repeatability.max_attempts_per_sample, + source_agreement_threshold=policy.repeatability.max_source_relative_difference, + observation_evidence_version=evidence_version, + ) + if repeats.exists(): + assessment = _assess_repeats(repeats, policy, frozen) + _write(output / "repeatability-assessment.json", assessment) + report["gates"]["repeatability"] = { + "status": _repeatability_gate_status(assessment, policy), + "directory": str(repeats), + "assessment": _identity(output / "repeatability-assessment.json"), + } + report["status"] = _status(report["gates"]) + _write(report_path, report) + print(f"Collection validation: {report['status']}; serving accuracy not assessed. Saved {report_path}") + return ( + 0 + if report["status"] == "passed" + or (report["status"] == "incomplete" and not args.execute and not args.repeatability_dir) + else 1 + ) + + +def _compare_holdout(saved: dict[str, Any], current: dict[str, Any]) -> None: + for key in ( + "status", + "policy", + "capture_boundaries", + "phases", + "predictions", + "native_query_coverage", + "worker_type", + "collection_phases", + ): + if saved.get(key) != current.get(key): + raise ValueError( + f"saved holdout {key} differs from native reassessment; " + "preserve it and use a fresh assessment directory" + ) + for key, value in saved["artifacts"].items(): + if key == "source": + for source in value.values(): + _checked(source) + else: + _checked(value) + + +def check_collection_report( + path: Path, policy: ValidationPolicy | None = None +) -> tuple[dict[str, Any], SupportRequest, Any, dict[str, Any]]: + """Recompute all mandatory collection gates from preserved native/formal data.""" + report = _json(path) + if report.get("schema_version") != "aisimulate-collection-validation/v1": + raise ValueError("unsupported collection validation report") + previous_policy = _policy(_checked(report["policy"])) + if policy is None: + policy = previous_policy + if previous_policy.repeatability != policy.repeatability or previous_policy.interpolation != policy.interpolation: + raise ValueError("collection validation policy differs; reassess with the selected policy") + inputs = report["inputs"] + request = SupportRequest.from_yaml(_checked(inputs["request"])) + root = Path(inputs["collection_directory"]) + for key in ("support_plan", "collection_plan", "checkpoint"): + _checked(inputs[key]) + plan, campaign, checkpoint, formal = _collection_source(request, root) + if formal != inputs["formal_data"]: + raise ValueError("source formal data changed since collection validation") + saved_selection = _json(_checked(report["selection"])) + frozen = _repeat_plan( + plan, + campaign, + checkpoint, + policy, + observation_evidence_version=saved_selection.get("observation_evidence_version", 1), + ) + if saved_selection != frozen: + raise ValueError("source runtime evidence or repeatability selection changed") + gates = report["gates"] + expected_execution = _execution_status(frozen) + if gates["point_validity"]["status"] != "passed" or gates["execution"]["status"] != expected_execution: + raise ValueError("saved collection integrity/execution gate differs from native evidence") + if gates["execution"].get("cells") != [{"cell_id": c["cell_id"], **c["execution"]} for c in frozen["cells"]]: + raise ValueError("saved execution inspection differs from native observations") + with tempfile.TemporaryDirectory(prefix="aisimulate-holdout-check-") as temporary: + holdout = _holdout( + request, root, Path(temporary) / "holdout", policy, _holdout_evidence(plan, campaign, frozen) + ) + _compare_holdout(_json(_checked(gates["interpolation"]["report"])), holdout) + if gates["interpolation"]["status"] != holdout["status"]: + raise ValueError("saved interpolation gate differs from native reassessment") + repeat = gates["repeatability"] + if "directory" in repeat: + assessment = _assess_repeats(Path(repeat["directory"]), policy, frozen) + if _json(_checked(repeat["assessment"])) != assessment or repeat["status"] != _repeatability_gate_status( + assessment, policy + ): + raise ValueError("saved repeatability gate differs from raw samples") + elif repeat["status"] != "incomplete": + raise ValueError("repeatability gate has no measurement evidence") + if report["status"] != _status(gates): + raise ValueError("collection validation status disagrees with required gates") + return report, request, plan, frozen + + +def _check_replay(path: Path, original: SupportRequest, collection: dict[str, Any], source_plan: Any) -> dict[str, Any]: + from .finalization import ( + _merge_resources, + _verify_collection, + finalization_manifest, + verify_finalized_data, + verify_finalized_quality, + ) + from .validation import _completion_issues, _fpm_artifacts, _prediction_config + + report = _json(path) + root = _checked(report["collection_plan"]).parent + request = SupportRequest.from_yaml(root / "request.yaml") + check_plan(request, root) + if request_id(request) != request_id(original): + manifest = finalization_manifest(request) + if ( + manifest is None + or manifest.get("source_request_id") != request_id(original) + or manifest.get("source_collection_plan_sha256") != source_plan.sha256 + or Path(manifest.get("source_directory", "")).resolve() + != Path(collection["inputs"]["collection_directory"]) + ): + raise ValueError("replay request is neither the collection request nor its finalized memory profile") + # Finalization may change resources/provenance, never model, topology, precision or collection settings. + payload = request.model_dump(mode="json") + payload["fpm_profile"]["provenance"] = original.fpm_profile.provenance + selected_index = original.fpm_profile.deployments.index(original.profile_deployment()) + payload["fpm_profile"]["deployments"][selected_index]["resources"] = ( + original.profile_deployment().resources.model_dump(mode="json") + ) + if payload != original.model_dump(mode="json"): + raise ValueError("finalized replay changed non-memory onboarding inputs") + for artifact in manifest["source_artifacts"]: + if hashlib.sha256(Path(artifact["path"]).read_bytes()).hexdigest() != artifact["sha256"]: + raise ValueError("finalization source evidence changed") + verify_finalized_data(request, root) + observations, verified_manifest, _, _ = _verify_collection( + original, + Path(collection["inputs"]["collection_directory"]), + memory_revision=manifest.get("memory_revision"), + ) + quality = verify_finalized_quality(manifest) + if quality is not None: + verified_manifest["collection_quality"] = quality + expected_resources = _merge_resources( + observations, + verified_manifest, + source_references=manifest.get("observation_provenance") == "source_references", + ) + if request.profile_deployment().resources.model_dump(mode="json") != expected_resources: + raise ValueError("finalized replay resources differ from verified collection memory") + request.profile_deployment().resources.require_memory() + trace = _checked(report["trace"]) + config_path = _checked(report["prediction_config"]) + if load_yaml(config_path) != _prediction_config(request, root, trace): + raise ValueError("replay configuration differs from the reviewed direct-FPM workload") + if report["collection_id"] != request_id(request) or report["fpm_artifacts"] != _fpm_artifacts(root): + raise ValueError("replay collection identity or formal data changed") + source_hashes = sorted((x["sha256"], x["size_bytes"]) for x in collection["inputs"]["formal_data"]) + if sorted((x["sha256"], x["size_bytes"]) for x in report["fpm_artifacts"]) != source_hashes: + raise ValueError("replay formal data differs from the validated collection") + if "prediction_evidence" not in report or "coverage_evidence" not in report: + raise KeyError("replay lacks bound prediction/coverage evidence; rerun onboard validate-fpm") + prediction_path = _checked(report["prediction_evidence"]) + coverage_path = _checked(report["coverage_evidence"]) + if str(prediction_path) != report["prediction_report"] or str(coverage_path) != report["coverage_report"]: + raise ValueError("replay evidence identities differ from the recorded paths") + prediction = _json(prediction_path) + coverage = _json(coverage_path) + issues = _completion_issues(prediction) + queries = coverage.get("queries", {}) + if ( + report.get("prediction_exit_code") != 0 + or coverage.get("status") != "covered" + or queries.get("unsupported") != 0 + or queries.get("measured", 0) + queries.get("interpolated", 0) <= 0 + or report.get("fpm_query_coverage") != coverage + ): + issues.append("native query coverage is incomplete or inconsistent") + return { + "status": "passed" if not issues else "incomplete", + "issues": issues, + "report": _identity(path), + "trace": _identity(trace), + "prediction_config": _identity(config_path), + "prediction_report": _identity(prediction_path), + "coverage_report": _identity(Path(report["coverage_report"])), + } + + +def _serving_expected(request: SupportRequest, plan: Any, frozen: dict[str, Any]) -> dict[str, Any]: + from collector.fpm_forward.repeatability import load_repeatability_deployment + + from .serving_validation import normalize_runtime_config + + deployment = request.profile_deployment() + workers = [worker for cell in frozen["cells"] for worker in cell["execution"]["observed_workers"]] + expected = { + "model": request.identity.model, + "model_revision": request.identity.model_revision, + "backend": request.identity.framework, + "backend_version": request.identity.framework_version, + "hardware": request.identity.gpu, + "parallelism": request.parallelism(), + "precision": { + "weights": deployment.gemm_quant_mode, + "fmha": deployment.fmha_quant_mode, + "kv": deployment.kv_cache_dtype, + }, + "scheduler": request.scheduler_limits(), + "context_length": request.search.context_length, + "gpu_memory_utilization": request.collection.memory_fraction, + } + for key in ("attention_groups", "graph_config"): + values = [worker.get(key) for worker in workers] + expected[key] = values[0] if values and all(item == values[0] for item in values) else None + configurations = [normalize_runtime_config(worker.get("resolved_config", {})) for worker in workers] + expected["runtime_config"] = ( + configurations[0] if configurations and all(item == configurations[0] for item in configurations) else None + ) + overrides = load_repeatability_deployment(frozen["source_campaign_dir"]) + image = plan.options.slurm_container_image or overrides.get("K8sConfig", {}).get("k8s_image", "") + match = re.search(r"(?:^|@)(sha256:[0-9a-f]{64})$", image) + expected["image_digest"] = match[1] if match else None + return expected + + +def _gate_failure(exc: Exception) -> dict[str, Any]: + return { + "status": "incomplete" if isinstance(exc, (FileNotFoundError, KeyError)) else "failed", + "issues": [str(exc)], + } + + +def validate_serving(args: argparse.Namespace) -> int: + from .serving_validation import ( + prepare_serving_validation, + run_serving_benchmark, + validate_serving_measurements, + ) + + if args.action == "run": + if not args.recipe or not args.aiperf_python: + raise ValueError("serving run requires --recipe and --aiperf-python; it launches the prepared benchmark") + receipt_path = run_serving_benchmark(Path(args.recipe), aiperf_python=Path(args.aiperf_python)) + receipt = _json(receipt_path) + print( + f"Serving benchmark: {receipt['status']} (producer exit {receipt['exit_code']}). Saved {receipt_path}; " + f"logs: {receipt_path.parent / 'aiperf.stdout.log'}, {receipt_path.parent / 'aiperf.stderr.log'}" + ) + return 0 if receipt["status"] == "completed" and receipt["exit_code"] == 0 else 1 + if not args.validation_output_dir or not args.collection_report or not args.replay_report: + raise ValueError( + "serving prepare/assess requires --validation-output-dir, --collection-report and --replay-report" + ) + collection_path, replay_path = Path(args.collection_report).resolve(), Path(args.replay_report).resolve() + output = Path(args.validation_output_dir).expanduser().resolve() + policy_path = args.policy + if policy_path is None and collection_path.is_file(): + policy_path = _json(collection_path).get("policy", {}).get("path") + policy = _policy(policy_path) + inputs = [collection_path.parent, replay_path.parent] + if args.recipe: + inputs.append(Path(args.recipe).parent) + for value in (policy_path, args.tokenizer, args.aiperf_python, args.execution_evidence, args.forward_evidence): + if value: + inputs.append(Path(value)) + for value in (collection_path, replay_path, args.recipe, args.execution_evidence, args.forward_evidence): + if value and Path(value).is_file(): + document = _json(Path(value)) + inputs.extend(_referenced_paths(document)) + # Replay may use a separate memory-finalized collection directory. + if value == replay_path and document.get("collection_plan", {}).get("path"): + inputs.append(Path(document["collection_plan"]["path"]).parent) + _fresh(output, *inputs) + _write(output / "policy.json", policy.model_dump(mode="json")) + combined: dict[str, Any] = { + "schema_version": "aisimulate-onboarding-validation/v1", + "status": "incomplete", + "accuracy": "not_qualified", + "policy": _identity(output / "policy.json"), + "scope": "selected representative collection coordinates and one complete matched Weka play", + "gates": {name: {"status": "incomplete"} for name in ("collection", "replay", "serving")}, + "limitations": [ + "This bounded assessment does not establish accuracy outside the evaluated configurations and workload." + ], + } + target = output / "validation.json" + _write(target, combined) + active_gate = "collection" + try: + collection, original, plan, frozen = check_collection_report(collection_path, policy) + combined["gates"]["collection"] = { + "status": collection["status"], + "report": _identity(collection_path), + "checks": collection["gates"], + } + active_gate = "replay" + replay = _check_replay(replay_path, original, collection, plan) + combined["gates"]["replay"] = replay + active_gate = "collection" + expected = _serving_expected(original, plan, frozen) + except (OSError, ValueError, KeyError, TypeError) as exc: + combined["gates"][active_gate] = _gate_failure(exc) + else: + try: + if args.action == "prepare": + if not args.endpoint or not args.tokenizer: + raise ValueError("serving prepare requires --endpoint and --tokenizer") + recipe_path = prepare_serving_validation( + trace=Path(replay["trace"]["path"]), + prediction_config=Path(replay["prediction_config"]["path"]), + prediction_report=Path(replay["prediction_report"]["path"]), + output=output / "serving", + endpoint=args.endpoint, + tokenizer=Path(args.tokenizer), + expected_execution=expected, + latency_p95_error=policy.serving.max_latency_p95_relative_error, + throughput_error=policy.serving.max_throughput_relative_error, + max_dispatch_delay_ms=policy.serving.max_dispatch_delay_ms, + max_dispatch_delay_fraction=policy.serving.max_dispatch_delay_fraction, + aiperf_python=Path(args.aiperf_python) if args.aiperf_python else None, + ) + combined["recipe"] = _identity(recipe_path) + combined["gates"]["serving"] = { + "status": "incomplete", + "issues": _json(recipe_path)["issues"] + ["matched serving measurement not assessed"], + } + else: + if not args.recipe or not args.execution_evidence: + raise ValueError("serving assess requires --recipe and --execution-evidence") + recipe = _json(Path(args.recipe)) + if ( + recipe["expected_execution"] != expected + or recipe["trace"] != replay["trace"] + or recipe["prediction_config"] != replay["prediction_config"] + or recipe["prediction_report"] != replay["prediction_report"] + ): + raise ValueError("serving recipe does not match validated collection execution and replay") + report_path = output / "serving-validation.json" + serving = validate_serving_measurements( + Path(args.recipe), + prediction_report=Path(replay["prediction_report"]["path"]), + execution_evidence_path=Path(args.execution_evidence), + output_report=report_path, + forward_evidence_path=Path(args.forward_evidence) if args.forward_evidence else None, + latency_p95_error=policy.serving.max_latency_p95_relative_error, + throughput_error=policy.serving.max_throughput_relative_error, + max_dispatch_delay_ms=policy.serving.max_dispatch_delay_ms, + max_dispatch_delay_fraction=policy.serving.max_dispatch_delay_fraction, + ) + combined["gates"]["serving"] = { + "status": serving["status"], + "report": _identity(report_path), + "metrics": serving.get("metrics", {}), + "forward": serving.get("forward", {"status": "unavailable"}), + "missing_evidence": serving["missing_evidence"], + "failures": serving["failures"], + } + except (OSError, ValueError, KeyError, TypeError) as exc: + combined["gates"]["serving"] = _gate_failure(exc) + combined["status"] = _status(combined["gates"]) + if combined["status"] == "passed": + combined["accuracy"] = "qualified_for_evaluated_scope" + _write(target, combined) + print(f"Onboarding validation: {combined['status']}; accuracy: {combined['accuracy']}. Saved {target}") + if ( + args.action == "prepare" + and combined.get("recipe") + and _json(Path(combined["recipe"]["path"]))["status"] == "ready" + ): + return 0 + return 0 if combined["status"] == "passed" else 1 + + +def add_quality_parsers(actions: Any) -> None: + collection = actions.add_parser( + "validate-collection", + help="Inspect collection and interpolation with editable quality policy; GPU repeats require --execute.", + ) + collection.add_argument( + "-c", "--config", required=True, help="Original reviewed collection request, before memory finalization." + ) + collection.add_argument("--output-dir", required=True, help="Original completed collection plan directory.") + collection.add_argument("--validation-output-dir", required=True, help="Separate fresh assessment directory.") + collection.add_argument( + "--policy", help="Editable validation policy JSON/YAML; omitted fields use documented defaults." + ) + collection.add_argument( + "--execute", + action="store_true", + help="Launch frozen full-grid repeats, or bounded diagnostics, using the original collector deployment.", + ) + collection.add_argument( + "--resume", + action="store_true", + help="Continue this exact policy and source campaign, preserving repeat attempts.", + ) + collection.add_argument( + "--retry-failed", + action="store_true", + help="Retry failed measurements within the frozen attempt budget; requires --execute --resume.", + ) + collection.add_argument( + "--repeatability-dir", + help="Reuse preserved repeat measurements in a fresh offline assessment; execution policy must match.", + ) + serving = actions.add_parser( + "validate-serving", + help="Prepare, run, or assess matched serving and combine mandatory onboarding validation gates.", + ) + serving.add_argument("--action", choices=("prepare", "run", "assess"), required=True) + serving.add_argument( + "--validation-output-dir", + help="Fresh directory for preparation or assessment; never replaces raw serving evidence.", + ) + serving.add_argument( + "--collection-report", help="Verified collection-validation.json; required for prepare/assess." + ) + serving.add_argument( + "--replay-report", help="Existing validate-fpm validation.json for the selected complete single-stream play." + ) + serving.add_argument( + "--policy", help="Validation policy JSON/YAML; defaults to the collection report's frozen policy." + ) + serving.add_argument("--recipe", help="Prepared recipe.json; required for run/assess.") + serving.add_argument("--endpoint", help="Caller-managed HTTP(S) serving endpoint; prepare never launches a server.") + serving.add_argument("--tokenizer", help="Local tokenizer directory for the pinned model; required for prepare.") + serving.add_argument( + "--aiperf-python", + help="Python with pinned AIPerf and target tokenizer; prepare freezes payloads, run sends traffic.", + ) + serving.add_argument( + "--execution-evidence", + help="Observed serving execution JSON, including actual runtime evidence sources; required for assess.", + ) + serving.add_argument( + "--forward-evidence", + help="Optional separately instrumented matched forward timings; absence is reported unavailable.", + ) diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 8167bbbae..505e8ec15 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -1108,6 +1108,47 @@ def validate_shape_dict(value: Any, label: str) -> None: raise ValueError(f"engine_integer_log_ranges.{name} must be positive integer [min, max] bounds") return self + @model_validator(mode="after") + def _validate_grouped_cache(self) -> SearchSpace: + if self.fpm_profile is None: + return self + from aisimulate.fpm_profile import load_fpm_profile + + profile = load_fpm_profile(self.fpm_profile) + if not any(deployment.resources.cache_layout == "grouped" for deployment in profile.deployments): + return self + from .search_space import _fpm_parallel_configs, _parallel_role + + for mode in self.deployment_mode: + configs = _fpm_parallel_configs(self, mode) + for role in ("agg",) if mode == "agg" else ("prefill", "decode"): + for shape in {_parallel_role(config, role).shape for config in configs}: + deployment = profile.select( + model=self.model_name, + worker_type="aggregated" if role == "agg" else role, + system=self.hardware_sku_for(role), + backend="vllm", + backend_version=self.requested_backend_version("vllm"), + tp_size=shape.tp, + pp_size=shape.pp, + attention_dp_size=shape.dp, + moe_tp_size=shape.moe_tp, + moe_ep_size=shape.moe_ep, + ) + if deployment.resources.cache_layout != "grouped": + continue + if role != "agg" or self.aic_nextn or self.speculation is not None: + raise ValueError("grouped FPM cache requires aggregated vLLM without speculative decoding") + if self.agg_enable_prefix_caching: + raise ValueError("grouped FPM cache requires agg_enable_prefix_caching=false for cold replay") + if self.agg_native_host_offload is not None: + raise ValueError("grouped FPM cache supports only HBM without host offload") + if self.agg_num_gpu_blocks is not None or self.agg_kv_bytes_per_token != "auto": + raise ValueError( + "grouped FPM cache requires profile groups and a byte budget, not scalar capacity" + ) + return self + class SweepConfig(BaseModel): """Sweep run-control.""" diff --git a/python/aisimulate/src/aisimulate/sweeper/deploy.py b/python/aisimulate/src/aisimulate/sweeper/deploy.py index 612d69072..3acea4a77 100644 --- a/python/aisimulate/src/aisimulate/sweeper/deploy.py +++ b/python/aisimulate/src/aisimulate/sweeper/deploy.py @@ -209,6 +209,7 @@ def _profile_kv_args(sample: dict[str, Any], role: str, backend_version: str) -> prefix = _role_prefix(role) return { "fpm_profile": sample["fpm_profile"], + "worker_type": "aggregated" if role == "agg" else role, "system": _role_hardware_sku(sample, role), "backend": sample["backend"], "backend_version": backend_version, diff --git a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py index 351b33525..d3ca843ec 100644 --- a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -115,6 +115,10 @@ def _resolve(self, request: ForwardPassPerfModelConfig, role: str) -> ForwardPas model: RustForwardPassPerfModel | None = None try: request_payload = RustForwardPassPerfModel.normalize_config(request) + if request_payload.get("fpm_profile") is not None: + from aisimulate_core.sdk.fpm_profile import _require_forward_pass_profile_memory + + _require_forward_pass_profile_memory(request_payload) cache_key = json.dumps(request_payload, sort_keys=True) cached = self._resolved.get(cache_key) if cached is not None: diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py index 49f65aa6a..2efce7941 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py @@ -52,6 +52,49 @@ def memory_fraction_kind(backend: str) -> str: return "of_free" if backend == "trtllm" else "of_total" +def estimate_grouped_cache_budget( + shape: ParallelShape, + *, + model_name: str, + hardware_sku: str, + backend: str, + backend_version: str, + fpm_profile: dict[str, Any], + worker_type: str = "aggregated", + context_length: int | None, + max_num_tokens: int, + max_batch_size: int, + memory_fraction: float, + systems_paths: list[str] | None = None, + model_controls: dict[str, str | int | bool] | None = None, + nextn: int = 0, + request_occupancy_tokens: int | None = None, +) -> dict[str, Any]: + """Use the canonical Rust grouped budget, including transient prefill pages.""" + return estimate_kv_cache( + model_name, + hardware_sku, + backend, + backend_version=backend_version, + max_num_tokens=max_num_tokens, + max_batch_size=max_batch_size, + context_length=context_length, + request_occupancy_tokens=request_occupancy_tokens, + memory_fraction_kind=memory_fraction_kind(backend), + memory_fraction_value=memory_fraction, + tp_size=shape.tp, + pp_size=shape.pp, + attention_dp_size=shape.dp, + moe_tp_size=shape.moe_tp, + moe_ep_size=shape.moe_ep, + nextn=nextn, + **(model_controls or {}), + systems_path=list(resolve_systems_paths(systems_paths)) if systems_paths is not None else None, + fpm_profile=fpm_profile, + worker_type=worker_type, + ) + + def resolve_backend_version( hardware_sku: str, backend: str, @@ -98,7 +141,9 @@ def estimate_kv_tokens( memory_fraction: float = DEFAULT_MEMORY_FRACTION, nextn: int = 0, fpm_profile: dict[str, Any] | None = None, + worker_type: str = "aggregated", model_controls: dict[str, str | int | bool] | None = None, + context_length: int | None = None, ) -> int | None: """Per-rank KV-cache capacity (in tokens) for ``shape``, or ``None`` when the shape leaves no KV budget (weights + activations already fill VRAM -> OOM). @@ -124,7 +169,8 @@ def estimate_kv_tokens( **(model_controls or {}), systems_path=(list(resolve_systems_paths(systems_paths)) if systems_paths is not None else None), allow_naive_fallback=False, - **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), + **({"fpm_profile": fpm_profile, "worker_type": worker_type} if fpm_profile is not None else {}), + **({"context_length": context_length} if context_length is not None else {}), ) except ValueError as exc: msg = str(exc) @@ -153,6 +199,7 @@ def feasible_shape_tokens( max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, fpm_profile: dict[str, Any] | None = None, + worker_type: str = "aggregated", model_controls: dict[str, str | int | bool] | None = None, nextn: int = 0, ) -> dict[ParallelShape, int]: @@ -176,7 +223,8 @@ def feasible_shape_tokens( max_num_tokens=max_num_tokens, max_batch_size=max_batch_size, memory_fraction=memory_fraction, - **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), + **({"fpm_profile": fpm_profile, "worker_type": worker_type} if fpm_profile is not None else {}), + **({"context_length": max_seq_len} if fpm_profile is not None else {}), **({"model_controls": model_controls} if model_controls else {}), **({"nextn": nextn} if nextn else {}), ) diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_load.py b/python/aisimulate/src/aisimulate/sweeper/kv_load.py index 00c9cb506..3b5a1bf08 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_load.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_load.py @@ -6,13 +6,13 @@ from __future__ import annotations import json -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import cache from typing import Any from .config import ENGINE_MODEL_CONTROL_FIELDS, Workload from .forward_pass_estimator import resolve_systems_paths -from .kv_estimate import estimate_kv_tokens +from .kv_estimate import DEFAULT_MEMORY_FRACTION, estimate_grouped_cache_budget, estimate_kv_tokens from .parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig @@ -28,6 +28,8 @@ class KVLoadResolution: concurrency: int concurrency_capacity: int role_capacity_tokens: dict[str, int] + role_capacity_bytes: dict[str, int] = field(default_factory=dict) + role_request_cache_bytes: dict[str, int] = field(default_factory=dict) @cache @@ -44,6 +46,8 @@ def _per_rank_capacity_tokens( memory_fraction: float, nextn: int, fpm_profile_json: str | None = None, + worker_type: str = "aggregated", + context_length: int | None = None, model_controls: tuple[tuple[str, str | int | bool], ...] = (), ) -> int: tokens = estimate_kv_tokens( @@ -57,7 +61,12 @@ def _per_rank_capacity_tokens( max_batch_size=max_batch_size, memory_fraction=memory_fraction, nextn=nextn, - **({"fpm_profile": json.loads(fpm_profile_json)} if fpm_profile_json is not None else {}), + **({"context_length": context_length} if context_length is not None else {}), + **( + {"fpm_profile": json.loads(fpm_profile_json), "worker_type": worker_type} + if fpm_profile_json is not None + else {} + ), **({"model_controls": dict(model_controls)} if model_controls else {}), ) if tokens is None: @@ -109,6 +118,8 @@ def _role_capacity_tokens( max_batch_size=int(sample[f"{role}_max_num_seqs"]), memory_fraction=float(sample[f"{role}_gpu_memory_utilization"]), fpm_profile_json=json.dumps(profile, sort_keys=True) if profile is not None else None, + worker_type="aggregated" if role == "agg" else role, + context_length=sample.get("context_length"), nextn=int(resolved.get("nextn", sample.get("aic_nextn")) or 0), model_controls=tuple( (name, resolved.get(name, sample.get(name))) @@ -124,6 +135,89 @@ def _role_capacity_tokens( return per_rank_usable_tokens * config.shape.dp * config.replicas +def _role_grouped_capacity( + sample: dict[str, Any], + *, + role: str, + config: ReplicaParallelConfig, + backend_version: str, + request_occupancy_tokens: int, +) -> tuple[int, int, int] | None: + resolved = sample.get("forward_pass_estimators", {}).get(role, {}).get("config") + if resolved is None: + timing = sample.get(f"{role}_timing_model") + if isinstance(timing, dict) and timing.get("type") == "external" and timing.get("provider") == "aic": + resolved = timing.get("config") + if not isinstance(resolved, dict): + raise ValueError(f"{role} external AIC timing config must be a mapping") + else: + resolved = {} + profile = resolved.get("fpm_profile", sample.get("fpm_profile")) + if profile is None: + return None + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + shape = config.shape + identity = { + "model": str(resolved.get("model", resolved.get("model_path", sample["model_name"]))), + "system": str(resolved.get("system", sample.get(f"{role}_hardware_sku") or sample["hardware_sku"])), + "backend": str(resolved.get("backend", sample["backend"])), + "backend_version": resolved.get("backend_version") or backend_version, + } + deployment = load_fpm_profile(profile).select( + **identity, + worker_type="aggregated" if role == "agg" else role, + tp_size=shape.tp, + pp_size=shape.pp, + attention_dp_size=shape.dp, + moe_tp_size=shape.moe_tp, + moe_ep_size=shape.moe_ep, + ) + if deployment.resources.cache_layout != "grouped": + return None + if role != "agg" or sample.get(f"{role}_num_gpu_blocks") is not None: + raise ValueError("grouped FPM cache requires aggregated replay without fixed scalar blocks") + roots = resolved.get("systems_paths") + if not roots and resolved.get("systems_path"): + roots = [resolved["systems_path"]] + controls = { + name: resolved.get(name, sample.get(name)) + for name in ENGINE_MODEL_CONTROL_FIELDS + if resolved.get(name, sample.get(name)) is not None + and not (name == "enable_eplb" and resolved.get(name, sample.get(name)) is False) + } + try: + budget = estimate_grouped_cache_budget( + shape, + model_name=identity["model"], + hardware_sku=identity["system"], + backend=identity["backend"], + backend_version=identity["backend_version"], + fpm_profile=profile, + worker_type="aggregated" if role == "agg" else role, + context_length=sample.get("context_length"), + request_occupancy_tokens=request_occupancy_tokens, + max_num_tokens=int(sample[f"{role}_max_num_batched_tokens"]), + max_batch_size=int(sample[f"{role}_max_num_seqs"]), + memory_fraction=float( + sample[f"{role}_gpu_memory_utilization"] + if sample.get(f"{role}_gpu_memory_utilization") is not None + else DEFAULT_MEMORY_FRACTION + ), + systems_paths=list(resolve_systems_paths(roots or sample.get("systems_paths"))), + model_controls=controls, + nextn=int(resolved.get("nextn", sample.get("aic_nextn")) or 0), + ) + except ValueError as exc: + if "no KV budget" in str(exc): + raise InfeasibleKVCapacity(str(exc)) from exc + raise + per_rank_bytes = budget["total_kv_size_bytes"] + per_request_bytes = budget["request_occupancy_cache_bytes"] + ranks = shape.dp * config.replicas + return per_rank_bytes * ranks, per_request_bytes, (per_rank_bytes // per_request_bytes) * ranks + + def resolve_kv_load( sample: dict[str, Any], *, @@ -153,20 +247,38 @@ def resolve_kv_load( else: raise TypeError(f"unsupported parallel config for KV load: {type(parallel_config).__name__}") - capacities = { - role: _role_capacity_tokens(sample, role=role, config=config, backend_version=backend_version) - for role, config in role_configs.items() - } expected_tokens_per_request = int(workload.isl) + int(workload.osl) // 2 if expected_tokens_per_request <= 0: raise InfeasibleKVCapacity( f"kv_load_ratio requires positive average tokens per request, got isl={workload.isl}, osl={workload.osl}" ) - concurrency_capacity = capacities[load_role] // expected_tokens_per_request + capacities = {} + capacity_bytes = {} + request_bytes = {} + request_capacities = {} + for role, config in role_configs.items(): + grouped = _role_grouped_capacity( + sample, + role=role, + config=config, + backend_version=backend_version, + request_occupancy_tokens=expected_tokens_per_request, + ) + if grouped is not None: + capacity_bytes[role], request_bytes[role], request_capacities[role] = grouped + else: + capacities[role] = _role_capacity_tokens( + sample, + role=role, + config=config, + backend_version=backend_version, + ) + request_capacities[role] = capacities[role] // expected_tokens_per_request + concurrency_capacity = request_capacities[load_role] if concurrency_capacity < 1: raise InfeasibleKVCapacity( - f"{load_role} KV capacity {capacities[load_role]} tokens cannot hold the " - f"estimated {expected_tokens_per_request} tokens per in-flight request" + f"{load_role} KV capacity cannot hold one request at " + f"{expected_tokens_per_request} average tokens per in-flight request" ) concurrency = max(1, int(float(ratio) * concurrency_capacity)) return KVLoadResolution( @@ -174,4 +286,6 @@ def resolve_kv_load( concurrency=concurrency, concurrency_capacity=concurrency_capacity, role_capacity_tokens=capacities, + role_capacity_bytes=capacity_bytes, + role_request_cache_bytes=request_bytes, ) diff --git a/python/aisimulate/src/aisimulate/sweeper/model_hw.py b/python/aisimulate/src/aisimulate/sweeper/model_hw.py index 8e9ee40ca..df053542f 100644 --- a/python/aisimulate/src/aisimulate/sweeper/model_hw.py +++ b/python/aisimulate/src/aisimulate/sweeper/model_hw.py @@ -133,8 +133,10 @@ def parallel_configs_for( role_runtime: dict[str, tuple[int, int, float] | tuple[int, int, float, int | None]] | None = None, systems_paths: list[str] | None = None, fpm_profile: dict[str, Any] | None = None, + worker_type: str = "aggregated", model_controls: dict[str, str | int | bool] | None = None, nextn: int = 0, + candidate_configs: list[ReplicaParallelConfig] | list[DisaggParallelConfig] | None = None, ) -> list[ReplicaParallelConfig] | list[DisaggParallelConfig]: """Resolve the model/hardware, then enumerate the parallel configs that fit the GPU budget and can hold a ``max_seq_len``-token sequence. @@ -164,14 +166,19 @@ def parallel_configs_for( seq_len = max_seq_len if max_seq_len is not None else profile.context_length if seq_len > profile.context_length: raise ValueError("max_seq_len exceeds FPM profile context_length") - configs = _profile_parallel_configs( - profile, - hardware_sku=hardware_sku, - backend=backend, - backend_version=backend_version, - gpu_budget=gpu_budget, - min_gpu_budget=min_gpu_budget, - deployment_mode=deployment_mode, + configs = ( + candidate_configs + if candidate_configs is not None + else _profile_parallel_configs( + profile, + hardware_sku=hardware_sku, + backend=backend, + backend_version=backend_version, + gpu_budget=gpu_budget, + min_gpu_budget=min_gpu_budget, + deployment_mode=deployment_mode, + worker_type=worker_type, + ) ) else: mh = resolve_model_hardware(model_name, hardware_sku, backend=backend, systems_paths=systems_paths) @@ -201,6 +208,7 @@ def parallel_configs_for( # KV-cache validity: keep configs whose every role-shape holds a max_seq_len sequence. def feasible_for(role: str, shapes): + profile_role = worker_type if role == "agg" else role runtime = (role_runtime or {}).get(role, (max_num_tokens, max_batch_size, memory_fraction)) if len(runtime) == 3: role_tokens, role_batch, role_memory = runtime @@ -211,10 +219,12 @@ def feasible_for(role: str, shapes): raise ValueError( "role_runtime values must be (tokens, batch, memory) or (tokens, batch, memory, fixed_tokens)" ) + grouped_shapes = set() if profile is not None: for shape in dict.fromkeys(shapes): deployment = profile.select( model=model_name, + worker_type=profile_role, system=hardware_sku, backend=backend, backend_version=backend_version, @@ -224,10 +234,47 @@ def feasible_for(role: str, shapes): moe_tp_size=shape.moe_tp, moe_ep_size=shape.moe_ep, ) + deployment.resources.require_memory() deployment.resources.validate_envelope(max_num_tokens=role_tokens, max_batch_size=role_batch) + if deployment.resources.runtime_memory is not None and fixed_tokens is not None: + raise ValueError("runtime FPM memory cannot use a fixed scalar token capacity") + if deployment.resources.cache_layout == "grouped": + grouped_shapes.add(shape) + grouped_feasible = set() + if grouped_shapes: + from .kv_estimate import estimate_grouped_cache_budget + + if deployment_mode != "agg" or nextn: + raise ValueError("grouped FPM cache requires aggregated replay without speculation") + if fixed_tokens is not None: + raise ValueError("grouped FPM cache cannot use a fixed scalar token capacity") + for shape in grouped_shapes: + try: + budget = estimate_grouped_cache_budget( + shape, + model_name=model_name, + hardware_sku=hardware_sku, + backend=backend, + backend_version=backend_version, + systems_paths=systems_paths, + fpm_profile=fpm_profile, + worker_type=profile_role, + context_length=seq_len, + max_num_tokens=role_tokens, + max_batch_size=role_batch, + memory_fraction=role_memory, + model_controls=model_controls, + ) + except ValueError as exc: + if "no KV budget" in str(exc): + continue + raise + if budget["total_kv_size_bytes"] >= budget["request_peak_cache_bytes"]: + grouped_feasible.add(shape) + shapes = [shape for shape in shapes if shape not in grouped_shapes] if fixed_tokens is not None: return {shape: fixed_tokens for shape in dict.fromkeys(shapes) if fixed_tokens > seq_len} - return feasible_shape_tokens( + linear_feasible = feasible_shape_tokens( shapes, model_name=model_name, hardware_sku=hardware_sku, @@ -238,10 +285,11 @@ def feasible_for(role: str, shapes): max_num_tokens=role_tokens, max_batch_size=role_batch, memory_fraction=role_memory, - **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), + **({"fpm_profile": fpm_profile, "worker_type": profile_role} if fpm_profile is not None else {}), **({"model_controls": model_controls} if model_controls else {}), **({"nextn": nextn} if nextn else {}), ) + return set(linear_feasible) | grouped_feasible if deployment_mode == "agg": feasible = feasible_for("agg", [c.shape for c in configs]) @@ -267,37 +315,44 @@ def _profile_parallel_configs( gpu_budget: int, min_gpu_budget: int | None, deployment_mode: str, + decode_hardware_sku: str | None = None, + worker_type: str = "aggregated", ): """Enumerate only declared deployment shapes; replicas do not change a cell.""" - deployments = [ - deployment - for deployment in profile.deployments - if deployment.system == hardware_sku - and deployment.backend == backend - and deployment.backend_version == backend_version - ] - if not deployments: - raise ValueError(f"no FPM deployment profile for {hardware_sku}/{backend}/{backend_version}") - replicas = [] - for deployment in deployments: - shape = ParallelShape( - tp=deployment.tp, - pp=deployment.pp, - dp=deployment.dp, - moe_tp=deployment.moe_tp, - moe_ep=deployment.moe_ep, - ) - replicas.extend( - ReplicaParallelConfig(shape=shape, replicas=count) - for count in range(1, gpu_budget // shape.gpus_per_worker + 1) - ) + + def replicas_for(system, role): + deployments = [ + deployment + for deployment in profile.deployments + if deployment.system == system + and deployment.backend == backend + and deployment.backend_version == backend_version + and deployment.worker_type in (None, role) + ] + replicas = [] + for deployment in deployments: + shape = ParallelShape( + tp=deployment.tp, + pp=deployment.pp, + dp=deployment.dp, + moe_tp=deployment.moe_tp, + moe_ep=deployment.moe_ep, + ) + replicas.extend( + ReplicaParallelConfig(shape=shape, replicas=count) + for count in range(1, gpu_budget // shape.gpus_per_worker + 1) + ) + return replicas + + replicas = replicas_for(hardware_sku, worker_type if deployment_mode == "agg" else "prefill") if deployment_mode == "agg": configs = replicas elif deployment_mode == "disagg": + decode_replicas = replicas_for(decode_hardware_sku or hardware_sku, "decode") configs = [ DisaggParallelConfig(prefill=prefill, decode=decode) for prefill in replicas - for decode in replicas + for decode in decode_replicas if prefill.total_gpus + decode.total_gpus <= gpu_budget ] else: diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index a2da1af57..4e3b97fb9 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -16,6 +16,7 @@ from pydantic import BaseModel from ..config.common import ENGINE_MODEL_CONTROL_FIELDS, is_active_engine_model_control +from ..fpm_profile import load_fpm_profile from ..power import POWER_FIELDS, normalize_power_summary from .provider import AdapterReplaySpec, JSONValue, RuntimeHookSpec @@ -260,6 +261,7 @@ class RunnerCapabilities: supported_engine_model_controls: tuple[str, ...] = () supports_mtp_expected_acceptance: bool = False supports_state_cache: bool = False + supports_grouped_kv_cache: bool = False def supports_backend_topology(self, backend: str, topology: str) -> bool: """Return whether a backend/topology pair is supported. @@ -316,7 +318,11 @@ def require_compatible(self, spec: ReplaySpec) -> None: raise ValueError("cached_prefix_tokens is unsupported for AFD") if not self.supports_cached_prefix_tokens: raise ValueError("runner does not support cached_prefix_tokens; use --stack engine") - for args in (deployment.agg_engine_args, deployment.prefill_engine_args, deployment.decode_engine_args): + for role, args in ( + ("aggregated", deployment.agg_engine_args), + ("prefill", deployment.prefill_engine_args), + ("decode", deployment.decode_engine_args), + ): if not isinstance(args, Mapping): continue rank = args.get("rank", args) @@ -327,6 +333,35 @@ def require_compatible(self, spec: ReplaySpec) -> None: if not isinstance(identity, Mapping): # The timing validator handles malformed identities; still inspect flat controls. identity = {} + profile = identity.get("fpm_profile") or rank.get("aic_fpm_profile") or rank.get("fpm_profile") + grouped_profile = False + if not self.supports_grouped_kv_cache and profile is not None: + # Match the runner's canonical and flat identity defaults before + # inspecting resources; another profile deployment is irrelevant. + selected = load_fpm_profile(profile).select( + model=identity.get("model", args.get("aic_model_path")), + system=identity.get("system", args.get("aic_system")), + backend=identity.get("backend", rank.get("aic_backend", rank.get("backend", deployment.backend))), + backend_version=identity.get( + "backend_version", + rank.get("aic_backend_version", rank.get("backend_version", deployment.backend_version)), + ), + worker_type=identity.get("worker_type", rank.get("worker_type", role)), + tp_size=identity.get("tp", args.get("tensor_parallel_size", args.get("aic_tp_size", 1))), + pp_size=identity.get("pp", rank.get("aic_pp_size", 1)), + attention_dp_size=identity.get( + "attention_dp", args.get("dp_size", args.get("aic_attention_dp_size", 1)) + ), + moe_tp_size=identity.get("moe_tp_size", rank.get("moe_tp_size", rank.get("aic_moe_tp_size"))), + moe_ep_size=identity.get("moe_ep_size", rank.get("moe_ep_size", rank.get("aic_moe_ep_size"))), + ) + grouped_profile = selected.resources.cache_layout == "grouped" + if not self.supports_grouped_kv_cache and ( + grouped_profile + or rank.get("kv_cache_groups") is not None + or rank.get("kv_cache_capacity_bytes") is not None + ): + raise ValueError("runner does not support grouped FPM caches; use --stack engine") active_controls = [ name for name in ENGINE_MODEL_CONTROL_FIELDS diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py index c1f8c146d..34e278776 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search.py +++ b/python/aisimulate/src/aisimulate/sweeper/search.py @@ -681,7 +681,11 @@ def _materialize_one( sample["kv_load_ratio"] = resolution.ratio sample["kv_load_concurrency_capacity"] = resolution.concurrency_capacity load_role = "decode" if sample["deployment_mode"] == "disagg" else "agg" - sample["kv_load_capacity_tokens"] = resolution.role_capacity_tokens[load_role] + if load_role in resolution.role_capacity_tokens: + sample["kv_load_capacity_tokens"] = resolution.role_capacity_tokens[load_role] + if load_role in resolution.role_capacity_bytes: + sample["kv_load_capacity_bytes"] = resolution.role_capacity_bytes[load_role] + sample["kv_load_request_cache_bytes"] = resolution.role_request_cache_bytes[load_role] for role, tokens in resolution.role_capacity_tokens.items(): sample[f"{role}_kv_capacity_tokens"] = tokens if concurrency is not None: diff --git a/python/aisimulate/src/aisimulate/sweeper/search_space.py b/python/aisimulate/src/aisimulate/sweeper/search_space.py index f1534b5df..06b3d5e51 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search_space.py +++ b/python/aisimulate/src/aisimulate/sweeper/search_space.py @@ -29,7 +29,13 @@ ) from .config import ENGINE_MODEL_CONTROL_FIELDS, SmartSearchConfig from .kv_estimate import NoPerfDatabase, resolve_backend_version -from .model_hw import ModelHardware, NoViableParallelConfig, parallel_configs_for, resolve_model_hardware +from .model_hw import ( + ModelHardware, + NoViableParallelConfig, + _profile_parallel_configs, + parallel_configs_for, + resolve_model_hardware, +) from .parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig from .replay import RunnerCapabilities @@ -164,6 +170,42 @@ def _parse_parallel_entry(entry: dict[str, Any], deployment_mode: str): ) +def _fpm_parallel_configs(search_space, deployment_mode: str): + """Apply declared topology domains before grouped compatibility or resource checks.""" + from aisimulate.fpm_profile import load_fpm_profile + + profile = load_fpm_profile(search_space.fpm_profile) + configs = _profile_parallel_configs( + profile, + hardware_sku=search_space.hardware_sku_for("agg" if deployment_mode == "agg" else "prefill"), + decode_hardware_sku=(search_space.hardware_sku_for("decode") if deployment_mode == "disagg" else None), + backend="vllm", + backend_version=search_space.requested_backend_version("vllm"), + gpu_budget=search_space.gpu_budget, + min_gpu_budget=search_space.min_gpu_budget, + deployment_mode=deployment_mode, + ) + raw_pinned = search_space.parallel_configs_by_mode.get(deployment_mode, search_space.parallel_configs) + if raw_pinned: + pinned = {_parse_parallel_entry(entry, deployment_mode) for entry in raw_pinned} + configs = [config for config in configs if config in pinned] + for role, entries in search_space.parallel_custom_configs_by_mode.get(deployment_mode, {}).items(): + choices = {_replica_from_dict(entry) for entry in entries} + configs = [config for config in configs if _parallel_role(config, role) in choices] + domains = search_space.parallel_independent_by_mode.get(deployment_mode, {}) + bounds = search_space.parallel_independent_log_ranges_by_mode.get(deployment_mode, {}) + return [ + config + for config in configs + if all( + bounds[name][0] <= value <= bounds[name][1] + if name in bounds + else domains.get(name) is None or value in domains[name] + for name, value in _parallel_leaf_values(config).items() + ) + ] + + def branch_knob_choices(search_space, deployment_mode: str) -> dict[str, list[Any]]: """Backend-owned atomic knobs for one deployment branch.""" names = _engine_knobs(deployment_mode) @@ -260,6 +302,7 @@ def _heterogeneous_disagg_configs( *, backend: str, max_seq_len: int | None, + candidate_configs: list[DisaggParallelConfig] | None = None, ) -> list[DisaggParallelConfig]: """Enumerate each P/D role against its own hardware, then apply the shared budget.""" role_hardware = {role: search_space.hardware_sku_for(role) for role in ("prefill", "decode")} @@ -285,6 +328,7 @@ def _heterogeneous_disagg_configs( hardware, gpu_budget=search_space.gpu_budget, deployment_mode="agg", + worker_type=role, backend=backend, backend_version=backend_version, min_gpu_budget=None, @@ -293,6 +337,11 @@ def _heterogeneous_disagg_configs( **_engine_memory_kwargs(search_space), **_estimator_root_kwargs(search_space, role), **({"fpm_profile": search_space.fpm_profile} if search_space.fpm_profile is not None else {}), + **( + {"candidate_configs": list(dict.fromkeys(getattr(config, role) for config in candidate_configs))} + if candidate_configs is not None + else {} + ), ) except (NoPerfDatabase, NoViableParallelConfig) as exc: raise type(exc)(f"{role} hardware_sku={hardware!r}: {exc}") from exc @@ -587,6 +636,7 @@ def enumerate_branches( custom_by_role = { role: tuple(_replica_from_dict(entry) for entry in entries) for role, entries in raw_custom.items() } + profile_configs = _fpm_parallel_configs(ss, deployment_mode) if ss.fpm_profile is not None else None def matches_custom(config: _ParallelConfig) -> bool: return all(_parallel_role(config, role) in choices for role, choices in custom_by_role.items()) @@ -607,6 +657,7 @@ def matches_custom(config: _ParallelConfig) -> bool: ss, backend=backend, max_seq_len=max_seq_len, + candidate_configs=profile_configs, ) else: legal = parallel_configs_for( @@ -622,6 +673,7 @@ def matches_custom(config: _ParallelConfig) -> bool: **_engine_memory_kwargs(ss), **_estimator_root_kwargs(ss, "agg" if deployment_mode == "agg" else "prefill"), **({"fpm_profile": ss.fpm_profile} if ss.fpm_profile is not None else {}), + **({"candidate_configs": profile_configs} if profile_configs is not None else {}), ) except (NoPerfDatabase, NoViableParallelConfig): continue # backend unusable for this mode -> drop it from the search diff --git a/python/aisimulate/src/aisimulate_core/_native.pyi b/python/aisimulate/src/aisimulate_core/_native.pyi index 6f6bf7699..a3d66e726 100644 --- a/python/aisimulate/src/aisimulate_core/_native.pyi +++ b/python/aisimulate/src/aisimulate_core/_native.pyi @@ -173,6 +173,8 @@ class RustForwardPassPerfModel: @staticmethod def normalize_config(config_json: str) -> str: ... @staticmethod + def estimate_cache_budget(config_json: str, budget_json: str) -> str: ... + @staticmethod def _normalize_fpm_options( options_json: str, fmha_quant_mode: str | None = None, @@ -192,6 +194,9 @@ class RustForwardPassPerfModel: def static_phase_latency(self, batch_size: int, input_tokens: int, output_tokens: int, prefill: bool) -> float: ... def static_phase_diagnostics(self, batch_size: int, context_length: int, prefix: int, prefill: bool) -> str: ... def diagnostics(self) -> str: ... + def fpm_query_coverage(self) -> str: ... + def predict_decode_latency_total(self, batch_size: int, total_past_kv_tokens: int) -> float: ... + def fpm_decode_kv_ceiling(self) -> int | None: ... def regression_store_diagnostics(self) -> str: ... def min_correction_factor(self) -> float | None: ... def max_correction_factor(self) -> float | None: ... diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index 39a0f9ac4..61cadd15b 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -1,12 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Explicit identity and resource bounds for decoder-only FPM simulation. +"""Explicit identity, cache geometry and memory evidence for FPM simulation. Profiles describe every rank of a deployment without constructing operations. -Resource values are declarations supplied by the caller, not GPU measurements -or automatically inferred checkpoint metadata. Their provenance must say how -they were obtained; CUDA graph reservations are supplied separately at runtime. +Memory may remain pending during collection or use complete declared non-KV +bytes or a runtime cache allocation. Provenance records how values were obtained. """ from __future__ import annotations @@ -14,7 +13,7 @@ import json from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_serializer, model_validator from . import quantization @@ -28,39 +27,133 @@ class _ProfileModel(BaseModel): model_config = ConfigDict(extra="forbid", strict=True) +class FpmCacheGroup(_ProfileModel): + """Rank-local physical cache pages for one attention or convolution group. + + Page bytes include all group layers and runtime padding. The retention + window does not replace the logical context used for FPM timing queries. + """ + + name: _Nonempty + kind: Literal["attention", "convolution"] + num_layers: Annotated[int, Field(gt=0, le=2**32 - 1)] + block_size_tokens: Annotated[int, Field(gt=0, le=2**32 - 1)] + page_size_bytes: Annotated[int, Field(gt=0, le=2**53)] + sliding_window: Annotated[int, Field(gt=0, le=2**32 - 1)] | None = None + + @model_validator(mode="after") + def _window(self) -> FpmCacheGroup: + if self.kind == "convolution" and self.sliding_window is None: + raise ValueError("convolution cache groups require a sliding_window") + return self + + +class FpmRuntimeMemoryProfile(_ProfileModel): + """Observed rank-local cache capacity for the exact runtime configuration. + + This capacity already accounts for runtime graph and workspace reservations. + It is not a non-KV memory limit and must not be rescaled to another setting. + """ + + kv_cache_bytes: Annotated[int, Field(gt=0, le=2**53)] + gpu_memory_utilization: Annotated[float, Field(gt=0, le=1, allow_inf_nan=False)] + max_model_len: Annotated[int, Field(gt=0, le=2**64 - 1)] + provenance: _Nonempty + + class FpmResourceProfile(_ProfileModel): - """Conservative rank-local bounds at a declared scheduler envelope. + """Rank-local cache geometry and optional memory evidence. ``max_num_tokens`` and ``max_batch_size`` apply to each attention-DP rank, - not the summed worker iteration used to query FPM timings. Overheads must - exclude the separately configured CUDA graph reservation. + not the summed worker iteration used to query FPM timings. Complete declared + non-KV bytes establish conservative bounds and exclude separate CUDA graph + reservations. Runtime capacity instead requires the exact recorded settings. """ - weights_bytes: _Bytes - activations_bytes: _Bytes - runtime_overhead_bytes: _Bytes - comm_overhead_bytes: _Bytes - kv_bytes_per_token: Annotated[int, Field(gt=0, le=2**53)] - cache_layout: Literal["linear"] + weights_bytes: _Bytes | None = None + activations_bytes: _Bytes | None = None + runtime_overhead_bytes: _Bytes | None = None + comm_overhead_bytes: _Bytes | None = None + runtime_memory: FpmRuntimeMemoryProfile | None = None + kv_bytes_per_token: Annotated[int, Field(gt=0, le=2**53)] | None = None + cache_layout: Literal["linear", "grouped"] + cache_groups: list[FpmCacheGroup] = Field(default_factory=list) max_num_tokens: _PositiveInt max_batch_size: _PositiveInt provenance: _Nonempty @model_validator(mode="after") def _exact_total(self) -> FpmResourceProfile: - if self.non_kv_bytes > 2**53: + declared = self._declared_bytes() + if sum(value for value in declared if value is not None) > 2**53: raise ValueError("total non-KV resource bytes must not exceed 2**53") + if self.runtime_memory is not None and any(value is not None for value in declared): + raise ValueError("runtime_memory cannot be combined with declared non-KV resource bytes") + if self.cache_layout == "linear": + if self.kv_bytes_per_token is None or self.cache_groups: + raise ValueError("linear cache requires kv_bytes_per_token and no cache_groups") + elif self.kv_bytes_per_token is not None or not self.cache_groups: + raise ValueError("grouped cache requires cache_groups without a scalar kv_bytes_per_token") + if len({group.name for group in self.cache_groups}) != len(self.cache_groups): + raise ValueError("cache group names must be unique") return self + @model_serializer(mode="wrap") + def _serialize_resources(self, handler): + # Do not add empty fields to saved linear profiles or change plan IDs. + result = handler(self) + if not self.cache_groups: + result.pop("cache_groups", None) + if self.kv_bytes_per_token is None: + result.pop("kv_bytes_per_token", None) + for name in ( + "weights_bytes", + "activations_bytes", + "runtime_overhead_bytes", + "comm_overhead_bytes", + "runtime_memory", + ): + if getattr(self, name) is None: + result.pop(name, None) + return result + + def _declared_bytes(self) -> tuple[int | None, ...]: + return self.weights_bytes, self.activations_bytes, self.runtime_overhead_bytes, self.comm_overhead_bytes + + @property + def memory_source(self) -> Literal["pending", "declared", "runtime"]: + if self.runtime_memory is not None: + return "runtime" + return "declared" if all(value is not None for value in self._declared_bytes()) else "pending" + + @property + def memory_ready(self) -> bool: + return self.memory_source != "pending" + + def require_memory(self) -> None: + if not self.memory_ready: + raise ValueError( + "FPM memory is pending runtime profiling; collect and finalize runtime memory or provide " + "all four declared non-KV resource values before simulation" + ) + @property def non_kv_bytes(self) -> int: - return self.weights_bytes + self.activations_bytes + self.runtime_overhead_bytes + self.comm_overhead_bytes + self.require_memory() + if self.runtime_memory is not None: + raise ValueError("runtime memory records cache capacity, not a non-KV byte breakdown") + return sum(value for value in self._declared_bytes() if value is not None) def validate_envelope(self, *, max_num_tokens: int, max_batch_size: int) -> None: for name, requested in (("max_num_tokens", max_num_tokens), ("max_batch_size", max_batch_size)): if isinstance(requested, bool) or not isinstance(requested, int) or requested <= 0: raise ValueError(f"{name} must be a positive integer, got {requested!r}") limit = getattr(self, name) + if self.runtime_memory is not None and requested != limit: + raise ValueError( + f"runtime memory requires the exact recorded rank-local {name}={limit}; " + "collect new runtime memory for changed scheduler settings" + ) if requested > limit: raise ValueError( f"FPM resource envelope exceeded: requested rank-local {name}={requested}, " @@ -69,7 +162,7 @@ def validate_envelope(self, *, max_num_tokens: int, max_batch_size: int) -> None class FpmDeploymentProfile(_ProfileModel): - """One exact hardware, runtime, topology, precision and resource identity.""" + """One hardware, runtime, topology, precision and optional serving-role identity.""" system: _Nonempty backend: Literal["vllm"] @@ -89,8 +182,17 @@ class FpmDeploymentProfile(_ProfileModel): attention_backend: _Nonempty = "auto" enable_wideep: Literal[False] = False enable_eplb: Literal[False] = False + # Absent in historical shared profiles; never infer role-specific evidence. + worker_type: Literal["prefill", "decode", "aggregated"] | None = None resources: FpmResourceProfile + @model_serializer(mode="wrap") + def _preserve_legacy_identity(self, handler): + result = handler(self) + if self.worker_type is None: + result.pop("worker_type", None) + return result + @field_validator("pp", "cp", mode="before") @classmethod def _integer_dimensions(cls, value: Any) -> int: @@ -199,7 +301,7 @@ def _pinned_revision(cls, value: str) -> str: @model_validator(mode="after") def _unique_deployments(self) -> FpmModelProfile: - identities = set() + identities: dict[tuple, set[str | None]] = {} for deployment in self.deployments: tp = deployment.dp == 1 and deployment.moe_ep == 1 and deployment.moe_tp == deployment.tp dep = deployment.tp == 1 and deployment.moe_tp == 1 and deployment.moe_ep == deployment.dp @@ -215,11 +317,16 @@ def _unique_deployments(self) -> FpmModelProfile: deployment.backend_version, deployment.parallel_tuple, ) - if identity in identities: + roles = identities.setdefault(identity, set()) + if deployment.worker_type in roles: + raise ValueError( + "duplicate FPM deployment identity; specify one precision/resource profile per topology and role" + ) + if roles and (None in roles or deployment.worker_type is None): raise ValueError( - "duplicate FPM deployment identity; specify one precision/resource profile per topology" + "ambiguous FPM deployment identity: cannot mix legacy shared and role-specific profiles" ) - identities.add(identity) + roles.add(deployment.worker_type) if deployment.moe_ep > 1 and (self.num_experts == 0 or self.num_experts % deployment.moe_ep): raise ValueError("num_experts must be positive and divisible by moe_ep for expert parallelism") return self @@ -237,7 +344,10 @@ def select( moe_tp_size: int | None = None, moe_ep_size: int | None = None, cp_size: int = 1, + worker_type: str = "aggregated", ) -> FpmDeploymentProfile: + if worker_type not in ("prefill", "decode", "aggregated"): + raise ValueError("worker_type must be prefill, decode, or aggregated") if model != self.model: raise ValueError(f"FPM profile model identity mismatch: requested {model!r}, profile={self.model!r}") shape = ( @@ -256,11 +366,13 @@ def select( and deployment.backend == backend and deployment.backend_version == backend_version and deployment.parallel_tuple == shape + and deployment.worker_type in (None, worker_type) ): return deployment raise ValueError( f"no matching FPM deployment profile for {model!r}, {system}/{backend}/{backend_version}, " - f"(tp,pp,dp,moe_tp,moe_ep,cp)={shape}; provide explicit resources and precision for that deployment" + f"(tp,pp,dp,moe_tp,moe_ep,cp)={shape}, worker_type={worker_type}; " + "provide explicit resources and precision for that deployment" ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 505320c8e..630b9da31 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -482,6 +482,7 @@ def compile_engine( forward_model: str | None = None, decoder_replay: bool = False, fpm_profile: dict | str | FpmModelProfile | None = None, + worker_type: str = "aggregated", fpm_interpolation: str | None = None, cp_size: int = 1, database_mode: str | None = None, @@ -579,6 +580,7 @@ def compile_engine( system=system, backend=backend, backend_version=literal_version, + worker_type=worker_type, tp_size=tp_size, pp_size=pp_size, attention_dp_size=attention_dp_size, @@ -619,6 +621,7 @@ def compile_engine( strict_provenance=strict_provenance, fpm_parquet_path=fpm_parquet_path, fpm_fmha_quant_mode=fpm_fmha_quant_mode, + worker_type=worker_type, fpm_options=fpm_config.options, ) return bytes(aisimulate_core.engine_spec_bincode_from_json(spec_json)) @@ -744,6 +747,8 @@ def compile_engine( fpm_profile=profile.model_dump_json(), estimator_config=json.dumps({"fpm_interpolation": {**fpm_config.options, "method": interpolation}}), ) + if deployment.worker_type is not None: + spec["engine"]["extra"]["worker_type"] = worker_type spec_json = json.dumps(spec) return bytes(aisimulate_core.engine_spec_bincode_from_json(spec_json)) @@ -761,6 +766,7 @@ def _direct_fpm_spec_json( strict_provenance: bool | None, fpm_parquet_path: str | None = None, fpm_fmha_quant_mode: str | None = None, + worker_type: str = "aggregated", fpm_options: dict | None = None, ) -> str: """Build whole-forward timing operations from metadata, without a graph.""" @@ -801,6 +807,7 @@ def _direct_fpm_spec_json( "transfer_policy": _transfer_policy_tokens(None, transfer_policy), "extra": { "fpm_profile": profile.model_dump_json(), + **({"worker_type": worker_type} if deployment.worker_type is not None else {}), "estimator_config": json.dumps({"fpm_interpolation": controls}), }, } @@ -813,7 +820,9 @@ def phase_op(phase: str) -> dict: "model_path": profile.model, "match_identity": identity, "original_fmha_quant_mode": deployment.fmha_quant_mode if fpm_fmha_quant_mode is not None else None, - "weight_bytes": deployment.resources.weights_bytes, + # Direct interpolation ignores this legacy operation field. + # Zero here is not an inferred resource or memory bound. + "weight_bytes": deployment.resources.weights_bytes or 0, "verify_width": 1, "sol_ops": [], "interpolation": "direct", diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_model_metadata.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_model_metadata.py new file mode 100644 index 000000000..8d032ee9e --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_model_metadata.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint precision reconciliation shared by FPM onboarding and collection. + +These functions do not change the general SDK's model-loading precedence. +""" + +from __future__ import annotations + +from typing import Any + +from .utils import ( + _attach_hf_quant_config, + _infer_quantization_fields, + _normalize_kv_cache_algo, + _normalize_quant_algo, +) + + +def inherit_decoder_precision(document: dict[str, Any], decoder: dict[str, Any]) -> tuple[dict[str, Any], list[str]]: + """Inherit a shared precision declaration only when the decoder lacks it.""" + raw = dict(decoder) + inherited = [] + for keys, inferred_keys in ( + (("dtype", "torch_dtype"), ()), + (("quantization_config", "hf_quant_config", "quant_algo"), ("quant_dynamic", "kv_cache_quant_algo")), + ): + if any(decoder.get(key) is not None for key in keys): + continue + for key in keys: + if document.get(key) is not None: + raw[key] = document[key] + inherited.append(key) + for key in inferred_keys: + if decoder.get(key) is None and document.get(key) is not None: + raw[key] = document[key] + inherited.append(key) + return raw, inherited + + +def attach_quantization_sidecar(raw: dict[str, Any], sidecar: dict[str, Any]) -> dict[str, Any]: + """Reconcile one effective decoder with its adjacent ModelOpt sidecar.""" + section = sidecar.get("quantization") + if section is not None and not isinstance(section, dict): + raise ValueError("hf_quant_config.json quantization must be an object") + for key in ("quant_algo", "quantization_algo", "kv_cache_quant_algo"): + if section and section.get(key) is not None and (not isinstance(section[key], str) or not section[key].strip()): + raise ValueError(f"hf_quant_config.json quantization.{key} must be a nonempty string") + if ( + section + and section.get("quant_algo") + and section.get("quantization_algo") + and _normalize_quant_algo(section["quant_algo"]) != _normalize_quant_algo(section["quantization_algo"]) + ): + raise ValueError("conflicting quant_algo and quantization_algo in hf_quant_config.json") + if raw.get("hf_quant_config") is not None and raw["hf_quant_config"] != sidecar: + raise ValueError("conflicting inline hf_quant_config and adjacent hf_quant_config.json") + # Validate the inline declaration independently: a repeated sidecar must not + # hide contradictory quantization_config fields through SDK precedence. + inline = _infer_quantization_fields({**raw, "hf_quant_config": None}) + merged = _attach_hf_quant_config(dict(raw), sidecar) + # Preserve even unknown/empty sidecars so they cannot imply BF16 weights. + merged["hf_quant_config"] = sidecar + inferred = _infer_quantization_fields(merged) + scheme = (raw.get("quantization_config") or {}).get("kv_cache_scheme") + if isinstance(scheme, str) and scheme.strip().upper() == "FP8": + inline.setdefault("kv_cache_quant_algo", "fp8") + elif isinstance(scheme, dict) and scheme.get("num_bits") == 8 and scheme.get("type") in {"float", "int"}: + inline.setdefault("kv_cache_quant_algo", "fp8" if scheme["type"] == "float" else "int8") + if "kv_cache_quant_algo" in inline: + inferred.setdefault("kv_cache_quant_algo", inline["kv_cache_quant_algo"]) + for key in ("quant_algo", "kv_cache_quant_algo"): + normalize = _normalize_quant_algo if key == "quant_algo" else _normalize_kv_cache_algo + after = normalize(inferred.get(key)) + for value in (raw.get(key), inline.get(key)): + before = normalize(value) + if key == "quant_algo" and (before in {"modelopt", "compressed-tensors"} or after == "mixed_precision"): + continue # Containers and per-layer layouts are not scalar algorithms. + if key == "kv_cache_quant_algo": + before = "bfloat16" if before == "none" else before + after = "bfloat16" if after == "none" else after + if before is not None and after is not None and before != after: + raise ValueError(f"conflicting {key} in model config and hf_quant_config.json: {before!r} != {after!r}") + dynamic = raw.get("quant_dynamic") + if dynamic is not None and type(dynamic) is not bool: + raise ValueError("quant_dynamic must be a boolean or null") + if dynamic is not None and inferred.get("quant_dynamic") is not None and dynamic != inferred["quant_dynamic"]: + raise ValueError("conflicting quant_dynamic in model config and quantization metadata") + for key, value in inferred.items(): + # Null is unknown. Container markers can be refined by the sidecar; + # actual conflicting scalar declarations were rejected above. + if merged.get(key) is None or (key == "quant_algo" and merged[key] in {"modelopt", "compressed-tensors"}): + merged[key] = value + return merged diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py index cdbc4779a..1ff25cd31 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py @@ -10,12 +10,38 @@ from pathlib import Path from aisimulate_core.fpm_profile import _MUTABLE_REFERENCES +from aisimulate_core.fpm_profile import FpmCacheGroup as FpmCacheGroup from aisimulate_core.fpm_profile import FpmDeploymentProfile as FpmDeploymentProfile from aisimulate_core.fpm_profile import FpmModelProfile as FpmModelProfile from aisimulate_core.fpm_profile import FpmResourceProfile as FpmResourceProfile +from aisimulate_core.fpm_profile import FpmRuntimeMemoryProfile as FpmRuntimeMemoryProfile from aisimulate_core.fpm_profile import load_fpm_profile as load_fpm_profile +def _require_forward_pass_profile_memory(config: dict) -> None: + """Check simulation readiness after canonical identity normalization. + + Standalone timing queries can use pending resources; simulation consumers + call this before loading timings so a new model gets a memory finalization + diagnostic even when collection has not produced any timing data yet. + """ + if config.get("fpm_profile") is None: + return + deployment = load_fpm_profile(config["fpm_profile"]).select( + model=config["model"], + system=config["system"], + backend=config["backend"], + backend_version=config["backend_version"], + worker_type=config.get("worker_type", "aggregated"), + tp_size=config["tp"], + pp_size=config["pp"], + attention_dp_size=config["attention_dp"], + moe_tp_size=config.get("moe_tp_size"), + moe_ep_size=config.get("moe_ep_size"), + ) + deployment.resources.require_memory() + + def _quantization_from_engine_config(config_json: str) -> dict[str, str]: """Restore exact profile modes after checking the native wire dtypes. @@ -33,6 +59,7 @@ def _quantization_from_engine_config(config_json: str) -> dict[str, str]: system=config["system_name"], backend=config["backend"], backend_version=config["backend_version"], + worker_type=config["extra"].get("worker_type", "aggregated"), tp_size=config["tp_size"], pp_size=config["pp_size"], attention_dp_size=1 if config.get("attention_dp_size") is None else config["attention_dp_size"], @@ -123,6 +150,7 @@ def _validate_forward_pass_profile(config_json: str) -> str: system=config["system"], backend=config["backend"], backend_version=version, + worker_type=config["worker_type"], tp_size=config["tp"], pp_size=config["pp"], attention_dp_size=config["attention_dp"], diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory.py b/python/aisimulate/src/aisimulate_core/sdk/memory.py index b6e5b31b2..15116bd2e 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory.py @@ -6,17 +6,19 @@ ``estimate_state_cache`` re-exports the shared config-driven state-footprint API. It is independent of the pool-capacity algorithms below. -The single source of truth for the rank-local KV-cache capacity estimate. The -Rust ``aisimulate_core::memory::estimate_kv_cache`` is a pure forwarder that -calls :func:`estimate_kv_cache` here and rebuilds its result, so all of the -math -- fraction + tolerance validation, the native AIC memory model, the naive -fallback, and the tolerance margin -- lives in this module. +This module implements scalar rank-local KV-cache capacity estimation. The +legacy Rust ``aisimulate_core::memory::estimate_kv_cache`` forwards here and +rebuilds the scalar result. Grouped FPM profiles delegate budget arithmetic, +page rounding and retention footprints to the canonical Rust +``ForwardPassPerfModelConfig::estimate_cache_budget`` API; they have no scalar +token capacity. Resource sources: - **FPM profile**: caller-declared rank-local resource bounds bypass model - construction and timing data. The same budget arithmetic below applies, - after checking the scheduler against the declared resource envelope. + construction and timing data. Linear profiles use the scalar arithmetic below; + grouped profiles use the native cache-budget API. Both check the scheduler + against the declared resource envelope. - **Native** (:class:`KVCacheEstimator`): ``from_request`` reuses AIC's full backend memory model (``BaseBackend._get_memory_usage`` plus the model's @@ -33,7 +35,7 @@ metadata to compute the weights or the per-token KV size (rather than guessing with a placeholder). -The budget math is factored into :func:`kv_cache_budget_bytes` so the deferred +The scalar budget math is factored into :func:`kv_cache_budget_bytes` so the deferred consolidation of ``InferenceSummary._check_and_set_kv_cache_oom`` onto a single formula is a small follow-up (tracked in #1208). """ @@ -65,7 +67,7 @@ _ONE_GIB = 1 << 30 _MAX_EXACT_BYTE_COUNT = 1 << 53 -# A KV byte-budget -> token-count inverse. Every estimation path produces one of +# A KV byte-budget -> token-count inverse. Every scalar estimation path produces one of # these (native: the model's ``get_kvcache_max_tokens``; naive: # ``NaiveKVCacheEstimator.get_kvcache_max_tokens``; the linear # ``_linear_tokens_from_bytes`` is the default for constant-per-token growth), so @@ -155,9 +157,9 @@ def _validate_cuda_graph_reservation(cuda_graph_reserved_bytes: int) -> None: # --------------------------------------------------------------------------- # -# Shared budget math. +# Shared scalar budget math. # -# The single formula both the estimate and (in a follow-up PR) the sweep's +# The scalar formula both the estimate and (in a follow-up PR) the sweep's # `InferenceSummary._check_and_set_kv_cache_oom` should use. The OOM check folds # `reserved` (block-allocator overhead) and `tolerance` into the fraction # multiplicatively; the estimate passes `reserved=tolerance=0` here and applies @@ -1046,15 +1048,19 @@ def estimate_kv_cache( allow_naive_fallback: bool = False, allow_hf_config_download: bool = False, fpm_profile: dict | str | FpmModelProfile | None = None, + worker_type: str = "aggregated", cp_size: int = 1, + context_length: int | None = None, + request_occupancy_tokens: int | None = None, ) -> dict[str, Any]: """Compute the KV-cache memory estimate (raw + optional tolerance margin). - This is the complete implementation: fraction + tolerance validation, the - native/naive budget math, AND the tolerance margin. The Rust - ``aisimulate_core::memory::estimate_kv_cache`` is a pure forwarder that - calls this function and rebuilds its result, so this is the single source of - truth for the estimate. + This function validates inputs and implements scalar native/naive budget + math and tolerance margins. Grouped and runtime-memory FPM profiles delegate + budget arithmetic to ``RustForwardPassPerfModel.estimate_cache_budget``. + Grouped token-capacity and bytes-per-token fields are ``None``. The legacy Rust + ``aisimulate_core::memory::estimate_kv_cache`` transport calls this function + and rebuilds scalar results; it rejects grouped results. Native path: :meth:`KVCacheEstimator.from_request` builds AIC's full backend memory model and :meth:`KVCacheEstimator.estimate` does the OfFree (TRT-LLM) / @@ -1085,18 +1091,26 @@ def estimate_kv_cache( model build is unsupported (default off -> the error propagates). allow_hf_config_download: allow the naive fallback to download a ``config.json`` from HuggingFace when it is not local / pre-cached. - fpm_profile: explicit FPM identity and conservative rank-local resource - bounds. This route reads only hardware specifications, requires the - requested scheduler envelope to fit the profile, and never builds - a model or falls back to naive estimation. CUDA graph bytes remain - a separate reservation in addition to declared profile overheads. + fpm_profile: explicit FPM identity and rank-local memory evidence. Pending + memory fails before estimation. Complete declared bytes require the + requested scheduler envelope to fit the profile; runtime capacity + requires its exact scheduler settings and memory fraction. This + route never constructs an analytical model or uses naive fallback. + Separate CUDA graph bytes apply only to declared profile overheads. + worker_type: canonical serving role selecting this deployment's resources. cp_size: profile context-parallel identity; currently only CP1 is supported. + context_length: optional context bound for a grouped profile's per-request + cache peak, also checked against runtime memory's ``max_model_len``. + Defaults to the profile context and does not change its byte budget. + request_occupancy_tokens: optional logical request length for a separate + native one-token decode occupancy estimate. Requires an FPM profile; + context and scheduler admission still use their original settings. Returns: A flat dict with ``total_gpu_capacity_bytes``, ``total_kv_size_bytes``, ``kv_size_per_token_bytes``, ``total_kv_size_tokens``, ``source`` (``"native"`` | ``"naive_fallback"`` | ``"profile"``), ``memory_breakdown`` (dict on the - native path, ``None`` on the fallback), and ``tolerance_adjusted`` (a dict + native/declared paths, ``None`` for runtime capacity or naive fallback), and ``tolerance_adjusted`` (a dict when ``tolerance_fraction`` is set, else ``None``). Raises: @@ -1115,6 +1129,8 @@ def estimate_kv_cache( # Validate the compute-side MTP depth before any fallback path. validate_nextn(nextn) + if request_occupancy_tokens is not None and fpm_profile is None: + raise ValueError("request_occupancy_tokens requires an FPM profile") if fpm_profile is not None: if nextn: raise ValueError("FPM profile resources support plain autoregressive execution; nextn must be 0") @@ -1126,6 +1142,7 @@ def estimate_kv_cache( system=system, backend=backend, backend_version=backend_version, + worker_type=worker_type, tp_size=tp_size, pp_size=pp_size, attention_dp_size=attention_dp_size, @@ -1145,12 +1162,49 @@ def estimate_kv_cache( wideep_num_slots=wideep_num_slots, ) resources = deployment.resources + resources.require_memory() resources.validate_envelope(max_num_tokens=max_num_tokens, max_batch_size=max_batch_size) - if resources.non_kv_bytes + cuda_graph_reserved_bytes > _MAX_EXACT_BYTE_COUNT: + if ( + resources.runtime_memory is None + and resources.non_kv_bytes + cuda_graph_reserved_bytes > _MAX_EXACT_BYTE_COUNT + ): raise ValueError("total FPM non-KV bytes including CUDA graphs must not exceed 2**53") capacity = gpu_memory_capacity_bytes_override if capacity is None: capacity = perf_database.load_system_spec(system, systems_paths=systems_path)["gpu"]["mem_capacity"] + if ( + resources.cache_layout == "grouped" + or resources.runtime_memory is not None + or request_occupancy_tokens is not None + ): + from aisimulate_core.sdk.rust_engine_step import RustForwardPassPerfModel + + return RustForwardPassPerfModel.estimate_cache_budget( + { + "model": model_path, + "system": system, + "backend": backend, + "backend_version": backend_version, + "worker_type": worker_type, + "tp": tp_size, + "pp": pp_size, + "attention_dp": attention_dp_size, + "moe_tp_size": deployment.moe_tp, + "moe_ep_size": deployment.moe_ep, + "fpm_profile": profile.model_dump(mode="json"), + }, + { + "total_gpu_capacity_bytes": capacity, + "memory_fraction_kind": memory_fraction_kind, + "memory_fraction_value": memory_fraction_value, + "max_num_tokens": max_num_tokens, + "max_batch_size": max_batch_size, + "context_length": context_length, + "request_occupancy_tokens": request_occupancy_tokens, + "cuda_graph_reserved_bytes": cuda_graph_reserved_bytes, + "tolerance_fraction": tolerance_fraction, + }, + ) estimate = KVCacheEstimator._estimate_from_breakdown( { "weights_bytes": resources.weights_bytes, @@ -1285,7 +1339,9 @@ def estimate_num_gpu_blocks( allow_hf_config_download: bool = False, diagnostics: dict[str, Any] | None = None, fpm_profile: dict | str | FpmModelProfile | None = None, + worker_type: str = "aggregated", cp_size: int = 1, + context_length: int | None = None, ) -> int: """Convert the KV-cache token capacity to a scheduler block count. @@ -1349,9 +1405,16 @@ def estimate_num_gpu_blocks( allow_naive_fallback=allow_naive_fallback, allow_hf_config_download=allow_hf_config_download, fpm_profile=fpm_profile, + worker_type=worker_type, cp_size=cp_size, + context_length=context_length, ) + if estimate.get("cache_layout") == "grouped": + raise ValueError( + "grouped FPM caches require cache_groups and the byte budget from estimate_kv_cache; " + "a scalar num_gpu_blocks or token capacity cannot represent windowed cache groups" + ) adjusted = estimate.get("tolerance_adjusted") if adjusted is not None: tokens = int(adjusted["total_kv_size_tokens"]) diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py index 8c090cd0f..b10ea9398 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py @@ -654,7 +654,7 @@ def _infer_quant_modes_from_raw_config(raw_config: dict, architecture: str | Non # TODO: support fp4 kv cache if kv_cache_algo == "fp8": overrides["kvcache_quant_mode"] = common.KVCacheQuantMode.fp8 - elif kv_cache_algo == "bfloat16": + elif kv_cache_algo in ("bfloat16", "none"): overrides["kvcache_quant_mode"] = common.KVCacheQuantMode.bfloat16 elif kv_cache_algo is not None: raise ValueError(f"Unsupported kv cache algorithm: {kv_cache_algo}") @@ -664,7 +664,12 @@ def _infer_quant_modes_from_raw_config(raw_config: dict, architecture: str | Non overrides["kvcache_quant_mode"] = common.KVCacheQuantMode.fp8 # FMHA quant mode - if quant_algo is not None and (quant_algo in ("fp8", "fp8_block", "nvfp4") or kv_cache_algo in ("fp8",)): + if kv_cache_algo == "none" and overrides["kvcache_quant_mode"] == common.KVCacheQuantMode.bfloat16: + # An explicit unquantized KV declaration uses the full-precision SDK + # default independently of weight quantization. Preserve native cache + # requirements applied above; missing/null metadata keeps legacy defaults. + overrides["fmha_quant_mode"] = common.FMHAQuantMode.bfloat16 + elif quant_algo is not None and (quant_algo in ("fp8", "fp8_block", "nvfp4") or kv_cache_algo in ("fp8",)): overrides["fmha_quant_mode"] = common.FMHAQuantMode.fp8 if kv_cache_algo is None or kv_cache_algo != "fp8": overrides["kvcache_quant_mode"] = common.KVCacheQuantMode.fp8 diff --git a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py index 1cdd5fc35..f4ea041f0 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -318,10 +318,28 @@ def normalize_config(config: ForwardPassPerfModelConfig | Mapping[str, Any]) -> return json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(_forward_pass_config_json(config))) + @staticmethod + def estimate_cache_budget( + config: ForwardPassPerfModelConfig | Mapping[str, Any], budget: Mapping[str, Any] + ) -> dict[str, Any]: + """Size explicit profile resources without a graph or timing data. + + Grouped caches expose a byte budget and per-request peak bound; their + scalar bytes-per-token and aggregate token capacities are unavailable. + Rust owns allocation rounding, retention windows and budget arithmetic. + """ + import aisimulate_core + + return json.loads( + aisimulate_core.RustForwardPassPerfModel.estimate_cache_budget( + _forward_pass_config_json(config), _json_dumps(dict(budget)) + ) + ) + def predict_prefill_latency(self, bs: int, isl: int, prefix: int = 0) -> float: - """Return latency in ms for a qualified homogeneous graph-prefill shape. + """Return uncorrected native prefill latency in milliseconds. - ``isl`` is the total sequence length including ``prefix``. The selected + ``isl`` is the total sequence length including ``prefix``. A selected graph profile admits only its measured integer shapes; Rust owns validation and prediction. This method returns no scheduler or energy estimate. """ @@ -394,6 +412,23 @@ def diagnostics(self) -> dict[str, Any]: """ return json.loads(self._inner.diagnostics()) + def fpm_query_coverage(self) -> dict[str, Any] | None: + """Bounded direct-FPM lookup evidence, or None when collection is disabled. + + Counts describe native lookups, including errors and mixed-pass + baselines. External timing-cache hits are not new lookups. A snapshot + does not establish replay completion or prediction accuracy. + """ + return json.loads(self._inner.fpm_query_coverage()) + + def predict_decode_latency_total(self, batch_size: int, total_past_kv_tokens: int) -> float: + """Uncorrected native decode timing at the exact past-KV batch total.""" + return self._inner.predict_decode_latency_total(batch_size, total_past_kv_tokens) + + def fpm_decode_kv_ceiling(self) -> int | None: + """Largest collected decode KV total; reading it does not record a lookup.""" + return self._inner.fpm_decode_kv_ceiling() + def regression_store_diagnostics(self) -> list[dict[str, Any]]: """Return each regression store's label, readiness, and retained count. diff --git a/python/aisimulate/tests/cross_package/test_core_public_api.py b/python/aisimulate/tests/cross_package/test_core_public_api.py index bfa3bad75..86ec7accb 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -108,6 +108,7 @@ def test_stable_function_signatures() -> None: "forward_model: 'str | None' = None, " "decoder_replay: 'bool' = False, " "fpm_profile: 'dict | str | FpmModelProfile | None' = None, " + "worker_type: 'str' = 'aggregated', " "fpm_interpolation: 'str | None' = None, cp_size: 'int' = 1, " "database_mode: 'str | None' = None, shared_layer: 'bool | None' = None, " "transfer_policy: 'str | list[str] | None' = None, " diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index 3fb0a8a37..5677bfbd5 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -38,6 +38,7 @@ "fpm_config", "fpm_dataset", "fpm_identity", + "fpm_model_metadata", "inference_summary", "memory", "state_memory", diff --git a/python/aisimulate/tests/unit/collector/fixtures/fpm_collection_plan_v10.json b/python/aisimulate/tests/unit/collector/fixtures/fpm_collection_plan_v10.json new file mode 100644 index 000000000..be47408ec --- /dev/null +++ b/python/aisimulate/tests/unit/collector/fixtures/fpm_collection_plan_v10.json @@ -0,0 +1,335 @@ +{ + "aic_revision": "4d702ff6b756b21e74b76f30273e9828fdde6861", + "backend": "vllm", + "backend_policies": [ + { + "admission_reason": "AIC automatic baseline for the selected model/backend", + "aic_fields": { + "attention_backend": null, + "enable_eplb": false, + "enable_wideep": false, + "moe_backend": null + }, + "expected_markers": {}, + "generator_overrides": {}, + "policy_id": "baseline_auto" + } + ], + "capability": { + "aic_database_version": "0.27.0", + "allow_pure_tp": true, + "architecture": "UnregisteredMoeForCausalLM", + "attention_kind": "moe_mha", + "attention_source": "dense_attention", + "dtype": { + "comm_quant_mode": "half", + "fmha_by_kv_dtype": { + "bfloat16": "bfloat16" + }, + "fmha_quant_mode": "bfloat16", + "fmha_resolution": "checkpoint_native", + "fmha_resolution_by_kv_dtype": { + "bfloat16": "checkpoint_native" + }, + "gemm_quant_mode": "bfloat16", + "kv_cache_dtypes": [ + "bfloat16" + ], + "moe_quant_mode": "bfloat16", + "native_kv_cache_dtype": "bfloat16" + }, + "is_moe": true, + "model_config": { + "payload": { + "architectures": [ + "UnregisteredMoeForCausalLM" + ], + "hidden_size": 128, + "max_position_embeddings": 8192, + "num_attention_heads": 8, + "num_experts": 8, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "torch_dtype": "bfloat16", + "vocab_size": 1024 + }, + "sha256": "b7e6b75d81c1b9eddba8326c963b8833350fbccb1b09b2a4c2581f0519119a73", + "source_kind": "explicit" + }, + "model_family": null, + "support_level": "bootstrap_template", + "support_reason": "architecture is not registered by AIC; using a conservative config-derived template", + "template_id": "generic:moe_mha", + "template_version": 1 + }, + "cells": [ + { + "backend_policy": { + "admission_reason": "AIC automatic baseline for the selected model/backend", + "aic_fields": { + "attention_backend": null, + "enable_eplb": false, + "enable_wideep": false, + "moe_backend": null + }, + "expected_markers": {}, + "generator_overrides": {}, + "policy_id": "baseline_auto" + }, + "cell_id": "fpm-d111aeb8823a9bee", + "kv_cache_dtype": "bfloat16", + "parallel_strategy": "tep", + "point_source": "dynamo_native_self_benchmark", + "resolved_dtypes": { + "comm_quant_mode": "half", + "fmha_quant_mode": "bfloat16", + "fmha_resolution": "checkpoint_native", + "gemm_quant_mode": "bfloat16", + "kvcache_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16" + }, + "topology": { + "cp": 1, + "dp": 1, + "moe_ep": 8, + "moe_tp": 1, + "pp": 1, + "tp": 8 + }, + "weight_quantization": "bfloat16", + "workload_kind": "prefill" + }, + { + "backend_policy": { + "admission_reason": "AIC automatic baseline for the selected model/backend", + "aic_fields": { + "attention_backend": null, + "enable_eplb": false, + "enable_wideep": false, + "moe_backend": null + }, + "expected_markers": {}, + "generator_overrides": {}, + "policy_id": "baseline_auto" + }, + "cell_id": "fpm-66274b2ba8ae195c", + "kv_cache_dtype": "bfloat16", + "parallel_strategy": "tep", + "point_source": "dynamo_native_self_benchmark", + "resolved_dtypes": { + "comm_quant_mode": "half", + "fmha_quant_mode": "bfloat16", + "fmha_resolution": "checkpoint_native", + "gemm_quant_mode": "bfloat16", + "kvcache_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16" + }, + "topology": { + "cp": 1, + "dp": 1, + "moe_ep": 8, + "moe_tp": 1, + "pp": 1, + "tp": 8 + }, + "weight_quantization": "bfloat16", + "workload_kind": "decode" + } + ], + "counts": { + "backend_policies": 1, + "candidate_topologies": 1, + "cells": 2, + "memory_rejected_topologies": 0, + "memory_unknown_topologies": 1, + "points": "runtime-determined", + "prefill_cudagraph_capture_sizes": null, + "prefill_new_token_axis_points": null, + "topologies": 1 + }, + "dtype_profile": { + "comm_quant_mode": "half", + "fmha_by_kv_dtype": { + "bfloat16": "bfloat16" + }, + "fmha_quant_mode": "bfloat16", + "fmha_resolution": "checkpoint_native", + "fmha_resolution_by_kv_dtype": { + "bfloat16": "checkpoint_native" + }, + "gemm_quant_mode": "bfloat16", + "kv_cache_dtypes": [ + "bfloat16" + ], + "moe_quant_mode": "bfloat16", + "native_kv_cache_dtype": "bfloat16" + }, + "fpm_profile": { + "architecture": "UnregisteredMoeForCausalLM", + "context_length": 8192, + "deployments": [ + { + "attention_backend": "auto", + "backend": "vllm", + "backend_version": "0.27.0", + "comm_quant_mode": "half", + "cp": 1, + "dp": 1, + "enable_eplb": false, + "enable_wideep": false, + "fmha_quant_mode": "bfloat16", + "gemm_quant_mode": "bfloat16", + "kv_cache_dtype": "bfloat16", + "moe_backend": "auto", + "moe_ep": 8, + "moe_quant_mode": "bfloat16", + "moe_tp": 1, + "pp": 1, + "resources": { + "cache_layout": "linear", + "kv_bytes_per_token": 60216, + "max_batch_size": 1024, + "max_num_tokens": 8192, + "provenance": "Independent test envelope: 100 GB weights plus three 1 GB overhead bounds." + }, + "system": "gb200", + "tp": 8 + } + ], + "model": "example/slurm-fpm-model", + "model_revision": "collector-test-snapshot", + "num_experts": 8, + "provenance": "Collector test declaration; not a measured memory qualification.", + "schema_version": 1 + }, + "generator_config_sha256": "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a", + "model_path": "example/slurm-fpm-model", + "options": { + "attention_backend": "auto", + "cp_sizes": null, + "dp_sizes": null, + "enable_eplb": "false", + "enable_wideep": "false", + "global_warmup_iterations": 5, + "gpu_counts": [ + 8 + ], + "gpu_memory_utilization": 0.9, + "kv_cache_dtypes": [ + "bfloat16" + ], + "max_gpus": 8, + "max_num_batched_tokens": 1024, + "max_num_seqs": 64, + "measurement_repeats": 1, + "moe_backend": "auto", + "moe_ep_sizes": null, + "moe_tp_sizes": null, + "parallel_axes": [], + "parallel_presets": [ + "tep" + ], + "point_source": "dynamo_native_self_benchmark", + "pp_sizes": null, + "prefill_sampling": { + "cudagraph_capture_size_count": null, + "cudagraph_capture_sizes": null, + "cudagraph_policy": "runtime", + "max_batch_size": 64, + "max_cudagraph_capture_size": null, + "max_isl": 1024, + "max_total_prefill_tokens": 1024, + "new_token_axis_point_count": null, + "new_token_axis_points": null, + "prefill_max_kv_read_token_samples": 19, + "prefill_max_new_token_samples": null + }, + "tp_sizes": null, + "vllm_max_model_len": 4096, + "warmup_repeats": 0, + "weight_quantizations": [] + }, + "point_generation": { + "coordinates": [ + "batch_size", + "total_prefill_tokens", + "total_kv_read_tokens" + ], + "method": "native_self_benchmark", + "owner": "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler", + "partition_policy": "balanced_v1", + "planned_point_count": null, + "point_admission": "dynamo_live_scheduler", + "precondition": "vllm_engine_initialized", + "prefill_sampling": { + "cudagraph_capture_size_count": null, + "cudagraph_capture_sizes": null, + "cudagraph_policy": "runtime", + "max_batch_size": 64, + "max_cudagraph_capture_size": null, + "max_isl": 1024, + "max_total_prefill_tokens": 1024, + "new_token_axis_point_count": null, + "new_token_axis_points": null, + "prefill_max_kv_read_token_samples": 19, + "prefill_max_new_token_samples": null + } + }, + "runtime_memory_policy": { + "async_scheduling": false, + "evidence_schema_version": 1, + "observation": "enabled", + "selected_vllm_version": "0.27.0", + "source": "vllm_initialization", + "supported_vllm_version": "0.27.0" + }, + "schema_name": "aic_fpm_collection_plan", + "schema_version": 10, + "sha256": "ea99845deee06804fd35a7a0999b0118446df1ac8ededbf7c813345360ba6bf7", + "system": "gb200", + "topologies": [ + { + "cp": 1, + "dp": 1, + "moe_ep": 8, + "moe_tp": 1, + "pp": 1, + "strategy": "tep", + "tp": 8 + } + ], + "topology_memory_admission": [ + { + "activation_envelope": { + "cuda_graph_reservation": "runtime_profiled_separately", + "max_batch_size": 1024, + "max_num_tokens": 8192, + "requested_prefill_tokens": 1024, + "scope": "rank_local" + }, + "disposition": "unknown", + "estimates": [ + { + "disposition": "unknown", + "estimated_non_kv_bytes": null, + "gpu_capacity_bytes": 198674743296, + "gpu_memory_budget_bytes": 178807268966, + "headroom_bytes": null, + "kv_cache_dtype": "bfloat16", + "provenance": "Independent test envelope: 100 GB weights plus three 1 GB overhead bounds.", + "reason": "memory capacity is determined by this collection worker's runtime initialization; no non-KV byte declaration is required" + } + ], + "reason": "AIC could not prove the topology is impossible; runtime verification is required", + "source": "fpm_profile_pending", + "topology": { + "cp": 1, + "dp": 1, + "moe_ep": 8, + "moe_tp": 1, + "pp": 1, + "tp": 8 + } + } + ] +} diff --git a/python/aisimulate/tests/unit/collector/test_fpm_cpu_affinity.py b/python/aisimulate/tests/unit/collector/test_fpm_cpu_affinity.py new file mode 100644 index 000000000..49b6a4b80 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_cpu_affinity.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +from collector.fpm_forward.cpu_affinity import inspect_cpu_affinity +from collector.fpm_forward.runtime import fpm_memory_observer as observer + +from .test_fpm_profile_collection import no_models_or_timing_data # noqa: F401 +from .test_fpm_repeatability import _write_campaign, campaign # noqa: F401 + +pytestmark = pytest.mark.unit + + +def _snapshot(host, pid, cpus): + return { + "hostname": host, + "pid": pid, + "observer_thread_id": pid, + "main_thread_allowed_cpus": cpus, + "threads": [{"tid": pid, "allowed_cpus": cpus}], + "thread_ids_before": [pid], + "thread_ids_after": [pid], + "thread_errors": [], + "errors": [], + "status": "observed", + } + + +def _campaign(tmp_path, *, tp=1, dp=8, cpus=16, bind="cores"): + plan = SimpleNamespace( + sha256="a" * 64, + options=SimpleNamespace(executor="slurm", slurm_cpus_per_task=cpus, slurm_cpu_bind=bind), + ) + cell = SimpleNamespace(cell_id="prefill", topology=SimpleNamespace(tp=tp, dp=dp, pp=1)) + collection = SimpleNamespace(collector_attempt_id="attempt") + provenance = {"plan_sha256": plan.sha256, "cell_id": cell.cell_id, "attempt_id": "attempt"} + + def write(kind, node, pid, dp_rank=None, tp_rank=None, mask=None): + rank = {"dp_rank": dp_rank, "tp_rank": tp_rank, "pp_rank": 0 if kind == "worker" else None} + suffix = f"-dp{dp_rank}-tp{tp_rank}-pp0" if kind == "worker" else f"-dp{dp_rank}" if kind == "scheduler" else "" + path = tmp_path / f"node{node}" / f"fpm-cpu-{kind}{suffix}.json" + path.parent.mkdir(exist_ok=True) + payload = { + "schema_name": "aisimulate_fpm_cpu_affinity", + "schema_version": 1, + "kind": kind, + "measurement": { + "launcher": "before_engine_start", + "worker": "after_warmup", + "scheduler": "scheduler_initialized", + }[kind], + "collector_provenance": provenance, + **rank, + **_snapshot(f"node{node}", pid, list(range(16)) if mask is None else mask), + } + if kind == "launcher": + payload.update(requested_cpus_per_task=cpus, cpu_bind=bind, local_gpu_count=4) + path.write_text(json.dumps(payload)) + return path + + for node in range(tp * dp // 4): + write("launcher", node, 100) + for d in range(dp): + for t in range(tp): + node, local = divmod(d * tp + t, 4) + write("worker", node, 200 + local, d, t) + node, local = divmod(d * tp, 4) + write("scheduler", node, 300 + local, d) + return plan, cell, collection + + +def _inspect(tmp_path, context): + plan, cell, collection = context + return inspect_cpu_affinity(cell, tmp_path, collection, plan=plan, expected_nodes=2) + + +def _change(path, transform): + payload = json.loads(path.read_text()) + transform(payload) + path.write_text(json.dumps(payload)) + + +@pytest.mark.parametrize("tp,dp", [(1, 8), (8, 1)]) +def test_shared_pools_and_identical_cpu_ids_on_different_hosts_are_valid(tmp_path, tp, dp): + report = _inspect(tmp_path, _campaign(tmp_path, tp=tp, dp=dp)) + assert report["status"] == "qualified" + assert report["policy_required"] is True + assert len(report["nodes"]) == 2 + assert not report["failures"] + for item in report["observations"]: + raw = Path(item["source"]["path"]).read_bytes() + assert item["source"]["sha256"] == hashlib.sha256(raw).hexdigest() + + +@pytest.mark.parametrize("masks", [[[0]] * 4, [[0], [0], [1, 2], [3]]]) +def test_dep_scheduler_masks_must_allow_distinct_logical_cpus(tmp_path, masks): + context = _campaign(tmp_path) + for dp, cpus in enumerate(masks): + path = tmp_path / "node0" / f"fpm-cpu-scheduler-dp{dp}.json" + _change(path, lambda value, cpus=cpus: value.update(_snapshot("node0", value["pid"], cpus))) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "cannot run on distinct logical CPUs" in str(report["failures"]) + + +def test_disjoint_scheduler_cpu_sets_are_allowed_but_not_required(tmp_path): + context = _campaign(tmp_path) + for dp in range(8): + node, local = divmod(dp, 4) + path = tmp_path / f"node{node}" / f"fpm-cpu-scheduler-dp{dp}.json" + _change(path, lambda value, local=local: value.update(_snapshot(value["hostname"], value["pid"], [local]))) + assert _inspect(tmp_path, context)["status"] == "qualified" + + +@pytest.mark.parametrize( + "change,expected", + [ + (lambda p: p["collector_provenance"].update(attempt_id="other"), "attempt"), + (lambda p: p.update(dp_rank=8), "rank"), + (lambda p: p.update(main_thread_allowed_cpus=[True]), "mask"), + (lambda p: p.update(main_thread_allowed_cpus=[1, 1]), "mask"), + (lambda p: p.update(hostname="unlaunched-host"), "launcher"), + (lambda p: p["threads"].append({"tid": 999, "allowed_cpus": [99]}), "thread"), + ], +) +def test_invalid_or_mismatched_runtime_evidence_is_not_qualified(tmp_path, change, expected): + context = _campaign(tmp_path) + _change(tmp_path / "node0/fpm-cpu-worker-dp0-tp0-pp0.json", change) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert expected in str(report["failures"]) + + +def test_observed_thread_outside_launcher_pool_is_detected(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-worker-dp0-tp0-pp0.json" + + def change(payload): + payload["thread_ids_before"].append(999) + payload["thread_ids_after"].append(999) + payload["threads"].append({"tid": 999, "allowed_cpus": [99]}) + + _change(path, change) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "outside the launcher CPU pool" in str(report["failures"]) + + +@pytest.mark.parametrize( + "change,expected", + [ + (lambda p: p.update(requested_cpus_per_task=8), "policy"), + (lambda p: p.update(cpu_bind="none"), "policy"), + (lambda p: p.update(local_gpu_count=8), "worker"), + (lambda p: p.update(_snapshot("node0", 100, [0])), "requested"), + ], +) +def test_launch_pool_and_requested_policy_are_independently_checked(tmp_path, change, expected): + context = _campaign(tmp_path) + _change(tmp_path / "node0/fpm-cpu-launcher.json", change) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert expected in str(report["failures"]) + + +@pytest.mark.parametrize("requested,observed", [(16, 16.0), (1, True)]) +def test_launcher_policy_requires_an_integer_even_when_numeric_value_matches(tmp_path, requested, observed): + context = _campaign(tmp_path, tp=8, dp=1, cpus=requested) + assert _inspect(tmp_path, context)["status"] == "qualified" + _change(tmp_path / "node0/fpm-cpu-launcher.json", lambda value: value.update(requested_cpus_per_task=observed)) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "policy" in str(report["failures"]) + + +@pytest.mark.parametrize("missing", ["worker-dp0-tp0-pp0", "scheduler-dp0", "launcher"]) +def test_missing_role_evidence_leaves_cpu_qualification_incomplete(tmp_path, missing): + context = _campaign(tmp_path) + (tmp_path / f"node0/fpm-cpu-{missing}.json").unlink() + report = _inspect(tmp_path, context) + assert report["status"] == "incomplete" + assert report["missing_evidence"] + + +def test_partial_thread_snapshot_cannot_claim_complete_affinity(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-worker-dp0-tp0-pp0.json" + + def change(payload): + payload["status"] = "partial" + payload["thread_ids_after"].append(999) + payload["errors"].append("thread list changed during observation") + + _change(path, change) + report = _inspect(tmp_path, context) + assert report["status"] == "incomplete" + assert "partial" in str(report["missing_evidence"]) + + +def test_legacy_plan_is_not_retroactively_given_verified_affinity(tmp_path): + context = _campaign(tmp_path, cpus=None, bind=None) + report = _inspect(tmp_path, context) + assert report["status"] == "incomplete" + assert report["policy_required"] is False + assert "frozen CPU policy" in str(report["missing_evidence"]) + + +def test_duplicate_rank_observation_is_rejected(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-worker-dp0-tp0-pp0.json" + copied = tmp_path / "copy" + copied.mkdir() + (copied / path.name).write_bytes(path.read_bytes()) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "duplicate" in str(report["failures"]) + + +def test_scheduler_host_cannot_hide_local_cpu_contention(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-scheduler-dp0.json" + _change(path, lambda payload: payload.update(_snapshot("node1", 999, list(range(16))))) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "scheduler host differs" in str(report["failures"]) + + +def test_different_rank_cannot_reuse_another_process_observation(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-scheduler-dp1.json" + _change(path, lambda payload: payload.update(_snapshot("node0", 300, list(range(16))))) + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "duplicate process" in str(report["failures"]) + + +def test_complete_snapshot_requires_observation_thread_membership(tmp_path): + context = _campaign(tmp_path) + _change(tmp_path / "node0/fpm-cpu-launcher.json", lambda payload: payload.update(observer_thread_id=999)) + assert _inspect(tmp_path, context)["status"] == "failed" + + +def test_json_duplicate_fields_do_not_overwrite_evidence_identity(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-launcher.json" + path.write_text(path.read_text()[:-1] + ', "pid": 999}') + report = _inspect(tmp_path, context) + assert report["status"] == "failed" + assert "duplicate" in str(report["failures"]) + + +def test_generated_node_count_cannot_be_replaced_by_observed_host_claims(tmp_path): + plan, cell, collection = _campaign(tmp_path) + report = inspect_cpu_affinity(cell, tmp_path, collection, plan=plan, expected_nodes=1) + assert report["status"] == "failed" + assert "generated" in str(report["failures"]) + + +def test_launcher_support_threads_must_stay_inside_its_main_thread_pool(tmp_path): + context = _campaign(tmp_path) + path = tmp_path / "node0/fpm-cpu-launcher.json" + + def change(payload): + payload["thread_ids_before"].append(101) + payload["thread_ids_after"].append(101) + payload["threads"].append({"tid": 101, "allowed_cpus": [99]}) + + _change(path, change) + assert "outside the launcher CPU pool" in str(_inspect(tmp_path, context)["failures"]) + + +@pytest.mark.parametrize("state", ["missing", "qualified", "failed"]) +def test_fresh_cpu_policy_gates_readiness_and_quality_without_discarding_native_rows( + campaign, # noqa: F811 + tmp_path, + monkeypatch, + state, +): + from collector.fpm_forward import repeatability, runner + + from aisimulate.support.collection_readiness import _cell_report + from aisimulate.support.validation_workflow import _execution_status + + source_plan, _, _ = campaign + plan = replace( + source_plan, + options=replace( + source_plan.options, + executor="slurm", + slurm_container_image="synthetic", + slurm_cpus_per_task=16, + slurm_cpu_bind="cores", + ), + ) + plan = repeatability._subset_plan( + plan, {"cell_id": plan.cells[0].cell_id, "benchmark_points": json.loads(plan.options.benchmark_points_json)} + ) + source, checkpoint = tmp_path / "fresh-cpu-policy", tmp_path / "fresh-checkpoint.json" + _write_campaign(source, checkpoint, plan) + cell = plan.cells[0] + cell_dir = source / "cells" / cell.cell_id + raw = cell_dir / "raw" + nodes = runner._expected_nodes(cell_dir / runner.FPM_MANIFEST_FILENAME) + local_gpus = cell.topology.total_gpus // nodes + node_directories = [raw / "pod" if node == 0 else raw / f"cpu-node{node}" for node in range(nodes)] + if state != "missing": + for node in range(nodes): + directory = node_directories[node] + directory.mkdir(exist_ok=True) + (directory / "collector-provenance.json").write_bytes((raw / "pod/collector-provenance.json").read_bytes()) + (directory / runner.POINTS_RECEIPT_FILENAME).write_bytes( + (raw / "pod" / runner.POINTS_RECEIPT_FILENAME).read_bytes() + ) + monkeypatch.setattr( + observer, "cpu_affinity_snapshot", lambda node=node: _snapshot(f"node{node}", 100, list(range(16))) + ) + observer.observe_cpu( + "launcher", + requested_cpus_per_task=16, + cpu_bind="cores", + local_gpu_count=local_gpus, + directory=directory, + ) + for dp in range(cell.topology.dp): + for tp in range(cell.topology.tp): + node, local = divmod(dp * cell.topology.tp + tp, local_gpus) + directory = node_directories[node] + mask = [99] if state == "failed" and (dp, tp) == (0, 0) else list(range(16)) + monkeypatch.setattr( + observer, + "cpu_affinity_snapshot", + lambda node=node, local=local, mask=mask: _snapshot(f"node{node}", 200 + local, mask), + ) + observer.observe_cpu("worker", dp_rank=dp, tp_rank=tp, pp_rank=0, directory=directory) + node, local = divmod(dp * cell.topology.tp, local_gpus) + monkeypatch.setattr( + observer, + "cpu_affinity_snapshot", + lambda node=node, local=local: _snapshot(f"node{node}", 300 + local, list(range(16))), + ) + observer.observe_cpu("scheduler", dp_rank=dp, directory=node_directories[node]) + before = {path: path.read_bytes() for path in source.rglob("*") if path.is_file()} + entry = json.loads(checkpoint.read_text())["cells"][cell.cell_id] + readiness = _cell_report(plan, cell, source, entry) + assert readiness["direct_eligible_points"] > 0 + assert readiness["execution"]["cpu_affinity"]["status"] == ("incomplete" if state == "missing" else state) + assert readiness["status"] == ("ready" if state == "qualified" else "blocked") + frozen = repeatability.freeze_repeatability_plan(plan, source, checkpoint, max_points_per_cell=6) + assert _execution_status(frozen) == {"missing": "incomplete", "qualified": "passed", "failed": "failed"}[state] + assert {path: path.read_bytes() for path in source.rglob("*") if path.is_file()} == before + + +def test_linux_snapshot_distinguishes_main_thread_from_observation_thread(tmp_path, monkeypatch): + tasks = tmp_path / "tasks" + tasks.mkdir() + for tid in (100, 101): + (tasks / str(tid)).mkdir() + monkeypatch.setattr(observer.os, "getpid", lambda: 100) + monkeypatch.setattr(observer.threading, "get_native_id", lambda: 101) + monkeypatch.setattr(observer.socket, "gethostname", lambda: "node0") + calls = [] + + def affinity(tid): + calls.append(tid) + return {0, 1, 2, 3} if tid == 100 else {2} + + monkeypatch.setattr(observer.os, "sched_getaffinity", affinity, raising=False) + payload = observer.cpu_affinity_snapshot(task_directory=tasks) + assert payload["status"] == "observed" + assert payload["main_thread_allowed_cpus"] == [0, 1, 2, 3] + assert payload["observer_thread_id"] == 101 + assert payload["threads"] == [{"tid": 100, "allowed_cpus": [0, 1, 2, 3]}, {"tid": 101, "allowed_cpus": [2]}] + assert 0 not in calls + + +def test_thread_disappearance_is_recorded_without_changing_affinity(tmp_path, monkeypatch): + tasks = tmp_path / "tasks" + tasks.mkdir() + for tid in (100, 101): + (tasks / str(tid)).mkdir() + monkeypatch.setattr(observer.os, "getpid", lambda: 100) + + def affinity(tid): + if tid == 101: + raise ProcessLookupError("thread exited") + return {0, 1} + + monkeypatch.setattr(observer.os, "sched_getaffinity", affinity, raising=False) + payload = observer.cpu_affinity_snapshot(task_directory=tasks) + assert payload["status"] == "partial" + assert payload["main_thread_allowed_cpus"] == [0, 1] + assert payload["thread_errors"][0]["tid"] == 101 + + +def test_unavailable_affinity_stays_unknown(monkeypatch): + monkeypatch.delattr(observer.os, "sched_getaffinity", raising=False) + payload = observer.cpu_affinity_snapshot() + assert payload["status"] == "unavailable" + assert payload["main_thread_allowed_cpus"] == [] + assert payload["errors"] + + +def test_sidecar_preserves_failed_snapshot_and_attempt_binding(tmp_path, monkeypatch): + provenance = {"plan_sha256": "a" * 64, "cell_id": "prefill", "attempt_id": "current"} + (tmp_path / "collector-provenance.json").write_text(json.dumps(provenance)) + snapshot = _snapshot("node0", 100, [0]) + monkeypatch.setattr(observer, "cpu_affinity_snapshot", lambda: snapshot) + result = observer.observe_cpu( + "launcher", requested_cpus_per_task=16, cpu_bind="cores", local_gpu_count=4, directory=tmp_path + ) + path = tmp_path / "fpm-cpu-launcher.json" + original = path.read_bytes() + assert result == json.loads(original) + assert result["collector_provenance"] == provenance + assert result["main_thread_allowed_cpus"] == [0] + assert result["requested_cpus_per_task"] == 16 + with pytest.raises(RuntimeError, match="duplicate"): + observer.observe_cpu( + "launcher", requested_cpus_per_task=16, cpu_bind="cores", local_gpu_count=4, directory=tmp_path + ) + assert path.read_bytes() == original diff --git a/python/aisimulate/tests/unit/collector/test_fpm_dense_synthetic_kv.py b/python/aisimulate/tests/unit/collector/test_fpm_dense_synthetic_kv.py new file mode 100644 index 000000000..93e2d3016 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_dense_synthetic_kv.py @@ -0,0 +1,536 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dense synthetic KV needs positive cache and eligibility evidence, not a label alone.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from dataclasses import replace +from pathlib import Path + +import pytest +from collector.fpm_forward.config import with_kv_warmup_defaults +from collector.fpm_forward.database import ( + aggregate_cell, + published_dense_synthetic_cells, + validate_formal_database_commit, + write_formal_database, +) +from collector.fpm_forward.model_capability import ResolvedModelConfig +from collector.fpm_forward.planner import _canonical_hash +from collector.fpm_forward.repeatability import load_repeatability_source +from collector.fpm_forward.runtime_memory import saved_plan_identity + +from aisimulate import main as cli +from aisimulate.support.collection_readiness import assess_readiness +from aisimulate.support.validation_workflow import _collection_source + +from ..test_onboard_collection_readiness import _change_decode +from ..test_onboard_finalization import build_completed_collection +from ..test_support_serving_validation import materialized_workload # noqa: F401 +from ..test_support_validation import validation_case # noqa: F401 +from ..test_support_validation_workflow import quality_case # noqa: F401 +from .test_fpm_forward import _synthetic_decode_cell +from .test_fpm_measurement_evidence import _add_measurement_protocol + +pytestmark = pytest.mark.unit + +_DENSE_SKIP = {"enabled": True, "warm_eligible": False, "skip_reason": "dense_model_content_insensitive"} +_FULL = { + "type": "FullAttentionSpec", + "dtype": "torch.bfloat16", + "block_size": 16, + "num_kv_heads": 2, + "head_size": 32, + "sliding_window": None, + "attention_chunk_size": None, + "non_causal": False, +} + + +def _engine(plan, cell, *, version=None, groups=None): + source = plan.capability.model_config.payload + return { + "model": { + "model_type": source["model_type"], + "architectures": source["architectures"], + "model": plan.model_path, + }, + "parallel": { + "tensor_parallel_size": cell.topology.tp, + "pipeline_parallel_size": cell.topology.pp, + "data_parallel_size": cell.topology.dp, + "prefill_context_parallel_size": cell.topology.cp, + "decode_context_parallel_size": cell.topology.cp, + "enable_expert_parallel": False, + "data_parallel_rank": 0, + }, + "kv_cache": {"groups": copy.deepcopy([_FULL] if groups is None else groups)}, + "versions": {"vllm": version or plan.capability.aic_database_version}, + } + + +def _dense_cell(tmp_path, *, groups=None): + plan, cell, directory = _synthetic_decode_cell( + tmp_path, + kvwarm=_DENSE_SKIP, + markers=("kvwarm_fake_fallback", "kvwarm_fake_fallback"), + parallel_strategy="tp", + ) + plan.capability.is_moe = False + plan.capability.attention_source = "dense_attention" + plan.capability.model_config = ResolvedModelConfig( + {"model_type": "llama", "architectures": ["LlamaForCausalLM"]}, + source_kind="explicit", + source_reference="test", + ) + plan.cells = (cell,) + path = directory / "raw/pod-0/benchmark.json" + payload = json.loads(path.read_text()) + payload["engine"] = _engine(plan, cell, groups=groups) + for row in payload["results"]: + row["kv_seed_regime"] = "fake_fallback" + path.write_text(json.dumps(payload)) + return plan, cell, directory, path + + +def _regimes(case): + plan, cell, directory, _path = case + return {row["kv_seed_regime"] for row in aggregate_cell(plan, cell, directory, expected_attempt_id="attempt")} + + +@pytest.mark.parametrize( + "groups", + [[_FULL], [_FULL, _FULL], [{"type": "UniformTypeKVCacheSpecs", "specs": [_FULL, _FULL]}]], + ids=["full", "multiple-full", "wrapped-full"], +) +def test_dense_full_attention_preserves_native_samples_and_becomes_direct_eligible(tmp_path, groups): + case = _dense_cell(tmp_path, groups=groups) + plan, cell, directory, path = case + before = path.read_bytes() + rows = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt") + assert {row["kv_seed_regime"] for row in rows} == {"skip:dense_model_content_insensitive"} + assert [row["latency_ms"] for row in rows] == [10.0, 20.0] + assert path.read_bytes() == before + assert {row["kv_seed_regime"] for row in json.loads(before)["results"]} == {"fake_fallback"} + + +@pytest.mark.parametrize( + "groups", + [ + [], + [None], + [{}], + [{"type": "MLAAttentionSpec"}], + [{"type": "SlidingWindowSpec"}], + [{"type": "MambaSpec"}], + [{**_FULL, "sliding_window": 512}], + [{**_FULL, "attention_chunk_size": 128}], + [{**_FULL, "non_causal": True}], + [{key: value for key, value in _FULL.items() if key != "sliding_window"}], + [{key: value for key, value in _FULL.items() if key != "attention_chunk_size"}], + [{key: value for key, value in _FULL.items() if key != "non_causal"}], + [_FULL, {"type": "ShortConvSpec"}], + [{"type": "UniformTypeKVCacheSpecs"}], + [{"type": "UniformTypeKVCacheSpecs", "specs": []}], + [{"type": "UniformTypeKVCacheSpecs", "specs": {"layer": _FULL}}], + [{"type": "UniformTypeKVCacheSpecs", "specs": [_FULL, {"type": "MambaSpec"}]}], + [{**_FULL, "specs": [{"type": "ShortConvSpec"}]}], + ], +) +def test_unknown_or_non_full_attention_cache_never_promotes_dense_skip(tmp_path, groups): + assert _regimes(_dense_cell(tmp_path, groups=groups)) == {"fake_fallback"} + + +@pytest.mark.parametrize( + ("section", "key", "value"), + [ + ("model", "model_type", "mamba"), + ("model", "architectures", ["DifferentForCausalLM"]), + ("model", "model", "another/checkpoint"), + ("parallel", "tensor_parallel_size", 8), + ("parallel", "tensor_parallel_size", True), + ("parallel", "pipeline_parallel_size", 2), + ("parallel", "prefill_context_parallel_size", 2), + ("parallel", "decode_context_parallel_size", 2), + ("parallel", "decode_context_parallel_size", True), + ("parallel", "enable_expert_parallel", True), + ("parallel", "enable_expert_parallel", 0), + ("parallel", "data_parallel_size", 2), + ("parallel", "data_parallel_rank", 1), + ("versions", "vllm", "0.99.0"), + ("kv_cache", "groups", None), + ], +) +def test_contradictory_or_missing_initialized_identity_stays_fake_fallback(tmp_path, section, key, value): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + payload["engine"][section][key] = value + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + del payload["engine"][section][key] + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + + +@pytest.mark.parametrize("missing", ["engine", "model", "parallel", "versions", "kv_cache"]) +def test_historical_or_incomplete_engine_evidence_is_not_promoted(tmp_path, missing): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + (payload if missing == "engine" else payload["engine"]).pop(missing) + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + + +def test_failed_engine_capture_cannot_establish_dense_eligibility(tmp_path): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + payload["engine"]["capture_error"] = "cache inspection failed" + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + + +@pytest.mark.parametrize("location", ["root", "text_config"]) +@pytest.mark.parametrize("field", ["num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts"]) +def test_expert_evidence_in_either_config_rejects_dense_skip(tmp_path, location, field): + case = _dense_cell(tmp_path) + capability = case[0].capability + source = capability.model_config.payload + target = source if location == "root" else source.setdefault("text_config", {}) + target[field] = 1 # Dynamo treats even one declared expert as a non-dense model. + capability.model_config = ResolvedModelConfig(source, source_kind="explicit", source_reference="test") + assert _regimes(case) == {"fake_fallback"} + + +@pytest.mark.parametrize("field,value", [("is_moe", True), ("attention_source", "mla_module")]) +def test_frozen_capability_must_agree_with_dense_full_attention(tmp_path, field, value): + case = _dense_cell(tmp_path) + setattr(case[0].capability, field, value) + assert _regimes(case) == {"fake_fallback"} + + +@pytest.mark.parametrize("field", ["is_moe", "attention_source", "model_config"]) +def test_unavailable_frozen_metadata_does_not_establish_dense_eligibility(tmp_path, field): + case = _dense_cell(tmp_path) + setattr(case[0].capability, field, None) + assert _regimes(case) == {"fake_fallback"} + + +def test_mounted_model_identity_requires_the_frozen_deployment_hash(tmp_path): + case = _dense_cell(tmp_path / "campaign/cells") + plan, _cell, directory, path = case + overrides = {"K8sConfig": {"k8s_model_path_in_pvc": "checkpoint", "k8s_pvc_mount_path": "/models"}} + plan.generator_config_sha256 = _canonical_hash(with_kv_warmup_defaults(overrides)) + generator = directory.parent.parent / "generator-overrides.json" + generator.write_text(json.dumps(overrides)) + for name in ("fpm_env.sh", "collector-runtime-env.sh"): + (directory / name).write_text("") + (directory / "generator-request.json").write_text( + json.dumps({"ServiceConfig": {"model_path": "/models/checkpoint"}}) + ) + (directory / "run.sh").write_text("engine_command=(python3 -m dynamo.vllm --model /models/checkpoint)\n") + payload = json.loads(path.read_text()) + payload["engine"]["model"]["model"] = "/models/checkpoint" + path.write_text(json.dumps(payload)) + assert _regimes(case) == {"skip:dense_model_content_insensitive"} + overrides["K8sConfig"]["k8s_model_path_in_pvc"] = "another-checkpoint" + generator.write_text(json.dumps(overrides)) + assert _regimes(case) == {"fake_fallback"} + + +def test_historical_saved_identity_retains_dense_evidence_without_reloading_sources(): + source = Path(__file__).parent / "fixtures/fpm_collection_plan_v10.json" + original = source.read_bytes() + payload = json.loads(original) + identity = saved_plan_identity(payload) + assert identity.sha256 == payload["sha256"] + assert identity.generator_config_sha256 == payload["generator_config_sha256"] + assert identity.capability.is_moe is payload["capability"]["is_moe"] + assert identity.capability.attention_source == payload["capability"]["attention_source"] + assert identity.capability.model_config.payload == payload["capability"]["model_config"]["payload"] + assert source.read_bytes() == original + + +@pytest.mark.parametrize( + "meta", + [ + {**_DENSE_SKIP, "enabled": False}, + {"enabled": True, "warm_eligible": True, "skip_reason": None}, + {"enabled": True, "warm_eligible": False, "skip_reason": "dataset_unavailable"}, + {"enabled": True, "warm_eligible": False, "skip_reason": "hybrid_state_layers_unsupported"}, + ], +) +def test_full_attention_does_not_convert_failed_or_disabled_warmup(tmp_path, meta): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + payload["kvwarm"] = meta + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + + +def test_dense_skip_with_real_kv_injection_is_rejected(tmp_path): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + for row, group in zip(payload["results"], payload["iteration_groups"], strict=True): + row["kv_seed_regime"] = "real_kv" + row["point"]["sample_reasons"] = ["kvwarm_real_kv"] + group["point"]["sample_reasons"] = ["kvwarm_real_kv"] + case[-1].write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="real KV.*warm-up"): + _regimes(case) + + +def test_unstamped_dense_point_is_not_promoted_from_skip_reason_alone(tmp_path): + case = _dense_cell(tmp_path) + payload = json.loads(case[-1].read_text()) + for row, group in zip(payload["results"], payload["iteration_groups"], strict=True): + row.pop("kv_seed_regime") + row["point"]["sample_reasons"] = [] + group["point"]["sample_reasons"] = [] + case[-1].write_text(json.dumps(payload)) + assert _regimes(case) == {"fake_fallback"} + + +def test_existing_moe_by_construction_and_legacy_contracts_are_unchanged(tmp_path): + plan, cell, directory = _synthetic_decode_cell( + tmp_path, + kvwarm={**_DENSE_SKIP, "skip_reason": "moe_tp_balanced_by_construction"}, + parallel_strategy="tp", + markers=("kvwarm_fake_fallback",), + ) + assert {row["kv_seed_regime"] for row in aggregate_cell(plan, cell, directory, expected_attempt_id="attempt")} == { + "skip:moe_tp_balanced_by_construction" + } + path = directory / "raw/pod-0/benchmark.json" + payload = json.loads(path.read_text()) + del payload["kvwarm"] + path.write_text(json.dumps(payload)) + assert {row["kv_seed_regime"] for row in aggregate_cell(plan, cell, directory, expected_attempt_id="attempt")} == { + "legacy" + } + + +def test_dense_publication_keeps_legacy_classification_until_fresh_publication(tmp_path): + import pyarrow.parquet as pq + + plan, cell, directory, _path = _dense_cell(tmp_path) + legacy_rows = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt", dense_synthetic_kv=False) + old_table, old_meta, _ = write_formal_database(plan, legacy_rows, systems_root=tmp_path / "old") + old_bytes = old_table.read_bytes(), old_meta.read_bytes() + committed = validate_formal_database_commit(old_table, old_meta, plan) + assert published_dense_synthetic_cells(committed) == frozenset() + assert pq.read_table(old_table).to_pylist() == legacy_rows + current_rows = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt") + with pytest.raises(ValueError, match="conflicting FPM database row"): + write_formal_database(plan, current_rows, systems_root=tmp_path / "old") + assert (old_table.read_bytes(), old_meta.read_bytes()) == old_bytes + + new_table, new_meta, _ = write_formal_database(plan, current_rows, systems_root=tmp_path / "new") + committed = validate_formal_database_commit(new_table, new_meta, plan) + assert published_dense_synthetic_cells(committed) == {cell.cell_id} + assert pq.read_table(new_table).to_pylist() == current_rows + assert [row["latency_ms"] for row in legacy_rows] == [row["latency_ms"] for row in current_rows] + assert (old_table.read_bytes(), old_meta.read_bytes()) == old_bytes + + +@pytest.mark.parametrize("decode_cp", [None, 1, 2]) +def test_formal_dense_publication_requires_observed_decode_context_parallel_size_one(tmp_path, decode_cp): + plan, cell, directory, path = _dense_cell(tmp_path) + # DCP2 is valid with TP2, but differs from this frozen CP1 collection. + cell = replace(cell, topology=replace(cell.topology, tp=2)) + plan.cells = (cell,) + payload = json.loads(path.read_text()) + parallel = payload["engine"]["parallel"] + parallel["tensor_parallel_size"] = 2 + if decode_cp is None: + parallel.pop("decode_context_parallel_size") + else: + parallel["decode_context_parallel_size"] = decode_cp + path.write_text(json.dumps(payload)) + raw = path.read_bytes() + rows = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt") + expected = "skip:dense_model_content_insensitive" if decode_cp == 1 else "fake_fallback" + assert {row["kv_seed_regime"] for row in rows} == {expected} + legacy = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt", dense_synthetic_kv=False) + assert {row["kv_seed_regime"] for row in legacy} == {"fake_fallback"} + assert [row["latency_ms"] for row in rows] == [row["latency_ms"] for row in legacy] + table, metadata, skipped = write_formal_database(plan, rows, systems_root=tmp_path / "published") + assert not skipped + committed = validate_formal_database_commit(table, metadata, plan) + assert published_dense_synthetic_cells(committed) == ({cell.cell_id} if decode_cp == 1 else set()) + assert path.read_bytes() == raw + + +@pytest.mark.parametrize("change", ["removed", "other-cell", "empty", "duplicate", "unknown-policy"]) +def test_new_dense_publication_requires_bound_classification_marker(tmp_path, change): + plan, cell, directory, _path = _dense_cell(tmp_path) + rows = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt") + table, meta, _ = write_formal_database(plan, rows, systems_root=tmp_path / "systems") + payload = json.loads(meta.read_text()) + if change == "removed": + payload.pop("kv_seed_classification") + elif change == "unknown-policy": + payload["kv_seed_classification"]["policy"] = "unknown" + else: + payload["kv_seed_classification"]["cells"] = { + "other-cell": ["different-cell"], + "empty": [], + "duplicate": [cell.cell_id, cell.cell_id], + }[change] + meta.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="KV seed classification"): + validate_formal_database_commit(table, meta, plan) + with pytest.raises(ValueError, match="KV seed classification"): + write_formal_database(plan, rows, systems_root=tmp_path / "systems") + + +def test_merged_database_does_not_upgrade_old_dense_cells(tmp_path): + import pyarrow.parquet as pq + + plan, cell, directory, _path = _dense_cell(tmp_path) + old = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt", dense_synthetic_kv=False) + table, meta, _ = write_formal_database(plan, old, systems_root=tmp_path / "systems") + new = aggregate_cell(plan, cell, directory, expected_attempt_id="attempt") + for row in new: + row["cell_id"] = "another-cell" + write_formal_database(plan, new, systems_root=tmp_path / "systems") + committed = validate_formal_database_commit(table, meta, plan) + assert published_dense_synthetic_cells(committed) == {"another-cell"} + assert [row for row in pq.read_table(table).to_pylist() if row["cell_id"] == cell.cell_id] == old + + +@pytest.mark.parametrize("observed", ["full", "convolution", "unreported"]) +def test_public_readiness_uses_corroborated_dense_regime(quality_case, observed): # noqa: F811 + case = quality_case + _change_decode(case, fake="all", kvwarm=_DENSE_SKIP) + for cell in case["plan"].cells: + if cell.workload_kind != "decode": + continue + for path in (case["campaign"] / "cells" / cell.cell_id / "raw").rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + if observed != "unreported": + groups = [_FULL] if observed == "full" else [_FULL, {"type": "ShortConvSpec"}] + payload["engine"] = _engine(case["plan"], cell, groups=groups) + _add_measurement_protocol(payload) + path.write_text(json.dumps(payload)) + before = {path: path.read_bytes() for path in case["root"].rglob("*") if path.is_file()} + report = assess_readiness(case["request"], case["root"], case["checkpoint"].parent) + decode = next(row for row in report["selected_cells"] if row["cell"]["workload_kind"] == "decode") + assert decode["formal_regime_counts"], decode["blockers"] + assert decode["native_regime_counts"] == {"fake_fallback": 12} + assert decode["direct_eligible_points"] == (12 if observed == "full" else 0) + assert decode["status"] == ("ready" if observed == "full" else "blocked") + assert decode["formal_regime_counts"] == { + "skip:dense_model_content_insensitive" if observed == "full" else "fake_fallback": 12 + } + assert "do not assess or upgrade an existing published table" in decode["classification_scope"] + assert all(rank["verified"] for rank in decode["rank_diagnostics"]) + assert {path: path.read_bytes() for path in case["root"].rglob("*") if path.is_file()} == before + + +@pytest.mark.parametrize("legacy", [True, False]) +def test_quality_source_revalidates_exact_old_or_new_published_classification(quality_case, legacy): # noqa: F811 + import pyarrow as pa + import pyarrow.parquet as pq + + case = quality_case + _change_decode(case, fake="all", kvwarm=_DENSE_SKIP) + all_rows = [] + for cell in case["plan"].cells: + directory = case["campaign"] / "cells" / cell.cell_id + if cell.workload_kind == "decode": + for path in (directory / "raw").rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + payload["engine"] = _engine(case["plan"], cell) + _add_measurement_protocol(payload) + path.write_text(json.dumps(payload)) + all_rows.extend( + aggregate_cell( + case["plan"], cell, directory, expected_attempt_id="synthetic-source", dense_synthetic_kv=not legacy + ) + ) + # Replace only fixture setup, before the publication under test exists. + for path in (case["root"] / "systems/data").rglob("fpm_forward_perf.*"): + path.unlink() + table, meta, _ = write_formal_database(case["plan"], all_rows, systems_root=case["root"] / "systems/data") + before = table.read_bytes(), meta.read_bytes() + source, _campaign, _checkpoint, _evidence = _collection_source(case["request"], case["root"]) + assert source.sha256 == case["plan"].sha256 + readiness = assess_readiness(case["request"], case["root"], case["checkpoint"].parent) + decode = next(row for row in readiness["selected_cells"] if row["cell"]["workload_kind"] == "decode") + assert decode["direct_eligible_points"] == 12 + assert "do not assess or upgrade an existing published table" in decode["classification_scope"] + assert (table.read_bytes(), meta.read_bytes()) == before + published_rows = pq.read_table(table).to_pylist() + assert {row["kv_seed_regime"] for row in published_rows if row["workload_kind"] == "decode"} == { + "fake_fallback" if legacy else "skip:dense_model_content_insensitive" + } + published_rows[0]["latency_ms"] *= 2 + pq.write_table(pa.Table.from_pylist(published_rows), table) + with pytest.raises(ValueError, match="does not match its commit record"): + _collection_source(case["request"], case["root"]) + changed_meta = json.loads(meta.read_text()) + changed_meta["parquet_sha256"] = hashlib.sha256(table.read_bytes()).hexdigest() + meta.write_text(json.dumps(changed_meta)) + with pytest.raises(ValueError, match="published formal rows differ"): + _collection_source(case["request"], case["root"]) + + +@pytest.mark.parametrize("legacy", [True, False]) +def test_public_finalization_preserves_old_or_new_dense_classification(tmp_path, legacy): + import pyarrow.parquet as pq + + _request, root = build_completed_collection(tmp_path) + campaign = next((root / "fpm-artifacts").iterdir()) + plan = load_repeatability_source(campaign) + rows = [] + for cell in plan.cells: + directory = campaign / "cells" / cell.cell_id + if cell.workload_kind == "decode": + path = directory / "raw/pod-0/benchmark.json" + payload = json.loads(path.read_text()) + payload["engine"] = _engine(plan, cell) + payload["kvwarm"] = _DENSE_SKIP + for row, group in zip(payload["results"], payload["iteration_groups"], strict=True): + row["kv_seed_regime"] = "fake_fallback" + row["point"]["sample_reasons"] = ["kvwarm_fake_fallback"] + group["point"]["sample_reasons"] = ["kvwarm_fake_fallback"] + path.write_text(json.dumps(payload)) + rows.extend( + aggregate_cell( + plan, + cell, + directory, + expected_attempt_id=f"attempt-{cell.workload_kind}", + dense_synthetic_kv=not legacy, + ) + ) + for path in (root / "systems/data").rglob("fpm_forward_perf.*"): + path.unlink() + table, _meta, _ = write_formal_database(plan, rows, systems_root=root / "systems/data") + before = {path: path.read_bytes() for path in root.rglob("*") if path.is_file()} + resolved = tmp_path / "resolved" + assert ( + cli.main( + [ + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(resolved), + ] + ) + == 0 + ) + assert {path: path.read_bytes() for path in root.rglob("*") if path.is_file()} == before + exported = next(resolved.rglob("fpm_forward_perf.parquet")) + assert exported.read_bytes() == table.read_bytes() + assert { + row["kv_seed_regime"] for row in pq.read_table(exported).to_pylist() if row["workload_kind"] == "decode" + } == {"fake_fallback" if legacy else "skip:dense_model_content_insensitive"} diff --git a/python/aisimulate/tests/unit/collector/test_fpm_exec.py b/python/aisimulate/tests/unit/collector/test_fpm_exec.py index 27090c792..6c2ee4e36 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_exec.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_exec.py @@ -891,7 +891,13 @@ def test_fpm_exec_consumes_only_contract_environment(): consumed = set(re.findall(r"\bFPM_[A-Z0-9_]+\b", script)) assert {"FPM_NODE_RANK", "FPM_MASTER_ADDR", "FPM_BENCHMARK_OUTPUT_PATH"} <= consumed - allowed = set(FPM_ENV_EXPORTED_VARS) | {"FPM_COMPLETION_BARRIER_TIMEOUT_SECONDS", "FPM_READINESS_TIMEOUT_SECONDS"} + allowed = set(FPM_ENV_EXPORTED_VARS) | { + "FPM_COMPLETION_BARRIER_TIMEOUT_SECONDS", + "FPM_READINESS_TIMEOUT_SECONDS", + "FPM_SLURM_CPUS_PER_TASK", + "FPM_SLURM_CPU_BIND", + "FPM_LOCAL_GPU_COUNT", + } assert consumed <= allowed, sorted(consumed - allowed) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_execution_protocol.py b/python/aisimulate/tests/unit/collector/test_fpm_execution_protocol.py new file mode 100644 index 000000000..344f120b8 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_execution_protocol.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Controls for observed execution, independent of planned capture expectations.""" + +import copy +import json + +import pytest +from collector.fpm_forward import execution_evidence, repeatability, runtime_probe +from collector.fpm_forward.measurement_evidence import compare_measurements, extract_measurement_evidence + +from .test_fpm_measurement_evidence import _add_execution_protocol, _artifacts, _edit, _remove_execution_protocol +from .test_fpm_profile_collection import no_models_or_timing_data # noqa: F401 +from .test_fpm_repeatability import _args, _fake_collector, campaign # noqa: F401 +from .test_fpm_runtime_probe import _inputs + +pytestmark = pytest.mark.unit + + +def _execution(payload, *, clock_offset=0.0): + _add_execution_protocol(payload, clock_offset=clock_offset) + + +def _pair(tmp_path, *, factor=1.0): + left_root, right_root = tmp_path / "source", tmp_path / "repeat" + left = _artifacts(left_root) + right = _artifacts(right_root, run="run-b", factor=factor) + _edit(left_root, _execution) + _edit(right_root, lambda payload: _execution(payload, clock_offset=9999)) + return left_root, right_root, left, right + + +def _compare(roots): + left_root, right_root, left_collection, right_collection = roots + left = extract_measurement_evidence(left_root, left_collection) + right = extract_measurement_evidence(right_root, right_collection) + key = next(iter(left["points"])) + return left, right, compare_measurements(left, right, point_key=key) + + +def test_observed_dispatch_keeps_slow_samples_and_process_clocks_separate(tmp_path): + left, right, comparison = _compare(_pair(tmp_path, factor=6)) + assert comparison["status"] == "comparable" + assert "observed_per_call_graph_dispatch" not in comparison["unobserved"] + assert "kv_cache_tensors" in comparison["unobserved"] + assert right["warmup_evidence"]["0"]["records"][0]["validation"]["status"] == "not_performed" + point = next(iter(right["points"].values())) + assert point["ranks"]["0"]["raw_fpms"][0]["wall_time"] == pytest.approx(0.06) + assert point["wall_time_seconds"] == pytest.approx(next(iter(left["points"].values()))["wall_time_seconds"] * 6) + + +def test_async_and_sync_observations_do_not_compare_even_when_launch_defaults_match(tmp_path): + roots = _pair(tmp_path) + _edit(roots[1], lambda payload: payload["engine"]["scheduler"].update(async_scheduling=True)) + assert "effective async_scheduling differs" in " ".join(_compare(roots)[2]["reasons"]) + assert _compare(roots)[2]["status"] == "mismatch" + + +def test_missing_scheduling_is_not_inferred_from_same_config(tmp_path): + roots = _pair(tmp_path) + _edit(roots[1], lambda payload: payload.pop("engine")) + comparison = _compare(roots)[2] + assert comparison["status"] == "unestablished" + assert "async_scheduling" in " ".join(comparison["reasons"]) + + +def test_older_observations_are_unknown_not_a_proven_protocol_mismatch(tmp_path): + roots = _pair(tmp_path) + + def older(payload): + payload["measurement_protocol"].pop("execution_evidence") + payload.pop("warmup_evidence") + for row in payload["results"]: + measurement = row["fpms"][0]["benchmark_measurement"] + measurement["preparation"].pop("warmup_records_before") + for raw in measurement["raw_fpms"]: + raw.pop("benchmark_sample") + + _edit(roots[0], older) + left, right, comparison = _compare(roots) + for key in left["points"]: + comparison = compare_measurements(left, right, point_key=key) + assert comparison["status"] == "unestablished", comparison + assert "measurement_protocol differs" not in comparison["reasons"] + + +def test_actual_dispatch_mismatch_is_detected_despite_equal_point_expectations(tmp_path): + roots = _pair(tmp_path) + _edit( + roots[1], + lambda payload: payload["results"][0]["fpms"][0]["benchmark_measurement"]["raw_fpms"][0]["benchmark_sample"][ + "cudagraph" + ].update(runtime_mode="FULL"), + ) + assert _compare(roots)[2]["status"] == "mismatch" + + +@pytest.mark.parametrize("both", [False, True]) +def test_unavailable_dispatch_remains_unknown_instead_of_eager(tmp_path, both): + roots = _pair(tmp_path) + + def unavailable(payload): + graph = payload["results"][0]["fpms"][0]["benchmark_measurement"]["raw_fpms"][0]["benchmark_sample"][ + "cudagraph" + ] + graph.update( + status="unavailable", runtime_mode=None, num_unpadded_tokens=None, num_padded_tokens=None, num_paddings=None + ) + + if both: + _edit(roots[0], unavailable) + _edit(roots[1], unavailable) + comparison = _compare(roots)[2] + assert comparison["status"] == "unestablished" + assert "observed_per_call_graph_dispatch" in comparison["unobserved"] + + +@pytest.mark.parametrize("change", ["clock", "index", "contract", "warmup", "padding", "retained"]) +def test_contradictory_sample_evidence_rejects(tmp_path, change): + roots = _pair(tmp_path) + + def mutate(payload): + retained = payload["results"][0]["fpms"][0] + sample = retained["benchmark_measurement"]["raw_fpms"][0]["benchmark_sample"] + if change == "clock": + sample["timing"]["end_monotonic"] += 1 + elif change == "index": + sample["forward_index"] = True + elif change == "contract": + payload["measurement_protocol"]["execution_evidence"]["timing_clock"] = "global" + elif change == "warmup": + retained["benchmark_measurement"]["preparation"]["warmup_records_before"] = 2 + elif change == "padding": + sample["cudagraph"]["num_paddings"] = 1 + else: + retained["benchmark_sample"] = copy.deepcopy(sample) + + _edit(roots[1], mutate) + with pytest.raises(ValueError): + _compare(roots) + + +def test_observed_preparation_change_affects_full_grid_but_not_cross_context_identity(tmp_path): + roots = _pair(tmp_path) + _edit( + roots[1], + lambda payload: payload["warmup_evidence"]["records"][0]["requested_shape"].update(prompt_lengths=[16]), + ) + left, right, comparison = _compare(roots) + assert comparison["status"] == "mismatch" + comparison = compare_measurements(left, right, point_key=next(iter(left["points"])), same_context=False) + assert comparison["identity_status"] == "comparable" + assert comparison["status"] == "unestablished" + + +def test_running_warmup_cannot_establish_completed_preparation(tmp_path): + roots = _pair(tmp_path) + for root in roots[:2]: + _edit( + root, + lambda payload: payload["warmup_evidence"]["records"][0].update(status="running", forward_index_end=None), + ) + left, _, comparison = _compare(roots) + assert comparison["status"] == "unestablished" + measurement = next(iter(left["points"].values()))["ranks"]["0"] + assert measurement["observed_execution"]["warmup_prefix_sha256"] is None + + +@pytest.mark.parametrize( + "mutation", ["future_count", "omitted_record", "late_forward", "early_forward", "missing_end", "count"] +) +def test_impossible_warmup_lifecycle_rejects(tmp_path, mutation): + roots = _pair(tmp_path) + + def change(payload): + history = payload["warmup_evidence"]["records"] + record = history[0] + if mutation == "future_count": + record["completed_points_before"] = 99 + elif mutation == "omitted_record": + payload["results"][0]["fpms"][0]["benchmark_measurement"]["preparation"]["warmup_records_before"] = 0 + elif mutation == "late_forward": + record.update(forward_index_start=4, forward_index_end=7) + elif mutation == "early_forward": + history.append({**copy.deepcopy(record), "kind": "eager_shape", "completed_points_before": 1}) + payload["results"][1]["fpms"][0]["benchmark_measurement"]["preparation"]["warmup_records_before"] = 2 + elif mutation == "missing_end": + record.update(status="failed", forward_index_end=None) + else: + record["observed_forward_count"] = None + + _edit(roots[1], change) + with pytest.raises(ValueError, match="warmup|preparation"): + _compare(roots) + + +@pytest.mark.parametrize("observed_attempt", [False, True]) +def test_failed_historical_warmup_is_visible_and_distinguishable(tmp_path, observed_attempt): + roots = _pair(tmp_path) + + def failed(payload): + record = payload["warmup_evidence"]["records"][0] + record.update(kind="eager_shape", status="failed", validation={"status": "failed", "reason": "shape_mismatch"}) + if not observed_attempt: + record.update(forward_index_start=None, forward_index_end=0, observed_forward_count=0) + record["validation"]["status"] = "not_performed" + + _edit(roots[1], failed) + assert _compare(roots)[2]["status"] == "mismatch" + _edit(roots[0], failed) + left, _, comparison = _compare(roots) + assert comparison["status"] == "comparable" + assert left["warmup_evidence"]["0"]["records"][0]["status"] == "failed" + + +def test_unreferenced_later_warmup_does_not_invalidate_completed_measurements(tmp_path): + roots = _pair(tmp_path) + + def unfinished(payload): + payload["warmup_evidence"]["records"].append( + { + "kind": "eager_shape", + "status": "running", + "validation": {"status": "not_performed"}, + "completed_points_before": 2, + "forward_index_start": 5, + "forward_index_end": None, + "observed_forward_count": 0, + } + ) + + _edit(roots[1], unfinished) + assert _compare(roots)[2]["status"] == "comparable" + + +@pytest.mark.parametrize( + "policy,eager,warning", [("runtime", False, False), ("explicit", False, True), ("explicit", True, False)] +) +def test_probe_graph_warning_is_early_actionable_and_does_not_reject_explicit( + tmp_path, monkeypatch, policy, eager, warning +): + launch, manifest, _ = _inputs(tmp_path, graph_policy=policy) + launch["collection"]["enforce_eager"] = eager + result = runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=tmp_path / "preview") + assert result["status"] == "preview" + messages = " ".join(result["configurations"]["tp4"]["diagnostics"]) + assert ("apply to prefill only" in messages) == warning + if warning: + assert "Compatibility is unresolved until observed" in messages + + +def test_legacy_frozen_assessment_keeps_exact_original_evidence(campaign, tmp_path, monkeypatch): # noqa: F811 + from .test_fpm_repeatability import _mutate_native + + _mutate_native(campaign[1], _remove_execution_protocol) + original_freeze = repeatability.freeze_repeatability_plan + + def old_freeze(*args, **kwargs): + return original_freeze(*args, **{**kwargs, "observation_evidence_version": 1}) + + monkeypatch.setattr(repeatability, "freeze_repeatability_plan", old_freeze) + _fake_collector(monkeypatch) + collect = repeatability.run_collection + + def historical_repeat(*args, **kwargs): + from pathlib import Path + + result = collect(*args, **kwargs) + _mutate_native(Path(kwargs["artifact_root"]), _remove_execution_protocol) + return result + + monkeypatch.setattr(repeatability, "run_collection", historical_repeat) + args = _args(campaign, tmp_path) + original_report = repeatability.run_repeatability(**args) + frozen_path = args["output_dir"] / repeatability.PLAN_FILENAME + frozen_bytes = frozen_path.read_bytes() + frozen = json.loads(frozen_bytes) + assert "observation_evidence_version" not in frozen + assert "scheduling" not in frozen["cells"][0]["execution"] + monkeypatch.setattr(repeatability, "freeze_repeatability_plan", original_freeze) + assessed = repeatability.assess_repeatability(frozen, original_report, cv_threshold=0.05) + assert assessed["status"] == original_report["status"] + resumed = repeatability.run_repeatability(**{**args, "resume": True}) + assert resumed["status"] == original_report["status"] + assert frozen_path.read_bytes() == frozen_bytes + + +def test_requested_default_does_not_claim_an_effective_async_value(): + assert execution_evidence._requested_scheduling({}) == {"source": "runtime_default", "async_scheduling": None} + assert execution_evidence._requested_scheduling({"--no-async-scheduling": True}) == { + "source": "generated_flag", + "async_scheduling": False, + } diff --git a/python/aisimulate/tests/unit/collector/test_fpm_forward.py b/python/aisimulate/tests/unit/collector/test_fpm_forward.py index 86b04ca73..e2528c698 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_forward.py @@ -5,6 +5,7 @@ import argparse import base64 +import copy import csv import hashlib import importlib.metadata @@ -17,7 +18,12 @@ import pytest from collector.fpm_forward.capabilities import resolve_model_capability -from collector.fpm_forward.config import FPMCollectionOptions, PrefillSamplingProfile, add_fpm_arguments +from collector.fpm_forward.config import ( + FPMCollectionOptions, + PrefillSamplingProfile, + add_fpm_arguments, + reject_fpm_arguments_without_fpm, +) from collector.fpm_forward.database import ( aggregate_cell, validate_formal_database_commit, @@ -43,6 +49,9 @@ "collector/fpm_forward/runner.py": b"runner-content", "collector/fpm_forward/runtime/fpm_exec.sh": b"runtime-content", "collector/fpm_forward/runtime/preflight.py": b"preflight-content", + "collector/fpm_forward/runtime/fpm_memory_observer.py": b"memory-observer-content", + "collector/fpm_forward/runtime/fpm_memory_worker.py": b"memory-worker-content", + "collector/fpm_forward/runtime/fpm_memory_scheduler.py": b"memory-scheduler-content", } @@ -152,7 +161,7 @@ def test_options_leave_point_generation_to_dynamo(): assert options.parallel_presets == ("auto",) assert options.to_dict()["point_source"] == "dynamo_native_self_benchmark" assert options.to_dict()["measurement_repeats"] == 1 - assert options.max_prefill_isl == 8192 + assert options.max_prefill_isl is None assert options.max_prefill_batch_size is None assert options.vllm_max_model_len == -1 assert options.prefill_sampling.max_total_prefill_tokens == 8192 @@ -184,6 +193,125 @@ def test_prefill_limits_expose_no_cli_aliases(): assert "--fpm-max-prefill-bs" not in help_text +def test_runtime_graph_policy_leaves_graph_dependent_axes_unresolved(): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + options = FPMCollectionOptions.from_args( + parser.parse_args(["--fpm-max-gpus", "4", "--fpm-prefill-cudagraph-policy", "runtime"]) + ) + + assert options.max_prefill_cudagraph_size is None + sampling = options.prefill_sampling.to_dict() + assert sampling["cudagraph_policy"] == "runtime" + for field in ( + "cudagraph_capture_sizes", + "cudagraph_capture_size_count", + "max_cudagraph_capture_size", + "new_token_axis_points", + "new_token_axis_point_count", + "prefill_max_new_token_samples", + ): + assert sampling[field] is None + assert sampling["prefill_max_kv_read_token_samples"] > 0 + + +def test_explicit_graph_policy_preserves_legacy_capture_settings_and_identity(): + legacy = FPMCollectionOptions.from_args(_args()) + explicit = FPMCollectionOptions.from_args(_args(fpm_prefill_cudagraph_policy="explicit")) + assert explicit.max_prefill_cudagraph_size == 2048 + assert explicit.to_dict() == legacy.to_dict() + assert "gpu_memory_utilization" not in legacy.to_dict() + + custom = FPMCollectionOptions.from_args(_args(fpm_max_prefill_cudagraph_size=1000)) + assert custom.prefill_cudagraph_policy == "explicit" + assert custom.prefill_sampling.cudagraph_capture_sizes[-1] == 1000 + + +def test_runtime_graph_policy_rejects_explicit_capture_size(): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + args = parser.parse_args( + [ + "--fpm-max-gpus", + "4", + "--fpm-prefill-cudagraph-policy", + "runtime", + "--fpm-max-prefill-cudagraph-size", + "512", + ] + ) + with pytest.raises(ValueError, match="runtime cannot use --fpm-max-prefill-cudagraph-size"): + FPMCollectionOptions.from_args(args) + + +@pytest.mark.parametrize("value", ["0", "-0.1", "1.01", "nan", "inf", "-inf", "true"]) +def test_gpu_memory_fraction_cli_rejects_invalid_values(value): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + with pytest.raises(SystemExit) as error: + parser.parse_args([f"--fpm-gpu-memory-utilization={value}"]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("value", [True, 0, -0.1, 1.01, float("nan"), float("inf")]) +def test_gpu_memory_fraction_options_reject_invalid_values(value): + with pytest.raises(ValueError, match="gpu-memory-utilization must be finite"): + FPMCollectionOptions.from_args(_args(fpm_gpu_memory_utilization=value)) + + +@pytest.mark.parametrize( + "arguments", + [ + ["--fpm-prefill-cudagraph-policy", "runtime"], + ["--fpm-gpu-memory-utilization", "0.9"], + ], +) +def test_graph_policy_and_memory_fraction_require_fpm_collection(arguments): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + args = parser.parse_args(arguments) + args.ops = ["gemm"] + with pytest.raises(ValueError, match="FPM-only arguments require --ops fpm_forward"): + reject_fpm_arguments_without_fpm(args) + + +@pytest.mark.parametrize("option", ["max-model-len", "max-num-batched-tokens", "max-num-seqs"]) +@pytest.mark.parametrize("value", ["0", "-1", "1.5"]) +def test_runtime_limit_cli_rejects_nonpositive_values_except_context_auto_fit(option, value): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + if option == "max-model-len" and value == "-1": + args = parser.parse_args([f"--fpm-{option}", value]) + assert args.fpm_max_model_len == -1 + return + with pytest.raises(SystemExit) as error: + parser.parse_args([f"--fpm-{option}", value]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("option", ["max-model-len", "max-num-batched-tokens", "max-num-seqs"]) +def test_runtime_limits_require_fpm_collection(option): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + args = parser.parse_args([f"--fpm-{option}", "64"]) + args.ops = ["gemm"] + with pytest.raises(ValueError, match="FPM-only arguments require --ops fpm_forward"): + reject_fpm_arguments_without_fpm(args) + + +@pytest.mark.parametrize( + "arguments, message", + [ + ({"fpm_max_num_batched_tokens": 1024, "fpm_max_prefill_isl": 2048}, "prefill-isl exceeds"), + ({"fpm_max_num_seqs": 8, "fpm_max_prefill_batch_size": 16}, "prefill-batch-size exceeds"), + ({"fpm_max_num_seqs": 8, "fpm_max_decode_batch_size": 16}, "decode-batch-size exceeds"), + ], +) +def test_phase_limits_cannot_exceed_shared_limits(arguments, message): + with pytest.raises(ValueError, match=message): + FPMCollectionOptions.from_args(_args(**arguments)) + + def test_prefill_sampling_profile_keeps_vllm_strides_and_exact_endpoint(): short = PrefillSamplingProfile.build(max_isl=1000, max_batch_size=16) @@ -309,6 +437,39 @@ def test_plan_contains_only_cell_matrix_and_native_point_contract(tmp_path, expl assert "runtime_overlay" not in payload +def test_kv_warmup_defaults_are_frozen_in_plan_identity(): + options = FPMCollectionOptions.from_args( + _args(fpm_parallel_axes=["dp", "moe_ep"], fpm_dp_sizes=[4], fpm_moe_ep_sizes=[4]) + ) + kwargs = { + "backend": "vllm", + "model_path": "nvidia/GLM-5.2-NVFP4", + "system": "b200_sxm", + "selected_ops": {"dsa_context_module", "dsa_generation_module"}, + "options": options, + } + default = build_collection_plan(**kwargs) + explicit_on = build_collection_plan( + **kwargs, + generator_overrides={ + "K8sConfig": { + "extra_env": [ + {"name": "DYN_BENCH_KV_WARMUP", "value": "on"}, + {"name": "DYN_BENCH_PREFILL_REAL_SEED", "value": "on"}, + ] + } + }, + ) + old_policy = build_collection_plan( + **kwargs, + generator_overrides={"K8sConfig": {"extra_env": [{"name": "DYN_BENCH_PREFILL_REAL_SEED", "value": "off"}]}}, + ) + assert default.aic_revision == explicit_on.aic_revision == old_policy.aic_revision + assert default.sha256 == explicit_on.sha256 + assert default.generator_config_sha256 != old_policy.generator_config_sha256 + assert default.sha256 != old_policy.sha256 + + def test_backend_policy_is_deeply_immutable(): source = {"nested": {"values": [1, 2]}} policy = BackendPolicy("baseline", source, {"runtime.mode": "FULL"}) @@ -446,6 +607,27 @@ def estimate(*_args, **kwargs): assert "fpm_forward: dropped 3/6 (topology, kv_dtype) cell groups (memory budget" in caplog.text +def test_model_memory_admission_uses_declared_gpu_fraction(monkeypatch): + monkeypatch.setattr( + "collector.fpm_forward.memory_admission.KVCacheEstimator.from_request", + lambda *_args, **_kwargs: SimpleNamespace(breakdown={"non_kv_bytes": 60, "gpu_memory_capacity_bytes": 100}), + ) + kwargs = { + "backend": "vllm", + "model_path": "nvidia/GLM-5.2-NVFP4", + "system": "b200_sxm", + "selected_ops": {"dsa_context_module", "dsa_generation_module"}, + } + admitted = build_collection_plan( + **kwargs, options=FPMCollectionOptions.from_args(_args(fpm_gpu_memory_utilization=0.9)) + ) + for decision in admitted.topology_memory_admission: + assert decision.estimates[0].to_dict()["gpu_memory_budget_bytes"] == 90 + assert decision.estimates[0].to_dict()["headroom_bytes"] == 30 + with pytest.raises(ValueError, match="memory admission rejected every FPM topology"): + build_collection_plan(**kwargs, options=FPMCollectionOptions.from_args(_args(fpm_gpu_memory_utilization=0.6))) + + def test_plan_identity_ignores_memory_estimator_error_text(monkeypatch): """Estimator failure diagnostics vary across runs and hosts (transient network errors, host paths); hashing them would spuriously invalidate @@ -2233,6 +2415,142 @@ def test_native_aggregation_rejects_native_coordinate_collision(tmp_path): aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") +def _prefill_cell_with_zero_kv_provenance_duplicates(tmp_path, *, seed_samples=3, canonical_first=False): + plan, cell, cell_dir = _synthetic_plan_and_cell(tmp_path) + count = seed_samples + 1 + canonical_id = 1 if canonical_first else count + for path in (cell_dir / "raw").glob("*/benchmark*.json"): + payload = json.loads(path.read_text()) + template = payload["iteration_groups"][0] + payload["results"] = [] + payload["iteration_groups"] = [] + for benchmark_id in range(1, count + 1): + ordinary = benchmark_id == canonical_id + group = copy.deepcopy(template) + group["benchmark_id"] = benchmark_id + point = group["point"] + point.update( + benchmark_id=benchmark_id, + total_kv_read_tokens=0, + sample_reasons=["post_capture"] + ([] if ordinary else ["prefill_real_seed"]), + partition=None, + rows=None, + ) + for item in group["rank_results"]: + fpm = item["fpms"][0] + fpm["counter_id"] = benchmark_id + fpm["scheduled_requests"]["sum_prefill_kv_tokens"] = 0 + # The ordinary sample is deliberately slower than the seeded + # duplicates: selection must not bias toward minimum latency. + fpm["wall_time"] = (0.010 if ordinary else 0.002) + item["dp_rank"] * 0.001 + group["wall_time"] = max(item["fpms"][0]["wall_time"] for item in group["rank_results"]) + payload["iteration_groups"].append(group) + payload["results"].append( + { + "point": point, + "kv_seed_regime": "not_applicable" if ordinary else "real_prefix", + "fpms": group["rank_results"][payload["dp"]["rank"]]["fpms"], + } + ) + payload["coverage"].update(expected_points=count, completed_points=count) + payload["timing"]["measured_iteration_seconds"] = sum( + group["wall_time"] for group in payload["iteration_groups"] + ) + path.write_text(json.dumps(payload)) + return plan, cell, cell_dir + + +@pytest.mark.parametrize("seed_samples", [3, 4]) +@pytest.mark.parametrize("canonical_first", [False, True]) +def test_zero_kv_prefill_duplicates_publish_unique_ordinary_sample(tmp_path, caplog, seed_samples, canonical_first): + import pyarrow.parquet as pq + + plan, cell, cell_dir = _prefill_cell_with_zero_kv_provenance_duplicates( + tmp_path, seed_samples=seed_samples, canonical_first=canonical_first + ) + raw_before = {path: path.read_bytes() for path in (cell_dir / "raw").rglob("*.json")} + + with caplog.at_level("INFO", logger="collector.fpm_forward.native_artifact"): + rows = aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") + + assert len(rows) == 1 + assert rows[0]["latency_ms"] == pytest.approx(11.0) + assert rows[0]["total_kv_read_tokens"] == 0 + assert rows[0]["kv_seed_regime"] == "n/a" + assert rows[0]["measurement_repeats"] == 1 + assert f"consolidated {seed_samples} zero-KV prefill provenance duplicate sample(s)" in caplog.text + assert "context-clamped duplicate" not in caplog.text + + parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=tmp_path / "systems") + assert skipped == () + assert pq.read_table(parquet).to_pylist() == rows + assert json.loads(metadata.read_text())["row_count"] == 1 + committed = validate_formal_database_commit(parquet, metadata, plan) + validate_formal_database_commit( + parquet, + metadata, + plan, + expected_attempt_ids={cell.cell_id: "attempt"}, + expected_cell_rows=committed["cell_rows"], + ) + write_formal_database(plan, rows, systems_root=tmp_path / "systems") + assert all(path.read_bytes() == original for path, original in raw_before.items()) + + +@pytest.mark.parametrize( + "conflict", + [ + "duplicate_ordinary", + "no_ordinary", + "missing_provenance", + "positive_kv", + "expected_cudagraph_mode", + "expected_capture_size", + "padding_tokens", + "partition", + "rows", + "sample_reasons", + "context_clamped", + ], +) +def test_prefill_provenance_consolidation_rejects_ambiguous_duplicates(tmp_path, conflict): + plan, cell, cell_dir = _prefill_cell_with_zero_kv_provenance_duplicates(tmp_path) + for path in (cell_dir / "raw").glob("*/benchmark*.json"): + payload = json.loads(path.read_text()) + for index, (row, group) in enumerate(zip(payload["results"], payload["iteration_groups"], strict=True)): + point = row["point"] + if conflict == "positive_kv": + point["total_kv_read_tokens"] = 128 + for item in group["rank_results"]: + item["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = 128 + row["fpms"] = group["rank_results"][payload["dp"]["rank"]]["fpms"] + elif conflict == "no_ordinary": + row["kv_seed_regime"] = "real_prefix" + point["sample_reasons"] = ["post_capture", "prefill_real_seed"] + elif index == 0: + if conflict == "duplicate_ordinary": + row["kv_seed_regime"] = "not_applicable" + point["sample_reasons"] = ["post_capture"] + elif conflict == "missing_provenance": + row.pop("kv_seed_regime") + elif conflict == "context_clamped": + point["sample_reasons"].append("context_clamped") + else: + point[conflict] = { + "expected_cudagraph_mode": "NONE", + "expected_capture_size": 512, + "padding_tokens": 255, + "partition": "different", + "rows": [{"tokens": 257}], + "sample_reasons": ["eager_tail", "prefill_real_seed"], + }[conflict] + group["point"] = point + path.write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match="unclamped samples share one key"): + aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") + + def test_backend_identity_defaults_record_auto_everywhere(): """v6: unspecified knobs record "auto" (engine decided), plumb nothing.""" diff --git a/python/aisimulate/tests/unit/collector/test_fpm_measurement_evidence.py b/python/aisimulate/tests/unit/collector/test_fpm_measurement_evidence.py new file mode 100644 index 000000000..18a4ec0f9 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_measurement_evidence.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import hashlib +import json +from dataclasses import replace + +import pytest +from collector.fpm_forward.measurement_evidence import compare_measurements, extract_measurement_evidence +from collector.fpm_forward.native_artifact import NativeCollection, NativePointMeasurement + +pytestmark = pytest.mark.unit + + +def _hash(value): + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +def _add_execution_protocol(payload, *, clock_offset=0.0): + # Synthetic observations of the documented Dynamo #15110 contract, not + # expectations reconstructed from an engine configuration. + payload["measurement_protocol"]["execution_evidence"] = { + "schema_version": 1, + "sample_field": "benchmark_sample", + "warmup_field": "warmup_evidence", + "forward_index_scope": "rank_process", + "timing_clock": "process_local_monotonic", + "cudagraph_source": "ModelRunnerOutput.cudagraph_stats", + "warmup_scope": "request_completion_and_existing_shape_validation", + } + payload["warmup_evidence"] = { + "status": "recorded", + "records": [ + { + "kind": "global", + "requested_shape": {"prompt_lengths": [8]}, + "status": "completed", + "forward_index_start": 0, + "forward_index_end": 3, + "completed_points_before": 0, + "observed_forward_count": 3, + "first_scheduled_requests": {"num_requests": 1}, + "last_scheduled_requests": {"num_requests": 1}, + "validation": {"status": "not_performed"}, + } + ], + } + forward_index = 3 + for index, row in enumerate(payload["results"]): + point = row["point"] + group = next(item for item in payload["iteration_groups"] if item["point"] == point) + fpms = [*row["fpms"], *(fpm for rank in group["rank_results"] for fpm in rank["fpms"])] + for fpm in fpms: + measurement = fpm["benchmark_measurement"] + measurement["preparation"]["warmup_records_before"] = 1 + for sample_index, raw in enumerate(measurement["raw_fpms"]): + start = clock_offset + measurement["dp_rank"] * 1000 + index + sample_index + tokens = point["total_prefill_tokens"] or point["batch_size"] + padded = point["expected_capture_size"] or tokens + raw["benchmark_sample"] = { + "sample_index": sample_index, + "forward_index": forward_index + sample_index, + "timing": { + "basis": "schedule_to_output" if sample_index == 0 else "inter_output", + "start_monotonic": start, + "end_monotonic": start + raw["wall_time"], + }, + "cudagraph": { + "status": "observed", + "runtime_mode": point["expected_cudagraph_mode"], + "num_unpadded_tokens": tokens, + "num_padded_tokens": padded, + "num_paddings": padded - tokens, + }, + } + forward_index += len(row["fpms"][0]["benchmark_measurement"]["raw_fpms"]) + + +def _remove_execution_protocol(payload): + payload["measurement_protocol"].pop("execution_evidence", None) + payload.pop("warmup_evidence", None) + fpms = [fpm for row in payload["results"] for fpm in row["fpms"]] + fpms.extend(fpm for group in payload["iteration_groups"] for rank in group["rank_results"] for fpm in rank["fpms"]) + for fpm in fpms: + measurement = fpm["benchmark_measurement"] + measurement["preparation"].pop("warmup_records_before", None) + for raw in measurement["raw_fpms"]: + raw.pop("benchmark_sample", None) + + +def _add_measurement_protocol(payload, *, seed="0", execution=True): + """Add the producer's observed metadata to a synthetic native rank artifact.""" + payload["measurement_protocol"] = { + "schema_version": 1, + "content_identity": "coordinate_rank_slot_v1", + "content_seed": seed, + "synthetic_content": "random", + "synthetic_pool_tag": "", + "prompt_hash_encoding": "uint32_le", + "independent_repetitions": 1, + "timing_metric": "scheduler_wall_time", + "input_evidence_scope": "injected_prompt_token_ids", + "unobserved": [ + "sampled_continuation_token_ids", + "kv_cache_tensors", + "recurrent_state_tensors", + "execution_history_equivalence", + ], + "preparation": { + "warmup_iterations": 5, + "prefill_real_seed": True, + "decode_real_kv_warmup": True, + "giant_kv_threshold": 32768, + "giant_kv_repeats": 3, + }, + } + for position, row in enumerate(payload["results"]): + point = row["point"] + identity = { + "point_type": point["point_type"], + "batch_size": point["batch_size"], + "total_prefill_tokens": point.get("total_prefill_tokens", 0), + "total_kv_read_tokens": point.get("total_kv_read_tokens", 0), + "partition": point.get("partition"), + "rows": point.get("rows"), + } + for rank in range(payload["dp"]["size"]): + fpm = next(group for group in payload["iteration_groups"] if group["point"] == point)["rank_results"][rank][ + "fpms" + ][0] + raw = {key: value for key, value in fpm.items() if key != "benchmark_measurement"} + requests = [{"slot": 0, "num_tokens": 4, "sha256": _hash([seed, rank, identity])}] + fpm["benchmark_measurement"] = { + "schema_version": 1, + "point_key": _hash(identity), + "dp_rank": rank, + "prompts": {"status": "recorded", "requests": requests, "sha256": _hash(requests)}, + "preparation": { + "grid_digest": payload["grid_digest"], + "completed_points_before": position, + "kv_seed_regime": row["kv_seed_regime"], + }, + "expected_internal_samples": 1, + "raw_fpms": [raw], + "estimate": {"method": "single_step", "raw_sample_indices": [0]}, + } + if rank == payload["dp"]["rank"]: + row["fpms"] = [copy.deepcopy(fpm)] + + if execution: + _add_execution_protocol(payload) + else: + _remove_execution_protocol(payload) + + +def _artifacts(tmp_path, *, run="run-a", seed="0", duplicate=False, order=(1, 2), factor=1.0): + tmp_path.mkdir(parents=True, exist_ok=True) + measurements = [] + for index, batch in enumerate((4, 8), 1): + point = { + "point_type": "prefill", + "benchmark_id": index, + "batch_size": batch, + "total_prefill_tokens": batch, + "total_kv_read_tokens": 0, + "expected_cudagraph_mode": "PIECEWISE", + "expected_capture_size": batch, + "padding_tokens": 0, + "sample_reasons": [], + } + measurements.append( + NativePointMeasurement( + point, tuple((rank, factor * (0.01 + rank * 0.001)) for rank in range(2)), "not_applicable" + ) + ) + if duplicate: + point = {**measurements[0].point, "benchmark_id": 3, "sample_reasons": ["prefill_real_seed"]} + measurements.append(NativePointMeasurement(point, ((0, 0.02), (1, 0.022)), "real_prefix")) + order = (*order, 3) + collection = NativeCollection( + points=tuple(measurements), + rank_timings=((0, 1.0, 0.03), (1, 1.0, 0.03)), + backend_version="0.27.0", + collector_attempt_id=run, + runtime_run_id=run, + runtime_grid_digest="native-grid", + ) + for rank in range(2): + results = [] + groups = [] + for index in order: + measurement = measurements[index - 1] + point = measurement.point + ranks = [ + { + "dp_rank": dp, + "fpms": [{"dp_rank": dp, "counter_id": index, "wall_time": wall, "scheduled_requests": {}}], + } + for dp, wall in measurement.rank_wall_times + ] + results.append({"point": point, "kv_seed_regime": measurement.kv_seed_regime, "fpms": ranks[rank]["fpms"]}) + groups.append({"point": point, "rank_results": ranks}) + payload = { + "artifact_type": "rank", + "dp": {"rank": rank, "size": 2}, + "run_id": run, + "grid_digest": "native-grid", + "config": {"mode": "prefill", "output_path": str(tmp_path / "benchmark.json")}, + "limits": {"block_size": 16, "max_model_len": 256000}, + "cudagraph": {"mode": "FULL_AND_PIECEWISE", "capture_sizes": [4, 8]}, + "recurrent_state": {"initialization": "unchanged", "policy": None, "uniform_bound": None}, + "engine": {"scheduler": {"async_scheduling": False}}, + "results": results, + "iteration_groups": groups, + } + _add_measurement_protocol(payload, seed=seed) + (tmp_path / f"benchmark-dp{rank}.json").write_text(json.dumps(payload)) + return collection + + +def _edit(root, action, rank=None): + paths = sorted(root.glob("benchmark-dp*.json")) if rank is None else [root / f"benchmark-dp{rank}.json"] + for path in paths: + payload = json.loads(path.read_text()) + action(payload) + path.write_text(json.dumps(payload)) + + +def _read_pair(tmp_path, **kwargs): + source = tmp_path / "source" + candidate = tmp_path / "candidate" + left = extract_measurement_evidence(source, _artifacts(source)) + right = extract_measurement_evidence(candidate, _artifacts(candidate, run="run-b", **kwargs)) + return left, right, next(iter(left["points"])) + + +def test_matching_conditions_remain_comparable_despite_slow_samples(tmp_path): + left, right, key = _read_pair(tmp_path, factor=6) + comparison = compare_measurements(left, right, point_key=key) + assert comparison["status"] == "comparable" + assert left["points"][key]["wall_time_seconds"] * 6 == right["points"][key]["wall_time_seconds"] + assert "kv_cache_tensors" in comparison["unobserved"] + assert "observed_per_call_graph_dispatch" not in comparison["unobserved"] + assert left["points"][key]["ranks"]["0"]["prompts"] != left["points"][key]["ranks"]["1"]["prompts"] + + +def test_publication_selection_preserves_ordinary_zero_kv_sample(tmp_path): + collection = _artifacts(tmp_path, duplicate=True) + evidence = extract_measurement_evidence(tmp_path, collection) + assert len(evidence["points"]) == 2 + selected = next(iter(evidence["points"].values())) + assert selected["point"]["benchmark_id"] == 1 + assert selected["kv_seed_regime"] == "not_applicable" + assert selected["wall_time_seconds"] == 0.011 + + +def test_legacy_and_missing_comparisons_remain_unestablished(tmp_path): + collection = _artifacts(tmp_path) + + def legacy(payload): + payload.pop("measurement_protocol") + for row in payload["results"]: + row["fpms"][0].pop("benchmark_measurement") + for group in payload["iteration_groups"]: + for rank in group["rank_results"]: + rank["fpms"][0].pop("benchmark_measurement") + + _edit(tmp_path, legacy) + evidence = extract_measurement_evidence(tmp_path, collection) + assert evidence["status"] == "unestablished" + comparison = compare_measurements(evidence, {}, point_key=next(iter(evidence["points"]))) + assert comparison["status"] == "unestablished" + assert compare_measurements({}, {}, point_key="absent")["status"] == "unestablished" + + +@pytest.mark.parametrize("mutation", ["missing", "invalid", "seed"]) +def test_mixed_or_conflicting_rank_protocol_rejected(tmp_path, mutation): + collection = _artifacts(tmp_path) + + def change(payload): + if mutation == "missing": + payload.pop("measurement_protocol") + elif mutation == "invalid": + payload["measurement_protocol"] = None + else: + payload["measurement_protocol"]["content_seed"] = "1" + + _edit(tmp_path, change, rank=1) + with pytest.raises(ValueError, match="measurement_protocol"): + extract_measurement_evidence(tmp_path, collection) + + +@pytest.mark.parametrize( + ("field", "value"), + [("schema_version", True), ("content_identity", "unknown"), ("independent_repetitions", 3), ("preparation", {})], +) +def test_invalid_shared_protocol_rejected(tmp_path, field, value): + collection = _artifacts(tmp_path) + _edit(tmp_path, lambda payload: payload["measurement_protocol"].update({field: value})) + with pytest.raises(ValueError, match="measurement_protocol"): + extract_measurement_evidence(tmp_path, collection) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("dp_rank", 4, "rank/point"), + ("point_key", "0" * 64, "rank/point"), + ("preparation", {"grid_digest": "wrong"}, "grid/order/KV"), + ("prompts", {"status": "recorded", "requests": [], "sha256": "a" * 64}, "no requests"), + ("expected_internal_samples", 0, "internal sample"), + ("estimate", {"method": "median", "raw_sample_indices": [0]}, "reduction"), + ("raw_fpms", [{"counter_id": 1, "dp_rank": 0, "wall_time": 99.0}], "retained timing"), + ], +) +def test_malformed_measurement_evidence_rejected(tmp_path, field, value, message): + collection = _artifacts(tmp_path) + _edit( + tmp_path, + lambda payload: payload["results"][0]["fpms"][0]["benchmark_measurement"].update({field: value}), + rank=0, + ) + with pytest.raises(ValueError, match=message): + extract_measurement_evidence(tmp_path, collection) + + +def test_prompt_hash_tampering_rejected(tmp_path): + collection = _artifacts(tmp_path) + _edit( + tmp_path, + lambda payload: payload["results"][0]["fpms"][0]["benchmark_measurement"]["prompts"].update(sha256="a" * 64), + ) + with pytest.raises(ValueError, match="aggregate hash"): + extract_measurement_evidence(tmp_path, collection) + + +@pytest.mark.parametrize("value", [999, None, True, "missing"]) +@pytest.mark.parametrize("raw", [True, False]) +def test_sample_counters_must_match_the_measured_point(tmp_path, raw, value): + collection = _artifacts(tmp_path) + + def change(payload): + fpms = [payload["results"][0]["fpms"][0]] + fpms.extend(rank["fpms"][0] for rank in payload["iteration_groups"][0]["rank_results"]) + for fpm in fpms: + samples = fpm["benchmark_measurement"]["raw_fpms"] if raw else [fpm] + for sample in samples: + if value == "missing": + sample.pop("counter_id") + else: + sample["counter_id"] = value + + _edit(tmp_path, change) + with pytest.raises(ValueError, match="counter"): + extract_measurement_evidence(tmp_path, collection) + + +def test_unavailable_prompts_and_missing_runtime_remain_unestablished(tmp_path): + collection = _artifacts(tmp_path) + + def missing(payload): + payload.pop("recurrent_state") + payload["results"][0]["fpms"][0]["benchmark_measurement"]["prompts"] = { + "status": "unavailable", + "sha256": None, + "requests": [], + } + + _edit(tmp_path, missing) + evidence = extract_measurement_evidence(tmp_path, collection) + assert evidence["status"] == "unestablished" + point = next(iter(evidence["points"].values())) + assert "injected prompt hashes unavailable" in " ".join(point["reasons"]) + assert "recurrent_state" in " ".join(evidence["reasons"]) + assert "preceding prompt evidence incomplete" in " ".join(list(evidence["points"].values())[1]["reasons"]) + + +@pytest.mark.parametrize("difference", ["seed", "prompts", "preparation", "regime", "ranks", "runtime", "estimator"]) +def test_different_recorded_conditions_are_mismatches(tmp_path, difference): + left, right, key = _read_pair(tmp_path) + point = right["points"][key] + if difference == "seed": + right["measurement_protocol"]["content_seed"] = "another" + elif difference == "prompts": + point["ranks"]["1"]["prompts"]["sha256"] = "b" * 64 + elif difference == "preparation": + right["measurement_protocol"]["preparation"]["warmup_iterations"] += 1 + elif difference == "regime": + point["regime"]["expected_capture_size"] = 16 + elif difference == "ranks": + point["ranks"].pop("1") + elif difference == "runtime": + right["runtime_identity"]["0"]["limits"]["block_size"] = 32 + else: + point["ranks"]["0"]["estimate"]["method"] = "last_step" + assert compare_measurements(left, right, point_key=key)["status"] == "mismatch" + + +def test_cross_grid_matching_prompts_do_not_prove_history_equivalence(tmp_path): + left, right, key = _read_pair(tmp_path) + right["runtime_grid_digest"] = "subset-grid" + right["points"][key]["ranks"]["0"]["preparation"]["grid_digest"] = "subset-grid" + assert compare_measurements(left, right, point_key=key)["status"] == "mismatch" + comparison = compare_measurements(left, right, point_key=key, same_context=False) + assert comparison["identity_status"] == "comparable" + assert comparison["status"] == "unestablished" + assert "cross-context" in " ".join(comparison["reasons"]) + + +def test_execution_order_changes_are_not_comparable(tmp_path): + left, right, key = _read_pair(tmp_path, order=(2, 1)) + result = compare_measurements(left, right, point_key=key) + assert result["status"] == "mismatch" + assert "preparation differs" in " ".join(result["reasons"]) + assert "execution_prefix_sha256 differs" in " ".join(result["reasons"]) + + +def test_contradictory_rank_execution_order_is_rejected(tmp_path): + collection = _artifacts(tmp_path) + + def reorder(payload): + payload["results"].reverse() + for position, row in enumerate(payload["results"]): + row["fpms"][0]["benchmark_measurement"]["preparation"]["completed_points_before"] = position + + _edit(tmp_path, reorder, rank=1) + with pytest.raises(ValueError, match="measurement order"): + extract_measurement_evidence(tmp_path, collection) + + +def test_internal_decode_samples_remain_one_launch_estimate(tmp_path): + collection = _artifacts(tmp_path) + + def adjacent(payload): + fpm = payload["results"][0]["fpms"][0] + measurement = fpm["benchmark_measurement"] + rank = payload["dp"]["rank"] + measurement.update( + expected_internal_samples=4, + raw_fpms=[ + {"counter_id": fpm["counter_id"], "dp_rank": rank, "wall_time": fpm["wall_time"] * scale} + for scale in (4, 0.8, 1, 3) + ], + estimate={"method": "adjacent_upper_median", "raw_sample_indices": [1, 2, 3]}, + ) + + _edit(tmp_path, adjacent) + _edit(tmp_path, _add_execution_protocol) + evidence = extract_measurement_evidence(tmp_path, collection) + point = next(iter(evidence["points"].values())) + assert evidence["measurement_protocol"]["independent_repetitions"] == 1 + assert point["wall_time_seconds"] == 0.011 + assert len(point["ranks"]["0"]["raw_fpms"]) == 4 + assert point["ranks"]["0"]["raw_fpms"][-1]["wall_time"] == 0.03 + + +def test_duplicate_launch_cannot_count_as_an_independent_repetition(tmp_path): + left, _, key = _read_pair(tmp_path) + assert compare_measurements(left, copy.deepcopy(left), point_key=key)["status"] == "mismatch" + + +def test_extraction_requires_same_native_samples_and_run(tmp_path): + collection = _artifacts(tmp_path) + with pytest.raises(ValueError, match="run/grid identity"): + extract_measurement_evidence(tmp_path, replace(collection, runtime_run_id="different")) + + +@pytest.mark.parametrize("reasons", [None, []]) +def test_zero_kv_duplicates_normalize_optional_reasons(reasons): + from collector.fpm_forward.native_artifact import _zero_kv_prefill_sample + + point = {"point_type": "prefill", "total_kv_read_tokens": 0} + ordinary = NativePointMeasurement(point, ((0, 0.01),), "not_applicable") + duplicate = NativePointMeasurement({**point, "sample_reasons": reasons}, ((0, 0.02),), "real_prefix") + assert _zero_kv_prefill_sample([ordinary, duplicate]) is ordinary + point["sample_reasons"] = None + assert _zero_kv_prefill_sample([ordinary, duplicate]) is ordinary diff --git a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py index 110a1190e..e2d015903 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy +import hashlib import importlib.metadata import json import sys @@ -165,6 +166,316 @@ def test_unknown_architecture_plans_and_renders_from_real_config(tmp_path, no_mo assert "--max-num-seqs 1024" in script +def _multimodal_config(): + # Independently authored geometry; the wrapper is the checkpoint identity, + # while the nested model type establishes the supported decoder layout. + return { + "architectures": ["ExampleMultimodalForConditionalGeneration"], + "model_type": "example_multimodal", + "text_config": { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 1024, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", + }, + "vision_config": {"hidden_size": 512, "num_hidden_layers": 8}, + } + + +@pytest.mark.parametrize("requested_kv_dtype", ["auto", "bfloat16"]) +def test_explicit_none_kv_sidecar_preserves_unquantized_capability( + tmp_path, no_models_or_timing_data, requested_kv_dtype +): + # Original synthetic metadata exercises the sentinel found in this checkpoint: + # https://huggingface.co/thinkingmachines/Inkling-NVFP4/blob/42a75a99a40eb2ba1e0717db6357a0bf15205044/hf_quant_config.json + document = _multimodal_config() + del document["text_config"]["architectures"] + document["text_config"].update(n_routed_experts=8, num_experts_per_tok=2) + quantization = {"quantization": {"quant_algo": "NVFP4", "kv_cache_quant_algo": "none"}} + config_path = tmp_path / "config.json" + quant_path = tmp_path / "hf_quant_config.json" + config_path.write_text(json.dumps(document)) + quant_path.write_text(json.dumps(quantization)) + source_bytes = (config_path.read_bytes(), quant_path.read_bytes()) + frozen = load_model_config("example/multimodal-checkpoint", explicit_config_path=str(config_path)) + + capability = capabilities.resolve_model_capability( + backend="vllm", + model_path="example/multimodal-checkpoint", + model_architecture=document["architectures"][0], + selected_ops=set(), + has_model_cases=False, + system="gb300", + requested_weight_quantizations=("nvfp4",), + requested_kv_cache_dtypes=(requested_kv_dtype,), + model_config_path=str(config_path), + database_version="0.27.0", + checkpoint_native_dtypes=True, + ) + + assert capability.support_level == "bootstrap_template" + assert capability.is_moe is True + assert capability.dtype.gemm_quant_mode == "nvfp4" + assert capability.dtype.moe_quant_mode == "nvfp4" + assert capability.dtype.native_kv_cache_dtype == "bfloat16" + assert capability.dtype.kv_cache_dtypes == ("bfloat16",) + assert capability.dtype.fmha_quant_mode == "bfloat16" + assert capability.dtype.fmha_by_kv_dtype == {"bfloat16": "bfloat16"} + assert capability.dtype.fmha_resolution == "checkpoint_native" + assert capability.model_config.payload == frozen.payload + assert capability.model_config.sha256 == frozen.sha256 + assert capability.model_config.payload["hf_quant_config"] == quantization + assert (config_path.read_bytes(), quant_path.read_bytes()) == source_bytes + + profile = _profile() + profile.update(model="example/multimodal-checkpoint", architecture=document["architectures"][0], num_experts=8) + for deployment in profile["deployments"]: + deployment.update( + system="gb300", backend_version="0.27.0", kv_cache_dtype="bfloat16", fmha_quant_mode="bfloat16" + ) + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + weight_quantizations=("nvfp4",), + kv_cache_dtypes=(requested_kv_dtype,), + ) + plan = _plan( + profile, + system="gb300", + model_config_path=str(config_path), + selected_ops=set(), + has_model_cases=False, + options=options, + ) + serialized = json.loads(json.dumps(plan.to_dict())) + assert serialized["dtype_profile"] == capability.dtype.to_dict() + assert serialized["capability"]["model_config"] == frozen.to_dict() + assert len(serialized["cells"]) == 4 + assert {cell["kv_cache_dtype"] for cell in serialized["cells"]} == {"bfloat16"} + assert {cell["resolved_dtypes"]["fmha_quant_mode"] for cell in serialized["cells"]} == {"bfloat16"} + assert (config_path.read_bytes(), quant_path.read_bytes()) == source_bytes + + +def _onboard_collection_plan(tmp_path, document, overrides=None): + from aisimulate.support.config_profile import derive_profile + from aisimulate.support.config_profile import load_model_config as load_onboarding_config + from aisimulate.support.fpm import fpm_cli_args + from aisimulate.support.schema import SupportRequest + + path = tmp_path / "config.json" + path.write_text(json.dumps(document)) + request = SupportRequest.model_validate( + { + "identity": { + "model": "example/multimodal-checkpoint", + "model_revision": "collector-test-snapshot", + "model_kind": "dense", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "NVLink", + }, + "search": {"context_length": 4096}, + "workload": {"input_tokens": 128, "concurrency": 8, "request_count": 8}, + } + ) + draft = derive_profile( + load_onboarding_config(path), + request, + { + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "bfloat16", + **(overrides or {}), + }, + ) + assert draft.profile is not None, draft.missing + assert hashlib.sha256(path.read_bytes()).hexdigest() in draft.profile.provenance + request = SupportRequest.model_validate({**request.model_dump(), "fpm_profile": draft.profile}) + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + args = cli._parser().parse_args(command[3:]) + return planner.build_collection_plan( + backend="vllm", + model_path=request.identity.model, + model_architecture=args.model_architecture, + model_config_path=str(path), + system="h200_sxm", + selected_ops={"attention_context", "attention_generation"}, + options=FPMCollectionOptions.from_args(args), + fpm_profile=draft.profile, + ) + + +@pytest.mark.parametrize("declare_text_architecture", [False, True]) +@pytest.mark.parametrize( + "wrapper_fields", + [ + {}, + {"n_routed_experts": 32}, + {"kv_lora_rank": 64}, + {"q_lora_rank": 32, "num_experts": 16, "sliding_window": 64, "multi_query": True, "hidden_size": 1024}, + ], +) +def test_config_onboarding_architecture_matches_multimodal_collection( + tmp_path, no_models_or_timing_data, declare_text_architecture, wrapper_fields +): + document = {**_multimodal_config(), **wrapper_fields} + if not declare_text_architecture: + del document["text_config"]["architectures"] + plan = _onboard_collection_plan(tmp_path, document) + expected = "LlamaForCausalLM" if declare_text_architecture else "ExampleMultimodalForConditionalGeneration" + assert plan.capability.architecture == plan.fpm_profile.architecture == expected + assert plan.capability.is_moe is False + assert plan.capability.attention_kind == "dense_gqa" + assert len(plan.cells) == 2 + assert "Text decoder only" in plan.fpm_profile.provenance + assert {decision["source"] for decision in plan.to_dict()["topology_memory_admission"]} == {"fpm_profile_pending"} + assert all(decision.disposition == "unknown" for decision in plan.topology_memory_admission) + parsed = plan.capability.model_config.parsed_payload() + assert parsed["hidden_size"] == 128 + assert parsed["num_experts"] == 0 + assert all(key not in parsed["raw_config"] for key in wrapper_fields if key != "hidden_size") + cell_dir = tmp_path / "rendered" + cell_dir.mkdir() + runner._render_cell(plan, plan.cells[0], cell_dir, {}) + assert "--model example/multimodal-checkpoint" in (cell_dir / "run.sh").read_text() + + +@pytest.mark.parametrize("outer_alias", ["quant_algo", "hf_quant_config"]) +@pytest.mark.parametrize( + "algorithm,expected_gemm,expected_moe", [("fp8", "fp8_static", "fp8"), ("nvfp4", "nvfp4", "nvfp4")] +) +@pytest.mark.parametrize("shared_precision", [False, True]) +@pytest.mark.parametrize("dtype_key", ["dtype", "torch_dtype"]) +def test_config_onboarding_precision_matches_multimodal_collection( + tmp_path, no_models_or_timing_data, outer_alias, algorithm, expected_gemm, expected_moe, shared_precision, dtype_key +): + document = _multimodal_config() + text = document["text_config"] + del text["torch_dtype"] + selected = document if shared_precision else text + selected[dtype_key] = "bfloat16" + selected["quantization_config"] = {"quant_method": algorithm, "kv_cache_scheme": {"type": "float", "num_bits": 8}} + other_dtype_key = "torch_dtype" if dtype_key == "dtype" else "dtype" + outer_algorithm = algorithm if shared_precision else ("nvfp4" if algorithm == "fp8" else "fp8") + document[outer_alias] = ( + outer_algorithm if outer_alias == "quant_algo" else {"quantization": {"quant_algo": outer_algorithm.upper()}} + ) + if shared_precision: + # Null aliases, like missing aliases, do not declare decoder precision. + text[dtype_key] = None + text["quantization_config"] = None + else: + document[other_dtype_key] = "float32" + document.update(quant_dynamic=True, kv_cache_quant_algo="int8") + plan = _onboard_collection_plan( + tmp_path, + document, + {"fmha_quant_mode": "fp8", "kv_cache_dtype": "fp8", "weights_bytes": 1024, "activations_bytes": 1024}, + ) + assert len(plan.cells) == 2 + assert plan.dtype_profile.gemm_quant_mode == expected_gemm + assert plan.dtype_profile.moe_quant_mode == expected_moe + assert plan.dtype_profile.native_kv_cache_dtype == "fp8" + evidence = plan.capability.model_config + frozen = evidence.to_dict() + assert frozen["payload"] == document + assert ( + frozen["sha256"] + == hashlib.sha256( + json.dumps(document, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode() + ).hexdigest() + ) + raw = evidence.parsed_payload()["raw_config"] + assert raw["quant_algo"] == algorithm + assert raw["kv_cache_quant_algo"] == "fp8" + assert "quant_dynamic" not in raw + assert raw[dtype_key] == "bfloat16" + assert other_dtype_key not in raw + if not shared_precision: + assert "hf_quant_config" not in raw + raw["quantization_config"]["quant_method"] = "mutated" + detached = evidence.effective_payload + detached["quantization_config"]["quant_method"] = "mutated" + (tmp_path / "config.json").unlink() + cell_dir = tmp_path / "rendered" + cell_dir.mkdir() + runner._render_cell(plan, plan.cells[0], cell_dir, {}) + assert "--model example/multimodal-checkpoint" in (cell_dir / "run.sh").read_text() + assert evidence.parsed_payload()["raw_config"]["quant_algo"] == algorithm + assert evidence.to_dict() == frozen + + +def test_registered_multimodal_parser_keeps_context_out_of_decoder_metadata(tmp_path, no_models_or_timing_data): + document = _multimodal_config() + document.update( + architectures=["Llama4ForConditionalGeneration"], + n_routed_experts=32, + kv_lora_rank=64, + hf_quant_config={"quantization": {"quant_algo": "NVFP4"}}, + quant_dynamic=True, + vision_config={ + "hidden_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_channels": 3, + "intermediate_size": 64, + "image_size": 16, + "patch_size": 4, + "pixel_shuffle_ratio": 0.5, + "projector_input_dim": 32, + "projector_output_dim": 16, + "vision_output_dim": 16, + }, + image_processor_config={"max_patches": 1, "resize_to_max_canvas": False, "add_global_tile": False}, + ) + del document["text_config"]["architectures"] + document["text_config"]["quantization_config"] = {"quant_method": "fp8", "kv_cache_scheme": "FP8"} + plan = _onboard_collection_plan( + tmp_path, + document, + {"fmha_quant_mode": "fp8", "kv_cache_dtype": "fp8", "weights_bytes": 1024, "activations_bytes": 1024}, + ) + assert plan.capability.architecture == "Llama4ForConditionalGeneration" + assert plan.capability.is_moe is False + assert plan.capability.attention_kind == "dense_gqa" + assert plan.dtype_profile.gemm_quant_mode == "fp8_static" + evidence = plan.capability.model_config + frozen = evidence.to_dict() + parsed = evidence.parsed_payload() + assert parsed["hidden_size"] == 128 + assert parsed["num_experts"] == 0 + assert parsed["extra_params"].vision_config.hidden_size == 16 + assert parsed["extra_params"].vision_config.max_num_tiles == 1 + assert parsed["raw_config"]["quant_algo"] == "fp8" + assert all( + key not in parsed["raw_config"] + for key in ("n_routed_experts", "kv_lora_rank", "vision_config", "image_processor_config", "hf_quant_config") + ) + (tmp_path / "config.json").unlink() + cell_dir = tmp_path / "rendered" + cell_dir.mkdir() + runner._render_cell(plan, plan.cells[0], cell_dir, {}) + assert "--model example/multimodal-checkpoint" in (cell_dir / "run.sh").read_text() + assert evidence.to_dict() == frozen + + +def test_unknown_decoder_quantization_is_not_replaced_by_wrapper_precision(tmp_path, no_models_or_timing_data): + document = _multimodal_config() + document["quantization_config"] = {"quant_method": "fp8"} + document["text_config"]["quant_algo"] = "unknown_decoder_quantization" + with pytest.raises(ValueError, match="Unsupported quant algorithm: unknown_decoder_quantization"): + _onboard_collection_plan( + tmp_path, + document, + {"gemm_quant_mode": "fp8_static", "moe_quant_mode": "fp8", "weights_bytes": 1024}, + ) + + def test_frozen_profile_content_invalidates_resume_without_changing_cell_ids(tmp_path, no_models_or_timing_data): source = _profile() first = _plan(source) @@ -184,6 +495,63 @@ def test_frozen_profile_content_invalidates_resume_without_changing_cell_ids(tmp runner._load_checkpoint(checkpoint, second, resume=True) +@pytest.mark.parametrize( + "controls", + [ + ["--fpm-prefill-cudagraph-policy", "runtime"], + ["--fpm-gpu-memory-utilization", "0.9"], + ], +) +def test_launch_policy_change_invalidates_frozen_campaign(tmp_path, no_models_or_timing_data, controls): + profile = _profile() + first = _plan(profile) + options = FPMCollectionOptions.from_args(cli._parser().parse_args([*_argv(profile), *controls])) + second = _plan(profile, options=options) + assert first.sha256 != second.sha256 + assert [cell.cell_id for cell in first.cells] == [cell.cell_id for cell in second.cells] + checkpoint = tmp_path / "checkpoint.json" + checkpoint.write_text(json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": first.sha256, "cells": {}})) + with pytest.raises(ValueError, match="checkpoint does not match the current frozen plan"): + runner._load_checkpoint(checkpoint, second, resume=True) + + +def test_runtime_graph_plan_defers_capture_axis_and_grid_counts(no_models_or_timing_data): + profile = _profile() + options = FPMCollectionOptions.from_args( + cli._parser().parse_args([*_argv(profile), "--fpm-prefill-cudagraph-policy", "runtime"]) + ) + payload = _plan(profile, options=options).to_dict() + assert payload["point_generation"]["prefill_sampling"]["cudagraph_policy"] == "runtime" + assert payload["point_generation"]["planned_point_count"] is None + assert payload["counts"]["prefill_cudagraph_capture_sizes"] is None + assert payload["counts"]["prefill_new_token_axis_points"] is None + assert payload["counts"]["points"] == "runtime-determined" + + +def test_profile_admission_honors_explicit_total_memory_fraction(monkeypatch, no_models_or_timing_data): + profile = _profile() + profile["deployments"][1]["resources"]["weights_bytes"] = 80_000_000_000 + monkeypatch.setattr(memory_admission, "load_system_spec", lambda _system: {"gpu": {"mem_capacity": 200e9}}) + options = FPMCollectionOptions.from_args( + cli._parser().parse_args([*_argv(profile), "--fpm-gpu-memory-utilization", "0.5"]) + ) + plan = _plan(profile, options=options) + assert {cell.parallel_strategy for cell in plan.cells} == {"tep"} + assert [decision.disposition for decision in plan.topology_memory_admission] == ["rejected", "admitted"] + estimates = [decision.estimates[0].to_dict() for decision in plan.topology_memory_admission] + assert [estimate["gpu_capacity_bytes"] for estimate in estimates] == [200_000_000_000] * 2 + assert [estimate["gpu_memory_budget_bytes"] for estimate in estimates] == [100_000_000_000] * 2 + assert [estimate["headroom_bytes"] for estimate in estimates] == [-3_000_000_000, 17_000_000_000] + + legacy = _plan(profile) + assert {cell.parallel_strategy for cell in legacy.cells} == {"dep", "tep"} + assert all( + "gpu_memory_budget_bytes" not in estimate.to_dict() + for decision in legacy.topology_memory_admission + for estimate in decision.estimates + ) + + @pytest.mark.parametrize("field,value", [("fmha_quant_mode", "bfloat16"), ("moe_backend", "flashinfer_cutlass")]) def test_profile_cannot_override_resolved_precision_or_backend(field, value, no_models_or_timing_data): profile = _profile() @@ -247,7 +615,9 @@ def test_resource_envelope_errors_do_not_fail_open(field, value, no_models_or_ti profile = _profile() profile["deployments"][0]["resources"][field] = value options = replace( - FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), max_prefill_batch_size=4 + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_prefill_isl=8192, + max_prefill_batch_size=4, ) with pytest.raises(ValueError, match="resource envelope exceeded"): _plan(profile, options=options) @@ -272,6 +642,9 @@ def test_profile_bounds_native_scheduling_before_cases_are_queued(phase, smoke, assert arguments[arguments.index("--max-num-batched-tokens") + 1] == "8192" assert arguments.count("--max-num-seqs") == 1 assert arguments[arguments.index("--max-num-seqs") + 1] == "1024" + # A profile's declared context is now authoritative; legacy profile + # callers previously used -1 and could auto-fit beyond that contract. + assert arguments[arguments.index("--max-model-len") + 1] == "8192" def test_profile_preserves_explicit_prefill_envelope(no_models_or_timing_data): @@ -289,6 +662,289 @@ def test_profile_preserves_explicit_prefill_envelope(no_models_or_timing_data): assert arguments[arguments.index("--max-num-seqs") + 1] == expected_batch +@pytest.mark.parametrize("with_profile", [False, True]) +@pytest.mark.parametrize("narrow_prefill", [False, True]) +@pytest.mark.parametrize("smoke", [False, True]) +@pytest.mark.parametrize( + "max_tokens,max_sequences,prefill_tokens,prefill_sequences", [(4096, 64, 1024, 4), (256, 256, 128, 128)] +) +def test_cli_runtime_limits_reach_both_rendered_workers( + tmp_path, + monkeypatch, + with_profile, + narrow_prefill, + smoke, + max_tokens, + max_sequences, + prefill_tokens, + prefill_sequences, +): + """Exercise parsing, plan resolution and real Generator output without launching GPUs.""" + profile = _profile() + argv = [ + *_argv(profile), + "--fpm-max-model-len", + "4096", + "--fpm-max-num-batched-tokens", + str(max_tokens), + "--fpm-max-num-seqs", + str(max_sequences), + ] + if with_profile: + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + argv.extend(["--fpm-model-profile", str(path)]) + if narrow_prefill: + argv.extend( + ["--fpm-max-prefill-isl", str(prefill_tokens), "--fpm-max-prefill-batch-size", str(prefill_sequences)] + ) + if smoke: + argv.append("--smoke") + + def render_without_launch(_args, resolved): + plan, overrides = resolved + for phase in ("prefill", "decode"): + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + target = tmp_path / phase + target.mkdir() + runner._render_cell(plan, cell, target, overrides, smoke=_args.smoke) + return [] + + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + monkeypatch.setattr(cli, "run_resolved", render_without_launch) + assert cli.main(argv) == 0 + for phase in ("prefill", "decode"): + script = (tmp_path / phase / "run.sh").read_text() + tokens, sequences = ( + (prefill_tokens, prefill_sequences) + if narrow_prefill and phase == "prefill" + else (max_tokens, max_sequences) + ) + for option, value in ( + ("--max-model-len", 4096), + ("--max-num-batched-tokens", tokens), + ("--max-num-seqs", sequences), + ): + assert script.count(option) == 1 + assert f"{option} {value}" in script + + +@pytest.mark.parametrize("smoke", [False, True]) +@pytest.mark.parametrize("memory_fraction", [0.9, 1.0]) +def test_cli_runtime_graph_policy_and_memory_fraction_reach_native_launch( + tmp_path, monkeypatch, smoke, memory_fraction +): + profile = _profile() + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + argv = [ + *_argv(profile), + "--fpm-model-profile", + str(path), + "--fpm-prefill-cudagraph-policy", + "runtime", + "--fpm-gpu-memory-utilization", + str(memory_fraction), + ] + if smoke: + argv.append("--smoke") + + def render_without_launch(_args, resolved): + plan, overrides = resolved + for phase in ("prefill", "decode"): + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + target = tmp_path / phase + target.mkdir() + runner._render_cell(plan, cell, target, overrides, smoke=_args.smoke) + return [] + + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + monkeypatch.setattr(cli, "run_resolved", render_without_launch) + assert cli.main(argv) == 0 + for phase in ("prefill", "decode"): + script = (tmp_path / phase / "run.sh").read_text() + saved = json.loads((tmp_path / phase / "generator-request.json").read_text()) + assert saved["params"]["agg"]["kv_cache_free_gpu_memory_fraction"] == memory_fraction + assert script.count("--gpu-memory-utilization") == 1 + assert f"--gpu-memory-utilization {memory_fraction}" in script + assert "--compilation-config" not in script + assert "--cudagraph-capture-sizes" not in script + assert "--max-cudagraph-capture-size" not in script + assert "--no-enable-prefix-caching" not in script # DEP keeps real KV warm-up. + if phase == "prefill": + assert "--no-async-scheduling" in script + assert "--prefill-max-kv-read-token-samples" in script + if smoke: + assert "--prefill-max-new-token-samples 2" in script + else: + assert "--prefill-max-new-token-samples" not in script + else: + assert "--no-async-scheduling" not in script + + +@pytest.mark.parametrize("policy,enforce_eager", [("explicit", False), ("runtime", False), ("explicit", True)]) +def test_smoke_preserves_runtime_launch_and_bounds_sampling(tmp_path, no_models_or_timing_data, policy, enforce_eager): + profile = _profile() + if enforce_eager: + # Eager collection is currently qualified only for V4.1. + profile.update(model="deepseek-ai/DeepSeek-V4.1-Flash", architecture="DeepseekV41ForCausalLM") + for deployment in profile["deployments"]: + deployment.update(gemm_quant_mode="fp8_block", moe_quant_mode="w4a8_mxfp4_mxfp8") + argv = [ + *_argv(profile), + "--fpm-max-num-batched-tokens", + "4096", + "--fpm-max-num-seqs", + "64", + "--fpm-prefill-cudagraph-policy", + policy, + ] + if policy == "explicit": + argv.extend(["--fpm-max-prefill-cudagraph-size", "512"]) + if enforce_eager: + argv.append("--fpm-enforce-eager") + options = FPMCollectionOptions.from_args(cli._parser().parse_args(argv)) + plan = _plan(profile, options=options) + for phase in ("prefill", "decode"): + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + capture_configs = {} + for smoke in (False, True): + target = tmp_path / f"{phase}-{'smoke' if smoke else 'full'}" + target.mkdir() + runner._render_cell(plan, cell, target, {}, smoke=smoke) + saved = json.loads((target / "generator-request.json").read_text()) + arguments = saved["params"]["agg"]["extra_cli_args"] + script = (target / "run.sh").read_text() + assert arguments.count("--enforce-eager") == int(enforce_eager) + assert ("--enforce-eager" in script) == enforce_eager + for option, value in ( + ("--benchmark-mode", phase), + ("--max-model-len", "8192"), + ("--max-num-batched-tokens", "4096"), + ("--max-num-seqs", "64"), + ): + assert arguments.count(option) == 1 + assert arguments[arguments.index(option) + 1] == value + assert script.count(option) == 1 + assert f"{option} {value}" in script + if phase == "prefill" and policy == "explicit" and not enforce_eager: + assert arguments.count("--compilation-config") == 1 + compilation = arguments[arguments.index("--compilation-config") + 1] + assert f"--compilation-config '{compilation}'" in script + capture_configs[smoke] = json.loads(compilation) + assert capture_configs[smoke]["max_cudagraph_capture_size"] == 512 + assert capture_configs[smoke]["cudagraph_capture_sizes"][-1] == 512 + else: + assert "--compilation-config" not in arguments + assert "--compilation-config" not in script + if smoke: + samples = ( + { + "--prefill-max-new-token-samples": "2", + "--prefill-max-kv-read-token-samples": "2", + "--prefix-max-batch-size-samples": "1", + } + if phase == "prefill" + else {"--decode-max-kv-read-token-samples": "2", "--decode-max-batch-size-samples": "2"} + ) + for option, value in samples.items(): + assert arguments.count(option) == 1 + assert arguments[arguments.index(option) + 1] == value + assert f"{option} {value}" in script + if capture_configs: + assert capture_configs[True] == capture_configs[False] + + +@pytest.mark.parametrize("smoke", [False, True]) +def test_omitted_prefill_defaults_stay_within_selected_profile_bounds(no_models_or_timing_data, smoke): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(max_num_tokens=4096, max_batch_size=32) + plan = _plan(profile) + assert plan.options.prefill_sampling.max_total_prefill_tokens == 4096 + for cell in plan.cells: + arguments = runner._cell_generator_overrides(plan, cell, {}, smoke=smoke)["params"]["agg"]["extra_cli_args"] + for option, value in ( + ("--max-model-len", "8192"), + ("--max-num-batched-tokens", "4096"), + ("--max-num-seqs", "32"), + ): + assert arguments.count(option) == 1 + assert arguments[arguments.index(option) + 1] == value + + +@pytest.mark.parametrize( + "with_profile,resource_limits,limits", + [ + (False, {}, ["--fpm-max-num-batched-tokens", "128", "--fpm-max-num-seqs", "256"]), + (True, {}, ["--fpm-max-num-batched-tokens", "128", "--fpm-max-num-seqs", "256"]), + ( + False, + {}, + ["--fpm-max-num-batched-tokens", "4096", "--fpm-max-num-seqs", "256", "--fpm-max-prefill-isl", "128"], + ), + (False, {}, ["--fpm-max-num-seqs", "16384"]), + (True, {}, ["--fpm-max-num-batched-tokens", "128"]), + (True, {}, ["--fpm-max-prefill-isl", "128"]), + (True, {"max_num_tokens": 128, "max_batch_size": 256}, []), + ( + True, + {"max_num_tokens": 128, "max_batch_size": 256}, + ["--fpm-max-prefill-batch-size", "8"], + ), + ], +) +@pytest.mark.parametrize("smoke", [False, True]) +def test_cli_rejects_scheduler_tokens_below_sequences_before_execution( + tmp_path, monkeypatch, capsys, with_profile, resource_limits, limits, smoke +): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(resource_limits) + argv = [*_argv(profile), *limits] + if with_profile: + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + argv.extend(["--fpm-model-profile", str(path)]) + if smoke: + argv.append("--smoke") + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + monkeypatch.setattr(cli, "run_resolved", lambda *_args: pytest.fail("collection execution started")) + + with pytest.raises(SystemExit) as error: + cli.main(argv) + + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_batched_tokens" in message + assert "max_num_seqs" in message + + +@pytest.mark.parametrize( + "limits, message", + [ + (["--fpm-max-model-len", "8193"], "context"), + (["--fpm-max-num-batched-tokens", "8192", "--fpm-max-prefill-isl", "1024"], "resource envelope exceeded"), + (["--fpm-max-num-seqs", "128", "--fpm-max-prefill-batch-size", "4"], "resource envelope exceeded"), + (["--fpm-max-prefill-isl", "8192"], "resource envelope exceeded"), + (["--fpm-max-prefill-batch-size", "128"], "resource envelope exceeded"), + (["--fpm-max-decode-batch-size", "128"], "resource envelope exceeded"), + ], +) +def test_cli_rejects_profile_limit_overshoots_before_execution(tmp_path, monkeypatch, capsys, limits, message): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(max_num_tokens=4096, max_batch_size=64) + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + monkeypatch.setattr(cli, "run_resolved", lambda *_args: pytest.fail("collection execution started")) + with pytest.raises(SystemExit) as error: + cli.main([*_argv(profile), "--fpm-model-profile", str(path), *limits]) + assert error.value.code == 2 + assert message in capsys.readouterr().err + + @pytest.mark.parametrize( ("expected_version", "actual_version"), [ @@ -368,3 +1024,25 @@ def test_decode_collection_respects_profile_scheduler_envelope(no_models_or_timi _plan(profile, options=options) else: assert _plan(profile, options=options).cells + + +def test_aggregated_profile_defaults_to_shared_token_limit(no_models_or_timing_data): + profile = _profile() + profile["context_length"] = 32768 + for deployment in profile["deployments"]: + deployment["resources"]["max_num_tokens"] = 16384 + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), worker_type="aggregated" + ) + plan = _plan(profile, options=options) + assert plan.options.max_prefill_isl == 16384 + + +def test_aggregated_profile_requires_shared_limit_for_different_bounds(no_models_or_timing_data): + profile = _profile() + profile["deployments"][0]["resources"]["max_num_tokens"] = 16384 + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), worker_type="aggregated" + ) + with pytest.raises(ValueError, match="explicit shared --fpm-max-num-batched-tokens"): + _plan(profile, options=options) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_repeatability.py b/python/aisimulate/tests/unit/collector/test_fpm_repeatability.py new file mode 100644 index 000000000..d2c5f4463 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_repeatability.py @@ -0,0 +1,1534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import json +import statistics +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +from collector.fpm_forward import cli, execution_evidence, repeatability, runner +from collector.fpm_forward.config import with_kv_warmup_defaults +from collector.fpm_forward.database import aggregate_cell +from collector.fpm_forward.native_artifact import ( + _expected_scheduled, + validate_native_collection, +) +from collector.fpm_forward.runtime import fpm_memory_observer as observer +from collector.fpm_forward.runtime_memory import validate_saved_plan + +from .test_fpm_measurement_evidence import _add_measurement_protocol, _remove_execution_protocol +from .test_fpm_profile_collection import _plan, _profile, no_models_or_timing_data # noqa: F401 +from .test_fpm_runner import _native_payload, _write_provenance +from .test_fpm_runtime_memory import _vllm_config + +pytestmark = pytest.mark.unit + + +def _point(coordinates, phase, index): + point = {"point_type": phase, "benchmark_id": index, "total_prefill_tokens": 0, **coordinates} + axis = point["total_prefill_tokens"] or point["batch_size"] + capture = 1 << (axis - 1).bit_length() if axis <= 16 else None + point.update( + expected_cudagraph_mode=("FULL" if phase == "decode" else "PIECEWISE") if capture else "NONE", + expected_capture_size=capture, + padding_tokens=capture - axis if capture else None, + sample_reasons=["kvwarm_real_kv"] + if phase == "decode" + else ["prefill_real_seed"] + if point["total_kv_read_tokens"] + else [], + ) + return point + + +def _write_campaign( + root, + checkpoint_path, + plan, + *, + attempt_id="source", + factor=1.0, + observed=True, + generator_overrides=None, + native_points=None, + execution_evidence=True, +): + deployment = generator_overrides or {} + root.mkdir(parents=True, exist_ok=True) + (root / "collection-plan.json").write_text(json.dumps(plan.to_dict())) + (root / "generator-overrides.json").write_text(json.dumps(with_kv_warmup_defaults(deployment))) + checkpoint = {"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": {}} + for cell in plan.cells: + directory = root / "cells" / cell.cell_id + raw = directory / "raw" / "pod" + raw.mkdir(parents=True, exist_ok=True) + runner._render_cell(plan, cell, directory, deployment) + generated_model = json.loads((directory / "generator-request.json").read_text())["ServiceConfig"]["model_path"] + _write_provenance( + raw / "collector-provenance.json", cell_id=cell.cell_id, plan_sha256=plan.sha256, attempt_id=attempt_id + ) + provenance = json.loads((raw / "collector-provenance.json").read_text()) + provenance["runtime"]["backend_version"] = plan.capability.aic_database_version + (raw / "collector-provenance.json").write_text(json.dumps(provenance)) + coords = ( + native_points[cell.workload_kind] + if native_points is not None + else json.loads(plan.options.benchmark_points_json)[cell.workload_kind] + if plan.options.benchmark_points_json + else [ + { + "batch_size": batch, + "total_prefill_tokens": batch if cell.workload_kind == "prefill" else 0, + "total_kv_read_tokens": batch * context, + } + for batch in (1, 2, 4, 8, 16, 32) + for context in (16, 1024) + ] + ) + points = [_point(item, cell.workload_kind, index + 1) for index, item in enumerate(coords)] + for rank in range(cell.topology.dp): + payload = _native_payload( + phase=cell.workload_kind, rank=rank, dp=cell.topology.dp, run_id=f"{attempt_id}-{cell.cell_id}" + ) + groups = [] + results = [] + for point in points: + rank_results = [ + { + "dp_rank": dp, + "fpms": [ + { + "counter_id": point["benchmark_id"], + "dp_rank": dp, + "wall_time": factor * (0.01 + dp * 0.001), + "scheduled_requests": _expected_scheduled(point), + } + ], + } + for dp in range(cell.topology.dp) + ] + groups.append( + { + "benchmark_id": point["benchmark_id"], + "point": point, + "expected_dp_ranks": list(range(cell.topology.dp)), + "complete": True, + "rank_results": rank_results, + "wall_time": max(row["fpms"][0]["wall_time"] for row in rank_results), + } + ) + results.append( + { + "point": point, + "kv_seed_regime": "real_kv" + if cell.workload_kind == "decode" + else "real_prefix" + if point["total_kv_read_tokens"] + else "not_applicable", + "fpms": rank_results[rank]["fpms"], + } + ) + payload.update( + results=results, + iteration_groups=groups, + coverage={"expected_points": len(points), "completed_points": len(points), "skipped_points": 0}, + cudagraph={"mode": "FULL_AND_PIECEWISE", "capture_sizes": [1, 2, 4, 8, 16]}, + timing={ + "benchmark_elapsed_seconds": 100.0, + "measured_iteration_seconds": sum(group["wall_time"] for group in groups), + }, + ) + payload.update( + limits={"block_size": 16, "max_model_len": plan.options.vllm_max_model_len}, + recurrent_state={"initialization": "unchanged", "policy": None, "uniform_bound": None}, + grid_digest=repeatability._canonical_hash(points), + ) + _add_measurement_protocol(payload, execution=execution_evidence) + (raw / f"benchmark-dp{rank}.json").write_text(json.dumps(payload)) + if observed: + for tp in range(cell.topology.tp): + execution = { + "schema_name": "aisimulate_fpm_runtime_execution", + "schema_version": 1, + "collector_provenance": provenance, + "backend_version": plan.capability.aic_database_version, + "dp_rank": rank, + "tp_rank": tp, + "pp_rank": 0, + "status": "observed", + "attention_groups": [ + { + "kv_cache_group_id": 0, + "backend_class": "example.ActualBackend", + "layer_names": ["layer.0"], + } + ], + "graph_config": { + "mode": "VLLM_COMPILE", + "cudagraph_mode": "FULL_AND_PIECEWISE", + "cudagraph_capture_sizes": [1, 2, 4, 8, 16], + "max_cudagraph_capture_size": 16, + }, + "resolved_config": _runtime_config(plan, cell), + } + execution["resolved_config"]["model_config"]["model"] = generated_model + (raw / f"fpm-execution-worker-dp{rank}-tp{tp}-pp0.json").write_text(json.dumps(execution)) + (raw / runner.POINTS_RECEIPT_FILENAME).write_text( + json.dumps( + { + "plan_sha256": plan.sha256, + "cell_id": cell.cell_id, + "attempt_id": attempt_id, + "sha256": plan.options.benchmark_points_sha256, + "phase": "after", + } + ) + ) + checkpoint["cells"][cell.cell_id] = {"status": "passed", "attempt_id": attempt_id} + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + checkpoint_path.write_text(json.dumps(checkpoint)) + + +def _runtime_config(plan, cell): + """Synthetic initialized vLLM state for the real generated test launch.""" + config = observer.execution_config(_vllm_config(tp=cell.topology.tp, dp=cell.topology.dp)) + config["model_config"].update( + model=plan.model_path, + revision=plan.fpm_profile.model_revision, + max_model_len=plan.options.vllm_max_model_len, + enforce_eager=plan.options.enforce_eager, + ) + if cell.gemm_quant_mode in {"fp8", "fp8_static", "fp8_block"}: + config["model_config"]["quantization"] = "fp8" + config["quantization_config"] = { + "type": "vllm.model_executor.layers.quantization.fp8.Fp8Config", + "activation_scheme": "static" if cell.gemm_quant_mode == "fp8_static" else "dynamic", + "weight_block_size": [128, 128] if cell.gemm_quant_mode == "fp8_block" else None, + } + elif cell.gemm_quant_mode in {"bfloat16", "half"}: + config["model_config"].update(dtype=cell.gemm_quant_mode, quantization=None) + config["quantization_config"] = None + config["cache_config"].update( + cache_dtype=cell.kv_cache_dtype, + gpu_memory_utilization=plan.options.gpu_memory_utilization or 0.9, + num_gpu_blocks_override=None, + enable_prefix_caching=cell.workload_kind == "decode" and runner._decode_prefix_caching_mode(cell) == "enabled", + ) + resources = plan.deployment_profile(cell).resources + config["scheduler_config"].update( + max_num_batched_tokens=plan.options.max_num_batched_tokens or resources.max_num_tokens, + max_num_seqs=plan.options.max_num_seqs or resources.max_batch_size, + ) + if cell.workload_kind == "prefill": + config["scheduler_config"]["max_num_batched_tokens"] = plan.options.prefill_sampling.max_total_prefill_tokens + config["scheduler_config"]["max_num_seqs"] = ( + plan.options.prefill_sampling.max_batch_size or config["scheduler_config"]["max_num_seqs"] + ) + config["parallel_config"].update( + enable_expert_parallel=cell.topology.moe_ep > 1, + enable_eplb=False, + data_parallel_size_local=cell.topology.dp, + data_parallel_external_lb=False, + decode_context_parallel_size=cell.topology.cp, + ) + config["compilation_config"].update(cudagraph_capture_sizes=[1, 2, 4, 8, 16], max_cudagraph_capture_size=16) + return config + + +@pytest.fixture +def campaign(tmp_path, no_models_or_timing_data): # noqa: F811 + profile = _profile() + for deployment in profile["deployments"]: + deployment["backend_version"] = "0.28.0" + plan = _plan(profile) + plan = replace( + plan, options=replace(plan.options, prefill_cudagraph_policy="runtime", max_prefill_cudagraph_size=None) + ) + coordinates = [ + {"batch_size": batch, "total_prefill_tokens": batch, "total_kv_read_tokens": batch * context} + for batch in (1, 2, 4, 8, 16, 32) + for context in (16, 1024) + ] + plan = repeatability._subset_plan( + plan, + { + "cell_id": plan.cells[0].cell_id, + "benchmark_points": {"schema_version": 3, "prefill": coordinates, "decode": []}, + }, + ) + source = tmp_path / "source" + checkpoint = tmp_path / "source-checkpoint.json" + _write_campaign(source, checkpoint, plan) + return plan, source, checkpoint + + +def _args(campaign, tmp_path, **kwargs): + plan, source, checkpoint = campaign + return { + "source_plan": plan, + "generator_overrides": {}, + "source_campaign_dir": source, + "source_checkpoint_path": checkpoint, + "output_dir": tmp_path / "validation", + **kwargs, + } + + +def _fake_collector(monkeypatch, *, factors=None, fail_at=None, duplicate=False, mutate=None): + calls = [] + + def collect(plan, **kwargs): + calls.append((plan, kwargs)) + assert kwargs["publish_database"] is False + assert kwargs["resume"] is False and kwargs["retry_failed"] is False + index = len(calls) + root = Path(kwargs["artifact_root"]) / plan.sha256[:16] + cp = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + _write_campaign( + root, + cp, + plan, + attempt_id="same" if duplicate else f"independent-{index}", + factor=(factors or [1.0])[((index - 1) % len(factors or [1.0]))], + generator_overrides=kwargs["generator_overrides"], + ) + if mutate: + mutate(root, index) + if index == fail_at: + cp.unlink() # No complete observation exists to retry or aggregate. + return [{"classification": "campaign_cell_failed", "error_message": "runtime failed"}] + return [] + + monkeypatch.setattr(repeatability, "run_collection", collect) + monkeypatch.setattr(repeatability, "_cell_runner", lambda *_args: SimpleNamespace(cleanup=lambda: None)) + return calls + + +def test_freeze_selects_validated_native_extremes_and_reports_unselected_capture_boundaries(campaign): + plan, source, checkpoint = campaign + frozen = repeatability.freeze_repeatability_plan( + plan, source, checkpoint, max_points_per_cell=6, comparison_mode="bounded" + ) + selected = frozen["cells"][0] + assert len(selected["points"]) == 6 < selected["source_point_count"] + assert {point["coordinates"]["batch_size"] for point in selected["points"]} >= {1, 32} + assert { + point["coordinates"]["total_kv_read_tokens"] / point["coordinates"]["batch_size"] + for point in selected["points"] + } == {16, 1024} + assert selected["selection_coverage"]["unselected_capture_boundaries"] == [2, 4, 8] + assert selected["execution"]["status"] == "qualified" + assert selected["execution"]["per_point_dispatch"] == "observed" + validate_saved_plan(repeatability._subset_plan(plan, selected).to_dict()) + assert frozen == repeatability.freeze_repeatability_plan( + plan, source, checkpoint, max_points_per_cell=6, comparison_mode="bounded" + ) + + +def test_too_small_budget_cannot_silently_drop_required_regimes(campaign): + with pytest.raises(ValueError, match="cannot cover.*uncovered"): + repeatability.freeze_repeatability_plan(*campaign, max_points_per_cell=1, comparison_mode="bounded") + + +def _add_zero_kv_duplicates(root, *, seed_samples=3): + for path in root.glob("cells/*/raw/*/benchmark-*.json"): + payload = json.loads(path.read_text()) + ordinary_row = payload["results"][0] + ordinary_group = payload["iteration_groups"][0] + assert ordinary_row["point"]["total_kv_read_tokens"] == 0 + assert ordinary_row["kv_seed_regime"] == "not_applicable" + for index in range(seed_samples): + group = copy.deepcopy(ordinary_group) + point = group["point"] + point["benchmark_id"] = index + 2 + point["sample_reasons"].append("prefill_real_seed") + group["benchmark_id"] = point["benchmark_id"] + for result in group["rank_results"]: + fpm = result["fpms"][0] + fpm["counter_id"] = point["benchmark_id"] + fpm["wall_time"] *= 0.5 + group["wall_time"] *= 0.5 + payload["iteration_groups"].append(group) + payload["results"].append( + { + "point": point, + "kv_seed_regime": "real_prefix", + "fpms": group["rank_results"][payload["dp"]["rank"]]["fpms"], + } + ) + count = len(payload["results"]) + payload["coverage"].update(expected_points=count, completed_points=count) + payload["timing"]["measured_iteration_seconds"] = sum( + group["wall_time"] for group in payload["iteration_groups"] + ) + payload["grid_digest"] = repeatability._canonical_hash([row["point"] for row in payload["results"]]) + _add_measurement_protocol(payload) + path.write_text(json.dumps(payload)) + + +@pytest.fixture +def zero_kv_campaign(campaign): + plan, source, checkpoint = campaign + plan = repeatability._subset_plan( + plan, + { + "cell_id": plan.cells[0].cell_id, + "benchmark_points": { + "schema_version": 3, + "prefill": [{"batch_size": 1, "total_prefill_tokens": 1, "total_kv_read_tokens": 0}], + "decode": [], + }, + }, + ) + _write_campaign(source, checkpoint, plan) + return plan, source, checkpoint + + +@pytest.mark.parametrize("seed_samples", [3, 4]) +def test_freeze_zero_kv_duplicates_uses_published_ordinary_sample(zero_kv_campaign, seed_samples): + plan, source, checkpoint = zero_kv_campaign + _add_zero_kv_duplicates(source, seed_samples=seed_samples) + original = runner._file_manifest(source) + frozen = repeatability.freeze_repeatability_plan(plan, source, checkpoint) + cell = frozen["cells"][0] + assert cell["source_point_count"] == seed_samples + 1 + assert len(cell["points"]) == 1 + point = cell["points"][0] + assert point["source_point"]["benchmark_id"] == 1 + assert point["source_point"]["sample_reasons"] == [] + assert point["kv_seed_regime"] == "not_applicable" + rows = aggregate_cell(plan, plan.cells[0], source / "cells" / cell["cell_id"], expected_attempt_id="source") + assert len(rows) == 1 + assert point["source_wall_time_seconds"] * 1000 == rows[0]["latency_ms"] + assert point["source_wall_time_seconds"] == max(value for _, value in point["source_rank_wall_times"]) + assert runner._file_manifest(source) == original + + +@pytest.mark.parametrize( + "conflict", + [ + "duplicate_ordinary", + "no_ordinary", + "missing_provenance", + "positive_kv", + "expected_cudagraph_mode", + "expected_capture_size", + "padding_tokens", + "partition", + "rows", + "sample_reasons", + "context_clamped", + ], +) +def test_freeze_rejects_conflicting_zero_kv_duplicates(zero_kv_campaign, conflict): + plan, source, checkpoint = zero_kv_campaign + _add_zero_kv_duplicates(source) + for path in source.glob("cells/*/raw/*/benchmark-*.json"): + payload = json.loads(path.read_text()) + for index, (row, group) in enumerate(zip(payload["results"], payload["iteration_groups"], strict=True)): + point = row["point"] + if conflict == "positive_kv": + point["total_kv_read_tokens"] = 128 + for result in group["rank_results"]: + result["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = 128 + row["fpms"] = group["rank_results"][payload["dp"]["rank"]]["fpms"] + elif conflict == "no_ordinary": + row["kv_seed_regime"] = "real_prefix" + point["sample_reasons"] = ["prefill_real_seed"] + elif index == 1: + if conflict == "duplicate_ordinary": + row["kv_seed_regime"] = "not_applicable" + point["sample_reasons"] = [] + elif conflict == "missing_provenance": + row.pop("kv_seed_regime") + elif conflict == "context_clamped": + point["sample_reasons"].append("context_clamped") + else: + point[conflict] = { + "expected_cudagraph_mode": "NONE", + "expected_capture_size": 2, + "padding_tokens": 1, + "partition": "different", + "rows": [{"tokens": 1}], + "sample_reasons": ["eager_tail", "prefill_real_seed"], + }[conflict] + group["point"] = point + path.write_text(json.dumps(payload)) + original = runner._file_manifest(source) + with pytest.raises(ValueError, match="unclamped samples share one key"): + repeatability.freeze_repeatability_plan(plan, source, checkpoint) + assert runner._file_manifest(source) == original + + +def test_zero_kv_source_consolidation_preserves_independent_fresh_samples(zero_kv_campaign, tmp_path, monkeypatch): + _add_zero_kv_duplicates(zero_kv_campaign[1]) + original = runner._file_manifest(zero_kv_campaign[1]) + calls = _fake_collector( + monkeypatch, factors=[0.99, 1.0, 1.01, 1.0, 1.0], mutate=lambda root, _index: _add_zero_kv_duplicates(root) + ) + args = _args(zero_kv_campaign, tmp_path) + report = repeatability.run_repeatability(**args) + assert report["status"] == "passed" and len(calls) == 5 + assert len(report["points"]) == 1 + point = report["points"][0] + assert point["sample_count"] == 5 + assert point["source_inclusive_sample_count"] == 6 + assert point["samples_seconds"] == pytest.approx( + [point["source_wall_time_seconds"] * factor for factor in (0.99, 1.0, 1.01, 1.0, 1.0)] + ) + assert runner._file_manifest(zero_kv_campaign[1]) == original + assert repeatability.run_repeatability(**args, resume=True) == report + + +def test_zero_kv_duplicates_in_fresh_repeat_are_not_counted_as_independent_samples( + zero_kv_campaign, tmp_path, monkeypatch +): + _add_zero_kv_duplicates(zero_kv_campaign[1]) + calls = _fake_collector(monkeypatch, mutate=lambda root, _index: _add_zero_kv_duplicates(root)) + report = repeatability.run_repeatability(**_args(zero_kv_campaign, tmp_path)) + assert report["status"] == "passed" and len(calls) == 5 + assert len(report["points"]) == 1 + assert report["points"][0]["sample_count"] == 5 + + +def test_five_independent_samples_keep_raw_data_and_pass_without_formal_publication(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, factors=[0.99, 1.0, 1.01, 1.0, 1.0]) + originals = runner._file_manifest(campaign[1]) + args = _args(campaign, tmp_path, max_points_per_cell=6) + report = repeatability.run_repeatability(**args) + assert report["status"] == "passed" and len(calls) == 5 + assert all(plan.options == campaign[0].options for plan, _ in calls) + assert len(report["points"]) == 12 + assert len({kwargs["artifact_root"] for _, kwargs in calls}) == 5 + assert runner._file_manifest(campaign[1]) == originals + assert all(point["sample_count"] == 5 and len(point["samples_seconds"]) == 5 for point in report["points"]) + values = report["points"][0]["samples_seconds"] + assert report["points"][0]["sample_cv"] == statistics.stdev(values) / statistics.mean(values) + assert report["points"][0]["sample_stddev_seconds"] == statistics.stdev(values) + assert report["repeatability"]["standard_deviation_denominator"] == "n - 1" + assert report["repeatability"]["time_unit"] == "seconds" + assert repeatability.run_repeatability(**args, resume=True) == report + assert len(calls) == 5 + + +def _mutate_workers(root, action): + for path in root.glob("cells/*/raw/*/fpm-execution-worker-*.json"): + payload = json.loads(path.read_text()) + action(payload) + path.write_text(json.dumps(payload)) + + +@pytest.mark.parametrize( + "section,field,value", + [ + ("model_config", "model", "other/checkpoint"), + ("model_config", "revision", "different-revision"), + ("model_config", "dtype", "torch.float32"), + ("model_config", "quantization", None), + ("model_config", "max_model_len", 4096), + ("model_config", "enforce_eager", True), + ("cache_config", "cache_dtype", "half"), + ("parallel_config", "tensor_parallel_size", 999), + ("scheduler_config", "max_num_batched_tokens", 4096), + ], +) +def test_identically_wrong_source_and_repeats_fail(campaign, tmp_path, monkeypatch, section, field, value): + def alter(payload): + payload["resolved_config"][section][field] = value + + _mutate_workers(campaign[1], alter) + calls = _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_workers(root, alter)) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "failed" + assert report["execution"]["status"] == "failed" + assert calls == [] + + +@pytest.mark.parametrize("missing", ["all", "model", "revision", "quantization", "scheduler", "graph"]) +def test_required_runtime_configuration_missing_is_incomplete(campaign, tmp_path, monkeypatch, missing): + def alter(payload): + config = payload["resolved_config"] + if missing == "all": + payload["resolved_config"] = None + elif missing == "scheduler": + del config["scheduler_config"] + elif missing == "graph": + del payload["graph_config"]["mode"] + else: + del config["model_config"][missing] + + _mutate_workers(campaign[1], alter) + _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_workers(root, alter)) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "incomplete" + + +@pytest.mark.parametrize("change", ["duplicate", "mode", "sizes", "maximum", "compile", "eager"]) +def test_invalid_initialized_graph_fails(campaign, tmp_path, monkeypatch, change): + def alter(payload): + graph = payload["graph_config"] + if change == "duplicate": + graph.update(cudagraph_mode="NONE", cudagraph_capture_sizes=[], max_cudagraph_capture_size=0) + return + if change == "mode": + graph["cudagraph_mode"] = "MADE_UP" + elif change == "sizes": + graph["cudagraph_capture_sizes"] = [1, 2, 2, 4, 8, 16] + elif change == "maximum": + graph["max_cudagraph_capture_size"] = 17 + elif change == "compile": + graph["mode"] = False + else: + payload["resolved_config"]["model_config"]["enforce_eager"] = True + payload["resolved_config"]["compilation_config"] = dict(graph) + + _mutate_workers(campaign[1], alter) + _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_workers(root, alter)) + assert repeatability.run_repeatability(**_args(campaign, tmp_path))["status"] == "failed" + + +def test_worker_graph_downgrade_preserves_scheduler_expectations(campaign, tmp_path, monkeypatch): + def alter(payload): + payload["graph_config"]["cudagraph_mode"] = "NONE" + payload["resolved_config"]["compilation_config"]["cudagraph_mode"] = "NONE" + + _mutate_workers(campaign[1], alter) + _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_workers(root, alter)) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "passed" + + +@pytest.mark.parametrize("evidence", ["correct", "missing", "wrong"]) +def test_explicit_local_checkpoint_uses_loaded_revision_evidence(campaign, tmp_path, monkeypatch, evidence): + plan, source, checkpoint = campaign + profile = json.loads(plan._fpm_profile_json) + profile["model_revision"] = "a" * 40 + deployment = {"K8sConfig": {"k8s_pvc_mount_path": "/models", "k8s_model_path_in_pvc": "checkpoint"}} + plan = replace( + plan, + _fpm_profile_json=json.dumps(profile), + generator_config_sha256=repeatability._canonical_hash(with_kv_warmup_defaults(deployment)), + ) + plan = repeatability._subset_plan( + plan, {"cell_id": plan.cells[0].cell_id, "benchmark_points": json.loads(plan.options.benchmark_points_json)} + ) + _write_campaign(source, checkpoint, plan, generator_overrides=deployment) + + def alter(payload): + model = payload["resolved_config"]["model_config"] + assert model["model"] == "/models/checkpoint" + model["revision"] = None + if evidence != "missing": + model["loaded_config_commit_hash"] = ("a" if evidence == "correct" else "b") * 40 + + _mutate_workers(source, alter) + calls = _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_workers(root, alter)) + result = repeatability.run_repeatability( + **_args((plan, source, checkpoint), tmp_path, generator_overrides=deployment) + ) + assert result["status"] == {"correct": "passed", "missing": "incomplete", "wrong": "failed"}[evidence] + if evidence == "wrong": + assert calls == [] + + +@pytest.mark.parametrize( + ("observed", "qualified"), + [ + ("/root/.cache/huggingface/hub/{repo}/snapshots/{revision}", True), + ("/custom/cache/{repo}/snapshots/{revision}", True), + ("/custom/cache/{repo}/snapshots/{revision}/config.json", False), + ("/custom/cache/models--other--GLM-5.2-NVFP4/snapshots/{revision}", False), + ("/custom/cache/models--nvidia--other/snapshots/{revision}", False), + ("/custom/cache/models--nvidia--GLM-5.2-NVFP4-extra/snapshots/{revision}", False), + ("/custom/cache/datasets--nvidia--GLM-5.2-NVFP4/snapshots/{revision}", False), + ("/custom/cache/{repo}/snapshot/{revision}", False), + ("/custom/cache/{repo}/refs/{revision}", False), + ("/custom/cache/{repo}/snapshots/main", False), + ("/custom/cache/{repo}/snapshots/{other_revision}", False), + ("/arbitrary/local-checkpoint/{revision}", False), + ("/arbitrary/local-checkpoint", False), + ("custom/cache/{repo}/snapshots/{revision}", False), + ("//custom/cache/{repo}/snapshots/{revision}", False), + ("/custom//cache/{repo}/snapshots/{revision}", False), + ("/custom/./cache/{repo}/snapshots/{revision}", False), + ("/custom/../cache/{repo}/snapshots/{revision}", False), + ("/custom/cache/{repo}/snapshots/{revision}/", False), + ("/custom/cache/{repo}/snapshots/{revision}/..", False), + ("/custom\\cache/{repo}/snapshots/{revision}", False), + ("/custom\x00cache/{repo}/snapshots/{revision}", False), + ], +) +def test_pinned_hub_model_accepts_only_its_canonical_snapshot_identity(campaign, observed, qualified): + plan, source, checkpoint = campaign + revision = "a" * 40 + profile = json.loads(plan._fpm_profile_json) + profile["model_revision"] = revision + plan = replace(plan, _fpm_profile_json=json.dumps(profile)) + _write_campaign(source, checkpoint, plan) + snapshot = observed.format( + repo=f"models--{plan.model_path.replace('/', '--')}", revision=revision, other_revision="b" * 40 + ) + + def alter(payload): + payload["resolved_config"]["model_config"].update(model=snapshot, loaded_config_commit_hash=revision) + + _mutate_workers(source, alter) + cell = plan.cells[0] + raw = source / "cells" / cell.cell_id / "raw" + collection = validate_native_collection(cell, raw, expected_plan_sha256=plan.sha256, expected_attempt_id="source") + result = execution_evidence.inspect_execution_evidence(cell, raw, collection, plan=plan) + assert result["status"] == ("qualified" if qualified else "failed"), result["failures"] + if not qualified: + assert "runtime model_config.model differs from generated --model" in result["failures"] + + +@pytest.mark.parametrize( + ("pin", "observed_revision", "loaded", "qualified"), + [ + ("a" * 40, "a" * 40, "a" * 40, True), + ("A" * 40, "A" * 40, "A" * 40, True), + ("a" * 40, None, "a" * 40, True), + ("a" * 40, "a" * 40, None, False), + ("a" * 40, None, None, False), + ("a" * 40, "a" * 40, "b" * 40, False), + ("a" * 40, "b" * 40, "a" * 40, False), + ("a" * 40, "a" * 40, "a" * 39, False), + ("a" * 40, "a" * 40, 123, False), + ("release-tag", "release-tag", "a" * 40, False), + ("a" * 39, "a" * 39, "a" * 40, False), + ], +) +def test_resolved_hub_snapshot_requires_pinned_loaded_revision(campaign, pin, observed_revision, loaded, qualified): + plan, source, checkpoint = campaign + profile = json.loads(plan._fpm_profile_json) + profile["model_revision"] = pin + plan = replace(plan, _fpm_profile_json=json.dumps(profile)) + _write_campaign(source, checkpoint, plan) + snapshot = f"/cache/models--{plan.model_path.replace('/', '--')}/snapshots/{'a' * 40}" + + def alter(payload): + model = payload["resolved_config"]["model_config"] + model.update(model=snapshot, revision=observed_revision) + if loaded is not None: + model["loaded_config_commit_hash"] = loaded + + _mutate_workers(source, alter) + cell = plan.cells[0] + raw = source / "cells" / cell.cell_id / "raw" + collection = validate_native_collection(cell, raw, expected_plan_sha256=plan.sha256, expected_attempt_id="source") + result = execution_evidence.inspect_execution_evidence(cell, raw, collection, plan=plan) + assert result["status"] == ("qualified" if qualified else "failed"), result["failures"] + + +@pytest.mark.parametrize("launch_revision", ["a" * 40, "b" * 40, "release-tag", ""]) +def test_resolved_hub_snapshot_does_not_override_generated_revision(campaign, launch_revision): + plan, source, checkpoint = campaign + revision = "a" * 40 + profile = json.loads(plan._fpm_profile_json) + profile["model_revision"] = revision + plan = replace(plan, _fpm_profile_json=json.dumps(profile)) + _write_campaign(source, checkpoint, plan) + _mutate_workers( + source, + lambda payload: payload["resolved_config"]["model_config"].update( + model=f"/cache/models--{plan.model_path.replace('/', '--')}/snapshots/{revision}", + loaded_config_commit_hash=revision, + ), + ) + cell = plan.cells[0] + directory = source / "cells" / cell.cell_id + script = directory / "run.sh" + command = script.read_text() + prefix = "engine_command=(python3 -m dynamo.vllm " + assert command.count(prefix) == 1 + script.write_text(command.replace(prefix, f"{prefix}--revision {launch_revision} ")) + collection = validate_native_collection( + cell, directory / "raw", expected_plan_sha256=plan.sha256, expected_attempt_id="source" + ) + result = execution_evidence.inspect_execution_evidence(cell, directory / "raw", collection, plan=plan) + assert result["status"] == ("qualified" if launch_revision == revision else "failed") + if launch_revision != revision: + assert "runtime model_config.model differs from generated --model" in result["failures"] + + +def test_explicit_checkpoint_mount_does_not_accept_a_hub_snapshot_alias(campaign): + plan, source, checkpoint = campaign + revision = "a" * 40 + profile = json.loads(plan._fpm_profile_json) + profile["model_revision"] = revision + deployment = {"K8sConfig": {"k8s_pvc_mount_path": "/models", "k8s_model_path_in_pvc": "checkpoint"}} + plan = replace( + plan, + _fpm_profile_json=json.dumps(profile), + generator_config_sha256=repeatability._canonical_hash(with_kv_warmup_defaults(deployment)), + ) + _write_campaign(source, checkpoint, plan, generator_overrides=deployment) + _mutate_workers( + source, + lambda payload: payload["resolved_config"]["model_config"].update( + model=f"/cache/models--{plan.model_path.replace('/', '--')}/snapshots/{revision}", + loaded_config_commit_hash=revision, + ), + ) + cell = plan.cells[0] + raw = source / "cells" / cell.cell_id / "raw" + collection = validate_native_collection(cell, raw, expected_plan_sha256=plan.sha256, expected_attempt_id="source") + result = execution_evidence.inspect_execution_evidence(cell, raw, collection, plan=plan) + assert result["status"] == "failed" + assert "runtime model_config.model differs from generated --model" in result["failures"] + + +@pytest.mark.parametrize("observed", [128, 256]) +def test_explicit_cache_sizing_is_checked_against_launch(campaign, observed): + plan, source, checkpoint = campaign + cell = plan.cells[0] + policy = repeatability.BackendPolicy( + "explicit-cache", {"params": {"agg": {"extra_cli_args": ["--num-gpu-blocks-override", "128"]}}}, {} + ) + plan = replace(plan, cells=(replace(cell, backend_policy=policy),), backend_policies=(policy,)) + plan = repeatability._subset_plan( + plan, {"cell_id": cell.cell_id, "benchmark_points": json.loads(plan.options.benchmark_points_json)} + ) + _write_campaign(source, checkpoint, plan) + _mutate_workers( + source, lambda payload: payload["resolved_config"]["cache_config"].update(num_gpu_blocks_override=observed) + ) + result = repeatability.freeze_repeatability_plan(plan, source, checkpoint) + assert result["cells"][0]["execution"]["status"] == ("qualified" if observed == 128 else "failed") + + +@pytest.mark.parametrize("last", [True, False]) +@pytest.mark.parametrize("matches", [True, False]) +def test_execution_only_boolean_overrides_use_last_flag(campaign, last, matches): + plan, source, checkpoint = campaign + positive, negative = "--async-scheduling", "--no-async-scheduling" + args = [positive, negative, positive] if last else [positive, negative] + policy = repeatability.BackendPolicy("boolean-override", {"params": {"agg": {"extra_cli_args": args}}}, {}) + cell = replace(plan.cells[0], backend_policy=policy) + plan = replace(plan, cells=(cell,), backend_policies=(policy,)) + plan = repeatability._subset_plan( + plan, {"cell_id": cell.cell_id, "benchmark_points": json.loads(plan.options.benchmark_points_json)} + ) + assert not runner._observe_runtime_memory(plan, cell) + _write_campaign(source, checkpoint, plan) + _mutate_workers( + source, + lambda payload: payload["resolved_config"]["scheduler_config"].update( + async_scheduling=last if matches else not last + ), + ) + result = repeatability.freeze_repeatability_plan(plan, source, checkpoint) + assert result["cells"][0]["execution"]["status"] == ("qualified" if matches else "failed") + + +def test_unrequested_explicit_cache_sizing_fails(campaign): + _mutate_workers( + campaign[1], lambda payload: payload["resolved_config"]["cache_config"].update(kv_cache_memory_bytes=1024) + ) + result = repeatability.freeze_repeatability_plan(*campaign) + assert result["cells"][0]["execution"]["status"] == "failed" + + +@pytest.mark.parametrize("missing", ["checkpoint", "native_rank", "worker", "launch"]) +def test_missing_sample_artifacts_stay_incomplete(campaign, tmp_path, monkeypatch, missing): + def mutate(root, _index): + if missing == "checkpoint": + (root.parent.parent / "checkpoint" / "fpm_forward.json").unlink() + elif missing == "native_rank": + next(root.glob("cells/*/raw/*/benchmark-dp*.json")).unlink() + elif missing == "worker": + next(root.glob("cells/*/raw/*/fpm-execution-worker*.json")).unlink() + else: + next(root.glob("cells/*/run.sh")).unlink() + + _fake_collector(monkeypatch, mutate=mutate) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "incomplete" + + +def test_backend_mismatch_fails_until_authorized_retry_and_preserves_history(campaign, tmp_path, monkeypatch): + def mutate(root, index): + if index == 1: + _mutate_workers(root, lambda payload: payload["attention_groups"][0].update(backend_class="other.Backend")) + + _fake_collector(monkeypatch, mutate=mutate) + args = _args(campaign, tmp_path) + report = repeatability.run_repeatability(**args) + assert report["status"] == "failed" + assert all(point["sample_stddev_seconds"] is None for point in report["points"]) + report = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert report["status"] == "passed" + attempts = report["samples"][campaign[0].cells[0].cell_id][0]["attempts"] + assert [attempt["status"] for attempt in attempts] == ["failed", "passed"] + assert attempts[0]["failure_kind"] == "validation_failed" + + +@pytest.mark.parametrize( + "flag,value", [("--kv-cache-memory-bytes", "1073741824"), ("--num-gpu-blocks-override", "128")] +) +def test_execution_observer_keeps_explicit_declared_cache_settings(campaign, flag, value): + cell = campaign[0].cells[0] + policy = repeatability.BackendPolicy("explicit", {"params": {"agg": {"extra_cli_args": [flag, value]}}}, {}) + cell = replace(cell, backend_policy=policy) + assert not runner._observe_runtime_memory(campaign[0], cell) + args = runner._cell_generator_overrides(campaign[0], cell, {})["params"]["agg"]["extra_cli_args"] + assert args[args.index(flag) + 1] == value + + +def test_execution_observer_owns_worker_class(campaign): + cell = campaign[0].cells[0] + policy = repeatability.BackendPolicy( + "explicit", {"params": {"agg": {"extra_cli_args": ["--worker-cls=Other"]}}}, {} + ) + with pytest.raises(ValueError, match="observer"): + runner._cell_generator_overrides(campaign[0], replace(cell, backend_policy=policy), {}) + + +def test_unstable_points_fail_and_can_be_reassessed_without_new_launches(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, factors=[0.5, 1.5]) + args = _args(campaign, tmp_path) + report = repeatability.run_repeatability(**args) + assert report["status"] == "failed" and report["repeatability"]["status"] == "unstable" + frozen = json.loads((args["output_dir"] / repeatability.PLAN_FILENAME).read_text()) + reassessed = repeatability.assess_repeatability(frozen, report, cv_threshold=0.75) + assert reassessed["status"] == "failed" # Source agreement remains a separate criterion. + reassessed = repeatability.assess_repeatability(frozen, report, cv_threshold=0.75, source_agreement_threshold=0.75) + assert reassessed["status"] == "passed" and len(calls) == 5 + assert report["status"] == "failed" + assert reassessed["assessment"]["cv_threshold"] == 0.75 + + +def test_stable_new_samples_do_not_accept_an_outlying_published_measurement(campaign, tmp_path, monkeypatch): + plan, source, checkpoint = campaign + _write_campaign(source, checkpoint, plan, factor=10.0) + _fake_collector(monkeypatch) + originals = runner._file_manifest(source) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "failed" + assert report["new_population"]["status"] == "qualified" + assert report["repeatability"]["status"] == "passed" + assert report["source_qualification"]["status"] == "failed" + assert all(point["sample_cv"] == 0 and point["source_inclusive_cv"] > 0.05 for point in report["points"]) + assert all(point["source_inclusive_sample_count"] == 6 for point in report["points"]) + assert runner._file_manifest(source) == originals + + +def test_failure_resume_preserves_failed_attempt_and_uses_new_directory(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, fail_at=2) + args = _args(campaign, tmp_path) + report = repeatability.run_repeatability(**args) + assert report["status"] == "incomplete" and len(calls) == 2 + assert all(point["sample_count"] == 1 and point["sample_stddev_seconds"] is None for point in report["points"]) + failed_raw = Path(calls[1][1]["artifact_root"]) + saved_failed = runner._file_manifest(failed_raw) + assert repeatability.run_repeatability(**args, resume=True) == report + assert len(calls) == 2 + report = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert report["status"] == "passed" and len(calls) == 6 + entries = report["samples"][campaign[0].cells[0].cell_id] + assert [item["status"] for item in entries[1]["attempts"]] == ["failed", "passed"] + assert runner._file_manifest(failed_raw) == saved_failed + + +@pytest.mark.parametrize("mutation", ["source", "sample", "policy", "deployment"]) +def test_resume_rejects_mutated_evidence_without_recollection(campaign, tmp_path, monkeypatch, mutation): + calls = _fake_collector(monkeypatch) + args = _args(campaign, tmp_path) + repeatability.run_repeatability(**args) + if mutation == "source": + path = next(campaign[1].glob("cells/*/raw/*/benchmark-*.json")) + path.write_text(path.read_text() + " ") + elif mutation == "sample": + path = next(args["output_dir"].glob("samples/**/raw/*/benchmark-*.json")) + path.write_text(path.read_text() + " ") + elif mutation == "policy": + args["samples"] = 6 + else: + args["generator_overrides"] = {"K8sConfig": {"k8s_image": "different"}} + with pytest.raises(ValueError, match="changed|differ"): + repeatability.run_repeatability(**args, resume=True) + assert len(calls) == 5 + + +def test_reused_attempt_is_not_an_independent_sample(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, duplicate=True) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert report["status"] == "failed" and len(calls) == 2 + entries = report["samples"][campaign[0].cells[0].cell_id] + assert "independent launch" in entries[1]["attempts"][-1]["error"] + + +@pytest.mark.parametrize("change", ["backend", "graph", "seed", "coordinate"]) +def test_repeat_changed_runtime_regime_never_enters_repeatability_statistics(campaign, tmp_path, monkeypatch, change): + def mutate(root, _index): + if change in {"backend", "graph"}: + for path in root.glob("cells/*/raw/*/fpm-execution*.json"): + value = json.loads(path.read_text()) + if change == "backend": + value["attention_groups"][0]["backend_class"] = "another.Backend" + else: + value["graph_config"]["cudagraph_mode"] = "NONE" + path.write_text(json.dumps(value)) + else: + for path in root.glob("cells/*/raw/*/benchmark-*.json"): + value = json.loads(path.read_text()) + if change == "seed": + value["results"][0]["kv_seed_regime"] = "fake_prefix" + value["results"][0]["point"]["sample_reasons"] = ["prefill_fake_prefix"] + value["iteration_groups"][0]["point"]["sample_reasons"] = ["prefill_fake_prefix"] + else: + value["results"][0]["point"]["batch_size"] += 1 + path.write_text(json.dumps(value)) + + calls = _fake_collector(monkeypatch, mutate=mutate) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert len(calls) == 1 and report["status"] == "failed" + assert all(point["sample_count"] == 0 for point in report["points"]) + + +def test_old_artifacts_without_execution_evidence_remain_incomplete(campaign): + plan, source, checkpoint = campaign + _mutate_native(source, _remove_execution_protocol) + for path in source.glob("cells/*/raw/*/fpm-execution*.json"): + path.unlink() + frozen = repeatability.freeze_repeatability_plan(plan, source, checkpoint) + evidence = frozen["cells"][0]["execution"] + assert evidence["status"] == "incomplete" + assert any("worker execution observations missing" in reason for reason in evidence["missing_evidence"]) + + +def test_execution_observer_preserves_mixed_selected_backends_and_graph_config(tmp_path, monkeypatch): + monkeypatch.setattr(observer, "RESULTS_DIR", tmp_path) + monkeypatch.setattr(observer.importlib.metadata, "version", lambda _name: "0.28.0") + (tmp_path / "collector-provenance.json").write_text("{}") + backend_a = type("BackendA", (), {"__module__": "actual.runtime"}) + backend_b = type("BackendB", (), {"__module__": "actual.runtime"}) + worker = SimpleNamespace( + vllm_config=_vllm_config(), + model_runner=SimpleNamespace( + attn_groups=[ + [ + SimpleNamespace(backend=backend_a, layer_names=["layer.0"], kv_cache_group_id=0), + SimpleNamespace(backend=backend_b, layer_names=["layer.1"], kv_cache_group_id=0), + ] + ] + ), + ) + worker.vllm_config.model_config.hf_config = SimpleNamespace(_commit_hash="a" * 40) + del worker.vllm_config.offload_config + observer.observe_execution(worker, dp_rank=0, tp_rank=0, pp_rank=0) + result = json.loads((tmp_path / "fpm-execution-worker-dp0-tp0-pp0.json").read_text()) + assert result["status"] == "observed" + assert [item["backend_class"] for item in result["attention_groups"]] == [ + "actual.runtime.BackendA", + "actual.runtime.BackendB", + ] + assert result["graph_config"]["cudagraph_mode"] == "FULL_AND_PIECEWISE" + assert result["resolved_config"]["compilation_config"] == result["graph_config"] + assert result["resolved_config"]["model_config"]["loaded_config_commit_hash"] == "a" * 40 + assert result["resolved_config"]["quantization_config"]["quant_method"] == "NVFP4" + assert "offload_config" not in result["resolved_config"] + assert result["per_point_dispatch"] == "unreported" + assert not list(tmp_path.glob("fpm-memory-*.json")) + with pytest.raises(RuntimeError, match="duplicate"): + observer.observe_execution(worker, dp_rank=0, tp_rank=0, pp_rank=0) + + +@pytest.mark.parametrize("version", ["0.27.0", "0.28.0", "0.29.0"]) +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_execution_only_observation_does_not_enable_memory_conversion(version): + profile = _profile() + for deployment in profile["deployments"]: + deployment["backend_version"] = version + plan = _plan(profile) + for cell in plan.cells: + args = runner._cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + assert runner._observe_runtime_memory(plan, cell) is False + if version in {"0.27.0", "0.28.0"}: + assert args[args.index("--worker-cls") + 1] == "fpm_memory_worker.FpmExecutionWorker" + assert args[args.index("--scheduler-cls") + 1] == "fpm_memory_scheduler.FpmExecutionInstrumentedScheduler" + else: + assert "--worker-cls" not in args + + +def test_cli_exposes_bounded_repeatability_defaults(): + args = cli._parser().parse_args(["--gpu", "gb300"]) + assert (args.repeatability_samples, args.repeatability_max_points, args.repeatability_cv_threshold) == (5, 12, 0.05) + + +def test_saved_loader_roundtrips_without_model_or_hardware_resolution(campaign): + plan, source, _checkpoint = campaign + assert repeatability.load_repeatability_source(source).to_dict() == plan.to_dict() + assert repeatability.load_repeatability_deployment(source) == with_kv_warmup_defaults({}) + + +def test_cli_preview_uses_saved_plan_without_model_or_gpu_arguments(campaign, tmp_path, monkeypatch, capsys): + plan, source, checkpoint = campaign + monkeypatch.setattr(cli, "build_collection_case_plan", lambda **kwargs: pytest.fail("re-resolved source plan")) + assert ( + cli.main( + [ + "--repeatability-source-campaign", + str(source), + "--repeatability-source-checkpoint", + str(checkpoint), + "--repeatability-output-dir", + str(tmp_path / "validation"), + "--plan-only", + ] + ) + == 0 + ) + assert json.loads(capsys.readouterr().out)["source_plan_sha256"] == plan.sha256 + assert not (tmp_path / "validation").exists() + + +def test_new_code_cannot_silently_remeasure_old_campaign(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch) + monkeypatch.setattr(repeatability.planner, "_git_revision", lambda: "different-code") + with pytest.raises(ValueError, match="checkout/code.*unqualified"): + repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert calls == [] + + +def test_interrupted_sample_is_checkpointed_without_deleting_evidence(campaign, tmp_path, monkeypatch): + def interrupted(plan, **kwargs): + root = Path(kwargs["artifact_root"]) + root.mkdir(parents=True) + (root / "runtime-failure.log").write_text("partial native output") + raise KeyboardInterrupt + + monkeypatch.setattr(repeatability, "run_collection", interrupted) + args = _args(campaign, tmp_path) + with pytest.raises(KeyboardInterrupt): + repeatability.run_repeatability(**args) + report = json.loads((args["output_dir"] / repeatability.REPORT_FILENAME).read_text()) + sample = report["samples"][campaign[0].cells[0].cell_id][0] + assert sample["attempts"][0]["status"] == "interrupted" + assert next(args["output_dir"].rglob("runtime-failure.log")).read_text() == "partial native output" + + +def test_receipt_mismatch_is_a_failed_sample(campaign, tmp_path, monkeypatch): + def mutate(root, _index): + path = next(root.glob(f"cells/*/raw/*/{runner.POINTS_RECEIPT_FILENAME}")) + value = json.loads(path.read_text()) + value["sha256"] = "wrong" + path.write_text(json.dumps(value)) + + _fake_collector(monkeypatch, mutate=mutate) + report = repeatability.run_repeatability(**_args(campaign, tmp_path)) + sample = report["samples"][campaign[0].cells[0].cell_id][0] + assert sample["attempts"][0]["status"] == "failed" + assert "receipt mismatch" in sample["attempts"][0]["error"] + + +def test_failed_source_point_is_never_selected(campaign): + plan, source, checkpoint = campaign + path = next(source.glob("cells/*/raw/*/benchmark-*.json")) + value = json.loads(path.read_text()) + value["timing_valid"] = False + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="invalid native terminal"): + repeatability.freeze_repeatability_plan(plan, source, checkpoint) + + +def _mutate_native(root, action): + for path in root.glob("cells/*/raw/*/benchmark-*.json"): + payload = json.loads(path.read_text()) + action(payload) + path.write_text(json.dumps(payload)) + + +def _remove_protocol(payload): + payload.pop("measurement_protocol") + for row in payload["results"]: + for fpm in row["fpms"]: + fpm.pop("benchmark_measurement", None) + for group in payload["iteration_groups"]: + for rank in group["rank_results"]: + for fpm in rank["fpms"]: + fpm.pop("benchmark_measurement", None) + + +def test_full_sweep_preserves_native_generation_and_assesses_every_coordinate(campaign, tmp_path, monkeypatch): + original, source, checkpoint = campaign + original = replace( + original, options=replace(original.options, benchmark_points_json=None, benchmark_points_sha256=None) + ) + plan = repeatability._subset_plan( + original, {"cell_id": original.cells[0].cell_id, "benchmark_points": {}}, comparison_mode="full_grid" + ) + _write_campaign(source, checkpoint, plan) + calls = _fake_collector(monkeypatch) + result = repeatability.run_repeatability(**_args((plan, source, checkpoint), tmp_path, max_points_per_cell=1)) + assert result["status"] == "passed" + assert len(result["points"]) == 12 + assert all(call.options == plan.options and call.options.benchmark_points_json is None for call, _ in calls) + aggregate = json.loads(Path(result["aggregate"]["path"]).read_text()) + assert aggregate["status"] == "qualified" + assert aggregate["formal_source_replaced"] is False + assert all(point["median_seconds"] == statistics.median(point["samples_seconds"]) for point in aggregate["points"]) + + +def test_bounded_population_never_qualifies_source_or_blends_source_timing(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, factors=[2.0]) + result = repeatability.run_repeatability( + **_args(campaign, tmp_path, comparison_mode="bounded", max_points_per_cell=6) + ) + assert len(calls) == 5 and len(result["points"]) == 6 + assert result["status"] == "incomplete" + assert result["new_population"] == {"status": "qualified", "scope": "bounded_diagnostic"} + assert result["source_qualification"]["status"] == "unestablished" + assert result["repeatability"]["status"] == "passed" + assert all(point["sample_cv"] == 0 and point["source_inclusive_cv"] > 0.05 for point in result["points"]) + assert all(point["source_comparison"]["context"] == "cross_context_unproven" for point in result["points"]) + + +def test_legacy_source_does_not_prevent_new_comparable_population(campaign, tmp_path, monkeypatch): + _mutate_native(campaign[1], _remove_protocol) + calls = _fake_collector(monkeypatch) + result = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert len(calls) == 5 + assert result["status"] == "incomplete" + assert result["new_population"]["status"] == "qualified" + assert result["source_qualification"]["status"] == "unestablished" + assert all(point["source_relative_difference"] is None for point in result["points"]) + + +def test_missing_new_protocol_preserves_raw_values_without_qualified_statistics(campaign, tmp_path, monkeypatch): + _fake_collector(monkeypatch, mutate=lambda root, _index: _mutate_native(root, _remove_protocol)) + result = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert result["new_population"]["status"] == "unestablished" + assert all(len(point["raw_samples_seconds"]) == 5 and point["samples_seconds"] == [] for point in result["points"]) + assert all(point["sample_cv"] is None for point in result["points"]) + + +def test_changed_prompt_seed_cannot_enter_same_population(campaign, tmp_path, monkeypatch): + def mutate(root, index): + if index == 2: + _mutate_native(root, lambda payload: _add_measurement_protocol(payload, seed="different")) + + calls = _fake_collector(monkeypatch, mutate=mutate) + result = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert len(calls) == 2 and result["status"] == "failed" + assert result["repeatability"]["status"] == "mismatch" + assert all(point["sample_cv"] is None and len(point["raw_samples_seconds"]) == 2 for point in result["points"]) + + +def test_valid_slow_outlier_is_retained_and_cannot_be_retried_away(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, factors=[1, 1, 3, 1, 1]) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + assert result["new_population"]["status"] == "unstable" + assert all(point["median_seconds"] == point["source_wall_time_seconds"] for point in result["points"]) + assert all(max(point["samples_seconds"]) == 3 * point["median_seconds"] for point in result["points"]) + assert repeatability.run_repeatability(**args, resume=True, retry_failed=True) == result + assert len(calls) == 5 + + +def test_retry_budget_is_frozen_and_finite(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch, fail_at=1) + args = _args(campaign, tmp_path, max_attempts_per_sample=1) + repeatability.run_repeatability(**args) + with pytest.raises(ValueError, match="retry budget exhausted"): + repeatability.run_repeatability(**args, resume=True, retry_failed=True) + with pytest.raises(ValueError, match="frozen plan.*changed"): + repeatability.run_repeatability(**{**args, "max_attempts_per_sample": 2}, resume=True, retry_failed=True) + assert len(calls) == 1 + + +def test_distinct_attempt_ids_do_not_make_reused_runtime_run_independent(campaign, tmp_path, monkeypatch): + def mutate(root, _index): + _mutate_native(root, lambda payload: payload.update(run_id="reused-runtime-run")) + + calls = _fake_collector(monkeypatch, mutate=mutate) + result = repeatability.run_repeatability(**_args(campaign, tmp_path)) + assert result["status"] == "failed" and len(calls) == 2 + attempt = result["samples"][campaign[0].cells[0].cell_id][1]["attempts"][0] + assert "runtime run" in attempt["error"] + + +def test_complete_measurement_survives_later_collection_error_and_is_not_replaced(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch) + collect = repeatability.run_collection + + def fail_after_collection(plan, **kwargs): + collect(plan, **kwargs) + return [{"classification": "teardown_error"}] + + monkeypatch.setattr(repeatability, "run_collection", fail_after_collection) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + assert result["status"] == "incomplete" + assert all(point["raw_sample_count"] == 1 for point in result["points"]) + assert repeatability.run_repeatability(**args, resume=True, retry_failed=True) == result + assert len(calls) == 1 + + +@pytest.mark.parametrize( + ("collector_status", "cleanup_error"), + [("failed", False), ("interrupted", False), ("running", False), ("failed", True)], +) +def test_complete_failed_native_observation_requires_cpu_recovery_without_replacement( + campaign, tmp_path, monkeypatch, collector_status, cleanup_error +): + calls = _fake_collector(monkeypatch, factors=[3, 1, 1, 1, 1]) + collect = repeatability.run_collection + cell_id = campaign[0].cells[0].cell_id + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: SimpleNamespace(cleanup=lambda: None)) + + def fail_after_native_collection(plan, **kwargs): + result = collect(plan, **kwargs) + if len(calls) == 1: + path = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + checkpoint = json.loads(path.read_text()) + entry = checkpoint["cells"][cell_id] + entry.update(status=collector_status, error_type="RuntimeError", error="log retrieval failed") + if cleanup_error: + entry["cleanup_error"] = "teardown timed out" + path.write_text(json.dumps(checkpoint)) + if collector_status == "interrupted": + raise KeyboardInterrupt("interrupted after native artifacts were copied") + return [{"classification": "campaign_cell_failed", "error_message": entry["error"]}] + return result + + monkeypatch.setattr(repeatability, "run_collection", fail_after_native_collection) + args = _args(campaign, tmp_path) + if collector_status == "interrupted": + with pytest.raises(KeyboardInterrupt, match="artifacts were copied"): + repeatability.run_repeatability(**args) + else: + repeatability.run_repeatability(**args) + result = repeatability.run_repeatability(**args, resume=True) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert len(calls) == 1 and result["status"] == "incomplete" + assert attempt["evidence"]["collector_status"] == collector_status + assert attempt["evidence"]["collector_error"]["error"] == "log retrieval failed" + assert all(point["raw_samples_seconds"] == [3 * point["source_wall_time_seconds"]] for point in result["points"]) + checkpoint_path = Path(calls[0][1]["checkpoint_dir"]) / "fpm_forward.json" + original_checkpoint = checkpoint_path.read_bytes() + + recover = repeatability._recover_completed_attempt + monkeypatch.setattr(repeatability, "_recover_completed_attempt", lambda *_args: None) + result = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert len(calls) == 1 and result["status"] == "incomplete" + assert attempt["postprocessing_attempts"][-1]["status"] == "failed" + assert "cannot be replaced" in attempt["postprocessing_attempts"][-1]["error"] + + monkeypatch.setattr(repeatability, "_recover_completed_attempt", recover) + result = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert len(calls) == 5 and result["new_population"]["status"] == "unstable" + assert all(point["raw_samples_seconds"][0] == 3 * point["source_wall_time_seconds"] for point in result["points"]) + assert all(point["sample_count"] == 5 and point["sample_cv"] > 0.05 for point in result["points"]) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert [item["status"] for item in attempt["postprocessing_attempts"]] == ["failed", "passed"] + assert attempt["postprocessing_attempts"][-1]["artifact_recovery"]["original_status"] == collector_status + assert all(len(sample["attempts"]) == 1 for sample in result["samples"][cell_id]) + assert checkpoint_path.read_bytes() == original_checkpoint + assert repeatability.run_repeatability(**args, resume=True, retry_failed=True) == result + assert len(calls) == 5 + + +@pytest.mark.parametrize("defect", ["missing_rank", "partial_point", "wrong_attempt"]) +def test_failed_checkpoint_with_invalid_native_artifacts_remains_retryable(campaign, tmp_path, monkeypatch, defect): + def invalidate(root, index): + if index != 1: + return + path = next(root.glob("cells/*/raw/*/benchmark-*.json")) + if defect == "missing_rank": + path.unlink() + elif defect == "partial_point": + payload = json.loads(path.read_text()) + payload["results"].pop() + path.write_text(json.dumps(payload)) + else: + path = next(root.glob("cells/*/raw/*/collector-provenance.json")) + payload = json.loads(path.read_text()) + payload["attempt_id"] = "unrelated-attempt" + path.write_text(json.dumps(payload)) + + calls = _fake_collector(monkeypatch, mutate=invalidate) + collect = repeatability.run_collection + cell_id = campaign[0].cells[0].cell_id + + def fail_partial(plan, **kwargs): + result = collect(plan, **kwargs) + if len(calls) == 1: + path = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + checkpoint = json.loads(path.read_text()) + checkpoint["cells"][cell_id].update(status="failed", error="incomplete artifact collection") + path.write_text(json.dumps(checkpoint)) + return [{"classification": "campaign_cell_failed"}] + return result + + monkeypatch.setattr(repeatability, "run_collection", fail_partial) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + assert all(point["raw_sample_count"] == 0 for point in result["points"]) + result = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert result["status"] == "passed" and len(calls) == 6 + attempts = result["samples"][cell_id][0]["attempts"] + assert [attempt["status"] for attempt in attempts] == ["failed", "passed"] + assert "evidence" not in attempts[0] + + +def test_failed_checkpoint_never_recovers_a_wrong_runtime_configuration(campaign, tmp_path, monkeypatch): + def wrong_runtime(root, _index): + for path in root.glob("cells/*/raw/*/fpm-execution*.json"): + payload = json.loads(path.read_text()) + payload["resolved_config"]["model_config"]["revision"] = "wrong-model-revision" + path.write_text(json.dumps(payload)) + + calls = _fake_collector(monkeypatch, mutate=wrong_runtime) + collect = repeatability.run_collection + cell_id = campaign[0].cells[0].cell_id + + def fail_after_wrong_runtime(plan, **kwargs): + collect(plan, **kwargs) + path = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + checkpoint = json.loads(path.read_text()) + checkpoint["cells"][cell_id].update(status="failed", error="log retrieval failed") + path.write_text(json.dumps(checkpoint)) + return [{"classification": "campaign_cell_failed"}] + + monkeypatch.setattr(repeatability, "run_collection", fail_after_wrong_runtime) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert result["status"] == "failed" and attempt["failure_kind"] == "validation_failed" + assert attempt["evidence"]["execution"]["status"] == "failed" + assert repeatability.run_repeatability(**args, resume=True, retry_failed=True) == result + assert len(calls) == 1 + + +@pytest.mark.parametrize("interrupted_cleanup", [False, True]) +def test_cleanup_failed_checkpoint_retains_slow_measurement_and_retries_only_cleanup( + campaign, tmp_path, monkeypatch, interrupted_cleanup +): + calls = _fake_collector(monkeypatch, factors=[3, 1, 1, 1, 1]) + collect = repeatability.run_collection + cell_id = campaign[0].cells[0].cell_id + + def fail_cleanup(plan, **kwargs): + result = collect(plan, **kwargs) + if len(calls) == 1: + path = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + checkpoint = json.loads(path.read_text()) + checkpoint["cells"][cell_id].update(status="cleanup_failed", cleanup_error="transport timeout") + path.write_text(json.dumps(checkpoint)) + return [{"classification": "resource_cleanup_failed", "error_message": "transport timeout"}] + return result + + monkeypatch.setattr(repeatability, "run_collection", fail_cleanup) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert attempt["failure_kind"] == "cleanup_failed" + assert attempt["evidence"]["collector_status"] == "cleanup_failed" + assert result["execution"]["status"] == result["status"] == "incomplete" + assert all(point["raw_samples_seconds"] == [3 * point["source_wall_time_seconds"]] for point in result["points"]) + assert repeatability.run_repeatability(**args, resume=True) == result + assert len(calls) == 1 + checkpoint_path = Path(calls[0][1]["checkpoint_dir"]) / "fpm_forward.json" + original_checkpoint = checkpoint_path.read_bytes() + + def cleanup_still_fails(): + (checkpoint_path.parent.parent / "cleanup.log").write_text("cleanup attempted") + if interrupted_cleanup: + raise KeyboardInterrupt("cleanup interrupted") + raise RuntimeError("cleanup still unavailable") + + monkeypatch.setattr(repeatability, "_cell_runner", lambda *_args: SimpleNamespace(cleanup=cleanup_still_fails)) + if interrupted_cleanup: + with pytest.raises(KeyboardInterrupt, match="cleanup interrupted"): + repeatability.run_repeatability(**args, resume=True, retry_failed=True) + result = json.loads((args["output_dir"] / repeatability.REPORT_FILENAME).read_text()) + else: + result = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert len(calls) == 1 and result["status"] == "incomplete" + failed_status = "interrupted" if interrupted_cleanup else "failed" + assert result["samples"][cell_id][0]["attempts"][0]["cleanup_attempts"][-1]["status"] == failed_status + monkeypatch.setattr(repeatability, "_cell_runner", lambda *_args: SimpleNamespace(cleanup=lambda: None)) + result = repeatability.run_repeatability(**args, resume=True, retry_failed=True) + assert len(calls) == 5 + assert result["new_population"]["status"] == "unstable" + assert all(point["sample_count"] == 5 and point["sample_cv"] > 0.05 for point in result["points"]) + assert all(point["raw_samples_seconds"][0] == 3 * point["source_wall_time_seconds"] for point in result["points"]) + attempt = result["samples"][cell_id][0]["attempts"][0] + assert [item["status"] for item in attempt["cleanup_attempts"]] == [failed_status, "passed"] + assert attempt["errors"][0]["classification"] == "resource_cleanup_failed" + assert checkpoint_path.read_bytes() == original_checkpoint + assert all(len(sample["attempts"]) == 1 for sample in result["samples"][cell_id]) + assert repeatability.run_repeatability(**args, resume=True, retry_failed=True) == result + assert len(calls) == 5 + + +@pytest.mark.parametrize("sample_count", [1, 5]) +def test_interrupted_aggregate_publication_recovers_completed_observation( + campaign, tmp_path, monkeypatch, sample_count +): + calls = _fake_collector(monkeypatch) + args = _args(campaign, tmp_path) + atomic = repeatability._atomic_json + interrupted = False + + def interrupt_report(path, payload): + nonlocal interrupted + if ( + Path(path).name == repeatability.REPORT_FILENAME + and payload.get("points") + and payload["points"][0].get("raw_sample_count") == sample_count + and not interrupted + ): + interrupted = True + raise KeyboardInterrupt("interrupted after aggregate publication") + return atomic(path, payload) + + monkeypatch.setattr(repeatability, "_atomic_json", interrupt_report) + with pytest.raises(KeyboardInterrupt, match="aggregate publication"): + repeatability.run_repeatability(**args) + saved = json.loads((args["output_dir"] / repeatability.REPORT_FILENAME).read_text()) + assert repeatability.file_evidence(Path(saved["aggregate"]["path"])) == saved["aggregate"] + assert len(calls) == sample_count + monkeypatch.setattr(repeatability, "_atomic_json", atomic) + result = repeatability.run_repeatability(**args, resume=True) + assert result["status"] == "passed" and len(calls) == 5 + assert all(len(sample["attempts"]) == 1 for sample in result["samples"][campaign[0].cells[0].cell_id]) + assert repeatability.file_evidence(Path(saved["aggregate"]["path"])) == saved["aggregate"] + assert repeatability.run_repeatability(**args, resume=True) == result + + +def test_immutable_aggregate_tampering_still_rejects_resume(campaign, tmp_path, monkeypatch): + calls = _fake_collector(monkeypatch) + args = _args(campaign, tmp_path) + result = repeatability.run_repeatability(**args) + aggregate = Path(result["aggregate"]["path"]) + payload = json.loads(aggregate.read_text()) + payload["points"][0]["median_seconds"] *= 0.5 + aggregate.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="aggregate artifact changed"): + repeatability.run_repeatability(**args, resume=True) + assert len(calls) == 5 + + +def test_legacy_report_stays_historical_and_unqualified_during_offline_inspection(campaign, tmp_path, monkeypatch): + _fake_collector(monkeypatch) + args = _args(campaign, tmp_path) + report = repeatability.run_repeatability(**args) + frozen = json.loads((args["output_dir"] / repeatability.PLAN_FILENAME).read_text()) + frozen["schema_version"] = 1 + frozen["sha256"] = repeatability._canonical_hash({key: value for key, value in frozen.items() if key != "sha256"}) + report["schema_version"] = 1 + report["plan_sha256"] = frozen["sha256"] + before = copy.deepcopy(report) + inspected = repeatability.assess_repeatability(frozen, report, cv_threshold=0.05) + assert inspected["schema_version"] == 1 and inspected["status"] == "incomplete" + assert inspected["historical_assessment"]["status"] == "passed" + assert inspected["source_qualification"]["status"] == "unestablished" + assert report == before diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runner.py b/python/aisimulate/tests/unit/collector/test_fpm_runner.py index 676839dc9..858d069c5 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_runner.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_runner.py @@ -1054,6 +1054,54 @@ def test_formal_prefill_metadata_records_candidate_axis_counts(): } +def test_runtime_graph_policy_metadata_does_not_invent_axis_counts(): + cell = _cell() + plan = _plan(cell) + plan.options.prefill_sampling = PrefillSamplingProfile.build( + max_isl=8192, max_batch_size=256, cudagraph_policy="runtime" + ) + assert _configured_sampling_metadata(plan, cell, smoke=False) == { + "prefill_cudagraph_policy": "runtime", + "prefill_cudagraph_capture_size_count": None, + "prefill_requested_new_token_axis_count": None, + "prefill_max_new_token_samples": None, + } + assert _configured_sampling_metadata(plan, cell, smoke=True) == { + "prefill_cudagraph_policy": "runtime", + "prefill_max_new_token_samples": 2, + } + + +@pytest.mark.parametrize("argument", ["--compilation-config", "--cudagraph-capture-sizes=[1,2]"]) +def test_runtime_graph_policy_rejects_backend_capture_overrides(argument): + cell = dataclasses.replace( + _cell(), backend_policy=BackendPolicy("conflict", {"params": {"agg": {"extra_cli_args": [argument]}}}, {}) + ) + plan = _plan(cell) + plan.options.prefill_sampling = PrefillSamplingProfile.build( + max_isl=8192, max_batch_size=256, cudagraph_policy="runtime" + ) + with pytest.raises(ValueError, match="runtime prefill CUDA-graph policy cannot use backend capture overrides"): + _cell_generator_overrides(plan, cell, {}) + + +@pytest.mark.parametrize( + "override", + [ + {"kv_cache_free_gpu_memory_fraction": 0.8}, + {"extra_cli_args": ["--gpu-memory-utilization", "0.8"]}, + {"extra_cli_args": ["--gpu-memory-utilization=0.8"]}, + {"extra_cli_args": ["--gpu-memory-utilization 0.8"]}, + ], +) +def test_backend_policy_cannot_change_collector_memory_fraction(override): + cell = dataclasses.replace(_cell(), backend_policy=BackendPolicy("conflict", {"params": {"agg": override}}, {})) + plan = _plan(cell) + plan.options.gpu_memory_utilization = 0.9 + with pytest.raises(ValueError, match="backend policy cannot .*--fpm-gpu-memory-utilization"): + _cell_generator_overrides(plan, cell, {}) + + def test_cell_render_adds_one_collector_owned_benchmark_timeout(): cell = _cell() plan = _plan(cell) @@ -1105,6 +1153,7 @@ def test_native_collection_validation_accepts_balanced_total_points(tmp_path): "measured_batch_size_axis_count": 1, "measured_kv_read_axis_count": 1, "measured_new_token_axis_count": 1, + "native_kv_seed_regime_counts": {"unreported": 1}, } second_rank = tmp_path / "pod-1" / "benchmark_dp1.json" @@ -1115,6 +1164,106 @@ def test_native_collection_validation_accepts_balanced_total_points(tmp_path): _validate_runtime_collection(cell, tmp_path) +@pytest.mark.parametrize( + ("phase", "regime", "stamp", "kv_tokens"), + [ + ("prefill", "real_prefix", "prefill_real_seed", 128), + ("prefill", "real_prefix", "prefill_real_seed", 0), + ("prefill", "fake_prefix", "prefill_fake_prefix", 128), + ("prefill", "not_applicable", None, 0), + ("decode", "real_kv", "kvwarm_real_kv", 128), + ("decode", "fake_fallback", "kvwarm_fake_fallback", 128), + ("decode", "skip:dense_model", None, 128), + ], +) +def test_native_seed_diagnostics_report_observed_regime(tmp_path, phase, regime, stamp, kv_tokens): + cell = _cell(phase=phase, strategy="dense_tp") + payload = _native_payload(phase=phase, rank=0, dp=1) + row = payload["results"][0] + row["kv_seed_regime"] = regime + row["point"]["total_kv_read_tokens"] = kv_tokens + if stamp: + row["point"]["sample_reasons"].append(stamp) + if phase == "prefill": + row["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = kv_tokens + if regime.startswith("skip:"): + payload["kvwarm"] = {"enabled": True, "warm_eligible": False, "skip_reason": "dense_model"} + pod = tmp_path / "pod-0" + pod.mkdir() + _write_provenance(pod / "collector-provenance.json", cell_id=cell.cell_id) + (pod / "benchmark.json").write_text(json.dumps(payload)) + + summary = _runtime_collection_summary(cell, tmp_path) + assert summary["native_kv_seed_regime_counts"] == {regime: 1} + if phase == "decode": + assert summary["native_kvwarm"] == payload["kvwarm"] + + +@pytest.mark.parametrize("regime", ["real_prefix", "fake_prefix", "not_applicable"]) +@pytest.mark.parametrize("kv_tokens", [0, 128]) +def test_native_seed_diagnostics_reject_contradictory_injection_evidence(tmp_path, regime, kv_tokens): + cell = _cell() + payload = _native_payload(phase="prefill", rank=0, dp=1) + payload["results"][0]["kv_seed_regime"] = regime + payload["results"][0]["point"]["total_kv_read_tokens"] = kv_tokens + payload["results"][0]["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = kv_tokens + if regime != "real_prefix": + payload["results"][0]["point"]["sample_reasons"].append("prefill_real_seed") + pod = tmp_path / "pod-0" + pod.mkdir() + _write_provenance(pod / "collector-provenance.json", cell_id=cell.cell_id) + (pod / "benchmark.json").write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="kv_seed_regime disagrees"): + _runtime_collection_summary(cell, tmp_path) + + +@pytest.mark.parametrize("kv_tokens", [0, 128]) +def test_native_seed_diagnostics_reject_rank_evidence_disagreement(tmp_path, kv_tokens): + cell = _cell(dp=2) + for rank in range(2): + path = tmp_path / f"pod-{rank}" / ("benchmark.json" if rank == 0 else "benchmark_dp1.json") + path.parent.mkdir() + payload = _native_payload(phase="prefill", rank=rank, dp=2) + payload["results"][0]["point"]["sample_reasons"].append("prefill_real_seed") + payload["results"][0]["point"]["total_kv_read_tokens"] = kv_tokens + payload["results"][0]["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = kv_tokens + if rank == 0: + payload["results"][0]["kv_seed_regime"] = "real_prefix" + _write_provenance(path.parent / "collector-provenance.json", cell_id=cell.cell_id) + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="DP ranks disagree on per-point KV seed regime"): + _runtime_collection_summary(cell, tmp_path) + + +@pytest.mark.parametrize( + ("phase", "regime", "stamp", "kv_tokens", "scheduled_kv_tokens", "error"), + [ + ("decode", "real_prefix", "prefill_real_seed", 128, 128, "requires prefill"), + ("decode", "fake_prefix", "prefill_fake_prefix", 128, 128, "requires prefill"), + ("prefill", "fake_prefix", "prefill_fake_prefix", 0, 0, "requires cached prefill"), + ("prefill", "real_prefix", "prefill_real_seed", -1, -1, "invalid totals"), + ("prefill", "real_prefix", "prefill_real_seed", 0, 128, "workload mismatch"), + ], +) +def test_native_prefix_seed_diagnostics_reject_invalid_workload( + tmp_path, phase, regime, stamp, kv_tokens, scheduled_kv_tokens, error +): + cell = _cell(phase=phase) + payload = _native_payload(phase=phase, rank=0, dp=1) + row = payload["results"][0] + row["kv_seed_regime"] = regime + row["point"]["sample_reasons"].append(stamp) + row["point"]["total_kv_read_tokens"] = kv_tokens + row["fpms"][0]["scheduled_requests"][f"sum_{phase}_kv_tokens"] = scheduled_kv_tokens + pod = tmp_path / "pod-0" + pod.mkdir() + _write_provenance(pod / "collector-provenance.json", cell_id=cell.cell_id) + (pod / "benchmark.json").write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match=error): + _runtime_collection_summary(cell, tmp_path) + + def test_native_collection_validation_accepts_execution_order_decoupled_from_ids(tmp_path): """KV warm-up reorders execution (batch/kv-depth descending, fake points last) and the engine writes results in execution order, so file order is @@ -1246,9 +1395,16 @@ def test_runtime_summaries_use_native_rank_artifacts_and_skip_merged(tmp_path): } -def test_run_collection_stages_no_explicit_scheduler_or_case_manifest(monkeypatch, tmp_path): +@pytest.mark.parametrize("pending_memory", [False, True]) +@pytest.mark.parametrize("validation_only", [False, True]) +def test_run_collection_stages_owned_runtime_files(monkeypatch, tmp_path, pending_memory, validation_only): cell = _cell() plan = _plan(cell) + if pending_memory: + plan.fpm_profile = True + plan.deployment_profile = lambda _cell: SimpleNamespace( + backend_version="0.27.0", resources=SimpleNamespace(memory_source="pending") + ) events = [] staged_names = [] @@ -1286,6 +1442,13 @@ def cleanup(self): monkeypatch.setattr(fpm_runner, "_render_cell", render_cell) monkeypatch.setattr(fpm_runner, "KubernetesCellRunner", FakeResource) + from collector.fpm_forward import database + + monkeypatch.setattr( + database, + "write_formal_database", + lambda *_args, **_kwargs: pytest.fail("validation/smoke published formal data"), + ) monkeypatch.setattr( fpm_runner, "_runtime_collection_summary", @@ -1306,8 +1469,9 @@ def cleanup(self): artifact_root=str(tmp_path / "artifacts"), resume=False, retry_failed=False, - smoke=True, + smoke=not validation_only, cell_limit=1, + publish_database=not validation_only, ) assert errors == [] @@ -1318,14 +1482,20 @@ def cleanup(self): assert events.index("cleanup") < events.index("apply") # Contract: the staged set includes the rendered runtime artifacts # plus the collector's own in-pod runtime and preflight. - assert set(staged_names) == {"run.sh", "fpm_env.sh", "collector-runtime-env.sh", "fpm_exec.sh", "preflight.py"} + expected_files = {"run.sh", "fpm_env.sh", "collector-runtime-env.sh", "fpm_exec.sh", "preflight.py"} + if pending_memory: + expected_files.update({"fpm_memory_observer.py", "fpm_memory_worker.py", "fpm_memory_scheduler.py"}) + assert set(staged_names) == expected_files assert "cases.json" not in staged_names assert "fpm_scheduler.py" not in staged_names assert "run_with_etcd.sh" not in staged_names - checkpoint = json.loads((checkpoint_dir / "fpm_forward_smoke.json").read_text()) - assert checkpoint["cells"][cell.cell_id]["prefill_max_new_token_samples"] == 2 + name = "fpm_forward.json" if validation_only else "fpm_forward_smoke.json" + checkpoint = json.loads((checkpoint_dir / name).read_text()) + if not validation_only: + assert checkpoint["cells"][cell.cell_id]["prefill_max_new_token_samples"] == 2 + assert "prefill_requested_new_token_axis_count" not in checkpoint["cells"][cell.cell_id] assert checkpoint["cells"][cell.cell_id]["measured_new_token_axis_count"] == 2 - assert "prefill_requested_new_token_axis_count" not in checkpoint["cells"][cell.cell_id] + assert "database" not in checkpoint def test_partial_formal_run_is_campaign_incomplete_not_database_failure(monkeypatch, tmp_path): @@ -1464,6 +1634,92 @@ def reject_rerun(*_args, **_kwargs): assert checkpoint["smoke"]["status"] == "passed" +@pytest.mark.parametrize("missing_decode", [False, True]) +def test_resume_publishes_saved_zero_kv_prefill_without_rerun(monkeypatch, tmp_path, missing_decode): + """Synthetic saved native results exercise recovery and formal publication on CPU.""" + import pyarrow.parquet as pq + + cell = _cell() + plan = _plan(cell) + decode = _cell(phase="decode") + if missing_decode: + plan.cells = (cell, decode) + artifact_root = tmp_path / "artifacts" + root = artifact_root / plan.sha256[:16] + raw = root / "cells" / cell.cell_id / "raw" / "pod-0" + raw.mkdir(parents=True) + _write_provenance( + raw / "collector-provenance.json", + cell_id=cell.cell_id, + plan_sha256=plan.sha256, + attempt_id="attempt-1", + ) + payload = _native_payload(phase="prefill", rank=0, dp=1) + row = payload["results"][0] + row["kv_seed_regime"] = "real_prefix" + row["point"]["sample_reasons"].append("prefill_real_seed") + row["point"]["total_kv_read_tokens"] = 0 + row["fpms"][0]["scheduled_requests"]["sum_prefill_kv_tokens"] = 0 + native_path = raw / "benchmark.json" + native_path.write_text(json.dumps(payload)) + native_bytes = native_path.read_bytes() + + records = { + cell.cell_id: { + "status": "failed", + "attempt_id": "attempt-1", + "error_type": "ValueError", + "error": "native prefix seed regime requires cached prefill", + } + } + if missing_decode: + records[decode.cell_id] = {"status": "failed", "attempt_id": "attempt-decode"} + checkpoint_dir = tmp_path / "checkpoints" + checkpoint_dir.mkdir() + checkpoint_path = checkpoint_dir / "fpm_forward.json" + checkpoint_path.write_text( + json.dumps({"schema": fpm_runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": records}) + ) + + def reject_cluster_work(*_args, **_kwargs): + raise AssertionError("saved prefill recovery must not launch GPU work") + + monkeypatch.setattr(fpm_runner, "_render_cell", reject_cluster_work) + monkeypatch.setattr(fpm_runner, "_cell_runner", reject_cluster_work) + errors = run_collection( + plan, + generator_overrides={}, + checkpoint_dir=str(checkpoint_dir), + artifact_root=str(artifact_root), + resume=True, + retry_failed=False, + database_root=str(tmp_path / "db"), + publish_partial=missing_decode, + ) + + assert [error["classification"] for error in errors] == (["campaign_incomplete"] if missing_decode else []) + checkpoint = json.loads(checkpoint_path.read_text()) + record = checkpoint["cells"][cell.cell_id] + assert record["status"] == "passed" + assert record["attempt_id"] == "attempt-1" + assert record["native_kv_seed_regime_counts"] == {"real_prefix": 1} + assert record["artifact_recovery"]["original_status"] == "failed" + assert record["artifact_recovery"]["error"] == records[cell.cell_id]["error"] + assert native_path.read_bytes() == native_bytes + assert json.loads((root / "run-manifest.json").read_text())["attempts"] == [] + database = checkpoint["database"] + assert database["status"] == "passed" + assert database["missing_cells"] == ([decode.cell_id] if missing_decode else []) + rows = pq.read_table(database["parquet"]).to_pylist() + assert len(rows) == 1 + assert rows[0]["total_kv_read_tokens"] == 0 + assert rows[0]["kv_seed_regime"] == "n/a" + assert rows[0]["latency_ms"] == pytest.approx(10.0) + assert rows[0]["collector_attempt_id"] == "attempt-1" + if missing_decode: + assert checkpoint["cells"][decode.cell_id] == records[decode.cell_id] + + def test_plain_resume_recovers_interrupted_attempt_without_rerun(monkeypatch, tmp_path): cell = _cell() plan = _plan(cell) @@ -1673,6 +1929,62 @@ def test_typed_generator_render_uses_collector_prefill_axis(tmp_path): subprocess.run(["bash", "-n", str(tmp_path / "fpm_env.sh")], check=True) +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +@pytest.mark.parametrize("explicit", [None, "off", "0", "false", "TRUE"]) +def test_rendered_kv_warmup_defaults_and_overrides_reach_startup_and_engine(tmp_path, phase, explicit): + cell = _cell(phase=phase) + base = {"K8sConfig": {"k8s_image": "nvcr.io/nvidia/ai-dynamo/vllm-runtime:test"}} + names = ("DYN_BENCH_KV_WARMUP", "DYN_BENCH_PREFILL_REAL_SEED") + if explicit is not None: + base["K8sConfig"]["extra_env"] = [{"name": name, "value": explicit} for name in names] + _render_cell(_plan(cell), cell, tmp_path, base) + + expected = explicit if explicit is not None else "on" + run_script = (tmp_path / "run.sh").read_text() + for name in names: + assert f"export {name}={expected}" in run_script + # Both executors stage this exact startup file before preflight; Slurm + # does not inherit Pod environment, and run.sh has not executed yet. + actual = subprocess.run( + [ + "bash", + "-c", + 'source "$1/fpm_env.sh"; source "$1/collector-runtime-env.sh"; ' + 'printf "%s\\n%s\\n%s\\n" "$FPM_BENCHMARK_MODE" "$DYN_BENCH_KV_WARMUP" ' + '"$DYN_BENCH_PREFILL_REAL_SEED"', + "bash", + str(tmp_path), + ], + check=True, + capture_output=True, + text=True, + ) + assert actual.stdout.splitlines() == [phase, expected, expected] + + +def test_backend_policy_kv_warmup_override_is_preserved(): + cell = dataclasses.replace( + _cell(), + backend_policy=BackendPolicy( + "explicit-seeding", + {"K8sConfig": {"extra_env": [{"name": "DYN_BENCH_PREFILL_REAL_SEED", "value": "off"}]}}, + {}, + ), + ) + overrides = _cell_generator_overrides(_plan(cell), cell, {}) + environment = {item["name"]: item["value"] for item in overrides["K8sConfig"]["extra_env"]} + assert environment["DYN_BENCH_PREFILL_REAL_SEED"] == "off" + assert environment["DYN_BENCH_KV_WARMUP"] == "on" + + +@pytest.mark.parametrize("name", ["DYN_BENCH_KV_WARMUP", "DYN_BENCH_PREFILL_REAL_SEED"]) +@pytest.mark.parametrize("value", [True, 1, "maybe", " on ", "${CHOICE}"]) +def test_kv_warmup_switches_require_explicit_boolean_strings(name, value): + cell = _cell() + with pytest.raises(ValueError, match=f"{name} must be a literal"): + _cell_generator_overrides(_plan(cell), cell, {"K8sConfig": {"extra_env": [{"name": name, "value": value}]}}) + + @pytest.mark.parametrize( ("model_path", "weight_quantization"), [ diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runtime_memory.py b/python/aisimulate/tests/unit/collector/test_fpm_runtime_memory.py new file mode 100644 index 000000000..9b155d153 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_runtime_memory.py @@ -0,0 +1,948 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.util +import json +import sys +from dataclasses import replace +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +from collector.fpm_forward import planner, runner, runtime_memory +from collector.fpm_forward.runtime import fpm_memory_observer as observer +from collector.fpm_forward.types import ParallelTopology + +from aisimulate_core.fpm_profile import FpmResourceProfile + +from .test_fpm_profile_collection import _plan, _profile, no_models_or_timing_data # noqa: F401 +from .test_fpm_runner import _cell, _native_payload, _write_provenance + +pytestmark = pytest.mark.unit + + +def _typed(module, name): + return type(name, (), {"__module__": module})() + + +class _Storage: + device = "cuda:0" + + def __init__(self, pointer, size): + self.pointer, self.size = pointer, size + + def data_ptr(self): + return self.pointer + + def nbytes(self): + return self.size + + +class _Tensor: + def __init__(self, storage): + self.storage = storage + + def untyped_storage(self): + return self.storage + + +def _vllm_config(*, tp=2, dp=2): + quant = _typed("vllm.model_executor.layers.quantization.modelopt", "ModelOptNvFp4Config") + quant.quant_method = "NVFP4" + return SimpleNamespace( + model_config=SimpleNamespace( + model="example/model", + revision="immutable", + dtype="torch.bfloat16", + quantization="modelopt_fp4", + max_model_len=4096, + enforce_eager=False, + ), + cache_config=SimpleNamespace( + cache_dtype="bfloat16", gpu_memory_utilization=0.9, kv_cache_memory_bytes=None, enable_prefix_caching=True + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=1024, max_num_seqs=64, async_scheduling=False), + parallel_config=SimpleNamespace( + tensor_parallel_size=tp, + pipeline_parallel_size=1, + data_parallel_size=dp, + enable_expert_parallel=True, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + compilation_config=SimpleNamespace( + mode="VLLM_COMPILE", + cudagraph_mode="FULL_AND_PIECEWISE", + cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32, 64], + max_cudagraph_capture_size=64, + static_forward_context={}, + ), + kernel_config=SimpleNamespace(moe_backend="auto"), + quant_config=quant, + offload_config=SimpleNamespace( + offload_backend="auto", + uva=SimpleNamespace(cpu_offload_gb=0), + prefetch=SimpleNamespace(offload_group_size=0, offload_num_in_group=1, offload_prefetch_step=1), + ), + ) + + +@pytest.fixture(autouse=True) +def fake_runtime_offloader(monkeypatch): + monkeypatch.setattr(observer, "_offloader_type", lambda: "vllm.model_executor.offloader.base.NoopOffloader") + + +def _initialized(*, blocks=100, free=98, packed=False): + """Two physical buffers, five logical layer views and three retention groups.""" + config = _vllm_config() + storages = [_Storage(100, blocks * (128 if packed else 64))] + storages.append(storages[0] if packed else _Storage(200, blocks * 64)) + tensors = [_Tensor(storage) for storage in storages] + layers = {} + groups = [] + for prefix, count, window, block in (("full", 2, None, 16), ("slide", 2, 32, 16), ("conv", 1, 4, 4)): + names = [f"{prefix}.{index}" for index in range(count)] + for index, name in enumerate(names): + module, cls = ( + ("vllm.models.inkling.nvidia.sconv_swa_attn", "InklingConvState") + if prefix == "conv" + else ("vllm.model_executor.layers.attention", "Attention") + ) + layer = _typed(module, cls) + layer.kv_cache = tensors[index] + layers[name] = layer + spec = _typed("vllm.v1.kv_cache_interface", "FullAttentionSpec" if window is None else "SlidingWindowSpec") + spec.block_size, spec.page_size_bytes, spec.sliding_window, spec.dtype = block, 64, window, "torch.bfloat16" + groups.append(SimpleNamespace(layer_names=names, kv_cache_spec=spec, is_eagle_group=False)) + config.compilation_config.static_forward_context = layers + cache_config = SimpleNamespace( + num_blocks=blocks, + kv_cache_groups=groups, + kv_cache_tensors=[ + SimpleNamespace( + size=storages[0].size, + shared_by=["full.0", "slide.0", "conv.0"], + block_stride=128 if packed else 0, + offset=0, + ), + SimpleNamespace( + size=storages[1].size, + shared_by=["full.1", "slide.1"], + block_stride=128 if packed else 0, + offset=64 if packed else 0, + ), + ], + ) + worker = SimpleNamespace( + vllm_config=config, + _fpm_cache_initialized=True, + _fpm_available_cache_bytes=blocks * 128 + 200, + model_runner=SimpleNamespace( + kv_cache_config=cache_config, + kv_caches=[layer.kv_cache for layer in layers.values()], + shared_kv_cache_layers={}, + ), + peak_activation_memory=123, + cudagraph_memory_estimate=456, + ) + pool = SimpleNamespace( + num_gpu_blocks=blocks, get_num_free_blocks=lambda: free, null_block=SimpleNamespace(is_null=True, block_id=0) + ) + manager = SimpleNamespace( + block_pool=pool, + watermark_blocks=0, + coordinator=SimpleNamespace(single_type_managers=[SimpleNamespace(block_pool=pool) for _group in groups]), + ) + scheduler = SimpleNamespace(vllm_config=config, kv_cache_manager=manager) + return worker, scheduler, cache_config + + +def _evidence(tmp_path, monkeypatch): + cell = replace( + _cell(dp=2), + topology=ParallelTopology(tp=2, pp=1, dp=2, moe_tp=1, moe_ep=4, cp=1), + kv_cache_dtype="bfloat16", + fmha_quant_mode="bfloat16", + ) + monkeypatch.setattr(observer.importlib.metadata, "version", lambda _name: "0.27.0") + for dp in range(2): + pod = tmp_path / f"pod-{dp}" + pod.mkdir() + _write_provenance(pod / "collector-provenance.json", cell_id=cell.cell_id) + provenance = json.loads((pod / "collector-provenance.json").read_text()) + provenance["runtime"]["backend_version"] = "0.27.0" + (pod / "collector-provenance.json").write_text(json.dumps(provenance)) + (pod / "benchmark.json").write_text(json.dumps(_native_payload(phase="prefill", rank=dp, dp=2))) + monkeypatch.setattr(observer, "RESULTS_DIR", pod) + for tp in range(2): + worker, scheduler, config = _initialized(blocks=100 - dp * 10, free=98 - dp * 10) + observer.observe("worker", worker, dp_rank=dp, tp_rank=tp, pp_rank=0) + observer.observe("scheduler", scheduler, dp_rank=dp, cache_config=config) + return cell + + +def _resolve(cell, root, **kwargs): + return runtime_memory.resolve_runtime_resources( + cell, + root, + expected_plan_sha256="plan-sha", + expected_attempt_id="attempt", + expected_backend_version="0.27.0", + expected_context_length=4096, + expected_max_num_tokens=1024, + expected_max_batch_size=64, + expected_gpu_memory_utilization=0.9, + **kwargs, + ) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_physical_grouped_capacity_deduplicates_aliases_and_charges_padding(packed): + worker, scheduler, config = _initialized(packed=packed) + observed = observer.worker_memory(worker) + assert observed["allocated_cache_bytes"] == 12_800 + assert len(observed["storages"]) == (1 if packed else 2) + pool = observer.scheduler_memory(scheduler, config) + assert pool["initial_free_blocks"] == 98 + assert pool["reserved_blocks"] == 2 + assert [group["kind"] for group in observed["groups"]] == ["attention", "attention", "convolution"] + assert [group["sliding_window"] for group in observed["groups"]] == [None, 32, 4] + + +def test_resolve_runtime_resources_requires_all_ranks_and_preserves_graph_evidence(tmp_path, monkeypatch): + cell = _evidence(tmp_path, monkeypatch) + result = _resolve(cell, tmp_path) + resources = FpmResourceProfile.model_validate(result) + assert resources.memory_source == "runtime" + assert resources.runtime_memory.kv_cache_bytes == 88 * 128 + assert [group.page_size_bytes for group in resources.cache_groups] == [128, 128, 128] + assert [group.num_layers for group in resources.cache_groups] == [2, 2, 1] + assert resources.activations_bytes is None + provenance = json.loads(resources.runtime_memory.provenance) + assert provenance["runtime_settings"]["compilation_config"]["cudagraph_capture_sizes"] == [1, 2, 4, 8, 16, 32, 64] + assert "enable_prefix_caching" not in provenance["runtime_settings"]["cache_config"] + assert len(provenance["artifacts"]) == 6 + assert provenance["artifacts"][-1]["evidence"]["cache"]["diagnostics"]["peak_activation_memory"] == 123 + + +@pytest.mark.parametrize("backend", ["uva", "prefetch"]) +def test_active_weight_offload_is_observed_and_cannot_resolve(tmp_path, monkeypatch, backend): + original = _initialized + + def offloaded(**kwargs): + worker, scheduler, cache = original(**kwargs) + offload = worker.vllm_config.offload_config + offload.offload_backend = backend + offload.uva.cpu_offload_gb = 8 if backend == "uva" else 0 + offload.prefetch.offload_group_size = 4 if backend == "prefetch" else 0 + return worker, scheduler, cache + + monkeypatch.setattr(sys.modules[__name__], "_initialized", offloaded) + monkeypatch.setattr(observer, "_offloader_type", lambda: f"vllm.model_executor.offloader.{backend}.Offloader") + cell = _evidence(tmp_path, monkeypatch) + for path in tmp_path.glob("*/fpm-memory-*.json"): + evidence = json.loads(path.read_text()) + assert evidence["status"] == "unresolved" + assert evidence["resolved_config"]["offload_config"]["offload_backend"] == backend + if evidence["kind"] == "worker": + assert backend in evidence["effective_offloader"] + with pytest.raises(ValueError, match="unresolved"): + _resolve(cell, tmp_path) + + +def _edit_memory(root, edit): + for path in root.glob("*/fpm-memory-*.json"): + evidence = json.loads(path.read_text()) + edit(evidence) + path.write_text(json.dumps(evidence)) + + +@pytest.mark.parametrize("change", ["fp8", "unknown_quantization", "async", "offload", "effective_offloader"]) +def test_resolver_rejects_consistent_but_unsupported_runtime_policy(tmp_path, monkeypatch, change): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + config = evidence["resolved_config"] + if change == "fp8": + config["model_config"]["quantization"] = "fp8" + config["quantization_config"] = { + "type": "vllm.model_executor.layers.quantization.fp8.Fp8Config", + "activation_scheme": "dynamic", + "weight_block_size": None, + } + elif change == "unknown_quantization": + config["model_config"]["quantization"] = "compressed-tensors" + elif change == "async": + config["scheduler_config"]["async_scheduling"] = True + elif change == "offload": + config["offload_config"]["uva"]["cpu_offload_gb"] = 8 + else: + evidence["effective_offloader"] = "vllm.model_executor.offloader.uva.UVAOffloader" + + _edit_memory(tmp_path, edit) + with pytest.raises(ValueError, match="precision|async_scheduling|offload"): + _resolve(cell, tmp_path) + + +def test_modelopt_nvfp4_allows_bfloat16_attention_and_shared_experts(tmp_path, monkeypatch): + cell = _evidence(tmp_path, monkeypatch) + cell = replace(cell, gemm_quant_mode="bfloat16", weight_quantization="bfloat16", moe_quant_mode="nvfp4") + resources = FpmResourceProfile.model_validate(_resolve(cell, tmp_path)) + assert resources.runtime_memory.kv_cache_bytes == 88 * 128 + settings = json.loads(resources.runtime_memory.provenance)["runtime_settings"] + assert settings["model_config"]["dtype"] == "torch.bfloat16" + assert settings["quantization_config"]["quant_method"] == "NVFP4" + + +@pytest.mark.parametrize( + "change", ["missing", "missing_mode", "wrong_type", "wrong_max", "unsorted", "eager_with_graphs"] +) +def test_missing_or_inconsistent_graph_evidence_cannot_finalize(tmp_path, monkeypatch, change): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + config = evidence["resolved_config"] + graph = config["compilation_config"] + if change == "missing": + graph.clear() + elif change == "missing_mode": + graph.pop("mode") + elif change == "wrong_type": + graph["cudagraph_capture_sizes"] = [True, 64] + elif change == "wrong_max": + graph["max_cudagraph_capture_size"] = 32 + elif change == "unsorted": + graph["cudagraph_capture_sizes"] = [64, 32] + else: + config["model_config"]["enforce_eager"] = True + + _edit_memory(tmp_path, edit) + with pytest.raises(ValueError, match="graph configuration"): + _resolve(cell, tmp_path) + + +@pytest.mark.parametrize("change", ["graph", "weights", "dtype"]) +def test_scheduler_must_agree_with_worker_memory_settings(tmp_path, monkeypatch, change): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + if evidence["kind"] != "scheduler": + return + config = evidence["resolved_config"] + if change == "graph": + config["compilation_config"].update( + cudagraph_mode="NONE", cudagraph_capture_sizes=[], max_cudagraph_capture_size=0 + ) + elif change == "weights": + config["model_config"]["quantization"] = "fp8" + else: + config["model_config"]["dtype"] = "torch.float16" + + _edit_memory(tmp_path, edit) + with pytest.raises(ValueError, match="precision|between scheduler and worker"): + _resolve(cell, tmp_path) + + +@pytest.mark.parametrize("eager", [False, True]) +def test_audited_none_graph_configuration_is_valid(tmp_path, monkeypatch, eager): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + config = evidence["resolved_config"] + config["model_config"]["enforce_eager"] = eager + config["compilation_config"].update( + mode=0 if eager else 3, cudagraph_mode="NONE", cudagraph_capture_sizes=[], max_cudagraph_capture_size=0 + ) + + _edit_memory(tmp_path, edit) + result = _resolve(cell, tmp_path) + assert result["runtime_memory"]["kv_cache_bytes"] == 88 * 128 + + +def test_audited_worker_graph_downgrade_preserves_scheduler_snapshot(tmp_path, monkeypatch): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + if evidence["kind"] == "worker": + graph = evidence["resolved_config"]["compilation_config"] + evidence["initial_compilation_config"] = dict(graph) + graph["cudagraph_mode"] = "PIECEWISE" + + _edit_memory(tmp_path, edit) + result = _resolve(cell, tmp_path) + provenance = json.loads(result["runtime_memory"]["provenance"]) + assert provenance["runtime_settings"]["compilation_config"]["cudagraph_mode"] == "PIECEWISE" + + +@pytest.mark.parametrize("declared_mode", ["FULL_AND_PIECEWISE", "FULL_DECODE_ONLY"]) +def test_formal_memory_checks_reviewed_initial_graph_before_validated_worker_downgrade( + tmp_path, monkeypatch, declared_mode +): + cell = _evidence(tmp_path, monkeypatch) + + def downgrade(record): + if record["kind"] == "worker": + graph = record["resolved_config"]["compilation_config"] + record["initial_compilation_config"] = dict(graph) + graph["cudagraph_mode"] = "PIECEWISE" + + _edit_memory(tmp_path, downgrade) + expected = {"cudagraph_mode": declared_mode, "max_cudagraph_capture_size": 64} + if declared_mode == "FULL_AND_PIECEWISE": + result = _resolve(cell, tmp_path, expected_compilation_config=expected) + assert result["runtime_memory"]["kv_cache_bytes"] == 88 * 128 + else: + with pytest.raises(ValueError, match="reviewed serving role"): + _resolve(cell, tmp_path, expected_compilation_config=expected) + + +@pytest.mark.parametrize("resolved_mode", ["PIECEWISE", "NONE"]) +def test_full_graph_cannot_fall_back_to_unsupported_attention_mode(tmp_path, monkeypatch, resolved_mode): + cell = _evidence(tmp_path, monkeypatch) + + def edit(evidence): + graph = evidence["resolved_config"]["compilation_config"] + graph["cudagraph_mode"] = "FULL" + if evidence["kind"] == "worker": + evidence["initial_compilation_config"] = dict(graph) + graph["cudagraph_mode"] = resolved_mode + + _edit_memory(tmp_path, edit) + with pytest.raises(ValueError, match="configurations differ between scheduler and worker"): + _resolve(cell, tmp_path) + + +@pytest.mark.parametrize( + "change,match", + [ + ("missing_worker", "worker rank evidence is incomplete"), + ("missing_scheduler", "scheduler rank evidence is incomplete"), + ("duplicate_worker", "duplicate runtime memory worker rank"), + ("stale", "different plan, attempt"), + ("wrong_context", "launch mismatch"), + ("wrong_dtype", "KV precision"), + ("wrong_gpu_fraction", "launch mismatch"), + ("graph_mismatch", "differ across worker ranks"), + ("unknown_spec", "unsupported runtime cache spec"), + ("nonintegral", "non-integral pool page"), + ("watermark", "pool capacity/reservations"), + ("second_pool", "pool capacity/reservations"), + ("duplicate_storage", "storage accounting"), + ("wrong_retention", "geometry"), + ("timing_failure", "terminal result"), + ], +) +def test_runtime_memory_rejects_incomplete_stale_mixed_or_unsupported_evidence(tmp_path, monkeypatch, change, match): + cell = _evidence(tmp_path, monkeypatch) + path = tmp_path / "pod-0/fpm-memory-worker-dp0-tp0-pp0.json" + payload = json.loads(path.read_text()) + scheduler_path = tmp_path / "pod-0/fpm-memory-scheduler-dp0.json" + scheduler = json.loads(scheduler_path.read_text()) + if change == "missing_worker": + path.unlink() + elif change == "missing_scheduler": + scheduler_path.unlink() + elif change == "duplicate_worker": + (tmp_path / "pod-0/fpm-memory-duplicate.json").write_text(path.read_text()) + elif change == "stale": + payload["collector_provenance"]["attempt_id"] = "earlier-attempt" + elif change == "wrong_context": + payload["resolved_config"]["model_config"]["max_model_len"] = 2048 + elif change == "wrong_dtype": + payload["resolved_config"]["cache_config"]["cache_dtype"] = "fp8" + elif change == "wrong_gpu_fraction": + payload["resolved_config"]["cache_config"]["gpu_memory_utilization"] = 0.8 + elif change == "graph_mismatch": + payload["resolved_config"]["compilation_config"]["cudagraph_capture_sizes"] = [1, 2, 4] + payload["resolved_config"]["compilation_config"]["max_cudagraph_capture_size"] = 4 + elif change == "unknown_spec": + payload["cache"]["groups"][0]["spec_type"] = "vllm.v1.kv_cache_interface.MambaSpec" + scheduler["cache"]["groups"][0]["spec_type"] = payload["cache"]["groups"][0]["spec_type"] + elif change == "nonintegral": + payload["cache"]["allocated_cache_bytes"] -= 1 + payload["cache"]["storages"][0]["size_bytes"] -= 1 + elif change == "watermark": + scheduler["cache"]["watermark_blocks"] = 1 + elif change == "second_pool": + scheduler["cache"]["pool_count"] = 2 + elif change == "duplicate_storage": + payload["cache"]["storages"].append(payload["cache"]["storages"][0]) + elif change == "wrong_retention": + payload["cache"]["groups"][1]["sliding_window"] = 64 + elif change == "timing_failure": + native = tmp_path / "pod-0/benchmark.json" + value = json.loads(native.read_text()) + value["status"] = "failed" + native.write_text(json.dumps(value)) + if change != "missing_worker": + path.write_text(json.dumps(payload)) + if change != "missing_scheduler": + scheduler_path.write_text(json.dumps(scheduler)) + with pytest.raises(ValueError, match=match): + _resolve(cell, tmp_path) + + +@pytest.mark.parametrize("change", ["unknown_spec", "alias", "missing_storage", "watermark", "pool", "version"]) +def test_unsupported_observation_records_unresolved_without_running_or_replacing_runtime(tmp_path, monkeypatch, change): + worker, scheduler, config = _initialized() + monkeypatch.setattr(observer, "RESULTS_DIR", tmp_path) + monkeypatch.setattr( + observer.importlib.metadata, "version", lambda _name: "0.29.0" if change == "version" else "0.27.0" + ) + _write_provenance(tmp_path / "collector-provenance.json", cell_id="test") + kind, owner = "worker", worker + if change == "unknown_spec": + config.kv_cache_groups[0].kv_cache_spec = SimpleNamespace() + elif change == "alias": + config.kv_cache_tensors[0].block_stride = 128 + config.kv_cache_tensors[1].block_stride = 128 + elif change == "missing_storage": + worker.model_runner.kv_caches = [] + elif change == "watermark": + kind, owner = "scheduler", scheduler + scheduler.kv_cache_manager.watermark_blocks = 2 + elif change == "pool": + kind, owner = "scheduler", scheduler + scheduler.kv_cache_manager.coordinator.single_type_managers[0].block_pool = object() + observer.observe( + kind, + owner, + dp_rank=0, + tp_rank=0 if kind == "worker" else None, + pp_rank=0 if kind == "worker" else None, + cache_config=config, + ) + payload = json.loads(next(tmp_path.glob("fpm-memory-*.json")).read_text()) + assert payload["status"] == "unresolved" + assert payload["error"] + assert not list(tmp_path.glob("*.tmp")) + + +@pytest.mark.parametrize("resolved_mode", ["PIECEWISE", "NONE"]) +@pytest.mark.parametrize("initialization_count", [1, 2]) +def test_worker_profile_time_graph_downgrade_resolves(tmp_path, monkeypatch, resolved_mode, initialization_count): + cell = _evidence(tmp_path, monkeypatch) + + class Worker: + def determine_available_memory(self): + self.vllm_config.compilation_config.cudagraph_mode = resolved_mode + return self._fpm_available_cache_bytes + + def initialize_from_config(self, config): + assert config is self.model_runner.kv_cache_config + + def compile_or_warm_up_model(self): + return "compiled" + + tp_group = SimpleNamespace(rank_in_group=0) + modules = { + "vllm.distributed": { + "get_pp_group": lambda: SimpleNamespace(rank_in_group=0), + "get_tp_group": lambda: tp_group, + }, + "vllm.v1.worker.gpu_worker": {"Worker": Worker}, + } + monkeypatch.setitem(sys.modules, "fpm_memory_observer", observer) + for name, values in modules.items(): + module = ModuleType(name) + module.__dict__.update(values) + monkeypatch.setitem(sys.modules, name, module) + spec = importlib.util.spec_from_file_location( + "fpm_memory_worker", Path(observer.__file__).with_name("fpm_memory_worker.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + for dp in range(2): + pod = tmp_path / f"pod-{dp}" + monkeypatch.setattr(observer, "RESULTS_DIR", pod) + for tp in range(2): + (pod / f"fpm-memory-worker-dp{dp}-tp{tp}-pp0.json").unlink() + initialized, _, cache = _initialized(blocks=100 - dp * 10, free=98 - dp * 10) + worker = module.FpmResourceWorker() + worker.__dict__.update(vars(initialized)) + worker._fpm_cache_initialized = False + worker.vllm_config.parallel_config.data_parallel_rank = dp + tp_group.rank_in_group = tp + for _ in range(initialization_count): + assert worker.determine_available_memory() == (100 - dp * 10) * 128 + 200 + worker.initialize_from_config(cache) + assert worker.compile_or_warm_up_model() == "compiled" + + result = _resolve(cell, tmp_path) + assert result["runtime_memory"]["kv_cache_bytes"] == 88 * 128 + provenance = json.loads(result["runtime_memory"]["provenance"]) + assert provenance["runtime_settings"]["compilation_config"]["cudagraph_mode"] == resolved_mode + for artifact in provenance["artifacts"]: + evidence = artifact["evidence"] + if evidence["kind"] == "worker": + assert evidence["initial_compilation_config"]["cudagraph_mode"] == "FULL_AND_PIECEWISE" + + +def test_worker_and_scheduler_wrappers_delegate_before_observation(monkeypatch, caplog): + events = [] + + class Worker: + def determine_available_memory(self): + events.append("profile") + return 1234 + + def initialize_from_config(self, config): + events.append(("initialize", config)) + + def compile_or_warm_up_model(self): + events.append("warmup") + return "compiled" + + class Scheduler: + def __init__(self, vllm_config, kv_cache_config, *args, **kwargs): + events.append(("scheduler", vllm_config, kv_cache_config, args, kwargs)) + self._fpm_dp_rank = 3 + + modules = { + "fpm_memory_observer": { + "observe": lambda *args, **kwargs: events.append(("observe", args, kwargs)), + "observe_execution": lambda *args, **kwargs: events.append(("observe_execution", args, kwargs)), + "observe_cpu": lambda *args, **kwargs: events.append(("observe_cpu", args, kwargs)), + "compilation_config": lambda _config: {"cudagraph_mode": "FULL"}, + }, + "vllm.distributed": { + "get_pp_group": lambda: SimpleNamespace(rank_in_group=0), + "get_tp_group": lambda: SimpleNamespace(rank_in_group=1), + }, + "vllm.v1.worker.gpu_worker": {"Worker": Worker}, + "dynamo.vllm.instrumented_scheduler": {"InstrumentedScheduler": Scheduler}, + } + for name, values in modules.items(): + module = ModuleType(name) + module.__dict__.update(values) + monkeypatch.setitem(sys.modules, name, module) + + def load(filename): + path = Path(observer.__file__).with_name(filename + ".py") + spec = importlib.util.spec_from_file_location(filename, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + worker_module = load("fpm_memory_worker") + worker = worker_module.FpmResourceWorker() + worker.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_index=3, data_parallel_rank=0)) + assert worker.determine_available_memory() == 1234 + worker.initialize_from_config("cache") + assert worker.compile_or_warm_up_model() == "compiled" + assert events[:3] == ["profile", ("initialize", "cache"), "warmup"] + assert events[3][2] == {"dp_rank": 3, "tp_rank": 1, "pp_rank": 0} + assert events[3][0] == "observe_cpu" + assert events[4][0] == "observe_execution" + assert events[5][0] == "observe" + assert worker._fpm_available_cache_bytes == 1234 + assert worker._fpm_cache_initialized is True + load("fpm_memory_scheduler").FpmResourceInstrumentedScheduler("config", "cache", "manager", block_size=16) + assert events[6] == ("scheduler", "config", "cache", ("manager",), {"block_size": 16}) + assert events[7] == ("observe_cpu", ("scheduler",), {"dp_rank": 3}) + assert events[8][2] == {"dp_rank": 3, "cache_config": "cache"} + before = len(events) + load("fpm_memory_scheduler").FpmExecutionInstrumentedScheduler("config", "cache") + assert events[before:] == [ + ("scheduler", "config", "cache", (), {}), + ("observe_cpu", ("scheduler",), {"dp_rank": 3}), + ] + + def fail(_self, *_args): + raise RuntimeError("real runtime failed") + + for method, args in ( + ("determine_available_memory", ()), + ("initialize_from_config", ("cache",)), + ("compile_or_warm_up_model", ()), + ): + with monkeypatch.context() as patch: + patch.setattr(Worker, method, fail) + before = len(events) + with pytest.raises(RuntimeError, match="real runtime failed"): + getattr(worker, method)(*args) + assert len(events) == before + + def missing_graph(_config): + raise AttributeError("compilation configuration is unavailable") + + monkeypatch.setattr(worker_module, "compilation_config", missing_graph) + worker = worker_module.FpmResourceWorker() + worker.vllm_config = SimpleNamespace() + assert worker.determine_available_memory() == 1234 + assert worker._fpm_initial_compilation_config is None + assert "FPM initial graph configuration is unavailable" in caplog.text + worker.initialize_from_config("cache") + assert worker._fpm_initial_compilation_config is None + + +def test_ordinary_serving_worker_observes_initialized_runtime_without_dynamo(tmp_path, monkeypatch): + provenance = { + "schema": "aisimulate-serving-validation/v1", + "run_id": "serving-123", + "purpose": "matched_serving_accuracy", + } + provenance_path = tmp_path / "serving-provenance.json" + provenance_path.write_text(json.dumps(provenance)) + output = tmp_path / "observations" + monkeypatch.setenv("FPM_EXECUTION_OUTPUT_DIR", str(output)) + monkeypatch.setenv("FPM_EXECUTION_PROVENANCE_FILE", str(provenance_path)) + monkeypatch.setattr(observer.importlib.metadata, "version", lambda _name: "0.28.0") + + class Worker: + def compile_or_warm_up_model(self): + self.vllm_config.compilation_config.cudagraph_mode = "PIECEWISE" + self.model_runner.attn_groups = [ + [ + SimpleNamespace( + backend=type("ResolvedBackend", (), {"__module__": "actual.runtime"}), + layer_names=["layer.0"], + kv_cache_group_id=0, + ) + ] + ] + return "compiled" + + monkeypatch.setitem(sys.modules, "fpm_memory_observer", observer) + monkeypatch.setitem(sys.modules, "dynamo.vllm.instrumented_scheduler", None) + for name, values in { + "vllm.distributed": { + "get_pp_group": lambda: SimpleNamespace(rank_in_group=0), + "get_tp_group": lambda: SimpleNamespace(rank_in_group=0), + }, + "vllm.v1.worker.gpu_worker": {"Worker": Worker}, + }.items(): + module = ModuleType(name) + module.__dict__.update(values) + monkeypatch.setitem(sys.modules, name, module) + spec = importlib.util.spec_from_file_location( + "fpm_memory_worker", Path(observer.__file__).with_name("fpm_memory_worker.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + worker = module.FpmExecutionWorker() + worker.vllm_config = _vllm_config(tp=1, dp=1) + worker.vllm_config.parallel_config.data_parallel_rank = 0 + worker.model_runner = SimpleNamespace(attn_groups=[]) + assert worker.compile_or_warm_up_model() == "compiled" + sidecar = output / "fpm-execution-worker-dp0-tp0-pp0.json" + result = json.loads(sidecar.read_text()) + assert result["status"] == "observed" + assert result["execution_provenance"] == provenance + assert result["provenance_source"]["path"] == str(provenance_path) + assert "collector_provenance" not in result + assert result["attention_groups"][0]["backend_class"] == "actual.runtime.ResolvedBackend" + assert result["graph_config"]["cudagraph_mode"] == "PIECEWISE" + assert not hasattr(worker, "_fpm_available_cache_bytes") + assert not list(output.glob("fpm-memory-*.json")) + original = sidecar.read_bytes() + with pytest.raises(RuntimeError, match="duplicate"): + worker.compile_or_warm_up_model() + assert sidecar.read_bytes() == original + # A second rank can share this run, but a different run cannot reuse it. + observer.observe_execution(worker, dp_rank=0, tp_rank=1, pp_rank=0) + provenance["run_id"] = "serving-456" + provenance_path.write_text(json.dumps(provenance)) + with pytest.raises(ValueError, match="different serving run"): + observer.observe_execution(worker, dp_rank=0, tp_rank=2, pp_rank=0) + + +@pytest.mark.parametrize( + "output,provenance", + [ + ("absolute", None), + (None, "absolute"), + ("relative", "absolute"), + ("absolute", "relative"), + ("/bad\npath", "absolute"), + ], +) +def test_serving_execution_observation_requires_paired_absolute_paths(tmp_path, monkeypatch, output, provenance): + for name, value in (("FPM_EXECUTION_OUTPUT_DIR", output), ("FPM_EXECUTION_PROVENANCE_FILE", provenance)): + monkeypatch.delenv(name, raising=False) + if value is not None: + monkeypatch.setenv(name, str(tmp_path) if value == "absolute" else value) + with pytest.raises(ValueError, match="supplied together|absolute path"): + observer.observe_execution(SimpleNamespace(), dp_rank=0, tp_rank=0, pp_rank=0) + + +@pytest.mark.parametrize("failure", ["purpose", "run_id", "collector", "orphan"]) +def test_serving_execution_observation_rejects_missing_or_reused_identity(tmp_path, monkeypatch, failure): + output = tmp_path / "observations" + output.mkdir() + provenance = {"run_id": "unique-run", "purpose": "matched_serving_accuracy"} + if failure in {"purpose", "run_id"}: + provenance.pop(failure) + elif failure == "collector": + (output / "collector-provenance.json").write_text("{}") + else: + (output / "fpm-execution-worker-dp0-tp0-pp0.json").write_text("{}") + provenance_path = tmp_path / "serving-provenance.json" + provenance_path.write_text(json.dumps(provenance)) + monkeypatch.setenv("FPM_EXECUTION_OUTPUT_DIR", str(output)) + monkeypatch.setenv("FPM_EXECUTION_PROVENANCE_FILE", str(provenance_path)) + with pytest.raises(ValueError, match="requires a serving run_id|fresh directory"): + observer.observe_execution(SimpleNamespace(), dp_rank=0, tp_rank=0, pp_rank=0) + + +def _pending_plan(version="0.27.0"): + profile = _profile() + for deployment in profile["deployments"]: + deployment["backend_version"] = version + for key in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + deployment["resources"].pop(key) + return _plan(profile) + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_pending_memory_plans_without_model_and_renders_audited_hooks(tmp_path): + plan = _pending_plan() + assert all(item.disposition == "unknown" for item in plan.topology_memory_admission) + assert all( + item.estimated_non_kv_bytes is None + for decision in plan.topology_memory_admission + for item in decision.estimates + ) + assert plan.to_dict()["runtime_memory_policy"]["async_scheduling"] is False + runtime_memory.validate_saved_plan(plan.to_dict()) + for cell in plan.cells: + output = tmp_path / cell.cell_id + output.mkdir() + runner._render_cell(plan, cell, output, {}) + script = (output / "run.sh").read_text() + assert "--worker-cls fpm_memory_worker.FpmResourceWorker" in script + assert "--scheduler-cls fpm_memory_scheduler.FpmResourceInstrumentedScheduler" in script + assert script.count("--no-async-scheduling") == 1 + + +@pytest.mark.parametrize( + "flag", + ["--worker-cls", "--scheduler-cls", "--kv-cache-memory-bytes", "--num-gpu-blocks-override", "--async-scheduling"], +) +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_pending_memory_rejects_conflicting_runtime_policy(flag): + plan = _pending_plan() + cell = plan.cells[0] + cell = replace( + cell, backend_policy=planner.BackendPolicy("conflict", {"params": {"agg": {"extra_cli_args": [flag]}}}, {}) + ) + with pytest.raises(ValueError, match="memory"): + runner._cell_generator_overrides(plan, cell, {}) + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_saved_plan_identity_reads_immutable_facts_without_model_resolution(): + plan = _plan(_profile()) + saved = plan.to_dict() + identity = runtime_memory.saved_plan_identity(saved) + assert identity.model_path == plan.model_path + assert identity.capability.aic_database_version == plan.capability.aic_database_version + assert identity.options.warmup_iterations == plan.options.warmup_iterations + assert runtime_memory.cell_from_dict(saved["cells"][0]) == plan.cells[0] + saved["options"]["global_warmup_iterations"] += 1 + with pytest.raises(ValueError, match="SHA-256"): + runtime_memory.saved_plan_identity(saved) + + +def test_saved_cell_rejects_lossy_extra_fields(): + payload = _cell().to_dict() + payload["unrecognized"] = True + with pytest.raises(ValueError, match="contract"): + runtime_memory.cell_from_dict(payload) + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_unaudited_runtime_keeps_native_timing_launch_and_pending_memory(): + plan = _pending_plan("0.25.1") + assert plan.to_dict()["runtime_memory_policy"]["observation"] == "unavailable_for_runtime" + for cell in plan.cells: + args = runner._cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + assert "--worker-cls" not in args + assert "--scheduler-cls" not in args + assert ("--no-async-scheduling" in args) == (cell.workload_kind == "prefill") + assert runner._observe_runtime_memory(plan, cell) is False + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_historical_v10_plan_keeps_hash_and_reaches_native_aggregation(tmp_path): + from collector.fpm_forward import database + + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY + + # Captured from the actual schema-10 producer at 4d702ff6b756b21e74b76f30273e9828fdde6861. + # Absolute source paths are immutable provenance; this reader must not open them. + fixture = Path(__file__).parent / "fixtures/fpm_collection_plan_v10.json" + original = fixture.read_bytes() + payload = json.loads(original) + frozen = runtime_memory.saved_plan_identity(payload) + assert frozen.sha256 == "ea99845deee06804fd35a7a0999b0118446df1ac8ededbf7c813345360ba6bf7" + assert [cell.cell_id for cell in frozen.cells] == ["fpm-d111aeb8823a9bee", "fpm-66274b2ba8ae195c"] + assert frozen.options.benchmark_points_json is None + rows = [] + for cell in frozen.cells: + unit = tmp_path / cell.cell_id / "raw/node0000" + unit.mkdir(parents=True) + provenance_path = unit / "collector-provenance.json" + _write_provenance(provenance_path, cell_id=cell.cell_id, plan_sha256=frozen.sha256, attempt_id="historical") + provenance = json.loads(provenance_path.read_text()) + provenance["runtime"]["backend_version"] = "0.27.0" + provenance_path.write_text(json.dumps(provenance)) + (unit / "benchmark.json").write_text(json.dumps(_native_payload(phase=cell.workload_kind, rank=0, dp=1))) + rows.extend(database.aggregate_cell(frozen, cell, tmp_path / cell.cell_id, expected_attempt_id="historical")) + assert len(rows) == 2 + for row in rows: + assert row["source_plan_sha256"] == frozen.sha256 + assert tuple(row[name] for name in EXECUTION_COLUMNS) == LEGACY_EXECUTION_IDENTITY + assert fixture.read_bytes() == original + + +@pytest.mark.parametrize("schema_version", [10, 11]) +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_saved_plan_rejects_cell_identity_from_other_schema(schema_version): + payload = _pending_plan().to_dict() + payload["schema_version"] = schema_version + if schema_version == 11: + for cell in payload["cells"]: + cell.pop("execution_identity") + cell.pop("input_text_sha256") + with pytest.raises(ValueError, match="execution identity does not match its schema"): + runtime_memory.validate_saved_plan(payload) + + +def test_saved_cell_preserves_complete_nonlegacy_execution_identity(): + payload = _cell().to_dict() + payload["execution_identity"] = { + "model_config_sha256": "a" * 64, + "execution_profile": "decoder_replay", + "engram_residency": "hbm_tp_sharded", + "input_modality": "text", + } + payload["input_text_sha256"] = "b" * 64 + assert runtime_memory.cell_from_dict(payload).to_dict() == payload + payload["execution_identity"].pop("engram_residency") + with pytest.raises(ValueError, match="invalid execution identity"): + runtime_memory.cell_from_dict(payload) + + +def test_saved_cell_rejects_partial_execution_identity(): + payload = _cell().to_dict() + payload.pop("input_text_sha256") + with pytest.raises(ValueError, match="incomplete execution identity"): + runtime_memory.cell_from_dict(payload) + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_historical_formal_publication_reports_unsupported_migration_without_changes(tmp_path): + from collector.fpm_forward import database + + parquet = tmp_path / "fpm_forward_perf.parquet" + metadata = tmp_path / "fpm_forward_perf.metadata.json" + parquet.write_bytes(b"historical sealed data stays unchanged") + metadata.write_text(json.dumps({"schema_name": "aic_fpm_forward_perf", "schema_version": 6})) + original = parquet.read_bytes(), metadata.read_bytes() + with pytest.raises(ValueError, match="historical schema-6.*Automatic migration is unsupported"): + database.validate_formal_database_commit(parquet, metadata, _pending_plan()) + assert (parquet.read_bytes(), metadata.read_bytes()) == original diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runtime_probe.py b/python/aisimulate/tests/unit/collector/test_fpm_runtime_probe.py new file mode 100644 index 000000000..eb475a57d --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_runtime_probe.py @@ -0,0 +1,1149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Launch-only probes use real rendering and synthetic executor evidence.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest + +from .test_fpm_profile_collection import no_models_or_timing_data # noqa: F401 + +pytestmark = [pytest.mark.unit, pytest.mark.usefixtures("no_models_or_timing_data")] + + +def _inputs(tmp_path, *, graph_policy="runtime"): + config = tmp_path / "config.json" + config.write_text( + json.dumps( + { + "architectures": ["ExternalMoeForCausalLM"], + "hidden_size": 128, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "num_hidden_layers": 5, + "max_position_embeddings": 8192, + "vocab_size": 1024, + "num_experts": 8, + "torch_dtype": "bfloat16", + } + ) + ) + source = tmp_path / "observer.py" + marker = tmp_path / "observer-was-imported" + source.write_text(f"from pathlib import Path\nPath({str(marker)!r}).touch()\nraise RuntimeError('never import')\n") + (tmp_path / "source-notes.txt").write_text("Synthetic source mapping fixture; no live verification.") + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps( + { + "schema_version": "aisimulate-runtime-instrumentation/v1", + "runtime": { + "framework": "vllm", + "version": "0.28.0", + "source_revision": "a" * 40, + "source_files": {"vllm/v1/core/block_pool.py": "b" * 64}, + }, + "files": ["observer.py", "source-notes.txt"], + "worker_class": "observer.ObservedWorker", + "scheduler_class": "observer.ObservedScheduler", + "observation_schema": "aisimulate-runtime-observation/v1", + "source_notes": "source-notes.txt", + } + ) + ) + launch = { + "identity": { + "model": "example/external-moe", + "model_revision": "pinned-model-revision", + "model_kind": "moe", + "framework": "vllm", + "framework_version": "0.28.0", + "gpu": "gb200", + "interconnect": "nvlink", + }, + "topology": {"tp": 4, "pp": 1, "dp": 1, "moe_tp": 1, "moe_ep": 4, "cp": 1}, + "precision": { + "gemm_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16", + "fmha_quant_mode": "bfloat16", + "kvcache_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "moe_backend": "auto", + "attention_backend": "auto", + "enable_wideep": False, + "enable_eplb": False, + }, + "collection": { + "max_model_len": 4096, + "max_num_batched_tokens": 1024, + "max_num_seqs": 64, + "gpu_memory_utilization": 0.85, + "prefill_cudagraph_policy": graph_policy, + "max_prefill_cudagraph_size": 512 if graph_policy == "explicit" else None, + }, + "model_config": {"path": str(config), "sha256": hashlib.sha256(config.read_bytes()).hexdigest()}, + "deployment": { + "executor": "slurm", + "image": "image@sha256:synthetic", + "container_mount": ["/cache:/cache"], + "cpus_per_task": 16, + "cpu_bind": "cores", + }, + } + return launch, manifest, marker + + +class _SyntheticExecutor: + def __init__(self, plan, cell, cell_dir, executions, *, fail_configuration=None): + self.plan, self.cell, self.cell_dir = plan, cell, cell_dir + self.executions = executions + self.fail_configuration = fail_configuration + + def cleanup(self): + pass + + def apply(self): + pass + + def wait_ready(self, expected_nodes): + return ["synthetic-node", *(f"synthetic-node-{index}" for index in range(1, expected_nodes))] + + def stage(self, pods, files): + self.files = {path.name: path for path in files} + assert "runtime-instrumentation.zip" in self.files + assert "runtime-probe-context.json" in self.files + + def prepare_attempt(self, pods, **kwargs): + self.identity = kwargs + + def execute(self, pods): + context = json.loads(self.files["runtime-probe-context.json"].read_text()) + self.executions.append((context["configuration"], context["phase"])) + for pod in pods: + raw = self.cell_dir / "raw" / pod + raw.mkdir(parents=True, exist_ok=True) + (raw / "synthetic-observation.json").write_text( + json.dumps({**context, "schema_version": "aisimulate-runtime-observation/v1", "synthetic": True}) + ) + if context["configuration"] == self.fail_configuration: + raise RuntimeError("synthetic runtime failure") + + def collect(self, pods, *, require_benchmark=True): + pass + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) +@pytest.mark.parametrize("graph_policy", ["runtime", "explicit"]) +def test_serving_role_probe_preserves_engine_settings_and_only_exercises_its_phases( + tmp_path, monkeypatch, worker_type, graph_policy +): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path, graph_policy=graph_policy) + launch["worker_type"] = worker_type + mode = {"prefill": "PIECEWISE", "decode": "FULL_DECODE_ONLY", "aggregated": "FULL_AND_PIECEWISE"}[worker_type] + launch["collection"]["cudagraph_mode"] = mode + if graph_policy == "runtime": + launch["collection"]["cudagraph_capture_sizes"] = [1, 2, 4, 8, 16, 32, 64] + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: pytest.fail("preview launched an executor")) + result = runtime_probe.probe_runtime({"selected": launch}, instrumentation=manifest, output_dir=tmp_path / "probe") + assert result["status"] == "preview", result + expected = {"prefill", "decode"} if worker_type == "aggregated" else {worker_type} + phases = result["configurations"]["selected"]["phases"] + assert set(phases) == expected + configs = [] + for phase in phases.values(): + request = json.loads(Path(phase["generator_request"]).read_text()) + args = request["params"]["agg"]["extra_cli_args"] + configs.append(json.loads(args[args.index("--compilation-config") + 1])) + assert args[args.index("--max-num-batched-tokens") + 1] == "1024" + assert args[args.index("--max-num-seqs") + 1] == "64" + assert all(config == configs[0] for config in configs) + assert configs[0]["cudagraph_mode"] == mode + assert configs[0]["max_cudagraph_capture_size"] == (512 if graph_policy == "explicit" else 64) + + +def test_disaggregated_probe_failure_keeps_other_role_resumable(tmp_path, monkeypatch): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + configurations = {role: {**copy.deepcopy(launch), "worker_type": role} for role in ("prefill", "decode")} + executions = [] + monkeypatch.setattr( + runner, + "_cell_runner", + lambda plan, cell, _manifest, directory: _SyntheticExecutor( + plan, cell, directory, executions, fail_configuration="prefill" + ), + ) + output = tmp_path / "probe" + result = runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output, execute=True) + assert result["status"] == "partial" + assert result["configurations"]["prefill"]["status"] == "failed" + assert result["configurations"]["decode"]["status"] == "completed" + assert executions == [("prefill", "prefill"), ("decode", "decode")] + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: pytest.fail("completed role launched again")) + resumed = runtime_probe.probe_runtime( + {"decode": configurations["decode"]}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + assert resumed["configurations"]["decode"]["resumed"] is True + + +@pytest.mark.parametrize( + "changes", + [ + {"cudagraph_capture_sizes": [8, 4]}, + {"cudagraph_capture_sizes": [4, 8], "max_cudagraph_capture_size": 16}, + {"max_cudagraph_capture_size": 2048}, + {"cudagraph_mode": "NONE", "cudagraph_capture_sizes": [4]}, + {"cudagraph_mode": "UNKNOWN"}, + ], +) +def test_invalid_role_graph_settings_fail_before_execution_without_blocking_other_role(tmp_path, monkeypatch, changes): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + invalid = {**copy.deepcopy(launch), "worker_type": "prefill"} + invalid["collection"].update(changes) + valid = {**copy.deepcopy(launch), "worker_type": "decode"} + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: pytest.fail("preview launched an executor")) + result = runtime_probe.probe_runtime( + {"prefill": invalid, "decode": valid}, instrumentation=manifest, output_dir=tmp_path / "probe" + ) + assert result["status"] == "partial" + assert result["configurations"]["prefill"]["status"] == "failed" + assert result["configurations"]["decode"]["status"] == "preview" + + +@pytest.mark.parametrize("graph_policy", ["runtime", "explicit"]) +def test_preview_requires_no_cache_geometry_and_never_imports_bundle(tmp_path, monkeypatch, graph_policy): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, marker = _inputs(tmp_path, graph_policy=graph_policy) + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: pytest.fail("preview launched an executor")) + result = runtime_probe.probe_runtime( + {"arbitrary-label": launch}, instrumentation=manifest, output_dir=tmp_path / "probe" + ) + + assert result["status"] == "preview" + assert not marker.exists() + preview = result["configurations"]["arbitrary-label"] + assert preview["minimum_gpus"] == 4 + assert set(preview["phases"]) == {"prefill", "decode"} + for phase, info in preview["phases"].items(): + params = json.loads(Path(info["generator_request"]).read_text()) + args = params["params"]["agg"]["extra_cli_args"] + assert args[args.index("--max-model-len") + 1] == "4096" + assert args[args.index("--max-num-batched-tokens") + 1] == "1024" + assert args[args.index("--max-num-seqs") + 1] == "64" + assert params["params"]["agg"]["kv_cache_free_gpu_memory_fraction"] == 0.85 + assert args[args.index("--revision") + 1] == launch["identity"]["model_revision"] + assert args.count("--revision") == 1 + assert args[args.index("--worker-cls") + 1] == "fpm_runtime_instrumentation.ObservedWorker" + assert "--no-async-scheduling" in args + batch_flag = "--prefix-max-batch-size-samples" if phase == "prefill" else "--decode-max-batch-size-samples" + context = json.loads((tmp_path / "probe" / info["launch_manifest"]["path"]).read_text()) + assert int(args[args.index(batch_flag) + 1]) == context["sampling"]["max_batch_size_samples"] == 2 + if phase == "prefill": + assert ("--compilation-config" in args) == (graph_policy == "explicit") + assert args[args.index("--prefill-max-new-token-samples") + 1] == "2" + if graph_policy == "explicit": + graph = json.loads(args[args.index("--compilation-config") + 1]) + assert graph["max_cudagraph_capture_size"] == 512 + + +@pytest.mark.parametrize("executor", ["kubernetes", "slurm"]) +def test_multinode_preview_labels_every_resource(tmp_path, monkeypatch, executor): + from collector.fpm_forward import runner, runtime_probe + from collector.fpm_forward.runtime_instrumentation import load_instrumentation + + launch, manifest, marker = _inputs(tmp_path) + launch["identity"]["gpu"] = "gb300" + if executor == "kubernetes": + launch["deployment"] = {"executor": executor, "image": launch["deployment"]["image"]} + launch["topology"].update(tp=8, dp=1, moe_tp=8, moe_ep=1) + dep = copy.deepcopy(launch) + dep["topology"].update(tp=1, dp=8, moe_tp=1, moe_ep=8) + configurations = {"tp8": launch, "dep8": dep} + monkeypatch.setattr(runner, "_cell_runner", lambda *_args: pytest.fail("preview launched an executor")) + output = tmp_path / "probe" + + result = runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output) + + assert result["status"] == "preview", result + assert not marker.exists() + bundle = load_instrumentation(manifest) + for name, facts in configurations.items(): + plan = runtime_probe.build_runtime_probe_plan(name, facts, bundle) + preview = result["configurations"][name] + assert preview["minimum_gpus"] == 8 + assert set(preview["phases"]) == {"prefill", "decode"} + for cell in plan.cells: + phase = preview["phases"][cell.workload_kind] + resource = output / phase["resource_manifest"]["path"] + documents = runner._manifest_documents(resource) + assert [document["kind"] for document in documents] == ["ComputeDomain", "LeaderWorkerSet"] + assert runner._expected_nodes(resource) == 2 + expected = { + "aiconfigurator.nvidia.com/owned-by": "fpm-forward-collector", + "aiconfigurator.nvidia.com/plan": plan.sha256[:16], + runner.FPM_CELL_LABEL: cell.cell_id, + } + for document in documents: + assert expected.items() <= document["metadata"]["labels"].items() + assert documents[0]["spec"]["numNodes"] == 0 + + +@pytest.mark.parametrize("flag", [["--revision", "another-model-revision"], ["--revision=another-model-revision"]]) +def test_runtime_revision_cannot_be_replaced_by_backend_policy(tmp_path, flag): + from dataclasses import replace + + from collector.fpm_forward import runner, runtime_probe + from collector.fpm_forward.planner import BackendPolicy + from collector.fpm_forward.runtime_instrumentation import load_instrumentation + + launch, manifest, _marker = _inputs(tmp_path) + plan = runtime_probe.build_runtime_probe_plan("worker", launch, load_instrumentation(manifest)) + cell = replace( + plan.cells[0], backend_policy=BackendPolicy("test", {"params": {"agg": {"extra_cli_args": flag}}}, {}) + ) + with pytest.raises(ValueError, match="selected model revision"): + runner._cell_generator_overrides(plan, cell, runtime_probe.probe_generator_overrides(plan.launch)) + + +def test_execute_covers_each_configuration_and_both_phases_then_resumes(tmp_path, monkeypatch): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, marker = _inputs(tmp_path) + # Fresh callers may omit the default CPU proposal. Preview freezes it, + # and resume must recover that same explicit identity from the saved plan. + launch["deployment"].pop("cpus_per_task") + launch["deployment"].pop("cpu_bind") + second = copy.deepcopy(launch) + second["topology"].update(tp=1, dp=4) + configurations = {"tp4": launch, "dep4": second} + executions = [] + monkeypatch.setattr( + runner, + "_cell_runner", + lambda plan, cell, _manifest, directory: _SyntheticExecutor(plan, cell, directory, executions), + ) + output = tmp_path / "probe" + runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output) + result = runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output, execute=True) + assert result["status"] == "completed" + assert executions == [(label, phase) for label in configurations for phase in ("prefill", "decode")] + assert not marker.exists() + index = json.loads((output / "observations.json").read_text()) + original = { + str(output / artifact["path"]): (output / artifact["path"]).read_bytes() + for entry in index["configurations"].values() + for phase in entry["attempts"][0]["phases"].values() + for artifact in phase["artifacts"] + if artifact["kind"] == "observation" + } + assert len(original) == 4 + for entry in index["configurations"].values(): + assert len(entry["attempts"]) == 1 + assert set(entry["attempts"][0]["phases"]) == {"prefill", "decode"} + runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output, execute=True, resume=True) + assert len(executions) == 4 + assert all(Path(path).read_bytes() == contents for path, contents in original.items()) + # Resume rechecks collector-owned launch bytes as well as observations. + altered = output / index["configurations"]["tp4"]["attempts"][0]["phases"]["prefill"]["artifacts"][0]["path"] + altered.write_bytes(altered.read_bytes() + b"\n") + runtime_probe.probe_runtime(configurations, instrumentation=manifest, output_dir=output, execute=True, resume=True) + assert len(executions) == 6 + updated = json.loads((output / "observations.json").read_text()) + assert len(updated["configurations"]["tp4"]["attempts"]) == 2 + assert len(updated["configurations"]["dep4"]["attempts"]) == 1 + + +def test_configuration_failure_preserves_successes_and_new_bundle_gets_new_attempt(tmp_path, monkeypatch): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _marker = _inputs(tmp_path) + second = copy.deepcopy(launch) + second["topology"].update(tp=1, dp=4) + executions = [] + monkeypatch.setattr( + runner, "_salvage_artifacts", lambda resource, _cell: resource.collect([], require_benchmark=False) + ) + monkeypatch.setattr( + runner, + "_cell_runner", + lambda plan, cell, _manifest, directory: _SyntheticExecutor( + plan, cell, directory, executions, fail_configuration="broken" + ), + ) + output = tmp_path / "probe" + result = runtime_probe.probe_runtime( + {"broken": launch, "good": second}, instrumentation=manifest, output_dir=output, execute=True + ) + assert result["status"] == "partial" + assert len(executions) == 4 + index = json.loads((output / "observations.json").read_text()) + first = index["configurations"]["good"]["attempts"][0] + preserved = (output / first["bundle"]["manifest"]).read_bytes() + (tmp_path / "observer.py").write_text("# revised synthetic observer\n") + runtime_probe.probe_runtime( + {"good": second}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + updated = json.loads((output / "observations.json").read_text()) + attempts = updated["configurations"]["good"]["attempts"] + assert len(attempts) == 2 + assert attempts[0] == first + assert attempts[0]["bundle"]["sha256"] != attempts[1]["bundle"]["sha256"] + assert (output / first["bundle"]["manifest"]).read_bytes() == preserved + + +@pytest.mark.parametrize( + "provided", [{"cpus_per_task": 32}, {"cpu_bind": "none"}, {"cpus_per_task": 8}, {"cpu_bind": "cores"}] +) +def test_probe_partial_resume_keeps_saved_cpu_policy_or_rejects_changes(tmp_path, monkeypatch, provided): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + launch["deployment"].update(cpus_per_task=32, cpu_bind="none") + executions = [] + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: _SyntheticExecutor(p, c, d, executions)) + output = tmp_path / "probe" + assert ( + runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True)[ + "status" + ] + == "completed" + ) + index_path = output / "observations.json" + before = index_path.read_bytes() + launch["deployment"].pop("cpus_per_task") + launch["deployment"].pop("cpu_bind") + launch["deployment"].update(provided) + result = runtime_probe.probe_runtime( + {"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + same = provided in ({"cpus_per_task": 32}, {"cpu_bind": "none"}) + assert result["status"] == ("completed" if same else "failed"), result + if same: + assert result["configurations"]["tp4"]["resumed"] is True + assert result["configurations"]["tp4"]["launch"]["deployment"]["cpus_per_task"] == 32 + assert result["configurations"]["tp4"]["launch"]["deployment"]["cpu_bind"] == "none" + else: + assert "launch settings changed" in str(result) + assert index_path.read_bytes() == before + assert len(executions) == 2 + + +@pytest.mark.parametrize("tamper", ["hash", "launch", "missing"]) +def test_probe_cpu_policy_inheritance_requires_verified_saved_context(tmp_path, monkeypatch, tamper): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + launch["deployment"].update(cpus_per_task=32, cpu_bind="none") + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: _SyntheticExecutor(p, c, d, [])) + output = tmp_path / "probe" + runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True) + index_path = output / "observations.json" + index = json.loads(index_path.read_text()) + entry = index["configurations"]["tp4"] + if tamper == "missing": + entry["attempts"][0]["phases"] = {} + else: + reference = entry["attempts"][0]["phases"]["prefill"]["launch_manifest"] + path = output / reference["path"] + context = json.loads(path.read_text()) + context["launch"]["deployment"]["cpu_bind"] = "cores" + runner._atomic_json(path, context) + if tamper == "launch": + reference.update(runner._file_metadata(path)) + runner._atomic_json(index_path, index) + before = index_path.read_bytes() + launch["deployment"].pop("cpu_bind") + monkeypatch.setattr(runner, "_cell_runner", lambda *args: pytest.fail("invalid CPU evidence reached an executor")) + result = runtime_probe.probe_runtime( + {"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + assert result["status"] == "failed", result + assert index_path.read_bytes() == before + + +def test_resume_without_an_archived_probe_uses_fresh_cpu_defaults(tmp_path, monkeypatch): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + launch["deployment"].pop("cpus_per_task") + launch["deployment"].pop("cpu_bind") + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: _SyntheticExecutor(p, c, d, [])) + result = runtime_probe.probe_runtime( + {"tp4": launch}, instrumentation=manifest, output_dir=tmp_path / "probe", execute=True, resume=True + ) + assert result["status"] == "completed", result + assert result["configurations"]["tp4"]["launch"]["deployment"]["cpus_per_task"] == 16 + assert result["configurations"]["tp4"]["launch"]["deployment"]["cpu_bind"] == "cores" + + +@pytest.mark.parametrize("completed", [True, False]) +def test_legacy_probe_reuse_is_read_only_and_new_execution_does_not_replace_attempt(tmp_path, monkeypatch, completed): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _ = _inputs(tmp_path) + executions = [] + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: _SyntheticExecutor(p, c, d, executions)) + output = tmp_path / "probe" + runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True) + # Build an explicitly historical fixture, retaining the old launch representation + # in its SHA-bound context. No CPU observations or policy are invented for it. + index_path = output / "observations.json" + index = json.loads(index_path.read_text()) + entry = index["configurations"]["tp4"] + for field in ("cpus_per_task", "cpu_bind"): + launch["deployment"].pop(field) + entry["launch"]["deployment"].pop(field) + for phase in entry["attempts"][0]["phases"].values(): + reference = phase["launch_manifest"] + path = output / reference["path"] + context = json.loads(path.read_text()) + context["launch"] = entry["launch"] + runner._atomic_json(path, context) + reference.update(runner._file_metadata(path)) + if not completed: + phase["status"] = "failed" + if not completed: + entry["attempts"][0]["status"] = "failed" + runner._atomic_json(index_path, index) + before = index_path.read_bytes() + attempts = sorted((output / "attempts").rglob("runtime-probe-context.json")) + monkeypatch.setattr(runner, "_cell_runner", lambda *args: pytest.fail("legacy resume reached an executor")) + result = runtime_probe.probe_runtime( + {"tp4": launch}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + assert result["status"] == ("completed" if completed else "failed"), result + if not completed: + assert "no frozen CPU policy" in str(result) + assert index_path.read_bytes() == before + assert sorted((output / "attempts").rglob("runtime-probe-context.json")) == attempts + + +def test_dense_tensor_parallel_probe_keeps_dense_moe_axes(tmp_path): + from collector.fpm_forward import runtime_probe + + launch, manifest, _marker = _inputs(tmp_path) + launch["identity"]["model_kind"] = "dense" + launch["topology"].update(moe_ep=1) + config = Path(launch["model_config"]["path"]) + payload = json.loads(config.read_text()) + payload.pop("num_experts") + config.write_text(json.dumps(payload)) + launch["model_config"]["sha256"] = hashlib.sha256(config.read_bytes()).hexdigest() + result = runtime_probe.probe_runtime({"dense": launch}, instrumentation=manifest, output_dir=tmp_path / "probe") + assert result["status"] == "preview" + assert result["configurations"]["dense"]["launch"]["topology"] == launch["topology"] + + +def test_probe_rejects_ignored_checkpoint_precision_and_resolves_explicit_graph_default(tmp_path): + from collector.fpm_forward import runtime_probe + + launch, manifest, _marker = _inputs(tmp_path, graph_policy="explicit") + launch["collection"]["max_prefill_cudagraph_size"] = None + wrong = copy.deepcopy(launch) + wrong["precision"]["gemm_quant_mode"] = "nvfp4" + result = runtime_probe.probe_runtime( + {"valid": launch, "wrong": wrong}, instrumentation=manifest, output_dir=tmp_path / "probe" + ) + assert result["status"] == "partial" + assert "checkpoint-native" in result["configurations"]["wrong"]["diagnostics"][0] + assert result["configurations"]["valid"]["launch"]["collection"]["max_prefill_cudagraph_size"] == 2048 + + +def _formal_cli_inputs(tmp_path): + from .test_fpm_profile_collection import _profile + + launch, manifest, _marker = _inputs(tmp_path) + profile = _profile() + profile.update( + model=launch["identity"]["model"], + model_revision=launch["identity"]["model_revision"], + architecture="ExternalMoeForCausalLM", + context_length=4096, + num_experts=8, + ) + profile["deployments"] = [profile["deployments"][0]] + deployment = profile["deployments"][0] + deployment.update(system="gb200", backend_version="0.28.0", kv_cache_dtype="bfloat16", **launch["topology"]) + deployment.update({key: value for key, value in launch["precision"].items() if key != "kvcache_quant_mode"}) + deployment["resources"].update(max_num_tokens=1024, max_batch_size=64) + profile_path = tmp_path / "profile.json" + profile_path.write_text(json.dumps(profile)) + launch_path = tmp_path / "launch.json" + launch_path.write_text(json.dumps(launch)) + argv = [ + "--model-path", + launch["identity"]["model"], + "--gpu", + "gb200", + "--fpm-model-profile", + str(profile_path), + "--fpm-model-config", + launch["model_config"]["path"], + "--fpm-max-gpus", + "4", + "--fpm-gpu-counts", + "4", + "--fpm-parallel-presets", + "tep", + "--fpm-kv-cache-dtypes", + "bfloat16", + "--fpm-executor", + "slurm", + "--fpm-slurm-container-image", + launch["deployment"]["image"], + "--fpm-slurm-container-mount", + "/cache:/cache", + "--fpm-max-model-len", + "4096", + "--fpm-max-num-batched-tokens", + "1024", + "--fpm-max-num-seqs", + "64", + "--fpm-gpu-memory-utilization", + "0.85", + "--fpm-prefill-cudagraph-policy", + "runtime", + ] + extra = [ + "--fpm-runtime-instrumentation", + str(manifest), + "--fpm-runtime-launch", + str(launch_path), + "--fpm-runtime-configuration", + "tp4", + ] + return argv, extra, launch + + +def test_formal_collector_cli_requires_all_instrumentation_inputs(tmp_path, capsys, monkeypatch): + from collector.fpm_forward import cli + + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + argv, extra, _launch = _formal_cli_inputs(tmp_path) + with pytest.raises(SystemExit) as error: + cli.main([*argv, *extra[:2], "--plan-only"]) + assert error.value.code == 2 + assert "required together" in capsys.readouterr().err + assert cli.main([*argv, *extra, "--plan-only"]) == 0 + plan = json.loads(capsys.readouterr().out) + assert plan["runtime_memory_policy"]["observation"] == "enabled" + assert plan["runtime_observation"]["configuration"] == "tp4" + assert plan["runtime_observation"]["launch"]["identity"]["framework_version"] == "0.28.0" + + +def test_repeatability_roundtrips_and_hashes_instrumented_formal_plan(tmp_path): + from collector.fpm_forward import cli, entry, repeatability + from collector.fpm_forward.runtime_memory import validate_saved_plan + from collector.model_cases import build_collection_case_plan + + argv, extra, _launch = _formal_cli_inputs(tmp_path) + args = cli._parser().parse_args([*argv, *extra]) + case_plan = build_collection_case_plan(backend="vllm", model_path=args.model_path, gpu_type=args.gpu) + plan, _overrides = entry.resolve_inputs(args, case_plan) + saved = plan.to_dict() + validate_saved_plan(saved) + (tmp_path / "collection-plan.json").write_text(json.dumps(saved)) + loaded = repeatability.load_repeatability_source(tmp_path) + assert loaded.to_dict() == saved + assert loaded.runtime_instrumentation.sha256 == plan.runtime_instrumentation.sha256 + selected = { + "cell_id": plan.cells[0].cell_id, + "benchmark_points": { + "schema_version": 3, + "prefill": [{"batch_size": 1, "total_prefill_tokens": 16, "total_kv_read_tokens": 0}], + "decode": [], + }, + } + subset = repeatability._subset_plan(loaded, selected) + validate_saved_plan(subset.to_dict()) + assert subset.to_dict()["runtime_observation"] == saved["runtime_observation"] + assert subset.sha256 != plan.sha256 + + +def test_formal_collection_instruments_resolved_memory_and_archives_phase_attempts(tmp_path, monkeypatch): + from collector.fpm_forward import cli, entry, runner + from collector.model_cases import build_collection_case_plan + + argv, extra, _launch = _formal_cli_inputs(tmp_path) + args = cli._parser().parse_args([*argv, *extra]) + case_plan = build_collection_case_plan(backend="vllm", model_path=args.model_path, gpu_type=args.gpu) + plan, overrides = entry.resolve_inputs(args, case_plan) + assert plan.deployment_profile(plan.cells[0]).resources.memory_source != "pending" + executions = [] + monkeypatch.setattr( + runner, + "_cell_runner", + lambda plan, cell, _manifest, directory: _SyntheticExecutor(plan, cell, directory, executions), + ) + monkeypatch.setattr(runner, "_runtime_collection_summary", lambda *_args, **_kwargs: {}) + errors = runner.run_collection( + plan, + generator_overrides=overrides, + checkpoint_dir=str(tmp_path / "checkpoint"), + artifact_root=str(tmp_path / "collection"), + resume=False, + retry_failed=False, + publish_database=False, + ) + assert errors == [] + root = tmp_path / "collection" / plan.sha256[:16] + index = json.loads((root / "runtime-observations.json").read_text()) + attempt = index["configurations"]["tp4"]["attempts"][0] + assert attempt["status"] == "captured" + assert set(attempt["phases"]) == {"prefill", "decode"} + contexts = [ + json.loads((root / phase["launch_manifest"]["path"]).read_text()) for phase in attempt["phases"].values() + ] + assert {context["attempt_id"] for context in contexts} == {attempt["attempt_id"]} + assert len({context["collector_attempt_id"] for context in contexts}) == 2 + assert all(len(items) == 1 for items in attempt["phase_attempts"].values()) + for path in (root / "cells").glob("*/generator-request.json"): + args = json.loads(path.read_text())["params"]["agg"]["extra_cli_args"] + assert args[args.index("--revision") + 1] == plan.runtime_launch["identity"]["model_revision"] + assert args.count("--revision") == 1 + + +def _sidecar_inputs(tmp_path): + launch, manifest, marker = _inputs(tmp_path) + sidecar = tmp_path / "hf_quant_config.json" + sidecar.write_text(json.dumps({"quantization": {"quant_algo": "NVFP4", "kv_cache_quant_algo": None}})) + launch["model_config"]["source_files"] = {sidecar.name: hashlib.sha256(sidecar.read_bytes()).hexdigest()} + launch["precision"].update(gemm_quant_mode="nvfp4", moe_quant_mode="nvfp4") + return launch, manifest, marker, sidecar + + +def test_probe_binds_raw_primary_and_adjacent_quantization_metadata(tmp_path): + from collector.fpm_forward import runtime_probe + + launch, manifest, marker, sidecar = _sidecar_inputs(tmp_path) + primary = Path(launch["model_config"]["path"]) + original = {"model-config.json": primary.read_bytes(), sidecar.name: sidecar.read_bytes()} + output = tmp_path / "probe" + result = runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=output) + assert result["status"] == "preview", result + assert primary.read_bytes() == original["model-config.json"] + assert not marker.exists() + for phase in result["configurations"]["tp4"]["phases"].values(): + context = json.loads((output / phase["launch_manifest"]["path"]).read_text()) + assert context["launch"]["model_config"] == launch["model_config"] + artifacts = {Path(a["path"]).name: a for a in phase["artifacts"]} + for name, content in original.items(): + assert (output / artifacts[name]["path"]).read_bytes() == content + assert artifacts[name]["sha256"] == hashlib.sha256(content).hexdigest() + + +@pytest.mark.parametrize("mutation", ["missing", "changed", "malformed", "unrecorded", "escape", "symlink", "bad-hash"]) +def test_probe_rejects_unbound_quantization_metadata(tmp_path, mutation): + from collector.fpm_forward import runtime_probe + + launch, manifest, _marker, sidecar = _sidecar_inputs(tmp_path) + # Keep native precision unchanged so rejection cannot be a precision mismatch. + launch["precision"].update(gemm_quant_mode="bfloat16", moe_quant_mode="bfloat16") + sidecar.write_text("{}") + launch["model_config"]["source_files"][sidecar.name] = hashlib.sha256(sidecar.read_bytes()).hexdigest() + if mutation == "missing": + sidecar.unlink() + elif mutation == "changed": + sidecar.write_text("{}\n") + elif mutation == "malformed": + sidecar.write_text("[]") + launch["model_config"]["source_files"][sidecar.name] = hashlib.sha256(sidecar.read_bytes()).hexdigest() + elif mutation == "unrecorded": + launch["model_config"].pop("source_files") + elif mutation == "escape": + launch["model_config"]["source_files"] = {"../hf_quant_config.json": "a" * 64} + elif mutation == "symlink": + sidecar.rename(tmp_path / "other.json") + (tmp_path / "other.json").write_text("{}\n") + sidecar.symlink_to(tmp_path / "other.json") + elif mutation == "bad-hash": + launch["model_config"]["source_files"][sidecar.name] = "not-a-digest" + result = runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=tmp_path / "probe") + assert result["status"] == "failed", result + + +def test_probe_preserves_logical_adjacency_for_hf_cache_symlinks(tmp_path): + from collector.fpm_forward import runtime_probe + + launch, manifest, _marker, sidecar = _sidecar_inputs(tmp_path) + primary = Path(launch["model_config"]["path"]) + blobs = tmp_path / "blobs" + blobs.mkdir() + for path, blob_name in ((primary, "primary-blob"), (sidecar, "quantization-blob")): + target = blobs / blob_name + path.rename(target) + path.symlink_to(target) + result = runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=tmp_path / "probe") + assert result["status"] == "preview", result + + +def test_probe_rechecks_sidecar_between_plan_and_stage(tmp_path, monkeypatch): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _marker, sidecar = _sidecar_inputs(tmp_path) + original_render = runner._render_cell + + def render_then_mutate(*args, **kwargs): + result = original_render(*args, **kwargs) + sidecar.write_bytes(sidecar.read_bytes() + b"\n") + return result + + monkeypatch.setattr(runner, "_render_cell", render_then_mutate) + result = runtime_probe.probe_runtime({"tp4": launch}, instrumentation=manifest, output_dir=tmp_path / "probe") + assert result["status"] == "failed" + assert any("SHA-256" in item for item in result["configurations"]["tp4"]["diagnostics"]) + + +@pytest.mark.parametrize("version", ["0.27.0", "0.28.0"]) +@pytest.mark.parametrize("pending", [False, True]) +def test_instrumented_saved_plan_hash_round_trip(tmp_path, version, pending): + from collector.fpm_forward import cli, entry, runtime_memory + from collector.model_cases import build_collection_case_plan + + argv, extra, _launch = _formal_cli_inputs(tmp_path) + for filename in ("profile.json", "manifest.json", "launch.json"): + path = tmp_path / filename + payload = json.loads(path.read_text().replace("0.28.0", version)) + if pending and filename == "profile.json": + resources = payload["deployments"][0]["resources"] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(field, None) + path.write_text(json.dumps(payload)) + args = cli._parser().parse_args([*argv, *extra]) + case_plan = build_collection_case_plan(backend="vllm", model_path=args.model_path, gpu_type=args.gpu) + plan, _overrides = entry.resolve_inputs(args, case_plan) + assert (plan.deployment_profile(plan.cells[0]).resources.memory_source == "pending") is pending + payload = plan.to_dict() + runtime_memory.validate_saved_plan(payload) + for field in ("runtime_observation", "runtime_memory_policy"): + changed = copy.deepcopy(payload) + changed[field]["bundle_sha256"] = "f" * 64 + with pytest.raises(ValueError, match="SHA-256"): + runtime_memory.validate_saved_plan(changed) + + +@pytest.mark.parametrize( + "mutation", + [ + "context", + "binding", + "manifest", + "redirect", + "ownership", + "compute_domain_ownership", + "compute_domain_missing_ownership", + ], +) +def test_probe_resume_rejects_changed_ownership_before_external_commands(tmp_path, monkeypatch, mutation): + import shutil + + import yaml + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _marker = _inputs(tmp_path) + launch["deployment"] = {"executor": "kubernetes", "image": "image@sha256:synthetic"} + auxiliary_ownership = mutation in {"compute_domain_ownership", "compute_domain_missing_ownership"} + if auxiliary_ownership: + launch["identity"]["gpu"] = "gb300" + launch["topology"].update(tp=8, dp=1, moe_tp=8, moe_ep=1) + configuration = "tp8" if auxiliary_ownership else "tp4" + real_runner = runner._cell_runner + + class Interrupted(_SyntheticExecutor): + def execute(self, pods): + super().execute(pods) + raise KeyboardInterrupt() + + monkeypatch.setattr(runner, "_salvage_artifacts", lambda *_args: None) + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: Interrupted(p, c, d, [])) + output = tmp_path / "probe" + with pytest.raises(KeyboardInterrupt): + runtime_probe.probe_runtime({configuration: launch}, instrumentation=manifest, output_dir=output, execute=True) + index_path = output / "observations.json" + index = json.loads(index_path.read_text()) + phase = index["configurations"][configuration]["attempts"][0]["phases"]["prefill"] + context = output / phase["launch_manifest"]["path"] + resource = context.parent / runner.FPM_MANIFEST_FILENAME + if mutation == "context": + context.write_bytes(context.read_bytes() + b"\n") + elif mutation == "binding": + payload = json.loads(context.read_text()) + payload["attempt_id"] = "unrelated-attempt" + context.write_text(json.dumps(payload)) + phase["launch_manifest"]["sha256"] = hashlib.sha256(context.read_bytes()).hexdigest() + elif mutation in {"manifest", "ownership"} or auxiliary_ownership: + documents = list(yaml.safe_load_all(resource.read_text())) + if auxiliary_ownership: + metadata = next(document for document in documents if document["kind"] == "ComputeDomain")["metadata"] + else: + metadata = runner._workload_document(documents)["metadata"] + metadata["name"] = "unrelated-production-worker" + if mutation == "compute_domain_missing_ownership": + for label in ( + "aiconfigurator.nvidia.com/owned-by", + "aiconfigurator.nvidia.com/plan", + runner.FPM_CELL_LABEL, + ): + metadata["labels"].pop(label) + else: + metadata["labels"][runner.FPM_CELL_LABEL] = "unrelated-cell" + resource.write_text(yaml.safe_dump_all(documents)) + if (mutation == "ownership" or auxiliary_ownership) and "resource_manifest" in phase: + phase["resource_manifest"]["sha256"] = hashlib.sha256(resource.read_bytes()).hexdigest() + else: + redirected = output / "redirected" + shutil.copytree(context.parent, redirected) + phase["launch_manifest"]["path"] = str((redirected / context.name).relative_to(output)) + index_path.write_text(json.dumps(index)) + monkeypatch.setattr(runner, "_cell_runner", real_runner) + monkeypatch.setattr(runner, "_kubectl_command", lambda: ["kubectl"]) + commands = [] + + def no_external(args, **kwargs): + commands.append(args) + raise RuntimeError("unexpected external command") + + monkeypatch.setattr(runner, "_run_command", no_external) + result = runtime_probe.probe_runtime( + {configuration: launch}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + assert result["status"] == "failed" + assert commands == [] + if auxiliary_ownership: + assert any("collector ownership" in item for item in result["configurations"][configuration]["diagnostics"]) + + +@pytest.mark.parametrize( + ("executor", "changed_bundle", "configuration"), + [ + ("kubernetes", False, "tp4"), + ("kubernetes", True, "tp4"), + ("slurm", False, "tp4"), + ("slurm", True, "tp4"), + ("kubernetes", False, "tp8"), + ("slurm", False, "dep8"), + ], +) +@pytest.mark.parametrize("saved_binding", ["current", "legacy", "older-loader"]) +def test_probe_resume_preserves_valid_interrupted_recovery( + tmp_path, monkeypatch, executor, changed_bundle, configuration, saved_binding +): + from collector.fpm_forward import runner, runtime_probe + + launch, manifest, _marker = _inputs(tmp_path) + if executor == "kubernetes": + launch["deployment"] = {"executor": executor, "image": launch["deployment"]["image"]} + if configuration == "tp8": + launch["identity"]["gpu"] = "gb300" + launch["topology"].update(tp=8, dp=1, moe_tp=8, moe_ep=1) + elif configuration == "dep8": + launch["identity"]["gpu"] = "gb300" + launch["topology"].update(tp=1, dp=8, moe_tp=1, moe_ep=8) + executions, cleanups, salvages = [], [], [] + + class InterruptedOnce(_SyntheticExecutor): + def cleanup(self): + cleanups.append(str(self.cell_dir)) + + def execute(self, pods): + super().execute(pods) + if len(executions) == 1: + raise KeyboardInterrupt() + + monkeypatch.setattr(runner, "_salvage_artifacts", lambda resource, _cell: salvages.append(str(resource.cell_dir))) + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: InterruptedOnce(p, c, d, executions)) + output = tmp_path / "probe" + with pytest.raises(KeyboardInterrupt): + runtime_probe.probe_runtime({configuration: launch}, instrumentation=manifest, output_dir=output, execute=True) + index_path = output / "observations.json" + index = json.loads(index_path.read_text()) + previous = index["configurations"][configuration]["attempts"][0] + reference = previous["phases"]["prefill"]["launch_manifest"] + saved_path = output / reference["path"] + context = json.loads(saved_path.read_text()) + if saved_binding == "legacy": + context.pop("instrumentation_binding") + elif saved_binding == "older-loader": + context["instrumentation_binding"]["loader_sha256"] = "f" * 64 + saved_path.write_text(json.dumps(context)) + reference.update(runner._file_metadata(saved_path)) + index_path.write_text(json.dumps(index)) + interrupted_dir = str((output / previous["phases"]["prefill"]["launch_manifest"]["path"]).parent) + if changed_bundle: + (tmp_path / "observer.py").write_text("# a revised observer retained alongside the interrupted bundle\n") + result = runtime_probe.probe_runtime( + {configuration: launch}, instrumentation=manifest, output_dir=output, execute=True, resume=True + ) + assert result["status"] == "completed", result + assert len(executions) == 3 + assert cleanups.count(interrupted_dir) == 3 + assert salvages.count(interrupted_dir) == 2 + attempts = json.loads((output / "observations.json").read_text())["configurations"][configuration]["attempts"] + previous = attempts[0] + assert any(item["kind"] == "observation" for item in previous["phases"]["prefill"]["recovery_artifacts"]) + for phase in attempts[-1]["phases"].values(): + from collector.fpm_forward.runtime_instrumentation import runtime_binding + + new_context = json.loads((output / phase["launch_manifest"]["path"]).read_text()) + assert new_context["instrumentation_binding"] == runtime_binding() + observations = [item for item in phase["artifacts"] if item["kind"] == "observation"] + assert len(observations) == (1 if configuration == "tp4" else 2) + + +@pytest.mark.parametrize("saved_status", ["running", "passed"]) +@pytest.mark.parametrize("archive_failure", [None, "copy", "index", "context", "stale"]) +def test_formal_resume_archives_recovered_native_evidence_before_skip( + tmp_path, monkeypatch, saved_status, archive_failure +): + from collector.fpm_forward import cli, entry, runner, runtime_probe + from collector.model_cases import build_collection_case_plan + + from .test_fpm_runner import _native_payload + + argv, extra, _launch = _formal_cli_inputs(tmp_path) + args = cli._parser().parse_args([*argv, *extra]) + case_plan = build_collection_case_plan(backend="vllm", model_path=args.model_path, gpu_type=args.gpu) + plan, overrides = entry.resolve_inputs(args, case_plan) + artifact_root = tmp_path / "collection" + root = artifact_root / plan.sha256[:16] + parent = "interrupted-parent" + frozen = runtime_probe.prepare_collection_observations(plan, root, parent) + prefill = next(cell for cell in plan.cells if cell.workload_kind == "prefill") + cell_dir = root / "cells" / prefill.cell_id + cell_dir.mkdir(parents=True) + runner._render_cell(plan, prefill, cell_dir, overrides) + context = runtime_probe.launch_context(plan, prefill, configuration=plan.runtime_configuration, attempt_id=parent) + context.update(collector_attempt_id="prefill-native-attempt", cell_id=prefill.cell_id) + runtime_probe.stage_runtime_instrumentation(frozen, cell_dir, context) + raw = cell_dir / "raw" / "synthetic-node" + raw.mkdir(parents=True) + + def native_files(cell, directory, attempt_id): + (directory / "collector-provenance.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell.cell_id, + "plan_sha256": plan.sha256, + "attempt_id": attempt_id, + "runtime": {"backend": "vllm", "backend_version": "0.28.0"}, + } + ) + ) + (directory / "benchmark.json").write_text(json.dumps(_native_payload(phase=cell.workload_kind, rank=0, dp=1))) + + native_files(prefill, raw, "prefill-native-attempt") + observation = raw / "synthetic-observation.json" + observation.write_text( + json.dumps({**context, "schema_version": "aisimulate-runtime-observation/v1", "synthetic": True}) + ) + original_observation = observation.read_bytes() + if archive_failure == "stale": + # A prior failed snapshot can predate the final observation write. + observation.unlink() + runtime_probe.record_collection_observations( + plan, root, parent, prefill, "prefill-native-attempt", cell_dir, "failed" + ) + observation.write_bytes(original_observation) + runner._runtime_collection_summary( + prefill, cell_dir / "raw", expected_plan_sha256=plan.sha256, expected_attempt_id="prefill-native-attempt" + ) + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + checkpoint_path = checkpoint_dir / "fpm_forward.json" + checkpoint_path.write_text( + json.dumps( + { + "schema": runner.CHECKPOINT_SCHEMA, + "plan_sha256": plan.sha256, + "runtime_observation_attempt_id": parent, + "cells": {prefill.cell_id: {"status": saved_status, "attempt_id": "prefill-native-attempt"}}, + } + ) + ) + executions = [] + + class NativeExecutor(_SyntheticExecutor): + def execute(self, pods): + super().execute(pods) + native_files(self.cell, self.cell_dir / "raw" / "synthetic-node", self.identity["attempt_id"]) + + monkeypatch.setattr(runner, "_cell_runner", lambda p, c, m, d: NativeExecutor(p, c, d, executions)) + original_copy = runtime_probe.shutil.copy2 + original_atomic = runner._atomic_json + context_path = cell_dir / runtime_probe.CONTEXT_FILENAME + original_context = context_path.read_bytes() + if archive_failure == "context": + corrupted = {**context, "collector_attempt_id": "unrelated-attempt"} + context_path.write_text(json.dumps(corrupted)) + if archive_failure == "copy": + + def interrupted_copy(source, target, *args, **kwargs): + if Path(source) == context_path: + raise OSError("synthetic observation snapshot interruption") + return original_copy(source, target, *args, **kwargs) + + monkeypatch.setattr(runtime_probe.shutil, "copy2", interrupted_copy) + if archive_failure == "index": + + def interrupted_index(path, payload): + if Path(path).name == "runtime-observations.json": + raise OSError("synthetic observation index interruption") + return original_atomic(path, payload) + + monkeypatch.setattr(runner, "_atomic_json", interrupted_index) + kwargs = dict( + generator_overrides=overrides, + checkpoint_dir=str(checkpoint_dir), + artifact_root=str(artifact_root), + resume=True, + retry_failed=True, + publish_database=False, + ) + errors = runner.run_collection(plan, **kwargs) + if archive_failure not in {None, "stale"}: + assert any(item["classification"] == "runtime_observation_archive_failed" for item in errors) + assert json.loads(checkpoint_path.read_text())["cells"][prefill.cell_id]["status"] == "failed" + assert observation.read_bytes() == original_observation + monkeypatch.setattr(runtime_probe.shutil, "copy2", original_copy) + monkeypatch.setattr(runner, "_atomic_json", original_atomic) + context_path.write_bytes(original_context) + assert runner.run_collection(plan, **kwargs) == [] + else: + assert errors == [] + index_path = root / "runtime-observations.json" + index = json.loads(index_path.read_text()) + attempt = index["configurations"]["tp4"]["attempts"][0] + assert set(attempt["phases"]) == {"prefill", "decode"} + assert attempt["status"] == "captured" + assert attempt["phases"]["prefill"]["collector_attempt_id"] == "prefill-native-attempt" + archived = [item for item in attempt["phases"]["prefill"]["artifacts"] if item["kind"] == "observation"] + assert len(archived) == 1 + assert (root / archived[0]["path"]).read_bytes() == original_observation + if archive_failure == "stale": + assert [phase["status"] for phase in attempt["phase_attempts"]["prefill"]] == ["failed", "passed"] + assert not any(item["kind"] == "observation" for item in attempt["phase_attempts"]["prefill"][0]["artifacts"]) + assert observation.read_bytes() == original_observation + checkpoint = json.loads(checkpoint_path.read_text()) + assert checkpoint["runtime_observations"] == str(index_path) + assert all(entry["status"] == "passed" for entry in checkpoint["cells"].values()) + assert executions == [("tp4", "decode")] + before = index_path.read_bytes() + assert runner.run_collection(plan, **kwargs) == [] + assert index_path.read_bytes() == before + assert executions == [("tp4", "decode")] diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py b/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py index 9e4f86cd8..f4cd3ff4b 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py @@ -144,6 +144,35 @@ def _run(staged: SimpleNamespace, *, timeout: int = 60, extra_env: dict[str, str ) +@pytest.mark.parametrize("status,mask", [("observed", [0]), ("unavailable", []), ("partial", list(range(16)))]) +def test_slurm_same_step_cpu_guard_preserves_failure_before_model_start(tmp_path, status, mask): + staged = _stage(tmp_path, run_script="touch model-started\n") + staged.script.write_text(staged.script.read_text().replace('Path("/results")', f"Path({str(staged.results)!r})")) + # The observer's Linux API is tested separately. This seam verifies the + # real wrapper enforces its result before etcd/preflight/model startup. + payload = {"status": status, "main_thread_allowed_cpus": mask} + (staged.workdir / "fpm_memory_observer.py").write_text( + "import json\ndef observe_cpu(kind, **kwargs):\n" + f" payload = {payload!r}\n" + " payload.update(kind=kind,requested=kwargs['requested_cpus_per_task'])\n" + " (kwargs['directory']/'fpm-cpu-launcher.json').write_text(json.dumps(payload))\n" + " return payload\n" + ) + result = _run( + staged, + extra_env={ + "FPM_SLURM_CPUS_PER_TASK": "16", + "FPM_SLURM_CPU_BIND": "cores", + "FPM_LOCAL_GPU_COUNT": "4", + }, + ) + assert result.returncode != 0 + assert "CPU affinity" in result.stderr + assert json.loads((staged.results / "fpm-cpu-launcher.json").read_text())["requested"] == 16 + assert not staged.etcd_trace.exists() + assert not (staged.workdir / "model-started").exists() + + def test_fpm_exec_starts_leader_etcd_before_preflight(): """The follower readiness probe budget only covers pod-exec skew when the leader's etcd starts before the unbounded vLLM/torch preflight import.""" @@ -153,6 +182,26 @@ def test_fpm_exec_starts_leader_etcd_before_preflight(): assert "time.monotonic() + float(sys.argv[2])" in script +@pytest.mark.parametrize("pending_memory", [False, True]) +@pytest.mark.parametrize("existing", [None, "", "/other/modules"]) +def test_memory_observer_pythonpath_is_available_before_engine_launch(tmp_path, pending_memory, existing): + staged = _stage(tmp_path, run_script="#!/bin/bash\n") + if pending_memory: + (staged.workdir / "fpm_memory_worker.py").write_text("# staged observer\n") + script = staged.script.read_text().split('etcd_endpoint="http://')[0] + script += '\nprintf "%s" "${PYTHONPATH:-}"\n' + env = dict(staged.env) + if existing is None: + env.pop("PYTHONPATH", None) + else: + env["PYTHONPATH"] = existing + result = subprocess.run(["bash", "-c", script], env=env, text=True, capture_output=True, check=True) + expected = existing or "" + if pending_memory: + expected = str(staged.workdir) + (":" + expected if expected else "") + assert result.stdout == expected + + def test_fpm_exec_leader_starts_etcd_and_cleanup_stops_it(tmp_path): output_path = tmp_path / "results" / "benchmark.json" staged = _stage( @@ -284,6 +333,62 @@ def test_preflight_rejects_runtime_without_kvwarm_capability(tmp_path, monkeypat assert audit["missing_methods"] == ["_kvwarm_warm_eligible"] +@pytest.mark.parametrize( + ("phase", "setting", "native_capability", "adapter", "accepted"), + [ + ("prefill", "on", False, False, False), + ("agg", "true", False, False, False), + ("prefill", "on", True, False, True), + ("prefill", "off", False, False, True), + ("prefill", "0", False, False, True), + ("decode", "on", False, False, True), + ("prefill", "on", False, True, True), + ], +) +def test_preflight_checks_requested_prefill_seeding_before_model_load( + tmp_path, monkeypatch, phase, setting, native_capability, adapter, accepted +): + import sys + from types import ModuleType + + from collector.fpm_forward.runtime import preflight + + def no_initialization(*_args, **_kwargs): + pytest.fail("preflight must inspect capabilities without constructing the model or scheduler") + + methods = {name: lambda self: None for name in preflight.GRAPH_AWARE_METHODS} + methods["__init__"] = no_initialization + if native_capability: + methods["_bench_realseed_on"] = lambda self: True + scheduler = type("InstrumentedScheduler", (), methods) + module = ModuleType("dynamo.vllm.instrumented_scheduler") + module.BenchmarkPoint = type( + "BenchmarkPoint", (), {"__dataclass_fields__": {name: object() for name in preflight.GRAPH_AWARE_FIELDS}} + ) + module.InstrumentedScheduler = scheduler + monkeypatch.setitem(sys.modules, "dynamo.vllm.instrumented_scheduler", module) + monkeypatch.setenv("FPM_BENCHMARK_MODE", phase) + monkeypatch.setenv("DYN_BENCH_PREFILL_REAL_SEED", setting) + monkeypatch.delenv("DYN_FPM_DSV41_REAL_KV", raising=False) + if adapter: + adapter_module = ModuleType("dsv41_scheduler") + adapter_module.DeepseekV41RealKVScheduler = scheduler + monkeypatch.setitem(sys.modules, "dsv41_scheduler", adapter_module) + monkeypatch.setenv("DYN_FPM_DSV41_REAL_KV", "1") + audit_path = tmp_path / "runtime-preflight.json" + monkeypatch.setattr(preflight, "_AUDIT_PATH", audit_path) + + if accepted: + preflight.main() + else: + with pytest.raises(RuntimeError, match="_bench_realseed_on"): + preflight.main() + audit = json.loads(audit_path.read_text()) + assert audit["status"] == ("passed" if accepted else "failed") + assert audit["missing_methods"] == ([] if accepted else ["_bench_realseed_on"]) + assert audit["prefill_real_seed_requested"] == (phase in {"prefill", "agg"} and setting in {"on", "true"}) + + def test_fpm_exec_propagates_fail_closed_env_source(tmp_path): """fpm_env.sh exits 2 on incomplete multinode discovery; sourcing must terminate fpm_exec.sh with the same status before any resource starts.""" diff --git a/python/aisimulate/tests/unit/collector/test_fpm_slurm.py b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py index 5eeacdb83..b600928b7 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_slurm.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py @@ -22,6 +22,7 @@ def runner(tmp_path, monkeypatch): monkeypatch.setenv("SLURM_JOB_ID", "1234") monkeypatch.setenv("SLURM_JOB_NODELIST", "test-node") + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", "16") manifest = tmp_path / "manifest.yaml" manifest.write_text( json.dumps( @@ -32,7 +33,16 @@ def runner(tmp_path, monkeypatch): } ) ) - return SlurmCellRunner(manifest, tmp_path, image="image@sha256:abc", mounts=("/cache:/cache",), total_gpus=4) + return SlurmCellRunner( + manifest, + tmp_path, + image="image@sha256:abc", + mounts=("/cache:/cache",), + total_gpus=4, + cpus_per_task=16, + cpu_bind="cores", + attention_tp=1, + ) def test_slurm_requires_scheduler_allocation(runner, monkeypatch): @@ -63,10 +73,123 @@ def command(args, **kwargs): assert (runner.cell_dir / "slurm-runtime" / source.name).read_text() == source.read_text() argv = commands[-1] assert "--jobid=1234" in argv and "--gpus-per-node=4" in argv + assert "--cpus-per-task=16" in argv and "--cpu-bind=cores" in argv assert "FPM_NODE_RANK=0" in argv and "FPM_MASTER_ADDR=test-node" in argv assert f"{runner.cell_dir}/raw/node0000:/results" in next(a for a in argv if a.startswith("--container-mounts=")) +@pytest.mark.parametrize("allocation", ["1", "8", "16(x0)", "sixteen", "16,16"]) +def test_slurm_rejects_small_or_malformed_allocated_cpu_pool(runner, monkeypatch, allocation): + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", allocation) + monkeypatch.setattr( + runner, + "_command", + lambda args, **kwargs: SimpleNamespace( + stdout="JobId=1234 JobState=RUNNING NodeList=test-node" if "job" in args else "test-node" + ), + ) + with pytest.raises(ValueError, match="CPU"): + runner.wait_ready(1) + assert runner.hosts == [] + + +def test_slurm_two_node_cpu_allocation_preserves_host_specific_pool(runner, monkeypatch): + runner.node_count = 2 + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", "32,16") + monkeypatch.setattr( + runner, + "_command", + lambda args, **kwargs: SimpleNamespace( + stdout="JobId=1234 JobState=RUNNING NodeList=node-a,node-b" if "job" in args else "node-a node-b" + ), + ) + assert runner.wait_ready(2) == ["node0000", "node0001"] + report = json.loads((runner.cell_dir / "slurm-cpu-allocation.json").read_text()) + assert report["allocated_cpus_per_node"] == {"node-a": 32, "node-b": 16} + + +def test_slurm_reads_allocated_cpu_ids_when_job_env_is_absent(runner, monkeypatch): + monkeypatch.delenv("SLURM_JOB_CPUS_PER_NODE") + + def command(args, **kwargs): + if "--details" in args: + return SimpleNamespace(stdout="JobId=1234\n Nodes=test-node CPU_IDs=32-39,64-71 Mem=0") + return SimpleNamespace( + stdout="JobId=1234 JobState=RUNNING NodeList=test-node" if "job" in args else "test-node" + ) + + monkeypatch.setattr(runner, "_command", command) + assert runner.wait_ready(1) == ["node0000"] + report = json.loads((runner.cell_dir / "slurm-cpu-allocation.json").read_text()) + assert report["allocated_cpus_per_node"] == {"test-node": 16} + assert "CPU_IDs" in report["source"] + + +def test_slurm_expands_grouped_cpu_allocation_by_host(runner, monkeypatch): + runner.node_count = 3 + monkeypatch.delenv("SLURM_JOB_CPUS_PER_NODE") + + def command(args, **kwargs): + if "--details" in args: + return SimpleNamespace( + stdout=("JobId=1234\n Nodes=node3 CPU_IDs=64-95 Mem=0\n Nodes=node[1-2] CPU_IDs=0-15 Mem=0") + ) + if "job" in args: + return SimpleNamespace(stdout="JobId=1234 JobState=RUNNING NodeList=node[1-3]") + return SimpleNamespace( + stdout={"node[1-3]": "node1 node2 node3", "node[1-2]": "node1 node2", "node3": "node3"}[args[-1]] + ) + + monkeypatch.setattr(runner, "_command", command) + assert runner.wait_ready(3) == ["node0000", "node0001", "node0002"] + report = json.loads((runner.cell_dir / "slurm-cpu-allocation.json").read_text()) + assert report["allocated_cpus_per_node"] == {"node1": 16, "node2": 16, "node3": 32} + + +@pytest.mark.parametrize("cpus,binding", [(0, "cores"), (True, "cores"), (16, "mask_cpu:1"), (1, "cores")]) +def test_slurm_rejects_invalid_policy_or_less_than_local_dp_schedulers(runner, cpus, binding): + with pytest.raises(ValueError, match="CPU|cpu"): + SlurmCellRunner( + runner.cell_dir / "manifest.yaml", + runner.cell_dir, + image=runner.image, + mounts=(), + total_gpus=4, + cpus_per_task=cpus, + cpu_bind=binding, + attention_tp=1, + ) + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX container entrypoint") +def test_slurm_container_entrypoint_ignores_shadowed_env_on_caller_path(runner, monkeypatch): + host_bin = runner.cell_dir / "host-bin" + host_bin.mkdir() + shadowed_env = host_bin / "env" + shadowed_env.write_text("#!/bin/sh\nexit 13\n") + shadowed_env.chmod(0o755) + monkeypatch.setenv("PATH", f"{host_bin}:{os.environ['PATH']}") + runner.hosts = ["test-node"] + + def command(args, *, timeout): + # Execute the emitted payload after the Slurm options so PATH lookup + # and environment injection remain real. + executable_index = next(index for index, arg in enumerate(args[1:], 1) if not arg.startswith("--")) + return subprocess.run(args[executable_index:], check=True, capture_output=True, text=True, timeout=timeout) + + monkeypatch.setattr(runner, "_command", command) + result = runner._exec( + "node0000", + [ + sys.executable, + "-c", + "import json, os; print(json.dumps([os.environ['FPM_NODE_RANK'], os.environ['FPM_MASTER_ADDR']]))", + ], + timeout=10, + ) + assert json.loads(result.stdout) == ["0", "test-node"] + + def test_slurm_cleanup_cancels_only_receipted_steps_and_verifies_exit(runner, monkeypatch): runner.owner_path.parent.mkdir(parents=True, exist_ok=True) runner.owner_path.write_text(json.dumps({"job_id": "1233", "step_name": runner.step_name})) @@ -579,6 +702,7 @@ def test_profile_campaign_observes_slurm_version_before_native_collection(tmp_pa commands = [] executed = [] monkeypatch.setenv("SLURM_JOB_ID", "1234") + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", "16") monkeypatch.setattr("collector.fpm_forward.slurm.shutil.which", lambda name: f"/fake/{name}") monkeypatch.setattr(campaign.time, "sleep", lambda _seconds: None) import importlib.metadata @@ -604,7 +728,7 @@ def command(args, **_kwargs): next(value.removesuffix(":/results") for value in mounts.split(",") if value.endswith(":/results")) ) assert raw.name == "node0000" - program = args[args.index("env") + 3 :] + program = args[args.index("/usr/bin/env") + 6 :] if program[:2] == ["python3", "-c"]: # Execute the emitted program with the container's result mount mapped to disk. with monkeypatch.context() as patch: diff --git a/python/aisimulate/tests/unit/collector/test_fpm_slurm_cpu_policy.py b/python/aisimulate/tests/unit/collector/test_fpm_slurm_cpu_policy.py new file mode 100644 index 000000000..1a9f3fb40 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm_cpu_policy.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Frozen CPU policy and legacy resume tests, without a scheduler or GPU.""" + +import json +from dataclasses import replace + +import pytest +from collector.fpm_forward import cli +from collector.fpm_forward.config import FPMCollectionOptions, resolve_slurm_cpu_policy +from collector.fpm_forward.repeatability import load_repeatability_source + +from .test_fpm_profile_collection import _argv, _plan, _profile + +pytestmark = pytest.mark.unit + + +def _options(*extra): + return FPMCollectionOptions.from_args( + cli._parser().parse_args( + [ + *_argv(_profile()), + "--fpm-executor", + "slurm", + "--fpm-slurm-container-image", + "runtime.sqsh", + *extra, + ] + ) + ) + + +def _archive(tmp_path, options): + plan = _plan(_profile(), options=options) + root = tmp_path / "artifacts" / plan.sha256[:16] + root.mkdir(parents=True) + (root / "collection-plan.json").write_text(json.dumps(plan.to_dict())) + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "fpm_forward.json").write_text(json.dumps({"plan_sha256": plan.sha256, "cells": {}})) + return plan, root, checkpoint + + +def test_new_slurm_plan_defaults_and_explicit_cpu_policy_are_frozen(tmp_path): + defaults = _options() + assert (defaults.slurm_cpus_per_task, defaults.slurm_cpu_bind) == (16, "cores") + explicit = _options("--fpm-slurm-cpus-per-task", "32", "--fpm-slurm-cpu-bind", "none") + plan, root, checkpoint = _archive(tmp_path, explicit) + assert load_repeatability_source(root).to_dict() == plan.to_dict() + assert resolve_slurm_cpu_policy(None, None, resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent) == ( + 32, + "none", + ) + assert resolve_slurm_cpu_policy(8, "cores", resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent) == ( + 8, + "cores", + ) + + +def test_legacy_slurm_plan_keeps_missing_policy_and_exact_bytes(tmp_path): + options = replace(_options(), slurm_cpus_per_task=None, slurm_cpu_bind=None) + plan, root, checkpoint = _archive(tmp_path, options) + before = (root / "collection-plan.json").read_bytes() + loaded = load_repeatability_source(root) + assert "slurm_cpus_per_task" not in loaded.to_dict()["options"] + assert loaded.sha256 == plan.sha256 + assert resolve_slurm_cpu_policy(None, None, resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent) == ( + None, + None, + ) + assert (root / "collection-plan.json").read_bytes() == before + + +@pytest.mark.parametrize("provided", [(32, None), (None, "none"), (8, None), (None, "cores")]) +def test_partial_collection_resume_inherits_only_omitted_cpu_fields(tmp_path, provided): + from collector.fpm_forward import runner + + options = _options("--fpm-slurm-cpus-per-task", "32", "--fpm-slurm-cpu-bind", "none") + plan, root, checkpoint = _archive(tmp_path, options) + checkpoint_path = checkpoint / "fpm_forward.json" + checkpoint_path.write_text( + json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": {}}) + ) + before = {path: path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} + cpus, binding = resolve_slurm_cpu_policy( + *provided, resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent + ) + assert (cpus, binding) == (provided[0] if provided[0] is not None else 32, provided[1] or "none") + resumed = _plan(_profile(), options=replace(options, slurm_cpus_per_task=cpus, slurm_cpu_bind=binding)) + if (cpus, binding) == (32, "none"): + assert runner._load_checkpoint(checkpoint_path, resumed, True)["plan_sha256"] == plan.sha256 + else: + with pytest.raises(ValueError, match="does not match the current frozen plan"): + runner._load_checkpoint(checkpoint_path, resumed, True) + assert {path: path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} == before + + +@pytest.mark.parametrize("provided", [(32, None), (None, "none")]) +def test_partial_resume_never_invents_missing_legacy_policy(tmp_path, provided): + options = replace(_options(), slurm_cpus_per_task=None, slurm_cpu_bind=None) + _, root, checkpoint = _archive(tmp_path, options) + cpus, binding = resolve_slurm_cpu_policy( + *provided, resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent + ) + assert (cpus, binding) == provided + with pytest.raises(ValueError, match="CPU policy"): + replace(options, slurm_cpus_per_task=cpus, slurm_cpu_bind=binding) + + +def test_partial_resume_requires_untampered_saved_plan(tmp_path): + _, root, checkpoint = _archive(tmp_path, _options()) + path = root / "collection-plan.json" + payload = json.loads(path.read_text()) + payload["options"]["slurm_cpu_bind"] = "none" + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="SHA-256"): + resolve_slurm_cpu_policy(16, None, resume=True, checkpoint_dir=checkpoint, artifact_root=root.parent) + + +@pytest.mark.parametrize( + "updates", + [ + {"slurm_cpus_per_task": True}, + {"slurm_cpus_per_task": 0}, + {"slurm_cpus_per_task": 1.5}, + {"slurm_cpu_bind": "mask_cpu:1"}, + {"slurm_cpu_bind": None}, + {"executor": "kubernetes"}, + ], +) +def test_invalid_or_partial_cpu_policy_is_rejected(updates): + with pytest.raises(ValueError, match="CPU|cpu"): + replace(_options(), **updates) + + +@pytest.mark.parametrize("flag,value", [("--fpm-slurm-cpus-per-task", "16"), ("--fpm-slurm-cpu-bind", "cores")]) +def test_cpu_options_are_rejected_on_kubernetes(flag, value): + with pytest.raises(ValueError, match="require --fpm-executor slurm"): + FPMCollectionOptions.from_args(cli._parser().parse_args([*_argv(_profile()), flag, value])) + + +def test_legacy_worker_retry_rejected_before_retained_artifacts_are_removed(tmp_path, monkeypatch): + from collector.fpm_forward import runner + + plan, root, checkpoint = _archive( + tmp_path, + replace(_options(), slurm_cpus_per_task=None, slurm_cpu_bind=None), + ) + raw = root / "cells" / plan.cells[0].cell_id / "raw" + raw.mkdir(parents=True) + preserved = raw / "incomplete-native-evidence.json" + preserved.write_text('{"diagnostic": "original failed attempt"}\n') + payload = { + "schema": runner.CHECKPOINT_SCHEMA, + "plan_sha256": plan.sha256, + "cells": {cell.cell_id: {"status": "failed", "attempt_id": "old-attempt"} for cell in plan.cells}, + } + (checkpoint / "fpm_forward.json").write_text(json.dumps(payload)) + monkeypatch.setattr(runner, "_cell_runner", lambda *a: pytest.fail("legacy retry launched or cleaned GPU steps")) + with pytest.raises(ValueError, match="no frozen CPU policy"): + runner.run_collection( + plan, + generator_overrides={}, + checkpoint_dir=str(checkpoint), + artifact_root=str(root.parent), + resume=True, + retry_failed=True, + publish_database=False, + ) + assert preserved.read_text() == '{"diagnostic": "original failed attempt"}\n' + assert json.loads((checkpoint / "fpm_forward.json").read_text()) == payload diff --git a/python/aisimulate/tests/unit/collector/test_fpm_slurm_profile.py b/python/aisimulate/tests/unit/collector/test_fpm_slurm_profile.py new file mode 100644 index 000000000..a9242e9ec --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm_profile.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU campaign coverage; the fake cluster emits explicitly synthetic timings.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import threading +from dataclasses import replace +from pathlib import Path + +import pytest +from collector.fpm_forward import cli, database, runner, runtime_memory +from collector.fpm_forward.config import FPMCollectionOptions +from collector.fpm_forward.runtime import fpm_memory_observer as observer + +from .test_fpm_cpu_affinity import _snapshot +from .test_fpm_profile_collection import _argv, _plan, _profile, no_models_or_timing_data # noqa: F401 +from .test_fpm_runner import _native_payload +from .test_fpm_runtime_memory import _initialized + +pytestmark = pytest.mark.unit + +_GENERATOR_INPUTS = { + "K8sConfig": { + "extra_env": [ + {"name": "FPM_READINESS_TIMEOUT_SECONDS", "value": "480"}, + {"name": "PYTHONPATH", "value": "/site-plugins"}, + ] + } +} + + +def _slurm_plan(tmp_path, *, points_file=None): + profile = _profile() + profile.update(model="example/slurm-fpm-model", architecture="UnregisteredMoeForCausalLM", num_experts=8) + profile["deployments"] = [profile["deployments"][1]] + deployment = profile["deployments"][0] + deployment.update( + system="gb200", + backend_version="0.27.0", + kv_cache_dtype="bfloat16", + fmha_quant_mode="bfloat16", + gemm_quant_mode="bfloat16", + moe_quant_mode="bfloat16", + ) + for key in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + deployment["resources"].pop(key) + config = { + "architectures": [profile["architecture"]], + "hidden_size": 128, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "num_hidden_layers": 5, + "max_position_embeddings": 8192, + "vocab_size": 1024, + "num_experts": 8, + "torch_dtype": "bfloat16", + } + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(config)) + args = cli._parser().parse_args( + [ + *_argv(profile), + "--fpm-executor", + "slurm", + "--fpm-slurm-container-image", + "image@sha256:synthetic", + "--fpm-slurm-container-mount", + "/cache:/cache", + "--fpm-prefill-cudagraph-policy", + "runtime", + "--fpm-gpu-memory-utilization", + "0.9", + "--fpm-max-model-len", + "4096", + "--fpm-max-num-batched-tokens", + "1024", + "--fpm-max-num-seqs", + "64", + *(["--fpm-benchmark-points-file", str(points_file)] if points_file is not None else []), + ] + ) + options = replace(FPMCollectionOptions.from_args(args), parallel_presets=("tep",), kv_cache_dtypes=("bfloat16",)) + return _plan( + profile, + options=options, + system="gb200", + model_config_path=str(config_path), + has_model_cases=False, + selected_ops={"attention_context", "attention_generation"}, + generator_overrides=_GENERATOR_INPUTS, + ) + + +class _SyntheticSlurm: + """Emulate the external Slurm/Pyxis boundary, retaining all real campaign code.""" + + def __init__(self, plan, monkeypatch): + self.plan = plan + self.commands = [] + self.executed = [] + self.steps = {"1234.99": "unrelated-job"} + self.sequence = 0 + self.version = "0.27.0" + self.missing_worker = False + self.interrupt_phase = None + self.lock = threading.Lock() + self.barrier = threading.Barrier(2) + self.patch = monkeypatch + monkeypatch.setenv("SLURM_JOB_ID", "1234") + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", "16(x2)") + monkeypatch.setattr("collector.fpm_forward.slurm.shutil.which", lambda name: f"/fake/{name}") + monkeypatch.setattr(runner, "_run_command", self.command) + # Only the artifact-settle delay is elided; no runtime is simulated by sleeping. + monkeypatch.setattr(runner.time, "sleep", lambda _seconds: None) + real_version = observer.importlib.metadata.version + monkeypatch.setattr( + observer.importlib.metadata, "version", lambda name: self.version if name == "vllm" else real_version(name) + ) + monkeypatch.setattr(observer, "_offloader_type", lambda: "vllm.model_executor.offloader.base.NoopOffloader") + + def command(self, args, **_kwargs): + with self.lock: + self.commands.append(args) + if args[0] == "scontrol": + text = "JobId=1234 JobState=RUNNING NodeList=node-a,node-b" if "job" in args else "node-a node-b" + return subprocess.CompletedProcess(args, 0, stdout=text, stderr="") + if args[0] == "squeue": + return subprocess.CompletedProcess( + args, 0, stdout="\n".join(f"{step}|{name}" for step, name in self.steps.items()), stderr="" + ) + if args[0] == "scancel": + del self.steps[args[1]] + return subprocess.CompletedProcess(args, 0, stdout="", stderr="") + assert args[0] == "srun", args + assert "--jobid=1234" in args and "--gpus-per-node=4" in args + assert "--cpus-per-task=16" in args and "--cpu-bind=cores" in args + assert "--container-image=image@sha256:synthetic" in args + mounts = next(value.split("=", 1)[1] for value in args if value.startswith("--container-mounts=")) + bindings = dict(value.rsplit(":", 1)[::-1] for value in mounts.split(",")) + assert bindings["/cache"] == "/cache" + raw = Path(bindings["/results"]) + stage = Path(bindings["/tmp/fpm-bench"]) + rank = int(next(value.split("=", 1)[1] for value in args if value.startswith("FPM_NODE_RANK="))) + assert raw.name == f"node{rank:04d}" + assert "FPM_MASTER_ADDR=node-a" in args + command = args[args.index("/usr/bin/env") + 6 :] + if command[:2] == ["python3", "-c"]: + # Interpret the actual preparation program with the Pyxis /results mount mapped locally. + with self.patch.context() as patch: + patch.setattr(sys, "argv", ["-c", *command[3:]]) + try: + exec(command[2].replace("/results", str(raw)), {}) + except RuntimeError as error: + raise subprocess.CalledProcessError(1, args, output="preparation", stderr=str(error)) from error + return subprocess.CompletedProcess(args, 0, stdout="", stderr="") + assert command == ["bash", "/tmp/fpm-bench/fpm_exec.sh"] + assert {"fpm_memory_observer.py", "fpm_memory_worker.py", "fpm_memory_scheduler.py"} <= { + item.name for item in stage.iterdir() + } + startup = stage / "collector-runtime-env.sh" + assert startup.is_file() + startup_probe = subprocess.run( + [ + "bash", + "-c", + 'set -Eeuo pipefail; source "$1"; source "$2"; ' + 'printf "%s|%s|%s|%s|%s" "$FPM_NODE_RANK" "$FPM_MASTER_ADDR" ' + '"$FPM_NODE_COUNT" "$FPM_READINESS_TIMEOUT_SECONDS" "$PYTHONPATH"', + "probe", + str(startup), + str(stage / "fpm_env.sh"), + ], + env={**os.environ, "FPM_NODE_RANK": str(rank), "FPM_MASTER_ADDR": "node-a"}, + capture_output=True, + text=True, + check=True, + ) + assert startup_probe.stdout == f"{rank}|node-a|2|480|/site-plugins" + launch = (stage / "run.sh").read_text() + assert "--worker-cls fpm_memory_worker.FpmResourceWorker" in launch + assert "--scheduler-cls fpm_memory_scheduler.FpmResourceInstrumentedScheduler" in launch + assert launch.count("--no-async-scheduling") == 1 + assert "--enforce-eager" not in launch and "DYN_FPM_DSV41_REAL_KV" not in startup.read_text() + assert "--compilation-config" not in launch and "fpm_points.json" not in {p.name for p in stage.iterdir()} + runtime = (stage / "fpm_exec.sh").read_text() + assert "collector-runtime-env.sh" in runtime and 'export PYTHONPATH="${workdir}' in runtime + cell = next(cell for cell in self.plan.cells if cell.cell_id == raw.parent.parent.name) + self.executed.append((cell.workload_kind, rank)) + self.sequence += 1 + step = f"1234.{self.sequence}" + self.steps[step] = next(value.split("=", 1)[1] for value in args if value.startswith("--job-name=")) + with self.patch.context() as patch: + patch.setattr(observer, "RESULTS_DIR", raw) + host = ("node-a", "node-b")[rank] + patch.setattr(observer, "cpu_affinity_snapshot", lambda: _snapshot(host, 100, list(range(16)))) + observer.observe_cpu( + "launcher", requested_cpus_per_task=16, cpu_bind="cores", local_gpu_count=4, directory=raw + ) + for tp in range(rank * 4, (rank + 1) * 4): + if self.missing_worker and tp == 7: + continue + worker, scheduler, cache = _initialized() + worker.vllm_config.parallel_config.tensor_parallel_size = 8 + worker.vllm_config.parallel_config.data_parallel_size = 1 + worker.vllm_config.quant_config = None + worker.vllm_config.model_config.quantization = None + worker.vllm_config.model_config.model = self.plan.model_path + worker.vllm_config.model_config.revision = self.plan.fpm_profile.model_revision + patch.setattr( + observer, "cpu_affinity_snapshot", lambda tp=tp: _snapshot(host, 200 + tp, list(range(16))) + ) + observer.observe_cpu("worker", dp_rank=0, tp_rank=tp, pp_rank=0, directory=raw) + observer.observe("worker", worker, dp_rank=0, tp_rank=tp, pp_rank=0) + if tp == 0: + patch.setattr(observer, "cpu_affinity_snapshot", lambda: _snapshot(host, 300, list(range(16)))) + observer.observe_cpu("scheduler", dp_rank=0, directory=raw) + observer.observe("scheduler", scheduler, dp_rank=0, cache_config=cache) + if rank == 0: + (raw / "benchmark.json").write_text( + json.dumps(_native_payload(phase=cell.workload_kind, rank=0, dp=1)) + ) + interrupt = cell.workload_kind == self.interrupt_phase + if not interrupt: + del self.steps[step] + if interrupt: + # Both nodes have emitted their complete attempt before the transport disconnects. + self.barrier.wait(timeout=5) + if rank == 1: + raise KeyboardInterrupt + return subprocess.CompletedProcess(args, 0, stdout="synthetic container completed", stderr="") + + +def _run(plan, tmp_path, *, resume=False): + return runner.run_collection( + plan, + generator_overrides=_GENERATOR_INPUTS, + checkpoint_dir=str(tmp_path / "checkpoints"), + artifact_root=str(tmp_path / "artifacts"), + database_root=str(tmp_path / "data"), + resume=resume, + retry_failed=False, + ) + + +def _checkpoint(tmp_path): + return json.loads((tmp_path / "checkpoints" / "fpm_forward.json").read_text()) + + +def _resolve(plan, cell, tmp_path): + entry = _checkpoint(tmp_path)["cells"][cell.cell_id] + return runtime_memory.resolve_runtime_resources( + cell, + Path(entry["artifact_dir"]) / "raw", + expected_plan_sha256=plan.sha256, + expected_attempt_id=entry["attempt_id"], + expected_backend_version="0.27.0", + expected_context_length=4096, + expected_max_num_tokens=1024, + expected_max_batch_size=64, + expected_gpu_memory_utilization=0.9, + expected_model_revision=plan.fpm_profile.model_revision, + ) + + +@pytest.mark.usefixtures("no_models_or_timing_data") +@pytest.mark.parametrize("missing_worker", [False, True]) +def test_slurm_profile_campaign_preserves_timing_and_resolves_only_complete_memory( + tmp_path, monkeypatch, missing_worker +): + plan = _slurm_plan(tmp_path) + cluster = _SyntheticSlurm(plan, monkeypatch) + cluster.missing_worker = missing_worker + assert _run(plan, tmp_path) == [] + checkpoint = _checkpoint(tmp_path) + assert {entry["status"] for entry in checkpoint["cells"].values()} == {"passed"} + publication = checkpoint["database"] + assert publication["status"] == "passed" + saved = runtime_memory.saved_plan_identity(plan.to_dict()) + database.validate_formal_database_commit(Path(publication["parquet"]), Path(publication["metadata"]), saved) + assert sorted(cluster.executed) == [("decode", 0), ("decode", 1), ("prefill", 0), ("prefill", 1)] + assert cluster.steps == {"1234.99": "unrelated-job"} + for cell in plan.cells: + if missing_worker: + with pytest.raises(ValueError, match="worker rank evidence is incomplete"): + _resolve(plan, cell, tmp_path) + else: + resolved = _resolve(plan, cell, tmp_path) + assert resolved["runtime_memory"]["kv_cache_bytes"] == 98 * 128 + assert len(resolved["cache_groups"]) == 3 + evidence = json.loads(resolved["runtime_memory"]["provenance"])["artifacts"] + assert {Path(item["path"]).parts[0] for item in evidence} == {"node0000", "node0001"} + before = runner._file_manifest(tmp_path / "artifacts") + commands = len(cluster.commands) + assert _run(plan, tmp_path, resume=True) == [] + assert len(cluster.commands) == commands + # Completed resume may refresh invocation metadata; it must retain raw runtime evidence. + after = runner._file_manifest(tmp_path / "artifacts") + assert {key: value for key, value in before.items() if "/raw/" in key} == { + key: value for key, value in after.items() if "/raw/" in key + } + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_slurm_profile_runtime_mismatch_preserves_observed_pin_without_benchmarking(tmp_path, monkeypatch): + plan = _slurm_plan(tmp_path) + cluster = _SyntheticSlurm(plan, monkeypatch) + cluster.version = "0.28.0" + errors = _run(plan, tmp_path) + assert len(errors) == 2 and all(error["classification"] == "campaign_cell_failed" for error in errors) + assert cluster.executed == [] + checkpoint = _checkpoint(tmp_path) + assert "database" not in checkpoint + for entry in checkpoint["cells"].values(): + cell_dir = Path(entry["artifact_dir"]) + provenance = json.loads((cell_dir / "raw/node0000/collector-provenance.json").read_text()) + assert provenance["runtime"]["backend_version"] == "0.28.0" + failures = list((cell_dir / "logs/transport-failures").glob("*/stderr.log")) + assert len(failures) == 1 and "FPM profile runtime mismatch" in failures[0].read_text() + assert cluster.steps == {"1234.99": "unrelated-job"} + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_slurm_profile_interrupt_salvages_memory_and_resumes_exact_completed_attempt(tmp_path, monkeypatch): + plan = _slurm_plan(tmp_path) + cluster = _SyntheticSlurm(plan, monkeypatch) + first = plan.cells[0] + cluster.interrupt_phase = first.workload_kind + with pytest.raises(KeyboardInterrupt): + _run(plan, tmp_path) + original = _checkpoint(tmp_path)["cells"][first.cell_id] + assert original["status"] == "interrupted" + before = runner._file_manifest(Path(original["artifact_dir"]) / "raw") + assert _resolve(plan, first, tmp_path)["runtime_memory"]["kv_cache_bytes"] == 98 * 128 + assert cluster.steps == {"1234.99": "unrelated-job"} + cancelled = [args[1] for args in cluster.commands if args[0] == "scancel"] + assert len(cancelled) == 2 and "1234.99" not in cancelled and "1234" not in cancelled + cluster.interrupt_phase = None + assert _run(plan, tmp_path, resume=True) == [] + checkpoint = _checkpoint(tmp_path) + recovered = checkpoint["cells"][first.cell_id] + assert recovered["status"] == "passed" and recovered["attempt_id"] == original["attempt_id"] + assert recovered["artifact_recovery"]["original_status"] == "interrupted" + assert runner._file_manifest(Path(recovered["artifact_dir"]) / "raw") == before + assert sorted(cluster.executed) == [("decode", 0), ("decode", 1), ("prefill", 0), ("prefill", 1)] + assert checkpoint["database"]["status"] == "passed" + + +@pytest.mark.usefixtures("no_models_or_timing_data") +def test_saved_slurm_plan_preserves_explicit_point_coverage_at_formal_consumer(tmp_path, monkeypatch): + plan = _slurm_plan(tmp_path) + _SyntheticSlurm(plan, monkeypatch) + assert _run(plan, tmp_path) == [] + publication = _checkpoint(tmp_path)["database"] + points = { + "schema_version": 1, + "prefill": [{"batch_size": 4, "total_prefill_tokens": 257, "total_kv_read_tokens": 512}], + "decode": [{"batch_size": 4, "total_kv_read_tokens": 512}], + } + path = tmp_path / "points.json" + path.write_text(json.dumps(points)) + explicit = _slurm_plan(tmp_path, points_file=path) + saved = runtime_memory.saved_plan_identity(explicit.to_dict()) + assert json.loads(saved.options.benchmark_points_json) == points + # Existing first-publisher data may have another plan identity; its actual + # coordinates must still satisfy the new saved manifest before reuse. + with pytest.raises(ValueError, match="does not cover requested coordinates"): + database.validate_formal_database_commit( + Path(publication["parquet"]), + Path(publication["metadata"]), + saved, + reused_cell_ids=[cell.cell_id for cell in saved.cells], + ) diff --git a/python/aisimulate/tests/unit/collector/test_runtime_adapter.py b/python/aisimulate/tests/unit/collector/test_runtime_adapter.py new file mode 100644 index 000000000..05ab9691f --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_runtime_adapter.py @@ -0,0 +1,682 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from collector.fpm_forward.bundled_instrumentation import bundled_instrumentation +from collector.fpm_forward.runtime_instrumentation import freeze_instrumentation, load_instrumentation + +pytestmark = pytest.mark.unit + + +def test_bundled_adapter_is_self_contained_and_unknown_version_needs_local_bundle(tmp_path): + bundle = bundled_instrumentation("0.27.0") + assert bundle.manifest["runtime"]["version"] == "0.27.0" + assert bundle.manifest["runtime"]["source_files"] + assert bundle.manifest["source_notes"] in bundle.files + frozen = freeze_instrumentation(bundle, tmp_path / "bundle") + assert load_instrumentation(frozen.manifest_path).sha256 == bundle.sha256 + assert bundled_instrumentation("0.28.0") is None + + +@pytest.mark.parametrize("version", ["0.27.0", "0.28.0"]) +def test_frozen_scheduler_export_preserves_native_benchmark_contract(tmp_path, monkeypatch, version): + import importlib + + _fake_campaign(tmp_path, monkeypatch, version=version) + frozen = load_instrumentation(tmp_path / "instrumentation/manifest.json") + scheduler_name = frozen.manifest["scheduler_class"] + # Both pinned Dynamo args.py consumers require this substring before importing + # the class. The immutable source mappings are in each bundle's source notes. + assert "InstrumentedScheduler" in scheduler_name + module, name = scheduler_name.rsplit(".", 1) + scheduler = getattr(importlib.import_module(module), name) + native = importlib.import_module("dynamo.vllm.instrumented_scheduler").InstrumentedScheduler + assert scheduler.__name__ == name + assert issubclass(scheduler, native) + + +def test_new_bundle_hooks_emit_observation_after_normal_warmup(tmp_path, monkeypatch): + import importlib + import json + import sys + from types import ModuleType, SimpleNamespace + + from .test_fpm_runtime_memory import _initialized + + bundle = freeze_instrumentation(bundled_instrumentation("0.27.0"), tmp_path / "bundle") + monkeypatch.syspath_prepend(str(bundle.root)) + modules = {} + for name in ( + "vllm", + "vllm.distributed", + "vllm.v1", + "vllm.v1.worker", + "vllm.v1.worker.gpu_worker", + "dynamo", + "dynamo.vllm", + "dynamo.vllm.instrumented_scheduler", + ): + modules[name] = ModuleType(name) + monkeypatch.setitem(sys.modules, name, modules[name]) + modules["dynamo.vllm.instrumented_scheduler"].InstrumentedScheduler = object + modules["vllm.distributed"].get_tp_group = lambda: SimpleNamespace(rank_in_group=0) + modules["vllm.distributed"].get_pp_group = lambda: SimpleNamespace(rank_in_group=0) + owner, _, cache = _initialized() + owner.vllm_config.parallel_config.data_parallel_rank = 0 + + class RuntimeWorker: + def __init__(self): + self.__dict__.update(owner.__dict__) + + def determine_available_memory(self): + return 123456 + + def initialize_from_config(self, config): + assert config is cache + return "cache initialized" + + def compile_or_warm_up_model(self): + return "normal warmup returned" + + modules["vllm.v1.worker.gpu_worker"].Worker = RuntimeWorker + output = tmp_path / "results" + output.mkdir() + monkeypatch.setenv("AISIMULATE_RUNTIME_OBSERVATION_DIR", str(output)) + context = tmp_path / "context.json" + context.write_text( + json.dumps( + { + "schema_version": "aisimulate-runtime-probe-launch/v1", + "attempt_id": "attempt-1", + "configuration": "tp2", + "phase": "prefill", + "bundle_sha256": bundle.sha256, + "launch": {"identity": {"gpu": "gb300"}, "deployment": {"image": "fake"}}, + } + ) + ) + monkeypatch.setenv("AISIMULATE_RUNTIME_CONTEXT", str(context)) + monkeypatch.setenv("AISIMULATE_RUNTIME_INSTRUMENTATION", str(bundle.manifest_path)) + module, cls = bundle.manifest["worker_class"].rsplit(".", 1) + worker = getattr(importlib.import_module(module), cls)() + assert worker.determine_available_memory() == 123456 + assert worker.initialize_from_config(cache) == "cache initialized" + assert worker.compile_or_warm_up_model() == "normal warmup returned" + records = [json.loads(path.read_text()) for path in output.glob("runtime-observation-*.json")] + assert len(records) == 1 + assert records[0]["lifecycle"] == { + "measurement": "after_warmup", + "cache_initialized": True, + "warmup_completed": True, + "capture_completed": True, + } + assert records[0]["cache"]["available_cache_bytes"] == 123456 + assert records[0]["unresolved_fields"] # Deliberately incomplete source/model/hardware fake. + + +@pytest.fixture(autouse=True) +def isolated_runtime_imports(): + import sys + + saved = { + name: module + for name, module in sys.modules.items() + if name == "fpm_memory_observer" or name == "instrumentation" or name.startswith("instrumentation.") + } + for name in saved: + del sys.modules[name] + yield + for name in list(sys.modules): + if name == "fpm_memory_observer" or name == "instrumentation" or name.startswith("instrumentation."): + del sys.modules[name] + sys.modules.update(saved) + + +def _fake_campaign(tmp_path, monkeypatch, *, packed=False, planar=False, version="0.28.0", runner_version=1): + """Synthetic vendor files and GPU objects; verifies mapping, never GPU support.""" + import hashlib + import importlib + import importlib.metadata + import json + import sys + from enum import Enum + from pathlib import Path + from types import ModuleType, SimpleNamespace + + from .test_fpm_runtime_memory import _Storage, _typed, _vllm_config + + runtime = Path(__file__).parents[3] / "collector/fpm_forward/runtime" + original = load_instrumentation(runtime / f"vllm-{version}{'.example' if version == '0.28.0' else ''}.json") + bundle = freeze_instrumentation(original, tmp_path / "instrumentation") + manifest = bundle.manifest + fake_site = tmp_path / "site-packages" + for name in manifest["runtime"]["source_files"]: + path = fake_site / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"# Synthetic runtime source for {name}\n") + manifest["runtime"]["source_files"][name] = hashlib.sha256(path.read_bytes()).hexdigest() + bundle.manifest_path.write_text(json.dumps(manifest)) + bundle = load_instrumentation(bundle.manifest_path) + model = tmp_path / "model" + model.mkdir() + model_config_path = model / "config.json" + model_config_path.write_text('{"model_type":"synthetic-test-model"}') + config = _vllm_config(tp=1, dp=1) + config.model_config.model = str(model) + config.model_config.revision = "d" * 40 + config.model_config.quantization = None + config.model_config.hf_config = SimpleNamespace(_commit_hash=None) + config.parallel_config.data_parallel_rank = 0 + config.parallel_config.enable_expert_parallel = False + config.quant_config = None + blocks, page = 10, 4096 + storage = _Storage(1024, blocks * page * (2 if packed else 1)) + + class Tensor: + def __init__(self, offset=0): + self.offset = offset + self.shape = (2, blocks, 16, 1, 64) if planar else (blocks, 2, 16, 1, 64) + + def untyped_storage(self): + return storage + + def element_size(self): + return 2 + + def storage_offset(self): + return self.offset // 2 + + def stride(self): + return (blocks * 1024, 1024, 64, 64, 1) if planar else (2048 * (2 if packed else 1), 1024, 64, 64, 1) + + class Backend: + @staticmethod + def get_name(): + return "FLASHINFER" + + @staticmethod + def get_kv_cache_block_dim(kernel, heads, head_size, *, cache_dtype_str): + assert (kernel, heads, head_size, cache_dtype_str) == (16, 1, 64, "auto") + return 1 if planar else 0 + + Backend.__module__ = "vllm.v1.attention.backends.flashinfer" + mode = Enum("KVQuantMode", ["NONE"]) + groups, layers, allocations, attn = [], {}, [], [] + for index in range(2 if packed else 1): + name = f"layer{index}" + spec = _typed("vllm.v1.kv_cache_interface", "FullAttentionSpec") + spec.block_size = spec.storage_block_size = 16 + spec.page_size_bytes, spec.sliding_window, spec.dtype = page, None, "torch.bfloat16" + spec.kv_quant_mode, spec.num_kv_heads, spec.head_size = mode.NONE, 1, 64 + layer = _typed("vllm.model_executor.layers.attention", "Attention") + layer.kv_cache = Tensor(page * index) + layers[name] = layer + groups.append(SimpleNamespace(layer_names=[name], kv_cache_spec=spec, is_eagle_group=False)) + allocations.append( + SimpleNamespace( + size=storage.size, shared_by=[name], block_stride=page * 2 if packed else 0, offset=page * index + ) + ) + attn.append([SimpleNamespace(layer_names=[name], kv_cache_spec=spec, kv_cache_group_id=index, backend=Backend)]) + cache = SimpleNamespace(num_blocks=blocks, kv_cache_groups=groups, kv_cache_tensors=allocations) + config.compilation_config.static_forward_context = layers + runner = _typed( + "vllm.v1.worker.gpu_model_runner" if runner_version == 1 else "vllm.v1.worker.gpu.model_runner", + "GPUModelRunner", + ) + runner.kv_cache_config = cache + runner.kv_caches = [layer.kv_cache for layer in layers.values()] + runner.shared_kv_cache_layers = {} + runner.attn_groups = attn + # The native initialize_kv_cache methods bind different fields. V2 never + # assigns the private V1 field, even though both classes have the same name. + setattr(runner, "_kernel_block_sizes" if runner_version == 1 else "kernel_block_sizes", [16] * len(groups)) + block_objects = [SimpleNamespace(block_id=index, is_null=index == 0) for index in range(blocks)] + pool = SimpleNamespace( + blocks=block_objects, + num_gpu_blocks=blocks, + null_block=block_objects[0], + get_num_free_blocks=lambda: blocks - 1, + free_block_queue=SimpleNamespace(get_all_free_blocks=lambda: block_objects[1:]), + ) + manager = SimpleNamespace( + block_pool=pool, + watermark_blocks=0, + coordinator=SimpleNamespace(single_type_managers=[SimpleNamespace(block_pool=pool) for group in groups]), + ) + + class RuntimeWorker: + fail = None + + def __init__(self): + self.vllm_config, self.model_runner, self.device = config, runner, "cuda:0" + + def determine_available_memory(self): + if self.fail == "profile": + raise RuntimeError("runtime profile failure") + return storage.size + page + + def initialize_from_config(self, supplied): + assert supplied is cache + if self.fail == "initialize": + raise RuntimeError("runtime initialize failure") + return "initialized" + + def compile_or_warm_up_model(self): + if self.fail == "warmup": + raise RuntimeError("runtime warmup failure") + return "warmed" + + class RuntimeScheduler: + def __init__(self, supplied_config, supplied_cache): + assert supplied_config is config and supplied_cache is cache + self.vllm_config, self.kv_cache_manager, self._fpm_dp_rank = config, manager, 0 + + for name in ( + "vllm", + "vllm._version", + "vllm.distributed", + "vllm.v1.worker.gpu_worker", + "vllm.model_executor.offloader", + "dynamo", + "dynamo.vllm.instrumented_scheduler", + "torch", + ): + module = ModuleType(name) + module.__path__ = [str(fake_site / name.replace(".", "/"))] + monkeypatch.setitem(sys.modules, name, module) + sys.modules["vllm._version"].__commit_id__ = "g" + manifest["runtime"]["source_revision"][:9] + sys.modules["vllm.distributed"].get_tp_group = lambda: SimpleNamespace(rank_in_group=0) + sys.modules["vllm.distributed"].get_pp_group = lambda: SimpleNamespace(rank_in_group=0) + sys.modules["vllm.v1.worker.gpu_worker"].Worker = RuntimeWorker + sys.modules["dynamo.vllm.instrumented_scheduler"].InstrumentedScheduler = RuntimeScheduler + sys.modules["vllm.model_executor.offloader"].get_offloader = lambda: _typed( + "vllm.model_executor.offloader.base", "NoopOffloader" + ) + sys.modules["torch"].cuda = SimpleNamespace( + get_device_properties=lambda _device: SimpleNamespace( + name="NVIDIA GB300", major=10, minor=3, total_memory=1_000_000_000 + ) + ) + old_version = importlib.metadata.version + monkeypatch.setattr(importlib.metadata, "version", lambda name: version if name == "vllm" else old_version(name)) + monkeypatch.syspath_prepend(str(bundle.root)) + module_name = bundle.manifest["worker_class"].rsplit(".", 1)[0] + adapter = importlib.import_module(module_name) + launch = { + "identity": { + "model": str(model), + "model_revision": config.model_config.revision, + "model_kind": "dense", + "framework": "vllm", + "framework_version": version, + "gpu": "gb300", + "interconnect": "nvlink", + "sm": 103, + }, + "topology": {"tp": 1, "pp": 1, "dp": 1, "moe_tp": 1, "moe_ep": 1, "cp": 1}, + "precision": { + "gemm_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16", + "fmha_quant_mode": "bfloat16", + "kvcache_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "moe_backend": "auto", + "attention_backend": "auto", + "enable_wideep": False, + "enable_eplb": False, + }, + "collection": { + "max_model_len": 4096, + "max_num_batched_tokens": 1024, + "max_num_seqs": 64, + "gpu_memory_utilization": 0.9, + "prefill_cudagraph_policy": "runtime", + "max_prefill_cudagraph_size": None, + "async_scheduling": False, + }, + "model_config": { + "path": str(model_config_path), + "sha256": hashlib.sha256(model_config_path.read_bytes()).hexdigest(), + }, + "deployment": {"executor": "slurm", "image": "synthetic/runtime@sha256:" + "c" * 64}, + } + attempt = { + "attempt_id": "attempt-1", + "bundle": {"manifest": str(bundle.manifest_path.relative_to(tmp_path)), "sha256": bundle.sha256}, + "phases": {}, + } + index = { + "schema_version": "aisimulate-runtime-observations/v1", + "configurations": {"worker": {"launch": launch, "active_attempt_id": "attempt-1", "attempts": [attempt]}}, + } + + def run(phase="prefill"): + directory = tmp_path / phase + directory.mkdir() + context = { + "schema_version": "aisimulate-runtime-probe-launch/v1", + "attempt_id": "attempt-1", + "configuration": "worker", + "phase": phase, + "bundle_sha256": bundle.sha256, + "launch": launch, + "expected_ranks": {"workers": [{"dp_rank": 0, "tp_rank": 0, "pp_rank": 0}], "schedulers": [{"dp_rank": 0}]}, + } + context_path = directory / "launch.json" + context_path.write_text(json.dumps(context)) + monkeypatch.setenv("AISIMULATE_RUNTIME_CONTEXT", str(context_path)) + monkeypatch.setenv("AISIMULATE_RUNTIME_INSTRUMENTATION", str(bundle.manifest_path)) + monkeypatch.setenv("AISIMULATE_RUNTIME_OBSERVATION_DIR", str(directory)) + worker = adapter.ObservedWorker() + assert worker.determine_available_memory() == storage.size + page + assert worker.initialize_from_config(cache) == "initialized" + assert worker.compile_or_warm_up_model() == "warmed" + adapter.ObservedInstrumentedScheduler(config, cache) + paths = sorted(directory.glob("runtime-observation-*.json")) + reference = lambda path: { + "path": str(path.relative_to(tmp_path)), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + attempt["phases"][phase] = { + "launch_manifest": reference(context_path), + "artifacts": [{**reference(path), "kind": "observation"} for path in paths], + } + (tmp_path / "observations.json").write_text(json.dumps(index)) + return [json.loads(path.read_text()) for path in paths] + + return SimpleNamespace( + run=run, + launch=launch, + config=config, + cache=cache, + pool=pool, + site=fake_site, + source=manifest, + adapter=adapter, + base=RuntimeWorker, + model_config=model_config_path, + root=tmp_path, + ) + + +@pytest.mark.parametrize( + "version,packed,planar", [("0.27.0", False, False), ("0.28.0", False, True), ("0.28.0", True, False)] +) +def test_runtime_records_import_complete_for_source_audited_layouts(tmp_path, monkeypatch, version, packed, planar): + from collector.fpm_forward.runtime_observations import validate_observations + + campaign = _fake_campaign(tmp_path, monkeypatch, version=version, packed=packed, planar=planar) + for phase in ("prefill", "decode"): + records = campaign.run(phase) + assert len(records) == 2 + assert all(record["unresolved_fields"] == [] for record in records), records + worker = next(record for record in records if record["kind"] == "worker") + assert worker["cache"]["layer_tensors"]["layer0"]["block_axis"] == (1 if planar else 0) + assert worker["model_config_sha256"] == campaign.launch["model_config"]["sha256"] + result = validate_observations(tmp_path / "observations.json", {"worker": campaign.launch})["worker"] + assert result["status"] == "complete", result["diagnostics"] + + +def test_bundled_lifecycle_through_verified_entrypoint_imports_both_phases(tmp_path, monkeypatch): + import importlib + import json + import os + import sys + from pathlib import Path + + from collector.fpm_forward.runtime_instrumentation import runtime_binding + from collector.fpm_forward.runtime_observations import validate_observations + from collector.fpm_forward.runtime_probe import observation_artifacts + + campaign = _fake_campaign(tmp_path, monkeypatch, version="0.27.0") + frozen = load_instrumentation(tmp_path / "instrumentation/manifest.json") + runtime = Path(__file__).parents[3] / "collector/fpm_forward/runtime" + monkeypatch.syspath_prepend(str(runtime)) + for phase in ("prefill", "decode"): + # Each native phase is a fresh process. Retain the fixture's fake runtime + # dependencies while clearing only adapters before invoking the entrypoint. + for name in list(sys.modules): + if name in {"fpm_memory_observer", "fpm_runtime_instrumentation", "instrumentation"} or name.startswith( + "instrumentation." + ): + monkeypatch.delitem(sys.modules, name) + original_adapter = campaign.adapter + entrypoint = importlib.import_module("fpm_runtime_instrumentation") + + def worker(): + context_path = Path(os.environ["AISIMULATE_RUNTIME_CONTEXT"]) + context = json.loads(context_path.read_text()) + context["instrumentation_binding"] = runtime_binding() + context_path.write_text(json.dumps(context)) + return entrypoint.ObservedWorker() + + # The fixture closes over its adapter module, so redirect those two + # exports to the runtime entrypoint rather than bypassing native hooks. + monkeypatch.setattr(original_adapter, "ObservedWorker", worker) + monkeypatch.setattr( + original_adapter, + "ObservedInstrumentedScheduler", + lambda *args: entrypoint.ObservedInstrumentedScheduler(*args), + ) + try: + records = campaign.run(phase) + assert all(record["instrumentation_binding"]["status"] == "verified" for record in records) + assert all(record["instrumentation_binding"]["files"] == frozen.files for record in records) + assert issubclass(entrypoint.ObservedWorker, campaign.base) + # Probe and formal collection use this same artifact enumerator. + references = observation_artifacts(tmp_path, tmp_path / phase) + assert any("runtime-instrumentation-imports-" in item["path"] for item in references) + index_path = tmp_path / "observations.json" + index = json.loads(index_path.read_text()) + index["configurations"]["worker"]["attempts"][0]["phases"][phase]["artifacts"] = references + index_path.write_text(json.dumps(index)) + finally: + if entrypoint._bundle is not None: + sys.meta_path.remove(entrypoint._bundle) + monkeypatch.delitem(sys.modules, "fpm_runtime_instrumentation") + result = validate_observations(tmp_path / "observations.json", {"worker": campaign.launch})["worker"] + assert result["status"] == "complete", result + assert result["provenance"]["instrumentation"]["import_binding"] == "verified" + + +@pytest.mark.parametrize("version", ["0.27.0", "0.28.0"]) +@pytest.mark.parametrize("runner_version", [1, 2]) +def test_native_runner_selection_uses_its_initialized_kernel_block_sizes( + tmp_path, monkeypatch, version, runner_version +): + from collector.fpm_forward.runtime_observations import validate_observations + + campaign = _fake_campaign(tmp_path, monkeypatch, version=version, runner_version=runner_version) + runner = campaign.adapter.ObservedWorker().model_runner + other_field = "kernel_block_sizes" if runner_version == 1 else "_kernel_block_sizes" + assert not hasattr(runner, other_field) + for phase in ("prefill", "decode"): + records = campaign.run(phase) + assert all(record["unresolved_fields"] == [] for record in records), records + worker = next(record for record in records if record["kind"] == "worker") + assert worker["cache"]["layer_tensors"]["layer0"]["kernel_block_size_tokens"] == 16 + assert worker["cache"]["model_runner"]["class"] == f"{type(runner).__module__}.GPUModelRunner" + result = validate_observations(tmp_path / "observations.json", {"worker": campaign.launch})["worker"] + assert result["status"] == "complete", result["diagnostics"] + + +@pytest.mark.parametrize("failure", ["unknown_runner", "unaudited_source", "missing_field"]) +def test_runner_mapping_never_guesses_a_compatible_field(tmp_path, monkeypatch, failure): + import importlib + + campaign = _fake_campaign(tmp_path, monkeypatch, runner_version=2) + runner = campaign.adapter.ObservedWorker().model_runner + if failure == "unknown_runner": + type(runner).__module__ = "vendor.unreviewed_runner" + elif failure == "unaudited_source": + campaign.source["runtime"]["source_files"].pop("vllm/v1/worker/gpu/model_runner.py", None) + else: + del runner.kernel_block_sizes + runner._kernel_block_sizes = [16] + observer = importlib.import_module("instrumentation.observer") + worker = campaign.adapter.ObservedWorker() + worker.determine_available_memory() + worker.initialize_from_config(campaign.cache) + with pytest.raises((ValueError, AttributeError)): + observer._layer_views( + worker, + observer.legacy.worker_memory(worker), + campaign.source["runtime"]["source_files"], + ) + + +def test_changed_vendor_source_is_unresolved_with_raw_measurements_retained(tmp_path, monkeypatch): + import hashlib + + campaign = _fake_campaign(tmp_path, monkeypatch) + source = campaign.site / "vllm/v1/worker/gpu_model_runner.py" + source.write_text("# Synthetic unreviewed vendor patch\n") + for record in campaign.run(): + assert "runtime" in record["unresolved_fields"] + assert "source hash differs" in record["error"] + assert ( + record["runtime"]["source_files"]["vllm/v1/worker/gpu_model_runner.py"] + == hashlib.sha256(source.read_bytes()).hexdigest() + ) + assert record["cache"]["num_blocks"] == 10 + + +@pytest.mark.parametrize("extra", [None, 8]) +def test_v028_extra_retention_is_not_silently_discarded(tmp_path, monkeypatch, extra): + campaign = _fake_campaign(tmp_path, monkeypatch) + spec = campaign.cache.kv_cache_groups[0].kv_cache_spec + spec.__class__.__name__ = "SlidingWindowSpec" + spec.sliding_window = 128 + if extra is not None: + spec.extra_retained_tokens = extra + for record in campaign.run(): + assert "cache.groups" in record["unresolved_fields"] + assert record["cache"]["groups"][0]["extra_retained_tokens"] == extra + + +def test_v027_absent_extra_retention_maps_to_audited_zero(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch, version="0.27.0") + spec = campaign.cache.kv_cache_groups[0].kv_cache_spec + spec.__class__.__name__ = "SlidingWindowSpec" + spec.sliding_window = 128 + for record in campaign.run(): + assert record["unresolved_fields"] == [] + assert record["cache"]["groups"][0]["extra_retained_tokens"] == 0 + + +@pytest.mark.parametrize("failure", ["profile", "initialize", "warmup"]) +def test_underlying_worker_exceptions_propagate_without_completion_evidence(tmp_path, monkeypatch, failure): + campaign = _fake_campaign(tmp_path, monkeypatch) + campaign.base.fail = failure + with pytest.raises(RuntimeError, match=f"runtime {failure} failure"): + campaign.run() + assert not list(tmp_path.rglob("runtime-observation-*.json")) + + +def test_loaded_config_hash_comes_from_actual_runtime_path(tmp_path, monkeypatch): + import hashlib + + campaign = _fake_campaign(tmp_path, monkeypatch) + campaign.model_config.write_text('{"model_type":"changed-runtime-config"}') + for record in campaign.run(): + assert record["model_config_sha256"] == hashlib.sha256(campaign.model_config.read_bytes()).hexdigest() + assert record["model_config_sha256"] != campaign.launch["model_config"]["sha256"] + + +def test_backend_without_source_mapping_is_unresolved(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch) + worker = campaign.adapter.ObservedWorker() + worker.model_runner.attn_groups[0][0].backend.__module__ = "vendor.unaudited_backend" + record = next(record for record in campaign.run() if record["kind"] == "worker") + assert "cache.layer_tensors" in record["unresolved_fields"] + assert "backend source" in record["error"] + + +def test_speculative_runtime_is_unresolved_without_erasing_raw_config(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch) + campaign.config.speculative_config = {"method": "synthetic-extra-tokens"} + record = next(record for record in campaign.run() if record["kind"] == "worker") + assert "cache.semantics" in record["unresolved_fields"] + assert record["resolved_config"]["speculative_config"] == campaign.config.speculative_config + assert record["cache"]["semantics"]["speculative"] is True + + +def test_non_hf_loaded_config_is_not_bound_to_incidental_config_json(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch) + campaign.config.model_config.config_format = "mistral" + for record in campaign.run(): + assert "model_config_sha256" in record["unresolved_fields"] + assert "model_config_sha256" not in record + + +def test_formal_runtime_hook_preserves_existing_execution_evidence(tmp_path, monkeypatch): + import json + + campaign = _fake_campaign(tmp_path, monkeypatch) + original_init = campaign.base.__init__ + + def initialized(self): + original_init(self) + (tmp_path / "prefill" / "collector-provenance.json").write_text(json.dumps({"cell_id": "synthetic"})) + + monkeypatch.setattr(campaign.base, "__init__", initialized) + records = campaign.run() + assert all(record["unresolved_fields"] == [] for record in records) + files = list((tmp_path / "prefill").glob("fpm-execution-worker-*.json")) + assert len(files) == 1 + evidence = json.loads(files[0].read_text()) + assert evidence["status"] == "observed" + assert evidence["attention_groups"][0]["backend_class"].startswith("vllm.v1.attention.backends.flashinfer.") + + +def test_unknown_cache_spec_keeps_its_raw_runtime_fields(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch) + spec = campaign.cache.kv_cache_groups[0].kv_cache_spec + spec.__class__.__name__ = "NewUnsupportedSpec" + records = campaign.run() + assert all("cache" in record["unresolved_fields"] for record in records) + assert all( + record["raw_cache_config"]["groups"][0]["spec_type"].endswith(".NewUnsupportedSpec") for record in records + ) + assert all(record["raw_cache_config"]["groups"][0]["page_size_bytes"] == 4096 for record in records) + + +def test_underlying_scheduler_exception_propagates_without_scheduler_record(tmp_path, monkeypatch): + import sys + + campaign = _fake_campaign(tmp_path, monkeypatch) + scheduler = sys.modules["dynamo.vllm.instrumented_scheduler"].InstrumentedScheduler + + def fail(*args, **kwargs): + raise RuntimeError("native scheduler initialization failed") + + monkeypatch.setattr(scheduler, "__init__", fail) + with pytest.raises(RuntimeError, match="native scheduler initialization failed"): + campaign.run() + assert not list(tmp_path.rglob("runtime-observation-scheduler-*.json")) + + +def test_model_sidecar_hashes_are_read_from_the_actual_loaded_config_directory(tmp_path, monkeypatch): + import hashlib + + campaign = _fake_campaign(tmp_path, monkeypatch) + sidecar = campaign.model_config.parent / "hf_quant_config.json" + sidecar.write_text('{"quantization":{"quant_algo":"NVFP4"}}') + campaign.launch["model_config"]["source_files"] = {"hf_quant_config.json": "0" * 64} + for record in campaign.run(): + assert record["model_config_source_files"] == { + "hf_quant_config.json": hashlib.sha256(sidecar.read_bytes()).hexdigest() + } + assert record["model_config_sha256"] == campaign.launch["model_config"]["sha256"] + + +def test_missing_model_sidecar_keeps_primary_config_evidence_unresolved(tmp_path, monkeypatch): + campaign = _fake_campaign(tmp_path, monkeypatch) + campaign.launch["model_config"]["source_files"] = {"hf_quant_config.json": "0" * 64} + for record in campaign.run(): + assert record["unresolved_fields"] + assert "hf_quant_config.json" in record["error"] + assert record["model_config_sha256"] == campaign.launch["model_config"]["sha256"] diff --git a/python/aisimulate/tests/unit/collector/test_runtime_instrumentation.py b/python/aisimulate/tests/unit/collector/test_runtime_instrumentation.py new file mode 100644 index 000000000..4d0c6d26a --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_runtime_instrumentation.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from collector.fpm_forward.runtime_instrumentation import freeze_instrumentation, load_instrumentation + +pytestmark = pytest.mark.unit + + +def _manifest(tmp_path: Path): + (tmp_path / "observer.py").write_text("raise AssertionError('observer must never execute on the host')\n") + (tmp_path / "source.md").write_text("Source mapping with bytes and lifecycle points.\n") + value = { + "schema_version": "aisimulate-runtime-instrumentation/v1", + "runtime": {"framework": "vllm", "version": "0.28.0", "source_revision": "a" * 40}, + "files": ["observer.py", "source.md"], + "source_notes": "source.md", + "worker_class": "observer.ObservedWorker", + "scheduler_class": "observer.ObservedScheduler", + "observation_schema": "aisimulate-runtime-observation/v1", + } + path = tmp_path / "manifest.yaml" + path.write_text(json.dumps(value)) + return path, value + + +def test_load_and_freeze_preserve_hashes_without_executing_observer(tmp_path): + path, value = _manifest(tmp_path) + bundle = load_instrumentation(path, expected_version="0.28.0") + frozen = freeze_instrumentation(bundle, tmp_path / "frozen") + assert frozen.sha256 == bundle.sha256 + assert frozen.manifest == value + assert frozen.manifest_path == tmp_path / "frozen" / "manifest.json" + assert frozen.files.keys() == {"observer.py", "source.md"} + (tmp_path / "observer.py").write_text("raise RuntimeError('changed')\n") + assert load_instrumentation(path).sha256 != frozen.sha256 + assert load_instrumentation(frozen.manifest_path).sha256 == frozen.sha256 + with pytest.raises(ValueError, match="changed"): + freeze_instrumentation(bundle, tmp_path / "another") + with pytest.raises(ValueError, match="fresh"): + freeze_instrumentation(frozen, tmp_path / "frozen") + + +@pytest.mark.parametrize( + "name", ["../observer.py", "/observer.py", "./observer.py", "x//observer.py", "x\\observer.py"] +) +def test_declared_paths_are_unambiguous_and_contained(tmp_path, name): + path, value = _manifest(tmp_path) + value["files"].append(name) + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="relative|normalized"): + load_instrumentation(path) + + +def test_symlinks_and_duplicate_modules_are_rejected(tmp_path): + path, value = _manifest(tmp_path) + (tmp_path / "alias.py").symlink_to("observer.py") + value["files"].append("alias.py") + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="symlink"): + load_instrumentation(path) + value["files"].remove("alias.py") + (tmp_path / "observer").mkdir() + (tmp_path / "observer" / "__init__.py").write_text("") + value["files"].append("observer/__init__.py") + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="ambiguous"): + load_instrumentation(path) + + +def test_version_unknown_keys_duplicate_keys_and_undeclared_modules_rejected(tmp_path): + path, value = _manifest(tmp_path) + with pytest.raises(ValueError, match="version"): + load_instrumentation(path, expected_version="0.27.0") + value["worker_class"] = "not_declared.Worker" + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="declared"): + load_instrumentation(path) + path.write_text('{"schema_version":"x", "schema_version":"y"}') + with pytest.raises(ValueError, match="duplicate"): + load_instrumentation(path) + path.write_text("schema_version: x\nschema_version: y\n") + with pytest.raises(ValueError, match="duplicate"): + load_instrumentation(path) diff --git a/python/aisimulate/tests/unit/collector/test_runtime_instrumentation_binding.py b/python/aisimulate/tests/unit/collector/test_runtime_instrumentation_binding.py new file mode 100644 index 000000000..e85e96eec --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_runtime_instrumentation_binding.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise the actual rendered entrypoints from a Slurm-like working directory.""" + +from __future__ import annotations + +import json +import os +import py_compile +import shutil +import subprocess +import sys +import zipfile + +import pytest +from collector.fpm_forward import runner, runtime_probe +from collector.fpm_forward.runtime_instrumentation import freeze_instrumentation, load_instrumentation + +from .test_fpm_runtime_probe import _inputs + +pytestmark = pytest.mark.unit + + +def _staged(tmp_path): + launch, manifest, _ = _inputs(tmp_path) + package = tmp_path / "adapter" + package.mkdir() + (package / "__init__.py").write_text("") + (package / "classes.py").write_text( + "import fpm_memory_observer as helper\n" + "from .constants import ENABLED\n" + "class ObservedWorker:\n" + " def marker(self): return helper.MARKER if ENABLED else 'disabled'\n" + "class ObservedInstrumentedScheduler(ObservedWorker): pass\n" + ) + (package / "constants.py").write_text("ENABLED = True\n") + (package / "other_runtime.py").write_text("raise RuntimeError('unused adapter must not be imported')\n") + (tmp_path / "fpm_memory_observer.py").write_text("MARKER = 'custom frozen helper'\n") + value = json.loads(manifest.read_text()) + value.update( + files=[ + "adapter/__init__.py", + "adapter/classes.py", + "adapter/constants.py", + "adapter/other_runtime.py", + "fpm_memory_observer.py", + "source-notes.txt", + ], + worker_class="adapter.classes.ObservedWorker", + scheduler_class="adapter.classes.ObservedInstrumentedScheduler", + ) + manifest.write_text(json.dumps(value)) + frozen = freeze_instrumentation(load_instrumentation(manifest), tmp_path / "frozen") + plan = runtime_probe.build_runtime_probe_plan("tp4", launch, frozen) + cell = plan.cells[0] + workdir = tmp_path / "slurm-workdir" + workdir.mkdir() + context = runtime_probe.launch_context(plan, cell, configuration="tp4", attempt_id="test-attempt") + for path in runtime_probe.stage_runtime_instrumentation(frozen, workdir, context): + if path.parent != workdir: + shutil.copy2(path, workdir / path.name) + with zipfile.ZipFile(workdir / "runtime-instrumentation.zip") as archive: + archive.extractall(workdir / "runtime-instrumentation") + (workdir / "fpm_memory_observer.py").write_text("MARKER = 'wrong Slurm cwd helper'\n") + args = runner._cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + classes = [args[args.index(flag) + 1] for flag in ("--worker-cls", "--scheduler-cls")] + env = { + **os.environ, + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHONPATH": os.pathsep.join([str(workdir / "runtime-instrumentation"), str(workdir)]), + "AISIMULATE_RUNTIME_CONTEXT": str(workdir / "runtime-probe-context.json"), + "AISIMULATE_RUNTIME_INSTRUMENTATION": str(workdir / "runtime-instrumentation" / "manifest.json"), + "AISIMULATE_RUNTIME_OBSERVATION_DIR": str(workdir / "results"), + } + return workdir, env, classes, frozen + + +def _run(workdir, env, classes, prefix=""): + script = prefix + "\nimport importlib, json\nclasses = " + repr(classes) + "\n" + script += ( + "values = []\n" + "for name in classes:\n" + " module, _, attr = name.rpartition('.')\n" + " cls = getattr(importlib.import_module(module), attr)\n" + " values.append(cls().marker())\n" + "print(json.dumps(values))\n" + ) + return subprocess.run([sys.executable, "-c", script], cwd=workdir, env=env, capture_output=True, text=True) + + +def test_rendered_worker_and_scheduler_use_frozen_helper_in_slurm_cwd(tmp_path): + workdir, env, classes, frozen = _staged(tmp_path) + result = _run(workdir, env, classes) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == ["custom frozen helper", "custom frozen helper"] + receipts = list((workdir / "results").glob("runtime-instrumentation-imports-*.json")) + assert len(receipts) == 1 + receipt = json.loads(receipts[0].read_text()) + assert receipt["status"] == "verified" + assert receipt["bundle_sha256"] == frozen.sha256 + assert receipt["files"] == frozen.files + assert "adapter.constants" in receipt["imports"] + assert "adapter.other_runtime" not in receipt["imports"] + helper = receipt["imports"]["fpm_memory_observer"] + assert helper["path"] == str(workdir / "runtime-instrumentation" / "fpm_memory_observer.py") + assert helper["sha256"] == frozen.files["fpm_memory_observer.py"] + + +@pytest.mark.parametrize("changed", ["fpm_memory_observer.py", "source-notes.txt"]) +def test_runtime_bundle_tampering_fails_before_adapter_import(tmp_path, changed): + workdir, env, classes, _ = _staged(tmp_path) + with (workdir / "runtime-instrumentation" / changed).open("a") as handle: + handle.write("\n# changed after staging\n") + result = _run(workdir, env, classes) + assert result.returncode != 0 + assert "bundle hash mismatch" in result.stderr + receipt = json.loads(next((workdir / "results").glob("runtime-instrumentation-imports-*.json")).read_text()) + assert receipt["status"] == "failed" + assert not receipt["imports"] + + +def test_wrong_helper_imported_before_runtime_entrypoint_is_rejected(tmp_path): + workdir, env, classes, _ = _staged(tmp_path) + result = _run(workdir, env, classes, "import fpm_memory_observer\n") + assert result.returncode != 0 + assert "already imported outside verified bundle loading" in result.stderr + + +def test_verified_source_is_executed_instead_of_unchecked_cached_bytecode(tmp_path): + workdir, env, classes, _ = _staged(tmp_path) + helper = workdir / "runtime-instrumentation" / "fpm_memory_observer.py" + original = helper.read_bytes() + helper.write_text("MARKER = 'wrong cached helper'\n") + py_compile.compile(str(helper), invalidation_mode=py_compile.PycInvalidationMode.UNCHECKED_HASH, doraise=True) + helper.write_bytes(original) + result = _run(workdir, env, classes) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == ["custom frozen helper", "custom frozen helper"] + + +def test_spawned_worker_and_scheduler_unpickling_use_verified_entrypoint(tmp_path): + workdir, env, classes, _ = _staged(tmp_path) + script = workdir / "spawn_check.py" + script.write_text( + "import importlib, multiprocessing\n" + "def child(cls, queue):\n" + " queue.put(cls().marker())\n" + "if __name__ == '__main__':\n" + " context = multiprocessing.get_context('spawn')\n" + f" for name in {classes!r}:\n" + " module, _, attr = name.rpartition('.')\n" + " cls = getattr(importlib.import_module(module), attr)\n" + " original = importlib.import_module('adapter.classes')\n" + " assert isinstance(cls(), original.ObservedWorker)\n" + " queue = context.Queue()\n" + " process = context.Process(target=child, args=(cls, queue))\n" + " process.start()\n" + " assert queue.get(timeout=10) == 'custom frozen helper'\n" + " process.join(10)\n" + " assert process.exitcode == 0\n" + ) + result = subprocess.run([sys.executable, str(script)], cwd=workdir, env=env, capture_output=True, text=True) + assert result.returncode == 0, result.stderr + receipts = [json.loads(p.read_text()) for p in (workdir / "results").glob("runtime-instrumentation-imports-*.json")] + assert len(receipts) == 3 + assert all(item["status"] == "verified" for item in receipts) + assert len({item["pid"] for item in receipts}) == 3 + + +def test_frozen_package_resources_use_standard_importlib_reader(tmp_path): + workdir, env, classes, _ = _staged(tmp_path) + root = workdir / "runtime-instrumentation" + (root / "adapter/data.txt").write_text("frozen package data") + (root / "adapter/classes.py").write_text( + "from importlib.resources import files\n" + "class ObservedWorker:\n" + " def marker(self): return files(__package__).joinpath('data.txt').read_text()\n" + "class ObservedInstrumentedScheduler(ObservedWorker): pass\n" + ) + manifest = root / "manifest.json" + value = json.loads(manifest.read_text()) + value["files"].append("adapter/data.txt") + manifest.write_text(json.dumps(value)) + frozen = load_instrumentation(manifest) + context_path = workdir / "runtime-probe-context.json" + context = json.loads(context_path.read_text()) + context["bundle_sha256"] = frozen.sha256 + context_path.write_text(json.dumps(context)) + result = _run(workdir, env, classes) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == ["frozen package data", "frozen package data"] + receipt = json.loads(next((workdir / "results").glob("runtime-instrumentation-imports-*.json")).read_text()) + assert receipt["files"]["adapter/data.txt"] == frozen.files["adapter/data.txt"] + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="requires a fork-capable platform") +def test_forked_worker_preserves_imports_with_its_own_receipt(tmp_path): + workdir, env, _classes, _ = _staged(tmp_path) + script = ( + "import multiprocessing, os, json\n" + "import fpm_runtime_instrumentation as binding\n" + "cls = binding.ObservedWorker\n" + "parent = binding.observed_binding()\n" + "def child(queue):\n" + " queue.put({'actual_pid': os.getpid(), 'binding': binding.observed_binding(), 'value': cls().marker()})\n" + "context = multiprocessing.get_context('fork')\n" + "queue = context.Queue()\n" + "process = context.Process(target=child, args=(queue,))\n" + "process.start()\n" + "child_result = queue.get(timeout=10)\n" + "process.join(10)\n" + "assert process.exitcode == 0\n" + "print(json.dumps({'parent': parent, 'child': child_result}))\n" + ) + result = subprocess.run([sys.executable, "-c", script], cwd=workdir, env=env, capture_output=True, text=True) + assert result.returncode == 0, result.stderr + value = json.loads(result.stdout) + parent, child = value["parent"], value["child"] + assert child["binding"]["pid"] == child["actual_pid"] != parent["pid"] + assert child["binding"]["imports"] == parent["imports"] + assert child["value"] == "custom frozen helper" + receipts = [ + json.loads(path.read_text()) for path in (workdir / "results").glob("runtime-instrumentation-imports-*.json") + ] + assert len(receipts) == 2 + assert next(item for item in receipts if item["pid"] == parent["pid"]) == parent + assert next(item for item in receipts if item["pid"] == child["actual_pid"]) == child["binding"] + + +@pytest.mark.parametrize( + "mutation", ["digest", "origin", "module-type", "classes-type", "missing-context", "pid", "hostname"] +) +def test_importer_requires_observed_binding_for_new_context_and_keeps_legacy_unverified(tmp_path, mutation): + from collector.fpm_forward.runtime_instrumentation import runtime_binding + from collector.fpm_forward.runtime_observations import validate_observations + + from .test_runtime_observations import _mutate, _write, observation_fixture + + path, launches = observation_fixture(tmp_path) + old = validate_observations(path, launches)["tp2"] + assert old["status"] == "complete" + assert old["provenance"]["instrumentation"]["import_binding"] == "unverified_legacy" + index = json.loads(path.read_text()) + attempt = index["configurations"]["tp2"]["attempts"][0] + for phase in attempt["phases"].values(): + reference = phase["launch_manifest"] + launch_path = path.parent / reference["path"] + context = json.loads(launch_path.read_text()) + context["instrumentation_binding"] = runtime_binding() + reference["sha256"] = _write(launch_path, context)["sha256"] + path.write_text(json.dumps(index)) + missing = validate_observations(path, launches)["tp2"] + assert missing["status"] == "incomplete" + assert "verified instrumentation imports" in str(missing["diagnostics"]) + + bundle = load_instrumentation(path.parent / attempt["bundle"]["manifest"]) + + def add_receipt(record): + # Synthetic observations exercise CPU validation only; runtime execution + # is independently covered by the rendered-entrypoint subprocess tests. + record["instrumentation_binding"] = { + **runtime_binding(), + **{key: record[key] for key in ("attempt_id", "configuration", "phase", "bundle_sha256")}, + "status": "verified", + "expected_bundle_sha256": bundle.sha256, + "hostname": "synthetic-host", + "pid": 123, + "root": "/frozen", + "files": bundle.files, + "imports": { + "observer": { + "file": "observer.py", + "path": "/frozen/observer.py", + "sha256": bundle.files["observer.py"], + } + }, + "classes": {key: bundle.manifest[key] for key in ("worker_class", "scheduler_class")}, + } + record["cpu_affinity_observation"] = {"status": "complete", "hostname": "synthetic-host", "pid": 123} + + _mutate(path, add_receipt) + new = validate_observations(path, launches)["tp2"] + assert new["status"] == "complete", new + assert new["provenance"]["instrumentation"]["import_binding"] == "verified" + if mutation == "missing-context": + index = json.loads(path.read_text()) + for phase in index["configurations"]["tp2"]["attempts"][0]["phases"].values(): + reference = phase["launch_manifest"] + launch_path = path.parent / reference["path"] + context = json.loads(launch_path.read_text()) + context.pop("instrumentation_binding") + reference["sha256"] = _write(launch_path, context)["sha256"] + path.write_text(json.dumps(index)) + else: + + def change(record): + evidence = record["instrumentation_binding"] + if mutation == "digest": + evidence["imports"]["observer"]["sha256"] = "c" * 64 + elif mutation == "origin": + evidence["imports"]["observer"]["path"] = "/cwd/observer.py" + elif mutation == "module-type": + evidence["imports"]["observer"] = [] + elif mutation == "pid": + evidence["pid"] = 124 + elif mutation == "hostname": + evidence["hostname"] = "different-host" + else: + evidence["classes"] = [] + + _mutate(path, change) + changed = validate_observations(path, launches)["tp2"] + assert changed["status"] == "incomplete" + assert "instrumentation" in str(changed["diagnostics"]) diff --git a/python/aisimulate/tests/unit/collector/test_runtime_observation_precisions.py b/python/aisimulate/tests/unit/collector/test_runtime_observation_precisions.py new file mode 100644 index 000000000..b050cd488 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_runtime_observation_precisions.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from collector.fpm_forward.runtime import fpm_memory_observer as observer +from collector.fpm_forward.runtime_observations import validate_observations + +from .test_fpm_runtime_memory import _typed +from .test_runtime_observations import _mutate, _replace_launch, observation_fixture + +pytestmark = pytest.mark.unit + + +def _precision_fixture(tmp_path, *, cache_dtype="fp8", storage_dtype="torch.uint8", item_size=1, precision="fp8"): + path, launches = observation_fixture(tmp_path, version="0.27.0") + launches["tp2"]["precision"]["kvcache_quant_mode"] = precision + _replace_launch(path, launches) + spec = _typed("vllm.v1.kv_cache_interface", "FullAttentionSpec") + spec.block_size, spec.page_size_bytes, spec.dtype = 16, 64, storage_dtype + observed_dtype = observer.cache_groups( + SimpleNamespace(kv_cache_groups=[SimpleNamespace(layer_names=["layer0"], kv_cache_spec=spec)]) + )[0]["dtype"] + + def change(record): + record["resolved_config"]["cache_config"]["cache_dtype"] = cache_dtype + for group in record["cache"]["groups"]: + group["dtype"] = observed_dtype + for tensor in record["cache"].get("layer_tensors", {}).values(): + tensor["shape"][-1] = 64 // item_size + tensor["stride_bytes"][-1] = item_size + tensor["element_size_bytes"] = item_size + + _mutate(path, change) + return path, launches + + +@pytest.mark.parametrize("cache_dtype", ["fp8", "fp8_e4m3"]) +def test_fp8_uint8_storage_import_uses_resolved_precision_and_preserves_physical_evidence(tmp_path, cache_dtype): + path, launches = _precision_fixture(tmp_path, cache_dtype=cache_dtype) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["resources"]["runtime_memory"]["kv_cache_bytes"] == 83 * 128 + provenance = result["provenance"] + assert provenance["runtime_settings"]["cache_config"]["cache_dtype"] == cache_dtype + assert len(provenance["artifacts"]) == 12 + for artifact in provenance["artifacts"]: + record = artifact["evidence"] + assert {group["dtype"] for group in record["cache"]["groups"]} == {"torch.uint8"} + for tensor in record["cache"].get("layer_tensors", {}).values(): + assert tensor["element_size_bytes"] == 1 + + +@pytest.mark.parametrize( + "cache_dtype,storage_dtype,precision", + [("bfloat16", "torch.bfloat16", "bfloat16"), ("float16", "torch.float16", "half")], +) +def test_unquantized_cache_storage_import_remains_supported(tmp_path, cache_dtype, storage_dtype, precision): + path, launches = _precision_fixture( + tmp_path, cache_dtype=cache_dtype, storage_dtype=storage_dtype, item_size=2, precision=precision + ) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["resources"]["runtime_memory"]["kv_cache_bytes"] == 83 * 128 + + +@pytest.mark.parametrize( + "cache_dtype", + [None, "auto", "bfloat16", "fp8_e5m2", "fp8_per_token_head", "int4_per_token_head", "nvfp4"], +) +def test_uint8_storage_does_not_establish_missing_or_other_semantic_precision(tmp_path, cache_dtype): + path, launches = _precision_fixture(tmp_path, cache_dtype=cache_dtype) + if cache_dtype is None: + _mutate(path, lambda record: record["resolved_config"]["cache_config"].pop("cache_dtype")) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert "KV precision" in " ".join(result["diagnostics"]) + + +def test_uint8_storage_is_not_accepted_for_validated_bf16_cache(tmp_path): + path, launches = _precision_fixture(tmp_path, cache_dtype="bfloat16", precision="bfloat16") + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert "cache group dtype" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("storage_dtype", ["torch.uint8", "fp8"]) +def test_fp8_storage_requires_one_byte_tensor_elements(tmp_path, storage_dtype): + path, launches = _precision_fixture(tmp_path, storage_dtype=storage_dtype, item_size=2) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert "tensor element size" in " ".join(result["diagnostics"]) + + +def test_fp8_semantics_cannot_relabel_bf16_storage(tmp_path): + path, launches = _precision_fixture(tmp_path, storage_dtype="torch.bfloat16", item_size=2) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert "cache group dtype" in " ".join(result["diagnostics"]) diff --git a/python/aisimulate/tests/unit/collector/test_runtime_observations.py b/python/aisimulate/tests/unit/collector/test_runtime_observations.py new file mode 100644 index 000000000..c5e6627f9 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_runtime_observations.py @@ -0,0 +1,811 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +from collector.fpm_forward import runtime_memory +from collector.fpm_forward.runtime import fpm_memory_observer as legacy +from collector.fpm_forward.runtime_instrumentation import load_instrumentation +from collector.fpm_forward.runtime_observations import validate_observations + +from aisimulate_core.fpm_profile import FpmResourceProfile + +from .test_fpm_runtime_memory import _vllm_config +from .test_runtime_instrumentation import _manifest + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("observed", [None, {}, {"hf_quant_config.json": "b" * 64}, {"hf_quant_config.json": "a" * 64}]) +def test_import_independently_verifies_model_config_source_hashes(tmp_path, observed): + path, launches = observation_fixture(tmp_path) + launches["tp2"]["model_config"]["source_files"] = {"hf_quant_config.json": "a" * 64} + _replace_launch(path, launches) + if observed is not None: + _mutate(path, lambda record: record.update(model_config_source_files=observed)) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == ("complete" if observed == {"hf_quant_config.json": "a" * 64} else "incomplete") + if result["status"] == "incomplete": + assert "source" in str(result["diagnostics"]) + + +def test_cpu_sidecars_are_archived_without_changing_memory_import(tmp_path): + from collector.fpm_forward.runtime_probe import observation_artifacts + + path, launches = observation_fixture(tmp_path) + baseline = validate_observations(path, launches)["tp2"] + cpu = tmp_path / "prefill/fpm-cpu-scheduler-dp0.json" + cpu.write_text( + json.dumps({"schema_name": "aisimulate_fpm_cpu_affinity", "schema_version": 1, "status": "unavailable"}) + ) + references = observation_artifacts(tmp_path, tmp_path / "prefill") + reference = next(item for item in references if item["path"] == "prefill/fpm-cpu-scheduler-dp0.json") + assert reference["kind"] == "runtime-artifact" + index = json.loads(path.read_text()) + index["configurations"]["tp2"]["attempts"][0]["phases"]["prefill"]["artifacts"].append(reference) + path.write_text(json.dumps(index)) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "complete" + assert result["resources"]["cache_groups"] == baseline["resources"]["cache_groups"] + assert ( + result["resources"]["runtime_memory"]["kv_cache_bytes"] + == baseline["resources"]["runtime_memory"]["kv_cache_bytes"] + ) + assert reference in result["provenance"]["runtime_artifacts"] + + +def _write(path, payload): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload)) + return {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} + + +def observation_fixture(root: Path, *, packed=False, planar=False, dense=False, tp=2, version="0.28.0"): + """Synthetic source, rank and runtime records; no GPU or runtime qualification.""" + root.mkdir(parents=True, exist_ok=True) + path, manifest = _manifest(root) + manifest["runtime"]["version"] = version + manifest["runtime"]["source_files"] = {"vllm/v1/core/block_pool.py": "b" * 64} + path.write_text(json.dumps(manifest)) + bundle = load_instrumentation(path) + model = root / "model.json" + model.write_text('{"model_type":"example"}') + launch = { + "identity": { + "model": "example/model", + "model_revision": "immutable", + "model_kind": "dense" if dense else "moe", + "framework": "vllm", + "framework_version": version, + "gpu": "gb300", + "interconnect": "nvlink", + "sm": 103, + }, + "topology": { + "tp": tp, + "pp": 1, + "dp": 1 if dense else 2, + "moe_tp": 1, + "moe_ep": 1 if dense else tp * 2, + "cp": 1, + }, + "precision": { + "gemm_quant_mode": "nvfp4", + "moe_quant_mode": "nvfp4", + "fmha_quant_mode": "bfloat16", + "kvcache_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "moe_backend": "auto", + "attention_backend": "auto", + "enable_wideep": False, + "enable_eplb": False, + }, + "collection": { + "max_model_len": 4096, + "max_num_batched_tokens": 1024, + "max_num_seqs": 64, + "gpu_memory_utilization": 0.9, + "prefill_cudagraph_policy": "runtime", + "max_prefill_cudagraph_size": None, + "async_scheduling": False, + }, + "model_config": {"path": str(model), "sha256": hashlib.sha256(model.read_bytes()).hexdigest()}, + "deployment": {"executor": "slurm", "image": "example/runtime@sha256:" + "c" * 64}, + } + attempt = { + "attempt_id": "attempt-001", + "bundle": {"manifest": "manifest.yaml", "sha256": bundle.sha256}, + "phases": {}, + } + index = { + "schema_version": "aisimulate-runtime-observations/v1", + "configurations": {"tp2": {"launch": launch, "active_attempt_id": "attempt-001", "attempts": [attempt]}}, + } + for phase in ("prefill", "decode"): + expected = { + "workers": [ + {"dp_rank": dp, "tp_rank": tp, "pp_rank": 0} + for dp in range(launch["topology"]["dp"]) + for tp in range(launch["topology"]["tp"]) + ], + "schedulers": [{"dp_rank": dp} for dp in range(launch["topology"]["dp"])], + } + context = { + "schema_version": "aisimulate-runtime-probe-launch/v1", + "configuration": "tp2", + "attempt_id": "attempt-001", + "phase": phase, + "bundle_sha256": bundle.sha256, + "launch": launch, + "expected_ranks": expected, + } + ref = _write(root / phase / "launch.json", context) + ref["path"] = str(Path(ref["path"]).relative_to(root)) + result = {"launch_manifest": ref, "artifacts": []} + attempt["phases"][phase] = result + for rank in expected["workers"] + [{**r, "tp_rank": None, "pp_rank": None} for r in expected["schedulers"]]: + worker = rank["tp_rank"] is not None + dp = rank["dp_rank"] + count = 100 - dp * 10 - (5 if phase == "decode" else 0) + groups = [ + { + "layer_names": [name], + "spec_type": "vllm.v1.kv_cache_interface.FullAttentionSpec", + "block_size_tokens": 16, + "spec_page_size_bytes": 64, + "sliding_window": None, + "extra_retained_tokens": 0, + "attention_chunk_size": None, + "dtype": "torch.bfloat16", + "pool_id": "pool", + } + for name in ("layer0", "layer1") + ] + # Different cache groups reuse one shared pool storage; a second + # slot represents padding or a simultaneously needed group layer. + if packed: + groups[0]["layer_names"].append("layer2") + cache = {"num_blocks": count, "groups": copy.deepcopy(groups), "pool_id": "pool"} + if worker: + cache.update( + { + "available_cache_bytes": count * 128 + 256, + "allocated_cache_bytes": count * 128, + "semantics": { + "allocation": "shared_block_pool", + "storage": "hbm", + "retention": "full_or_window", + "prefix_reuse": False, + "offload": False, + "speculative": False, + }, + "storages": [ + {"storage_id": "s0", "device": "cuda:0", "pointer": 1024, "size_bytes": count * 128} + ], + "tensor_allocations": [ + { + "storage_id": "s0", + "size": count * 128, + "shared_by": ["layer0", "layer1"], + "block_stride": 128 if packed else 0, + "offset": 0, + } + ], + "layer_tensors": {}, + } + ) + if packed: + cache["tensor_allocations"].append( + { + "storage_id": "s0", + "size": count * 128, + "shared_by": ["layer2"], + "block_stride": 128, + "offset": 64, + } + ) + for group in cache["groups"]: + group["kind"] = "attention" + group["layer_classes"] = dict.fromkeys( + group["layer_names"], "vllm.model_executor.layers.attention.Attention" + ) + for name in group["layer_names"]: + cache["layer_tensors"][name] = { + "storage_id": "s0", + "storage_offset_bytes": 64 if name == "layer2" else 0, + "shape": [2, count, 16] if planar else [count, 32], + "stride_bytes": [count * 32, 32, 2] if planar else [128, 2], + "element_size_bytes": 2, + "block_axis": 1 if planar else 0, + "kernel_block_size_tokens": 16, + } + else: + cache.update( + { + "initial_free_blocks": count - 2, + "reserved_blocks": 2, + "null_block_id": 0, + "watermark_blocks": 0, + "pool_count": 1, + "permanent_reserved_block_ids": [0, 1], + "group_pool_ids": ["pool"] * len(groups), + } + ) + record = { + "schema_version": "aisimulate-runtime-observation/v1", + **{k: context[k] for k in ("configuration", "attempt_id", "phase", "bundle_sha256")}, + "kind": "worker" if worker else "scheduler", + **rank, + "identity": launch["identity"], + "model_config_sha256": launch["model_config"]["sha256"], + "runtime": {**manifest["runtime"], "image": launch["deployment"]["image"]}, + "launch": launch, + "hardware": { + "gpu": "gb300", + "sm": 103, + "device_name": "NVIDIA GB300", + "total_memory_bytes": 1000000000, + }, + "resolved_config": legacy.resolved_config(_vllm_config(dp=launch["topology"]["dp"])), + "cache": cache, + "lifecycle": { + "measurement": "after_warmup", + "cache_initialized": True, + "warmup_completed": True, + "capture_completed": True, + } + if worker + else {"measurement": "scheduler_initialized", "cache_initialized": True}, + "effective_offloader": "vllm.model_executor.offloader.base.NoopOffloader", + "unresolved_fields": [], + } + record["resolved_config"]["parallel_config"]["enable_expert_parallel"] = not dense + record["resolved_config"]["parallel_config"]["tensor_parallel_size"] = tp + ref = _write( + root / phase / (f"worker-{dp}-{rank['tp_rank']}.json" if worker else f"scheduler-{dp}.json"), record + ) + ref["path"] = str(Path(ref["path"]).relative_to(root)) + result["artifacts"].append({**ref, "kind": "observation"}) + _write(root / "observations.json", index) + return root / "observations.json", {"tp2": launch} + + +def _mutate(index_path, edit): + index = json.loads(index_path.read_text()) + attempt = index["configurations"]["tp2"]["attempts"][0] + for phase in attempt["phases"].values(): + for artifact in phase["artifacts"]: + path = index_path.parent / artifact["path"] + record = json.loads(path.read_text()) + edit(record) + path.write_text(json.dumps(record)) + artifact["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + index_path.write_text(json.dumps(index)) + + +def _replace_launch(index_path, launches): + launch = launches["tp2"] + index = json.loads(index_path.read_text()) + index["configurations"]["tp2"]["launch"] = launch + for phase in index["configurations"]["tp2"]["attempts"][0]["phases"].values(): + path = index_path.parent / phase["launch_manifest"]["path"] + context = json.loads(path.read_text()) + context["launch"] = launch + path.write_text(json.dumps(context)) + phase["launch_manifest"]["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + index_path.write_text(json.dumps(index)) + _mutate(index_path, lambda r: r.update(launch=launch, identity=launch["identity"])) + + +@pytest.mark.parametrize("worker_type,mode", [("prefill", "PIECEWISE"), ("decode", "FULL_DECODE_ONLY")]) +def test_role_observations_keep_their_own_graph_settings_and_capacity(tmp_path, worker_type, mode): + path, launches = observation_fixture(tmp_path) + launch = launches["tp2"] + launch["worker_type"] = worker_type + launch["collection"]["cudagraph_mode"] = mode + _replace_launch(path, launches) + index = json.loads(path.read_text()) + phases = index["configurations"]["tp2"]["attempts"][0]["phases"] + phases.pop("decode" if worker_type == "prefill" else "prefill") + path.write_text(json.dumps(index)) + _mutate(path, lambda record: record["resolved_config"]["compilation_config"].update(cudagraph_mode=mode)) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "complete", result + assert result["resources"]["runtime_memory"]["kv_cache_bytes"] == (88 if worker_type == "prefill" else 83) * 128 + assert {value["phase"] for value in result["provenance"]["rank_capacities"]} == {worker_type} + # A correctly hashed worker record still cannot contradict its role's + # accepted engine settings. Other roles are not substitutes for evidence. + _mutate(path, lambda record: record["resolved_config"]["compilation_config"].update(cudagraph_mode="FULL")) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert "reviewed serving role" in str(result["diagnostics"]) + + +def test_aggregated_generic_graph_settings_are_checked_in_both_phases(tmp_path): + path, launches = observation_fixture(tmp_path) + launch = launches["tp2"] + launch["worker_type"] = "aggregated" + launch["collection"].update(cudagraph_mode="FULL_AND_PIECEWISE", cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32, 64]) + _replace_launch(path, launches) + _mutate( + path, + lambda record: record["resolved_config"]["compilation_config"].update( + cudagraph_mode="FULL_AND_PIECEWISE", + cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32, 64], + max_cudagraph_capture_size=64, + ), + ) + assert validate_observations(path, launches)["tp2"]["status"] == "complete" + _mutate( + path, + lambda record: record["resolved_config"]["compilation_config"].update( + cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32], max_cudagraph_capture_size=32 + ) + if record["phase"] == "decode" + else None, + ) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert "reviewed serving role" in str(result["diagnostics"]) + + +@pytest.mark.parametrize("packed,planar", [(False, False), (True, False), (False, True)]) +def test_import_recomputes_storage_and_pool_capacity_and_keeps_every_rank(tmp_path, packed, planar): + path, launch = observation_fixture(tmp_path, packed=packed, planar=planar) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + resources = FpmResourceProfile.model_validate(result["resources"]) + assert resources.runtime_memory.kv_cache_bytes == 83 * 128 + assert [group.page_size_bytes for group in resources.cache_groups] == [128, 128] + assert len(result["provenance"]["rank_capacities"]) == 8 + assert len(result["provenance"]["artifacts"]) == 12 + + +@pytest.mark.parametrize("kernel_subdivision", [False, True]) +def test_aliased_layers_cannot_assign_same_bytes_to_different_pool_blocks(tmp_path, kernel_subdivision): + path, launch = observation_fixture(tmp_path) + + def collide(record): + if record["kind"] == "worker": + view = record["cache"]["layer_tensors"]["layer1"] + if kernel_subdivision: + view.update( + shape=[record["cache"]["num_blocks"] * 2, 16], + stride_bytes=[32, 2], + kernel_block_size_tokens=8, + ) + else: + view["stride_bytes"] = [64, 2] + + _mutate(path, collide) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert "pool block" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("planar", [False, True]) +def test_aliases_preserve_pool_ownership_after_kernel_subdivision_and_permutation(tmp_path, planar): + path, launch = observation_fixture(tmp_path, planar=planar) + + def equivalent_ownership(record): + if record["kind"] == "worker": + count = record["cache"]["num_blocks"] + record["cache"]["layer_tensors"]["layer1"].update( + shape=[8, 2, count * 2] if planar else [count * 2, 16], + stride_bytes=[2, count * 32, 16] if planar else [64, 2], + block_axis=2 if planar else 0, + kernel_block_size_tokens=8, + ) + + _mutate(path, equivalent_ownership) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["resources"]["runtime_memory"]["kv_cache_bytes"] == 83 * 128 + + +@pytest.mark.parametrize("maximum", [16, 64]) +def test_explicit_graph_launch_must_match_initialized_capture_sizes(tmp_path, maximum): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["collection"].update(prefill_cudagraph_policy="explicit", max_prefill_cudagraph_size=maximum) + _replace_launch(path, launch) + # The synthetic runtime captured powers of two through 64. An explicit + # limit of 16 is ignored; explicit 64 also requires the intervening sizes + # emitted by the collector's existing prefill sampling policy. + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert "explicit prefill graph" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize( + "maximum,tokens,sizes", + [ + (16, 1024, [1, 2, 4, 8, 16]), + (64, 1024, [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64]), + (64, 32, [1, 2, 4, 8, 16, 24, 32]), + ], +) +@pytest.mark.parametrize("fallback", [None, "PIECEWISE", "NONE"]) +def test_explicit_graph_capture_contract_allows_verified_worker_mode_downgrades( + tmp_path, maximum, tokens, sizes, fallback +): + path, launch = observation_fixture(tmp_path) + sequences = min(tokens, 64) + launch["tp2"]["collection"].update( + prefill_cudagraph_policy="explicit", + max_prefill_cudagraph_size=maximum, + max_num_batched_tokens=tokens, + max_num_seqs=sequences, + ) + _replace_launch(path, launch) + + def captured(record): + config = record["resolved_config"] + config["scheduler_config"].update(max_num_batched_tokens=tokens, max_num_seqs=sequences) + config["compilation_config"].update(cudagraph_capture_sizes=sizes, max_cudagraph_capture_size=sizes[-1]) + if fallback and record["kind"] == "worker": + record["initial_compilation_config"] = copy.deepcopy(config["compilation_config"]) + config["compilation_config"]["cudagraph_mode"] = fallback + + _mutate(path, captured) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["provenance"]["runtime_settings"]["compilation_config"]["cudagraph_mode"] == ( + fallback or "FULL_AND_PIECEWISE" + ) + + +def test_explicit_prefill_does_not_impose_capture_override_on_decode(tmp_path): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["collection"].update(prefill_cudagraph_policy="explicit", max_prefill_cudagraph_size=16) + _replace_launch(path, launch) + _mutate( + path, + lambda r: r["resolved_config"]["compilation_config"].update( + cudagraph_capture_sizes=[1, 2, 4, 8, 16], max_cudagraph_capture_size=16 + ) + if r["phase"] == "prefill" + else None, + ) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + # Decode legitimately keeps runtime defaults. Its different graph memory + # still makes the two phases incompatible for a shared resource profile. + assert "across phases/ranks" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("declared_eager", [False, True]) +def test_explicit_graph_override_is_omitted_only_for_declared_eager_execution(tmp_path, declared_eager): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["collection"].update(prefill_cudagraph_policy="explicit", max_prefill_cudagraph_size=16) + if declared_eager: + launch["tp2"]["collection"]["enforce_eager"] = True + _replace_launch(path, launch) + + def eager(record): + record["resolved_config"]["model_config"]["enforce_eager"] = True + record["resolved_config"]["compilation_config"].update( + mode="NONE", cudagraph_mode="NONE", cudagraph_capture_sizes=[], max_cudagraph_capture_size=0 + ) + + _mutate(path, eager) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == ("complete" if declared_eager else "incomplete"), result["diagnostics"] + + +def test_deployed_checkpoint_path_keeps_public_identity_and_loaded_config_binding(tmp_path): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["deployment"].update(executor="kubernetes", model_cache="models:/cache:checkpoint") + _replace_launch(path, launch) + _mutate(path, lambda r: r["resolved_config"]["model_config"].update(model="/cache/checkpoint")) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["provenance"]["launch"]["identity"]["model"] == "example/model" + assert result["provenance"]["runtime_settings"]["model_config"]["model"] == "/cache/checkpoint" + + +@pytest.mark.parametrize( + "executor,model_cache,cache_root", + [ + ("slurm", None, "/root/.cache/huggingface/hub"), + ("slurm", None, "/custom/cache"), + ("kubernetes", "models", "/cache"), + ], +) +def test_cached_hf_snapshot_keeps_public_identity_and_loaded_config_binding( + tmp_path, executor, model_cache, cache_root +): + path, launches = observation_fixture(tmp_path) + launch = launches["tp2"] + revision = "a" * 40 + launch["identity"]["model_revision"] = revision + launch["deployment"].update(executor=executor, model_cache=model_cache) + _replace_launch(path, launches) + snapshot = f"{cache_root}/models--example--model/snapshots/{revision}" + _mutate( + path, + lambda record: record["resolved_config"]["model_config"].update( + model=snapshot, revision=revision, loaded_config_commit_hash=revision + ), + ) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert result["provenance"]["launch"]["identity"]["model"] == "example/model" + assert result["provenance"]["runtime_settings"]["model_config"]["model"] == snapshot + + +@pytest.mark.parametrize( + "corruption", + [ + "repo", + "snapshot_revision", + "symbolic_revision", + "local_path", + "relative_path", + "nested_path", + "parent_path", + "blob_path", + "explicit_checkpoint", + "hash", + "sidecar", + "revision", + "loaded_revision", + "identity", + ], +) +def test_cached_hf_snapshot_does_not_weaken_model_identity(tmp_path, corruption): + path, launches = observation_fixture(tmp_path) + launch = launches["tp2"] + revision = "main" if corruption == "symbolic_revision" else "a" * 40 + launch["identity"]["model_revision"] = revision + launch["model_config"]["source_files"] = {"hf_quant_config.json": "b" * 64} + if corruption == "explicit_checkpoint": + launch["deployment"].update(executor="kubernetes", model_cache="models:/cache:checkpoint") + _replace_launch(path, launches) + + def mismatch(record): + model = record["resolved_config"]["model_config"] + snapshot = f"/cache/models--example--model/snapshots/{revision}" + if corruption == "repo": + snapshot = snapshot.replace("models--example--model", "models--another--model") + elif corruption == "snapshot_revision": + snapshot = snapshot.replace(revision, "c" * 40) + elif corruption == "local_path": + snapshot = "/unexpected/checkpoint" + elif corruption == "relative_path": + snapshot = snapshot.lstrip("/") + elif corruption == "nested_path": + snapshot += "/checkpoint" + elif corruption == "parent_path": + snapshot = "/cache/../" + snapshot.lstrip("/") + elif corruption == "blob_path": + snapshot = snapshot.replace("/snapshots/", "/blobs/") + model.update(model=snapshot, revision=revision, loaded_config_commit_hash=revision) + record["model_config_source_files"] = {"hf_quant_config.json": "b" * 64} + if corruption == "hash": + record["model_config_sha256"] = "f" * 64 + elif corruption == "sidecar": + record["model_config_source_files"]["hf_quant_config.json"] = "f" * 64 + elif corruption == "revision": + model["revision"] = "f" * 40 + elif corruption == "loaded_revision": + model["loaded_config_commit_hash"] = "f" * 40 + elif corruption == "identity": + record["identity"]["model"] = "another/model" + + _mutate(path, mismatch) + result = validate_observations(path, launches)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + + +@pytest.mark.parametrize("corruption", ["path", "hash", "revision", "missing_mount", "parent_subpath", "slurm"]) +def test_deployed_checkpoint_mapping_does_not_weaken_model_identity(tmp_path, corruption): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["deployment"].update(executor="kubernetes", model_cache="models:/cache:checkpoint") + if corruption == "missing_mount": + launch["tp2"]["deployment"]["model_cache"] = "models::checkpoint" + elif corruption == "parent_subpath": + launch["tp2"]["deployment"]["model_cache"] = "models:/cache:../checkpoint" + elif corruption == "slurm": + launch["tp2"]["deployment"]["executor"] = "slurm" + _replace_launch(path, launch) + + def mismatch(record): + model = record["resolved_config"]["model_config"] + model["model"] = "/cache/another" if corruption == "path" else "/cache/checkpoint" + if corruption == "hash": + record["model_config_sha256"] = "f" * 64 + elif corruption == "revision": + model["revision"] = "another" + + _mutate(path, mismatch) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + + +def test_cache_only_pvc_preserves_public_model_path(tmp_path): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["deployment"].update(executor="kubernetes", model_cache="models") + _replace_launch(path, launch) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + + +def test_observed_expert_parallelism_cannot_certify_partial_tensor_expert_split(tmp_path): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["topology"].update(moe_tp=2, moe_ep=2) + _replace_launch(path, launch) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert "expert topology" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("enabled,moe_tp,moe_ep", [(True, 1, 4), (False, 4, 1)]) +def test_expert_axes_follow_resolved_ep_setting_and_attention_worker_width(tmp_path, enabled, moe_tp, moe_ep): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["topology"].update(moe_tp=moe_tp, moe_ep=moe_ep) + _replace_launch(path, launch) + _mutate(path, lambda r: r["resolved_config"]["parallel_config"].update(enable_expert_parallel=enabled)) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert len(result["provenance"]["rank_capacities"]) == 8 + + +@pytest.mark.parametrize( + "mutation,diagnostic", + [ + (lambda r: r.update(attempt_id="another"), "attempt"), + (lambda r: r["runtime"].update(source_revision="e" * 40), "build"), + (lambda r: r["runtime"]["source_files"].update({"vllm/v1/core/block_pool.py": "e" * 64}), "build"), + (lambda r: r["resolved_config"]["model_config"].update(quantization=None), "precision"), + (lambda r: r["lifecycle"].update(warmup_completed=False) if r["kind"] == "worker" else None, "warmup"), + (lambda r: r["cache"]["groups"][0].update(extra_retained_tokens=3), "retention"), + (lambda r: r["cache"].update(group_pool_ids=["other", "pool"]) if r["kind"] == "scheduler" else None, "pool"), + (lambda r: r["cache"]["storages"][0].update(size_bytes=1) if r["kind"] == "worker" else None, "storage"), + ( + lambda r: r["cache"]["layer_tensors"]["layer0"].update(stride_bytes=[2, 2]) + if r["kind"] == "worker" + else None, + "overlap", + ), + ( + lambda r: r["cache"]["tensor_allocations"][0].update(shared_by=["layer0"]) + if r["kind"] == "worker" + else None, + "layer", + ), + ( + lambda r: r["cache"].update(permanent_reserved_block_ids=[0]) if r["kind"] == "scheduler" else None, + "reservation", + ), + (lambda r: r["cache"]["semantics"].update(offload=True) if r["kind"] == "worker" else None, "semantics"), + ], +) +def test_import_independently_rejects_corrupt_or_unsupported_evidence(tmp_path, mutation, diagnostic): + path, launch = observation_fixture(tmp_path) + _mutate(path, mutation) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["resources"] is None + assert diagnostic in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("mode", ["missing", "duplicate", "hash", "phase", "graph"]) +def test_rank_phase_hash_and_graph_coverage_is_required(tmp_path, mode): + path, launch = observation_fixture(tmp_path) + index = json.loads(path.read_text()) + phases = index["configurations"]["tp2"]["attempts"][0]["phases"] + if mode == "missing": + phases["prefill"]["artifacts"].pop() + elif mode == "duplicate": + phases["prefill"]["artifacts"].append(phases["prefill"]["artifacts"][0]) + elif mode == "hash": + phases["prefill"]["artifacts"][0]["sha256"] = "f" * 64 + elif mode == "phase": + phases.pop("decode") + path.write_text(json.dumps(index)) + if mode == "graph": + _mutate( + path, + lambda r: r["resolved_config"]["compilation_config"].update(cudagraph_mode="PIECEWISE") + if r["phase"] == "decode" + else None, + ) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert result["diagnostics"] + + +def test_other_configurations_keep_successful_result_and_legacy_guard_survives(tmp_path): + path, launch = observation_fixture(tmp_path) + result = validate_observations(path, {**launch, "missing": launch["tp2"]}) + assert result["tp2"]["status"] == "complete", result + assert result["missing"]["status"] == "incomplete" + with pytest.raises(ValueError, match="currently supports vLLM 0.27.0"): + runtime_memory.resolve_runtime_resources( + None, + tmp_path, + expected_plan_sha256="p", + expected_attempt_id="a", + expected_backend_version="0.28.0", + expected_context_length=4096, + expected_max_num_tokens=1024, + expected_max_batch_size=64, + expected_gpu_memory_utilization=0.9, + ) + + +def test_dense_tp_configuration_does_not_require_moe_axes_to_equal_worker_width(tmp_path): + path, launch = observation_fixture(tmp_path, dense=True) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "complete", result["diagnostics"] + assert len(result["provenance"]["rank_capacities"]) == 4 + + +def test_packaged_hardware_is_checked_when_request_omits_optional_sm(tmp_path): + path, launch = observation_fixture(tmp_path) + launch["tp2"]["identity"].pop("sm") + index = json.loads(path.read_text()) + index["configurations"]["tp2"]["launch"] = launch["tp2"] + for phase in index["configurations"]["tp2"]["attempts"][0]["phases"].values(): + context_path = path.parent / phase["launch_manifest"]["path"] + context = json.loads(context_path.read_text()) + context["launch"] = launch["tp2"] + context_path.write_text(json.dumps(context)) + phase["launch_manifest"]["sha256"] = hashlib.sha256(context_path.read_bytes()).hexdigest() + path.write_text(json.dumps(index)) + + def wrong_hardware(record): + record["identity"] = launch["tp2"]["identity"] + record["launch"] = launch["tp2"] + record["hardware"].update(sm=90, device_name="NVIDIA H100 80GB HBM3") + + _mutate(path, wrong_hardware) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert "hardware" in " ".join(result["diagnostics"]) + + +@pytest.mark.parametrize("change", ["missing", "different"]) +def test_loaded_model_config_hash_is_independently_required(tmp_path, change): + path, launch = observation_fixture(tmp_path) + _mutate( + path, lambda r: r.pop("model_config_sha256") if change == "missing" else r.update(model_config_sha256="f" * 64) + ) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert "model configuration" in " ".join(result["diagnostics"]) + + +def test_changed_bundle_code_invalidates_saved_observations_without_execution(tmp_path): + path, launch = observation_fixture(tmp_path) + (tmp_path / "observer.py").write_text("raise RuntimeError('must never execute')\n") + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert "bundle hash" in " ".join(result["diagnostics"]) + + +def test_invented_observer_verdict_does_not_authorize_invalid_physical_evidence(tmp_path): + path, launch = observation_fixture(tmp_path, packed=True) + + def corrupt(record): + record["passed"] = True + record["status"] = "validated" + if record["kind"] == "worker": + record["cache"]["tensor_allocations"][1]["offset"] = 0 + record["cache"]["layer_tensors"]["layer2"]["storage_offset_bytes"] = 0 + + _mutate(path, corrupt) + result = validate_observations(path, launch)["tp2"] + assert result["status"] == "incomplete" + assert "overlap" in " ".join(result["diagnostics"]) diff --git a/python/aisimulate/tests/unit/generator/test_fpm_artifacts.py b/python/aisimulate/tests/unit/generator/test_fpm_artifacts.py index 9b5994cab..bca678e79 100644 --- a/python/aisimulate/tests/unit/generator/test_fpm_artifacts.py +++ b/python/aisimulate/tests/unit/generator/test_fpm_artifacts.py @@ -396,9 +396,14 @@ def _pod_template_labels(artifacts: dict[str, str], kind: str) -> list[dict]: @pytest.mark.parametrize("kind", ["Pod", "LeaderWorkerSet", "PodCliqueSet"]) -def test_fpm_cell_label_reaches_every_pod_template(kind): +def test_fpm_ownership_labels_reach_every_resource_and_pod_template(kind): params = _params() - params["K8sConfig"]["fpm_resource_labels"] = {FPM_CELL_LABEL: "glm52-fpm-cell-0"} + ownership = { + "aiconfigurator.nvidia.com/owned-by": "fpm-forward-collector", + "aiconfigurator.nvidia.com/plan": "a" * 16, + FPM_CELL_LABEL: "glm52-fpm-cell-0", + } + params["K8sConfig"]["fpm_resource_labels"] = ownership if kind != "Pod": params["NodeConfig"].update({"system_name": "gb200", "num_gpus_per_node": 4}) params["WorkerConfig"]["agg_gpus_per_worker"] = 8 @@ -406,11 +411,49 @@ def test_fpm_cell_label_reaches_every_pod_template(kind): if kind == "PodCliqueSet": params["K8sConfig"]["fpm_orchestrator"] = "grove" - label_sets = _pod_template_labels(_render(params), kind) + artifacts = _render(params) + label_sets = [document["metadata"]["labels"] for document in _k8s_documents(artifacts)] + label_sets.extend(_pod_template_labels(artifacts, kind)) assert label_sets for labels in label_sets: - assert labels[FPM_CELL_LABEL] == "glm52-fpm-cell-0" + assert ownership.items() <= labels.items() + + +def test_fpm_compute_domain_preserves_metadata_and_source_document(monkeypatch): + from aisimulate.generator.builders import fpm_builder + from aisimulate.generator.builders.dgd_model import ComputeDomainDoc + + params = _params() + params["NodeConfig"].update({"system_name": "gb200", "num_gpus_per_node": 4}) + params["WorkerConfig"]["agg_gpus_per_worker"] = 8 + params["params"]["agg"].update({"gpus_per_worker": 8, "tensor_parallel_size": 8}) + params["K8sConfig"]["fpm_resource_labels"] = {FPM_CELL_LABEL: "glm52-fpm-cell-0"} + original_build = fpm_builder.build_dgd + sources = [] + + def build_with_metadata(*args, **kwargs): + documents = original_build(*args, **kwargs) + domain = next(document for document in documents if isinstance(document, ComputeDomainDoc)) + domain.metadata_extra = { + "labels": {"infrastructure.example.com/owner": "gpu-team"}, + "annotations": {"infrastructure.example.com/description": "reserved fabric"}, + "finalizers": ["resource.nvidia.com/domain"], + } + sources.append((domain, domain.to_dict())) + return documents + + monkeypatch.setattr(fpm_builder, "build_dgd", build_with_metadata) + artifacts = _render(params) + domain = _k8s_document(artifacts, "ComputeDomain") + workload = _k8s_document(artifacts, "LeaderWorkerSet") + assert len(sources) == 1 + source, before = sources[0] + expected = copy.deepcopy(before) + expected["metadata"]["labels"].update(workload["metadata"]["labels"]) + assert domain == expected + assert source.to_dict() == before + assert domain["metadata"]["name"] != workload["metadata"]["name"] @pytest.mark.parametrize( @@ -921,6 +964,10 @@ def test_fpm_multinode_worker_emits_keepalive_leaderworkerset_and_rank_aware_scr "metadata": { "name": "glm52-fpm-agg-compute-domain", "namespace": "default", + "labels": { + "app.kubernetes.io/name": "glm52-fpm-agg", + "app.kubernetes.io/component": "fpm-resource", + }, }, "spec": { "channel": { @@ -981,6 +1028,10 @@ def test_fpm_multinode_gb200_grove_emits_compute_domain_and_keepalive_podcliques "metadata": { "name": "glm52-fpm-agg-compute-domain", "namespace": "default", + "labels": { + "app.kubernetes.io/name": "glm52-fpm-agg", + "app.kubernetes.io/component": "fpm-resource", + }, }, "spec": { "channel": { diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py b/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py index 8cc8f117a..affc3479d 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py @@ -256,7 +256,7 @@ def test_migration_rejects_unknown_and_invalid_controls(saved, options, reason): ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "prefill", options) -@pytest.mark.parametrize("options", [{"text_only": False}, {"unrecorded_quant_modes": []}]) +@pytest.mark.parametrize("options", [{"text_only": False}, {"unrecorded_quant_modes": []}, {"collect_coverage": False}]) def test_explicit_skipped_defaults_migrate_without_becoming_null(options): config = ForwardPassPerfModelConfig.from_legacy_engine_config( _legacy_config({"fpm_interpolation": options}), "prefill" diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py b/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py index 63cfd2f9a..2601b057a 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py @@ -141,6 +141,65 @@ def test_direct_constructor_rejects_missing_genuine_phase(phase_roots, missing_p assert root in message +@pytest.mark.parametrize("worker_type", ["prefill", "decode"]) +@pytest.mark.parametrize("missing_rows", ["absent", "fake"]) +def test_single_role_requires_only_its_own_genuine_phase(phase_roots, worker_type, missing_rows): + config, make_root = phase_roots + other = "decode" if worker_type == "prefill" else "prefill" + root = make_root("single-role", **{other: missing_rows}) + config.update(worker_type=worker_type, systems_paths=[root]) + config["fpm_profile"]["deployments"][0]["worker_type"] = worker_type + queries = { + "prefill": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}, + "decode": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}, + } + for request in (config, json.loads(json.dumps(config))): + model = RustForwardPassPerfModel.best_available(request) + try: + assert model.diagnostics()["readiness"] == "ready" + assert model.estimate_forward_pass_time_ms({"scheduled_requests": queries[worker_type]}) == ( + 2.0 if worker_type == "prefill" else 3.0 + ) + with pytest.raises(PerfDataNotAvailableError): + model.estimate_forward_pass_time_ms({"scheduled_requests": queries[other]}) + finally: + model.close() + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode"]) +def test_single_role_does_not_accept_other_roles_measurements(phase_roots, worker_type): + config, make_root = phase_roots + root = make_root("wrong-role", **{worker_type: "absent"}) + config.update(worker_type=worker_type, systems_paths=[root]) + config["fpm_profile"]["deployments"][0]["worker_type"] = worker_type + with pytest.raises(PerfDataNotAvailableError, match=f"direct FPM {worker_type}"): + RustForwardPassPerfModel.best_available(config) + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode"]) +def test_single_role_pins_root_with_its_phase_without_merging_other_root(phase_roots, worker_type): + config, make_root = phase_roots + other = "decode" if worker_type == "prefill" else "prefill" + wrong = make_root("wrong", **{worker_type: "absent"}) + matching = make_root("matching", **{other: "absent"}) + config.update(worker_type=worker_type, systems_paths=[wrong, matching]) + config["fpm_profile"]["deployments"][0]["worker_type"] = worker_type + model = RustForwardPassPerfModel.best_available(config) + try: + provenance = model.diagnostics()["provenance"] + assert provenance["selected_systems_root"] == matching + assert provenance["config"]["systems_paths"] == [matching] + absent_workload = ( + {"num_decode_requests": 1, "sum_decode_kv_tokens": 1} + if other == "decode" + else {"num_prefill_requests": 1, "sum_prefill_tokens": 1} + ) + with pytest.raises(PerfDataNotAvailableError): + model.estimate_forward_pass_time_ms({"scheduled_requests": absent_workload}) + finally: + model.close() + + @pytest.mark.parametrize("missing_phase", ["prefill", "decode"]) @pytest.mark.parametrize("rows", ["absent", "fake"]) @pytest.mark.parametrize("fallback", ["deny", "allow"]) diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index e50c457c1..92ace2255 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -579,6 +579,70 @@ def test_dense_tp_does_not_require_expert_tensor_parallelism(profile_dict): assert load_fpm_profile(profile_dict).deployments[0].parallel_tuple == (2, 1, 1, 1, 1, 1) +def test_role_profiles_preserve_separate_resources_and_legacy_wire_identity(profile_dict): + legacy = load_fpm_profile(profile_dict).model_dump_json() + assert "worker_type" not in legacy + assert load_fpm_profile(legacy).model_dump_json() == legacy + prefill = profile_dict["deployments"][0] + prefill["worker_type"] = "prefill" + decode = copy.deepcopy(prefill) + decode["worker_type"] = "decode" + decode["resources"].update(weights_bytes=180, max_batch_size=128) + profile_dict["deployments"].append(decode) + profile = load_fpm_profile(load_fpm_profile(profile_dict).model_dump_json()) + identity = dict( + model=profile.model, system="test_gpu", backend="vllm", backend_version="0.25.1", tp_size=2, moe_tp_size=2 + ) + assert profile.select(**identity, worker_type="prefill").resources.weights_bytes == 100 + selected = profile.select(**identity, worker_type="decode") + assert selected.resources.weights_bytes == 180 + assert selected.resources.max_batch_size == 128 + with pytest.raises(ValueError, match="worker_type=aggregated"): + profile.select(**identity) + for role in ("prefill", "decode", "aggregated"): + assert load_fpm_profile(legacy).select(**identity, worker_type=role).worker_type is None + + +@pytest.mark.parametrize("other_role", [None, "prefill"]) +def test_role_profiles_reject_ambiguous_or_duplicate_identity(profile_dict, other_role): + prefill = profile_dict["deployments"][0] + prefill["worker_type"] = "prefill" + other = copy.deepcopy(prefill) + other["worker_type"] = other_role + profile_dict["deployments"].append(other) + with pytest.raises(ValidationError, match="ambiguous|duplicate"): + load_fpm_profile(profile_dict) + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) +def test_direct_profile_compilation_and_saved_config_preserve_role(profile_dict, direct_compile, worker_type): + profile_dict["deployments"][0]["worker_type"] = worker_type + spec = direct_compile(profile_dict, worker_type=worker_type) + assert spec["engine"]["extra"]["worker_type"] == worker_type + migrated = ForwardPassPerfModelConfig.from_legacy_engine_config(spec["engine"], worker_type).to_dict() + assert migrated["worker_type"] == worker_type + assert migrated["fpm_profile"]["deployments"][0]["worker_type"] == worker_type + different_role = "decode" if worker_type == "prefill" else "prefill" + with pytest.raises(ValueError, match="worker_type conflicts"): + ForwardPassPerfModelConfig.from_legacy_engine_config(spec["engine"], different_role) + + +def test_same_topology_role_profiles_preserve_distinct_precision_in_native_bridge(profile_dict, direct_compile): + prefill = profile_dict["deployments"][0] + prefill["worker_type"] = "prefill" + decode = copy.deepcopy(prefill) + decode.update( + worker_type="decode", gemm_quant_mode="fp8_block", fmha_quant_mode="bfloat16", kv_cache_dtype="bfloat16" + ) + profile_dict["deployments"].append(decode) + for role, expected in (("prefill", prefill), ("decode", decode)): + spec = direct_compile(profile_dict, worker_type=role) + restored = ForwardPassPerfModelConfig.from_legacy_engine_config(spec["engine"], role).to_dict() + assert restored["gemm_quant_mode"] == expected["gemm_quant_mode"] + assert restored["fmha_quant_mode"] == expected["fmha_quant_mode"] + assert restored["kvcache_quant_mode"] == expected["kv_cache_dtype"] + + @pytest.mark.parametrize("kwargs", [{"moe_tp_size": 0}, {"moe_ep_size": False}, {"cp_size": True}]) def test_select_rejects_zero_or_boolean_dimensions(profile_dict, kwargs): profile = load_fpm_profile(profile_dict) @@ -625,12 +689,325 @@ def test_profile_memory_uses_declared_bounds_without_graph_or_timing_data(profil json.dumps(result) # The result is replayable; it contains no callables. +def test_legacy_profile_remains_a_complete_declaration(profile_dict): + resources = load_fpm_profile(profile_dict).deployments[0].resources + assert resources.memory_source == "declared" + assert resources.memory_ready + assert resources.non_kv_bytes == 200 + resources.require_memory() + + def test_profile_memory_capacity_override_does_not_read_hardware(profile_memory, monkeypatch): monkeypatch.setattr(memory.perf_database, "load_system_spec", _fail_graph) result = profile_memory(gpu_memory_capacity_bytes_override=2000) assert result["total_kv_size_bytes"] == 1400 +@pytest.fixture +def grouped_profile(profile_dict): + resources = profile_dict["deployments"][0]["resources"] + resources.pop("kv_bytes_per_token") + resources["cache_layout"] = "grouped" + resources["cache_groups"] = [ + { + "name": "global", + "kind": "attention", + "num_layers": 2, + "block_size_tokens": 16, + "page_size_bytes": 32, + "sliding_window": None, + }, + { + "name": "local", + "kind": "attention", + "num_layers": 4, + "block_size_tokens": 16, + "page_size_bytes": 16, + "sliding_window": 32, + }, + { + "name": "conv", + "kind": "convolution", + "num_layers": 6, + "block_size_tokens": 4, + "page_size_bytes": 8, + "sliding_window": 4, + }, + ] + return profile_dict + + +def test_grouped_profile_memory_keeps_byte_budget_and_context_bound(grouped_profile, profile_memory): + result = profile_memory(max_num_tokens=16, context_length=64, cuda_graph_reserved_bytes=20, tolerance_fraction=0.1) + assert result["total_kv_size_bytes"] == 580 + assert result["tolerance_adjusted"]["total_kv_size_bytes"] == 522 + assert result["total_kv_size_tokens"] is None + assert result["kv_size_per_token_bytes"] is None + assert result["tolerance_adjusted"]["total_kv_size_tokens"] is None + # Full: four pages * 32B; local: <=four aligned pages *16B; + # conv: <=six pages spanning retained history plus the 16-token chunk *8B. + assert result["request_peak_cache_bytes"] == 240 + assert result["cache_groups"] == grouped_profile["deployments"][0]["resources"]["cache_groups"] + assert result["memory_breakdown"]["cuda_graph_reserved_bytes"] == 20 + with pytest.raises(ValueError, match="context_length"): + profile_memory(context_length=4097) + json.dumps(result) + + +def test_grouped_profile_cannot_be_reduced_to_scalar_blocks(grouped_profile, profile_memory): + with pytest.raises(ValueError, match="scalar num_gpu_blocks"): + memory.estimate_num_gpu_blocks( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + scheduler_block_size=16, + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + tp_size=2, + moe_tp_size=2, + fpm_profile=grouped_profile, + ) + + +@pytest.mark.parametrize("capacity", [True, 1000.0, 0, 2**53 + 1]) +def test_grouped_profile_requires_exact_positive_capacity(grouped_profile, profile_memory, capacity): + with pytest.raises(ValueError, match="(cache budget|GPU capacity)"): + profile_memory(gpu_memory_capacity_bytes_override=capacity) + + +def test_grouped_profile_roundtrip_compiles_without_graph(grouped_profile, direct_compile): + spec = direct_compile(grouped_profile) + saved = json.loads(spec["engine"]["extra"]["fpm_profile"]) + assert ( + saved["deployments"][0]["resources"]["cache_groups"] + == grouped_profile["deployments"][0]["resources"]["cache_groups"] + ) + assert "kv_bytes_per_token" not in saved["deployments"][0]["resources"] + assert _normalize(_request(saved))["fpm_profile"] == saved + + +def test_linear_profile_serialization_omits_new_group_fields(profile_dict): + resource = load_fpm_profile(profile_dict).model_dump(mode="json")["deployments"][0]["resources"] + assert resource == profile_dict["deployments"][0]["resources"] + + +@pytest.fixture +def runtime_profile(profile_dict): + resources = profile_dict["deployments"][0]["resources"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(name, None) + resources["runtime_memory"] = { + "kv_cache_bytes": 600, + "gpu_memory_utilization": 0.8, + "max_model_len": 2048, + "provenance": "Synthetic worker initialization observation; not a silicon qualification.", + } + return profile_dict + + +def test_runtime_memory_selects_exact_role_capacity_and_scheduler(runtime_profile, profile_memory): + prefill = runtime_profile["deployments"][0] + prefill["worker_type"] = "prefill" + decode = copy.deepcopy(prefill) + decode["worker_type"] = "decode" + decode["resources"]["max_batch_size"] = 128 + decode["resources"]["runtime_memory"]["kv_cache_bytes"] = 720 + runtime_profile["deployments"].append(decode) + assert profile_memory(worker_type="prefill", context_length=1024)["total_kv_size_bytes"] == 600 + assert profile_memory(worker_type="decode", max_batch_size=128, context_length=1024)["total_kv_size_bytes"] == 720 + with pytest.raises(ValueError, match="exact recorded"): + profile_memory(worker_type="decode", context_length=1024) + with pytest.raises(ValueError, match="worker_type=aggregated"): + profile_memory(context_length=1024) + + +@pytest.mark.parametrize("keep_weights", [False, True]) +def test_pending_memory_is_saved_without_fabricated_bytes(profile_dict, profile_memory, direct_compile, keep_weights): + resources = profile_dict["deployments"][0]["resources"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + if name != "weights_bytes" or not keep_weights: + resources.pop(name) + profile = load_fpm_profile(profile_dict) + parsed = profile.deployments[0].resources + assert parsed.memory_source == "pending" + assert not parsed.memory_ready + assert profile.model_dump(mode="json")["deployments"][0]["resources"] == resources + assert _normalize(_request(profile_dict))["fpm_profile"] == profile.model_dump(mode="json") + assert direct_compile(profile_dict)["context_ops"][0]["FpmForward"]["weight_bytes"] == (100 if keep_weights else 0) + with pytest.raises(ValueError, match="pending runtime profiling"): + profile_memory(allow_naive_fallback=True) + with pytest.raises(ValueError, match="pending runtime profiling"): + _ = parsed.non_kv_bytes + with pytest.raises(ValueError, match="pending runtime profiling"): + RustForwardPassPerfModel.estimate_cache_budget( + _request(profile_dict), + { + "total_gpu_capacity_bytes": 1000, + "memory_fraction_kind": "of_total", + "memory_fraction_value": 0.8, + "max_num_tokens": 8192, + "max_batch_size": 256, + }, + ) + + +def test_runtime_linear_memory_uses_observation_and_canonical_budget(runtime_profile, profile_memory, monkeypatch): + parsed = load_fpm_profile(runtime_profile).deployments[0].resources + assert parsed.memory_source == "runtime" + assert parsed.memory_ready + parsed.require_memory() + assert parsed.model_dump(mode="json") == runtime_profile["deployments"][0]["resources"] + with pytest.raises(ValueError, match="not a non-KV byte breakdown"): + _ = parsed.non_kv_bytes + monkeypatch.setattr(memory.KVCacheEstimator, "_estimate_from_breakdown", _fail_graph) + result = profile_memory(context_length=1024, gpu_memory_capacity_bytes_override=2000, tolerance_fraction=0.1) + # The observed 600-byte pool remains fixed on a larger modeled device. + # Ten bytes/token gives 60 raw tokens and 54 after a 10% safety margin. + assert result["total_kv_size_bytes"] == 600 + assert result["total_kv_size_tokens"] == 60 + assert result["tolerance_adjusted"]["total_kv_size_bytes"] == 540 + assert result["tolerance_adjusted"]["total_kv_size_tokens"] == 54 + assert result["memory_breakdown"] is None + assert result["source"] == "profile" + + +def test_runtime_grouped_memory_keeps_layout_and_transient_pages(grouped_profile, runtime_profile, profile_memory): + # Reuse the hand-derived 240-byte per-request grouped footprint above. + runtime_profile["deployments"][0]["resources"]["max_num_tokens"] = 16 + result = profile_memory(max_num_tokens=16, context_length=64) + assert result["total_kv_size_bytes"] == 600 + assert result["request_peak_cache_bytes"] == 240 + assert result["memory_breakdown"] is None + assert result["total_kv_size_tokens"] is None + assert result["cache_groups"] == grouped_profile["deployments"][0]["resources"]["cache_groups"] + + +@pytest.mark.parametrize( + "kwargs,reason", + [ + ({"max_num_tokens": 8191}, "exact recorded"), + ({"max_num_tokens": 8193}, "exact recorded"), + ({"max_batch_size": 255}, "exact recorded"), + ({"max_batch_size": 257}, "exact recorded"), + ({"memory_fraction_value": 0.9}, "gpu_memory_utilization"), + ({"context_length": 2049}, "max_model_len"), + ({"cuda_graph_reserved_bytes": 1}, "graph reservations"), + ({"gpu_memory_capacity_bytes_override": 599}, "smaller than"), + ({"gpu_memory_capacity_bytes_override": 749}, "smaller than"), + ], +) +def test_runtime_memory_rejects_changed_settings(runtime_profile, profile_memory, kwargs, reason): + with pytest.raises(ValueError, match=reason): + profile_memory(**{"context_length": 2048, **kwargs}) + + +@pytest.mark.parametrize( + "change", + [ + {"kv_cache_bytes": 0}, + {"kv_cache_bytes": True}, + {"kv_cache_bytes": 2**53 + 1}, + {"gpu_memory_utilization": 0}, + {"gpu_memory_utilization": True}, + {"gpu_memory_utilization": float("nan")}, + {"gpu_memory_utilization": float("inf")}, + {"gpu_memory_utilization": 1.1}, + {"max_model_len": 0}, + {"max_model_len": True}, + {"provenance": " "}, + {"unknown": 1}, + ], +) +def test_runtime_memory_schema_rejects_invalid_evidence(runtime_profile, change): + runtime_profile["deployments"][0]["resources"]["runtime_memory"].update(change) + with pytest.raises(ValidationError): + load_fpm_profile(runtime_profile) + with pytest.raises(ValueError): + _normalize(_request(runtime_profile)) + + +def test_runtime_memory_byte_boundary_is_exact(runtime_profile, profile_memory): + runtime_profile["deployments"][0]["resources"]["runtime_memory"].update( + kv_cache_bytes=2**53, gpu_memory_utilization=1.0 + ) + result = profile_memory(context_length=2048, gpu_memory_capacity_bytes_override=2**53, memory_fraction_value=1.0) + assert result["total_kv_size_bytes"] == 2**53 + assert result["memory_breakdown"] is None + + +@pytest.mark.parametrize( + "field", ["weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"] +) +def test_runtime_memory_cannot_mix_in_declared_bytes(runtime_profile, field): + runtime_profile["deployments"][0]["resources"][field] = 0 + with pytest.raises(ValidationError, match="cannot be combined"): + load_fpm_profile(runtime_profile) + + +def test_runtime_scalar_blocks_propagate_context_bound(runtime_profile): + kwargs = { + "scheduler_block_size": 4, + "max_num_tokens": 8192, + "max_batch_size": 256, + "memory_fraction_kind": "of_total", + "memory_fraction_value": 0.8, + "tp_size": 2, + "moe_tp_size": 2, + "gpu_memory_capacity_bytes_override": 1000, + "fpm_profile": runtime_profile, + } + assert ( + memory.estimate_num_gpu_blocks( + "test/unknown-decoder", "test_gpu", "vllm", "0.25.1", context_length=1024, **kwargs + ) + == 15 + ) + with pytest.raises(ValueError, match="max_model_len"): + memory.estimate_num_gpu_blocks( + "test/unknown-decoder", "test_gpu", "vllm", "0.25.1", context_length=2049, **kwargs + ) + + +@pytest.mark.parametrize( + "change", + [ + lambda r: r.update(kv_bytes_per_token=1), + lambda r: r.update(cache_groups=[]), + lambda r: r["cache_groups"][1].update(name="global"), + lambda r: r["cache_groups"][0].update(block_size_tokens=True), + lambda r: r["cache_groups"][0].update(block_size_tokens=0), + lambda r: r["cache_groups"][0].update(page_size_bytes=2**53 + 1), + lambda r: r["cache_groups"][0].update(per_layer_page_size=1), + lambda r: r["cache_groups"][2].update(sliding_window=None), + ], +) +def test_grouped_profile_rejects_ambiguous_or_invalid_page_contract(grouped_profile, change): + change(grouped_profile["deployments"][0]["resources"]) + with pytest.raises(ValidationError): + load_fpm_profile(grouped_profile) + + +def test_grouped_budget_canonical_identity_and_bounds(grouped_profile): + config = _request(grouped_profile) + budget = { + "total_gpu_capacity_bytes": 1000, + "memory_fraction_kind": "of_total", + "memory_fraction_value": 0.8, + "max_num_tokens": 1, + "max_batch_size": 1, + "context_length": 64, + } + estimate = RustForwardPassPerfModel.estimate_cache_budget(config, budget) + assert estimate["total_kv_size_bytes"] == 600 + assert estimate["request_peak_cache_bytes"] == 192 # 4*32 + 3*16 + 2*8. + config["system"] = "another_gpu" + with pytest.raises(ValueError, match="matching FPM deployment"): + RustForwardPassPerfModel.estimate_cache_budget(config, budget) + + def test_profile_block_budget_preserves_resource_provenance(profile_dict, profile_memory): diagnostics = {} blocks = memory.estimate_num_gpu_blocks( @@ -1080,12 +1457,15 @@ def measured_profile_roots(profile_dict, tmp_path, monkeypatch): def make_root(name, kind): root = tmp_path / name root.mkdir() - packaged = Path(aisimulate_core.__file__).parent / "systems/h200_sxm.yaml" - (root / packaged.name).write_bytes(packaged.read_bytes()) identity = load_fpm_profile(profile_dict).deployments[0].model_dump(mode="json", exclude={"resources"}) + system = identity["system"] + packaged = Path(aisimulate_core.__file__).parent / "systems" / f"{system}.yaml" + (root / packaged.name).write_bytes(packaged.read_bytes()) coordinates = [] - if kind in {"genuine", "mixed"}: + if kind in {"genuine", "mixed", "coverage"}: coordinates.extend([("prefill", 0, "real_kv", 2.0), ("decode", 1, "real_kv", 3.0)]) + if kind == "coverage": + coordinates.append(("prefill", 16, "real_kv", 4.0)) if kind in {"fake", "mixed"}: coordinates.append(("decode", 64, "fake_fallback", 99.0)) rows = [ @@ -1104,7 +1484,7 @@ def make_root(name, kind): } for phase, kv, regime, latency in coordinates ] - path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path = root / "data" / system / "vllm/0.25.1/fpm_forward_perf.parquet" path.parent.mkdir(parents=True) pq.write_table(pa.Table.from_pylist(rows), path) path.with_suffix(".metadata.json").write_text( @@ -1114,7 +1494,7 @@ def make_root(name, kind): "schema_version": 6, "coordinate_system": "iteration_totals_balanced_v1", "measurement_policy": "dynamo_native_single_sample_v1", - "system": "h200_sxm", + "system": system, "backend": "vllm", "backend_version": "0.25.1", "row_count": len(rows), @@ -1127,6 +1507,169 @@ def make_root(name, kind): return make_root +@pytest.mark.parametrize("memory_state", ["pending", "runtime"]) +def test_direct_timings_work_without_declared_memory(profile_dict, measured_profile_roots, memory_state): + resources = profile_dict["deployments"][0]["resources"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(name) + if memory_state == "runtime": + resources["runtime_memory"] = { + "kv_cache_bytes": 600, + "gpu_memory_utilization": 0.8, + "max_model_len": 4096, + "provenance": "Synthetic initialization observation, independent of timing fixture.", + } + config = _request(profile_dict, "direct", systems_paths=[measured_profile_roots(memory_state, "genuine")]) + model = RustForwardPassPerfModel.best_available(config) + try: + # The synthetic prefill cell has exactly 2ms at this coordinate. + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}} + ) + == 2.0 + ) + saved = model.diagnostics()["provenance"]["config"] + assert saved["fpm_profile"]["deployments"][0]["resources"] == resources + assert saved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + finally: + model.close() + with pytest.raises(ValueError, match="registered"): + _normalize(_request(profile_dict, "sol")) + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) +def test_native_profile_role_survives_compilation_and_saved_canonical_config( + profile_dict, measured_profile_roots, worker_type +): + root = measured_profile_roots(worker_type, "genuine") + profile_dict["deployments"][0]["worker_type"] = worker_type + request = _request(profile_dict, "direct", worker_type=worker_type, systems_paths=[root]) + for config in (request, json.loads(json.dumps(_normalize(request)))): + model = RustForwardPassPerfModel.best_available(config) + try: + workload = ( + {"num_decode_requests": 1, "sum_decode_kv_tokens": 1} + if worker_type == "decode" + else {"num_prefill_requests": 1, "sum_prefill_tokens": 1} + ) + assert model.estimate_forward_pass_time_ms({"scheduled_requests": workload}) == ( + 3.0 if worker_type == "decode" else 2.0 + ) + saved = model.diagnostics()["provenance"]["config"] + assert saved["worker_type"] == worker_type + assert saved["fpm_profile"]["deployments"][0]["worker_type"] == worker_type + finally: + model.close() + + +def test_direct_coverage_is_native_persists_errors_and_round_trips(profile_dict, measured_profile_roots): + root = measured_profile_roots("coverage", "mixed") + config = _request(profile_dict, "direct", systems_paths=[root]) + plain = RustForwardPassPerfModel.best_available(config) + assert plain.fpm_query_coverage() is None + assert ( + "collect_coverage" not in plain.diagnostics()["provenance"]["config"]["estimator_config"]["fpm_interpolation"] + ) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + model = RustForwardPassPerfModel.best_available(config) + assert model.predict_prefill_latency(bs=1, isl=1, prefix=0) == plain.predict_prefill_latency(1, 1, 0) == 2.0 + assert model.predict_decode_latency_total(1, 1) == plain.predict_decode_latency_total(1, 1) == 3.0 + assert model.fpm_decode_kv_ceiling() == 1 # The fake KV=64 row is not coverage. + for _ in range(2): + with pytest.raises(PerfDataNotAvailableError, match="genuine FPM points"): + model.predict_decode_latency_total(1, 64) + report = model.fpm_query_coverage() + assert report["counting_unit"] == "native_lookup_resolutions" + assert report["queries"] == {"measured": 2, "interpolated": 0, "unsupported": 2} + assert len(report["gaps"]) == 1 + gap = report["gaps"][0] + assert gap["occurrences"] == 2 + assert gap["coordinates"] == {"batch_size": 1.0, "total_kv_read_tokens": 64.0} + assert gap["cell_identity"]["attention_backend"] == "future_attention" + assert gap["model_path"] == profile_dict["model"] + saved = json.loads(json.dumps(model.diagnostics()["provenance"]["config"])) + assert saved["estimator_config"]["fpm_interpolation"]["collect_coverage"] is True + reloaded = RustForwardPassPerfModel.best_available(saved) + assert reloaded.fpm_query_coverage()["queries"] == {"measured": 0, "interpolated": 0, "unsupported": 0} + assert reloaded.predict_decode_latency_total(1, 1) == 3.0 + assert model.fpm_query_coverage() == report + + +@pytest.mark.parametrize( + ("field", "value"), + [("estimation_mode", "auto"), ("estimation_mode", "op_level"), ("fallback_policy", "allow")], +) +def test_direct_coverage_rejects_estimator_fallback(profile_dict, field, value): + config = _request(profile_dict, "direct", **{field: value}) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + with pytest.raises(ValueError, match="collect_coverage"): + _normalize(config) + + +@pytest.mark.parametrize( + ("model_id", "system", "tp", "dp", "moe_tp", "moe_ep", "gemm", "fmha"), + [ + ("MiniMaxAI/MiniMax-M2.7", "h200_sxm", 4, 1, 4, 1, "fp8_block", "bfloat16"), + ("nvidia/GLM-5.2-NVFP4", "b200_sxm", 1, 8, 1, 8, "nvfp4", "fp8"), + ("nvidia/GLM-5.2-NVFP4", "b200_sxm", 8, 1, 1, 8, "nvfp4", "fp8"), + ], + ids=["tp4", "dep8", "tep8"], +) +def test_coverage_wrappers_preserve_existing_engine_queries_on_external_cells( + profile_dict, measured_profile_roots, model_id, system, tp, dp, moe_tp, moe_ep, gemm, fmha +): + import aisimulate_core + + # Synthetic external timings exercise TP/DEP/TEP routing parity, not model + # accuracy or GPU memory qualification. No bundled FPM data is required. + profile_dict.update( + model=model_id, + architecture="MiniMaxM2ForCausalLM" if model_id == "MiniMaxAI/MiniMax-M2.7" else "GlmMoeDsaForCausalLM", + ) + profile_dict["deployments"][0].update( + system=system, + tp=tp, + dp=dp, + moe_tp=moe_tp, + moe_ep=moe_ep, + gemm_quant_mode=gemm, + moe_quant_mode=gemm, + fmha_quant_mode=fmha, + attention_backend="auto", + ) + root = measured_profile_roots("external", "coverage") + config = _request(profile_dict, "direct", systems_paths=[root]) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + covered = RustForwardPassPerfModel.best_available(config) + original = aisimulate_core.AicEngine.from_spec( + engine.compile_engine( + model_id, + system, + "vllm", + "0.25.1", + tp_size=tp, + attention_dp_size=dp, + moe_tp_size=moe_tp, + moe_ep_size=moe_ep, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + systems_path=root, + ) + ) + for prefix, latency in ((0, 2.0), (16, 4.0)): + assert covered.predict_prefill_latency(1, 1 + prefix, prefix) == original.predict_prefill_latency( + 1, 1 + prefix, prefix + ) + assert covered.predict_prefill_latency(1, 1 + prefix, prefix) == latency + assert covered.predict_decode_latency_total(1, 1) == original.predict_decode_latency_total(1, 1) == 3.0 + assert covered.predict_decode_latency_total(0, 0) == original.predict_decode_latency_total(0, 0) == 0.0 + assert covered.fpm_decode_kv_ceiling() == original.fpm_decode_kv_ceiling() + report = covered.fpm_query_coverage() + assert report["queries"] == {"measured": 5, "interpolated": 0, "unsupported": 0} + + @pytest.mark.parametrize("fallback", ["deny", "allow"]) def test_direct_fake_only_cell_is_unavailable(profile_dict, measured_profile_roots, fallback): root = measured_profile_roots("fake", "fake") diff --git a/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py b/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py index da626b0f5..8edded8a2 100644 --- a/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py +++ b/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py @@ -74,7 +74,7 @@ def graph_model(): @pytest.mark.parametrize("call,expected,native", CASES) def test_canonical_direct_prefill_matches_frozen_operator_and_native_evidence(graph_model, call, expected, native): - result = graph_model.predict_prefill_latency(*call) + result = graph_model.predict_prefill_latency(bs=call[0], isl=call[1], prefix=call[2]) assert result == pytest.approx(expected, rel=1e-12, abs=1e-10) assert abs(result / native - 1.0) <= 0.15 assert graph_model.predict_prefill_latency(*call) == result diff --git a/python/aisimulate/tests/unit/sdk/test_utils.py b/python/aisimulate/tests/unit/sdk/test_utils.py index cf3a757ce..0b6f76c8b 100644 --- a/python/aisimulate/tests/unit/sdk/test_utils.py +++ b/python/aisimulate/tests/unit/sdk/test_utils.py @@ -1352,6 +1352,88 @@ def test_constraints_default_to_95_percent_ttft_when_not_provided(self): assert derived_pair[1] == pytest.approx((1000 - 950) / (50 - 1)) +class TestKVCacheQuantInference: + @pytest.mark.parametrize("source", ["quantization_config", "hf_quant_config"]) + @pytest.mark.parametrize("kv_algo", ["none", " NONE "]) + @pytest.mark.parametrize( + "weight_algo,gemm,moe", + [("NVFP4", "nvfp4", "nvfp4"), ("FP8", "fp8_static", "fp8"), (None, None, None)], + ) + def test_explicit_none_is_distinct_from_missing_metadata(self, source, kv_algo, weight_algo, gemm, moe): + from aisimulate.sdk.models import _infer_quant_modes_from_raw_config + from aisimulate.sdk.utils import _attach_inferred_quant_fields + + quant = {"quant_algo": weight_algo, "kv_cache_quant_algo": kv_algo} + raw = {source: {"quantization": quant} if source == "hf_quant_config" else quant} + _attach_inferred_quant_fields(raw) + modes = _infer_quant_modes_from_raw_config(raw) + + assert raw["kv_cache_quant_algo"] == "none" + assert quant["kv_cache_quant_algo"] == kv_algo + assert modes["kvcache_quant_mode"] == common.KVCacheQuantMode.bfloat16 + assert modes["fmha_quant_mode"] == common.FMHAQuantMode.bfloat16 + assert modes.get("gemm_quant_mode") == (common.GEMMQuantMode[gemm] if gemm else None) + assert modes.get("moe_quant_mode") == (common.MoEQuantMode[moe] if moe else None) + + @pytest.mark.parametrize("weight_algo", [None, "NVFP4"]) + @pytest.mark.parametrize( + "kv_fields,normalized_kv", + [ + ({}, None), + ({"kv_cache_quant_algo": None}, None), + ({"kv_cache_quant_algo": "FP8"}, "fp8"), + ({"kv_cache_quant_algo": "E4M3"}, "fp8"), + ({"kv_cache_quant_algo": "bfloat16"}, "bfloat16"), + ({"kv_cache_quant_algo": "FP16"}, "bfloat16"), + ], + ) + def test_legacy_defaults_and_supported_algorithms_are_unchanged(self, weight_algo, kv_fields, normalized_kv): + from aisimulate.sdk.models import _infer_quant_modes_from_raw_config + from aisimulate.sdk.utils import _attach_inferred_quant_fields + + raw = _attach_inferred_quant_fields({"quantization_config": {"quant_algo": weight_algo, **kv_fields}}) + modes = _infer_quant_modes_from_raw_config(raw) + + assert raw.get("kv_cache_quant_algo") == normalized_kv + if weight_algo: + assert modes["kvcache_quant_mode"] == common.KVCacheQuantMode.fp8 + assert modes["fmha_quant_mode"] == common.FMHAQuantMode.fp8 + else: + assert modes.get("kvcache_quant_mode") == ( + common.KVCacheQuantMode[normalized_kv] if normalized_kv else None + ) + assert "fmha_quant_mode" not in modes + + @pytest.mark.parametrize("kv_algo", ["not-an-algorithm", "null", ["none"], {"algorithm": "none"}, True]) + @pytest.mark.parametrize("source", ["quantization_config", "hf_quant_config"]) + def test_unknown_or_malformed_kv_algorithms_still_fail(self, kv_algo, source): + from aisimulate.sdk.models import _infer_quant_modes_from_raw_config + from aisimulate.sdk.utils import _attach_inferred_quant_fields + + quant = {"quant_algo": "NVFP4", "kv_cache_quant_algo": kv_algo} + raw = _attach_inferred_quant_fields({source: {"quantization": quant} if source == "hf_quant_config" else quant}) + with pytest.raises(ValueError, match="Unsupported kv cache algorithm"): + _infer_quant_modes_from_raw_config(raw) + + @pytest.mark.parametrize("architecture", ["DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"]) + @pytest.mark.parametrize("weight_algo", [None, "NVFP4"]) + @pytest.mark.parametrize("kv_algo", [None, "none", "FP8", "bfloat16"]) + def test_deepseek_native_cache_requirement_is_preserved(self, architecture, weight_algo, kv_algo): + from aisimulate.sdk.models import _infer_quant_modes_from_raw_config + from aisimulate.sdk.utils import _attach_inferred_quant_fields + + raw = _attach_inferred_quant_fields( + {"quantization_config": {"quant_algo": weight_algo, "kv_cache_quant_algo": kv_algo}} + ) + modes = _infer_quant_modes_from_raw_config(raw, architecture) + + assert modes["kvcache_quant_mode"] == common.KVCacheQuantMode.fp8 + if weight_algo: + assert modes["fmha_quant_mode"] == common.FMHAQuantMode.fp8 + else: + assert "fmha_quant_mode" not in modes + + class TestParseCompressedTensorsQuant: """Tests for parse_compressed_tensors_quant and its _categorize_ignore_pattern helper.""" diff --git a/python/aisimulate/tests/unit/test_fpm_config_imports.py b/python/aisimulate/tests/unit/test_fpm_config_imports.py index 04a59b74f..f9b8a7229 100644 --- a/python/aisimulate/tests/unit/test_fpm_config_imports.py +++ b/python/aisimulate/tests/unit/test_fpm_config_imports.py @@ -8,6 +8,7 @@ import json import subprocess import sys +from copy import deepcopy from pathlib import Path import pytest @@ -122,7 +123,9 @@ def supervised_child(*args, **kwargs): @pytest.mark.parametrize("command", ["predict", "recommend"]) -@pytest.mark.parametrize("case", ["no-profile", "profile", "unknown-quant", "boolean-topology", "wrong-model"]) +@pytest.mark.parametrize( + "case", ["no-profile", "profile", "mixed-profile", "unknown-quant", "boolean-topology", "wrong-model"] +) def test_supervised_profile_validation_is_runtime_free(command: str, case: str, tmp_path: Path) -> None: """Exercise the real parser and envelope validation before the child starts.""" raw = _config() @@ -131,6 +134,15 @@ def test_supervised_profile_validation_is_runtime_free(command: str, case: str, error = "" if case == "no-profile": del raw["engine"]["fpm_profile"] + elif case == "mixed-profile": + grouped = deepcopy(raw["engine"]["fpm_profile"]["deployments"][0]) + grouped["tp"] = 2 + grouped["resources"]["cache_layout"] = "grouped" + grouped["resources"].pop("kv_bytes_per_token") + grouped["resources"]["cache_groups"] = [ + {"name": "full", "kind": "attention", "num_layers": 1, "block_size_tokens": 64, "page_size_bytes": 64} + ] + raw["engine"]["fpm_profile"]["deployments"].append(grouped) elif case == "unknown-quant": raw["engine"]["fpm_profile"]["deployments"][0]["fmha_quant_mode"] = "unsupported" error = "unknown FPM fmha_quant_mode" diff --git a/python/aisimulate/tests/unit/test_onboard_checkpoint.py b/python/aisimulate/tests/unit/test_onboard_checkpoint.py new file mode 100644 index 000000000..33b2a73ed --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_checkpoint.py @@ -0,0 +1,862 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Session pause/resume through the public CLI without GPU execution.""" + +from __future__ import annotations + +import io +import json +import os +import shlex +import subprocess +import sys +from copy import deepcopy +from pathlib import Path + +import pytest +import yaml + +import aisimulate.main as cli +from aisimulate.support import checkpoint +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def request_payload(): + """Synthetic resources are declared test bounds, not measured silicon.""" + profile = { + "schema_version": 1, + "model": "example/checkpoint-test", + "model_revision": "revision-123", + "architecture": "SyntheticDecoder", + "context_length": 16384, + "num_experts": 0, + "provenance": "Synthetic checkpoint test", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.27.0", + "tp": 1, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "bfloat16", + "resources": { + "weights_bytes": 1000, + "activations_bytes": 100, + "runtime_overhead_bytes": 100, + "comm_overhead_bytes": 0, + "kv_bytes_per_token": 128, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Synthetic declared bounds", + }, + } + ], + } + return SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "dense", + "framework_version": "0.27.0", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": {"context_length": 16384}, + "collection": {"prefill_cudagraph_policy": "runtime"}, + "fpm_profile": profile, + } + ).model_dump(mode="json", exclude_none=True) + + +def _save(path, capsys, monkeypatch, patch=None, revision=None, accept=(), expected=0): + args = ["onboard", "checkpoint", "--file", str(path)] + if patch is not None: + monkeypatch.setattr(sys, "stdin", io.StringIO(json.dumps(patch))) + args += ["--update", "-"] + if revision is not None: + args += ["--expect-revision", str(revision)] + for name in accept: + args += ["--accept-profile", name] + assert cli.main(args) == expected + return json.loads(capsys.readouterr().out) + + +def _resume(path, capsys, expected=0): + assert cli.main(["onboard", "resume", "--checkpoint", str(path)]) == expected + return json.loads(capsys.readouterr().out) + + +def test_role_graph_edit_invalidates_only_its_accepted_configuration(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + configurations = {} + for role, mode in (("prefill", "PIECEWISE"), ("decode", "FULL_DECODE_ONLY")): + payload = deepcopy(request_payload) + payload["worker_type"] = role + payload["collection"]["cudagraph_mode"] = mode + payload["fpm_profile"]["deployments"][0]["worker_type"] = role + configurations[role] = {"draft_request": payload} + first = _save(path, capsys, monkeypatch, {"configurations": configurations}, accept=("prefill", "decode")) + assert all(value["profile_accepted"] for value in first["configurations"].values()) + original_decode = first["state"]["configurations"]["decode"] + changed = _save( + path, + capsys, + monkeypatch, + {"configurations": {"prefill": {"draft_request": {"collection": {"max_cudagraph_capture_size": 512}}}}}, + revision=first["state"]["revision"], + ) + assert not changed["configurations"]["prefill"]["profile_accepted"] + assert changed["configurations"]["decode"]["profile_accepted"] + assert changed["state"]["configurations"]["decode"] == original_decode + resumed = _resume(path, capsys) + assert resumed["state"]["configurations"]["prefill"]["draft_request"]["worker_type"] == "prefill" + assert ( + resumed["state"]["configurations"]["prefill"]["draft_request"]["collection"]["max_cudagraph_capture_size"] + == 512 + ) + + +def _subprocess(*args, cwd, stdin=None): + # The installed public console entry point bypasses supervision for onboarding. + return subprocess.run( + [str(Path(sys.executable).parent / "aisimulate"), "onboard", *args], + cwd=cwd, + input=stdin, + capture_output=True, + text=True, + timeout=30, + env={**os.environ, "PYTHONPATH": str(Path(cli.__file__).parents[1])}, + ) + + +def test_empty_stage_one_and_research_resume_in_new_process_and_cwd(tmp_path): + path = tmp_path / "session" / "onboarding-checkpoint.json" + created = _subprocess("checkpoint", "--file", str(path), cwd=tmp_path) + assert created.returncode == 0, created.stderr + assert json.loads(created.stdout)["state"]["revision"] == 1 + marker = tmp_path / "must-not-execute" + patch = { + "research": {"runtime": {"source": "pinned/source/path", "confidence": "unknown", "alternatives": ["a", "b"]}}, + "pending_questions": ["Choose framework version"], + "progress": {"stage": 2, "next_action": f"touch {marker}"}, + } + updated = _subprocess( + "checkpoint", + "--file", + str(path), + "--update", + "-", + "--expect-revision", + "1", + cwd=tmp_path, + stdin=json.dumps(patch), + ) + assert updated.returncode == 0, updated.stderr + resumed = _subprocess("resume", "--checkpoint", str(path), cwd=tmp_path.parent) + assert resumed.returncode == 0, resumed.stderr + result = json.loads(resumed.stdout) + assert result["state"]["research"] == patch["research"] + assert result["state"]["pending_questions"] == patch["pending_questions"] + assert result["verification"]["executed"] is False + assert not marker.exists() + assert sorted(p.name for p in path.parent.iterdir()) == [ + ".onboarding-checkpoint.json.lock", + "onboarding-checkpoint.json", + ] + + +def test_partial_approval_and_affected_changes(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + first = {"inputs": {"precision": "bf16"}, "draft_request": request_payload} + second = {"inputs": {"precision": "alternative"}, "draft_request": {"identity": {"model": "still investigating"}}} + saved = _save( + path, capsys, monkeypatch, {"configurations": {"bf16-tp1": first, "other-tp4": second}}, accept=["bf16-tp1"] + ) + assert saved["configurations"]["bf16-tp1"]["profile_accepted"] + assert not saved["configurations"]["other-tp4"]["profile_accepted"] + original_acceptance = saved["state"]["configurations"]["bf16-tp1"]["acceptance"] + notes = _save(path, capsys, monkeypatch, {"research": {"note": "Paused after reviewing first profile"}}, 1) + assert notes["state"]["configurations"]["bf16-tp1"]["acceptance"] == original_acceptance + changed_second = _save( + path, capsys, monkeypatch, {"configurations": {"other-tp4": {"inputs": {"precision": "fp8"}}}}, 2 + ) + assert changed_second["configurations"]["bf16-tp1"]["profile_accepted"] + changed_first = _save( + path, + capsys, + monkeypatch, + {"configurations": {"bf16-tp1": {"draft_request": {"collection": {"gpu_memory_utilization": 0.8}}}}}, + 3, + ) + assert not changed_first["configurations"]["bf16-tp1"]["profile_accepted"] + assert changed_first["state"]["configurations"]["bf16-tp1"]["acceptance"] is None + _save(path, capsys, monkeypatch, revision=4, accept=["bf16-tp1"]) + shared = _save(path, capsys, monkeypatch, {"inputs": {"runtime_build": "different"}}, 5) + assert not shared["configurations"]["bf16-tp1"]["profile_accepted"] + assert shared["state"]["configurations"]["bf16-tp1"]["progress"]["stage"] == 3 + assert not (tmp_path / "request.yaml").exists() + + +@pytest.mark.parametrize( + "edit", + [ + {"image": "registry.example/fpm:replacement"}, + {"container_mount": ["/shared/other checkpoint:/models:ro"]}, + ], +) +def test_slurm_options_survive_pause_and_require_fresh_collection_after_edit( + tmp_path, capsys, monkeypatch, request_payload, edit +): + from collector.fpm_forward import cli as collector_cli + from collector.fpm_forward.config import FPMCollectionOptions + from collector.fpm_forward.planner import build_collection_plan + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + from aisimulate.support.fpm import fpm_cli_args + from aisimulate.support.schema import FPMDeployment + + path = tmp_path / "onboarding-checkpoint.json" + request = tmp_path / "request.yaml" + request.write_text(yaml.safe_dump(request_payload)) + original = tmp_path / "collection" + assert cli.main(["onboard", "plan", "-c", str(request), "--output-dir", str(original)]) == 0 + capsys.readouterr() + collector = original / "fpm-checkpoint/fpm_forward.json" + collector.parent.mkdir() + deployment = { + "executor": "slurm", + "image": "registry.example/fpm@sha256:" + "a" * 64, + "container_mount": ["/shared/model cache:/models:ro", "/shared/cache:/root/.cache/huggingface"], + "transport": "ib", + } + saved_request = SupportRequest.model_validate(request_payload) + args = collector_cli._parser().parse_args( + fpm_cli_args(saved_request, output_dir=original, plan_only=True, deployment=FPMDeployment(**deployment))[3:] + ) + model_config = tmp_path / "model-config.json" + model_config.write_text( + json.dumps( + { + "architectures": [saved_request.fpm_profile.architecture], + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "vocab_size": 1024, + "max_position_embeddings": 16384, + "torch_dtype": "bfloat16", + } + ) + ) + frozen = build_collection_plan( + backend=saved_request.identity.framework, + model_path=saved_request.identity.model, + system=saved_request.identity.gpu, + selected_ops={"attention_context", "attention_generation"}, + options=FPMCollectionOptions.from_args(args), + fpm_profile=saved_request.fpm_profile, + model_config_path=str(model_config), + generator_overrides=collector_cli._load_generator_overrides(args), + ) + campaign = original / "fpm-artifacts" / frozen.sha256[:16] + campaign.mkdir(parents=True) + (campaign / "collection-plan.json").write_text(json.dumps(frozen.to_dict())) + collector.write_text(json.dumps({"schema": CHECKPOINT_SCHEMA, "plan_sha256": frozen.sha256, "cells": {}})) + original_bytes = {file: file.read_bytes() for file in original.rglob("*") if file.is_file()} + _save( + path, + capsys, + monkeypatch, + { + "configurations": { + "worker": { + "inputs": {"collection_deployment": deployment}, + "draft_request": request_payload, + "progress": {"stage": 5, "next_action": "Inspect the existing allocation, then resume collection."}, + "artifacts": { + "plan": {"path": "collection/support-plan.json", "scope": "collection"}, + "collector": { + "path": "collection/fpm-checkpoint/fpm_forward.json", + "kind": "collector_checkpoint", + "scope": "collection", + }, + }, + } + } + }, + accept=["worker"], + ) + + paused_bytes = path.read_bytes() + resumed = _resume(path, capsys) + assert path.read_bytes() == paused_bytes + assert resumed["configurations"]["worker"]["profile_accepted"] + assert resumed["verification"]["executed"] is False + saved_options = resumed["state"]["configurations"]["worker"]["inputs"]["collection_deployment"] + assert saved_options == deployment + + def preview(options, root, *, resume=False): + command = ["onboard", "collect-fpm", "-c", str(request), "--output-dir", str(root)] + for name in ("executor", "image", "transport"): + command.extend(("--" + name, options[name])) + for mount in options["container_mount"]: + command.extend(("--container-mount", mount)) + if resume: + command.append("--resume") + assert cli.main(command) == 0 + return shlex.split(capsys.readouterr().out) + + command = preview(saved_options, original, resume=True) + assert command[command.index("--fpm-executor") + 1] == "slurm" + assert command[command.index("--fpm-slurm-container-image") + 1] == deployment["image"] + assert [ + command[index + 1] for index, value in enumerate(command) if value == "--fpm-slurm-container-mount" + ] == deployment["container_mount"] + assert "--resume" in command + + updated = _save( + path, + capsys, + monkeypatch, + {"configurations": {"worker": {"inputs": {"collection_deployment": edit}}}}, + 1, + expected=2, + ) + assert updated["saved"] + assert not updated["configurations"]["worker"]["profile_accepted"] + assert updated["state"]["configurations"]["worker"]["progress"]["stage"] == 3 + assert {issue["artifact"] for issue in updated["integrity_issues"]} == {"plan", "collector"} + _resume(path, capsys, expected=2) + + replacement = tmp_path / "collection-v2" + assert cli.main(["onboard", "plan", "-c", str(request), "--output-dir", str(replacement)]) == 0 + capsys.readouterr() + recovered = _save( + path, + capsys, + monkeypatch, + { + "research": {"deployment_edit": "Reviewed changed deployment; preserve the earlier campaign."}, + "configurations": { + "worker": { + "artifacts": { + "plan": {"archived": True}, + "collector": {"archived": True}, + "plan_v2": {"path": "collection-v2/support-plan.json", "scope": "collection"}, + } + } + }, + }, + 2, + accept=["worker"], + ) + assert recovered["configurations"]["worker"]["profile_accepted"] + assert len(recovered["archived_artifacts"]) == 2 + final = _resume(path, capsys) + current = final["state"]["configurations"]["worker"]["inputs"]["collection_deployment"] + replacement_command = preview(current, replacement) + assert replacement_command[replacement_command.index("--fpm-slurm-container-image") + 1] == current["image"] + assert [ + replacement_command[index + 1] + for index, value in enumerate(replacement_command) + if value == "--fpm-slurm-container-mount" + ] == current["container_mount"] + assert replacement_command[replacement_command.index("--fpm-artifact-root") + 1] == str( + replacement / "fpm-artifacts" + ) + assert all(file.read_bytes() == content for file, content in original_bytes.items()) + + +def test_validation_changes_preserve_collection_identity_and_approval(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + for name in ("plan.json", "validation.json"): + (tmp_path / name).write_text("{}") + patch = { + "configurations": { + "one": { + "draft_request": request_payload, + "progress": {"stage": 6, "status": "complete"}, + "artifacts": { + "plan": {"path": "plan.json", "scope": "collection"}, + "validation": {"path": "validation.json", "scope": "validation"}, + }, + } + } + } + saved = _save(path, capsys, monkeypatch, patch, accept=["one"]) + assert saved["configurations"]["one"]["profile_accepted"] + updated = _save(path, capsys, monkeypatch, {"validation_inputs": {"trace": "new trace"}}, 1, expected=2) + assert updated["saved"] + assert updated["configurations"]["one"]["profile_accepted"] + assert [issue["artifact"] for issue in updated["integrity_issues"]] == ["validation"] + assert updated["state"]["configurations"]["one"]["progress"]["status"] == "in_progress" + assert (tmp_path / "plan.json").read_text() == "{}" + (tmp_path / "validation-v2.json").write_text('{"trace":"new trace"}') + recovered = _save( + path, + capsys, + monkeypatch, + { + "configurations": { + "one": { + "artifacts": { + "validation": {"archived": True}, + "validation_v2": {"path": "validation-v2.json", "scope": "validation"}, + } + } + } + }, + 2, + ) + original = saved["state"]["configurations"]["one"] + current = recovered["state"]["configurations"]["one"] + assert current["acceptance"] == original["acceptance"] + assert current["artifacts"]["plan"] == original["artifacts"]["plan"] + assert current["artifacts"]["validation"] == {**original["artifacts"]["validation"], "archived": True} + assert _resume(path, capsys)["configurations"]["one"]["effective_status"] == "accepted" + + +def test_input_and_immutable_drift_survive_note_save_while_collector_can_advance( + tmp_path, capsys, monkeypatch, request_payload +): + source = tmp_path / "config.json" + source.write_text('{"model":1}') + collector = tmp_path / "collector.json" + collector.write_text('{"cells":{}}') + path = tmp_path / "checkpoint.json" + saved = _save( + path, + capsys, + monkeypatch, + { + "artifacts": {"source": {"path": str(source)}}, + "configurations": { + "one": { + "draft_request": request_payload, + "artifacts": { + "collector": { + "path": "collector.json", + "kind": "collector_checkpoint", + "scope": "collection", + } + }, + } + }, + }, + accept=["one"], + ) + digest = saved["state"]["artifacts"]["source"]["sha256"] + collector.write_text('{"cells":{"one":"passed"}}') + assert _resume(path, capsys)["configurations"]["one"]["profile_accepted"] + source.write_text('{"model":2}') + saved = _save( + path, + capsys, + monkeypatch, + {"research": {"note": "do not refresh hashes"}, "artifacts": {"source": {"path": "config.json"}}}, + 1, + expected=2, + ) + assert saved["state"]["artifacts"]["source"]["sha256"] == digest + assert not saved["configurations"]["one"]["profile_accepted"] + assert saved["integrity_issues"][0]["artifact"] == "source" + assert json.loads(collector.read_text()) == {"cells": {"one": "passed"}} + collector.unlink() + resumed = _resume(path, capsys, expected=2) + assert {issue["artifact"] for issue in resumed["integrity_issues"]} == {"source", "collector"} + + +def test_exact_request_and_profile_references(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + (tmp_path / "request.yaml").write_text(yaml.safe_dump(request_payload)) + (tmp_path / "profile.json").write_text(json.dumps(request_payload["fpm_profile"])) + patch = { + "configurations": { + "one": { + "draft_request": request_payload, + "artifacts": { + "request": {"path": "request.yaml", "kind": "request", "scope": "collection"}, + "profile": {"path": "profile.json", "kind": "profile", "scope": "collection"}, + }, + } + } + } + assert _save(path, capsys, monkeypatch, patch, accept=["one"])["configurations"]["one"]["profile_accepted"] + wrong = deepcopy(request_payload) + wrong["collection"]["gpu_memory_utilization"] = 0.7 + (tmp_path / "wrong.yaml").write_text(yaml.safe_dump(wrong)) + result = _save( + path, + capsys, + monkeypatch, + {"configurations": {"one": {"artifacts": {"request": {"path": "wrong.yaml"}}}}}, + 1, + expected=2, + ) + assert "does not match" in result["integrity_issues"][0]["detail"] + assert "agent-reported" in result["verification"]["scope"] + + +def test_archive_superseded_outputs_retains_history_and_allows_current_work( + tmp_path, capsys, monkeypatch, request_payload +): + path = tmp_path / "checkpoint.json" + (tmp_path / "request-v1.yaml").write_text(yaml.safe_dump(request_payload)) + (tmp_path / "plan-v1.json").write_text('{"plan":"old"}') + (tmp_path / "collector-v1.json").write_text('{"cells":{"one":"passed"}}') + old_refs = { + "request_v1": {"path": "request-v1.yaml", "kind": "request", "scope": "collection"}, + "plan_v1": {"path": "plan-v1.json", "scope": "collection"}, + "collector_v1": {"path": "collector-v1.json", "kind": "collector_checkpoint", "scope": "collection"}, + } + initial = _save( + path, + capsys, + monkeypatch, + {"configurations": {"one": {"draft_request": request_payload, "artifacts": old_refs}}}, + accept=["one"], + ) + snapshots = initial["state"]["configurations"]["one"]["artifacts"] + changed = deepcopy(request_payload) + changed["collection"]["gpu_memory_utilization"] = 0.8 + (tmp_path / "request-v2.yaml").write_text(yaml.safe_dump(changed)) + (tmp_path / "plan-v2.json").write_text('{"plan":"current"}') + current = _save( + path, + capsys, + monkeypatch, + { + "configurations": { + "one": { + "draft_request": changed, + "artifacts": { + "request_v2": {"path": "request-v2.yaml", "kind": "request", "scope": "collection"}, + "plan_v2": {"path": "plan-v2.json", "scope": "collection"}, + }, + } + } + }, + 1, + accept=["one"], + expected=2, + ) + assert current["configurations"]["one"]["profile_accepted"] + assert current["configurations"]["one"]["effective_status"] == "needs_attention" + assert {issue["artifact"] for issue in current["integrity_issues"]} == set(old_refs) + archived = _save( + path, + capsys, + monkeypatch, + {"configurations": {"one": {"artifacts": {name: {"archived": True} for name in old_refs}}}}, + 2, + ) + assert archived["configurations"]["one"]["effective_status"] == "accepted" + assert {ref["artifact"] for ref in archived["archived_artifacts"]} == set(old_refs) + assert "Historical references only" in archived["verification"]["archived_artifacts"] + for name, original in snapshots.items(): + assert archived["state"]["configurations"]["one"]["artifacts"][name] == {**original, "archived": True} + assert (tmp_path / original["path"]).is_file() + restored = _save( + path, + capsys, + monkeypatch, + {"configurations": {"one": {"artifacts": {"plan_v1": {"archived": False}}}}}, + 3, + expected=2, + ) + assert [issue["artifact"] for issue in restored["integrity_issues"]] == ["plan_v1"] + assert "stale" in restored["integrity_issues"][0]["detail"] + _save( + path, + capsys, + monkeypatch, + {"configurations": {"one": {"artifacts": {"plan_v1": {"archived": True}}}}}, + 4, + ) + for original in snapshots.values(): + (tmp_path / original["path"]).unlink() + resumed = _resume(path, capsys) + assert resumed["configurations"]["one"]["effective_status"] == "accepted" + assert resumed["archived_artifacts"] == archived["archived_artifacts"] + assert (tmp_path / "request-v2.yaml").is_file() + assert (tmp_path / "plan-v2.json").is_file() + + +@pytest.mark.parametrize("shared", [True, False]) +def test_retiring_source_requires_new_acceptance(tmp_path, capsys, monkeypatch, request_payload, shared): + path = tmp_path / "checkpoint.json" + source = tmp_path / "config.json" + source.write_text('{"version":1}') + patch = { + "configurations": { + name: {"draft_request": request_payload, "progress": {"stage": 6, "status": "complete"}} + for name in ("one", "two") + } + } + owner = patch if shared else patch["configurations"]["one"] + owner["artifacts"] = {"source": {"path": "config.json"}} + initial = _save(path, capsys, monkeypatch, patch, accept=["one", "two"]) + original = initial["state"] if shared else initial["state"]["configurations"]["one"] + retired_ref = {"artifacts": {"source": {"archived": True}}} + update = retired_ref if shared else {"configurations": {"one": retired_ref}} + source.unlink() + retired = _save(path, capsys, monkeypatch, update, 1) + assert not retired["configurations"]["one"]["profile_accepted"] + assert retired["state"]["configurations"]["one"]["acceptance"] is None + assert retired["state"]["configurations"]["one"]["progress"]["stage"] == 3 + assert retired["configurations"]["two"]["profile_accepted"] is not shared + owner = retired["state"] if shared else retired["state"]["configurations"]["one"] + assert owner["artifacts"]["source"] == {**original["artifacts"]["source"], "archived": True} + assert retired["archived_artifacts"] == [ + {"configuration": None if shared else "one", "artifact": "source", "path": "config.json"} + ] + accepted = _save(path, capsys, monkeypatch, revision=2, accept=["one"]) + assert accepted["configurations"]["one"]["profile_accepted"] + + +@pytest.mark.parametrize("invalid", ["true", 1, [], {}]) +def test_invalid_archive_field_preserves_checkpoint(tmp_path, capsys, monkeypatch, invalid): + path = tmp_path / "checkpoint.json" + (tmp_path / "config.json").write_text("{}") + _save(path, capsys, monkeypatch, {"artifacts": {"source": {"path": "config.json"}}}) + before = path.read_bytes() + with pytest.raises(SystemExit) as error: + _save(path, capsys, monkeypatch, {"artifacts": {"source": {"archived": invalid}}}, 1) + assert error.value.code == 2 + capsys.readouterr() + assert path.read_bytes() == before + malformed = json.loads(before) + malformed["artifacts"]["source"]["archived"] = invalid + path.write_text(json.dumps(malformed)) + before = path.read_bytes() + with pytest.raises(SystemExit) as error: + _resume(path, capsys) + assert error.value.code == 2 + capsys.readouterr() + assert path.read_bytes() == before + + +def test_saved_artifact_without_archive_field_remains_current(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + source = tmp_path / "config.json" + source.write_text("{}") + _save( + path, + capsys, + monkeypatch, + { + "artifacts": {"source": {"path": "config.json"}}, + "configurations": {"one": {"draft_request": request_payload}}, + }, + accept=["one"], + ) + legacy = json.loads(path.read_text()) + del legacy["artifacts"]["source"]["archived"] + path.write_text(json.dumps(legacy)) + before = path.read_bytes() + assert _resume(path, capsys)["configurations"]["one"]["profile_accepted"] + assert path.read_bytes() == before + source.write_text('{"changed":true}') + resumed = _resume(path, capsys, expected=2) + assert not resumed["configurations"]["one"]["profile_accepted"] + assert resumed["integrity_issues"][0]["artifact"] == "source" + + +@pytest.mark.parametrize( + "patch,accept", + [ + ({"schema_version": "unknown"}, []), + ({"configurations": []}, []), + ({"configurations": {"one": {"acceptance": {"sha256": "0" * 64, "revision": 1}}}}, []), + ({"configurations": {"one": {"draft_request": {}}}}, ["one"]), + ({"configurations": {"one": {"progress": {"stage": 7}}}}, []), + ({"artifacts": {"missing": {"path": "missing.json"}}}, []), + ({"artifacts": {"self": {"path": "checkpoint.json"}}}, []), + ({"artifacts": {"self": {"path": ".checkpoint.json.lock"}}}, []), + ({"artifacts": {"fake": {"path": "missing.json", "sha256": "0" * 64}}}, []), + ], +) +def test_invalid_updates_preserve_previous_checkpoint(tmp_path, capsys, monkeypatch, patch, accept): + path = tmp_path / "checkpoint.json" + _save(path, capsys, monkeypatch) + before = path.read_bytes() + with pytest.raises(SystemExit) as error: + _save(path, capsys, monkeypatch, patch, 1, accept=accept) + assert error.value.code == 2 + assert path.read_bytes() == before + capsys.readouterr() + + +def test_stale_writer_and_interrupted_atomic_save_preserve_checkpoint(tmp_path, capsys, monkeypatch): + path = tmp_path / "checkpoint.json" + _save(path, capsys, monkeypatch) + _save(path, capsys, monkeypatch, {"research": {"first": "saved"}}, 1) + before = path.read_bytes() + with pytest.raises(SystemExit): + _save(path, capsys, monkeypatch, {"research": {"lost": "stale"}}, 1) + capsys.readouterr() + assert path.read_bytes() == before + + def interrupt(*args): + raise KeyboardInterrupt + + monkeypatch.setattr(checkpoint.os, "replace", interrupt) + monkeypatch.setattr(sys, "stdin", io.StringIO('{"research":{"second":"never saved"}}')) + assert cli.main(["onboard", "checkpoint", "--file", str(path), "--update", "-", "--expect-revision", "2"]) == 130 + assert path.read_bytes() == before + assert sorted(p.name for p in tmp_path.iterdir()) == [".checkpoint.json.lock", "checkpoint.json"] + + +def test_two_processes_cannot_overwrite_same_revision(tmp_path): + path = tmp_path / "checkpoint.json" + assert _subprocess("checkpoint", "--file", str(path), cwd=tmp_path).returncode == 0 + command = [ + str(Path(sys.executable).parent / "aisimulate"), + "onboard", + "checkpoint", + "--file", + str(path), + "--update", + "-", + "--expect-revision", + "1", + ] + env = {**os.environ, "PYTHONPATH": str(Path(cli.__file__).parents[1])} + processes = [ + subprocess.Popen( + command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env + ) + for _ in range(2) + ] + for index, process in enumerate(processes): + process.stdin.write(json.dumps({"research": {"writer": index}})) + process.stdin.close() + process.stdin = None + results = [process.communicate(timeout=30) for process in processes] + assert sorted(process.returncode for process in processes) == [0, 2], results + assert json.loads(path.read_text())["revision"] == 2 + + +@pytest.mark.parametrize( + "payload", + [ + "{", + "[]", + '{"revision":1}', + '{"schema_version":"future","revision":1}', + '{"schema_version":"aisimulate-onboarding-checkpoint/v1","revision":1,"revision":2}', + ], +) +def test_invalid_saved_checkpoints_are_not_replaced(tmp_path, capsys, monkeypatch, payload): + path = tmp_path / "checkpoint.json" + path.write_text(payload) + for action in (["resume", "--checkpoint"], ["checkpoint", "--file"]): + with pytest.raises(SystemExit) as error: + cli.main(["onboard", *action, str(path)]) + assert error.value.code == 2 + assert path.read_text() == payload + capsys.readouterr() + + +def test_checkpoint_and_lock_symlink_collisions_do_not_replace_files(tmp_path, capsys, monkeypatch): + victim = tmp_path / "victim.json" + victim.write_text("keep me") + path = tmp_path / "checkpoint.json" + path.symlink_to(victim) + with pytest.raises(SystemExit): + _save(path, capsys, monkeypatch) + capsys.readouterr() + assert victim.read_text() == "keep me" + path.unlink() + (tmp_path / ".checkpoint.json.lock").symlink_to(victim) + with pytest.raises(SystemExit): + _save(path, capsys, monkeypatch) + capsys.readouterr() + assert victim.read_text() == "keep me" + assert not path.exists() + + +def test_one_session_keeps_same_topology_precision_variants_independent(tmp_path, capsys, monkeypatch, request_payload): + path = tmp_path / "checkpoint.json" + alternate = deepcopy(request_payload) + alternate["fpm_profile"]["deployments"][0]["kv_cache_dtype"] = "fp8" + alternate["fpm_profile"]["deployments"][0]["resources"]["kv_bytes_per_token"] = 64 + configurations = { + "bf16-tp1": {"draft_request": request_payload}, + "fp8-kv-tp1": {"draft_request": alternate}, + } + saved = _save(path, capsys, monkeypatch, {"configurations": configurations}, accept=list(configurations)) + assert all(config["profile_accepted"] for config in saved["configurations"].values()) + assert saved["verification"]["executed"] is False + # Recording a valid precision profile does not establish collector compatibility. + updated = _save( + path, + capsys, + monkeypatch, + { + "configurations": { + "fp8-kv-tp1": { + "draft_request": { + "fpm_profile": { + "deployments": [ + { + **alternate["fpm_profile"]["deployments"][0], + "resources": { + **alternate["fpm_profile"]["deployments"][0]["resources"], + "weights_bytes": 2000, + }, + } + ] + }, + } + } + } + }, + 1, + ) + assert updated["configurations"]["bf16-tp1"]["profile_accepted"] + assert not updated["configurations"]["fp8-kv-tp1"]["profile_accepted"] + + +def test_locked_session_fails_promptly_and_preserves_previous_revision(tmp_path, capsys, monkeypatch): + import fcntl + + path = tmp_path / "checkpoint.json" + _save(path, capsys, monkeypatch) + before = path.read_bytes() + with (tmp_path / ".checkpoint.json.lock").open("r+") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + with pytest.raises(SystemExit) as error: + _save(path, capsys, monkeypatch, {"research": {"note": "busy"}}, 1) + assert error.value.code == 2 + assert "another checkpoint operation" in capsys.readouterr().err + assert path.read_bytes() == before + + +def test_existing_update_requires_revision_and_nan_is_rejected(tmp_path, capsys, monkeypatch): + path = tmp_path / "checkpoint.json" + _save(path, capsys, monkeypatch) + before = path.read_bytes() + with pytest.raises(SystemExit): + _save(path, capsys, monkeypatch, {"research": {"note": "no revision"}}) + assert "--expect-revision 1" in capsys.readouterr().err + for invalid in ('{"research":{"value":NaN}}', '{"research":{"value":1e1000}}'): + monkeypatch.setattr(sys, "stdin", io.StringIO(invalid)) + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "checkpoint", "--file", str(path), "--update", "-", "--expect-revision", "1"]) + assert error.value.code == 2 + capsys.readouterr() + assert path.read_bytes() == before diff --git a/python/aisimulate/tests/unit/test_onboard_collection_readiness.py b/python/aisimulate/tests/unit/test_onboard_collection_readiness.py new file mode 100644 index 000000000..79f8ece26 --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_collection_readiness.py @@ -0,0 +1,582 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Readiness checks consume actual native artifacts without claiming silicon accuracy.""" + +from __future__ import annotations + +# ruff: noqa: F811 +import hashlib +import json +import shutil +from dataclasses import replace + +import pytest +from collector.fpm_forward import entry, planner, runner + +from aisimulate import main as cli +from aisimulate.support import fpm +from aisimulate.support.collection_readiness import assess_readiness, resume_without_workers + +from .collector.test_fpm_measurement_evidence import _add_measurement_protocol +from .test_support_serving_validation import materialized_workload # noqa: F401 +from .test_support_validation import validation_case # noqa: F401 +from .test_support_validation_workflow import quality_case # noqa: F401 + +pytestmark = pytest.mark.unit + + +def _read(case): + return assess_readiness(case["request"], case["root"], case["checkpoint"].parent) + + +def _decode(report): + return next(item for item in report["selected_cells"] if item["cell"]["workload_kind"] == "decode") + + +def _change_decode(case, *, fake=None, kvwarm="unchanged"): + for cell in case["plan"].cells: + if cell.workload_kind != "decode": + continue + for path in (case["campaign"] / "cells" / cell.cell_id / "raw").rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + if kvwarm is None: + payload.pop("kvwarm", None) + elif kvwarm != "unchanged": + payload["kvwarm"] = kvwarm + if fake is not None: + for index, row in enumerate(payload["results"]): + if fake == "all" or index == 0: + row["kv_seed_regime"] = "fake_fallback" + row["point"]["sample_reasons"] = ["kvwarm_fake_fallback"] + payload["iteration_groups"][index]["point"]["sample_reasons"] = ["kvwarm_fake_fallback"] + _add_measurement_protocol(payload) + path.write_text(json.dumps(payload)) + + +def test_public_readiness_revalidates_without_model_resolution_or_writes(quality_case, monkeypatch, capsys): + case = quality_case + before = {path: path.read_bytes() for path in case["root"].rglob("*") if path.is_file()} + monkeypatch.setattr(planner, "build_collection_plan", lambda **kwargs: pytest.fail("must not resolve models")) + assert ( + cli.main( + [ + "onboard", + "collect-fpm", + "-c", + str(case["root"] / "request.yaml"), + "--output-dir", + str(case["root"]), + "--check-readiness", + ] + ) + == 0 + ) + report = json.loads(capsys.readouterr().out) + assert report["ready_for_full_collection"] + assert {item["cell"]["workload_kind"] for item in report["selected_cells"]} == {"prefill", "decode"} + assert _decode(report)["formal_regime_counts"] == {"real_kv": 12} + assert report["campaigns"][0]["runtime_identity"]["collector_revision"] == "synthetic-workflow-revision" + assert {path: path.read_bytes() for path in case["root"].rglob("*") if path.is_file()} == before + _change_decode(case, fake="all") + assert not _read(case)["ready_for_full_collection"] + decode = _decode(_read(case)) + assert decode["direct_eligible_points"] == 0 + assert decode["native_regime_counts"] == {"fake_fallback": 12} + assert decode["formal_regime_counts"] == {"fake_fallback": 12} + assert "unchanged" in " ".join(decode["blockers"]) + + +def test_mixed_coverage_and_optional_observation_gaps_do_not_claim_complete_qualification(quality_case): + case = quality_case + _change_decode(case, fake="one") + for path in case["campaign"].rglob("fpm-execution-worker-*.json"): + path.unlink() + report = _read(case) + assert report["ready_for_full_collection"] + decode = _decode(report) + assert decode["coverage"] == "partial" + assert decode["direct_eligible_points"] == 11 + assert decode["execution"]["status"] == "incomplete" + assert decode["execution"]["missing_evidence"] + assert "not complete query coverage" in report["scope"] + + +@pytest.mark.parametrize( + "kvwarm", + [ + None, + {"enabled": True, "warm_eligible": True, "skip_reason": "contradiction"}, + {"enabled": True, "warm_eligible": False, "skip_reason": "hybrid_state_layers_unsupported"}, + ], +) +def test_missing_or_contradictory_warm_protocol_blocks_with_raw_diagnostics(quality_case, kvwarm): + case = quality_case + _change_decode(case, fake="all", kvwarm=kvwarm) + report = _read(case) + assert not report["ready_for_full_collection"] + decode = _decode(report) + assert decode["status"] == "blocked" + assert decode["rank_diagnostics"][0]["kvwarm"] == kvwarm + if kvwarm is not None and kvwarm["warm_eligible"]: + assert not decode["rank_diagnostics"][0]["verified"] + assert decode["blockers"] + + +def test_execution_contradiction_blocks_but_does_not_discard_measurements(quality_case): + case = quality_case + path = next(case["campaign"].rglob("fpm-execution-worker-*.json")) + payload = json.loads(path.read_text()) + payload["resolved_config"]["model_config"]["revision"] = "wrong-revision" + path.write_text(json.dumps(payload)) + report = _read(case) + assert not report["ready_for_full_collection"] + failed = next(item for item in report["selected_cells"] if item["status"] == "blocked") + assert failed["native_regime_counts"] + assert failed["execution"]["failures"] + + +@pytest.mark.parametrize("value", [None, [], ["bad"], 42]) +@pytest.mark.parametrize( + ("filename", "field", "diagnostic"), + [ + ("fpm-execution-worker-*.json", None, "runtime execution evidence must be an object"), + ("fpm-execution-worker-*.json", "collector_provenance", "collector_provenance must be an object"), + ("generator-request.json", None, "generated request must be an object"), + ("generator-request.json", "ServiceConfig", "generated ServiceConfig must be an object"), + ("generator-overrides.json", "K8sConfig", "archived K8sConfig must be an object"), + ], +) +def test_public_readiness_blocks_malformed_optional_evidence_without_writes( + quality_case, capsys, filename, field, diagnostic, value +): + case = quality_case + assert _read(case)["ready_for_full_collection"] + decode = next(cell for cell in case["plan"].cells if cell.workload_kind == "decode") + directory = ( + case["campaign"] if filename == "generator-overrides.json" else case["campaign"] / "cells" / decode.cell_id + ) + path = next(directory.rglob(filename)) + payload = json.loads(path.read_text()) + if field is None: + payload = value + else: + payload[field] = value + path.write_text(json.dumps(payload)) + before = {path: hashlib.sha256(path.read_bytes()).hexdigest() for path in case["root"].rglob("*") if path.is_file()} + assert ( + cli.main( + [ + "onboard", + "collect-fpm", + "-c", + str(case["root"] / "request.yaml"), + "--output-dir", + str(case["root"]), + "--check-readiness", + ] + ) + == 1 + ) + report = json.loads(capsys.readouterr().out) + assert report["status"] == "blocked" + assert not report["ready_for_full_collection"] + assert diagnostic in " ".join(report["blockers"]) + if filename != "generator-overrides.json": + assert _decode(report)["status"] == "blocked" + assert _decode(report)["native_regime_counts"] == {"real_kv": 12} + assert ( + next(item for item in report["selected_cells"] if item["cell"]["workload_kind"] == "prefill")["status"] + == "ready" + ) + assert { + path: hashlib.sha256(path.read_bytes()).hexdigest() for path in case["root"].rglob("*") if path.is_file() + } == before + + +def test_later_formal_attempt_cannot_reuse_older_good_smoke(quality_case): + case = quality_case + shutil.copytree(case["campaign"], case["campaign"] / "smoke", ignore=shutil.ignore_patterns("smoke")) + smoke_checkpoint = case["checkpoint"].with_name("fpm_forward_smoke.json") + smoke_checkpoint.write_bytes(case["checkpoint"].read_bytes()) + checkpoint = json.loads(case["checkpoint"].read_text()) + for value in checkpoint["cells"].values(): + value.update(status="failed", attempt_id="later-failing-attempt") + case["checkpoint"].write_text(json.dumps(checkpoint)) + report = _read(case) + assert not report["ready_for_full_collection"] + assert all(item["campaign"] == "formal" for item in report["selected_cells"]) + assert all("attempt mismatch" in " ".join(item["blockers"]) for item in report["selected_cells"]) + + +def test_resolved_smoke_defaults_cover_all_cells_without_assuming_order(quality_case, monkeypatch): + case = quality_case + # A generic plan can have multiple prefill cells before any decode cell. + plan = replace(case["plan"], cells=case["plan"].cells + case["plan"].cells) + monkeypatch.setattr(entry, "resolve_run_inputs", lambda *_: (plan, {})) + command = fpm.fpm_cli_args(case["request"], output_dir=case["root"], plan_only=False, smoke=True) + args, _ = fpm._resolve_execution(command) + assert args.limit == 4 + args, _ = fpm._resolve_execution([*command, "--limit", "1"]) + assert args.limit == 1 + + +@pytest.mark.parametrize("quality_case", [False, True], indirect=True, ids=["ordinary", "instrumented"]) +@pytest.mark.parametrize("state", ["passed", "failed", "cleanup_failed"]) +def test_recovery_bypass_uses_real_runner_with_every_resource_launch_forbidden(quality_case, monkeypatch, state): + case = quality_case + _change_decode(case, fake="all") + checkpoint = json.loads(case["checkpoint"].read_text()) + checkpoint.pop("database", None) + # Model an unpublished attempt, so recovery can commit its native rows. + shutil.rmtree(case["root"] / "systems/data") + (case["root"] / "systems/data").mkdir() + for value in checkpoint["cells"].values(): + value["status"] = state + if state == "cleanup_failed": + value["cleanup_error"] = "preserved unresolved teardown" + case["checkpoint"].write_text(json.dumps(checkpoint)) + assert resume_without_workers(case["plan"], case["checkpoint"]) + monkeypatch.setattr(runner, "_cell_runner", lambda *_: pytest.fail("recovery must not touch resources")) + monkeypatch.setattr(runner, "_render_cell", lambda *_: pytest.fail("recovery must not launch a cell")) + assert not _read(case)["ready_for_full_collection"] + args = ( + __import__("collector.fpm_forward.cli", fromlist=["_parser"]) + ._parser() + .parse_args(fpm.fpm_cli_args(case["request"], output_dir=case["root"], plan_only=False, resume=True)[3:]) + ) + monkeypatch.setattr( + fpm, + "_resolve_execution", + lambda _: (args, (case["plan"], {"K8sConfig": {"k8s_image": "example/runtime@sha256:" + "a" * 64}})), + ) + assert fpm.run_fpm(case["request"], output_dir=case["root"], execute=True, resume=True) == 1 + report = json.loads((case["root"] / "fpm-readiness.json").read_text()) + assert not report["ready_for_full_collection"] + assert report["recovery_only"] + if not case["plan"].runtime_instrumentation and state != "cleanup_failed": + assert report["collector_exit_status"] == 0 + assert _decode(report)["formal_regime_counts"] == {"fake_fallback": 12} + + +@pytest.mark.parametrize("state", ["running", "interrupted", "pending", "missing", "malformed"]) +def test_unsafe_or_incomplete_resume_does_not_bypass_readiness(quality_case, monkeypatch, state): + case = quality_case + _change_decode(case, fake="all") + checkpoint = json.loads(case["checkpoint"].read_text()) + key = next(iter(checkpoint["cells"])) + if state == "missing": + del checkpoint["cells"][key] + elif state == "malformed": + checkpoint["cells"][key] = "passed" + else: + checkpoint["cells"][key]["status"] = state + case["checkpoint"].write_text(json.dumps(checkpoint)) + assert not resume_without_workers(case["plan"], case["checkpoint"]) + args = ( + __import__("collector.fpm_forward.cli", fromlist=["_parser"]) + ._parser() + .parse_args(fpm.fpm_cli_args(case["request"], output_dir=case["root"], plan_only=False, resume=True)[3:]) + ) + monkeypatch.setattr(fpm, "_resolve_execution", lambda _: (args, (case["plan"], {}))) + monkeypatch.setattr(entry, "run_resolved", lambda *_: pytest.fail("blocked before runner")) + assert fpm.run_fpm(case["request"], output_dir=case["root"], execute=True, resume=True) == 1 + + +def test_fresh_full_collection_requires_saved_readiness(quality_case, monkeypatch): + case = quality_case + shutil.rmtree(case["campaign"]) + case["checkpoint"].unlink() + shutil.rmtree(case["root"] / "systems/data") + (case["root"] / "systems/data").mkdir() + args = ( + __import__("collector.fpm_forward.cli", fromlist=["_parser"]) + ._parser() + .parse_args(fpm.fpm_cli_args(case["request"], output_dir=case["root"], plan_only=False)[3:]) + ) + monkeypatch.setattr(fpm, "_resolve_execution", lambda _: (args, (case["plan"], {}))) + monkeypatch.setattr(entry, "run_resolved", lambda *_: pytest.fail("must smoke before full collection")) + assert fpm.run_fpm(case["request"], output_dir=case["root"], execute=True) == 1 + + +def test_approved_legacy_tp_skip_retains_formal_consumer_eligibility(quality_case): + case = quality_case + _change_decode( + case, + fake="all", + kvwarm={"enabled": True, "warm_eligible": False, "skip_reason": "moe_tp_balanced_by_construction"}, + ) + report = _read(case) + assert report["ready_for_full_collection"] + decode = _decode(report) + assert decode["native_regime_counts"] == {"fake_fallback": 12} + assert decode["formal_regime_counts"] == {"skip:moe_tp_balanced_by_construction": 12} + assert decode["direct_eligible_points"] == 12 + + +@pytest.mark.parametrize("value", [[], 42, None, {"cells": []}]) +def test_malformed_checkpoint_reports_blocker_without_crashing(quality_case, value): + case = quality_case + case["checkpoint"].write_text(json.dumps(value)) + report = _read(case) + assert not report["ready_for_full_collection"] + assert report["blockers"] + assert not resume_without_workers(case["plan"], case["checkpoint"]) + + +def test_malformed_native_artifact_preserves_independent_phase_result(quality_case): + case = quality_case + decode = next(cell for cell in case["plan"].cells if cell.workload_kind == "decode") + path = next((case["campaign"] / "cells" / decode.cell_id).rglob("benchmark-dp*.json")) + path.write_text("[]") + report = _read(case) + assert not report["ready_for_full_collection"] + assert ( + next(item for item in report["selected_cells"] if item["cell"]["workload_kind"] == "prefill")["status"] + == "ready" + ) + + +def test_consistent_but_wrong_observed_runtime_version_is_blocked(quality_case): + case = quality_case + for path in case["campaign"].rglob("collector-provenance.json"): + payload = json.loads(path.read_text()) + payload["runtime"]["backend_version"] = "0.99.0" + path.write_text(json.dumps(payload)) + report = _read(case) + assert not report["ready_for_full_collection"] + assert all(item["status"] == "blocked" for item in report["selected_cells"]) + + +def _bounded_source(case, tmp_path, *, explicit=False): + from aisimulate.support.plan import create_plan + from aisimulate.support.schema import SupportRequest + + request_payload = case["request"].model_dump(mode="json") + if explicit: + request_payload["collection"].update(prefill_cudagraph_policy="explicit", max_prefill_cudagraph_size=16) + request = SupportRequest.model_validate(request_payload) + root = tmp_path / "bounded" + create_plan(request, root) + model_config = tmp_path / "bounded-model-config.json" + model_config.write_text(json.dumps(case["plan"].capability.model_config.payload)) + overrides = {"K8sConfig": {"k8s_image": "example/runtime@sha256:" + "a" * 64}} + plan = planner.build_collection_plan( + backend="vllm", + model_path=request.identity.model, + system=request.identity.gpu, + selected_ops=set(), + model_architecture=request.fpm_profile.architecture, + model_config_path=str(model_config), + fpm_profile=request.fpm_profile, + generator_overrides=overrides, + options=replace( + case["plan"].options, + benchmark_points_json=None, + benchmark_points_sha256=None, + prefill_cudagraph_policy=request.collection.prefill_cudagraph_policy, + max_prefill_cudagraph_size=request.collection.max_prefill_cudagraph_size, + ), + ) + return request, root, plan, overrides + + +def _write_bounded_native(plan, root, overrides, *, smoke): + from collector.fpm_forward.config import with_kv_warmup_defaults + + from .collector.test_fpm_runner import _native_payload, _write_provenance + + campaign = root / "fpm-artifacts" / plan.sha256[:16] + if smoke: + campaign /= "smoke" + campaign.mkdir(parents=True, exist_ok=True) + (campaign / "collection-plan.json").write_text(json.dumps(plan.to_dict())) + (campaign / "generator-overrides.json").write_text(json.dumps(with_kv_warmup_defaults(overrides))) + entries = {} + for cell in plan.cells: + directory = campaign / "cells" / cell.cell_id + raw = directory / "raw/pod" + raw.mkdir(parents=True) + runner._render_cell(plan, cell, directory, overrides, smoke=smoke) + provenance = raw / "collector-provenance.json" + _write_provenance(provenance, cell_id=cell.cell_id, plan_sha256=plan.sha256, attempt_id="bounded") + payload = json.loads(provenance.read_text()) + payload["runtime"]["backend_version"] = plan.capability.aic_database_version + provenance.write_text(json.dumps(payload)) + for rank in range(cell.topology.dp): + native = _native_payload(phase=cell.workload_kind, rank=rank, dp=cell.topology.dp) + regime, stamp = ( + ("real_kv", "kvwarm_real_kv") + if cell.workload_kind == "decode" + else ("real_prefix", "prefill_real_seed") + ) + native["results"][0]["kv_seed_regime"] = regime + native["results"][0]["point"]["sample_reasons"] = [stamp] + (raw / f"benchmark-dp{rank}.json").write_text(json.dumps(native)) + entries[cell.cell_id] = {"status": "passed", "attempt_id": "bounded"} + checkpoint = root / "fpm-checkpoint" / ("fpm_forward_smoke.json" if smoke else "fpm_forward.json") + checkpoint.parent.mkdir(exist_ok=True) + checkpoint.write_text( + json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": entries}) + ) + return campaign + + +def test_public_smoke_covers_both_phases_then_matching_full_collection_can_start(quality_case, tmp_path, monkeypatch): + request, root, plan, overrides = _bounded_source(quality_case, tmp_path) + calls = [] + monkeypatch.setattr(entry, "resolve_run_inputs", lambda *_: (plan, overrides)) + + def synthetic_collection(args, resolved): + calls.append((args.smoke, args.limit)) + _write_bounded_native(resolved[0], root, resolved[1], smoke=args.smoke) + if not args.smoke: + return runner.run_collection( + resolved[0], + generator_overrides=resolved[1], + checkpoint_dir=str(root / "fpm-checkpoint"), + artifact_root=str(root / "fpm-artifacts"), + database_root=str(root / "systems/data"), + resume=True, + retry_failed=False, + ) + return [] + + monkeypatch.setattr(entry, "run_resolved", synthetic_collection) + command = ["onboard", "collect-fpm", "-c", str(root / "request.yaml"), "--output-dir", str(root), "--execute"] + assert cli.main([*command, "--smoke"]) == 0 + assert calls == [(True, len(plan.cells))] + report = json.loads((root / "fpm-readiness.json").read_text()) + assert report["ready_for_full_collection"] + assert {item["cell"]["workload_kind"] for item in report["selected_cells"]} == {"prefill", "decode"} + assert cli.main(command) == 0 + assert calls == [(True, len(plan.cells)), (False, None)] + assert all( + item["campaign"] == "formal" for item in json.loads((root / "fpm-readiness.json").read_text())["selected_cells"] + ) + + +def test_old_explicit_smoke_with_missing_or_changed_capture_flags_is_blocked(quality_case, tmp_path): + import re + import shlex + + request, root, plan, overrides = _bounded_source(quality_case, tmp_path, explicit=True) + campaign = _write_bounded_native(plan, root, overrides, smoke=True) + assert assess_readiness(request, root, root / "fpm-checkpoint")["ready_for_full_collection"] + prefill = next(cell for cell in plan.cells if cell.workload_kind == "prefill") + path = campaign / "cells" / prefill.cell_id / "run.sh" + script = path.read_text() + command = re.search(r"^engine_command=\((.*)\)$", script, re.MULTILINE) + argv = shlex.split(command[1]) + index = argv.index("--compilation-config") + for config in (None, json.dumps({"cudagraph_capture_sizes": [1], "max_cudagraph_capture_size": 1})): + edited = list(argv) + if config is None: + del edited[index : index + 2] + else: + edited[index + 1] = config + path.write_text(script.replace(command[0], "engine_command=(" + shlex.join(edited) + ")")) + report = assess_readiness(request, root, root / "fpm-checkpoint") + assert not report["ready_for_full_collection"] + assert "explicit graph configuration" in " ".join(report["blockers"]) + + +@pytest.mark.parametrize( + ("explicit", "phase", "graph", "ready"), + [ + pytest.param( + True, "prefill", {"capture_sizes": [1, 2, 4, 8], "max_capture_size": 8}, False, id="contradictory-native" + ), + pytest.param(True, "prefill", {"capture_sizes": [1, 2, 4, 8]}, False, id="capture-list-only"), + pytest.param(True, "prefill", {"max_capture_size": 8}, False, id="maximum-only"), + pytest.param( + True, + "prefill", + { + "capture_sizes": [1, 2, 4, 8, 16], + "max_capture_size": 16, + "prefill_capture_sizes": [], + "decode_capture_sizes": [1, 2], + }, + True, + id="matching-full-captures-with-filtered-phase-lists", + ), + pytest.param(True, "prefill", None, True, id="missing-optional"), + pytest.param( + False, "prefill", {"capture_sizes": [1, 2, 4, 8], "max_capture_size": 8}, True, id="runtime-policy" + ), + pytest.param(True, "decode", {"capture_sizes": [1, 2, 4, 8], "max_capture_size": 8}, True, id="decode-default"), + ], +) +def test_public_readiness_checks_native_explicit_prefill_graph_evidence( + quality_case, tmp_path, capsys, explicit, phase, graph, ready +): + request, root, plan, overrides = _bounded_source(quality_case, tmp_path, explicit=explicit) + campaign = _write_bounded_native(plan, root, overrides, smoke=True) + assert assess_readiness(request, root, root / "fpm-checkpoint")["ready_for_full_collection"] + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + path = next((campaign / "cells" / cell.cell_id / "raw").rglob("benchmark-dp*.json")) + payload = json.loads(path.read_text()) + if graph is not None: + payload["cudagraph"] = graph + else: + payload.pop("cudagraph", None) + path.write_text(json.dumps(payload)) + before = {path: hashlib.sha256(path.read_bytes()).hexdigest() for path in root.rglob("*") if path.is_file()} + assert cli.main( + [ + "onboard", + "collect-fpm", + "-c", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--check-readiness", + ] + ) == (0 if ready else 1) + report = json.loads(capsys.readouterr().out) + assert report["ready_for_full_collection"] == ready + prefill = next(item for item in report["selected_cells"] if item["cell"]["workload_kind"] == "prefill") + assert prefill["execution"]["observed_workers"] == [] + if not ready: + assert prefill["status"] == "blocked" + assert "native prefill graph configuration differs" in " ".join(prefill["blockers"]) + assert _decode(report)["status"] == "ready" + assert {path: hashlib.sha256(path.read_bytes()).hexdigest() for path in root.rglob("*") if path.is_file()} == before + + +def test_missing_evidence_cannot_be_replaced_by_a_cached_ready_report(quality_case, monkeypatch): + case = quality_case + (case["root"] / "fpm-readiness.json").write_text(json.dumps({"status": "ready", "ready_for_full_collection": True})) + decode = next(cell for cell in case["plan"].cells if cell.workload_kind == "decode") + shutil.rmtree(case["campaign"] / "cells" / decode.cell_id / "raw") + report = _read(case) + assert not report["ready_for_full_collection"] + assert _decode(report)["status"] == "incomplete" + + +def test_malformed_formal_cell_overrides_good_smoke_even_when_null(quality_case): + case = quality_case + shutil.copytree(case["campaign"], case["campaign"] / "smoke", ignore=shutil.ignore_patterns("smoke")) + case["checkpoint"].with_name("fpm_forward_smoke.json").write_bytes(case["checkpoint"].read_bytes()) + checkpoint = json.loads(case["checkpoint"].read_text()) + decode = next(cell for cell in case["plan"].cells if cell.workload_kind == "decode") + checkpoint["cells"][decode.cell_id] = None + case["checkpoint"].write_text(json.dumps(checkpoint)) + report = _read(case) + assert not report["ready_for_full_collection"] + assert _decode(report)["campaign"] == "formal" + + +@pytest.mark.parametrize("regime", [42, ["real_kv"], {"regime": "real_kv"}]) +def test_malformed_native_regime_remains_json_reportable(quality_case, regime): + case = quality_case + decode = next(cell for cell in case["plan"].cells if cell.workload_kind == "decode") + path = next((case["campaign"] / "cells" / decode.cell_id).rglob("benchmark-dp*.json")) + payload = json.loads(path.read_text()) + payload["results"][0]["kv_seed_regime"] = regime + path.write_text(json.dumps(payload)) + report = _read(case) + assert not report["ready_for_full_collection"] + assert _decode(report)["rank_diagnostics"][0]["regime_counts"]["invalid"] == 1 + assert json.loads(json.dumps(report, sort_keys=True))["status"] == "blocked" diff --git a/python/aisimulate/tests/unit/test_onboard_finalization.py b/python/aisimulate/tests/unit/test_onboard_finalization.py new file mode 100644 index 000000000..d3ce4e4fc --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_finalization.py @@ -0,0 +1,848 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public onboarding with independent synthetic initialization/native evidence.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import subprocess +import sys +from dataclasses import replace +from pathlib import Path + +import pytest + +import aisimulate.main as cli +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.checkpoint import save_checkpoint +from aisimulate.support.finalization import finalize +from aisimulate.support.plan import check_plan +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +def _write(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value)) + + +def _native(phase): + """Authored schema-v2 timings, deliberately no production observer helpers.""" + points, iterations = [], [] + for index, kv in enumerate([0] if phase == "prefill" else [8, 32], 1): + prefill = 16 if phase == "prefill" else 0 + point = { + "point_type": phase, + "benchmark_id": index, + "total_prefill_tokens": prefill, + "total_kv_read_tokens": kv, + "batch_size": 1, + "expected_cudagraph_mode": "PIECEWISE" if prefill else "FULL", + "expected_capture_size": prefill or 1, + "padding_tokens": 0, + "sample_reasons": ["capture"] if prefill else ["capture", "kvwarm_real_kv"], + } + fpm = { + "counter_id": index, + "dp_rank": 0, + "wall_time": 0.001 * index, + "scheduled_requests": { + "num_prefill_requests": 1 if prefill else 0, + "sum_prefill_tokens": prefill, + "sum_prefill_kv_tokens": kv if prefill else 0, + "num_decode_requests": 0 if prefill else 1, + "sum_decode_kv_tokens": 0 if prefill else kv, + }, + } + points.append({"point": point, "fpms": [fpm]}) + iterations.append( + { + "benchmark_id": index, + "point": point, + "expected_dp_ranks": [0], + "complete": True, + "wall_time": fpm["wall_time"], + "rank_results": [{"dp_rank": 0, "fpms": [fpm]}], + } + ) + result = { + "schema_version": 2, + "artifact_type": "rank", + "status": "complete", + "valid": True, + "usable": True, + "timing_valid": True, + "run_id": f"synthetic-{phase}", + "grid_digest": f"synthetic-{phase}-grid", + "config": {"mode": phase}, + "coverage": {"expected_points": len(points), "completed_points": len(points), "skipped_points": 0}, + "dp": {"rank": 0, "size": 1}, + "results": points, + "iteration_groups": iterations, + "skipped_points": [], + "missing_phases": [], + "timing": { + "benchmark_elapsed_seconds": 1.0, + "measured_iteration_seconds": sum(item["wall_time"] for item in iterations), + }, + } + if phase == "decode": + result["kvwarm"] = {"enabled": True, "warm_eligible": True, "skip_reason": None} + return result + + +def _memory(provenance, phase, model, revision): + count = 1000 if phase == "prefill" else 900 + group = { + "layer_names": ["model.layers.0.attn", "model.layers.1.attn"], + "spec_type": "vllm.v1.kv_cache_interface.FullAttentionSpec", + "block_size_tokens": 16, + "spec_page_size_bytes": 512, + "sliding_window": None, + "dtype": "torch.bfloat16", + } + settings = { + "model_config": { + "model": model, + "revision": revision, + "dtype": "torch.bfloat16", + "quantization": None, + "max_model_len": 128, + "enforce_eager": False, + }, + "cache_config": { + "cache_dtype": "bfloat16", + "gpu_memory_utilization": 0.9, + "kv_cache_memory_bytes": None, + "enable_prefix_caching": phase == "decode", + }, + "scheduler_config": {"max_num_batched_tokens": 64, "max_num_seqs": 4, "async_scheduling": False}, + "parallel_config": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "enable_expert_parallel": False, + "decode_context_parallel_size": 1, + "prefill_context_parallel_size": 1, + }, + "compilation_config": { + "mode": "VLLM_COMPILE", + "cudagraph_mode": "FULL_AND_PIECEWISE", + "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64], + "max_cudagraph_capture_size": 64, + }, + "kernel_config": {"moe_backend": "auto"}, + "quantization_config": None, + "offload_config": { + "offload_backend": "auto", + "uva": {"cpu_offload_gb": 0}, + "prefetch": {"offload_group_size": 0, "offload_num_in_group": 1, "offload_prefetch_step": 1}, + }, + } + common = { + "schema_name": "aisimulate_fpm_runtime_memory", + "schema_version": 1, + "collector_provenance": provenance, + "backend_version": "0.27.0", + "status": "resolved", + "dp_rank": 0, + "resolved_config": settings, + } + worker = { + **common, + "kind": "worker", + "tp_rank": 0, + "pp_rank": 0, + "effective_offloader": "vllm.model_executor.offloader.base.NoopOffloader", + "cache": { + "num_blocks": count, + "allocated_cache_bytes": count * 1024, + "available_cache_bytes": count * 1024 + 512, + "storages": [ + {"device": "cuda:0", "pointer": 100, "size_bytes": count * 512}, + {"device": "cuda:0", "pointer": 200, "size_bytes": count * 512}, + ], + "groups": [ + { + **group, + "kind": "attention", + "layer_classes": dict.fromkeys( + group["layer_names"], "vllm.model_executor.layers.attention.Attention" + ), + } + ], + }, + } + scheduler = { + **common, + "kind": "scheduler", + "tp_rank": None, + "pp_rank": None, + "cache": { + "num_blocks": count, + "initial_free_blocks": count - 1, + "reserved_blocks": 1, + "null_block_id": 0, + "watermark_blocks": 0, + "pool_count": 1, + "groups": [group], + }, + } + return worker, scheduler + + +def _prepare_collection(tmp_path: Path, extra_init_args: tuple[str, ...] = ()) -> tuple[SupportRequest, Path]: + source = tmp_path / "config.json" + _write( + source, + { + "_name_or_path": "example/runtime-memory-test", + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "vocab_size": 64, + "max_position_embeddings": 128, + "torch_dtype": "bfloat16", + }, + ) + # Keep checkpoint metadata independently available after the onboarding + # input is removed; native profile construction verifies its architecture. + checkpoint = tmp_path / "checkpoint" + _write(checkpoint / "config.json", json.loads(source.read_text())) + overrides = tmp_path / "overrides.json" + _write( + overrides, + { + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "bfloat16", + "max_num_tokens": 64, + "max_batch_size": 4, + }, + ) + original = tmp_path / "draft.yaml" + assert ( + cli.main( + [ + "onboard", + "init", + "--model-config", + str(source), + "--model", + str(checkpoint), + "--resource-overrides", + str(overrides), + "--model-revision", + "synthetic-immutable-revision", + "--framework-version", + "0.27.0", + "--gpu", + "h200_sxm", + "--interconnect", + "nvswitch", + "--tensor-parallel", + "1", + "--input-tokens", + "16", + "--output-tokens", + "4", + "--concurrency", + "1", + "--request-count", + "1", + "--context-length", + "128", + "--output", + str(original), + *extra_init_args, + ] + ) + == 0 + ) + root = (tmp_path / "collection").resolve() + assert cli.main(["onboard", "plan", "--config", str(original), "--output-dir", str(root)]) == 0 + request = SupportRequest.from_yaml(root / "request.yaml") + assert request.profile_deployment().resources.memory_source == "pending" + return request, root + + +def build_completed_collection( + tmp_path: Path, *, extra_init_args: tuple[str, ...] = (), plan_changes: dict | None = None +) -> tuple[SupportRequest, Path]: + """Reusable CPU-only fixture for source and installed-wheel CLI checks. + + Timings and cache allocations are synthetic test declarations, not measured + silicon. The production validator/resolver and formal publisher run unchanged. + """ + from collector.fpm_forward import cli as collector_cli + from collector.fpm_forward.config import FPMCollectionOptions + from collector.fpm_forward.database import aggregate_cell, write_formal_database + from collector.fpm_forward.planner import build_collection_plan + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + request, root = _prepare_collection(tmp_path, extra_init_args) + source = tmp_path / "config.json" + manifest = json.loads((root / "support-plan.json").read_text()) + args = collector_cli._parser().parse_args(manifest["fpm"]["plan_command"][3:]) + plan = build_collection_plan( + backend="vllm", + model_path=request.identity.model, + model_architecture=args.model_architecture, + model_config_path=str(source), + system="h200_sxm", + selected_ops={"attention_context", "attention_generation"}, + options=replace(FPMCollectionOptions.from_args(args), **(plan_changes or {})), + fpm_profile=request.fpm_profile, + ) + artifact_root = root / "fpm-artifacts" / plan.sha256[:16] + _write(artifact_root / "collection-plan.json", plan.to_dict()) + entries, rows = {}, [] + for cell in plan.cells: + directory = artifact_root / "cells" / cell.cell_id + _write(directory / "cell.json", cell.to_dict()) + raw = directory / "raw/pod-0" + provenance = { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell.cell_id, + "plan_sha256": plan.sha256, + "attempt_id": f"attempt-{cell.workload_kind}", + "runtime": {"backend": "vllm", "backend_version": "0.27.0"}, + } + _write(raw / "collector-provenance.json", provenance) + _write(raw / "benchmark.json", _native(cell.workload_kind)) + worker, scheduler = _memory( + provenance, cell.workload_kind, request.identity.model, request.identity.model_revision + ) + sampling = plan.options.prefill_sampling + if sampling.cudagraph_policy == "explicit": + worker["resolved_config"]["compilation_config"].update( + cudagraph_capture_sizes=list(sampling.cudagraph_capture_sizes), + max_cudagraph_capture_size=sampling.max_cudagraph_capture_size, + ) + _write(raw / "fpm-memory-worker-dp0-tp0-pp0.json", worker) + _write(raw / "fpm-memory-scheduler-dp0.json", scheduler) + rows.extend(aggregate_cell(plan, cell, directory, expected_attempt_id=provenance["attempt_id"])) + entries[cell.cell_id] = {"status": "passed", "attempt_id": provenance["attempt_id"]} + parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=root / "systems/data") + assert not skipped + _write( + root / "fpm-checkpoint/fpm_forward.json", + { + "schema": CHECKPOINT_SCHEMA, + "plan_sha256": plan.sha256, + "cells": entries, + "database": { + "status": "passed", + "parquet": str(parquet), + "metadata": str(metadata), + "row_count": len(rows), + "published_cells": len(plan.cells), + "plan_cells": len(plan.cells), + "missing_cells": [], + "skipped_first_publisher_wins": [], + }, + }, + ) + # Finalization reuses saved onboarding inputs; the separate checkpoint + # config remains available for the canonical architecture check. + source.unlink() + (tmp_path / "overrides.json").unlink() + return request, root + + +def _snapshot(root): + return {path.relative_to(root): path.read_bytes() for path in root.rglob("*") if path.is_file()} + + +@pytest.mark.parametrize( + "plan_changes", + [{}, {"max_prefill_cudagraph_size": 32}, {"max_prefill_isl": 32}, {"warmup_iterations": 2}], +) +@pytest.mark.parametrize( + "executor,cpu_policy,checkpoint_subdir", + [ + ("kubernetes", None, None), + ("slurm", None, None), + ("slurm", (16, "cores"), None), + ("slurm", (32, "none"), None), + ("slurm", (24, "cores"), None), + ("slurm", None, "nested"), + ("slurm", (32, "none"), "nested"), + ], +) +def test_public_finalize_binds_reviewed_collection_options( + tmp_path, capsys, plan_changes, executor, cpu_policy, checkpoint_subdir +): + deployment = ( + { + "executor": "slurm", + "slurm_container_image": "image@sha256:synthetic", + "slurm_container_mounts": ("/cache:/cache",), + } + if executor == "slurm" + else {} + ) + if cpu_policy is not None: + deployment.update(slurm_cpus_per_task=cpu_policy[0], slurm_cpu_bind=cpu_policy[1]) + _request, root = build_completed_collection( + tmp_path, + extra_init_args=("--prefill-cudagraph-policy", "explicit", "--max-prefill-cudagraph-size", "64"), + plan_changes={**deployment, **plan_changes}, + ) + if checkpoint_subdir is not None: + checkpoint = root / "fpm-checkpoint" / "fpm_forward.json" + nested = checkpoint.parent / checkpoint_subdir / checkpoint.name + nested.parent.mkdir() + checkpoint.rename(nested) + target = tmp_path / "resolved" + before = _snapshot(root) + try: + status = cli.main( + [ + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + except SystemExit as error: + status = error.code + assert status == (2 if plan_changes else 0) + assert _snapshot(root) == before + if plan_changes: + assert "saved collection options differ" in capsys.readouterr().err + assert not target.exists() + else: + assert SupportRequest.from_yaml(target / "request.yaml").collection.max_prefill_cudagraph_size == 64 + + +@pytest.mark.parametrize("tamper", [None, "image", "mounts", "executor", "worker"]) +def test_public_slurm_collect_to_finalize_preserves_frozen_deployment(tmp_path, monkeypatch, capsys, tamper): + """Real CLI, Slurm campaign and finalizer; only cluster/runtime output is synthetic.""" + from collector.fpm_forward import runner + from collector.fpm_forward.runtime import fpm_memory_observer as observer + + from .collector.test_fpm_cpu_affinity import _snapshot as cpu_snapshot + + request, root = _prepare_collection(tmp_path) + image = "registry.example/fpm@sha256:" + "a" * 64 + commands = [] + + def cluster_command(args, **_kwargs): + commands.append(args) + if args[0] == "scontrol": + output = "JobId=1234 JobState=RUNNING NodeList=node-a" if "job" in args else "node-a" + return subprocess.CompletedProcess(args, 0, stdout=output, stderr="") + if args[0] == "squeue": + return subprocess.CompletedProcess(args, 0, stdout="1234.99|unrelated-job", stderr="") + assert args[0] == "srun", args + assert "--jobid=1234" in args and "--gpus-per-node=1" in args + assert "--cpus-per-task=16" in args and "--cpu-bind=cores" in args + assert f"--container-image={image}" in args + mounts = next(value.split("=", 1)[1] for value in args if value.startswith("--container-mounts=")) + assert mounts.startswith("/cache:/cache,/models:/models,") + raw = Path(next(value.rsplit(":", 1)[0] for value in mounts.split(",") if value.endswith(":/results"))) + command = args[args.index("/usr/bin/env") + 6 :] + if command[:2] == ["python3", "-c"]: + with monkeypatch.context() as patch: + patch.setattr(sys, "argv", ["-c", *command[3:]]) + patch.setattr(importlib.metadata, "version", lambda _name: "0.27.0") + exec(command[2].replace("/results", str(raw)), {}) + else: + assert command == ["bash", "/tmp/fpm-bench/fpm_exec.sh"] + cell = json.loads((raw.parent.parent / "cell.json").read_text()) + provenance = json.loads((raw / "collector-provenance.json").read_text()) + _write(raw / "benchmark.json", _native(cell["workload_kind"])) + worker, scheduler = _memory( + provenance, cell["workload_kind"], request.identity.model, request.identity.model_revision + ) + _write(raw / "fpm-memory-worker-dp0-tp0-pp0.json", worker) + _write(raw / "fpm-memory-scheduler-dp0.json", scheduler) + with monkeypatch.context() as patch: + for kind, pid, coordinates in ( + ("launcher", 100, {"requested_cpus_per_task": 16, "cpu_bind": "cores", "local_gpu_count": 1}), + ("worker", 200, {"dp_rank": 0, "tp_rank": 0, "pp_rank": 0}), + ("scheduler", 300, {"dp_rank": 0}), + ): + patch.setattr( + observer, "cpu_affinity_snapshot", lambda pid=pid: cpu_snapshot("node-a", pid, list(range(16))) + ) + observer.observe_cpu(kind, directory=raw, **coordinates) + return subprocess.CompletedProcess(args, 0, stdout="synthetic Slurm runtime", stderr="") + + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + monkeypatch.setenv("SLURM_JOB_ID", "1234") + monkeypatch.setenv("SLURM_JOB_CPUS_PER_NODE", "16") + monkeypatch.setattr("collector.fpm_forward.slurm.shutil.which", lambda name: f"/fake/{name}") + monkeypatch.setattr(runner, "_run_command", cluster_command) + command = [ + "onboard", + "collect-fpm", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--executor", + "slurm", + "--image", + image, + "--container-mount", + "/cache:/cache", + "--container-mount", + "/models:/models", + "--execute", + ] + assert cli.main([*command, "--smoke"]) == 0 + smoke = json.loads((root / "fpm-readiness.json").read_text()) + assert smoke["ready_for_full_collection"] + assert {item["cell"]["workload_kind"] for item in smoke["selected_cells"]} == {"prefill", "decode"} + assert cli.main(command) == 0 + collection_path = next(root.glob("fpm-artifacts/*/collection-plan.json")) + payload = json.loads(collection_path.read_text()) + assert payload["options"]["executor"] == "slurm" + assert payload["options"]["slurm_container_image"] == image + assert payload["options"]["slurm_container_mounts"] == ["/cache:/cache", "/models:/models"] + if tamper == "worker": + next(root.glob("fpm-artifacts/*/cells/*/raw/node0000/fpm-memory-worker*.json")).unlink() + elif tamper is not None: + key, value = { + "image": ("slurm_container_image", "changed-image"), + "mounts": ("slurm_container_mounts", ["/different:/cache"]), + "executor": ("executor", "kubernetes"), + }[tamper] + payload["options"][key] = value + _write(collection_path, payload) + (tmp_path / "config.json").unlink() + before, calls = _snapshot(root), list(commands) + target = tmp_path / "resolved" + try: + status = cli.main( + [ + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + except SystemExit as error: + status = error.code + assert status == (2 if tamper else 0) + assert _snapshot(root) == before + assert commands == calls # Finalization never starts or reconnects to the executor. + if tamper: + diagnostic = "worker rank evidence is incomplete" if tamper == "worker" else "SHA-256" + assert diagnostic in capsys.readouterr().err + assert not target.exists() + else: + resolved = SupportRequest.from_yaml(target / "request.yaml") + assert resolved.profile_deployment().resources.runtime_memory.kv_cache_bytes == 899 * 1024 + check_plan(resolved, target) + + +@pytest.mark.parametrize( + "revision,loaded,accepted", + [ + ("2" * 40, None, False), + (None, "2" * 40, False), + ("main", "1" * 40, True), + (None, None, True), + ("1" * 40, "1" * 40, True), + ], +) +def test_public_finalize_checks_positive_immutable_revision_evidence(tmp_path, capsys, revision, loaded, accepted): + _request, root = build_completed_collection(tmp_path, extra_init_args=("--model-revision", "1" * 40)) + for path in root.glob("fpm-artifacts/*/cells/*/raw/*/fpm-memory-*.json"): + evidence = json.loads(path.read_text()) + model = evidence["resolved_config"]["model_config"] + model.update(model="/mnt/pvc/local-checkpoint", revision=revision) + if loaded is not None: + model["loaded_config_commit_hash"] = loaded + _write(path, evidence) + target = tmp_path / "resolved" + before = _snapshot(root) + try: + status = cli.main( + [ + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + except SystemExit as error: + status = error.code + assert status == (0 if accepted else 2) + assert _snapshot(root) == before + if accepted: + assert SupportRequest.from_yaml(target / "request.yaml").identity.model_revision == "1" * 40 + else: + assert "contradicts the reviewed immutable revision" in capsys.readouterr().err + assert not target.exists() + + +def test_public_finalize_preserves_collection_and_resolves_ordinary_configs(tmp_path, capsys): + from aisimulate_core.sdk.memory import estimate_kv_cache + + request, root = build_completed_collection(tmp_path) + checkpoint = tmp_path / "session.json" + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp1": {"draft_request": request.model_dump(mode="json")}}}, + expected_revision=None, + accept=["tp1"], + ) + assert state.configurations["tp1"].acceptance is not None + before, session = _snapshot(root), checkpoint.read_bytes() + target = tmp_path / "resolved" + assert ( + cli.main( + [ + "onboard", + "finalize", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + == 0 + ) + assert _snapshot(root) == before + assert checkpoint.read_bytes() == session + resolved = SupportRequest.from_yaml(target / "request.yaml") + assert resolved.profile_deployment().resources.memory_source == "runtime" + resources = resolved.profile_deployment().resources + assert resources.runtime_memory.kv_cache_bytes == 899 * 1024 + assert resources.weights_bytes is resources.activations_bytes is None + assert len(resources.cache_groups) == 1 + assert resources.cache_groups[0].page_size_bytes == 1024 + assert resources.cache_groups[0].num_layers == 2 + manifest = json.loads((target / "finalization.json").read_text()) + assert manifest["review_status"] == "requires_review" + assert len(manifest["cell_observations"]) == 2 + for item in manifest["formal_data"]: + copied = target / item["relative_path"] + assert copied.read_bytes() == (root / item["relative_path"]).read_bytes() + assert hashlib.sha256(copied.read_bytes()).hexdigest() == item["sha256"] + assert sorted(path.name for path in (target / "systems/data").rglob("*") if path.is_file()) == [ + "fpm_forward_perf.metadata.json", + "fpm_forward_perf.parquet", + ] + check_plan(resolved, target) + for config in ( + CorePredictionConfig.from_yaml(target / "predict/pilot.yaml"), + CoreRecommendationConfig.from_yaml(target / "recommend/pilot.yaml"), + ): + assert config.engine.fpm_profile == resolved.fpm_profile + assert not config.engine.workers.aggregated.kv_cache.prefix_caching + assert [str(path) for path in config.engine.systems_paths] == [str(target / "systems")] + estimate = estimate_kv_cache( + resolved.identity.model, + "h200_sxm", + "vllm", + backend_version="0.27.0", + tp_size=1, + pp_size=1, + attention_dp_size=1, + moe_tp_size=1, + moe_ep_size=1, + max_num_tokens=64, + max_batch_size=4, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + context_length=128, + fpm_profile=resolved.fpm_profile, + ) + assert estimate["total_kv_size_bytes"] == 899 * 1024 + assert json.loads(resolved.fpm_profile.provenance)["memory_source"] == "runtime" + for command in ("predict", "recommend"): + assert ( + cli.main( + [ + command, + "--config", + str(target / command / "pilot.yaml"), + "--output-dir", + str(tmp_path / f"{command}-result"), + "--format", + "json", + ] + ) + == 0 + ) + trace = tmp_path / "trace.jsonl" + _write( + trace, + { + "id": "synthetic-play", + "models": ["trace/source-model"], + "block_size": 16, + "hash_id_scope": "local", + "requests": [{"t": 0.0, "type": "s", "model": "trace/source-model", "in": 16, "out": 4, "hash_ids": [1]}], + }, + ) + assert ( + cli.main( + [ + "onboard", + "validate-fpm", + "--config", + str(target / "request.yaml"), + "--output-dir", + str(target), + "--trace", + str(trace), + "--validation-output-dir", + str(tmp_path / "validation"), + ] + ) + == 0 + ) + assert json.loads((tmp_path / "validation/validation.json").read_text())["status"] == "covered" + assert "Review the new profile" in capsys.readouterr().out + + +@pytest.mark.parametrize( + "change,diagnostic", + [ + ("missing_worker", "worker rank evidence is incomplete"), + ("attempt", "different plan, attempt"), + ("phase_graph", "incompatible resolved launch settings"), + ("layout", "incompatible runtime cache layouts"), + ("plan", "SHA-256"), + ("partial", "complete formal publication"), + ("formal", "commit record"), + ("formal_row", "differ from the verified native collection"), + ], +) +def test_finalize_refuses_unverified_or_incompatible_collection(tmp_path, change, diagnostic): + request, root = build_completed_collection(tmp_path) + if change in {"missing_worker", "attempt", "phase_graph", "layout"}: + worker_path = next(root.glob("fpm-artifacts/*/cells/*/raw/pod-0/fpm-memory-worker*.json")) + worker = json.loads(worker_path.read_text()) + if change == "missing_worker": + worker_path.unlink() + elif change == "attempt": + worker["collector_provenance"]["attempt_id"] = "stale-attempt" + _write(worker_path, worker) + elif change == "phase_graph": + worker["resolved_config"]["compilation_config"]["cudagraph_capture_sizes"] = [1, 2] + worker["resolved_config"]["compilation_config"]["max_cudagraph_capture_size"] = 2 + _write(worker_path, worker) + scheduler_path = worker_path.with_name("fpm-memory-scheduler-dp0.json") + scheduler = json.loads(scheduler_path.read_text()) + scheduler["resolved_config"]["compilation_config"] = worker["resolved_config"]["compilation_config"] + _write(scheduler_path, scheduler) + else: + worker["cache"]["groups"][0]["block_size_tokens"] = 8 + _write(worker_path, worker) + scheduler_path = worker_path.with_name("fpm-memory-scheduler-dp0.json") + scheduler = json.loads(scheduler_path.read_text()) + scheduler["cache"]["groups"][0]["block_size_tokens"] = 8 + _write(scheduler_path, scheduler) + elif change == "plan": + path = next(root.glob("fpm-artifacts/*/collection-plan.json")) + payload = json.loads(path.read_text()) + payload["options"]["max_num_seqs"] = 9 + _write(path, payload) + elif change == "partial": + path = root / "fpm-checkpoint/fpm_forward.json" + payload = json.loads(path.read_text()) + payload["database"]["missing_cells"] = ["unfinished"] + _write(path, payload) + elif change == "formal": + path = next(root.glob("systems/data/**/fpm_forward_perf.parquet")) + path.write_bytes(path.read_bytes() + b"changed") + else: + import pyarrow as pa + import pyarrow.parquet as pq + + path = next(root.glob("systems/data/**/fpm_forward_perf.parquet")) + rows = pq.read_table(path).to_pylist() + rows[0]["latency_ms"] *= 2 + pq.write_table(pa.Table.from_pylist(rows), path) + metadata = path.with_name("fpm_forward_perf.metadata.json") + payload = json.loads(metadata.read_text()) + payload["parquet_sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(metadata, payload) + before = _snapshot(root) + with pytest.raises(ValueError, match=diagnostic): + finalize(request, root, tmp_path / "resolved") + assert _snapshot(root) == before + assert not (tmp_path / "resolved").exists() + + +def test_finalized_data_changes_fail_normal_saved_plan_validation(tmp_path): + request, root = build_completed_collection(tmp_path) + target = tmp_path / "resolved" + finalize(request, root, target) + resolved = SupportRequest.from_yaml(target / "request.yaml") + data = next(target.glob("systems/data/**/fpm_forward_perf.parquet")) + data.write_bytes(data.read_bytes() + b"tampered") + with pytest.raises(ValueError, match="finalized FPM data changed"): + check_plan(resolved, target) + + +@pytest.mark.parametrize("target_kind", ["original", "inside", "ancestor", "existing"]) +def test_finalize_requires_new_separate_output(tmp_path, target_kind): + request, root = build_completed_collection(tmp_path) + existing = tmp_path / "existing" + existing.mkdir() + (existing / "keep").write_text("untouched") + target = {"original": root, "inside": root / "new", "ancestor": tmp_path, "existing": existing}[target_kind] + before = _snapshot(tmp_path) + with pytest.raises(ValueError, match="separate|must be new"): + finalize(request, root, target) + assert _snapshot(tmp_path) == before + + +def test_pending_profile_collection_plan_checkpoint_and_replay_gate(tmp_path, capsys): + request, root = build_completed_collection(tmp_path) + check_plan(request, root) + plan = json.loads((root / "support-plan.json").read_text()) + assert plan["resources"] == {"memory_source": "pending", "simulation_ready": False} + assert "Synchronous" in plan["fpm"]["scheduling_policy"] + assert "finalize" in json.loads((root / "commands.json").read_text()) + _write(tmp_path / "trace.json", []) + with pytest.raises(SystemExit) as error: + cli.main( + [ + "onboard", + "validate-fpm", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--trace", + str(tmp_path / "trace.json"), + "--validation-output-dir", + str(tmp_path / "validation"), + ] + ) + assert error.value.code == 2 + assert "pending" in capsys.readouterr().err + assert not (tmp_path / "validation").exists() diff --git a/python/aisimulate/tests/unit/test_onboard_finalization_quality.py b/python/aisimulate/tests/unit/test_onboard_finalization_quality.py new file mode 100644 index 000000000..2562c7b00 --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_finalization_quality.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Finalization exports rechecked quality, separately from serving accuracy.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from collector.fpm_forward import cli as collector_cli +from collector.fpm_forward import planner, repeatability +from collector.fpm_forward.config import FPMCollectionOptions +from collector.fpm_forward.database import aggregate_cell, write_formal_database + +from aisimulate import main as cli +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support import finalization, validation_workflow +from aisimulate.support.plan import check_plan, create_plan +from aisimulate.support.schema import SupportRequest + +from .collector.test_fpm_repeatability import _runtime_config, _write_campaign +from .test_onboard_finalization import _memory, _prepare_collection, _write, build_completed_collection + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def assessed_collection(tmp_path, monkeypatch): + monkeypatch.setattr(planner, "_git_revision", lambda: "synthetic-quality-test-revision") + original, _ = _prepare_collection(tmp_path) + payload = original.model_dump(mode="json") + payload["search"]["context_length"] = 2048 + payload["fpm_profile"]["context_length"] = 2048 + payload["fpm_profile"]["deployments"][0]["resources"]["max_batch_size"] = 32 + original = SupportRequest.model_validate(payload) + root = tmp_path / "quality-source" + saved = create_plan(original, root) + args = collector_cli._parser().parse_args(saved["fpm"]["plan_command"][3:]) + plan = planner.build_collection_plan( + backend="vllm", + model_path=original.identity.model, + system=original.identity.gpu, + selected_ops=set(), + options=FPMCollectionOptions.from_args(args), + model_architecture=original.fpm_profile.architecture, + model_config_path=str(tmp_path / "config.json"), + fpm_profile=original.fpm_profile, + ) + campaign = root / "fpm-artifacts" / plan.sha256[:16] + checkpoint = root / "fpm-checkpoint/fpm_forward.json" + points = { + phase: [ + { + "batch_size": batch, + "total_prefill_tokens": batch if phase == "prefill" else 0, + "total_kv_read_tokens": batch * context, + } + for batch in (1, 2, 4, 8, 16, 32) + for context in (16, 64, 256, 1024) + ] + for phase in ("prefill", "decode") + } + _write_campaign(campaign, checkpoint, plan, attempt_id="source-quality", native_points=points) + rows = [] + for cell in plan.cells: + directory = campaign / "cells" / cell.cell_id + _write(directory / "cell.json", cell.to_dict()) + raw = directory / "raw/pod" + provenance = json.loads((raw / "collector-provenance.json").read_text()) + worker, scheduler = _memory( + provenance, cell.workload_kind, original.identity.model, original.identity.model_revision + ) + for record in (worker, scheduler): + record["resolved_config"] = { + **_runtime_config(plan, cell), + "offload_config": record["resolved_config"]["offload_config"], + } + _write(raw / "fpm-memory-worker-dp0-tp0-pp0.json", worker) + _write(raw / "fpm-memory-scheduler-dp0.json", scheduler) + rows.extend(aggregate_cell(plan, cell, directory, expected_attempt_id="source-quality")) + parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=root / "systems/data") + assert not skipped + saved = json.loads(checkpoint.read_text()) + saved["database"] = { + "status": "passed", + "missing_cells": [], + "skipped_first_publisher_wins": [], + "plan_cells": 2, + "published_cells": 2, + "parquet": str(parquet), + "metadata": str(metadata), + } + _write(checkpoint, saved) + case = { + "request": original, + "root": root, + "plan": plan, + "repeat_factor": 1.0, + "execution_evidence": True, + "calls": [], + } + + def synthesize_repeat(sample_plan, **kwargs): + case["calls"].append(kwargs) + _write_campaign( + Path(kwargs["artifact_root"]) / sample_plan.sha256[:16], + Path(kwargs["checkpoint_dir"]) / "fpm_forward.json", + sample_plan, + attempt_id=f"repeat-{len(case['calls'])}", + factor=case["repeat_factor"], + generator_overrides=kwargs["generator_overrides"], + native_points=points, + execution_evidence=case["execution_evidence"], + ) + return 0 + + monkeypatch.setattr(repeatability, "run_collection", synthesize_repeat) + case["validation_args"] = [ + "onboard", + "validate-collection", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--validation-output-dir", + str(tmp_path / "quality"), + ] + case["report"] = tmp_path / "quality/collection-validation.json" + return case + + +def _assess(case, status): + if status == "not_assessed": + return + case["repeat_factor"] = 1.5 if status == "failed" else 1.0 + assert cli.main([*case["validation_args"], *([] if status == "incomplete" else ["--execute"])]) == ( + 1 if status == "failed" else 0 + ) + assert json.loads(case["report"].read_text())["status"] == status + + +def _finalize_args(case, target, *, include_report=True): + return [ + "onboard", + "finalize", + "--config", + str(case["root"] / "request.yaml"), + "--output-dir", + str(case["root"]), + "--resolved-output-dir", + str(target), + *(["--collection-report", str(case["report"])] if include_report else []), + ] + + +@pytest.mark.parametrize("mutation", ["stale_report", "forged_status"]) +def test_checkpoint_acceptance_verifies_quality_without_prior_runtime_probe(assessed_collection, tmp_path, mutation): + case = assessed_collection + _assess(case, "failed" if mutation == "forged_status" else "passed") + target = tmp_path / "resolved" + finalization.finalize(case["request"], case["root"], target, collection_report=case["report"]) + request = SupportRequest.from_yaml(target / "request.yaml") + payload = request.model_dump(mode="json") + manifest = finalization.finalization_manifest(request) + assert "runtime_probe" not in manifest + if mutation == "stale_report": + case["report"].write_text(case["report"].read_text() + "\n") + else: + manifest["collection_quality"]["status"] = "passed" + payload["fpm_profile"]["deployments"][0]["resources"]["runtime_memory"]["provenance"] = json.dumps(manifest) + profile = json.loads(payload["fpm_profile"]["provenance"]) + profile["collection_quality"] = manifest["collection_quality"] + payload["fpm_profile"]["provenance"] = json.dumps(profile) + patch = tmp_path / "checkpoint-update.json" + _write(patch, {"configurations": {"tp1": {"draft_request": payload}}}) + checkpoint = tmp_path / "checkpoint.json" + with pytest.raises(SystemExit, match="2"): + cli.main( + ["onboard", "checkpoint", "--file", str(checkpoint), "--update", str(patch), "--accept-profile", "tp1"] + ) + assert not checkpoint.exists() + + +@pytest.mark.parametrize("status", ["passed", "failed", "incomplete", "not_assessed", "historical"]) +def test_checkpoint_acceptance_and_resume_preserve_verified_exploratory_quality( + assessed_collection, tmp_path, capsys, status +): + case = assessed_collection + _assess(case, "not_assessed" if status == "historical" else status) + target = tmp_path / "resolved" + finalization.finalize( + case["request"], + case["root"], + target, + collection_report=case["report"] if status not in {"not_assessed", "historical"} else None, + ) + request = SupportRequest.from_yaml(target / "request.yaml") + payload = request.model_dump(mode="json") + manifest = finalization.finalization_manifest(request) + assert "runtime_probe" not in manifest + if status == "historical": + manifest.pop("collection_quality") + payload["fpm_profile"]["deployments"][0]["resources"]["runtime_memory"]["provenance"] = json.dumps(manifest) + profile = json.loads(payload["fpm_profile"]["provenance"]) + profile.pop("collection_quality") + payload["fpm_profile"]["provenance"] = json.dumps(profile) + patch = tmp_path / "checkpoint-update.json" + _write(patch, {"configurations": {"tp1": {"draft_request": payload}}}) + checkpoint = tmp_path / "checkpoint.json" + capsys.readouterr() + assert ( + cli.main( + ["onboard", "checkpoint", "--file", str(checkpoint), "--update", str(patch), "--accept-profile", "tp1"] + ) + == 0 + ) + capsys.readouterr() + assert cli.main(["onboard", "resume", "--checkpoint", str(checkpoint)]) == 0 + resumed = json.loads(capsys.readouterr().out) + assert resumed["configurations"]["tp1"]["profile_accepted"] + assert resumed["integrity_issues"] == [] + + +@pytest.mark.parametrize("status", ["not_assessed", "incomplete", "failed", "passed"]) +def test_public_finalization_exports_quality_without_claiming_serving_accuracy(assessed_collection, tmp_path, status): + case = assessed_collection + _assess(case, status) + source = {p: p.read_bytes() for p in case["root"].rglob("*") if p.is_file()} + target = tmp_path / "resolved" + assert cli.main(_finalize_args(case, target, include_report=status != "not_assessed")) == 0 + assert source == {p: p.read_bytes() for p in source} + manifest = json.loads((target / "finalization.json").read_text()) + quality = manifest["collection_quality"] + assert quality["status"] == status + assert quality["accuracy"] == "not_assessed" + if status == "not_assessed": + assert quality == {"status": "not_assessed", "accuracy": "not_assessed"} + else: + assert quality["report"] == validation_workflow._identity(case["report"]) + assert quality["policy"] == validation_workflow._identity(case["report"].parent / "policy.json") + assert set(quality["gates"]) == {"point_validity", "execution", "repeatability", "interpolation"} + assert "cell_observations" not in quality + request = SupportRequest.from_yaml(target / "request.yaml") + assert json.loads(request.fpm_profile.provenance)["collection_quality"] == quality + check_plan(request, target) + for config in ( + CorePredictionConfig.from_yaml(target / "predict/pilot.yaml"), + CoreRecommendationConfig.from_yaml(target / "recommend/pilot.yaml"), + ): + assert json.loads(config.engine.fpm_profile.provenance)["collection_quality"] == quality + assert config.engine.fpm_profile == request.fpm_profile + + +@pytest.mark.parametrize("changed", ["status", "policy", "sample", "request", "directory"]) +def test_finalization_rejects_stale_or_unrelated_quality(assessed_collection, tmp_path, changed): + case = assessed_collection + _assess(case, "passed") + report = json.loads(case["report"].read_text()) + if changed == "status": + report["status"] = "failed" + elif changed == "policy": + path = case["report"].parent / "policy.json" + path.write_text(path.read_text() + "\n") + elif changed == "sample": + path = next((case["report"].parent / "repeatability").rglob("benchmark-dp0.json")) + path.write_text(path.read_text() + "\n") + elif changed == "request": + payload = case["request"].model_dump(mode="json") + payload["workload"]["request_count"] += 1 + alternate = tmp_path / "different-request.yaml" + alternate.write_text(json.dumps(payload)) + report["inputs"]["request"] = validation_workflow._identity(alternate) + else: + report["inputs"]["collection_directory"] = str(tmp_path / "another-collection") + _write(case["report"], report) + target = tmp_path / "rejected" + with pytest.raises(SystemExit, match="2"): + cli.main(_finalize_args(case, target)) + assert not target.exists() + + +@pytest.mark.parametrize("changed", ["report", "policy", "sample", "source"]) +def test_saved_plan_rechecks_bound_quality_evidence(assessed_collection, tmp_path, changed): + case = assessed_collection + _assess(case, "passed") + target = tmp_path / "resolved" + finalization.finalize(case["request"], case["root"], target, collection_report=case["report"]) + request = SupportRequest.from_yaml(target / "request.yaml") + path = { + "report": case["report"], + "policy": case["report"].parent / "policy.json", + "source": case["root"] / "request.yaml", + }.get(changed) + if changed == "sample": + path = next((case["report"].parent / "repeatability").rglob("benchmark-dp0.json")) + path.write_text(path.read_text() + "\n") + with pytest.raises(ValueError, match="changed"): + check_plan(request, target) + + +def test_finalization_uses_the_reports_frozen_nondefault_policy(assessed_collection, tmp_path): + case = assessed_collection + policy = tmp_path / "reviewed-policy.json" + _write(policy, {"repeatability": {"samples": 2, "max_cv": 0.025}}) + assert cli.main([*case["validation_args"], "--execute", "--policy", str(policy)]) == 0 + target = tmp_path / "resolved" + assert cli.main(_finalize_args(case, target)) == 0 + quality = json.loads((target / "finalization.json").read_text())["collection_quality"] + assert quality["status"] == "passed" + assert quality["policy"] == validation_workflow._identity(case["report"].parent / "policy.json") + assert len(case["calls"]) == 4 + + +@pytest.mark.parametrize("target_status", ["passed", "failed", "incomplete"]) +def test_missing_report_cannot_be_relabelled_as_assessed(tmp_path, target_status): + original, root = build_completed_collection(tmp_path) + target = tmp_path / "resolved" + finalization.finalize(original, root, target) + payload = SupportRequest.from_yaml(target / "request.yaml").model_dump(mode="json") + resource = payload["fpm_profile"]["deployments"][0]["resources"] + manifest = json.loads(resource["runtime_memory"]["provenance"]) + manifest["collection_quality"]["status"] = target_status + resource["runtime_memory"]["provenance"] = json.dumps(manifest) + profile = json.loads(payload["fpm_profile"]["provenance"]) + profile["collection_quality"] = manifest["collection_quality"] + payload["fpm_profile"]["provenance"] = json.dumps(profile) + altered = SupportRequest.model_validate(payload) + with pytest.raises(ValueError, match="collection quality differs"): + finalization.verify_finalized_data(altered, target) + + +def test_historical_finalization_remains_unassessed_without_rewriting_its_manifest(tmp_path): + original, root = build_completed_collection(tmp_path) + observations, manifest, formal, _ = finalization._verify_collection(original, root) + assert "collection_quality" not in manifest + payload = original.model_dump(mode="json") + payload["fpm_profile"]["deployments"][0]["resources"] = finalization._merge_resources(observations, manifest) + historical = SupportRequest.model_validate(payload) + target = tmp_path / "historical" + create_plan(historical, target) + for relative, content in formal.items(): + path = target / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + before = (target / "request.yaml").read_bytes() + check_plan(historical, target) + assert finalization.verify_finalized_quality(finalization.finalization_manifest(historical)) is None + assert (target / "request.yaml").read_bytes() == before + + +def test_quality_report_stays_bound_to_original_collection_with_accepted_memory_revision(tmp_path, capsys): + from aisimulate.support.plan import request_id + from aisimulate.support.runtime import verify_runtime_profile + + from .test_onboard_runtime import _reviewed_capacity_revision + + checkpoint, original, _, root, memory_config = _reviewed_capacity_revision(tmp_path, capsys) + revised = SupportRequest.from_yaml(memory_config) + assessment = tmp_path / "quality" + assert ( + cli.main( + [ + "onboard", + "validate-collection", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--validation-output-dir", + str(assessment), + ] + ) + == 0 + ) + report = assessment / "collection-validation.json" + assert json.loads(report.read_text())["status"] == "incomplete" + before = {p: p.read_bytes() for p in root.rglob("*") if p.is_file()} + accepted = checkpoint.read_bytes() + target = tmp_path / "resolved" + finalization.finalize(original, root, target, memory_config=memory_config, collection_report=report) + resolved = SupportRequest.from_yaml(target / "request.yaml") + manifest = finalization.finalization_manifest(resolved) + assert manifest["collection_quality"]["status"] == "incomplete" + assert manifest["source_request_id"] == request_id(original) + assert manifest["memory_revision"]["request_id"] == request_id(revised) + assert ( + resolved.profile_deployment().resources.runtime_memory.kv_cache_bytes + == revised.profile_deployment().resources.runtime_memory.kv_cache_bytes + ) + check_plan(resolved, target) + verify_runtime_profile(resolved) + assert before == {p: p.read_bytes() for p in before} + assert checkpoint.read_bytes() == accepted diff --git a/python/aisimulate/tests/unit/test_onboard_model_config.py b/python/aisimulate/tests/unit/test_onboard_model_config.py new file mode 100644 index 000000000..662c2aadb --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_model_config.py @@ -0,0 +1,1695 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public local-config onboarding, including saved-request replay.""" + +from __future__ import annotations + +import builtins +import hashlib +import json +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml + +import aisimulate.main as cli +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.fpm import fpm_cli_args +from aisimulate.support.plan import request_id +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + +# Synthetic geometry and explicit illustrative bounds, not checkpoint measurements. +_CONFIG = { + "_name_or_path": "example/local-decoder", + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "max_position_embeddings": 32768, + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 1024, + "torch_dtype": "bfloat16", + "tie_word_embeddings": False, +} +_OVERRIDES = { + "weights_bytes": 2 * 1024**3, + "activations_bytes": 32 * 1024**2, + "runtime_overhead_bytes": 16 * 1024**2, + "comm_overhead_bytes": 16 * 1024**2, + "kv_bytes_per_token": 4096, + "cache_layout": "linear", + "max_num_tokens": 4096, + "max_batch_size": 8, + "gemm_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "bfloat16", + "provenance": "Synthetic CLI test bounds; not measured.", +} +_IDENTITY = { + "model_revision": "checkpoint-revision-123", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", +} +_PILOT = { + "tensor_parallel": 1, + "input_tokens": 1024, + "output_tokens": 128, + "concurrency": 1, + "ttft_ms": 1000, + "tpot_ms": 100, +} + + +def test_validation_changes_preserve_profile_and_collection_command(tmp_path, capsys): + source, resources = _files(tmp_path) + requests = [] + for name, validation in ( + ("small", {}), + ( + "large", + {"input_tokens": 4096, "output_tokens": 256, "concurrency": 16, "request_count": 32, "ttft_ms": 5000}, + ), + ): + output = tmp_path / f"{name}.yaml" + assert cli.main(_args(output, source, resources, **validation)) == 0 + requests.append(SupportRequest.from_yaml(output)) + first, second = requests + assert first.workload != second.workload + assert first.fpm_profile == second.fpm_profile + assert request_id(first) == request_id(second) + assert fpm_cli_args(first, output_dir=tmp_path / "plan", plan_only=True) == fpm_cli_args( + second, output_dir=tmp_path / "plan", plan_only=True + ) + assert first.scheduler_limits() == {"max_batched_tokens": 4096, "max_sequences": 8} + + +@pytest.mark.parametrize("role,graph", [("prefill", "PIECEWISE"), ("decode", "FULL_DECODE_ONLY")]) +def test_role_init_exports_independent_collection_config_without_claiming_replay(tmp_path, role, graph): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + args = _args(output, source, resources, worker_type=role, cudagraph_mode=graph) + args += ["--cudagraph-capture-sizes", "1", "2", "4", "8", "--max-cudagraph-capture-size", "8"] + assert cli.main(args) == 0 + request = SupportRequest.from_yaml(output) + assert request.worker_type == request.profile_deployment().worker_type == role + assert request.collection_phases == (role,) + assert request.collection.cudagraph_capture_sizes == [1, 2, 4, 8] + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + saved = SupportRequest.from_yaml(root / "request.yaml") + assert saved.model_dump(mode="json") == request.model_dump(mode="json") + artifact = yaml.safe_load((root / "worker.yaml").read_text()) + assert artifact["worker_type"] == role + assert artifact["serving_mode"] == artifact["engine"]["mode"] == "disaggregated" + assert set(artifact["engine"]["workers"]) == {role} + worker = artifact["engine"]["workers"][role] + from aisimulate.config.engine import WorkerPredictionConfig + + parsed_worker = WorkerPredictionConfig.model_validate(worker) + assert parsed_worker.timing.systems_paths == [str(root / "systems")] + assert worker["timing"]["estimation_mode"] == "fpm_interpolation" + assert worker["timing"]["fallback_policy"] == "deny" + assert artifact["engine"]["systems_paths"] == [str(root / "systems")] + assert artifact["engine"]["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert artifact["collection"]["cudagraph_mode"] == graph + commands = json.loads((root / "commands.json").read_text()) + assert commands["predict"] == commands["recommend"] == [] + assert not (root / "predict/pilot.yaml").exists() + command = commands["fpm_plan_local"] + assert command[command.index("--fpm-worker-type") + 1] == role + assert command[command.index("--fpm-cudagraph-mode") + 1] == graph + plan = json.loads((root / "support-plan.json").read_text()) + assert plan["fpm"]["collection_phases"] == [role] + assert plan["replay"]["status"] == "not_assessed" + from aisimulate.support.validation import _prediction_config + + with pytest.raises(ValueError, match="aggregated replay only"): + _prediction_config(request, root, tmp_path / "trace.json") + + +def test_explicit_aggregated_role_retains_executable_prediction_and_recommendation(tmp_path): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + assert ( + cli.main(_args(output, source, resources, worker_type="aggregated", cudagraph_mode="FULL_AND_PIECEWISE")) == 0 + ) + request = SupportRequest.from_yaml(output) + assert request.collection_phases == ("prefill", "decode") + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + prediction = CorePredictionConfig.model_validate(yaml.safe_load((root / "predict/pilot.yaml").read_text())) + recommendation = CoreRecommendationConfig.model_validate( + yaml.safe_load((root / "recommend/pilot.yaml").read_text()) + ) + assert prediction.engine.mode == recommendation.engine.mode == "aggregated" + assert prediction.engine.fpm_profile.deployments[0].worker_type == "aggregated" + assert not (root / "worker.yaml").exists() + + +def test_explicit_role_does_not_relabel_a_legacy_shared_profile(tmp_path): + source, resources = _files(tmp_path) + output = tmp_path / "legacy.yaml" + assert cli.main(_args(output, source, resources)) == 0 + legacy = SupportRequest.from_yaml(output) + payload = legacy.model_dump(mode="json") + assert "worker_type" not in payload + assert "worker_type" not in payload["fpm_profile"]["deployments"][0] + with pytest.raises(ValueError, match="do not relabel a legacy shared profile"): + SupportRequest.model_validate({**payload, "worker_type": "decode"}) + assert SupportRequest.model_validate(payload).model_dump(mode="json") == payload + + +def test_guided_role_and_graph_review_edits_are_saved(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + _terminal( + monkeypatch, + [ + "edit", + "worker_type", + "decode", + "edit", + "cudagraph_mode", + "FULL_DECODE_ONLY", + "edit", + "cudagraph_capture_sizes", + "1 2 4 8", + "accept", + ], + ) + assert ( + cli.main( + _args(output, source, resources, worker_type="prefill", cudagraph_mode="PIECEWISE") + ["--interactive"] + ) + == 0 + ) + request = SupportRequest.from_yaml(output) + assert request.worker_type == request.profile_deployment().worker_type == "decode" + assert request.collection.cudagraph_mode == "FULL_DECODE_ONLY" + assert request.collection.cudagraph_capture_sizes == [1, 2, 4, 8] + + +def test_guided_graph_overrides_can_be_cleared_before_changing_to_none(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + _terminal(monkeypatch, ["edit", "max_cudagraph_capture_size", "clear", "edit", "cudagraph_mode", "NONE", "accept"]) + args = _args(output, source, resources, worker_type="decode", cudagraph_mode="FULL", max_cudagraph_capture_size=8) + assert cli.main([*args, "--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + assert request.collection.cudagraph_mode == "NONE" + assert "max_cudagraph_capture_size" not in request.collection.model_dump(mode="json") + + +def test_multiple_roles_on_same_topology_keep_distinct_requests_and_collection_limits(tmp_path): + source, _ = _files(tmp_path) + selections = tmp_path / "roles.yaml" + selections.write_text( + yaml.safe_dump( + [ + { + "worker_type": "prefill", + "tensor_parallel": 1, + "resource_overrides": _OVERRIDES, + "collection": {"cudagraph_mode": "PIECEWISE", "max_num_tokens": 4096, "max_batch_size": 8}, + }, + { + "worker_type": "decode", + "tensor_parallel": 1, + "resource_overrides": {**_OVERRIDES, "max_num_tokens": 2048, "max_batch_size": 4}, + "collection": { + "cudagraph_mode": "FULL_DECODE_ONLY", + "max_num_tokens": 2048, + "max_batch_size": 4, + "gpu_memory_utilization": 0.92, + }, + }, + ] + ) + ) + root = tmp_path / "onboarding" + args = _args(root, source, tensor_parallel=None) + args[args.index("--output")] = "--output-dir" + assert cli.main([*args, "--parallel-configs", str(selections)]) == 0 + index = json.loads((root / "onboarding.json").read_text()) + assert [entry["worker_type"] for entry in index["configurations"]] == ["prefill", "decode"] + requests = [SupportRequest.from_yaml(entry["request"]) for entry in index["configurations"]] + assert requests[0].parallelism() == requests[1].parallelism() + assert request_id(requests[0]) != request_id(requests[1]) + assert requests[0].scheduler_limits() == {"max_batched_tokens": 4096, "max_sequences": 8} + assert requests[1].scheduler_limits() == {"max_batched_tokens": 2048, "max_sequences": 4} + assert requests[1].collection.memory_fraction == 0.92 + assert [Path(entry["request"]).parent.name.split("-")[0] for entry in index["configurations"]] == [ + "prefill", + "decode", + ] + + +@pytest.mark.parametrize("interactive", [False, True]) +@pytest.mark.parametrize( + "shared", + [ + {"cudagraph_mode": "FULL_AND_PIECEWISE"}, + {"cudagraph_capture_sizes": [1, 2, 4]}, + {"max_cudagraph_capture_size": 4}, + { + "cudagraph_mode": "FULL_AND_PIECEWISE", + "cudagraph_capture_sizes": [1, 2, 4], + "max_cudagraph_capture_size": 4, + }, + ], +) +def test_parallel_roles_accept_shared_graph_settings_without_a_global_role(tmp_path, monkeypatch, shared, interactive): + source, _ = _files(tmp_path) + selections = tmp_path / "roles.yaml" + selections.write_text( + yaml.safe_dump( + [ + {"worker_type": "prefill", "tensor_parallel": 1, "resource_overrides": _OVERRIDES}, + { + "worker_type": "decode", + "tensor_parallel": 1, + "resource_overrides": _OVERRIDES, + "collection": {"cudagraph_mode": "FULL_DECODE_ONLY"}, + }, + ] + ) + ) + root = tmp_path / "onboarding" + args = _args(root, source, tensor_parallel=None) + args[args.index("--output")] = "--output-dir" + for name, value in shared.items(): + args += ["--" + name.replace("_", "-"), *map(str, value if isinstance(value, list) else [value])] + if interactive: + _terminal(monkeypatch, ["accept", "accept"]) + args.append("--interactive") + assert cli.main([*args, "--parallel-configs", str(selections)]) == 0 + index = json.loads((root / "onboarding.json").read_text()) + requests = [SupportRequest.from_yaml(entry["request"]) for entry in index["configurations"]] + assert [request.worker_type for request in requests] == ["prefill", "decode"] + assert [request.profile_deployment().worker_type for request in requests] == ["prefill", "decode"] + for name, value in shared.items(): + assert getattr(requests[0].collection, name) == value + assert getattr(requests[1].collection, name) == ("FULL_DECODE_ONLY" if name == "cudagraph_mode" else value) + assert requests[1].collection.cudagraph_mode == "FULL_DECODE_ONLY" + + +@pytest.mark.parametrize("roleless_position", [0, 1]) +@pytest.mark.parametrize( + "graph_flags", + [ + ["--cudagraph-mode", "FULL_AND_PIECEWISE"], + ["--cudagraph-capture-sizes", "1", "2", "4"], + ["--max-cudagraph-capture-size", "4"], + ], +) +def test_shared_graph_settings_reject_any_roleless_parallel_entry(tmp_path, capsys, roleless_position, graph_flags): + source, _ = _files(tmp_path) + entries = [ + {"worker_type": role, "tensor_parallel": 1, "resource_overrides": _OVERRIDES} for role in ("prefill", "decode") + ] + entries[roleless_position].pop("worker_type") + selections = tmp_path / "roles.yaml" + selections.write_text(yaml.safe_dump(entries)) + root = tmp_path / "onboarding" + args = _args(root, source, tensor_parallel=None) + args[args.index("--output")] = "--output-dir" + with pytest.raises(SystemExit) as error: + cli.main([*args, "--parallel-configs", str(selections), *graph_flags]) + assert error.value.code == 2 + assert "generic CUDA graph settings require an explicit worker_type" in capsys.readouterr().err + assert not root.exists() + + +@pytest.mark.parametrize( + "graph_flags,collection,error_message", + [ + ( + ["--cudagraph-mode", "NONE", "--max-cudagraph-capture-size", "4"], + {}, + "cudagraph_mode=NONE rejects CUDA graph capture settings", + ), + ( + ["--cudagraph-capture-sizes", "4", "2"], + {}, + "cudagraph_capture_sizes must be a nonempty, increasing list of distinct sizes", + ), + ( + ["--max-cudagraph-capture-size", "4"], + {"prefill_cudagraph_policy": "explicit"}, + "generic CUDA graph capture settings require prefill_cudagraph_policy=runtime", + ), + ], +) +def test_parallel_entries_validate_merged_graph_settings(tmp_path, capsys, graph_flags, collection, error_message): + source, resources = _files(tmp_path) + selections = tmp_path / "roles.yaml" + selections.write_text(yaml.safe_dump([{"worker_type": "decode", "tensor_parallel": 1, "collection": collection}])) + root = tmp_path / "onboarding" + args = _args(root, source, resources, tensor_parallel=None) + args[args.index("--output")] = "--output-dir" + with pytest.raises(SystemExit) as error: + cli.main([*args, "--parallel-configs", str(selections), *graph_flags]) + assert error.value.code == 2 + assert error_message in capsys.readouterr().err + assert not root.exists() + + +def test_legacy_parallel_entry_preserves_single_request_serialization_and_hash(tmp_path): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + assert cli.main(_args(output, source, resources)) == 0 + original = SupportRequest.from_yaml(output) + selections = tmp_path / "parallel.yaml" + selections.write_text(yaml.safe_dump([{"tensor_parallel": 1, "resource_overrides": _OVERRIDES}])) + root = tmp_path / "onboarding" + args = _args(root, source, tensor_parallel=None) + args[args.index("--output")] = "--output-dir" + assert cli.main([*args, "--parallel-configs", str(selections)]) == 0 + index = json.loads((root / "onboarding.json").read_text()) + request = SupportRequest.from_yaml(index["configurations"][0]["request"]) + assert request.model_dump(mode="json") == original.model_dump(mode="json") + assert request_id(request) == request_id(original) + payload = request.model_dump(mode="json") + assert "worker_type" not in payload + assert "worker_type" not in payload["fpm_profile"]["deployments"][0] + assert ( + not {"cudagraph_mode", "cudagraph_capture_sizes", "max_cudagraph_capture_size"} & payload["collection"].keys() + ) + + +@pytest.mark.parametrize( + "changes", + [ + {"cudagraph_mode": "PIECEWISE"}, + {"worker_type": "decode", "cudagraph_capture_sizes": [4, 2]}, + {"worker_type": "decode", "cudagraph_capture_sizes": [2, 2]}, + {"worker_type": "decode", "cudagraph_capture_sizes": [1, 2], "max_cudagraph_capture_size": 4}, + {"worker_type": "decode", "cudagraph_capture_sizes": [1, 2], "prefill_cudagraph_policy": "explicit"}, + {"worker_type": "decode", "max_cudagraph_capture_size": 2, "cudagraph_mode": "NONE"}, + ], +) +def test_invalid_role_graph_settings_fail_before_saving(tmp_path, changes): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + sizes = changes.get("cudagraph_capture_sizes") + args = _args(output, source, resources, **{k: v for k, v in changes.items() if k != "cudagraph_capture_sizes"}) + if sizes is not None: + args += ["--cudagraph-capture-sizes", *map(str, sizes)] + with pytest.raises(SystemExit): + cli.main(args) + assert not output.exists() + + +@pytest.mark.parametrize("model_context,expected", [(32768, 32768), (262144, 256000)]) +def test_model_context_and_reference_cap_define_initial_runtime_context(tmp_path, model_context, expected): + source, resources = _files(tmp_path, config={**_CONFIG, "max_position_embeddings": model_context}) + output = tmp_path / "request.yaml" + assert cli.main(_args(output, source, resources)) == 0 + request = SupportRequest.from_yaml(output) + assert request.search.context_length == expected + assert request.fpm_profile.context_length == model_context + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + assert command[command.index("--fpm-max-model-len") + 1] == str(expected) + assert command[command.index("--fpm-max-prefill-isl") + 1] == "4096" + + +def test_guided_collection_edit_reestimates_resources_before_acceptance(tmp_path, monkeypatch): + source, resources = _files( + tmp_path, + config={**_CONFIG, "hidden_size": 4096, "num_attention_heads": 32, "intermediate_size": 8192}, + overrides={"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"}, + ) + output = tmp_path / "request.yaml" + prompts = _terminal( + monkeypatch, + [ + "edit", + "max_num_tokens", + "4096", + "edit", + "runtime_context_length", + "8192", + "edit", + "max_prefill_cudagraph_size", + "512", + "accept", + ], + ) + assert cli.main(_args(output, source, resources, max_num_tokens=16384) + ["--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + profile = request.profile_deployment() + assert request.search.context_length == 8192 + assert request.collection.max_num_tokens == profile.resources.max_num_tokens == 4096 + assert request.collection.max_prefill_cudagraph_size == 512 + assert profile.resources.memory_source == "pending" + assert json.loads(profile.resources.provenance)["fields"]["activations_bytes"]["value"] == 2 * 4096 * 4096 * 11 + assert request.scheduler_limits() == {"max_batched_tokens": 4096, "max_sequences": 256} + assert not any("Input tokens" in prompt or "Concurrent requests" in prompt for prompt in prompts) + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + for option, expected in ( + ("--fpm-max-model-len", "8192"), + ("--fpm-max-num-batched-tokens", "4096"), + ("--fpm-max-num-seqs", "256"), + ("--fpm-max-prefill-cudagraph-size", "512"), + ): + assert command[command.index(option) + 1] == expected + + +def test_conflicting_resource_and_collection_bounds_fail_before_saving(tmp_path): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + with pytest.raises(SystemExit): + cli.main(_args(output, source, resources, max_num_tokens=8192)) + assert not output.exists() + + +@pytest.mark.parametrize( + "options,policy,limit", + [ + ({}, "runtime", None), + ({"prefill_cudagraph_policy": "runtime"}, "runtime", None), + ({"max_prefill_cudagraph_size": 512}, "explicit", 512), + ({"prefill_cudagraph_policy": "explicit"}, "explicit", 2048), + ({"prefill_cudagraph_policy": "explicit", "max_prefill_cudagraph_size": 1024}, "explicit", 1024), + ], +) +def test_new_runtime_policy_is_saved_and_preserved_in_collection_preview(tmp_path, options, policy, limit): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + assert cli.main(_args(output, source, resources, **options)) == 0 + request = SupportRequest.from_yaml(output) + assert request.collection.prefill_cudagraph_policy == policy + assert request.collection.gpu_memory_utilization == 0.9 + assert request.collection_settings()["max_prefill_cudagraph_size"] == limit + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + command = json.loads((root / "commands.json").read_text())["fpm_plan_local"] + assert command[command.index("--fpm-prefill-cudagraph-policy") + 1] == policy + assert command[command.index("--fpm-gpu-memory-utilization") + 1] == "0.9" + if policy == "runtime": + assert "--fpm-max-prefill-cudagraph-size" not in command + assert "deferred" in request.collection_settings()["sources"]["max_prefill_cudagraph_size"] + + +@pytest.mark.parametrize( + "options", + [ + {"prefill_cudagraph_policy": "runtime", "max_prefill_cudagraph_size": 512}, + {"gpu_memory_utilization": "nan"}, + {"gpu_memory_utilization": "inf"}, + {"gpu_memory_utilization": 0}, + {"gpu_memory_utilization": -0.1}, + {"gpu_memory_utilization": 1.1}, + ], +) +def test_invalid_runtime_settings_leave_request_unchanged(tmp_path, options): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + output.write_text("original request\n") + with pytest.raises(SystemExit): + cli.main(_args(output, source, resources, **options) + ["--overwrite"]) + assert output.read_text() == "original request\n" + + +def test_review_reset_to_runtime_clears_capture_override_and_saves_edited_memory(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + _terminal( + monkeypatch, + [ + "edit", + "prefill_cudagraph_policy", + "invalid", + "runtime", + "edit", + "gpu_memory_utilization", + "nan", + "1.1", + "0.75", + "accept", + ], + ) + assert cli.main(_args(output, source, resources, max_prefill_cudagraph_size=1024) + ["--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + assert request.collection.prefill_cudagraph_policy == "runtime" + assert request.collection.max_prefill_cudagraph_size is None + assert request.collection.gpu_memory_utilization == 0.75 + command = fpm_cli_args(request, output_dir=tmp_path / "collection", plan_only=True) + assert "--fpm-max-prefill-cudagraph-size" not in command + assert command[command.index("--fpm-gpu-memory-utilization") + 1] == "0.75" + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_memory_fraction_survives_plan_compilation_and_replay_validation(tmp_path, grouped): + from aisimulate.compiler import prediction_to_replay_spec + from aisimulate.support.validation import _prediction_config + + overrides = dict(_OVERRIDES) + if grouped: + overrides.pop("kv_bytes_per_token") + overrides.pop("cache_layout") + overrides["cache_block_sizes"] = {"sliding_attention": 16} + source, resources = _files( + tmp_path, config={**_CONFIG, **({"sliding_window": 128} if grouped else {})}, overrides=overrides + ) + output = tmp_path / "request.yaml" + fraction = 0.73 + assert cli.main(_args(output, source, resources, model=str(tmp_path), gpu_memory_utilization=fraction)) == 0 + request = SupportRequest.from_yaml(output) + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + plan = json.loads((root / "support-plan.json").read_text()) + budget = plan["resources"] + assert budget["total_kv_size_bytes"] == int(budget["total_gpu_capacity_bytes"] * fraction) - sum( + budget["memory_breakdown"].values() + ) + prediction = CorePredictionConfig.from_yaml(root / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(root / "recommend/pilot.yaml") + replay = CorePredictionConfig.model_validate(_prediction_config(request, root, tmp_path / "trace.jsonl")) + for config in (prediction, recommendation, replay): + assert config.engine.workers.aggregated.kv_cache.capacity.memory_fraction == fraction + if grouped: + assert not config.engine.workers.aggregated.kv_cache.prefix_caching + # Host RAM controls are independent of the simulated GPU budget. + assert config.execution.resources.available_memory_fraction == 0.9 + for config in (prediction, replay): + spec = prediction_to_replay_spec(config) + args = spec.backend_deployment.agg_engine_args + assert args["gpu_memory_utilization"] == fraction + + +def test_communication_override_is_retained_with_collection_compatibility_guidance(tmp_path, monkeypatch, capsys): + source, resources = _files(tmp_path, overrides={**_OVERRIDES, "comm_quant_mode": "int8"}) + output = tmp_path / "request.yaml" + _terminal(monkeypatch, ["accept"]) + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + deployment = SupportRequest.from_yaml(output).profile_deployment() + assert deployment.comm_quant_mode == "int8" + source = json.loads(deployment.resources.provenance)["fields"]["comm_quant_mode"]["source"] + assert "user override" in source + assert "collector supports half" in source + assert "new collection will reject a mismatched identity" in capsys.readouterr().out + + +@pytest.mark.parametrize( + "route,limits", + [ + (route, limits) + for route in ("identifier", "profile", "config") + for limits in ({"max_num_tokens": 128, "max_batch_size": 256}, {"max_num_tokens": 128}) + ] + + [("profile", {}), ("config", {})], +) +def test_invalid_resolved_scheduler_cannot_be_saved_or_planned(tmp_path, capsys, route, limits): + source, resources = _files(tmp_path, overrides={**_OVERRIDES, "max_num_tokens": 8192, "max_batch_size": 256}) + baseline = tmp_path / "baseline.yaml" + assert cli.main(_args(baseline, source, resources)) == 0 + payload = yaml.safe_load(baseline.read_text()) + profile = payload["fpm_profile"] + output = tmp_path / "invalid.yaml" + if route == "config": + resources.write_text(yaml.safe_dump({**_OVERRIDES, "max_num_tokens": 128, "max_batch_size": 256})) + command = _args(output, source, resources, **limits) + else: + command = _args(output, source, None, model=_CONFIG["_name_or_path"], model_kind="dense", **limits) + command.remove(str(source)) + command.remove("--model-config") + if route == "identifier": + command.extend(["--context-length", "4096"]) + if route == "profile": + if not limits: + profile["deployments"][0]["resources"]["max_num_tokens"] = 128 + profile_path = tmp_path / "profile.json" + profile_path.write_text(json.dumps(profile)) + command.extend(["--fpm-profile", str(profile_path)]) + + with pytest.raises(SystemExit) as error: + cli.main(command) + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_tokens" in message and "max_batch_size" in message + assert not output.exists() + + payload["collection"] = limits + if route == "identifier": + payload.pop("fpm_profile") + elif not limits: + payload["fpm_profile"]["deployments"][0]["resources"]["max_num_tokens"] = 128 + request_path = tmp_path / "invalid-saved.yaml" + request_path.write_text(yaml.safe_dump(payload)) + original = request_path.read_bytes() + plan = tmp_path / "invalid-plan" + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "plan", "--config", str(request_path), "--output-dir", str(plan)]) + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_tokens" in message and "max_batch_size" in message + assert request_path.read_bytes() == original + assert not plan.exists() + + +@pytest.mark.parametrize("suggest_parallel", [False, True]) +def test_config_small_token_budget_uses_reviewed_sequence_bound(tmp_path, capsys, suggest_parallel): + overrides = {**_OVERRIDES, "max_num_tokens": 128} + if suggest_parallel: + # Automatic topology selection requires shared inputs rather than rank-local byte overrides. + overrides = {key: value for key, value in overrides.items() if not key.endswith("_bytes")} + overrides.pop("kv_bytes_per_token") + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "request.yaml" + command = _args(output, source, resources, max_num_tokens=128, tensor_parallel=None if suggest_parallel else 1) + if suggest_parallel: + command.append("--suggest-parallel") + + assert cli.main(command) == 0 + + if suggest_parallel: + assert json.loads(capsys.readouterr().out)["default"] is not None + assert not output.exists() + else: + request = SupportRequest.from_yaml(output) + assert request.scheduler_limits() == {"max_batched_tokens": 128, "max_sequences": 8} + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + command = fpm_cli_args(request, output_dir=root, plan_only=True) + assert command[command.index("--fpm-max-num-batched-tokens") + 1] == "128" + assert command[command.index("--fpm-max-num-seqs") + 1] == "8" + + +def test_guided_config_intake_needs_no_validation_lengths_concurrency_or_slas(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["accept"]) + omitted = dict.fromkeys(("input_tokens", "output_tokens", "concurrency", "ttft_ms", "tpot_ms")) + assert cli.main(_args(output, source, resources, **omitted) + ["--interactive"]) == 0 + assert prompts == ["Review action (accept/edit/cancel): "] + saved = SupportRequest.from_yaml(output) + assert saved.workload.concurrency == 1 + assert saved.scheduler_limits()["max_sequences"] == 8 + + +def test_grouped_review_collects_runtime_blocks_then_edits_pages_before_save(tmp_path, monkeypatch, capsys): + config = {**_CONFIG, "sliding_window": 128, "layer_types": ["full_attention", "sliding_attention"]} + overrides = {key: value for key, value in _OVERRIDES.items() if key not in {"kv_bytes_per_token", "cache_layout"}} + source, resources = _files(tmp_path, config=config, overrides=overrides) + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_bytes(source.read_bytes()) + output = tmp_path / "new" / "request.yaml" + edited_groups = [ + { + "name": "full_attention", + "kind": "attention", + "num_layers": 1, + "block_size_tokens": 16, + "page_size_bytes": 8192, + }, + { + "name": "sliding_attention", + "kind": "attention", + "num_layers": 1, + "block_size_tokens": 8, + "page_size_bytes": 2048, + "sliding_window": 128, + }, + ] + answers = iter( + [ + '{"full_attention": 16}', + '{"sliding_attention": 8}', + "edit", + "cache_groups", + json.dumps(edited_groups), + "accept", + ] + ) + prompts = [] + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def answer(prompt): + assert not output.parent.exists(), "grouped profile was saved before acceptance" + prompts.append(prompt) + return next(answers) + + monkeypatch.setattr(builtins, "input", answer) + assert cli.main(_args(output, source, resources, model=str(checkpoint)) + ["--interactive"]) == 0 + saved = SupportRequest.from_yaml(output) + cache = saved.profile_deployment().resources + assert cache.cache_layout == "grouped" + assert cache.kv_bytes_per_token is None + assert cache.cache_groups[0].page_size_bytes == 8192 + assert cache.cache_groups[1].sliding_window == 128 + provenance = json.loads(cache.provenance)["fields"] + assert provenance["cache_block_sizes"]["value"] == {"full_attention": 16, "sliding_attention": 8} + assert "user override" in provenance["cache_groups"]["source"] + assert "config geometry reference" in provenance["cache_groups"]["source"] + assert sum(prompt.startswith("cache_block_sizes") for prompt in prompts) == 2 + assert "minimum packed tensor estimates" in capsys.readouterr().out + source.unlink() + resources.unlink() + plan = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(plan)]) == 0 + prediction = CorePredictionConfig.from_yaml(plan / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan / "recommend/pilot.yaml") + assert prediction.engine.fpm_profile == recommendation.engine.fpm_profile == saved.fpm_profile + assert not prediction.engine.workers.aggregated.kv_cache.prefix_caching + assert not recommendation.engine.workers.aggregated.kv_cache.prefix_caching + planned = json.loads((plan / "support-plan.json").read_text()) + assert planned["resources"]["total_kv_size_tokens"] is None + assert planned["resources"]["request_peak_cache_bytes"] > 0 + + +def test_grouped_automatic_intake_defers_rank_resources_until_after_topology(tmp_path, monkeypatch): + source, resources = _files( + tmp_path, + config={**_CONFIG, "sliding_window": 128}, + overrides={"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"}, + ) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["1", '{"sliding_attention": 16}', "accept"]) + assert cli.main(_args(output, source, resources, tensor_parallel=None) + ["--interactive"]) == 0 + assert prompts[0].startswith("Choose topology") + assert prompts[1].startswith("cache_block_sizes") + assert SupportRequest.from_yaml(output).profile_deployment().resources.cache_layout == "grouped" + + +def _files(tmp_path, *, config=None, overrides=None, suffix="yaml"): + source = tmp_path / "config.json" + source.write_text(json.dumps(_CONFIG if config is None else config)) + resource = tmp_path / f"resources.{suffix}" + values = _OVERRIDES if overrides is None else overrides + resource.write_text(json.dumps(values) if suffix == "json" else yaml.safe_dump(values)) + return source, resource + + +def _multimodal_config(text=None): + return { + "_name_or_path": "example/multimodal-checkpoint", + "architectures": ["ExampleMultimodalForConditionalGeneration"], + "model_type": "example_multimodal", + "hidden_size": 4096, + "num_hidden_layers": 32, + "text_config": dict(_CONFIG if text is None else text), + "vision_config": {"hidden_size": 512, "num_hidden_layers": 8}, + "audio_config": {"hidden_size": 256, "num_hidden_layers": 4}, + } + + +def _args(output: Path, source: Path, resource: Path | None = None, **changes) -> list[str]: + options = {**_IDENTITY, **_PILOT, **changes} + command = ["onboard", "init", "--model-config", str(source), "--output", str(output)] + if resource is not None: + command += ["--resource-overrides", str(resource)] + return command + [ + part + for name, value in options.items() + if value is not None + for part in ("--" + name.replace("_", "-"), str(value)) + ] + + +def _terminal(monkeypatch, answers=()): + prompts = [] + remaining = iter(answers) + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def answer(prompt): + prompts.append(prompt) + try: + value = next(remaining) + except StopIteration: + pytest.fail(f"unexpected prompt: {prompt}") + if isinstance(value, BaseException): + raise value + return value + + monkeypatch.setattr(builtins, "input", answer) + return prompts + + +def test_help_documents_local_config_and_override_options(capsys): + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "init", "--help"]) + assert result.value.code == 0 + text = capsys.readouterr().out + assert "--model-config" in text + assert "--resource-overrides" in text + + +@pytest.mark.parametrize("suffix", ["json", "yaml"]) +def test_scripted_config_embeds_complete_profile_and_sources(tmp_path, monkeypatch, capsys, suffix): + _terminal(monkeypatch) + source, resources = _files(tmp_path, suffix=suffix) + output = tmp_path / "request.yaml" + + assert cli.main(_args(output, source, resources)) == 0 + + request = SupportRequest.from_yaml(output) + profile = request.fpm_profile + assert request.identity.model == _CONFIG["_name_or_path"] + assert request.identity.model_kind == "dense" + assert request.identity.model_revision == _IDENTITY["model_revision"] + assert profile.architecture == "LlamaForCausalLM" + assert profile.context_length == 32768 + assert request.search.context_length == 32768 + assert profile.num_experts == 0 + assert profile.deployments[0].resources.weights_bytes == _OVERRIDES["weights_bytes"] + provenance = profile.provenance + profile.deployments[0].resources.provenance + assert hashlib.sha256(source.read_bytes()).hexdigest() in provenance + assert "weights_bytes" in provenance + assert "override" in provenance.lower() + assert _OVERRIDES["provenance"] in provenance + transcript = capsys.readouterr().out + assert "weights_bytes" in transcript + assert "source" in transcript.lower() + assert "estimate" in transcript.lower() + + +def test_explicit_identity_and_pilot_context_override_config_hints(tmp_path, monkeypatch): + _terminal(monkeypatch) + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + + assert cli.main(_args(output, source, resources, model="deployment/model-alias", context_length=2048)) == 0 + + request = SupportRequest.from_yaml(output) + assert request.identity.model == request.fpm_profile.model == "deployment/model-alias" + assert request.search.context_length == 2048 + assert request.fpm_profile.context_length == 32768 + + +@pytest.mark.parametrize("multimodal", [False, True]) +def test_guided_review_requires_explicit_accept_before_creating_output(tmp_path, monkeypatch, capsys, multimodal): + source, resources = _files(tmp_path, config=_multimodal_config() if multimodal else None) + output = tmp_path / "new" / "request.yaml" + prompts = [] + answers = iter(["", "save", "accept"]) + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def respond(prompt): + prompts.append(prompt) + assert not output.parent.exists(), "setup wrote output before explicit acceptance" + return next(answers) + + monkeypatch.setattr(builtins, "input", respond) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert prompts == ["Review action (accept/edit/cancel): "] * 3 + transcript = capsys.readouterr().out + assert "Review FPM profile before saving:" in transcript + review = transcript.split("Review FPM profile before saving:", 1)[1] + assert "Text decoder only" in review + assert "encoders, projectors, preprocessing" in review + assert "full multimodal latency is not modeled" in review + assert "weights_bytes: 2147483648" in transcript + assert "source: user override" in transcript + assert "per rank" in transcript + assert "checkpoint-revision-123" in transcript + assert "h200_sxm" in transcript + assert SupportRequest.from_yaml(output).profile_deployment().resources.weights_bytes == 2 * 1024**3 + + +def test_guided_review_edits_file_values_and_provenance_then_replays_saved_profile(tmp_path, monkeypatch, capsys): + source, resources = _files(tmp_path) + original_overrides = resources.read_bytes() + output = tmp_path / "request.yaml" + note = "Synthetic replacement bounds for CLI validation; not measured." + prompts = _terminal( + monkeypatch, + [ + "edit", + "not_a_field", + "weights_bytes", + "many", + "0.1 B", + "-1", + "3 GiB", + "edit", + "weights_bytes", + "4 GiB", + "edit", + "kv_cache_dtype", + "auto", + "fp8", + "edit", + "max_num_tokens", + "16384", + "edit", + "provenance", + note, + "accept", + ], + ) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + request = SupportRequest.from_yaml(output) + deployment = request.profile_deployment() + assert deployment.resources.weights_bytes == 4 * 1024**3 + assert deployment.kv_cache_dtype == "fp8" + assert deployment.resources.max_num_tokens == 16384 + # Explicit bounds remain user declarations when their dependencies change. + assert deployment.resources.kv_bytes_per_token == _OVERRIDES["kv_bytes_per_token"] + assert deployment.resources.activations_bytes == _OVERRIDES["activations_bytes"] + provenance = json.loads(request.fpm_profile.provenance)["fields"] + assert provenance["weights_bytes"]["value"] == 4 * 1024**3 + assert provenance["provenance"]["value"] == note + assert all("user override" in provenance[field]["source"] for field in ("weights_bytes", "kv_bytes_per_token")) + assert resources.read_bytes() == original_overrides + assert sum(prompt.startswith("weights_bytes") for prompt in prompts) == 5 + transcript = capsys.readouterr().out + assert "Choose an editable field" in transcript + assert "whole number of bytes" in transcript + review = transcript.rsplit("Review FPM profile before saving:", 1)[1] + assert "weights_bytes: 4294967296" in review + assert "kv_cache_dtype: fp8" in review + assert note in review + assert _OVERRIDES["provenance"] not in review + + source.unlink() + resources.unlink() + plan = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(plan)]) == 0 + prediction = CorePredictionConfig.from_yaml(plan / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan / "recommend/pilot.yaml") + assert prediction.engine.fpm_profile == recommendation.engine.fpm_profile == request.fpm_profile + assert json.loads((plan / "fpm-model-profile.json").read_text()) == request.fpm_profile.model_dump(mode="json") + + +def test_guided_review_recomputes_inferred_kv_and_activation_bounds(tmp_path, monkeypatch, capsys): + config = { + **_CONFIG, + "hidden_size": 4096, + "intermediate_size": 14336, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 32000, + } + source, resources = _files( + tmp_path, + config=config, + overrides={"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"}, + ) + output = tmp_path / "request.yaml" + _terminal( + monkeypatch, + ["edit", "kv_cache_dtype", "fp8", "edit", "max_num_tokens", "4096", "edit", "max_batch_size", "8", "accept"], + ) + + assert cli.main(_args(output, source, resources, tensor_parallel=4, concurrency=4) + ["--interactive"]) == 0 + + request = SupportRequest.from_yaml(output) + deployment = request.profile_deployment() + assert deployment.parallel_tuple == (4, 1, 1, 1, 1, 1) + assert deployment.resources.kv_bytes_per_token == 16384 # 2 K/V * 32 layers * 2 local heads * 128 * 1 byte + assert deployment.resources.activations_bytes is None + assert json.loads(deployment.resources.provenance)["fields"]["activations_bytes"]["value"] == 167772160 + assert deployment.resources.max_num_tokens == 4096 + assert deployment.resources.max_batch_size == 8 + provenance = json.loads(request.fpm_profile.provenance)["fields"] + assert "exact linear tensor geometry" in provenance["kv_bytes_per_token"]["source"] + assert "estimate:" in provenance["activations_bytes"]["source"] + assert "max_num_tokens=4096" in provenance["activations_bytes"]["source"] + assert "max_batch_size=8" in provenance["activations_bytes"]["source"] + transcript = capsys.readouterr().out + assert "kv_bytes_per_token: 32768" in transcript + assert "activations_bytes: 335544320" in transcript + review = transcript.rsplit("Review FPM profile before saving:", 1)[1] + assert "kv_bytes_per_token: 16384" in review + assert "activations_bytes: 167772160" in review + + +@pytest.mark.parametrize( + "field,value,error", + [ + ("weights_bytes", str(2**53), "total non-KV resource bytes"), + ("architecture", "UnknownForCausalLM", "conflicts with the source config"), + ("num_experts", "4", "conflicts with the source config"), + ("context_length", "65536", "exceeds the source config context limit"), + ("context_length", "8192", "search.context_length exceeds"), + ], +) +def test_guided_review_rejects_inconsistent_edits_and_retains_complete_prior_request( + tmp_path, monkeypatch, capsys, field, value, error +): + source, resources = _files(tmp_path) + baseline = tmp_path / "baseline.yaml" + assert cli.main(_args(baseline, source, resources)) == 0 + output = tmp_path / "request.yaml" + _terminal(monkeypatch, ["edit", field, value, "accept"]) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert SupportRequest.from_yaml(output) == SupportRequest.from_yaml(baseline) + transcript = capsys.readouterr().out + assert "Edit rejected:" in transcript + assert error in transcript + assert "Previous profile values retained" in transcript + + +def test_guided_review_keeps_unavailable_activation_memory_pending(tmp_path, monkeypatch, capsys): + source, resources = _files( + tmp_path, overrides={name: value for name, value in _OVERRIDES.items() if name != "activations_bytes"} + ) + output = tmp_path / "new" / "request.yaml" + prompts = [] + answers = iter(["edit", "fmha_quant_mode", "fp8", "accept"]) + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def respond(prompt): + assert not output.parent.exists(), "staged dependency must not write a partial request" + prompts.append(prompt) + return next(answers) + + monkeypatch.setattr(builtins, "input", respond) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert len(prompts) == 4 + assert not any(prompt.startswith("activations_bytes") for prompt in prompts) + assert prompts[-1] == "Review action (accept/edit/cancel): " + deployment = SupportRequest.from_yaml(output).profile_deployment() + assert deployment.fmha_quant_mode == "fp8" + assert deployment.resources.activations_bytes is None + assert deployment.resources.memory_source == "pending" + review = capsys.readouterr().out.rsplit("Review FPM profile before saving:", 1)[1] + assert "fmha_quant_mode: fp8" in review + assert "No activation or non-KV byte estimate is required" in review + + +def test_guided_review_rolls_back_edit_and_new_dependency_together_when_invalid(tmp_path, monkeypatch, capsys): + overrides = {name: value for name, value in _OVERRIDES.items() if name != "activations_bytes"} + overrides["weights_bytes"] = ( + 2**53 - 2 * 1024**3 - overrides["runtime_overhead_bytes"] - overrides["comm_overhead_bytes"] + ) + source, resources = _files(tmp_path, overrides=overrides) + baseline = tmp_path / "baseline.yaml" + assert cli.main(_args(baseline, source, resources)) == 0 + output = tmp_path / "request.yaml" + _terminal(monkeypatch, ["edit", "activations_bytes", "3 GiB", "accept"]) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert SupportRequest.from_yaml(output) == SupportRequest.from_yaml(baseline) + assert "total non-KV resource bytes" in capsys.readouterr().out + + +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize( + "answers", + [["cancel"], ["edit", "weights_bytes", "3 GiB", "cancel"]], + ids=["review", "after-edit"], +) +def test_guided_review_explicit_cancel_leaves_output_untouched(tmp_path, monkeypatch, existing, answers): + source, resources = _files(tmp_path) + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("preserve previous request\n") + _terminal(monkeypatch, answers) + + assert cli.main(_args(output, source, resources) + ["--interactive", "--overwrite"]) == 130 + + if existing: + assert output.read_text() == "preserve previous request\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + + +@pytest.mark.parametrize("interruption", [KeyboardInterrupt, EOFError]) +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize( + "answers", + [[], ["edit"], ["edit", "weights_bytes"], ["edit", "fmha_quant_mode", "fp8"]], + ids=["review", "field", "value", "dependency"], +) +def test_guided_review_interruption_leaves_output_untouched(tmp_path, monkeypatch, interruption, existing, answers): + source, resources = _files( + tmp_path, overrides={name: value for name, value in _OVERRIDES.items() if name != "activations_bytes"} + ) + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("preserve previous request\n") + _terminal(monkeypatch, [*answers, interruption()]) + + assert cli.main(_args(output, source, resources) + ["--interactive", "--overwrite"]) == 130 + + if existing: + assert output.read_text() == "preserve previous request\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + + +def test_scripted_identity_failure_lists_all_missing_identity_without_input(tmp_path, monkeypatch, capsys): + _terminal(monkeypatch) + source, _ = _files(tmp_path) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "init", "--model-config", str(source), "--output", str(output)]) + + assert result.value.code == 2 + errors = capsys.readouterr().err + for option in ("--model-revision", "--framework-version", "--gpu", "--interconnect"): + assert option in errors + assert "--resource-overrides" in errors + assert "identity" in errors.lower() + assert not output.parent.exists() + + +def test_scripted_profile_failure_lists_every_unresolved_field_without_input(tmp_path, monkeypatch, capsys): + _terminal(monkeypatch) + unknown = {"_name_or_path": "example/unknown", "architectures": ["UnknownForCausalLM"]} + source, _ = _files(tmp_path, config=unknown) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, model_kind="dense", context_length=4096)) + + assert result.value.code == 2 + errors = capsys.readouterr().err + for field in ( + "context_length", + "num_experts", + "kv_bytes_per_token", + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kv_cache_dtype", + ): + assert field in errors + assert "--resource-overrides" in errors + assert "JSON" in errors and "YAML" in errors + assert not output.parent.exists() + + +def test_guided_config_asks_only_unresolved_inputs_and_recovers_invalid_answers(tmp_path, monkeypatch, capsys): + overrides = {key: value for key, value in _OVERRIDES.items() if key != "fmha_quant_mode"} + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["invalid-precision", "bfloat16", "accept"]) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert len(prompts) == 3 + assert all("fmha_quant_mode" in prompt for prompt in prompts[:-1]) + assert SupportRequest.from_yaml(output).fpm_profile.deployments[0].fmha_quant_mode == "bfloat16" + assert "invalid-precision" in capsys.readouterr().out + + +@pytest.mark.parametrize("option", ["model_revision", "framework_version"]) +@pytest.mark.parametrize("missing_precision", [False, True]) +def test_guided_final_profile_identity_correction_preserves_resource_answers( + tmp_path, monkeypatch, option, missing_precision +): + overrides = {key: value for key, value in _OVERRIDES.items() if not missing_precision or key != "fmha_quant_mode"} + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "request.yaml" + answers = (["bfloat16"] if missing_precision else []) + [_IDENTITY[option], "accept"] + prompts = _terminal(monkeypatch, answers) + + assert cli.main(_args(output, source, resources, **{option: "unknown"}) + ["--interactive"]) == 0 + + assert len(prompts) == len(answers) + if missing_precision: + assert "fmha_quant_mode" in prompts[0] + assert ("Pinned model revision" if option == "model_revision" else "Pinned vLLM version") in prompts[-2] + request = SupportRequest.from_yaml(output) + assert getattr(request.identity, option) == _IDENTITY[option] + assert request.profile_deployment().fmha_quant_mode == "bfloat16" + assert request.profile_deployment().resources.weights_bytes == _OVERRIDES["weights_bytes"] + + +@pytest.mark.parametrize("missing_precision", [False, True]) +def test_guided_composite_resource_error_does_not_repeat_unrelated_precision_prompt( + tmp_path, monkeypatch, capsys, missing_precision +): + overrides = { + key: value + for key, value in _OVERRIDES.items() + if key != "activations_bytes" and (not missing_precision or key != "fmha_quant_mode") + } + overrides["weights_bytes"] = 2**53 - overrides["runtime_overhead_bytes"] - overrides["comm_overhead_bytes"] + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "new" / "request.yaml" + prompts = _terminal(monkeypatch, ["bfloat16"] if missing_precision else []) + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, resources) + ["--interactive"]) + + assert result.value.code == 2 + assert "total non-KV resource bytes must not exceed 2**53" in capsys.readouterr().err + assert len(prompts) == int(missing_precision) + if missing_precision: + assert "fmha_quant_mode" in prompts[0] + assert not output.parent.exists() + + +def test_scripted_late_dense_architecture_rejects_moe_config_without_output(tmp_path, monkeypatch, capsys): + _terminal(monkeypatch) + config = {key: value for key, value in _CONFIG.items() if key not in ("architectures", "model_type")} + config["num_experts"] = 4 + source, resources = _files(tmp_path, config=config, overrides={**_OVERRIDES, "architecture": "LlamaForCausalLM"}) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, resources)) + + assert result.value.code == 2 + assert "dense architecture" in capsys.readouterr().err + assert not output.parent.exists() + + +@pytest.mark.parametrize("field", ["num_experts_per_tok", "first_k_dense_replace"]) +def test_real_cli_accepts_optional_null_moe_metadata_with_late_expert_count(tmp_path, field): + config = {**_CONFIG, "architectures": ["MixtralForCausalLM"], "model_type": "mixtral", field: None} + source, resources = _files(tmp_path, config=config, overrides={**_OVERRIDES, "num_experts": 4}) + output = tmp_path / "request.yaml" + + result = subprocess.run( + [sys.executable, "-m", "aisimulate", *_args(output, source, resources)], + input="", + text=True, + capture_output=True, + timeout=30, + ) + + assert result.returncode == 0, result.stderr + request = SupportRequest.from_yaml(output) + assert request.identity.model_kind == "moe" + assert request.fpm_profile.num_experts == 4 + assert request.profile_deployment().resources.weights_bytes == _OVERRIDES["weights_bytes"] + + +def test_real_cli_multimodal_intake_preserves_text_scope_and_original_source(tmp_path): + source, resources = _files(tmp_path, config=_multimodal_config()) + source_hash = hashlib.sha256(source.read_bytes()).hexdigest() + output = tmp_path / "request.yaml" + result = subprocess.run( + [sys.executable, "-m", "aisimulate", *_args(output, source, resources)], + input="", + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert "Text decoder only" in result.stdout + assert "encoders, projectors, preprocessing" in result.stdout + assert "full multimodal latency is not modeled" in result.stdout + assert "Review action" not in result.stdout + request = SupportRequest.from_yaml(output) + assert request.identity.model == "example/multimodal-checkpoint" + assert request.fpm_profile.architecture == "LlamaForCausalLM" + provenance = json.loads(request.fpm_profile.provenance) + assert provenance["config_sha256"] == source_hash + assert "Text decoder only" in provenance["config_notes"]["modeling_scope"] + assert "text_config" in provenance["config_notes"]["decoder_config"] + assert provenance["config_notes"]["hidden_size"] == "config hidden_size=128" + + +def test_real_cli_unknown_multimodal_decoder_lists_missing_resource_bounds_without_writing(tmp_path): + text = {key: value for key, value in _CONFIG.items() if key not in ("architectures", "model_type")} + text["num_experts"] = 0 + source, _ = _files(tmp_path, config=_multimodal_config(text)) + output = tmp_path / "new" / "request.yaml" + result = subprocess.run( + [sys.executable, "-m", "aisimulate", *_args(output, source)], + input="", + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 2, result.stderr + assert "Text decoder only" in result.stdout + for field in ("kv_bytes_per_token", "cache_layout"): + assert field in result.stderr + assert "--resource-overrides" in result.stderr + assert "unsupported encoder, multimodal" not in result.stderr + assert not output.parent.exists() + + +def test_bundled_modelopt_string_cache_config_creates_complete_request(tmp_path, monkeypatch): + _terminal(monkeypatch) + config_path = ( + Path(__file__).parents[2] / "src/aisimulate_core/model_configs/Qwen--Qwen3-32B-FP8-Static-PerTensor_config.json" + ) + source, resources = _files( + tmp_path, + config=json.loads(config_path.read_text()), + overrides={"weights_bytes": 40 * 1024**3, "fmha_quant_mode": "bfloat16", "comm_quant_mode": "half"}, + ) + output = tmp_path / "request.yaml" + + assert cli.main(_args(output, source, resources, model="review/qwen3-config")) == 0 + + profile = SupportRequest.from_yaml(output).fpm_profile + assert profile.architecture == "Qwen3ForCausalLM" + assert profile.num_experts == 0 + deployment = profile.deployments[0] + assert deployment.gemm_quant_mode == "fp8_static" + assert deployment.moe_quant_mode == deployment.kv_cache_dtype == "fp8" + assert deployment.fmha_quant_mode == "bfloat16" + assert deployment.resources.kv_bytes_per_token == 2 * 64 * 8 * 128 + + +def test_guided_missing_identity_is_asked_before_profile_fields(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["checkpoint-revision-123", "0.25.1", "h200_sxm", "nvswitch", "accept"]) + + assert cli.main(_args(output, source, resources, **dict.fromkeys(_IDENTITY)) + ["--interactive"]) == 0 + + assert len(prompts) == 5 + assert "Pinned model revision" in prompts[0] + assert not any("Model name" in prompt or "Model kind" in prompt for prompt in prompts) + + +def test_guided_unknown_model_metadata_bounds_pilot_before_ordinary_prompts(tmp_path, monkeypatch): + config = {**_CONFIG, "architectures": ["UnknownForCausalLM"], "model_type": "unknown"} + del config["max_position_embeddings"] + source, resources = _files(tmp_path, config=config) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["invalid", "2048", "-1", "0", "accept"]) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert len(prompts) == 5 + assert all("context_length" in prompt for prompt in prompts[:2]) + assert all("num_experts" in prompt for prompt in prompts[2:-1]) + request = SupportRequest.from_yaml(output) + assert request.search.context_length == request.fpm_profile.context_length == 2048 + assert request.identity.model_kind == "dense" + + +@pytest.mark.parametrize( + "changes,answers", + [ + ({"model_kind": "moe"}, ["dense"]), + ({"context_length": 65536}, ["8192"]), + ({"tensor_parallel": 8}, ["4"]), + ({"model_revision": "unknown"}, ["checkpoint-revision-123"]), + ], +) +def test_guided_setup_can_correct_request_conflicts_with_config_metadata(tmp_path, monkeypatch, changes, answers): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, [*answers, "accept"]) + + assert cli.main(_args(output, source, resources, **changes) + ["--interactive"]) == 0 + + assert len(prompts) == len(answers) + 1 + request = SupportRequest.from_yaml(output) + assert request.identity.model_kind == "dense" + assert request.identity.model_revision == "checkpoint-revision-123" + assert request.search.context_length <= request.fpm_profile.context_length + assert request.search.tensor_parallel <= 4 + + +def test_guided_precision_answers_resolve_supported_memory_without_redundant_prompts(tmp_path, monkeypatch): + source, _ = _files(tmp_path) + output = tmp_path / "request.yaml" + prompts = [] + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + answers = {"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"} + + def respond(prompt): + prompts.append(prompt) + if prompt == "Review action (accept/edit/cancel): ": + return "accept" + name = prompt.split(" ")[0] + assert name in answers, f"unexpected prompt for a derived input: {prompt}" + return answers[name] + + monkeypatch.setattr(builtins, "input", respond) + + assert cli.main(_args(output, source) + ["--interactive"]) == 0 + + assert len(prompts) == 3 + resources = SupportRequest.from_yaml(output).profile_deployment().resources + assert resources.kv_bytes_per_token == 512 + assert resources.max_num_tokens == 8192 + assert resources.max_batch_size == 256 + + +def test_guided_memory_accepts_exact_units_and_reprompts_invalid_quantities(tmp_path, monkeypatch, capsys): + config = {**_CONFIG, "architectures": ["UnknownForCausalLM"], "model_type": "unknown", "num_local_experts": 0} + overrides = {key: value for key, value in _OVERRIDES.items() if key != "weights_bytes"} + source, resources = _files(tmp_path, config=config, overrides=overrides) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["edit", "weights_bytes", "many", "0.1 B", "-1", "2 GiB", "accept"]) + + assert cli.main(_args(output, source, resources) + ["--interactive"]) == 0 + + assert len(prompts) == 7 + assert all("weights_bytes" in prompt and "per rank" in prompt for prompt in prompts[2:-1]) + assert SupportRequest.from_yaml(output).profile_deployment().resources.weights_bytes == 2 * 1024**3 + assert "whole number of bytes" in capsys.readouterr().out + + +@pytest.mark.parametrize("interruption", [KeyboardInterrupt, EOFError]) +@pytest.mark.parametrize("existing", [False, True]) +def test_resource_prompt_cancellation_writes_no_partial_request(tmp_path, monkeypatch, interruption, existing): + overrides = {key: value for key, value in _OVERRIDES.items() if key != "fmha_quant_mode"} + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("preserve previous request\n") + _terminal(monkeypatch, [interruption()]) + + assert cli.main(_args(output, source, resources) + ["--interactive", "--overwrite"]) == 130 + + if existing: + assert output.read_text() == "preserve previous request\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "overrides", + [ + {"weights_bytes": -1}, + {"weights_bytes": True}, + {"weights_bytes": 1.5}, + {"weights_bytes": "2 GiB"}, + {"fmha_quant_mode": "auto"}, + {"not_a_profile_field": 1}, + {"resources": {"weights_bytes": 1}}, + {"max_num_tokens": 0}, + {"cache_layout": "recurrent"}, + ], +) +def test_invalid_overrides_fail_before_prompts_or_output(tmp_path, monkeypatch, overrides): + _terminal(monkeypatch) + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, resources) + ["--interactive"]) + + assert result.value.code == 2 + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "contents", ['{"weights_bytes": 1, "weights_bytes": 2}', "weights_bytes: 1\nweights_bytes: 2\n"] +) +def test_duplicate_override_fields_are_rejected_instead_of_discarded(tmp_path, monkeypatch, capsys, contents): + _terminal(monkeypatch) + source, resources = _files(tmp_path) + resources.write_text(contents) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, resources) + ["--interactive"]) + + assert result.value.code == 2 + assert "duplicate resource override" in capsys.readouterr().err + assert not output.parent.exists() + + +def test_config_interactive_requires_a_terminal_without_reading_input(tmp_path, monkeypatch, capsys): + _terminal(monkeypatch) + monkeypatch.setattr(sys.stdin, "isatty", lambda: False) + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source, resources) + ["--interactive"]) + + assert result.value.code == 2 + assert "requires a terminal" in capsys.readouterr().err + assert not output.exists() + + +@pytest.mark.parametrize("contents", ["{broken", "[]", "model_type: llama\n"]) +def test_model_config_requires_a_local_json_mapping(tmp_path, monkeypatch, contents): + _terminal(monkeypatch) + source = tmp_path / "config.json" + source.write_text(contents) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(_args(output, source) + ["--interactive"]) + + assert result.value.code == 2 + assert not output.parent.exists() + + +@pytest.mark.parametrize("cache_type", [[], {}, True, 8, 1.5]) +@pytest.mark.parametrize("existing", [False, True]) +def test_real_cli_rejects_malformed_cache_type_without_writing_request(tmp_path, cache_type, existing): + config = {**_CONFIG, "quantization_config": {"kv_cache_scheme": {"num_bits": 8, "type": cache_type}}} + source, resources = _files(tmp_path, config=config) + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("preserve previous request\n") + + result = subprocess.run( + [sys.executable, "-m", "aisimulate", *_args(output, source, resources), "--overwrite"], + input="", + text=True, + capture_output=True, + timeout=30, + ) + + assert result.returncode == 2, result.stderr + assert "quantization_config.kv_cache_scheme.type" in result.stderr + assert "Traceback" not in result.stderr + if existing: + assert output.read_text() == "preserve previous request\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "options,expected", + [ + (["--resource-overrides", "missing.yaml"], "requires --model-config"), + (["--model-config", "missing.json", "--fpm-profile", "missing.yaml"], "not allowed"), + ], +) +def test_config_route_option_relationships_are_rejected_before_reading_files( + tmp_path, monkeypatch, capsys, options, expected +): + _terminal(monkeypatch) + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "init", *options, "--output", str(output), "--interactive"]) + + assert result.value.code == 2 + assert expected in capsys.readouterr().err + assert not output.parent.exists() + + +@pytest.mark.parametrize("multimodal", [False, True]) +@pytest.mark.parametrize("interactive", [False, True]) +@pytest.mark.parametrize( + "tp,dp,moe_tp,moe_ep,preset", + [(4, 1, 4, 1, "pure_tp"), (1, 8, 1, 8, "dep"), (8, 1, 1, 8, "tep")], +) +def test_config_init_to_plan_replay_preserves_full_topology_without_source_files( + tmp_path, monkeypatch, capsys, tp, dp, moe_tp, moe_ep, preset, multimodal, interactive +): + config = { + **_CONFIG, + "architectures": ["ExampleMoeForCausalLM"], + "model_type": "example_moe", + "num_local_experts": 8, + } + if multimodal: + config = _multimodal_config(config) + source, resources = _files(tmp_path, config=config) + output = tmp_path / "request.yaml" + real_import = builtins.__import__ + + def no_model_or_download_import(name, *args, **kwargs): + assert not name.startswith(("collector", "huggingface_hub", "transformers")), name + assert ".sdk.models" not in name, name + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_model_or_download_import) + prompts = _terminal(monkeypatch, ["accept"] if interactive else []) + assert ( + cli.main( + _args( + output, + source, + resources, + tensor_parallel=tp, + attention_data_parallel=dp, + moe_tensor_parallel=moe_tp, + moe_expert_parallel=moe_ep, + ) + + (["--interactive"] if interactive else []) + ) + == 0 + ) + source.unlink() + resources.unlink() + request = SupportRequest.from_yaml(output) + assert prompts == (["Review action (accept/edit/cancel): "] if interactive else []) + assert "Collection GPUs required: " + str(tp * dp) in capsys.readouterr().out + assert not {"gpu_count", "node_count", "gpus_per_node"} & request.identity.model_dump().keys() + assert "max_candidates" not in request.search.model_dump() + assert "Text decoder only" in json.loads(request.fpm_profile.provenance)["config_notes"]["modeling_scope"] + assert request.parallel_preset == preset + assert request.profile_deployment().parallel_tuple == (tp, 1, dp, moe_tp, moe_ep, 1) + plan = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(plan)]) == 0 + saved_plan = json.loads((plan / "support-plan.json").read_text()) + assert saved_plan["fpm"]["collection_gpus_required"] == tp * dp + command = saved_plan["fpm"]["plan_command"] + assert command[command.index("--fpm-gpu-counts") + 1] == str(tp * dp) + assert command[command.index("--fpm-parallel-presets") + 1] == preset + assert SupportRequest.from_yaml(plan / "request.yaml") == request + prediction = CorePredictionConfig.from_yaml(plan / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan / "recommend/pilot.yaml") + assert recommendation.optimization.constraints.max_candidate_gpus == tp * dp + assert len(recommendation.engine.workers.aggregated.parallelism.preset) == 1 + assert recommendation.engine.workers.aggregated.parallelism.preset[0].replicas == 1 + assert prediction.engine.fpm_profile == recommendation.engine.fpm_profile == request.fpm_profile + for generated, path in ( + (prediction, plan / "predict/pilot.yaml"), + (recommendation, plan / "recommend/pilot.yaml"), + ): + assert generated.engine.systems_paths == [str(plan / "systems")] + timing = generated.engine.workers.aggregated.timing + assert timing.estimation_mode == "fpm_interpolation" + assert timing.fallback_policy == "deny" + assert timing.estimator_config["fpm_interpolation"]["method"] == "direct" + saved_engine = yaml.safe_load(path.read_text())["engine"] + assert saved_engine["systems_paths"] == generated.engine.systems_paths + assert "systems_path" not in saved_engine + assert saved_engine["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + saved_timing = saved_engine["workers"]["aggregated"]["timing"] + assert saved_timing == { + "type": "default", + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": {"fpm_interpolation": {"method": "direct"}}, + } + assert prediction.engine.workers.aggregated.parallelism.tensor == tp + assert prediction.engine.workers.aggregated.parallelism.attention_data == dp + assert prediction.engine.workers.aggregated.parallelism.moe_tensor == moe_tp + assert prediction.engine.workers.aggregated.parallelism.moe_expert == moe_ep + assert json.loads((plan / "fpm-model-profile.json").read_text()) == request.fpm_profile.model_dump(mode="json") diff --git a/python/aisimulate/tests/unit/test_onboard_runtime.py b/python/aisimulate/tests/unit/test_onboard_runtime.py new file mode 100644 index 000000000..85a31307e --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_runtime.py @@ -0,0 +1,1318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public onboarding from synthetic, independently validated runtime evidence.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +from collector.fpm_forward.runtime_probe import normalize_probe_launch + +import aisimulate.main as cli +from aisimulate.support.checkpoint import _load, report, save_checkpoint +from aisimulate.support.schema import SupportRequest + +from .collector.test_runtime_observations import observation_fixture + +pytestmark = pytest.mark.unit + + +def _write(path, value): + path.write_text(json.dumps(value)) + + +def _campaign(tmp_path, *, capacity_multiplier=1, context_length=4096, tp=2, cpu_policy=True, worker_type=None): + index, launches = observation_fixture(tmp_path / "probe", dense=True, tp=tp) + launch = launches["tp2"] + if cpu_policy: + cpus, binding = cpu_policy if isinstance(cpu_policy, tuple) else (16, "cores") + launch["deployment"].update(cpus_per_task=cpus, cpu_bind=binding) + model_path = Path(launch["model_config"]["path"]) + _write( + model_path, + { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", + "hidden_size": 128, + "vocab_size": 1024, + "num_hidden_layers": 2, + "num_attention_heads": max(4, tp), + "num_key_value_heads": max(4, tp), + "intermediate_size": 256, + "quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "kv_cache_quant_algo": "none"}, + }, + ) + checkpoint_model = tmp_path / "model-checkpoint" + checkpoint_model.mkdir() + (checkpoint_model / "config.json").write_bytes(model_path.read_bytes()) + launch["identity"]["model"] = str(checkpoint_model) + launch["model_config"]["sha256"] = hashlib.sha256(model_path.read_bytes()).hexdigest() + launch["collection"]["max_model_len"] = context_length + if worker_type is not None: + launch["worker_type"] = worker_type + launch["collection"]["cudagraph_mode"] = { + "prefill": "PIECEWISE", + "decode": "FULL_DECODE_ONLY", + "aggregated": "FULL_AND_PIECEWISE", + }[worker_type] + launch = normalize_probe_launch(launch) + document = json.loads(index.read_text()) + document["configurations"]["tp2"]["launch"] = launch + if worker_type in {"prefill", "decode"}: + document["configurations"]["tp2"]["attempts"][0]["phases"].pop( + "decode" if worker_type == "prefill" else "prefill" + ) + for phase in document["configurations"]["tp2"]["attempts"][0]["phases"].values(): + for ref in [phase["launch_manifest"], *phase["artifacts"]]: + path = index.parent / ref["path"] + value = json.loads(path.read_text()) + value["launch"] = launch + if ref.get("kind") == "observation": + value["identity"] = launch["identity"] + value["resolved_config"]["model_config"]["model"] = str(checkpoint_model) + value["model_config_sha256"] = launch["model_config"]["sha256"] + value["resolved_config"]["model_config"]["max_model_len"] = context_length + if worker_type is not None: + value["resolved_config"]["compilation_config"]["cudagraph_mode"] = launch["collection"][ + "cudagraph_mode" + ] + cache = value["cache"] + cache["num_blocks"] *= capacity_multiplier + if value["kind"] == "worker": + for field in ("available_cache_bytes", "allocated_cache_bytes"): + cache[field] *= capacity_multiplier + for storage in cache["storages"]: + storage["size_bytes"] *= capacity_multiplier + for allocation in cache["tensor_allocations"]: + allocation["size"] *= capacity_multiplier + for view in cache["layer_tensors"].values(): + view["shape"][view["block_axis"]] *= capacity_multiplier + else: + cache["initial_free_blocks"] = cache["num_blocks"] - cache["reserved_blocks"] + _write(path, value) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + precision = dict(launch["precision"]) + precision["kv_cache_dtype"] = precision.pop("kvcache_quant_mode") + draft = { + **({"worker_type": worker_type} if worker_type is not None else {}), + "identity": launch["identity"], + "workload": {"input_tokens": 64, "output_tokens": 16, "concurrency": 1, "request_count": 1}, + "search": { + "tensor_parallel": tp, + "attention_data_parallel": 1, + "moe_tensor_parallel": 1, + "moe_expert_parallel": 1, + "context_length": context_length, + }, + "collection": { + "max_num_tokens": 1024, + "max_batch_size": 64, + "gpu_memory_utilization": 0.9, + "prefill_cudagraph_policy": "runtime", + **({"cudagraph_mode": launch["collection"]["cudagraph_mode"]} if worker_type is not None else {}), + }, + } + checkpoint = tmp_path / "checkpoint.json" + save_checkpoint( + checkpoint, + patch={ + "inputs": {"model_config": str(model_path)}, + "configurations": { + name: { + "draft_request": copy.deepcopy(draft), + "inputs": {"precision": precision, "collection_deployment": launch["deployment"]}, + } + for name in ("tp2", "incomplete") + }, + }, + expected_revision=None, + accept=[], + ) + return checkpoint, index, launch + + +def _import(checkpoint, index, output, *, configurations=()): + args = [ + "onboard", + "import-observations", + "--checkpoint", + str(checkpoint), + "--observations", + str(index), + "--output-dir", + str(output), + ] + for name in configurations: + args.extend(("--configuration", name)) + return cli.main(args) + + +def test_import_keeps_partial_results_and_memory_is_independent_of_timing(tmp_path, capsys): + checkpoint, index, _ = _campaign(tmp_path) + assert _import(checkpoint, index, tmp_path / "drafts") == 1 + result = json.loads(capsys.readouterr().out) + assert result["configurations"]["tp2"]["status"] == "complete" + assert result["configurations"]["incomplete"]["status"] == "incomplete" + state = _load(checkpoint) + request = SupportRequest.model_validate(state.configurations["tp2"].draft_request) + assert request.profile_deployment().resources.runtime_memory.kv_cache_bytes == 93 * 128 + assert state.configurations["tp2"].acceptance is None + assert not state.configurations["incomplete"].draft_request.get("fpm_profile") + assert state.configurations["incomplete"].progress.status == "blocked" + assert state.configurations["incomplete"].progress.blockers == result["configurations"]["incomplete"]["diagnostics"] + assert not list((tmp_path / "drafts").rglob("*.parquet")) + assert not report(state, checkpoint)["integrity_issues"] + previous = next(item for item in state.configurations["tp2"].history if item["event"] == "draft_replaced") + assert previous["draft_request"]["search"]["tensor_parallel"] == 2 + + +def test_legacy_memory_import_keeps_missing_cpu_policy_and_source_bytes(tmp_path, capsys): + checkpoint, index, launch = _campaign(tmp_path, cpu_policy=False) + preserved = {path: path.read_bytes() for path in index.parent.rglob("*") if path.is_file()} + assert "cpus_per_task" not in launch["deployment"] + assert _import(checkpoint, index, tmp_path / "legacy-drafts", configurations=["tp2"]) == 0 + assert all(path.read_bytes() == contents for path, contents in preserved.items()) + state = _load(checkpoint) + assert "cpus_per_task" not in state.configurations["tp2"].inputs["collection_deployment"] + request = SupportRequest.model_validate(state.configurations["tp2"].draft_request) + assert request.profile_deployment().resources.runtime_memory.kv_cache_bytes == 93 * 128 + assert not report(state, checkpoint)["integrity_issues"] + + +def test_snapshot_acceptance_survives_unrelated_probe_progress_but_detects_raw_tamper(tmp_path, capsys): + checkpoint, index, _ = _campaign(tmp_path) + assert _import(checkpoint, index, tmp_path / "drafts", configurations=["tp2"]) == 0 + capsys.readouterr() + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + document = json.loads(index.read_text()) + document["configurations"]["unrelated"] = {"attempts": []} + _write(index, document) + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + snapshot = next((tmp_path / "drafts").rglob("worker-0-0.json")) + snapshot.write_text("{}") + status = report(state, checkpoint) + assert not status["configurations"]["tp2"]["profile_accepted"] + assert status["integrity_issues"] + + +def test_probe_public_preview_does_not_need_geometry_or_execute_bundle(tmp_path, capsys, monkeypatch): + checkpoint, index, _ = _campaign(tmp_path) + from collector.fpm_forward import runtime_probe + + calls = [] + + def probe(configurations, **kwargs): + calls.append((configurations, kwargs)) + return {"status": "preview", "configurations": {key: {"status": "preview"} for key in configurations}} + + monkeypatch.setattr(runtime_probe, "probe_runtime", probe) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 0 + ) + capsys.readouterr() + facts, options = calls[0] + assert set(facts) == {"tp2", "incomplete"} + assert all(value["topology"]["tp"] == 2 for value in facts.values()) + assert all("cache_groups" not in str(value) for value in facts.values()) + assert options["execute"] is False + assert facts["tp2"]["collection"]["prefill_cudagraph_policy"] == "runtime" + + +@pytest.mark.parametrize("flag,value", [("--cpus-per-task", "32"), ("--cpu-bind", "none")]) +@pytest.mark.parametrize("checkpoint_policy", [True, False]) +def test_public_probe_partial_resume_recovers_saved_policy_before_validation( + tmp_path, capsys, monkeypatch, flag, value, checkpoint_policy +): + from collector.fpm_forward import runtime_probe + + checkpoint, index, launch = _campaign(tmp_path, cpu_policy=(32, "none")) + if checkpoint_policy: + assert _import(checkpoint, index, tmp_path / "drafts", configurations=["tp2"]) == 0 + capsys.readouterr() + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + acceptance = state.configurations["tp2"].acceptance + else: + state = _load(checkpoint) + save_checkpoint( + checkpoint, + patch={ + "configurations": { + "tp2": {"inputs": {"collection_deployment": {"cpus_per_task": None, "cpu_bind": None}}} + } + }, + expected_revision=state.revision, + accept=[], + ) + calls = [] + + def reused(configurations, **kwargs): + calls.append((configurations, kwargs)) + return {"status": "completed", "configurations": {"tp2": {"status": "completed", "resumed": True}}} + + monkeypatch.setattr(runtime_probe, "probe_runtime", reused) + original = _load(checkpoint).configurations["tp2"].draft_request + before = {path: path.read_bytes() for path in index.parent.rglob("*") if path.is_file()} + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(index.parent), + "--execute", + "--resume", + flag, + value, + ] + ) + == 0 + ) + capsys.readouterr() + assert len(calls) == 1 + assert calls[0][0]["tp2"]["deployment"] == launch["deployment"] + state = _load(checkpoint) + assert state.configurations["tp2"].draft_request == original + assert state.configurations["tp2"].inputs["collection_deployment"] == launch["deployment"] + if checkpoint_policy: + assert state.configurations["tp2"].acceptance == acceptance + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + assert all(path.read_bytes() == contents for path, contents in before.items()) + + +def test_real_public_preview_parses_bundle_without_executing_python(tmp_path, capsys): + checkpoint, index, _ = _campaign(tmp_path) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 0 + ) + result = json.loads(capsys.readouterr().out) + assert result["status"] == "preview" + assert set(result["configurations"]["tp2"]["phases"]) == {"prefill", "decode"} + + +@pytest.mark.parametrize("workload", [None, {"input_tokens": 1024, "output_tokens": 128}]) +def test_launch_only_preview_preserves_workload_without_validating_it(tmp_path, capsys, workload): + checkpoint, index, _ = _campaign(tmp_path, context_length=512) + state = _load(checkpoint) + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": {"workload": workload, "search": {"context_length": 512}}}}}, + expected_revision=state.revision, + accept=[], + ) + draft = copy.deepcopy(state.configurations["tp2"].draft_request) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 0 + ) + result = json.loads(capsys.readouterr().out) + assert set(result["configurations"]["tp2"]["phases"]) == {"prefill", "decode"} + assert _load(checkpoint).configurations["tp2"].draft_request == draft + with pytest.raises(ValueError, match="input and output tokens"): + SupportRequest.model_validate(draft) + # Import reaches real simulation validation, retaining the user's pending + # workload rather than inventing a smaller synthetic request. + assert _import(checkpoint, index, tmp_path / "drafts", configurations=["tp2"]) == 1 + result = json.loads(capsys.readouterr().out) + assert "input and output tokens" in str(result["configurations"]["tp2"]["diagnostics"]) + assert _load(checkpoint).configurations["tp2"].draft_request == draft + + +@pytest.mark.parametrize( + ("section", "field", "value", "diagnostic"), + [ + ("search", "context_length", 0, "context_length"), + ("search", "attention_data_parallel", 2, "topology"), + ("collection", "max_num_tokens", 1, "max_num_tokens"), + ("identity", "model_revision", "main", "pinned revision"), + ], +) +def test_launch_only_preview_still_rejects_invalid_launch_choices(tmp_path, capsys, section, field, value, diagnostic): + checkpoint, index, _ = _campaign(tmp_path) + state = _load(checkpoint) + save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": {"workload": None, section: {field: value}}}}}, + expected_revision=state.revision, + accept=[], + ) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 1 + ) + result = json.loads(capsys.readouterr().out) + assert diagnostic in str(result["configurations"]["tp2"]) + + +def test_probe_preserves_symlinked_config_directory_for_adjacent_sources(tmp_path, capsys): + from aisimulate.support.runtime import checkpoint_launch + + checkpoint, index, launch = _campaign(tmp_path) + snapshot = tmp_path / "hf-snapshot" + snapshot.mkdir() + model = snapshot / "config.json" + model.symlink_to(launch["model_config"]["path"]) + blob = tmp_path / "sidecar-blob" + _write(blob, {"quantization": {"quant_algo": "NVFP4", "kv_cache_quant_algo": "none"}}) + (snapshot / "hf_quant_config.json").symlink_to(blob) + state = _load(checkpoint) + state, _ = save_checkpoint( + checkpoint, patch={"inputs": {"model_config": str(model)}}, expected_revision=state.revision, accept=[] + ) + _, facts, _ = checkpoint_launch(state, "tp2", checkpoint) + assert facts["model_config"]["path"] == str(model) + assert facts["model_config"]["sha256"] == launch["model_config"]["sha256"] + assert facts["model_config"]["source_files"] == { + "hf_quant_config.json": hashlib.sha256(blob.read_bytes()).hexdigest() + } + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 0 + ) + result = json.loads(capsys.readouterr().out) + assert set(result["configurations"]["tp2"]["phases"]) == {"prefill", "decode"} + + +def test_incomplete_legacy_profile_can_be_saved_incrementally(tmp_path): + checkpoint = tmp_path / "checkpoint.json" + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"worker": {"draft_request": {"fpm_profile": {"model": "example/model"}}}}}, + expected_revision=None, + accept=[], + ) + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"worker": {"draft_request": {"fpm_profile": {"model_revision": "abc123"}}}}}, + expected_revision=state.revision, + accept=[], + ) + assert state.revision == 2 + assert state.configurations["worker"].draft_request == { + "fpm_profile": {"model": "example/model", "model_revision": "abc123"} + } + assert state.configurations["worker"].acceptance is None + + +@pytest.mark.parametrize("remove_profile", [False, True]) +def test_changed_scheduler_can_be_saved_and_probed_before_fresh_import(tmp_path, capsys, remove_profile): + checkpoint, index, request, _ = _imported(tmp_path, capsys) + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + accepted = state.configurations["tp2"].acceptance.model_dump(mode="json") + draft = {"collection": {"max_num_tokens": 2048}} + if remove_profile: + draft["fpm_profile"] = None + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": draft}}}, + expected_revision=state.revision, + accept=[], + ) + saved = state.configurations["tp2"] + assert saved.draft_request["collection"]["max_num_tokens"] == 2048 + assert saved.acceptance is None + assert saved.history[-1]["acceptance"] == accepted + if not remove_profile: + assert saved.draft_request["fpm_profile"] == request.fpm_profile.model_dump(mode="json", exclude_none=True) + # A second edit to an already incomplete request cannot rewrite the + # observations or discard source hashes in the imported provenance. + edited = copy.deepcopy(saved.draft_request) + runtime = edited["fpm_profile"]["deployments"][0]["resources"]["runtime_memory"] + provenance = json.loads(runtime["provenance"]) + provenance["observed_resources"]["runtime_memory"]["kv_cache_bytes"] = 1 + provenance["source_artifacts"] = [] + runtime["provenance"] = json.dumps(provenance) + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": edited}}}, + expected_revision=state.revision, + accept=[], + ) + assert state.configurations["tp2"].draft_request["fpm_profile"] == saved.draft_request["fpm_profile"] + with pytest.raises(ValueError): + save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 0 + ) + capsys.readouterr() + # A separate synthetic probe captures the selected scheduler envelope. + fresh = tmp_path / "fresh-probe" + shutil.copytree(index.parent, fresh) + fresh_index = fresh / index.name + document = json.loads(fresh_index.read_text()) + entry = document["configurations"]["tp2"] + entry["launch"]["collection"]["max_num_batched_tokens"] = 2048 + for phase in entry["attempts"][0]["phases"].values(): + for ref in [phase["launch_manifest"], *phase["artifacts"]]: + path = fresh / ref["path"] + record = json.loads(path.read_text()) + record["launch"] = entry["launch"] + if "resolved_config" in record: + record["resolved_config"]["scheduler_config"]["max_num_batched_tokens"] = 2048 + _write(path, record) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(fresh_index, document) + assert _import(checkpoint, fresh_index, tmp_path / "new-drafts", configurations=["tp2"]) == 0 + capsys.readouterr() + state = _load(checkpoint) + assert state.configurations["tp2"].acceptance is None + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + + +def test_failed_reimport_keeps_accepted_profile_and_historical_error_evidence(tmp_path, capsys): + checkpoint, index, _, _ = _imported(tmp_path, capsys) + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + previous = copy.deepcopy(state.configurations["tp2"]) + valid = index.read_bytes() + document = json.loads(valid) + document["configurations"]["tp2"]["attempts"] = [] + _write(index, document) + failed = tmp_path / "failed-drafts" + assert _import(checkpoint, index, failed, configurations=["tp2"]) == 1 + result = json.loads(capsys.readouterr().out)["configurations"]["tp2"] + state = _load(checkpoint) + saved = state.configurations["tp2"] + assert saved.acceptance == previous.acceptance + assert saved.draft_request == previous.draft_request + assert saved.progress.status == "blocked" + assert saved.progress.blockers == result["diagnostics"] + assert result["diagnostics"] + refs = {key: ref for key, ref in saved.artifacts.items() if key not in previous.artifacts} + assert refs and all(ref.archived and ref.sha256 for ref in refs.values()) + assert saved.history[-1]["validation"] == result["validation"] + assert saved.history[-1]["sha256"] == hashlib.sha256(Path(result["validation"]).read_bytes()).hexdigest() + assert cli.main(["onboard", "resume", "--checkpoint", str(checkpoint)]) == 0 + resumed = json.loads(capsys.readouterr().out) + assert resumed["configurations"]["tp2"]["profile_accepted"] + assert resumed["configurations"]["tp2"]["reported_progress"]["blockers"] == result["diagnostics"] + # Changing failed-attempt files cannot invalidate the independently accepted snapshot. + for ref in refs.values(): + (checkpoint.parent / ref.path).write_text("changed failed historical evidence") + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + assert report(state, checkpoint)["configurations"]["tp2"]["reported_progress"]["blockers"] == result["diagnostics"] + index.write_bytes(valid) + assert _import(checkpoint, index, tmp_path / "replacement-drafts", configurations=["tp2"]) == 0 + capsys.readouterr() + saved = _load(checkpoint).configurations["tp2"] + assert saved.acceptance is None + assert any(item.get("acceptance") == previous.acceptance.model_dump(mode="json") for item in saved.history) + + +def test_changed_inputs_archive_acceptance_even_when_draft_is_unchanged(tmp_path, capsys): + checkpoint, _, _, _ = _imported(tmp_path, capsys) + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + previous = copy.deepcopy(state.configurations["tp2"]) + state, _ = save_checkpoint( + checkpoint, + patch={ + "configurations": { + "tp2": { + "inputs": {"collection_deployment": {"image": "changed@sha256:" + "d" * 64}}, + "progress": { + "status": "blocked", + "blockers": ["Fresh image needs compatible observations."], + "next_action": "Run the selected image's probe.", + }, + } + } + }, + expected_revision=state.revision, + accept=[], + ) + saved = state.configurations["tp2"] + assert saved.draft_request == previous.draft_request + assert saved.acceptance is None + history = saved.history[-1] + assert history["acceptance"] == previous.acceptance.model_dump(mode="json") + assert history["draft_request"] == previous.draft_request + assert history["inputs"] == previous.inputs + assert saved.progress.status == "blocked" + assert saved.progress.blockers == ["Fresh image needs compatible observations."] + + +def test_unknown_geometry_does_not_mask_missing_explicit_topology(tmp_path, capsys): + checkpoint, index, _ = _campaign(tmp_path) + state = _load(checkpoint) + save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": {"search": {"tensor_parallel": None}}}}}, + expected_revision=state.revision, + accept=[], + ) + assert ( + cli.main( + [ + "onboard", + "probe-runtime", + "--checkpoint", + str(checkpoint), + "--configuration", + "tp2", + "--instrumentation", + str(index.parent / "manifest.yaml"), + "--output-dir", + str(tmp_path / "preview"), + ] + ) + == 1 + ) + result = json.loads(capsys.readouterr().out) + assert "tensor_parallel" in str(result["configurations"]["tp2"]["diagnostics"]) + + +def _imported(tmp_path, capsys, *, capacity_multiplier=1, tp=2, worker_type=None): + checkpoint, index, _ = _campaign(tmp_path, capacity_multiplier=capacity_multiplier, tp=tp, worker_type=worker_type) + assert _import(checkpoint, index, tmp_path / "drafts", configurations=["tp2"]) == 0 + result = json.loads(capsys.readouterr().out) + state = _load(checkpoint) + request = SupportRequest.model_validate(state.configurations["tp2"].draft_request) + return checkpoint, index, request, result["configurations"]["tp2"] + + +def test_profile_edit_preserves_measured_values_and_requires_fresh_probe_for_settings(tmp_path, capsys): + from aisimulate.support.runtime import runtime_probe_manifest, verify_runtime_profile + + checkpoint, _, request, _ = _imported(tmp_path, capsys) + state = _load(checkpoint) + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + payload = request.model_dump(mode="json", exclude_none=True) + resources = payload["fpm_profile"]["deployments"][0]["resources"] + observed = resources["runtime_memory"]["kv_cache_bytes"] + resources["runtime_memory"]["kv_cache_bytes"] = observed - 128 + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": payload}}}, + expected_revision=state.revision, + accept=[], + ) + changed = SupportRequest.model_validate(state.configurations["tp2"].draft_request) + manifest = runtime_probe_manifest(changed) + assert manifest["observed_resources"]["runtime_memory"]["kv_cache_bytes"] == observed + assert manifest["user_overrides"]["runtime_memory.kv_cache_bytes"]["value"] == observed - 128 + assert state.configurations["tp2"].acceptance is None + assert state.configurations["tp2"].history[-1]["acceptance"] is not None + verify_runtime_profile(changed) + payload = changed.model_dump(mode="json", exclude_none=True) + payload["collection"]["prefill_cudagraph_policy"] = "explicit" + state, _ = save_checkpoint( + checkpoint, + patch={"configurations": {"tp2": {"draft_request": payload}}}, + expected_revision=state.revision, + accept=[], + ) + with pytest.raises(ValueError, match="fresh probe"): + save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + + +def test_plan_carries_explicit_formal_observer_flags_and_execution_requires_acceptance(tmp_path, capsys, monkeypatch): + from aisimulate.support.runtime import runtime_probe_manifest + + checkpoint, _, request, files = _imported(tmp_path, capsys, capacity_multiplier=10) + root = tmp_path / "collection" + assert ( + cli.main(["onboard", "plan", "--config", files["request"], "--output-dir", str(root), "--format", "json"]) == 0 + ) + capsys.readouterr() + commands = json.loads((root / "commands.json").read_text()) + preview = commands["fpm_plan_local"] + assert "--fpm-runtime-instrumentation" in preview + assert "--fpm-runtime-launch" in preview + assert preview[preview.index("--fpm-runtime-configuration") + 1] == "tp2" + execute = commands["fpm_run_local"][1:] + assert "--executor" in execute and "slurm" in execute + with pytest.raises(SystemExit, match="2"): + cli.main(execute) + assert "acceptance" in capsys.readouterr().err + state = _load(checkpoint) + save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + from types import SimpleNamespace + + from collector.fpm_forward import cli as collector_cli + from collector.fpm_forward import entry + + from aisimulate.support import collection_readiness, fpm + + commands = [] + + def resolve(command): + commands.append(command[3:]) + return collector_cli._parser().parse_args(command[3:]), (SimpleNamespace(cells=()), {}) + + monkeypatch.setattr(fpm, "_resolve_execution", resolve) + monkeypatch.setattr(collection_readiness, "assess_readiness", lambda *a, **k: {"ready_for_full_collection": True}) + + def fake_collect(argv): + assert "--fpm-runtime-instrumentation" in argv + manifest = runtime_probe_manifest(request) + target = root / "observed" + shutil.copytree(Path(manifest["observations_index"]).parent, target) + (root / "fpm-checkpoint").mkdir(exist_ok=True) + bindings = _formal_bindings(target / "observations.json") + metadata = root / "synthetic-publication.metadata.json" + _write(metadata, {"backend_version": request.identity.framework_version}) + bindings["database"] = {"status": "passed", "metadata": str(metadata)} + _write( + root / "fpm-checkpoint" / "fpm_forward.json", + {"runtime_observations": str(target / "observations.json"), **bindings}, + ) + return 0 + + monkeypatch.setattr(entry, "run_resolved", lambda *_: fake_collect(commands[-1]) or []) + assert cli.main(execute) == 0 + comparison = json.loads((root / "runtime-compatibility.json").read_text()) + assert comparison["compatibility"]["status"] == "compatible" + assert comparison["compatibility"]["accepted_capacity_bytes"] == 948 * 128 + + +def _formal_bindings(index, cells=None): + """Synthetic formal contexts bound to distinct native cell attempts.""" + document = json.loads(index.read_text()) + entry = document["configurations"]["tp2"] + entry["active_attempt_id"] = "formal-parent" + attempt = entry["attempts"][0] + attempt["attempt_id"] = "formal-parent" + entries = {} + for phase, result in attempt["phases"].items(): + cell_id = cells[phase] if cells else f"cell-{phase}" + attempt_id = f"native-{phase}" + result.update(cell_id=cell_id, collector_attempt_id=attempt_id) + entries[cell_id] = {"status": "passed", "attempt_id": attempt_id} + for ref in [result["launch_manifest"], *result["artifacts"]]: + path = index.parent / ref["path"] + value = json.loads(path.read_text()) + value.update(attempt_id="formal-parent", cell_id=cell_id, collector_attempt_id=attempt_id) + _write(path, value) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + return {"runtime_observation_attempt_id": "formal-parent", "cells": entries} + + +def _completed_runtime_collection(tmp_path, capsys, *, tp=2, worker_type=None): + from collector.fpm_forward.cli import _parser + from collector.fpm_forward.config import FPMCollectionOptions + from collector.fpm_forward.database import aggregate_cell, write_formal_database + from collector.fpm_forward.entry import _load_generator_overrides + from collector.fpm_forward.planner import build_collection_plan + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + from aisimulate.support.runtime import runtime_probe_manifest + + from .test_onboard_finalization import _native + + checkpoint, _, request, files = _imported(tmp_path, capsys, capacity_multiplier=10, tp=tp, worker_type=worker_type) + state = _load(checkpoint) + save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + root = tmp_path / "collection" + assert cli.main(["onboard", "plan", "--config", files["request"], "--output-dir", str(root)]) == 0 + capsys.readouterr() + args = _parser().parse_args(json.loads((root / "commands.json").read_text())["fpm_plan_local"][3:]) + manifest = runtime_probe_manifest(request) + snapshot = Path(manifest["observations_index"]) + launch = json.loads(snapshot.read_text())["configurations"]["tp2"]["launch"] + plan = build_collection_plan( + backend="vllm", + model_path=request.identity.model, + model_architecture=request.fpm_profile.architecture, + model_config_path=launch["model_config"]["path"], + system=request.identity.gpu, + selected_ops={"attention_context", "attention_generation"}, + options=FPMCollectionOptions.from_args(args), + fpm_profile=request.fpm_profile, + generator_overrides=_load_generator_overrides(args), + runtime_instrumentation=args.fpm_runtime_instrumentation, + runtime_launch=launch, + runtime_configuration="tp2", + ) + artifact_root = root / "fpm-artifacts" / plan.sha256[:16] + artifact_root.mkdir(parents=True) + _write(artifact_root / "collection-plan.json", plan.to_dict()) + observed = artifact_root / "observations" + shutil.copytree(snapshot.parent, observed) + bindings = _formal_bindings( + observed / "observations.json", {cell.workload_kind: cell.cell_id for cell in plan.cells} + ) + rows = [] + for cell in plan.cells: + directory = artifact_root / "cells" / cell.cell_id + raw = directory / "raw/pod-0" + raw.mkdir(parents=True) + _write(directory / "cell.json", cell.to_dict()) + provenance = { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell.cell_id, + "plan_sha256": plan.sha256, + "attempt_id": f"native-{cell.workload_kind}", + "runtime": {"backend": "vllm", "backend_version": "0.28.0"}, + } + _write(raw / "collector-provenance.json", provenance) + _write(raw / "benchmark.json", _native(cell.workload_kind)) + rows.extend(aggregate_cell(plan, cell, directory, expected_attempt_id=provenance["attempt_id"])) + parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=root / "systems/data") + assert not skipped + (root / "fpm-checkpoint").mkdir(exist_ok=True) + _write( + root / "fpm-checkpoint/fpm_forward.json", + { + "schema": CHECKPOINT_SCHEMA, + "plan_sha256": plan.sha256, + **bindings, + "runtime_observations": str(observed / "observations.json"), + "database": { + "status": "passed", + "parquet": str(parquet), + "metadata": str(metadata), + "published_cells": len(plan.cells), + "plan_cells": len(plan.cells), + "missing_cells": [], + "skipped_first_publisher_wins": [], + }, + }, + ) + return checkpoint, request, files, root + + +def test_formal_finalize_validates_new_observations_and_preserves_verified_timing_pair(tmp_path, capsys): + from aisimulate.support.finalization import finalization_manifest + + checkpoint, request, files, root = _completed_runtime_collection(tmp_path, capsys) + target = tmp_path / "finalized" + assert ( + cli.main( + [ + "onboard", + "finalize", + "--config", + files["request"], + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + == 0 + ) + capsys.readouterr() + resolved = SupportRequest.from_yaml(target / "request.yaml") + evidence = finalization_manifest(resolved) + assert evidence["runtime_compatibility"]["status"] == "compatible" + assert len(evidence["formal_data"]) == 2 + assert list((target / "systems/data").rglob("fpm_forward_perf.parquet")) + assert _load(checkpoint).configurations["tp2"].acceptance is not None + # Finalization produces another draft. Replacing the imported request keeps + # the accepted probe in history and does not treat the new profile as accepted. + state = _load(checkpoint) + artifacts = {key: {"archived": True} for key in state.configurations["tp2"].artifacts} + artifacts["final-request"] = {"path": str(target / "request.yaml"), "kind": "request"} + state, _ = save_checkpoint( + checkpoint, + patch={ + "configurations": { + "tp2": {"draft_request": resolved.model_dump(mode="json", exclude_none=True), "artifacts": artifacts} + } + }, + expected_revision=state.revision, + accept=[], + ) + assert state.configurations["tp2"].acceptance is None + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + + +@pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) +def test_serving_role_round_trip_import_collection_and_finalization(tmp_path, capsys, worker_type): + import pyarrow.parquet as pq + from collector.fpm_forward.repeatability import load_repeatability_source + + from aisimulate.support.finalization import finalization_manifest, finalize + from aisimulate.support.plan import check_plan + from aisimulate.support.runtime import verify_runtime_profile + + checkpoint, request, _, root = _completed_runtime_collection(tmp_path, capsys, worker_type=worker_type) + assert request.worker_type == request.profile_deployment().worker_type == worker_type + assert report(_load(checkpoint), checkpoint)["configurations"]["tp2"]["profile_accepted"] + plan_path = next((root / "fpm-artifacts").glob("*/collection-plan.json")) + frozen = load_repeatability_source(plan_path.parent) + assert frozen.options.worker_type == worker_type + assert {cell.workload_kind for cell in frozen.cells} == set(request.collection_phases) + target = tmp_path / "finalized" + finalize(request, root, target) + resolved = SupportRequest.from_yaml(target / "request.yaml") + assert resolved.worker_type == resolved.profile_deployment().worker_type == worker_type + assert resolved.collection.cudagraph_mode == request.collection.cudagraph_mode + assert ( + resolved.profile_deployment().resources.runtime_memory.kv_cache_bytes + == request.profile_deployment().resources.runtime_memory.kv_cache_bytes + ) + evidence = finalization_manifest(resolved) + assert evidence["runtime_compatibility"]["status"] == "compatible" + assert len(evidence["formal_data"]) == 2 + parquet = next((target / "systems/data").rglob("fpm_forward_perf.parquet")) + assert set(pq.read_table(parquet).column("workload_kind").to_pylist()) == set(request.collection_phases) + verify_runtime_profile(resolved) + check_plan(resolved, target) + + +def _reviewed_capacity_revision(tmp_path, capsys): + checkpoint, original, files, root = _completed_runtime_collection(tmp_path, capsys) + collection = json.loads((root / "fpm-checkpoint/fpm_forward.json").read_text()) + index = Path(collection["runtime_observations"]) + document = json.loads(index.read_text()) + for phase in document["configurations"]["tp2"]["attempts"][0]["phases"].values(): + for ref in phase["artifacts"]: + path = index.parent / ref["path"] + record = json.loads(path.read_text()) + if record["kind"] == "scheduler": + record["cache"]["initial_free_blocks"] -= 1 + record["cache"]["reserved_blocks"] += 1 + record["cache"]["permanent_reserved_block_ids"].append(2) + _write(path, record) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + assert _import(checkpoint, index, tmp_path / "revised", configurations=["tp2"]) == 0 + revised = Path(json.loads(capsys.readouterr().out)["configurations"]["tp2"]["request"]) + state = _load(checkpoint) + save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + return checkpoint, original, files, root, revised + + +def test_finalization_reuses_original_collection_with_accepted_observed_capacity_revision(tmp_path, capsys): + from aisimulate.support.finalization import finalization_manifest, finalize + from aisimulate.support.plan import check_plan, request_id + from aisimulate.support.runtime import verify_runtime_profile + + checkpoint, original, files, root, memory_config = _reviewed_capacity_revision(tmp_path, capsys) + revised = SupportRequest.from_yaml(memory_config) + before = {path: path.read_bytes() for path in root.rglob("*") if path.is_file()} + accepted = checkpoint.read_bytes() + with pytest.raises(ValueError, match="smaller than the accepted bound"): + finalize(original, root, tmp_path / "unreviewed") + with pytest.raises(ValueError, match="different request identity"): + check_plan(revised, root) + target = tmp_path / "resolved-revision" + assert ( + cli.main( + [ + "onboard", + "finalize", + "--config", + files["request"], + "--output-dir", + str(root), + "--memory-config", + str(memory_config), + "--resolved-output-dir", + str(target), + ] + ) + == 0 + ) + capsys.readouterr() + resolved = SupportRequest.from_yaml(target / "request.yaml") + manifest = finalization_manifest(resolved) + assert manifest["source_request_id"] == request_id(original) + assert manifest["memory_revision"]["request_id"] == request_id(revised) + assert resolved.profile_deployment().resources.runtime_memory.kv_cache_bytes == ( + revised.profile_deployment().resources.runtime_memory.kv_cache_bytes + ) + assert before == {path: path.read_bytes() for path in before} + assert checkpoint.read_bytes() == accepted + check_plan(resolved, target) + verify_runtime_profile(resolved) + state = _load(checkpoint) + artifacts = {key: {"archived": True} for key in state.configurations["tp2"].artifacts} + artifacts["resolved-request"] = {"path": str(target / "request.yaml"), "kind": "request"} + state, _ = save_checkpoint( + checkpoint, + patch={ + "configurations": { + "tp2": { + "draft_request": resolved.model_dump(mode="json", exclude_none=True), + "artifacts": artifacts, + } + } + }, + expected_revision=state.revision, + accept=[], + ) + assert state.configurations["tp2"].acceptance is None + state, _ = save_checkpoint(checkpoint, patch={}, expected_revision=state.revision, accept=["tp2"]) + assert report(state, checkpoint)["configurations"]["tp2"]["profile_accepted"] + memory_config.write_text(memory_config.read_text() + "\n") + with pytest.raises(ValueError, match="memory revision.*changed"): + verify_runtime_profile(resolved) + + +@pytest.mark.parametrize( + ("change", "diagnostic"), + [ + ("runtime", "non-capacity"), + ("layout", "non-capacity"), + ("workload", "non-capacity"), + ("capacity_increase", "must lower"), + ("capacity_override", "without capacity overrides"), + ("forged_observation", "saved observed resources differ"), + ("missing_attempts", "complete formal collection attempt bindings"), + ("wrong_attempt", "different collection attempt"), + ], +) +def test_capacity_revision_rejects_unrelated_changes_and_unverified_evidence(tmp_path, capsys, change, diagnostic): + from aisimulate.support.runtime import verify_collection_runtime + + _, original, _, root, memory_config = _reviewed_capacity_revision(tmp_path, capsys) + revised = SupportRequest.from_yaml(memory_config) + resources = revised.profile_deployment().resources + evidence = json.loads(resources.runtime_memory.provenance) + checkpoint = json.loads((root / "fpm-checkpoint/fpm_forward.json").read_text()) + index = Path(checkpoint["runtime_observations"]) + if change == "runtime": + revised.collection.gpu_memory_utilization = 0.8 + elif change == "layout": + resources.cache_groups[0].page_size_bytes += 128 + elif change == "workload": + revised.workload.request_count += 1 + elif change == "capacity_increase": + resources.runtime_memory.kv_cache_bytes = ( + original.profile_deployment().resources.runtime_memory.kv_cache_bytes + 128 + ) + elif change == "capacity_override": + resources.runtime_memory.kv_cache_bytes -= 128 + evidence["user_overrides"] = { + "runtime_memory.kv_cache_bytes": {"value": resources.runtime_memory.kv_cache_bytes} + } + elif change == "forged_observation": + evidence["observed_resources"]["runtime_memory"]["kv_cache_bytes"] -= 128 + elif change == "missing_attempts": + checkpoint = None + elif change == "wrong_attempt": + checkpoint["runtime_observation_attempt_id"] = "unrelated" + resources.runtime_memory.provenance = json.dumps(evidence) + with pytest.raises(ValueError, match=diagnostic): + verify_collection_runtime(original, index, collection_checkpoint=checkpoint, memory_request=revised) + + +def test_capacity_revision_requires_acceptance_and_exact_formal_observations(tmp_path, capsys): + from aisimulate.support.finalization import finalize + from aisimulate.support.runtime import verify_collection_runtime + + checkpoint, original, _, root, memory_config = _reviewed_capacity_revision(tmp_path, capsys) + formal = json.loads((root / "fpm-checkpoint/fpm_forward.json").read_text()) + source = Path(formal["runtime_observations"]) + unrelated = tmp_path / "other-attempt" + shutil.copytree(source.parent, unrelated) + index = unrelated / source.name + document = json.loads(index.read_text()) + attempt = document["configurations"]["tp2"]["attempts"][0] + attempt["attempt_id"] = "unrelated-parent" + document["configurations"]["tp2"]["active_attempt_id"] = attempt["attempt_id"] + for phase in attempt["phases"].values(): + for ref in [phase["launch_manifest"], *phase["artifacts"]]: + path = index.parent / ref["path"] + record = json.loads(path.read_text()) + record["attempt_id"] = attempt["attempt_id"] + _write(path, record) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + assert _import(checkpoint, index, tmp_path / "unrelated-revision", configurations=["tp2"]) == 0 + revised_path = json.loads(capsys.readouterr().out)["configurations"]["tp2"]["request"] + revised = SupportRequest.from_yaml(revised_path) + with pytest.raises(ValueError, match="explicit acceptance"): + finalize(original, root, tmp_path / "unaccepted", memory_config=revised_path) + assert not (tmp_path / "unaccepted").exists() + with pytest.raises(ValueError, match="same complete formal collection observations"): + verify_collection_runtime(original, source, collection_checkpoint=formal, memory_request=revised) + # The previously reviewed request is no longer the checkpoint's accepted + # draft, even though its immutable evidence is still available. + with pytest.raises(ValueError, match="explicit acceptance"): + finalize(original, root, tmp_path / "superseded", memory_config=memory_config) + + +def test_large_runtime_provenance_reaches_supervised_recommendation_and_keeps_sources(tmp_path, capsys): + from aisimulate.support.finalization import _merge_resources, _verify_collection, finalization_manifest + from aisimulate.support.runtime import verify_runtime_profile + + # The fixture's arbitrary configuration label is independent of topology. + checkpoint, original, files, root = _completed_runtime_collection(tmp_path, capsys, tp=8) + before = {path: path.read_bytes() for path in root.rglob("*") if path.is_file()} + accepted = checkpoint.read_bytes() + target = tmp_path / "resolved" + assert ( + cli.main( + [ + "onboard", + "finalize", + "--config", + files["request"], + "--output-dir", + str(root), + "--resolved-output-dir", + str(target), + ] + ) + == 0 + ) + capsys.readouterr() + assert before == {path: path.read_bytes() for path in before} + assert checkpoint.read_bytes() == accepted + resolved = SupportRequest.from_yaml(target / "request.yaml") + assert resolved.profile_deployment().tp == 8 + manifest = finalization_manifest(resolved) + assert manifest == json.loads((target / "finalization.json").read_text()) + assert manifest["observation_provenance"] == "source_references" + observation = json.loads(manifest["cell_observations"][0]["runtime_memory"]["provenance"]) + assert len(observation["rank_capacities"]) == 16 + assert len(observation["artifacts"]) == 18 + assert all(set(reference) == {"path", "sha256"} for reference in observation["artifacts"]) + for reference in manifest["source_artifacts"]: + assert hashlib.sha256(Path(reference["path"]).read_bytes()).hexdigest() == reference["sha256"] + output = tmp_path / "recommendation" + result = subprocess.run( + [ + sys.executable, + "-m", + "aisimulate", + "recommend", + "--config", + str(target / "recommend/pilot.yaml"), + "--output-dir", + str(output), + "--format", + "json", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 1, result.stderr + recommendation = json.loads((output / "recommendation.json").read_text()) + assert len(recommendation["candidates"]) == 1 + candidate = recommendation["candidates"][0] + assert candidate["status"] == "failed" + assert "FPM direct" in candidate["reason"] and "unsupported" in candidate["reason"] + assert "bounded checkpoint budget" not in result.stderr + events = [json.loads(line) for line in (output / "execution-events.jsonl").read_text().splitlines()] + assert any(event["event"] == "candidate_completed" for event in events) + verify_runtime_profile(resolved) + observations, legacy_manifest, _, _ = _verify_collection(original, root) + legacy = resolved.model_dump(mode="json") + legacy["fpm_profile"]["deployments"][0]["resources"] = _merge_resources(observations, legacy_manifest) + review = tmp_path / "resolved-checkpoint.json" + state, _ = save_checkpoint( + review, + patch={ + "configurations": { + "legacy": {"draft_request": legacy}, + "references": {"draft_request": resolved.model_dump(mode="json")}, + } + }, + expected_revision=None, + accept=["legacy", "references"], + ) + assert all(config["profile_accepted"] for config in report(state, review)["configurations"].values()) + assert SupportRequest.model_validate(state.configurations["legacy"].draft_request).model_dump(mode="json") == legacy + changed = resolved.model_dump(mode="json") + resources = changed["fpm_profile"]["deployments"][0]["resources"] + forged = copy.deepcopy(manifest) + forged["source_artifacts"][0]["sha256"] = "0" * 64 + resources["runtime_memory"]["provenance"] = json.dumps(forged, sort_keys=True, separators=(",", ":")) + with pytest.raises(ValueError, match="resources differ"): + verify_runtime_profile(SupportRequest.model_validate(changed)) + for value in (None, "future-format", {}): + forged["observation_provenance"] = value + resources["runtime_memory"]["provenance"] = json.dumps(forged) + with pytest.raises(ValueError, match="unsupported finalized observation provenance format"): + verify_runtime_profile(SupportRequest.model_validate(changed)) + # The compact representation must still bind every immutable raw record. + source = next(path for path in before if path.name == "worker-0-0.json") + source.write_text(source.read_text() + "\n") + with pytest.raises(ValueError, match="changed|hash|incomplete"): + verify_runtime_profile(resolved) + assert all(not config["profile_accepted"] for config in report(state, review)["configurations"].values()) + + +def test_formal_capacity_drop_and_geometry_change_are_rejected(tmp_path, capsys): + from aisimulate.support.runtime import verify_collection_runtime + + _, index, request, _ = _imported(tmp_path, capsys) + result = verify_collection_runtime(request, index) + assert result["compatibility"]["status"] == "compatible" + document = json.loads(index.read_text()) + for phase in document["configurations"]["tp2"]["attempts"][0]["phases"].values(): + for ref in phase["artifacts"]: + path = index.parent / ref["path"] + record = json.loads(path.read_text()) + if record["kind"] == "scheduler": + record["cache"]["initial_free_blocks"] -= 1 + record["cache"]["reserved_blocks"] += 1 + record["cache"]["permanent_reserved_block_ids"].append(2) + _write(path, record) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + with pytest.raises(ValueError, match="smaller than the accepted bound"): + verify_collection_runtime(request, index) + for phase in document["configurations"]["tp2"]["attempts"][0]["phases"].values(): + for ref in phase["artifacts"]: + path = index.parent / ref["path"] + record = json.loads(path.read_text()) + cache = record["cache"] + if record["kind"] == "scheduler": + cache["initial_free_blocks"] += 1 + cache["reserved_blocks"] -= 1 + cache["permanent_reserved_block_ids"].remove(2) + else: + cache["layer_tensors"]["changed_layer"] = cache["layer_tensors"].pop("layer0") + for allocation in cache["tensor_allocations"]: + allocation["shared_by"] = [ + "changed_layer" if name == "layer0" else name for name in allocation["shared_by"] + ] + for group in cache["groups"]: + group["layer_names"] = ["changed_layer" if name == "layer0" else name for name in group["layer_names"]] + if "layer_classes" in group and "layer0" in group["layer_classes"]: + group["layer_classes"]["changed_layer"] = group["layer_classes"].pop("layer0") + _write(path, record) + ref["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() + _write(index, document) + with pytest.raises(ValueError, match="geometry differs"): + verify_collection_runtime(request, index) + + +def test_malformed_configuration_does_not_discard_an_independent_valid_draft(tmp_path, capsys): + checkpoint, index, _ = _campaign(tmp_path) + document = json.loads(index.read_text()) + document["configurations"]["incomplete"] = {"attempts": "malformed"} + _write(index, document) + assert _import(checkpoint, index, tmp_path / "drafts") == 1 + result = json.loads(capsys.readouterr().out) + assert result["configurations"]["tp2"]["status"] == "complete" + assert "list of objects" in str(result["configurations"]["incomplete"]["diagnostics"]) + + +def test_import_optimistic_lock_keeps_concurrent_checkpoint_update(tmp_path, capsys, monkeypatch): + import aisimulate.support.runtime as runtime + + checkpoint, index, _ = _campaign(tmp_path) + actual_draft = runtime._draft + + def concurrent_edit(*args, **kwargs): + result = actual_draft(*args, **kwargs) + state = _load(checkpoint) + save_checkpoint( + checkpoint, patch={"research": {"concurrent": "preserved"}}, expected_revision=state.revision, accept=[] + ) + return result + + monkeypatch.setattr(runtime, "_draft", concurrent_edit) + with pytest.raises(SystemExit, match="2"): + _import(checkpoint, index, tmp_path / "drafts", configurations=["tp2"]) + assert "checkpoint revision" in capsys.readouterr().err + assert _load(checkpoint).research["concurrent"] == "preserved" + assert not _load(checkpoint).configurations["tp2"].draft_request.get("fpm_profile") + + +def test_collection_inputs_without_probe_do_not_reverify_finalized_quality(monkeypatch): + from aisimulate.support import runtime + + monkeypatch.setattr(runtime, "runtime_probe_manifest", lambda request: None) + monkeypatch.setattr(runtime, "verify_runtime_profile", lambda request: pytest.fail("unnecessary full verification")) + assert runtime.runtime_collection_inputs(None, None) == ([], None) diff --git a/python/aisimulate/tests/unit/test_onboard_topology.py b/python/aisimulate/tests/unit/test_onboard_topology.py new file mode 100644 index 000000000..7a52320d4 --- /dev/null +++ b/python/aisimulate/tests/unit/test_onboard_topology.py @@ -0,0 +1,1278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Topology suggestions through onboarding's actual CLI and saved-plan consumers.""" + +from __future__ import annotations + +import builtins +import hashlib +import json +import shlex +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml + +import aisimulate.main as cli +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.fpm_profile import load_fpm_profile +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + +# Original synthetic decoder geometries, not downloaded checkpoints or measurements. +_SMALL = { + "_name_or_path": "example/synthetic-decoder", + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "vocab_size": 64, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", +} +_LARGE = { + **_SMALL, + "hidden_size": 8192, + "intermediate_size": 28672, + "num_hidden_layers": 80, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 32768, +} +_MOE = {**_SMALL, "architectures": ["MixtralForCausalLM"], "model_type": "mixtral", "num_local_experts": 4} +_PRECISION = {"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"} +_OPTIONS = { + "model_revision": "synthetic-checkpoint-123", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + "tokenizer_revision": "synthetic-tokenizer-456", + "chat_template_revision": "synthetic-template-789", + "input_tokens": 1024, + "output_tokens": 128, + "concurrency": 2, + "request_count": 3, + "context_length": 4096, + "ttft_ms": 1000, + "tpot_ms": 100, + "seed": 123, +} + + +def _inputs(tmp_path, config=None, overrides=None): + source = tmp_path / "config.json" + source.write_text(json.dumps(_SMALL if config is None else config)) + resources = tmp_path / "resources.yaml" + resources.write_text(yaml.safe_dump(_PRECISION if overrides is None else overrides)) + return source, resources + + +def _args(source, resources, output, **changes): + options = {**_OPTIONS, **changes} + command = ["onboard", "init", "--model-config", str(source), "--output", str(output)] + if resources is not None: + command += ["--resource-overrides", str(resources)] + return command + [ + part + for name, value in options.items() + if value is not None + for part in ("--" + name.replace("_", "-"), str(value)) + ] + + +def _real_cli(args): + return subprocess.run( + [sys.executable, "-m", "aisimulate", *args], input="", text=True, capture_output=True, timeout=30 + ) + + +def _terminal(monkeypatch, output, answers): + remaining = iter(answers) + prompts = [] + original = output.read_bytes() if output.exists() else None + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def respond(prompt): + prompts.append(prompt) + if original is None: + assert not output.parent.exists(), "interactive setup wrote before acceptance" + else: + assert output.read_bytes() == original, "interactive setup replaced output before acceptance" + try: + answer = next(remaining) + except StopIteration: + pytest.fail(f"unexpected prompt: {prompt}") + if isinstance(answer, BaseException): + raise answer + return answer + + monkeypatch.setattr(builtins, "input", respond) + return prompts + + +@pytest.mark.parametrize("config,gpu,width", [(_SMALL, "h200_sxm", 1), (_LARGE, "h200_sxm", 2), (_LARGE, "gb200", 1)]) +def test_real_cli_default_uses_model_geometry_and_hardware_then_embeds_exact_profile(tmp_path, config, gpu, width): + source, resources = _inputs(tmp_path, config) + output = tmp_path / "new" / "request.yaml" + result = _real_cli(_args(source, resources, output, gpu=gpu)) + assert result.returncode == 0, result.stderr + request = SupportRequest.from_yaml(output) + assert request.worker_gpus == width + assert request.profile_deployment().parallel_tuple == (width, 1, 1, 1, 1, 1) + assert request.identity.model_revision == _OPTIONS["model_revision"] + assert request.identity.tokenizer_revision == _OPTIONS["tokenizer_revision"] + assert request.identity.chat_template_revision == _OPTIONS["chat_template_revision"] + assert request.workload.concurrency == 2 + assert request.workload.request_count == 3 + assert request.search.seed == 123 + assert request.search.context_length == 4096 + assert f"Selected tp with {width} collection GPUs" in result.stdout + assert "estimated_fit" in result.stdout + assert "not a performance ranking" in result.stdout + assert not {"gpu_count", "node_count", "gpus_per_node"} & request.identity.model_dump().keys() + assert "max_candidates" not in request.search.model_dump() + assert list(output.parent.iterdir()) == [output] + + +@pytest.mark.parametrize("role", [None, "aggregated", "prefill", "decode"]) +@pytest.mark.parametrize("preview", [False, True]) +def test_real_cli_grouped_topology_uses_native_role_estimate_and_preserves_saved_role(tmp_path, role, preview): + source, resources = _inputs( + tmp_path, + {**_SMALL, "sliding_window": 128}, + {**_PRECISION, "cache_block_sizes": {"sliding_attention": 16}, "max_num_tokens": 256, "max_batch_size": 8}, + ) + output = tmp_path / "new" / "request.yaml" + command = _args(source, resources, output, model=str(tmp_path), worker_type=role) + result = _real_cli(command + (["--suggest-parallel"] if preview else [])) + assert result.returncode == 0, result.stderr + if preview: + report = json.loads(result.stdout) + candidate = report["default"] + assert candidate["status"] == "estimated_fit" + assert candidate["required_gpus"] == 1 + fields = candidate["resolved_profile_fields"] + assert fields["cache_layout"] == "grouped" + non_kv = sum( + fields[name] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes") + ) + assert non_kv < candidate["estimated_required_bytes"] < candidate["memory_budget_bytes"] + assert not output.parent.exists() + return + + request = SupportRequest.from_yaml(output) + assert request.worker_type == request.profile_deployment().worker_type == role + assert request.worker_gpus == 1 + assert request.profile_deployment().resources.cache_layout == "grouped" + assert request.profile_deployment().resources.memory_source == "pending" + plan_root = tmp_path / "plan" + result = _real_cli(["onboard", "plan", "--config", str(output), "--output-dir", str(plan_root)]) + assert result.returncode == 0, result.stderr + assert SupportRequest.from_yaml(plan_root / "request.yaml") == request + assert load_fpm_profile((plan_root / "fpm-model-profile.json").read_text()) == request.fpm_profile + plan = json.loads((plan_root / "support-plan.json").read_text()) + command = plan["fpm"]["plan_command"] + if role is not None: + assert plan["fpm"]["worker_type"] == role + assert plan["fpm"]["collection_phases"] == list(request.collection_phases) + assert command[command.index("--fpm-worker-type") + 1] == role + else: + assert "worker_type" not in request.model_dump(mode="json") + assert "worker_type" not in request.profile_deployment().model_dump(mode="json") + assert "--fpm-worker-type" not in command + if role in {"prefill", "decode"}: + artifact = yaml.safe_load((plan_root / "worker.yaml").read_text()) + assert artifact["worker_type"] == role + assert set(artifact["engine"]["workers"]) == {role} + assert artifact["engine"]["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert not (plan_root / "predict/pilot.yaml").exists() + else: + prediction = CorePredictionConfig.from_yaml(plan_root / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan_root / "recommend/pilot.yaml") + assert prediction.engine.mode == recommendation.engine.mode == "aggregated" + assert prediction.engine.fpm_profile == recommendation.engine.fpm_profile == request.fpm_profile + assert not (plan_root / "worker.yaml").exists() + + +@pytest.mark.parametrize("output_kind", ["existing", "directory", "missing", "file_parent"]) +def test_real_cli_preview_is_json_and_never_validates_or_writes_output_target(tmp_path, output_kind): + source, resources = _inputs(tmp_path, _LARGE) + existing = tmp_path / "existing-request.yaml" + existing.write_text("preserve previous request\n") + targets = { + "existing": existing, + "directory": tmp_path, + "missing": tmp_path / "new" / "request.yaml", + "file_parent": existing / "request.yaml", + } + before = {path.relative_to(tmp_path): path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} + result = _real_cli( + _args(source, resources, targets[output_kind], model="deployment/actual-model") + ["--suggest-parallel"] + ) + assert result.returncode == 0, result.stderr + report = json.loads(result.stdout) + assert report["default"]["topology"]["tensor_parallel"] == 2 + assert report["identity"]["model"] == "deployment/actual-model" + assert report["identity"]["model_revision"] == _OPTIONS["model_revision"] + assert report["identity"]["framework_version"] == _OPTIONS["framework_version"] + assert "workload" not in report + assert report["collection"]["max_sequences"] == 256 + assert report["context_length"] == 4096 + assert report["config_sha256"] == hashlib.sha256(source.read_bytes()).hexdigest() + assert {path.relative_to(tmp_path): path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} == before + assert not (tmp_path / "new").exists() + + +@pytest.mark.parametrize("gpu,domain,width", [("h200_sxm", "node", 8), ("gb200", "rack", 72)]) +def test_real_cli_preview_exposes_architecture_domain_and_moe_alternatives(tmp_path, gpu, domain, width): + source, resources = _inputs(tmp_path, _MOE) + output = tmp_path / "new" / "request.yaml" + result = _real_cli(_args(source, resources, output, gpu=gpu) + ["--suggest-parallel"]) + assert result.returncode == 0, result.stderr + report = json.loads(result.stdout) + assert report["hardware"]["fast_domain"] == domain + assert report["hardware"]["fast_domain_gpus"] == width + assert {candidate["family"] for candidate in report["candidates"]} == {"pure_tp", "dep", "tep"} + assert len(report["candidates"]) == 6 + for candidate in report["candidates"]: + if candidate["family"] in {"dep", "tep"}: + assert candidate["status"] == "needs_inputs" + assert "comm_overhead_bytes" not in candidate["missing"] + assert "memory" in " ".join(candidate["reasons"]).lower() + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "config,resources,missing", + [ + (_SMALL, False, "fmha_quant_mode"), + ], +) +def test_real_cli_missing_facts_preview_has_no_default_and_normal_init_fails_before_write( + tmp_path, config, resources, missing +): + source, overrides = _inputs(tmp_path, config) + output = tmp_path / "new" / "request.yaml" + command = _args(source, overrides if resources else None, output) + preview = _real_cli([*command, "--suggest-parallel"]) + assert preview.returncode == 0, preview.stderr + report = json.loads(preview.stdout) + assert report["default"] is None + assert all(item["status"] == "needs_inputs" and missing in item["missing"] for item in report["candidates"]) + result = _real_cli(command) + assert result.returncode == 2 + assert "No fully assessed topology default" in result.stderr + assert missing in result.stderr + assert "--tensor-parallel" in result.stderr + assert "--resource-overrides" in result.stderr + assert "Selected" not in result.stdout + assert not output.parent.exists() + + +@pytest.mark.parametrize("preview", [False, True]) +def test_flat_rank_bounds_require_an_explicit_topology_without_transferring_bytes(tmp_path, preview): + source, resources = _inputs(tmp_path, overrides={**_PRECISION, "weights_bytes": 9 * 1024**3}) + output = tmp_path / "new" / "request.yaml" + command = _args(source, resources, output) + result = _real_cli(command + (["--suggest-parallel"] if preview else [])) + assert result.returncode == 2 + assert "rank-local byte overrides" in result.stderr + assert "weights_bytes" in result.stderr + assert "exact topology" in result.stderr + assert not output.parent.exists() + explicit = _real_cli(command + ["--tensor-parallel", "2"]) + assert explicit.returncode == 0, explicit.stderr + request = SupportRequest.from_yaml(output) + assert request.worker_gpus == 2 + assert request.profile_deployment().resources.weights_bytes is None + assert ( + json.loads(request.profile_deployment().resources.provenance)["fields"]["weights_bytes"]["value"] == 9 * 1024**3 + ) + assert "Suggested worker topologies" not in explicit.stdout + + +@pytest.mark.parametrize( + "option", ["--tensor-parallel", "--attention-data-parallel", "--moe-tensor-parallel", "--moe-expert-parallel"] +) +def test_any_explicit_topology_option_preserves_existing_route(tmp_path, option): + source, resources = _inputs(tmp_path, overrides={**_PRECISION, "weights_bytes": 2 * 1024**3}) + output = tmp_path / "request.yaml" + result = _real_cli(_args(source, resources, output) + [option, "1"]) + assert result.returncode == 0, result.stderr + profile = SupportRequest.from_yaml(output).profile_deployment().resources + assert profile.memory_source == "pending" + assert json.loads(profile.provenance)["fields"]["weights_bytes"]["value"] == 2 * 1024**3 + assert "Suggested worker topologies" not in result.stdout + + +def test_explicit_topology_may_exceed_automatic_hardware_domain(tmp_path): + source, resources = _inputs( + tmp_path, + _LARGE, + {**_PRECISION, "activations_bytes": 70 * 1024**2, "comm_overhead_bytes": 256 * 1024**2}, + ) + output = tmp_path / "request.yaml" + result = _real_cli(_args(source, resources, output, tensor_parallel=16)) + assert result.returncode == 0, result.stderr + request = SupportRequest.from_yaml(output) + assert request.worker_gpus == request.profile_deployment().tp == 16 + assert request.profile_deployment().resources.comm_overhead_bytes is None + assert ( + json.loads(request.profile_deployment().resources.provenance)["fields"]["comm_overhead_bytes"]["value"] + == 256 * 1024**2 + ) + assert "Suggested worker topologies" not in result.stdout + + +def test_no_candidate_in_declared_domain_fails_without_a_tp1_fallback(tmp_path): + source, resources = _inputs(tmp_path, _LARGE) + output = tmp_path / "new" / "request.yaml" + result = _real_cli(_args(source, resources, output, interconnect="none")) + assert result.returncode == 2 + assert "No eligible topology" in result.stderr + assert "Rejected tp, 1 GPUs" in result.stderr + assert "lower bound" in result.stderr + assert "Selected" not in result.stdout + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "extra,expected", + [ + (["--interactive"], "cannot be combined"), + (["--profile", "onboarding"], "cannot be combined"), + (["--tensor-parallel", "1"], "cannot be combined"), + (["--attention-data-parallel", "1"], "cannot be combined"), + (["--moe-tensor-parallel", "1"], "cannot be combined"), + (["--moe-expert-parallel", "1"], "cannot be combined"), + (["--fpm-profile", "missing.yaml"], "not allowed"), + ], +) +def test_preview_rejects_incompatible_options_before_loading_files(tmp_path, extra, expected): + output = tmp_path / "request.yaml" + output.write_text("keep existing\n") + result = _real_cli(_args(tmp_path / "absent.json", None, output) + ["--suggest-parallel", "--overwrite", *extra]) + assert result.returncode == 2 + assert expected in result.stderr + assert output.read_text() == "keep existing\n" + assert list(tmp_path.iterdir()) == [output] + + +def test_preview_requires_a_model_config_and_real_target_identity(tmp_path): + output = tmp_path / "new" / "request.yaml" + result = _real_cli(["onboard", "init", "--suggest-parallel", "--output", str(output)]) + assert result.returncode == 2 + assert "requires --model-config" in result.stderr + source, resources = _inputs(tmp_path) + result = _real_cli( + _args(source, resources, output, model_revision=None, framework_version=None) + ["--suggest-parallel"] + ) + assert result.returncode == 2 + assert "--model-revision" in result.stderr + assert "--framework-version" in result.stderr + assert not result.stdout + assert not output.parent.exists() + + +@pytest.mark.parametrize( + "updates,extra,expected", + [ + ({}, ["--context-length", "8192"], "context_length"), + ({"weights_bytes": True}, [], "weights_bytes"), + ({"unknown_field": 3}, [], "unknown"), + ({}, ["--tensor-parallel", "3"], "num_attention_heads"), + ], +) +def test_invalid_inputs_preserve_output_before_any_write(tmp_path, updates, extra, expected): + source, resources = _inputs(tmp_path, overrides={**_PRECISION, **updates}) + output = tmp_path / "request.yaml" + output.write_text("original\n") + result = _real_cli(_args(source, resources, output) + ["--overwrite", *extra]) + assert result.returncode == 2 + assert expected in result.stderr + assert output.read_text() == "original\n" + + +def test_guided_shared_precision_then_default_skips_blind_tp_prompt_and_retains_review(monkeypatch, tmp_path, capsys): + source, _ = _inputs(tmp_path) + output = tmp_path / "new" / "request.yaml" + prompts = _terminal(monkeypatch, output, ["invalid", "bfloat16", "bfloat16", "", "accept"]) + assert cli.main(_args(source, None, output) + ["--interactive"]) == 0 + assert prompts[0].startswith("fmha_quant_mode") + assert prompts[2].startswith("kv_cache_dtype") + assert prompts[3] == "Choose topology 1-2 [1] (or cancel): " + assert prompts[-1] == "Review action (accept/edit/cancel): " + assert not any("Attention tensor-parallel" in prompt or "available" in prompt for prompt in prompts) + assert SupportRequest.from_yaml(output).worker_gpus == 1 + assert "collector FPM identity default" in capsys.readouterr().out + + +@pytest.mark.parametrize( + "choice,preset,parallel", + [("2", "pure_tp", (2, 1, 1, 2, 1, 1)), ("3", "dep", (1, 1, 2, 1, 2, 1)), ("5", "tep", (2, 1, 1, 1, 2, 1))], +) +def test_selected_moe_tuple_reaches_plan_collector_and_ordinary_runtime_configs( + monkeypatch, tmp_path, capsys, choice, preset, parallel +): + source, resources = _inputs(tmp_path, _MOE) + output = tmp_path / "new" / "request.yaml" + answers = [choice, "accept"] + prompts = _terminal(monkeypatch, output, answers) + assert cli.main(_args(source, resources, output) + ["--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + assert request.profile_deployment().parallel_tuple == parallel + assert request.parallel_preset == preset + assert prompts[0].startswith("Choose topology") + assert not any("overhead_bytes" in prompt for prompt in prompts) + assert request.profile_deployment().resources.memory_source == "pending" + source.unlink() + resources.unlink() + plan_root = tmp_path / "plan" + result = _real_cli(["onboard", "plan", "--config", str(output), "--output-dir", str(plan_root), "--format", "json"]) + assert result.returncode == 0, result.stderr + plan = json.loads((plan_root / "support-plan.json").read_text()) + assert plan["search"]["candidate_count"] == 1 + assert plan["fpm"]["collection_gpus_required"] == request.worker_gpus == 2 + command = plan["fpm"]["plan_command"] + assert command[command.index("--fpm-gpu-counts") + 1] == "2" + assert command[command.index("--fpm-parallel-presets") + 1] == preset + assert command[command.index("--fpm-model-profile") + 1] == str(plan_root / "fpm-model-profile.json") + collector = _real_cli(["onboard", "collect-fpm", "--config", str(output), "--output-dir", str(plan_root)]) + assert collector.returncode == 0, collector.stderr + assert shlex.split(collector.stdout) == command + prediction = CorePredictionConfig.from_yaml(plan_root / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan_root / "recommend/pilot.yaml") + assert prediction.engine.fpm_profile == recommendation.engine.fpm_profile == request.fpm_profile + for config in (prediction, recommendation): + worker = config.engine.workers.aggregated + assert worker.timing.estimation_mode == "fpm_interpolation" + assert worker.timing.fallback_policy == "deny" + assert worker.timing.estimator_config["fpm_interpolation"]["method"] == "direct" + assert prediction.engine.workers.aggregated.parallelism.model_dump(exclude_none=True) == request.parallelism() + assert ( + recommendation.engine.workers.aggregated.parallelism.preset[0].model_dump(exclude_none=True) + == request.parallelism() + ) + assert recommendation.optimization.constraints.max_candidate_gpus == 2 + assert not (plan_root / "fpm-checkpoint").exists() + assert not list((plan_root / "systems/data").glob("**/*.parquet")) + + +def test_guided_no_default_requires_number_then_keeps_memory_pending(monkeypatch, tmp_path, capsys): + source, resources = _inputs(tmp_path, {**_SMALL, "quantization_config": {"quant_method": "fp8"}}) + output = tmp_path / "new" / "request.yaml" + prompts = _terminal(monkeypatch, output, ["", "0", "2", "accept"]) + assert cli.main(_args(source, resources, output) + ["--interactive"]) == 0 + assert prompts[:3] == ["Choose topology 1-2 (or cancel): "] * 3 + assert prompts[3].startswith("Review action") + request = SupportRequest.from_yaml(output) + assert request.worker_gpus == 2 + assert request.profile_deployment().resources.weights_bytes is None + assert request.profile_deployment().resources.memory_source == "pending" + assert "Memory fit is unresolved" in capsys.readouterr().out + + +@pytest.mark.parametrize("interruption", ["cancel", EOFError(), KeyboardInterrupt()]) +@pytest.mark.parametrize("existing", [False, True]) +def test_guided_selection_cancel_and_interrupt_preserve_output(monkeypatch, tmp_path, interruption, existing): + source, resources = _inputs(tmp_path) + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("keep existing\n") + _terminal(monkeypatch, output, [interruption]) + assert cli.main(_args(source, resources, output) + ["--interactive", "--overwrite"]) == 130 + if existing: + assert output.read_text() == "keep existing\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + + +def test_selected_topology_survives_late_identity_correction(monkeypatch, tmp_path): + # Unknown framework versions are corrected before the selected profile is accepted. + source, resources = _inputs(tmp_path, {**_SMALL, "quantization_config": {"quant_method": "fp8"}}) + output = tmp_path / "new" / "request.yaml" + _terminal(monkeypatch, output, [_OPTIONS["framework_version"], "2", "accept"]) + assert cli.main(_args(source, resources, output, framework_version="unknown") + ["--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + assert request.worker_gpus == request.profile_deployment().tp == 2 + assert request.identity.framework_version == _OPTIONS["framework_version"] + + +@pytest.mark.parametrize("nested", [False, True]) +def test_cold_supervised_preview_does_not_import_execution_or_model_code(tmp_path, nested): + source, resources = _inputs(tmp_path) + if nested: + raw = json.loads(source.read_text()) + source.write_text( + json.dumps( + {"_name_or_path": raw.pop("_name_or_path"), "torch_dtype": raw.pop("torch_dtype"), "text_config": raw} + ) + ) + output = tmp_path / "new" / "request.yaml" + command = _args(source, resources, output) + ["--suggest-parallel"] + script = """ +import importlib.abc +import sys +blocked = ('aiconfigurator', 'aiconfigurator_core', 'aisimulate.sdk', 'aisimulate_core.sdk', + 'aisimulate_core._native', 'aisimulate._native', 'aisimulate._runtime', + 'collector', 'huggingface_hub', 'transformers', 'numpy', 'pandas', 'pyarrow', 'torch') +class Guard(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if any(fullname == name or fullname.startswith(name + '.') for name in blocked): + raise AssertionError('unexpected execution import: ' + fullname) +sys.meta_path.insert(0, Guard()) +from aisimulate.supervision import main +code = main(sys.argv[1:]) +assert not any(name == prefix or name.startswith(prefix + '.') for name in sys.modules for prefix in blocked) +core = {name for name in sys.modules if name == 'aisimulate_core' or name.startswith('aisimulate_core.')} +assert core <= {'aisimulate_core', 'aisimulate_core.fpm_profile', 'aisimulate_core.quantization'}, core +raise SystemExit(code) +""" + result = subprocess.run([sys.executable, "-c", script, *command], text=True, capture_output=True, timeout=30) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout)["default"]["required_gpus"] == 1 + assert not output.parent.exists() + + +def _directory_args(source, resources, root, **changes): + args = _args(source, resources, root, **changes) + args[args.index("--output")] = "--output-dir" + return args + + +def _directory_terminal(monkeypatch, root, answers): + remaining = iter(answers) + prompts = [] + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def respond(prompt): + prompts.append(prompt) + assert not root.exists() or not list(root.iterdir()), "configuration published before all acceptances" + try: + answer = next(remaining) + except StopIteration: + pytest.fail(f"unexpected prompt: {prompt}") + if isinstance(answer, BaseException): + raise answer + return answer + + monkeypatch.setattr(builtins, "input", respond) + return prompts + + +def _saved_configurations(root): + index = json.loads((root / "onboarding.json").read_text()) + assert index["schema_version"] == "aisimulate-onboarding/v1" + entries = index["configurations"] + assert len({entry["request_id"] for entry in entries}) == len(entries) + requests = [] + for entry in entries: + path = Path(entry["request"]) + assert path.parent.parent == root + request = SupportRequest.from_yaml(path) + assert request.fpm_profile == load_fpm_profile(Path(entry["fpm_profile"]).read_text()) + assert len(request.fpm_profile.deployments) == 1 + assert entry["collection_gpus_required"] == request.worker_gpus + assert shlex.split(entry["plan_command"]) == [ + "aisimulate", + "onboard", + "plan", + "--config", + str(path), + "--output-dir", + entry["collection_dir"], + ] + requests.append(request) + return entries, requests + + +def _consume_configurations(root): + entries, requests = _saved_configurations(root) + for entry, request in zip(entries, requests, strict=True): + plan_root = Path(entry["collection_dir"]) + assert plan_root == Path(entry["request"]).parent / "collection" + assert cli.main(shlex.split(entry["plan_command"])[1:]) == 0 + plan = json.loads((plan_root / "support-plan.json").read_text()) + assert plan["request_id"] == entry["request_id"] + assert plan["search"]["candidate_count"] == 1 + assert plan["fpm"]["collection_gpus_required"] == request.worker_gpus + assert plan["fpm"]["runtime_limits"]["context_length"] == request.search.context_length + command = plan["fpm"]["plan_command"] + for flag, expected in ( + ("--fpm-gpu-counts", str(request.worker_gpus)), + ("--fpm-parallel-presets", request.parallel_preset), + ("--fpm-max-model-len", str(request.search.context_length)), + ("--fpm-model-profile", str(plan_root / "fpm-model-profile.json")), + ("--fpm-database-root", str(plan_root / "systems/data")), + ("--checkpoint-dir", str(plan_root / "fpm-checkpoint")), + ): + assert command[command.index(flag) + 1] == expected + preview = _real_cli(["onboard", "collect-fpm", "--config", entry["request"], "--output-dir", str(plan_root)]) + assert preview.returncode == 0, preview.stderr + assert shlex.split(preview.stdout) == command + prediction = CorePredictionConfig.from_yaml(plan_root / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(plan_root / "recommend/pilot.yaml") + for config in (prediction, recommendation): + assert config.engine.fpm_profile == request.fpm_profile + assert config.engine.context_length == request.search.context_length + assert [str(path) for path in config.engine.systems_paths] == [str(plan_root / "systems")] + worker = config.engine.workers.aggregated + assert worker.timing.estimation_mode == "fpm_interpolation" + assert worker.timing.fallback_policy == "deny" + assert worker.timing.estimator_config["fpm_interpolation"]["method"] == "direct" + assert worker.kv_cache.capacity.memory_fraction == request.collection.memory_fraction + if request.profile_deployment().resources.cache_layout == "grouped": + assert not worker.kv_cache.prefix_caching + assert prediction.engine.workers.aggregated.parallelism.model_dump(exclude_none=True) == request.parallelism() + assert ( + recommendation.engine.workers.aggregated.parallelism.preset[0].model_dump(exclude_none=True) + == request.parallelism() + ) + assert recommendation.optimization.constraints.max_candidate_gpus == request.worker_gpus + assert not (plan_root / "fpm-checkpoint").exists() + assert not list((plan_root / "systems/data").glob("**/*.parquet")) + return requests + + +def test_guided_multiple_tp_dep_tep_profiles_reach_independent_consumers(monkeypatch, tmp_path): + source, resources = _inputs(tmp_path, _MOE) + root = tmp_path / "new collection" / "profiles" + prompts = _directory_terminal( + monkeypatch, + root, + [ + "bfloat16", + "half", + "bfloat16", + "1,1", + "7", + "x", + "2,,3", + "2, 3,5", + "accept", + "accept", + "accept", + ], + ) + assert cli.main(_directory_args(source, None, root) + ["--interactive"]) == 0 + assert sum(prompt.startswith("Choose topologies") for prompt in prompts) == 5 + assert sum(prompt.startswith("fmha_quant_mode") for prompt in prompts) == 1 + assert sum(prompt.startswith("Review action") for prompt in prompts) == 3 + source.unlink() + resources.unlink() + requests = _consume_configurations(root) + assert [request.parallel_preset for request in requests] == ["pure_tp", "dep", "tep"] + assert [request.profile_deployment().parallel_tuple for request in requests] == [ + (2, 1, 1, 2, 1, 1), + (1, 1, 2, 1, 2, 1), + (2, 1, 1, 1, 2, 1), + ] + assert all(request.profile_deployment().resources.memory_source == "pending" for request in requests) + assert not any("overhead_bytes" in prompt for prompt in prompts) + assert ( + requests[0].profile_deployment().resources.kv_bytes_per_token + < requests[1].profile_deployment().resources.kv_bytes_per_token + ) + + +@pytest.mark.parametrize("file_format", ["json", "yaml"]) +def test_parallel_configuration_file_preserves_rank_bounds_and_precision(tmp_path, file_format): + source, resources = _inputs(tmp_path, _MOE) + entries = [ + {"tensor_parallel": 2, "resource_overrides": {"weights_bytes": 1024**3}}, + { + "tensor_parallel": 1, + "attention_data_parallel": 2, + "moe_expert_parallel": 2, + "resource_overrides": {"weights_bytes": 2 * 1024**3, "comm_overhead_bytes": 16 * 1024**2}, + }, + { + "tensor_parallel": 2, + "moe_tensor_parallel": 1, + "moe_expert_parallel": 2, + "resource_overrides": { + "weights_bytes": 3 * 1024**3, + "comm_overhead_bytes": 32 * 1024**2, + "kv_cache_dtype": "fp8", + }, + }, + ] + choices = tmp_path / f"parallel.{file_format}" + choices.write_text(json.dumps(entries) if file_format == "json" else yaml.safe_dump(entries)) + root = tmp_path / "profiles" + root.mkdir() # A new empty directory is accepted, but never a collected-data directory. + result = _real_cli(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) + assert result.returncode == 0, result.stderr + assert "Suggested worker topologies" not in result.stdout + source.unlink() + resources.unlink() + choices.unlink() + requests = _consume_configurations(root) + assert all(request.profile_deployment().resources.memory_source == "pending" for request in requests) + assert [ + json.loads(request.profile_deployment().resources.provenance)["fields"]["weights_bytes"]["value"] + for request in requests + ] == [1024**3, 2 * 1024**3, 3 * 1024**3] + assert [request.profile_deployment().kv_cache_dtype for request in requests] == ["bfloat16", "bfloat16", "fp8"] + + +@pytest.mark.parametrize( + "model_context,shared_context,entry_contexts,runtime_context,expected", + [ + (4096, None, [None, 2048], None, [4096, 2048]), + (8192, 6144, [None, 7168], None, [6144, 7168]), + (524288, None, [262144, 131072], None, [256000, 131072]), + (4096, None, [None, 2048], 2048, [2048, 2048]), + (524288, None, [None, 400000], 300000, [300000, 300000]), + ], + ids=["model-default", "merged-default", "agentx-cap", "explicit-fitting", "explicit-above-default-cap"], +) +def test_parallel_context_limits_reach_saved_plan_consumers( + tmp_path, model_context, shared_context, entry_contexts, runtime_context, expected +): + shared = {**_PRECISION} + if shared_context is not None: + shared["context_length"] = shared_context + source, resources = _inputs(tmp_path, {**_SMALL, "max_position_embeddings": model_context}, shared) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps( + [ + { + "tensor_parallel": tp, + "resource_overrides": {} if context is None else {"context_length": context}, + } + for tp, context in enumerate(entry_contexts, 1) + ] + ) + ) + root = tmp_path / "profiles" + result = _real_cli( + _directory_args(source, resources, root, context_length=runtime_context) + ["--parallel-configs", str(choices)] + ) + assert result.returncode == 0, result.stderr + source.unlink() + resources.unlink() + choices.unlink() + requests = _consume_configurations(root) + assert [request.search.context_length for request in requests] == expected + assert [request.fpm_profile.context_length for request in requests] == [ + context if context is not None else shared_context or model_context for context in entry_contexts + ] + + +def test_explicit_context_exceeding_later_profile_cap_publishes_nothing(tmp_path, capsys): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps([{"tensor_parallel": 1}, {"tensor_parallel": 2, "resource_overrides": {"context_length": 2048}}]) + ) + root = tmp_path / "new" / "profiles" + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root, context_length=4096) + ["--parallel-configs", str(choices)]) + assert error.value.code == 2 + assert "search.context_length exceeds the configured profile context_length" in capsys.readouterr().err + assert not root.parent.exists() + + +def test_implicit_context_defaults_survive_candidate_identity_corrections(monkeypatch, tmp_path): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps([{"tensor_parallel": 1, "resource_overrides": {"context_length": 2048}}, {"tensor_parallel": 2}]) + ) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["0.25.1", "accept", "0.25.2", "accept"]) + assert ( + cli.main( + _directory_args(source, resources, root, context_length=None, framework_version="unknown") + + ["--parallel-configs", str(choices), "--interactive"] + ) + == 0 + ) + assert sum(prompt.startswith("Pinned vLLM version") for prompt in prompts) == 2 + assert not any(prompt.startswith("Runtime per-request context") for prompt in prompts) + source.unlink() + resources.unlink() + choices.unlink() + requests = _consume_configurations(root) + assert [request.search.context_length for request in requests] == [2048, 4096] + assert [request.identity.framework_version for request in requests] == ["0.25.1", "0.25.2"] + + +def test_explicit_context_correction_stays_with_its_candidate(monkeypatch, tmp_path): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps([{"tensor_parallel": 1, "resource_overrides": {"context_length": 2048}}, {"tensor_parallel": 2}]) + ) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["2048", "accept", "accept"]) + assert ( + cli.main( + _directory_args(source, resources, root, context_length=4096) + + ["--parallel-configs", str(choices), "--interactive"] + ) + == 0 + ) + assert sum(prompt.startswith("Runtime per-request context") for prompt in prompts) == 1 + source.unlink() + resources.unlink() + choices.unlink() + requests = _consume_configurations(root) + assert [request.search.context_length for request in requests] == [2048, 4096] + + +def test_parallel_profiles_reach_real_collector_plan_only_without_gpu_work(tmp_path): + source, resources = _inputs(tmp_path, _MOE) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps( + [ + {"tensor_parallel": 2}, + { + "tensor_parallel": 1, + "attention_data_parallel": 2, + "moe_expert_parallel": 2, + "resource_overrides": {"comm_overhead_bytes": 16 * 1024**2}, + }, + { + "tensor_parallel": 2, + "moe_tensor_parallel": 1, + "moe_expert_parallel": 2, + "resource_overrides": {"comm_overhead_bytes": 32 * 1024**2}, + }, + ] + ) + ) + root = tmp_path / "profiles" + assert cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) == 0 + entries, requests = _saved_configurations(root) + collector_ids = set() + for entry, request in zip(entries, requests, strict=True): + assert cli.main(shlex.split(entry["plan_command"])[1:]) == 0 + plan_root = Path(entry["collection_dir"]) + command = json.loads((plan_root / "support-plan.json").read_text())["fpm"]["plan_command"] + # Runtime collector metadata still comes from the checkpoint/config; + # the emitted inline profile supplies the reviewed resource bounds. + result = subprocess.run( + [sys.executable, *command[1:], "--fpm-model-config", str(source)], + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + collected_plan = json.loads(result.stdout) + assert collected_plan["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert collected_plan["counts"]["topologies"] == 1 + assert {cell["parallel_strategy"] for cell in collected_plan["cells"]} == {request.parallel_preset} + assert {item["source"] for item in collected_plan["topology_memory_admission"]} == {"fpm_profile_pending"} + collector_ids.add(collected_plan["sha256"]) + assert not (plan_root / "fpm-checkpoint").exists() + assert not (plan_root / "fpm-artifacts").exists() + assert not list((plan_root / "systems/data").glob("**/*.parquet")) + assert len(collector_ids) == 3 + + +def test_interactive_file_resolves_shared_fields_once_and_retains_entry_precision(monkeypatch, tmp_path): + source, _ = _inputs(tmp_path) + choices = tmp_path / "parallel.json" + choices.write_text( + json.dumps( + [ + {"tensor_parallel": 1, "resource_overrides": {"kv_cache_dtype": "fp8"}}, + {"tensor_parallel": 2}, + ] + ) + ) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["bfloat16", "bfloat16", "accept", "accept"]) + assert cli.main(_directory_args(source, None, root) + ["--parallel-configs", str(choices), "--interactive"]) == 0 + assert [prompt.split(" (")[0] for prompt in prompts[:2]] == ["fmha_quant_mode", "kv_cache_dtype"] + assert len(prompts) == 4 + _, requests = _saved_configurations(root) + assert [request.profile_deployment().kv_cache_dtype for request in requests] == ["fp8", "bfloat16"] + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_directory_singleton_preserves_default_or_explicit_topology(monkeypatch, tmp_path, explicit): + source, resources = _inputs(tmp_path) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["accept"] if explicit else ["", "accept"]) + flags = ["--tensor-parallel", "2"] if explicit else [] + assert cli.main(_directory_args(source, resources, root) + ["--interactive", *flags]) == 0 + _, requests = _saved_configurations(root) + assert len(requests) == 1 + assert requests[0].worker_gpus == (2 if explicit else 1) + assert sum(prompt.startswith("Choose topologies") for prompt in prompts) == (0 if explicit else 1) + + +def test_guided_directory_without_default_requires_selection_and_preserves_pending_memory(monkeypatch, tmp_path): + source, resources = _inputs(tmp_path, {**_SMALL, "quantization_config": {"quant_method": "fp8"}}) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["", "1,2", "accept", "accept"]) + assert cli.main(_directory_args(source, resources, root) + ["--interactive"]) == 0 + assert sum(prompt.startswith("Choose topologies") for prompt in prompts) == 2 + assert not any(prompt.startswith("weights_bytes") for prompt in prompts) + _, requests = _saved_configurations(root) + assert all(request.profile_deployment().resources.memory_source == "pending" for request in requests) + + +def test_parallel_file_accepts_explicit_topology_outside_shortlist(tmp_path): + source, resources = _inputs(tmp_path, _LARGE) + choices = tmp_path / "parallel.yaml" + choices.write_text( + yaml.safe_dump( + [ + { + "tensor_parallel": 16, + "resource_overrides": {"activations_bytes": 70 * 1024**2, "comm_overhead_bytes": 256 * 1024**2}, + } + ] + ) + ) + root = tmp_path / "profiles" + assert cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) == 0 + assert _saved_configurations(root)[1][0].worker_gpus == 16 + + +def test_grouped_profile_edits_and_collection_limits_stay_with_their_candidate(monkeypatch, tmp_path): + source, resources = _inputs( + tmp_path, + {**_SMALL, "sliding_window": 128, "layer_types": ["full_attention", "sliding_attention"]}, + {**_PRECISION, "cache_block_sizes": {"full_attention": 16, "sliding_attention": 8}}, + ) + choices = tmp_path / "parallel.json" + choices.write_text(json.dumps([{"tensor_parallel": 1}, {"tensor_parallel": 2}])) + root = tmp_path / "profiles" + _directory_terminal( + monkeypatch, + root, + [ + "edit", + "cache_block_sizes", + '{"full_attention": 32, "sliding_attention": 16}', + "edit", + "max_num_tokens", + "2048", + "edit", + "runtime_context_length", + "2048", + "edit", + "max_prefill_cudagraph_size", + "512", + "edit", + "gpu_memory_utilization", + "0.75", + "accept", + "accept", + ], + ) + assert ( + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices), "--interactive"]) == 0 + ) + source.unlink() + resources.unlink() + choices.unlink() + first, second = _consume_configurations(root) + assert [request.search.context_length for request in (first, second)] == [2048, 4096] + assert [request.scheduler_limits()["max_batched_tokens"] for request in (first, second)] == [2048, 8192] + assert [request.collection_settings()["max_prefill_cudagraph_size"] for request in (first, second)] == [512, None] + assert [request.collection.prefill_cudagraph_policy for request in (first, second)] == ["explicit", "runtime"] + assert [request.collection.gpu_memory_utilization for request in (first, second)] == [0.75, 0.9] + first_cache, second_cache = (request.profile_deployment().resources for request in (first, second)) + assert [group.block_size_tokens for group in first_cache.cache_groups] == [32, 16] + assert [group.block_size_tokens for group in second_cache.cache_groups] == [16, 8] + assert first_cache.cache_groups[0].page_size_bytes != second_cache.cache_groups[0].page_size_bytes + assert "user override" not in json.loads(second_cache.provenance)["fields"]["cache_groups"]["source"] + + +def test_grouped_per_entry_pages_do_not_replace_other_topology_estimates(tmp_path): + source, resources = _inputs( + tmp_path, + {**_SMALL, "sliding_window": 128}, + {**_PRECISION, "cache_block_sizes": {"sliding_attention": 16}}, + ) + groups = [ + { + "name": "sliding_attention", + "kind": "attention", + "num_layers": 2, + "block_size_tokens": 32, + "page_size_bytes": 8192, + "sliding_window": 128, + } + ] + choices = tmp_path / "parallel.yaml" + choices.write_text( + yaml.safe_dump( + [ + {"tensor_parallel": 1, "resource_overrides": {"cache_groups": groups}}, + {"tensor_parallel": 2}, + ] + ) + ) + root = tmp_path / "profiles" + assert cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) == 0 + source.unlink() + first, second = _consume_configurations(root) + assert first.profile_deployment().resources.cache_groups[0].page_size_bytes == 8192 + assert first.profile_deployment().resources.cache_groups[0].block_size_tokens == 32 + assert second.profile_deployment().resources.cache_groups[0].block_size_tokens == 16 + assert second.profile_deployment().resources.cache_groups[0].page_size_bytes < 8192 + + +@pytest.mark.parametrize("interruption", ["cancel", EOFError(), KeyboardInterrupt()]) +@pytest.mark.parametrize("existing", [False, True]) +def test_later_profile_cancellation_publishes_nothing(monkeypatch, tmp_path, interruption, existing): + source, resources = _inputs(tmp_path) + root = tmp_path / "new" / "profiles" + if existing: + root.mkdir(parents=True) + prompts = _directory_terminal(monkeypatch, root, ["1,2", "accept", interruption]) + assert cli.main(_directory_args(source, resources, root) + ["--interactive"]) == 130 + assert sum(prompt.startswith("Review action") for prompt in prompts) == 2 + if existing: + assert not list(root.iterdir()) + else: + assert not root.parent.exists() + + +def test_late_identity_corrections_and_pending_profiles_are_independent(monkeypatch, tmp_path): + source, resources = _inputs(tmp_path, {**_SMALL, "quantization_config": {"quant_method": "fp8"}}) + choices = tmp_path / "parallel.json" + choices.write_text(json.dumps([{"tensor_parallel": 1}, {"tensor_parallel": 2}])) + root = tmp_path / "profiles" + prompts = _directory_terminal(monkeypatch, root, ["0.25.1", "accept", "0.25.2", "accept"]) + assert ( + cli.main( + _directory_args(source, resources, root, framework_version="unknown") + + ["--parallel-configs", str(choices), "--interactive"] + ) + == 0 + ) + _, requests = _saved_configurations(root) + assert [request.worker_gpus for request in requests] == [1, 2] + assert [request.identity.framework_version for request in requests] == ["0.25.1", "0.25.2"] + assert all(request.profile_deployment().resources.memory_source == "pending" for request in requests) + assert sum(prompt.startswith("Pinned vLLM version") for prompt in prompts) == 2 + + +def test_topology_correction_cannot_publish_duplicate_configurations(monkeypatch, tmp_path, capsys): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.yaml" + choices.write_text(yaml.safe_dump([{"tensor_parallel": 1}, {"tensor_parallel": 3}])) + root = tmp_path / "new" / "profiles" + _directory_terminal(monkeypatch, root, ["accept", "1", "accept"]) + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices), "--interactive"]) + assert error.value.code == 2 + assert "duplicate resolved parallel configuration" in capsys.readouterr().err + assert not root.parent.exists() + + +@pytest.mark.parametrize( + "payload,expected", + [ + ([], "nonempty"), + ({"tensor_parallel": 1}, "list"), + ([1], "object"), + ([{}], "requires tensor_parallel"), + ([{"tensor_parallel": True}], "positive integer"), + ([{"tensor_parallel": "2"}], "positive integer"), + ([{"tensor_parallel": 2.0}], "positive integer"), + ([{"tensor_parallel": 0}], "positive integer"), + ([{"tensor_parallel": 1, "attention_data_parallel": None}], "positive integer"), + ([{"tensor_parallel": 1, "name": "../escape"}], "unknown fields"), + ([{"tensor_parallel": 1, "resource_overrides": None}], "must be a mapping"), + ([{"tensor_parallel": 1, "resource_overrides": {"weights_bytes": True}}], "weights_bytes"), + ([{"tensor_parallel": 1, "resource_overrides": {"unknown": 4}}], "unknown resource"), + ( + [ + {"tensor_parallel": 2}, + { + "tensor_parallel": 2, + "attention_data_parallel": 1, + "moe_tensor_parallel": 2, + "moe_expert_parallel": 1, + }, + ], + "duplicate resolved", + ), + ([{"tensor_parallel": 2, "attention_data_parallel": 2}], "complete TP, DEP, or TEP"), + ([{"tensor_parallel": 1, 7: "invalid", "unknown": 3}], "field names must be strings"), + ], +) +def test_invalid_parallel_files_leave_no_partial_output(tmp_path, capsys, payload, expected): + source, resources = _inputs(tmp_path, _MOE) + choices = tmp_path / "parallel.yaml" + choices.write_text(yaml.safe_dump(payload)) + root = tmp_path / "new" / "profiles" + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) + assert error.value.code == 2 + assert expected in capsys.readouterr().err + assert not root.parent.exists() + + +@pytest.mark.parametrize("payload", ["- tensor_parallel: 1\n tensor_parallel: 2\n", "[{"]) +def test_duplicate_keys_and_malformed_parallel_files_fail_cleanly(tmp_path, capsys, payload): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.yaml" + choices.write_text(payload) + root = tmp_path / "profiles" + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) + assert error.value.code == 2 + expected = "duplicate" if "tensor_parallel" in payload else "malformed" + assert expected in capsys.readouterr().err + assert not root.exists() + + +@pytest.mark.parametrize( + "field,value", + [ + ("weights_bytes", 1024**3), + ("activations_bytes", 1024), + ("runtime_overhead_bytes", 1024), + ("comm_overhead_bytes", 1024), + ("kv_bytes_per_token", 1024), + ( + "cache_groups", + [ + { + "name": "full_attention", + "kind": "attention", + "num_layers": 2, + "block_size_tokens": 16, + "page_size_bytes": 4096, + } + ], + ), + ], +) +def test_parallel_configurations_reject_shared_rank_local_resources(tmp_path, capsys, field, value): + source, resources = _inputs(tmp_path, overrides={**_PRECISION, field: value}) + choices = tmp_path / "parallel.json" + choices.write_text(json.dumps([{"tensor_parallel": 1}, {"tensor_parallel": 2}])) + root = tmp_path / "profiles" + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) + assert error.value.code == 2 + assert "cannot share rank-local resource overrides" in capsys.readouterr().err + assert not root.exists() + + +@pytest.mark.parametrize( + "extra", + [ + ["--output", "request.yaml"], + ["--overwrite"], + ["--suggest-parallel"], + ["--parallel-configs", "absent.yaml", "--tensor-parallel", "1"], + ["--parallel-configs", "absent.yaml", "--attention-data-parallel", "1"], + ["--parallel-configs", "absent.yaml", "--moe-tensor-parallel", "1"], + ["--parallel-configs", "absent.yaml", "--moe-expert-parallel", "1"], + ], +) +def test_directory_option_conflicts_fail_before_loading_inputs(tmp_path, capsys, extra): + root = tmp_path / "profiles" + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(tmp_path / "absent.json", None, root) + extra) + assert error.value.code == 2 + expected = "not allowed" if "--output" in extra else "cannot be combined" + assert expected in capsys.readouterr().err + assert not root.exists() + + +@pytest.mark.parametrize("extra", [[], ["--model-config", "absent.json"], ["--output-dir", "absent"]]) +def test_parallel_file_requires_both_config_and_directory(capsys, extra): + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "init", "--parallel-configs", "absent.yaml", *extra]) + assert error.value.code == 2 + assert "requires --model-config and --output-dir" in capsys.readouterr().err + + +def test_empty_directory_argument_cannot_fall_back_to_single_file_output(monkeypatch, tmp_path, capsys): + source, resources = _inputs(tmp_path) + monkeypatch.chdir(tmp_path) + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, "")) + assert error.value.code == 2 + assert "nonempty path" in capsys.readouterr().err + assert not (tmp_path / "support-request.yaml").exists() + + +@pytest.mark.parametrize("kind", ["nonempty", "file", "symlink", "dangling", "parent_symlink"]) +def test_directory_path_conflicts_preserve_existing_contents(tmp_path, capsys, kind): + source, resources = _inputs(tmp_path) + protected = tmp_path / "collected" + protected.mkdir() + (protected / "fpm_forward_perf.parquet").write_bytes(b"existing collected data") + root = tmp_path / "profiles" + if kind == "nonempty": + root = protected + elif kind == "file": + root.write_text("existing file") + elif kind in {"symlink", "dangling"}: + root.symlink_to(protected if kind == "symlink" else tmp_path / "absent", target_is_directory=True) + else: + link = tmp_path / "linked" + link.symlink_to(protected, target_is_directory=True) + root = link / "profiles" + before = {path: path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root)) + assert error.value.code == 2 + expected = "symlinks" if "symlink" in kind or kind == "dangling" else "new or empty" + assert expected in capsys.readouterr().err + assert {path: path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} == before + + +@pytest.mark.parametrize("failure", ["occupied", "symlink", "serialization"]) +def test_directory_publication_failure_leaves_no_partial_profiles(monkeypatch, tmp_path, capsys, failure): + source, resources = _inputs(tmp_path) + choices = tmp_path / "parallel.json" + choices.write_text(json.dumps([{"tensor_parallel": 1}, {"tensor_parallel": 2}])) + root = tmp_path / "profiles" + protected = tmp_path / "collected" + protected.mkdir() + marker = protected / "fpm_forward_perf.parquet" + marker.write_bytes(b"existing collected data") + write_text = Path.write_text + + def interrupt_publication(path, *args, **kwargs): + if path.name == "onboarding.json": + if failure == "occupied": + root.mkdir() + (root / "collected.txt").write_bytes(b"new external data") + elif failure == "symlink": + root.symlink_to(protected, target_is_directory=True) + if failure == "serialization" and path.name == "fpm-profile.json" and path.parent.name.startswith("tp2-"): + raise OSError("simulated profile write failure") + return write_text(path, *args, **kwargs) + + monkeypatch.setattr(Path, "write_text", interrupt_publication) + with pytest.raises(SystemExit) as error: + cli.main(_directory_args(source, resources, root) + ["--parallel-configs", str(choices)]) + assert error.value.code == 2 + assert marker.read_bytes() == b"existing collected data" + if failure == "occupied": + assert [(path.name, path.read_bytes()) for path in root.iterdir()] == [("collected.txt", b"new external data")] + elif failure == "symlink": + assert root.is_symlink() + else: + assert not root.exists() + assert not list(tmp_path.glob(".profiles-*")) diff --git a/python/aisimulate/tests/unit/test_support_cli.py b/python/aisimulate/tests/unit/test_support_cli.py index 4333cb712..ce1ef4652 100644 --- a/python/aisimulate/tests/unit/test_support_cli.py +++ b/python/aisimulate/tests/unit/test_support_cli.py @@ -16,6 +16,7 @@ from types import SimpleNamespace import pytest +import yaml import aisimulate.main as cli from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig @@ -29,7 +30,6 @@ "model_kind": "dense", "framework_version": "0.24.0", "gpu": "h200_sxm", - "gpu_count": 4, "interconnect": "nvswitch", } _PILOT = { @@ -44,7 +44,7 @@ def _init_args(output: Path, *, full: bool = False, **changes) -> list[str]: - options = {**_REQUIRED, **(_PILOT if full else {}), **changes} + options = {**_REQUIRED, "context_length": 16384, **(_PILOT if full else {}), **changes} return ["onboard", "init", "--output", str(output)] + [ part for name, value in options.items() @@ -95,6 +95,79 @@ def test_help_explains_model_fpm_and_target_hardware_scope(capsys, command) -> N assert "aisimulate" in output +def test_init_help_explains_collection_requirements_and_runtime_budgets(capsys) -> None: + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "init", "--help"]) + assert result.value.code == 0 + output = " ".join(capsys.readouterr().out.split()) + assert "Collection GPUs are derived from the selected topology" in output + assert "Set deployment replicas and GPU budgets in ordinary predict/recommend configs" in output + for obsolete in ("--gpu-count", "--node-count", "--gpus-per-node", "--max-candidates"): + assert obsolete not in output + + +@pytest.mark.parametrize("option", ["gpu-count", "node-count", "gpus-per-node", "max-candidates"]) +def test_obsolete_onboarding_options_fail_without_writing(tmp_path, monkeypatch, capsys, option): + output = tmp_path / "request.yaml" + output.write_text("existing request\n") + prompts = _terminal(monkeypatch) + with pytest.raises(SystemExit) as result: + cli.main(_init_args(output) + [f"--{option}", "4", "--interactive", "--overwrite"]) + assert result.value.code == 2 + assert f"unrecognized arguments: --{option}" in capsys.readouterr().err + assert prompts == [] + assert output.read_text() == "existing request\n" + assert list(tmp_path.iterdir()) == [output] + + +@pytest.mark.parametrize("action", ["plan", "collect-fpm"]) +def test_legacy_saved_request_requires_new_plan_and_preserves_old_data(tmp_path, monkeypatch, capsys, action): + from aisimulate.support.plan import create_plan + + def unexpected_launch(*args, **kwargs): + pytest.fail("request migration must not launch collection") + + monkeypatch.setitem(sys.modules, "collector.fpm_forward.cli", SimpleNamespace(main=unexpected_launch)) + request = tmp_path / "request.yaml" + assert cli.main(_init_args(request, tensor_parallel=4)) == 0 + current = SupportRequest.from_yaml(request) + old_root = tmp_path / "old-plan" + create_plan(current, old_root) + old_request = yaml.safe_load((old_root / "request.yaml").read_text()) + old_request["identity"].update(gpu_count=16, node_count=2, gpus_per_node=8) + old_request["search"]["max_candidates"] = 2 + (old_root / "request.yaml").write_text(yaml.safe_dump(old_request)) + old_plan = json.loads((old_root / "support-plan.json").read_text()) + old_plan["request_id"] = "onboarding-legacy-request" + (old_root / "support-plan.json").write_text(json.dumps(old_plan)) + (old_root / "systems/data/existing.parquet").write_bytes(b"collected timing data") + before = {p.relative_to(old_root): p.read_bytes() for p in old_root.rglob("*") if p.is_file()} + args = ["onboard", action, "--config", str(old_root / "request.yaml"), "--output-dir", str(old_root)] + args += ["--overwrite"] if action == "plan" else ["--execute"] + + with pytest.raises(SystemExit) as result: + cli.main(args) + assert result.value.code == 2 + assert "Copy the request, remove these fields" in capsys.readouterr().err + assert {p.relative_to(old_root): p.read_bytes() for p in old_root.rglob("*") if p.is_file()} == before + + # Migrate a copy and create a new plan; the old plan and timings retain their identity. + for name in ("gpu_count", "node_count", "gpus_per_node"): + del old_request["identity"][name] + del old_request["search"]["max_candidates"] + migrated = tmp_path / "migrated.yaml" + migrated.write_text(yaml.safe_dump(old_request)) + assert SupportRequest.from_yaml(migrated) == current + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "plan", "--config", str(migrated), "--output-dir", str(old_root), "--overwrite"]) + assert result.value.code == 2 + assert "choose a new output directory" in capsys.readouterr().err + new_root = tmp_path / "new-plan" + assert cli.main(["onboard", "plan", "--config", str(migrated), "--output-dir", str(new_root)]) == 0 + assert SupportRequest.from_yaml(new_root / "request.yaml") == current + assert {p.relative_to(old_root): p.read_bytes() for p in old_root.rglob("*") if p.is_file()} == before + + @pytest.mark.parametrize("action", ["init", "plan", "collect-fpm"]) def test_retired_support_command_is_rejected_before_dispatch(tmp_path, monkeypatch, capsys, action) -> None: def unexpected_dispatch(*args, **kwargs): @@ -118,26 +191,54 @@ def unexpected_dispatch(*args, **kwargs): assert not list(tmp_path.iterdir()) +@pytest.mark.parametrize("search", [123, True, [], None, [["context_length", 4096]]]) +def test_v1_request_rejects_malformed_search_without_writing_plan(tmp_path, capsys, search): + source = tmp_path / "request.yaml" + source.write_text( + yaml.safe_dump({"schema_version": "aisimulate-support-request/v1", "identity": _REQUIRED, "search": search}) + ) + original = source.read_bytes() + output = tmp_path / "plan" + + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "plan", "--config", str(source), "--output-dir", str(output)]) + + assert result.value.code == 2 + assert "search" in capsys.readouterr().err + assert source.read_bytes() == original + assert not output.exists() + + def test_guided_and_scripted_setup_produce_the_same_request(tmp_path, monkeypatch, capsys) -> None: guided = tmp_path / "guided request.yaml" scripted = tmp_path / "scripted.yaml" prompts = _terminal( monkeypatch, - ["example/unintegrated-model", "revision-123", "dense", "0.24.0", "h200_sxm", "4", "nvswitch", "2"] + [""] * 6, + ["example/unintegrated-model", "revision-123", "dense", "0.24.0", "h200_sxm", "nvswitch", "2"] + + ["16384"] + + [""] * 6, ) assert cli.main(["onboard", "init", "--interactive", "--output", str(guided)]) == 0 - assert cli.main(_init_args(scripted, tensor_parallel=2)) == 0 + assert cli.main(_init_args(scripted, tensor_parallel=2, context_length=16384)) == 0 assert SupportRequest.from_yaml(guided) == SupportRequest.from_yaml(scripted) - assert any("Input tokens per request [1024]" in prompt for prompt in prompts) - assert any("Target time to first token (ms) [1000.0]" in prompt for prompt in prompts) + assert any("Runtime per-request context limit" in prompt for prompt in prompts) + assert not any( + "input tokens" in prompt.lower() or "output tokens" in prompt.lower() or "concurrent" in prompt.lower() + for prompt in prompts + ) + assert not any( + "time to first token" in prompt.lower() or "time per output token" in prompt.lower() for prompt in prompts + ) request = SupportRequest.from_yaml(guided) assert request.workload.request_count == 4 - assert request.identity.gpus_per_node == 4 + assert request.worker_gpus == 2 + assert not any("node" in prompt.lower() or "available" in prompt.lower() for prompt in prompts) assert request.identity.aisimulate_revision is None assert request.identity.tokenizer_revision is None output = capsys.readouterr().out + assert "Collection GPUs required: 2" in output assert "accuracy is not assessed" in output assert "are unchecked" in output @@ -146,7 +247,7 @@ def test_supplied_options_skip_prompts_and_onboarding_alias_is_optional(tmp_path prompts = _terminal(monkeypatch) guided = tmp_path / "guided.yaml" scripted = tmp_path / "scripted.yaml" - options = {"request_count": 8, "seed": 123, "max_candidates": 2, "objective": "goodput"} + options = {"request_count": 8, "seed": 123, "objective": "goodput"} assert cli.main(_init_args(guided, full=True, **options) + ["--interactive", "--profile", "onboarding"]) == 0 assert cli.main(_init_args(scripted, full=True, **options)) == 0 @@ -159,10 +260,10 @@ def test_guided_setup_recovers_numeric_input_and_shared_field_validation(tmp_pat output = tmp_path / "request.yaml" prompts = _terminal(monkeypatch, ["many", "4", "revision-123"]) - assert cli.main(_init_args(output, full=True, model_revision="main", gpu_count=None) + ["--interactive"]) == 0 + assert cli.main(_init_args(output, full=True, model_revision="main", tensor_parallel=None) + ["--interactive"]) == 0 request = SupportRequest.from_yaml(output) - assert request.identity.gpu_count == 4 + assert request.worker_gpus == 4 assert request.identity.model_revision == "revision-123" assert len(prompts) == 3 transcript = capsys.readouterr().out @@ -173,10 +274,9 @@ def test_guided_setup_recovers_numeric_input_and_shared_field_validation(tmp_pat @pytest.mark.parametrize( ("changes", "answers", "expected"), [ - ({"tensor_parallel": 8}, ["unknown-option", "--tensor-parallel", "2"], "tensor_parallel"), + ({"attention_data_parallel": 2}, ["unknown-option", "--attention-data-parallel", "1"], "TP, DEP, or TEP"), ({"context_length": 100}, ["context-length", "4096"], "context_length"), ({"concurrency": 8}, ["request-count", "8"], "request_count"), - ({"node_count": 2, "gpus_per_node": 4}, ["gpu-count", "8"], "gpu_count"), ], ) def test_guided_setup_recovers_cross_field_validation( @@ -197,14 +297,11 @@ def test_guided_setup_recovers_cross_field_validation( [ {"model": None}, {"model_revision": "latest"}, - {"gpu_count": 0}, {"ttft_ms": "nan"}, {"request_count": 0}, {"context_length": 100}, - {"tensor_parallel": 8}, - {"max_candidates": 3}, + {"tensor_parallel": 0}, {"objective": "invented"}, - {"node_count": 2}, ], ) def test_scripted_setup_uses_shared_validation_and_never_writes_invalid_requests(tmp_path, capsys, changes) -> None: @@ -375,6 +472,7 @@ def unexpected_launch(*args, **kwargs): plan = json.loads((output / "support-plan.json").read_text()) assert summary["request_id"] == plan["request_id"] assert summary["candidate_count"] == 1 + assert summary["collection_gpus_required"] == plan["fpm"]["collection_gpus_required"] == 2 prediction = CorePredictionConfig.from_yaml(plan["outputs"]["prediction_configs"][0]) recommendation = CoreRecommendationConfig.from_yaml(plan["outputs"]["recommendation_configs"][0]) assert prediction.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" @@ -393,13 +491,10 @@ def unexpected_launch(*args, **kwargs): def test_explicit_execute_forwards_diagnostic_options_and_exit_status(tmp_path, monkeypatch) -> None: - calls = [] - - def collector_main(argv): - calls.append(argv) - return 7 + from .test_support_plan import _mock_collector_execution - monkeypatch.setitem(sys.modules, "collector.fpm_forward.cli", SimpleNamespace(main=collector_main)) + calls = [] + _mock_collector_execution(monkeypatch, calls) request_path = tmp_path / "request.yaml" output = tmp_path / "plan" assert cli.main(_init_args(request_path, tensor_parallel=2)) == 0 @@ -421,7 +516,7 @@ def collector_main(argv): "--resume", ] ) - == 7 + == 1 ) assert len(calls) == 1 @@ -431,7 +526,7 @@ def collector_main(argv): assert calls[0][calls[0].index("--limit") + 1] == "1" -def _local_collection_command(tmp_path, **changes): +def _local_collection_command(tmp_path, *, with_profile=False, **changes): model = tmp_path / "model with spaces" model.mkdir() (model / "config.json").write_text( @@ -452,11 +547,77 @@ def _local_collection_command(tmp_path, **changes): ) request_path = tmp_path / "request.yaml" root = tmp_path / "plan" - assert cli.main(_init_args(request_path, model=model, **changes)) == 0 + init = _init_args(request_path, model=model, **changes) + if with_profile: + overrides = tmp_path / "resources.json" + overrides.write_text(json.dumps({"fmha_quant_mode": "bfloat16", "kv_cache_dtype": "bfloat16"})) + init.extend( + ( + "--model-config", + str(model / "config.json"), + "--resource-overrides", + str(overrides), + "--tensor-parallel", + "1", + ) + ) + assert cli.main(init) == 0 assert cli.main(["onboard", "plan", "-c", str(request_path), "--output-dir", str(root)]) == 0 return ["onboard", "collect-fpm", "-c", str(request_path), "--output-dir", str(root), "--execute"] +@pytest.mark.parametrize("with_profile", [False, True], ids=["generic", "profile"]) +def test_plan_guides_both_collection_executors_without_launching_workers(tmp_path, monkeypatch, capsys, with_profile): + def unexpected_launch(*args, **kwargs): + pytest.fail("initialization, planning and collection preview must not launch workers") + + monkeypatch.setattr(subprocess, "run", unexpected_launch) + monkeypatch.setitem(sys.modules, "collector.fpm_forward.cli", SimpleNamespace(main=unexpected_launch)) + command = _local_collection_command(tmp_path, with_profile=with_profile) + request = SupportRequest.from_yaml(tmp_path / "request.yaml") + assert (request.fpm_profile is not None) == with_profile + plan = json.loads((tmp_path / "plan/support-plan.json").read_text()) + runtime = next(item for item in plan["prerequisites"] if item["id"] == "runtime") + assert runtime["status"] == "not_checked" + detail = runtime["detail"] + for prerequisite in ( + "Kubernetes", + "caller-owned sbatch/salloc Slurm allocation", + "Pyxis/Enroot", + "explicit --image", + "same absolute path", + "does not download a pinned checkpoint", + ): + assert prerequisite in detail + if with_profile: + assert "observed worker vLLM version" in detail + assert "profile's literal backend version before benchmarking on either executor" in detail + assert "do not verify the loaded checkpoint weights" in detail + else: + assert "does not enforce model/tokenizer revision or vLLM version declarations" in detail + + command.remove("--execute") + capsys.readouterr() + assert cli.main([*command, "--executor", "kubernetes"]) == 0 + assert "--fpm-executor" not in shlex.split(capsys.readouterr().out) + assert cli.main([*command, "--executor", "slurm", "--image", "runtime.sqsh"]) == 0 + preview = shlex.split(capsys.readouterr().out) + assert preview[preview.index("--fpm-executor") + 1] == "slurm" + assert preview[preview.index("--fpm-slurm-container-image") + 1] == "runtime.sqsh" + assert "--plan-only" in preview + assert not (tmp_path / "plan/fpm-artifacts").exists() + assert not (tmp_path / "plan/fpm-checkpoint").exists() + + +@pytest.fixture +def timing_ready(monkeypatch): + # These tests isolate deployment forwarding and frozen-plan resume guards; + # native readiness behavior is exercised with real artifacts separately. + from aisimulate.support import collection_readiness + + monkeypatch.setattr(collection_readiness, "assess_readiness", lambda *a, **k: {"ready_for_full_collection": True}) + + def _synthetic_collector(monkeypatch, version): """Exercise real collection, provenance validation, and publication without GPUs.""" from collector.fpm_forward import planner, runner @@ -573,7 +734,7 @@ def cleanup(self): ], ) def test_formal_collection_checks_published_version_and_preserves_evidence_on_resume( - tmp_path, monkeypatch, capsys, smoke, declared, observed + tmp_path, monkeypatch, capsys, smoke, declared, observed, timing_ready ): executions = _synthetic_collector(monkeypatch, observed) command = _local_collection_command(tmp_path, framework_version=declared) @@ -605,13 +766,18 @@ def test_formal_collection_checks_published_version_and_preserves_evidence_on_re if declared != observed: # A stale matching directory must not hide this campaign's different runtime. (root / "systems/data/h200_sxm/vllm" / declared).mkdir() + readiness_path = root / "fpm-readiness.json" + initial_readiness = json.loads(readiness_path.read_text()) preserved = { - path: path.read_bytes() for path in root.rglob("*") if path.is_file() and path.name not in {"run-manifest.json"} + path: path.read_bytes() + for path in root.rglob("*") + if path.is_file() and path.name != "run-manifest.json" and path != readiness_path } initial_executions = executions.copy() assert cli.main([*command, "--resume"]) == expected assert executions == initial_executions + assert json.loads(readiness_path.read_text()) == {**initial_readiness, "recovery_only": not smoke} assert all(path.read_bytes() == contents for path, contents in preserved.items()) if not smoke: # A completed formal checkpoint remains resumable after raw-artifact reclamation. @@ -622,20 +788,28 @@ def test_formal_collection_checks_published_version_and_preserves_evidence_on_re assert metadata.read_bytes() == preserved[metadata] -def test_formal_collection_requires_publication_evidence(tmp_path, monkeypatch, capsys): +def test_formal_collection_requires_publication_evidence(tmp_path, monkeypatch, capsys, timing_ready): from collector.fpm_forward import runner command = _local_collection_command(tmp_path) - monkeypatch.setattr(runner, "run_collection", lambda *args, **kwargs: []) + + def incomplete_publication(*args, **kwargs): + checkpoint = tmp_path / "plan/fpm-checkpoint/fpm_forward.json" + checkpoint.parent.mkdir(exist_ok=True) + checkpoint.write_text("{}") + return [] + + monkeypatch.setattr(runner, "run_collection", incomplete_publication) capsys.readouterr() assert cli.main(command) == 1 stderr = capsys.readouterr().err assert "aisimulate onboard collect-fpm failed:" in stderr + assert "did not record a completed FPM database publication" in stderr assert "Traceback" not in stderr -def test_deployment_options_reach_frozen_collector_plan(tmp_path, monkeypatch): +def test_deployment_options_reach_frozen_collector_plan(tmp_path, monkeypatch, timing_ready): from collector.fpm_forward import runner calls = [] @@ -677,6 +851,129 @@ def test_deployment_options_reach_frozen_collector_plan(tmp_path, monkeypatch): assert kwargs["database_root"] == str(tmp_path / "plan/systems/data") +def _record_synthetic_publication(root: Path) -> None: + """Complete mocked formal collection's publication contract without GPU data.""" + metadata = root / "synthetic-publication.metadata.json" + version = SupportRequest.from_yaml(root / "request.yaml").identity.framework_version + metadata.write_text(json.dumps({"backend_version": version})) + checkpoint = root / "fpm-checkpoint/fpm_forward.json" + checkpoint.parent.mkdir(exist_ok=True) + payload = json.loads(checkpoint.read_text()) if checkpoint.exists() else {} + payload["database"] = {"status": "passed", "metadata": str(metadata)} + checkpoint.write_text(json.dumps(payload)) + + +def test_slurm_deployment_preview_and_execution_reach_frozen_collector_plan( + tmp_path, monkeypatch, capsys, timing_ready +): + from collector.fpm_forward import runner + + calls = [] + + def capture_collection(plan, **kwargs): + calls.append((plan, kwargs)) + _record_synthetic_publication(tmp_path / "plan") + return [] + + monkeypatch.setattr(runner, "run_collection", capture_collection) + image = "registry.example/fpm@sha256:" + "a" * 64 + mounts = ["/shared/model cache:/models:ro", "/shared/huggingface:/root/.cache/huggingface"] + command = [ + *_local_collection_command(tmp_path), + "--executor", + "slurm", + "--dynamo-version", + "1.2.0", + "--image", + image, + "--transport", + "ib", + ] + for mount in mounts: + command.extend(("--container-mount", mount)) + capsys.readouterr() + + assert cli.main([argument for argument in command if argument != "--execute"]) == 0 + preview = shlex.split(capsys.readouterr().out) + assert preview[preview.index("--fpm-executor") + 1] == "slurm" + assert preview[preview.index("--fpm-slurm-container-image") + 1] == image + assert preview[preview.index("--fpm-slurm-cpus-per-task") + 1] == "16" + assert preview[preview.index("--fpm-slurm-cpu-bind") + 1] == "cores" + assert [ + preview[index + 1] for index, value in enumerate(preview) if value == "--fpm-slurm-container-mount" + ] == mounts + assert "--generator-set" not in preview + assert "--plan-only" in preview + assert calls == [] + + assert cli.main(command) == 0 + [(plan, kwargs)] = calls + assert plan.options.executor == "slurm" + assert plan.options.slurm_container_image == image + assert plan.options.slurm_container_mounts == tuple(mounts) + assert plan.options.slurm_cpus_per_task == 16 + assert plan.options.slurm_cpu_bind == "cores" + assert kwargs["generator_overrides"] == { + "generator_dynamo_version": "1.2.0", + "K8sConfig": {"transport": "ib"}, + } + assert kwargs["artifact_root"] == str(tmp_path / "plan/fpm-artifacts") + assert kwargs["database_root"] == str(tmp_path / "plan/systems/data") + + +def test_explicit_kubernetes_executor_preserves_default_collector_preview(tmp_path, capsys): + command = [argument for argument in _local_collection_command(tmp_path) if argument != "--execute"] + capsys.readouterr() + assert cli.main(command) == 0 + default = capsys.readouterr().out + assert cli.main([*command, "--executor", "kubernetes"]) == 0 + assert capsys.readouterr().out == default + assert "--fpm-executor" not in default + assert "--fpm-slurm" not in default + + +@pytest.mark.parametrize( + "options,diagnostic", + [ + (["--executor", "slurm"], "--executor slurm requires --image"), + (["--executor", "slurm", "--image", "runtime.sqsh", "--namespace", "example"], "--namespace"), + (["--executor", "slurm", "--image", "runtime.sqsh", "--model-cache", "model-cache"], "--model-cache"), + ( + ["--executor", "slurm", "--image", "runtime.sqsh", "--image-pull-secret", "registry"], + "--image-pull-secret", + ), + (["--container-mount", "/shared:/models"], "--container-mount requires --executor slurm"), + (["--cpus-per-task", "16"], "require --executor slurm"), + (["--cpu-bind", "cores"], "require --executor slurm"), + (["--executor", "slurm", "--image", "runtime.sqsh", "--cpus-per-task", "0"], "greater than 0"), + *[ + (["--executor", "slurm", "--image", "runtime.sqsh", "--container-mount", mount], "container mounts") + for mount in ( + "", + " ", + "/one:/models,/two:/cache", + "/one\n:/models", + "/one\r:/models", + "/one\x00:/models", + "/one\t:/models", + "/one\x7f:/models", + ) + ], + ], +) +def test_executor_options_fail_before_collection_for_incompatible_inputs(tmp_path, capsys, options, diagnostic): + command = _local_collection_command(tmp_path) + capsys.readouterr() + with pytest.raises(SystemExit) as error: + cli.main([*command, *options]) + assert error.value.code == 2 + output = capsys.readouterr() + assert diagnostic in output.err + assert "Traceback" not in output.err + assert not (tmp_path / "plan/fpm-checkpoint").exists() + assert not (tmp_path / "plan/fpm-artifacts").exists() + + @pytest.mark.parametrize( "option,value", [ @@ -717,7 +1014,9 @@ def test_deployment_options_reject_invalid_or_unrestricted_inputs(tmp_path, caps ("--image-pull-secret", "first", "second"), ], ) -def test_deployment_resume_requires_the_same_frozen_identity(tmp_path, monkeypatch, capsys, option, first, changed): +def test_deployment_resume_requires_the_same_frozen_identity( + tmp_path, monkeypatch, capsys, option, first, changed, timing_ready +): from collector.fpm_forward import planner, runner command = [*_local_collection_command(tmp_path), "--smoke", option, first] @@ -731,6 +1030,9 @@ def checkpoint_only(plan, **kwargs): plans.append(plan) checkpoint.parent.mkdir(exist_ok=True) checkpoint.write_text(json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": {}})) + campaign = Path(kwargs["artifact_root"]) / plan.sha256[:16] + campaign.mkdir(parents=True, exist_ok=True) + (campaign / "collection-plan.json").write_text(json.dumps(plan.to_dict())) return [] monkeypatch.setattr(planner, "_git_revision", lambda: "test-source-revision") @@ -754,10 +1056,62 @@ def checkpoint_only(plan, **kwargs): assert checkpoint.read_bytes() == saved +@pytest.mark.parametrize( + "changed_options", + [ + ["--image", "registry.example/fpm:changed"], + ["--container-mount", "/shared/checkpoint:/models:ro"], + ["--executor", "kubernetes"], + ["--cpus-per-task", "32"], + ["--cpu-bind", "none"], + ], +) +def test_slurm_collection_resume_rejects_changed_deployment( + tmp_path, monkeypatch, capsys, changed_options, timing_ready +): + from collector.fpm_forward import planner, runner, slurm + + command = [*_local_collection_command(tmp_path), "--executor", "slurm", "--image", "registry.example/fpm:pinned"] + checkpoint = tmp_path / "plan/fpm-checkpoint/fpm_forward.json" + plans = [] + actual_run = runner.run_collection + + def checkpoint_only(plan, **kwargs): + # A synthetic empty checkpoint exercises the public frozen-plan guard; + # no Slurm allocation, GPU execution or silicon measurements are used. + plans.append(plan) + checkpoint.parent.mkdir(exist_ok=True) + checkpoint.write_text(json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": {}})) + campaign = Path(kwargs["artifact_root"]) / plan.sha256[:16] + campaign.mkdir(parents=True, exist_ok=True) + (campaign / "collection-plan.json").write_text(json.dumps(plan.to_dict())) + _record_synthetic_publication(tmp_path / "plan") + return [] + + monkeypatch.setattr(planner, "_git_revision", lambda: "test-source-revision") + monkeypatch.setattr(runner, "run_collection", checkpoint_only) + assert cli.main(command) == 0 + saved = checkpoint.read_bytes() + assert cli.main([*command, "--resume"]) == 0 + assert plans[0].sha256 == plans[1].sha256 + + def unexpected_launch(*args, **kwargs): + pytest.fail("changed deployment must fail before launching GPU work") + + monkeypatch.setattr(runner, "run_collection", actual_run) + monkeypatch.setattr(runner, "KubernetesCellRunner", unexpected_launch) + monkeypatch.setattr(slurm, "SlurmCellRunner", unexpected_launch) + capsys.readouterr() + assert cli.main([*command, "--resume", *changed_options]) == 1 + assert "checkpoint does not match the current frozen plan" in capsys.readouterr().err + assert checkpoint.read_bytes() == saved + + @pytest.mark.parametrize("stage", ["input", "execution"]) @pytest.mark.parametrize("error_type", [RuntimeError, ValueError, OSError, KeyboardInterrupt]) -def test_collector_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, stage, error_type): - from collector.fpm_forward import cli as collector_cli +def test_collector_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, stage, error_type, timing_ready): + from collector import model_cases + from collector.fpm_forward import entry request_path = tmp_path / "request.yaml" root = tmp_path / "plan" @@ -768,9 +1122,11 @@ def test_collector_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, ca def fail(*args, **kwargs): raise error_type("test collection failure") - monkeypatch.setattr(collector_cli, "build_collection_case_plan", lambda **kwargs: SimpleNamespace(model_path=None)) - monkeypatch.setattr(collector_cli, "resolve_run_inputs", fail if stage == "input" else lambda *args: None) - monkeypatch.setattr(collector_cli, "run_resolved", fail) + monkeypatch.setattr(model_cases, "build_collection_case_plan", lambda **kwargs: SimpleNamespace(model_path=None)) + monkeypatch.setattr( + entry, "resolve_run_inputs", fail if stage == "input" else lambda *args: (SimpleNamespace(cells=()), {}) + ) + monkeypatch.setattr(entry, "run_resolved", fail) command = ["onboard", "collect-fpm", "-c", str(request_path), "--output-dir", str(root), "--execute"] if error_type is KeyboardInterrupt: @@ -785,26 +1141,46 @@ def fail(*args, **kwargs): assert "Traceback" not in stderr if stage == "execution" and error_type is not KeyboardInterrupt: assert "aisimulate onboard collect-fpm failed: test collection failure" in stderr + report = json.loads((root / "fpm-readiness.json").read_text()) + assert report["execution_error"] == "test collection failure" + assert report["collector_exit_status"] is None +@pytest.mark.parametrize( + "module", ["collector.fpm_forward.cli", "collector.fpm_forward.entry", "collector.fpm_forward.runner"] +) @pytest.mark.parametrize("error_type", [ImportError, KeyboardInterrupt]) -def test_collector_import_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, error_type): +def test_collector_import_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, error_type, module): + from collector.fpm_forward import entry + command = _local_collection_command(tmp_path) real_import = builtins.__import__ + failures = [] + + def unexpected_launch(*args, **kwargs): + pytest.fail("a collector import failure must not launch GPU work") def fail_collector_import(name, *args, **kwargs): - if name == "collector.fpm_forward.cli": + if name == module: + failures.append(name) raise error_type("collector dependency unavailable") return real_import(name, *args, **kwargs) + monkeypatch.setattr(entry, "run_resolved", unexpected_launch) monkeypatch.setattr(builtins, "__import__", fail_collector_import) capsys.readouterr() assert cli.main(command) == (130 if error_type is KeyboardInterrupt else 1) + assert failures == [module] stderr = capsys.readouterr().err assert "Traceback" not in stderr if error_type is ImportError: assert "aisimulate onboard collect-fpm failed: collector dependency unavailable" in stderr + report = tmp_path / "plan/fpm-readiness.json" + if error_type is KeyboardInterrupt or module != "collector.fpm_forward.cli": + assert not report.exists() + else: + assert json.loads(report.read_text())["execution_error"] == "collector dependency unavailable" def test_malformed_resumed_checkpoint_fails_cleanly_without_launching_collection(tmp_path, capsys): @@ -953,3 +1329,10 @@ def test_installed_module_scripted_init_works_without_a_terminal(tmp_path) -> No assert result.returncode == 0, result.stderr assert SupportRequest.from_yaml(output).identity.model == _REQUIRED["model"] assert "next:" in result.stdout + + +def test_scripted_init_requires_context_evidence(tmp_path, capsys): + with pytest.raises(SystemExit) as error: + cli.main(_init_args(tmp_path / "request.yaml", context_length=None)) + assert error.value.code == 2 + assert "--context-length is required" in capsys.readouterr().err diff --git a/python/aisimulate/tests/unit/test_support_config_profile.py b/python/aisimulate/tests/unit/test_support_config_profile.py new file mode 100644 index 000000000..c93b8059e --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_config_profile.py @@ -0,0 +1,921 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Adapted and modified from vLLM at commit 98dff2a81d747d1dba01a47f939f48c3526d4206: +# https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py + +"""Config onboarding uses declared geometry, without analytical model construction.""" + +from __future__ import annotations + +import builtins +import hashlib +import json +from pathlib import Path + +import pytest + +from aisimulate.support.config_profile import derive_profile, load_model_config +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +def _config(tmp_path, **updates): + # Original small synthetic geometry makes every tensor count hand-checkable. + raw = { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "vocab_size": 64, + "max_position_embeddings": 2048, + "torch_dtype": "bfloat16", + "tie_word_embeddings": False, + "attention_bias": False, + "mlp_bias": False, + } + raw.update(updates) + path = tmp_path / "config.json" + path.write_text(json.dumps(raw), encoding="utf-8") + return load_model_config(path) + + +def _request(kind="dense", **topology): + return SupportRequest.model_validate( + { + "identity": { + "model": "example/configured-model", + "model_revision": "checkpoint-123", + "model_kind": kind, + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "NVLink", + }, + "search": {"context_length": 2048, **topology}, + "workload": {"concurrency": 2}, + } + ) + + +def _runtime(**updates): + return {"kv_cache_dtype": "bfloat16", "fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", **updates} + + +def test_config_hash_and_identity_are_content_facts_not_checkpoint_revision(tmp_path): + config = _config(tmp_path, _name_or_path="source/repository", auto_map={"AutoConfig": "do_not_execute.Config"}) + assert config.sha256 == hashlib.sha256((tmp_path / "config.json").read_bytes()).hexdigest() + assert config.suggestions == { + "model": "source/repository", + "architecture": "LlamaForCausalLM", + "model_kind": "dense", + "context_length": 2048, + "num_experts": 0, + } + draft = derive_profile(config, _request(), _runtime()) + assert draft.profile is not None + assert draft.profile.model == "example/configured-model" + assert draft.profile.model_revision == "checkpoint-123" + assert config.sha256 in draft.profile.provenance + assert config.sha256 in draft.profile.deployments[0].resources.provenance + assert "auto_map" in config.notes + + +@pytest.mark.parametrize("kind", ["dense", "moe"]) +@pytest.mark.parametrize("declared_architecture", [False, True]) +def test_nested_text_geometry_and_architecture_are_independent_of_wrapper_and_encoders( + tmp_path, kind, declared_architecture +): + text = _config(tmp_path).raw + if kind == "moe": + text.update(architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) + architecture = text["architectures"][0] + baseline = derive_profile(_config(tmp_path, **text), _request(kind), _runtime()) + if not declared_architecture: + del text["architectures"] + config = _config( + tmp_path, + architectures=["ExampleMultimodalForConditionalGeneration"], + model_type="example_multimodal", + _name_or_path="example/original-checkpoint", + hidden_size=1024, + n_layer=99, + max_position_embeddings=8192, + num_experts=64, + text_config={**text, "_name_or_path": "example/decoder-base"}, + vision_config={"hidden_size": 512, "num_hidden_layers": 20, "num_experts": 8}, + audio_config={"hidden_size": 256, "num_hidden_layers": 10}, + ) + + draft = derive_profile(config, _request(kind), _runtime()) + + published_architecture = architecture if declared_architecture else "ExampleMultimodalForConditionalGeneration" + assert draft.resolved == {**baseline.resolved, "architecture": published_architecture} + assert draft.profile.architecture == published_architecture + assert config.decoder_architecture == architecture + assert config.suggestions["model"] == "example/original-checkpoint" + assert config.sha256 == hashlib.sha256((tmp_path / "config.json").read_bytes()).hexdigest() + notes = json.loads(draft.profile.provenance)["config_notes"] + assert "text_config" in notes["decoder_config"] + assert f"decoder architecture={architecture}" in notes["decoder_architecture"] + assert "Text decoder only" in notes["modeling_scope"] + assert "full multimodal latency is not modeled" in notes["modeling_scope"] + assert json.loads(draft.profile.deployments[0].resources.provenance)["config_notes"] == notes + + +def test_nested_config_inherits_only_shared_precision_metadata_and_original_identity(tmp_path): + text = _config(tmp_path).raw + del text["torch_dtype"] + config = _config( + tmp_path, + _name_or_path="example/quantized-checkpoint", + auto_map={"AutoConfig": "do_not_execute.Config"}, + text_config={**text, "_name_or_path": "example/decoder-base"}, + quantization_config={"quant_method": "modelopt", "quant_algo": "NVFP4", "kv_cache_scheme": "FP8"}, + ) + draft = derive_profile(config, _request(), _runtime(weights_bytes=1024)) + assert draft.profile is not None + assert draft.resolved["gemm_quant_mode"] == draft.resolved["moe_quant_mode"] == "nvfp4" + assert config.raw["torch_dtype"] == "bfloat16" + assert config.raw["quantization_config"]["quant_algo"] == "NVFP4" + assert config.suggestions["model"] == "example/quantized-checkpoint" + assert "torch_dtype, quantization_config inherited" in config.notes["shared_metadata"] + assert "auto_map" in config.notes + assert config.sha256 in draft.profile.provenance + + +def test_nested_precision_metadata_takes_precedence_as_a_group(tmp_path): + text = _config(tmp_path).raw + del text["torch_dtype"] + config = _config( + tmp_path, + text_config={**text, "dtype": "float16", "quantization_config": {"quant_method": "fp8"}}, + quantization_config={"quant_method": "modelopt", "quant_algo": "NVFP4"}, + ) + draft = derive_profile(config, _request()) + assert draft.resolved["gemm_quant_mode"] == "fp8_static" + assert draft.resolved["moe_quant_mode"] == "fp8" + assert "torch_dtype" not in config.raw + assert "config tensor dtype=float16" in config.notes["dtype"] + assert "shared_metadata" not in config.notes + + +@pytest.mark.parametrize("metadata", [{"hf_quant_config": {"format": "custom"}}, {"quant_algo": "CUSTOM"}]) +def test_shared_unrecognized_quantization_prevents_bfloat16_storage_assumptions(tmp_path, metadata): + config = _config(tmp_path, text_config=_config(tmp_path).raw, **metadata) + draft = derive_profile(config, _request()) + assert {"gemm_quant_mode", "moe_quant_mode"} <= draft.missing.keys() + assert {"weights_bytes"}.isdisjoint(draft.resolved) + for key, value in metadata.items(): + assert config.raw[key] == value + assert key in config.notes["shared_metadata"] + + +@pytest.mark.parametrize("metadata", [{"hf_quant_config": {"format": "custom"}}, {"quant_algo": "CUSTOM"}]) +def test_nested_unrecognized_quantization_is_not_replaced_by_shared_metadata(tmp_path, metadata): + config = _config( + tmp_path, + text_config={**_config(tmp_path).raw, **metadata}, + quantization_config={"quant_method": "fp8"}, + ) + draft = derive_profile(config, _request()) + assert {"gemm_quant_mode", "moe_quant_mode"} <= draft.missing.keys() + assert {"weights_bytes"}.isdisjoint(draft.resolved) + assert "quantization_config" not in config.raw + + +@pytest.mark.parametrize("model_type", [None, "example_custom_text"]) +def test_unknown_nested_decoder_retains_wrapper_identity_and_requires_resource_bounds(tmp_path, model_type): + text = _config(tmp_path).raw + for key in ("architectures", "model_type", "max_position_embeddings"): + del text[key] + text.update(model_max_length=4096, n_routed_experts=4, local_layer_ids=[0], sliding_window_size=128) + if model_type: + text["model_type"] = model_type + config = _config( + tmp_path, + architectures=["ExampleMultimodalForConditionalGeneration"], + model_type="example_multimodal", + text_config=text, + ) + draft = derive_profile(config, _request("moe"), _runtime()) + assert draft.resolved["architecture"] == "ExampleMultimodalForConditionalGeneration" + assert draft.resolved["context_length"] == 4096 + assert draft.resolved["num_experts"] == 4 + assert {"cache_block_sizes"} <= draft.missing.keys() + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + assert draft.resolved["cache_layout"] == "grouped" + assert "wrapper architecture" in draft.sources["architecture"] + assert "assumption" in draft.sources["gemm_quant_mode"] + with pytest.raises(ValueError, match="unsupported linear cache"): + derive_profile(config, _request("moe"), _runtime(kv_bytes_per_token=64, cache_layout="linear")) + complete = derive_profile( + config, + _request("moe"), + _runtime( + weights_bytes=1024, + activations_bytes=2048, + cache_block_sizes={"full_attention": 16, "sliding_attention": 16}, + ), + ) + assert complete.profile is not None + assert len(complete.profile.deployments[0].resources.cache_groups) == 2 + + +def test_missing_text_geometry_does_not_fall_back_to_wrapper_dimensions(tmp_path): + config = _config(tmp_path, text_config={"model_type": "llama", "model_max_length": 2048}) + draft = derive_profile(config, _request(), _runtime()) + assert {"kv_bytes_per_token"} <= draft.missing.keys() + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + assert "hidden_size" not in config.raw + assert "weights_bytes" not in draft.missing + + +def test_known_wrapper_does_not_establish_unknown_nested_decoder_resource_layout(tmp_path): + text = {key: value for key, value in _config(tmp_path).raw.items() if key not in ("architectures", "model_type")} + config = _config(tmp_path, text_config=text) + draft = derive_profile(config, _request(), _runtime()) + assert draft.resolved["architecture"] == "LlamaForCausalLM" + assert config.decoder_architecture is None + assert {"num_experts", "kv_bytes_per_token", "cache_layout"} <= draft.missing.keys() + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + + +@pytest.mark.parametrize("text", [None, {}, [], "decoder", [{"hidden_size": 16}], {"text_config": {"hidden_size": 16}}]) +def test_malformed_or_ambiguous_text_sections_have_actionable_diagnostics(tmp_path, text): + with pytest.raises(ValueError, match="text_config.*decoder.*--fpm-profile"): + _config(tmp_path, text_config=text) + + +def test_nested_architecture_ambiguity_is_not_hidden_by_wrapper_fallback(tmp_path): + text = {**_config(tmp_path).raw, "architectures": ["LlamaForCausalLM", "MixtralForCausalLM"]} + with pytest.raises(ValueError, match="exactly one decoder architecture"): + _config(tmp_path, text_config=text) + + +def test_flat_multimodal_metadata_does_not_change_decoder_estimates(tmp_path): + baseline = derive_profile(_config(tmp_path), _request(), _runtime()) + config = _config( + tmp_path, + vision_config={"hidden_size": 512, "num_hidden_layers": 20}, + audio_config={"hidden_size": 256, "num_hidden_layers": 10}, + ) + draft = derive_profile(config, _request(), _runtime()) + assert draft.resolved == baseline.resolved + assert "Text decoder only" in json.loads(draft.profile.provenance)["config_notes"]["modeling_scope"] + + +@pytest.mark.parametrize("tp,weights,kv", [(1, 13472, 64), (2, 6816, 32), (4, 3744, 32)]) +def test_dense_tensor_counts_include_replicated_norms_and_kv_heads(tmp_path, tp, weights, kv): + # TP1: 2 embedding/head matrices * 64*16 + 2 layers * + # (Q/O 2*16*16 + K/V 2*16*8 + MLP 3*16*32 + norms 2*16) + final norm 16. + draft = derive_profile(_config(tmp_path), _request(tensor_parallel=tp), _runtime()) + assert draft.missing == {} + resources = draft.profile.deployments[0].resources + assert resources.weights_bytes is None + assert draft.resolved["weights_bytes"] == weights + assert resources.kv_bytes_per_token == kv + assert resources.max_num_tokens == 8192 + assert resources.max_batch_size == 256 + assert "estimate" in draft.sources["weights_bytes"] + assert "exact" in draft.sources["kv_bytes_per_token"] + assert "estimate" in draft.sources["runtime_overhead_bytes"] + assert "sha256" in draft.sources["runtime_overhead_bytes"] + + +def test_weight_tying_biases_and_qk_norm_have_independent_counts(tmp_path): + request = _request(tensor_parallel=2) + base = derive_profile(_config(tmp_path), request, _runtime()).resolved["weights_bytes"] + tied = derive_profile(_config(tmp_path, tie_word_embeddings=True), request, _runtime()) + assert tied.resolved["weights_bytes"] == base - 64 // 2 * 16 * 2 + biased = derive_profile(_config(tmp_path, attention_bias=True, mlp_bias=True), request, _runtime()) + # Per layer: QKV (8+4+4), output (16), gated MLP (16+16+16), all two-byte. + assert biased.resolved["weights_bytes"] == base + 2 * (16 + 16 + 48) * 2 + qwen = derive_profile( + _config(tmp_path, architectures=["Qwen3ForCausalLM"], model_type="qwen3"), request, _runtime() + ) + assert qwen.resolved["weights_bytes"] == base + 2 * (2 * 4) * 2 + + +@pytest.mark.parametrize("tp,expected_bias_bytes", [(1, 128), (2, 64), (4, 48)]) +def test_qwen3_attention_bias_counts_qkv_without_output_bias(tmp_path, tp, expected_bias_bytes): + request = _request(tensor_parallel=tp) + facts = {"architectures": ["Qwen3ForCausalLM"], "model_type": "qwen3"} + base = derive_profile(_config(tmp_path, **facts), request, _runtime()) + biased = derive_profile(_config(tmp_path, **facts, attention_bias=True), request, _runtime()) + # Two layers, two-byte Q/K/V biases: widths (16,8,8), (8,4,4), or (4,4,4). + # The output projection is bias-free; KV heads replicate at TP4. + assert biased.resolved["weights_bytes"] - base.resolved["weights_bytes"] == expected_bias_bytes + + +@pytest.mark.parametrize( + "topology,weights,kv,parallel", + [ + ({"tensor_parallel": 2}, 16288, 32, (2, 1, 1, 2, 1, 1)), + ( + {"attention_data_parallel": 2, "moe_tensor_parallel": 1, "moe_expert_parallel": 2}, + 19872, + 64, + (1, 1, 2, 1, 2, 1), + ), + ({"tensor_parallel": 2, "moe_tensor_parallel": 1, "moe_expert_parallel": 2}, 16288, 32, (2, 1, 1, 1, 2, 1)), + ], +) +def test_moe_resources_are_rank_local_for_tp_dep_and_tep(tmp_path, topology, weights, kv, parallel): + config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) + draft = derive_profile(config, _request("moe", **topology), _runtime(comm_overhead_bytes=321)) + assert draft.missing == {} + deployment = draft.profile.deployments[0] + assert deployment.parallel_tuple == parallel + assert deployment.resources.weights_bytes is None + assert draft.resolved["weights_bytes"] == weights + assert deployment.resources.kv_bytes_per_token == kv + assert deployment.resources.comm_overhead_bytes is None + assert draft.resolved["comm_overhead_bytes"] == 321 + assert deployment.resources.memory_source == "pending" + + +def test_activation_estimate_uses_declared_scheduler_envelope(tmp_path): + config = _config(tmp_path, hidden_size=4096, intermediate_size=8192, num_attention_heads=32, head_dim=128) + small = derive_profile(config, _request(tensor_parallel=2), _runtime(max_num_tokens=16)) + large = derive_profile(config, _request(tensor_parallel=2), _runtime(max_num_tokens=4096)) + assert small.resolved["activations_bytes"] == 70 * 1024 * 1024 + assert large.resolved["activations_bytes"] == 2 * 4096 * 4096 * 6.5 + assert "estimate" in large.sources["activations_bytes"] + assert "max_num_tokens=4096" in large.sources["activations_bytes"] + + +def test_mha_and_weight_padding_have_separate_supported_domains(tmp_path): + draft = derive_profile(_config(tmp_path, num_key_value_heads=4, vocab_size=65), _request(), _runtime()) + assert draft.resolved["kv_bytes_per_token"] == 128 + assert "weights_bytes" not in draft.resolved + assert "weights_bytes" not in draft.missing + assert draft.profile.deployments[0].resources.memory_source == "pending" + + +@pytest.mark.parametrize("field", ["gemm_quant_mode", "kv_cache_dtype", "weights_bytes", "activations_bytes"]) +def test_explicit_overrides_survive_derivation_and_are_recorded_per_field(tmp_path, field): + value = "fp8" if "mode" in field or "dtype" in field else 123456 + draft = derive_profile(_config(tmp_path), _request(), _runtime(**{field: value})) + assert draft.resolved[field] == value + assert "user override" in draft.sources[field] + + +def test_hardware_defaults_do_not_reuse_tp_communication_for_dep(tmp_path): + config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) + request = _request("moe", attention_data_parallel=2, moe_tensor_parallel=1, moe_expert_parallel=2) + draft = derive_profile(config, request, _runtime()) + assert "comm_overhead_bytes" not in draft.resolved + assert "comm_overhead_bytes" not in draft.missing + assert draft.resolved["runtime_overhead_bytes"] == 3758096384 + + +def test_unknown_hardware_has_no_fallback_reservations(tmp_path): + request = _request() + request.identity.gpu = "custom_gpu" + draft = derive_profile(_config(tmp_path), request, _runtime()) + assert {"comm_overhead_bytes", "runtime_overhead_bytes"}.isdisjoint(draft.resolved) + + +def test_dynamic_kv_scales_need_explicit_accounting(tmp_path): + config = _config( + tmp_path, + quantization_config={ + "quant_method": "fp8", + "kv_cache_scheme": {"num_bits": 8, "type": "float", "dynamic": True}, + }, + ) + draft = derive_profile(config, _request()) + assert draft.resolved["kv_cache_dtype"] == "fp8" + assert "kv_bytes_per_token" in draft.missing + assert "scales" in draft.missing["kv_bytes_per_token"] + + +@pytest.mark.parametrize("cache_type", [[], {}, True, 8, 1.5, "", "\x00"]) +def test_malformed_kv_cache_type_is_rejected_at_config_load(tmp_path, cache_type): + with pytest.raises(ValueError, match=r"quantization_config\.kv_cache_scheme\.type"): + _config(tmp_path, quantization_config={"kv_cache_scheme": {"num_bits": 8, "type": cache_type}}) + + +@pytest.mark.parametrize( + "scheme,dtype", + [ + ("FP8", "fp8"), + ({"num_bits": 8, "type": "float"}, "fp8"), + ({"num_bits": 8, "type": "int"}, "int8"), + ({"num_bits": 8, "type": None}, None), + ({"num_bits": 8}, None), + ({"num_bits": 8, "type": "unknown"}, None), + ], +) +def test_kv_cache_type_preserves_known_precision_and_unresolved_metadata(tmp_path, scheme, dtype): + config = _config(tmp_path, quantization_config={"kv_cache_scheme": scheme}) + draft = derive_profile(config, _request()) + if dtype is None: + assert "kv_cache_dtype" not in draft.resolved + assert "kv_cache_dtype" in draft.missing + else: + assert draft.resolved["kv_cache_dtype"] == dtype + assert draft.resolved["kv_bytes_per_token"] == 32 + + +def test_weight_quantization_does_not_select_runtime_attention_or_kv_dtype(tmp_path): + config = _config(tmp_path, quantization_config={"quant_method": "fp8", "weight_block_size": [128, 128]}) + draft = derive_profile(config, _request()) + assert draft.resolved["gemm_quant_mode"] == "fp8_block" + assert draft.resolved["moe_quant_mode"] == "fp8_block" + assert {"fmha_quant_mode", "kv_cache_dtype"} <= draft.missing.keys() + assert {"weights_bytes"}.isdisjoint(draft.resolved) + assert draft.resolved["comm_quant_mode"] == "half" + assert "collector FPM identity default" in draft.sources["comm_quant_mode"] + assert draft.profile is None + assert "weights_bytes" not in draft.missing + + +@pytest.mark.parametrize( + "quant,gemm,moe", + [ + ({"quant_method": "fp8"}, "fp8_static", "fp8"), + ({"quant_method": "fp8", "activation_scheme": "static"}, "fp8_static", "fp8"), + ({"quant_method": "fp8", "activation_scheme": "dynamic"}, "fp8", "fp8"), + ( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "config_groups": {"g": {"input_activations": {"dynamic": True}}}, + }, + "fp8", + "fp8", + ), + ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8_block", "fp8_block"), + ({"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["self_attn*"]}, "nvfp4", "nvfp4"), + ], +) +def test_config_quantization_uses_existing_fpm_checkpoint_identity(tmp_path, quant, gemm, moe): + from aisimulate.support.fpm import fpm_cli_args + from aisimulate_core.sdk.models.helpers import _infer_quant_modes_from_raw_config + from aisimulate_core.sdk.utils import _infer_quantization_fields + + config = _config(tmp_path, quantization_config=quant) + draft = derive_profile(config, _request(), _runtime(weights_bytes=1024)) + assert draft.resolved["gemm_quant_mode"] == gemm + assert draft.resolved["moe_quant_mode"] == moe + # Compare the real collector's identity helper, not a duplicate test formula. + existing = _infer_quant_modes_from_raw_config({**config.raw, **_infer_quantization_fields(config.raw)}) + assert existing["gemm_quant_mode"].name == gemm + assert existing["moe_quant_mode"].name == moe + request = SupportRequest.model_validate({**_request().model_dump(), "fpm_profile": draft.profile}) + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + assert command[command.index("--fpm-weight-quantizations") + 1] == gemm + assert draft.profile.deployments[0].match_identity()[:2] == [gemm, moe] + assert "checkpoint/FPM" in draft.sources["gemm_quant_mode"] + + +def test_unknown_quantization_never_falls_back_to_residual_bfloat16(tmp_path): + config = _config(tmp_path, quantization_config={"quant_method": "custom", "quant_algo": "custom"}) + draft = derive_profile(config, _request()) + assert {"gemm_quant_mode", "moe_quant_mode"} <= set(draft.missing) + assert {"weights_bytes"}.isdisjoint(draft.resolved) + + +def test_missing_fields_are_complete_and_unknown_layout_can_use_explicit_resources(tmp_path): + config = _config(tmp_path, architectures=["NewDecoderForCausalLM"], model_type="new_decoder") + assert "num_experts" not in config.suggestions + draft = derive_profile(config, _request()) + assert {"num_experts", "kv_bytes_per_token", "cache_layout"} <= set(draft.missing) + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + complete = derive_profile( + config, + _request(), + _runtime( + num_experts=0, weights_bytes=1000, activations_bytes=2000, kv_bytes_per_token=32, cache_layout="linear" + ), + ) + assert complete.profile is not None + assert "user override" in complete.sources["weights_bytes"] + + +@pytest.mark.parametrize( + "updates,match", + [ + ({"hidden_size": True}, "hidden_size"), + ({"num_hidden_layers": 2.0}, "num_hidden_layers"), + ({"num_attention_heads": 0}, "num_attention_heads"), + ({"architectures": ["LlamaForCausalLM", "OtherForCausalLM"]}, "architectures"), + ({"n_layer": 4}, "conflicting"), + ({"model_max_length": 4096}, "conflicting context_length"), + ({"dtype": "float32"}, "conflicting"), + ({"num_key_value_heads": 3}, "num_key_value_heads"), + ({"attention_bias": "false"}, "attention_bias"), + ({"model_type": "qwen3"}, "model_type"), + ({"quantization_config": []}, "quantization_config"), + ({"num_mtp_modules": -1}, "num_mtp_modules"), + ( + {"quantization_config": {"config_groups": {"g": {"weights": {"dynamic": "false"}}}}}, + "dynamic", + ), + ], +) +def test_corrupt_or_ambiguous_config_fails_before_profile_overrides(tmp_path, updates, match): + with pytest.raises(ValueError, match=match): + _config(tmp_path, **updates) + + +@pytest.mark.parametrize("raw", ['{"hidden_size":16,"hidden_size":32}', '{"hidden_size":NaN}', "[]"]) +def test_json_must_be_finite_object_with_unique_keys(tmp_path, raw): + path = tmp_path / "config.json" + path.write_text(raw) + with pytest.raises(ValueError): + load_model_config(path) + + +@pytest.mark.parametrize( + "overrides", + [ + {"weights_byte": 7}, + {"weights_bytes": True}, + {"weights_bytes": 7.0}, + {"weights_bytes": -1}, + {"weights_bytes": 2**53 + 1}, + {"max_num_tokens": 0}, + {"kv_bytes_per_token": 0}, + {"num_experts": -1}, + {"cache_layout": "hybrid"}, + {"gemm_quant_mode": "fp16"}, + {"kv_cache_dtype": "auto"}, + {"provenance": " "}, + {"context_length": "2048"}, + {"weights_bytes": 2**53, "activations_bytes": 1}, + ], +) +def test_overrides_are_strict_even_before_identity_is_complete(tmp_path, overrides): + with pytest.raises(ValueError): + derive_profile(_config(tmp_path), None, overrides) + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize( + "updates", + [ + {"sliding_window": 128}, + {"layer_types": ["full_attention", "linear_attention"]}, + {"is_encoder_decoder": True}, + {"ssm_cfg": {"d_state": 16}}, + {"kv_lora_rank": 8}, + ], +) +def test_known_incompatible_state_cannot_be_papered_over_with_overrides(tmp_path, updates, nested): + if nested: + updates = {"text_config": {**_config(tmp_path).raw, **updates}} + with pytest.raises(ValueError, match="unsupported"): + config = _config(tmp_path, **updates) + derive_profile(config, _request(), _runtime(weights_bytes=1, kv_bytes_per_token=1, cache_layout="linear")) + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("preview", [False, True]) +@pytest.mark.parametrize("explicit_resources", [False, True]) +@pytest.mark.parametrize( + "facts,architecture,experts,match", + [ + ({"num_experts": 4}, "LlamaForCausalLM", 4, "dense architecture"), + ({}, "MixtralForCausalLM", 0, "MoE architecture"), + ({"kv_lora_rank": 8}, "LlamaForCausalLM", 0, "cache/projection modifiers"), + ({}, "MambaForCausalLM", 0, "recurrent/hybrid"), + ({"num_experts_per_tok": 8}, "MixtralForCausalLM", 4, "routed expert count"), + ], +) +def test_late_architecture_and_expert_inputs_cannot_bypass_structural_guards( + tmp_path, preview, explicit_resources, facts, architecture, experts, match, nested +): + config = _config(tmp_path, architectures=None, model_type="unknown", **facts) + if nested: + config = _config(tmp_path, architectures=None, model_type="unknown_wrapper", text_config=config.raw) + overrides = _runtime(architecture=architecture, num_experts=experts) + if explicit_resources: + overrides.update(weights_bytes=1, kv_bytes_per_token=1, cache_layout="linear") + with pytest.raises(ValueError, match=match): + derive_profile(config, None if preview else _request("moe" if experts else "dense"), overrides) + # Supplying the same effective facts in the original config also fails. + with pytest.raises(ValueError, match=match): + _config(tmp_path, **{**facts, "architectures": [architecture], "model_type": "unknown", "num_experts": experts}) + + +@pytest.mark.parametrize("field", ["num_experts_per_tok", "first_k_dense_replace"]) +@pytest.mark.parametrize("late_experts", [False, True]) +@pytest.mark.parametrize("preview", [False, True]) +def test_optional_null_moe_metadata_matches_absent_metadata(tmp_path, field, late_experts, preview): + facts = {"architectures": ["MixtralForCausalLM"], "model_type": "mixtral"} + overrides = _runtime() + if late_experts: + overrides["num_experts"] = 4 + else: + facts["num_local_experts"] = 4 + request = None if preview else _request("moe") + absent = derive_profile(_config(tmp_path, **facts), request, overrides) + config = _config(tmp_path, **facts, **{field: None}) + + draft = derive_profile(config, request, overrides) + + assert config.raw[field] is None + assert field not in config.notes + assert draft.resolved == absent.resolved + assert draft.missing == absent.missing + assert draft.resolved["num_experts"] == 4 + assert (draft.profile is None) == preview + + +@pytest.mark.parametrize("field", ["num_experts_per_tok", "first_k_dense_replace"]) +@pytest.mark.parametrize("value", [True, False, -1, 1.5, "1"]) +def test_optional_moe_metadata_rejects_noninteger_or_negative_values(tmp_path, field, value): + with pytest.raises(ValueError, match=field): + _config(tmp_path, **{field: value}) + + +@pytest.mark.parametrize( + "facts,match", + [ + ({"num_experts_per_tok": 0}, "num_experts_per_tok"), + ({"num_experts_per_tok": 5}, "routed expert count"), + ({"first_k_dense_replace": 3}, "num_hidden_layers"), + ], +) +def test_optional_moe_metadata_preserves_count_validation(tmp_path, facts, match): + with pytest.raises(ValueError, match=match): + _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4, **facts) + + +def test_bundled_modelopt_string_kv_scheme_preserves_usable_geometry(): + path = ( + Path(__file__).parents[2] / "src/aisimulate_core/model_configs/Qwen--Qwen3-32B-FP8-Static-PerTensor_config.json" + ) + config = load_model_config(path) + draft = derive_profile(config, _request(tensor_parallel=4)) + assert config.raw["quantization_config"]["kv_cache_scheme"] == "FP8" + assert config.suggestions["architecture"] == "Qwen3ForCausalLM" + assert draft.resolved["gemm_quant_mode"] == "fp8_static" + assert draft.resolved["moe_quant_mode"] == draft.resolved["kv_cache_dtype"] == "fp8" + assert draft.resolved["kv_bytes_per_token"] == 2 * 64 * 2 * 128 + assert {"fmha_quant_mode"} <= draft.missing.keys() + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + + +def test_bundled_custom_layout_keeps_null_geometry_unresolved(): + path = ( + Path(__file__).parents[2] + / "src/aisimulate_core/model_configs/nvidia--Llama-3_3-Nemotron-Super-49B-v1_config.json" + ) + config = load_model_config(path) + draft = derive_profile(config, _request(), _runtime()) + assert config.suggestions["architecture"] == "DeciLMForCausalLM" + assert config.suggestions["context_length"] == 131072 + assert config.raw["intermediate_size"] is None + assert "intermediate_size" not in config.notes + assert {"kv_bytes_per_token", "num_experts"} <= draft.missing.keys() + assert {"weights_bytes", "activations_bytes"}.isdisjoint(draft.resolved) + + +def test_selected_topology_and_context_must_match_structural_config(tmp_path): + config = _config(tmp_path) + with pytest.raises(ValueError, match="attention heads"): + derive_profile(config, _request(tensor_parallel=8), _runtime(weights_bytes=7, kv_bytes_per_token=7)) + with pytest.raises(ValueError, match="num_experts"): + derive_profile(config, _request(), _runtime(num_experts=2)) + with pytest.raises(ValueError, match="context_length"): + derive_profile(config, _request(), _runtime(context_length=4096)) + + +@pytest.mark.parametrize( + "filename,architecture,context,quant,kv", + [ + ("MiniMaxAI--MiniMax-M2.7_config.json", "MiniMaxM2ForCausalLM", 204800, "fp8_block", None), + ("nvidia--MiniMax-M2.7-NVFP4_config.json", "MiniMaxM2ForCausalLM", 196608, "nvfp4", "fp8"), + ("zai-org--GLM-5.2_config.json", "GlmMoeDsaForCausalLM", 1048576, "bfloat16", None), + ("nvidia--GLM-5.2-NVFP4_config.json", "GlmMoeDsaForCausalLM", 1048576, "nvfp4", "fp8"), + ], +) +def test_real_minimax_and_glm_extract_independent_facts_without_naive_storage_rules( + filename, architecture, context, quant, kv +): + path = Path(__file__).parents[2] / "src/aisimulate_core/model_configs" / filename + config = load_model_config(path) + draft = derive_profile(config, _request("moe", tensor_parallel=4)) + assert config.suggestions["architecture"] == architecture + assert config.suggestions["context_length"] == context + assert config.suggestions["num_experts"] == 256 + assert config.suggestions["model_kind"] == "moe" + assert draft.resolved["gemm_quant_mode"] == quant + assert draft.resolved.get("kv_cache_dtype") == kv + assert "fmha_quant_mode" in draft.missing + assert "weights_bytes" not in draft.resolved + assert "weights_bytes" not in draft.missing + if architecture == "GlmMoeDsaForCausalLM": + assert "kv_bytes_per_token" in draft.missing + assert "DSA" in draft.missing["kv_bytes_per_token"] + else: + finished_kv = derive_profile(config, _request("moe", tensor_parallel=4), {"kv_cache_dtype": "bfloat16"}) + assert finished_kv.resolved["kv_bytes_per_token"] == 62 * 2 * 2 * 128 * 2 + + +@pytest.mark.parametrize("nested", [False, True]) +def test_config_route_does_not_import_model_construction_or_remote_code(tmp_path, monkeypatch, nested): + original = builtins.__import__ + + def checked(name, *args, **kwargs): + assert not name.startswith(("transformers", "huggingface_hub", "collector")), name + assert ".sdk.models" not in name, name + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", checked) + config = _config(tmp_path) + if nested: + config = _config(tmp_path, text_config=config.raw, vision_config={"hidden_size": 512}) + draft = derive_profile(config, _request(), _runtime()) + assert draft.profile is not None + + +@pytest.mark.parametrize("mixed", [False, True]) +def test_windowed_cache_derives_separate_group_pages_without_linear_rate(tmp_path, mixed): + updates = {"sliding_window": 128} + if mixed: + updates["layer_types"] = ["full_attention", "sliding_attention"] + config = _config(tmp_path, **updates) + missing = derive_profile(config, _request(tensor_parallel=2), _runtime()) + assert missing.resolved["cache_layout"] == "grouped" + assert "cache_block_sizes" in missing.missing + assert "kv_bytes_per_token" not in missing.missing + blocks = {"sliding_attention": 8, **({"full_attention": 16} if mixed else {})} + draft = derive_profile(config, _request(tensor_parallel=2), _runtime(cache_block_sizes=blocks)) + assert draft.profile is not None + resource = draft.profile.deployments[0].resources + assert resource.kv_bytes_per_token is None + groups = {group.name: group for group in resource.cache_groups} + assert groups["sliding_attention"].page_size_bytes == (128 if mixed else 256) + assert groups["sliding_attention"].sliding_window == 128 + if mixed: + assert groups["full_attention"].page_size_bytes == 256 + assert groups["full_attention"].sliding_window is None + serialized = draft.profile.model_dump(mode="json") + assert "kv_bytes_per_token" not in serialized["deployments"][0]["resources"] + assert type(draft.profile).model_validate(serialized) == draft.profile + assert "minimum packed" in draft.sources["cache_groups"] + + +def _synthetic_inkling(tmp_path, **updates): + # Original small geometry. Expected packing/TP constraints follow the pinned + # Apache-2.0 vLLM implementation (copyright vLLM contributors): + # https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py + # Modified for metadata-only tests; no checkpoint or upstream fixture copied. + text = { + "hidden_size": 16, + "num_hidden_layers": 8, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 4, + "local_layer_ids": [0, 1, 2, 4, 5, 6], + "swa_num_attention_heads": 8, + "swa_num_key_value_heads": 4, + "swa_head_dim": 4, + "sliding_window_size": 128, + "use_sconv": True, + "sconv_kernel_size": 4, + "model_max_length": 4096, + "n_routed_experts": 8, + "torch_dtype": "bfloat16", + **updates, + } + return _config( + tmp_path, architectures=["InklingForConditionalGeneration"], model_type="inkling_mm_model", text_config=text + ) + + +@pytest.mark.parametrize("dtype,attention_factor", [("bfloat16", 2), ("fp8", 1)]) +@pytest.mark.parametrize( + "topology", + [ + {"tensor_parallel": 2}, + {"tensor_parallel": 2, "attention_data_parallel": 1, "moe_tensor_parallel": 1, "moe_expert_parallel": 2}, + {"tensor_parallel": 1, "attention_data_parallel": 8, "moe_tensor_parallel": 1, "moe_expert_parallel": 8}, + ], +) +def test_inkling_mixed_kv_and_bfloat16_convolution_are_rank_local(tmp_path, dtype, attention_factor, topology): + config = _synthetic_inkling(tmp_path) + draft = derive_profile( + config, + _request("moe", **topology), + _runtime( + kv_cache_dtype=dtype, + weights_bytes=4096, + activations_bytes=1024, + comm_overhead_bytes=0, + cache_block_sizes={"full_attention": 16, "sliding_attention": 8}, + ), + ) + assert draft.profile is not None + groups = {group.name: group for group in draft.profile.deployments[0].resources.cache_groups} + rank_factor = 2 // topology["tensor_parallel"] + assert groups["full_attention"].page_size_bytes == 256 * attention_factor * rank_factor + assert groups["sliding_attention"].page_size_bytes == 768 * attention_factor * rank_factor + assert groups["full_convolution"].page_size_bytes == 512 * rank_factor + assert groups["sliding_convolution"].page_size_bytes == 1536 * rank_factor + assert groups["full_convolution"].sliding_window == groups["sliding_convolution"].sliding_window == 4 + assert groups["full_attention"].num_layers == 2 + assert groups["sliding_attention"].num_layers == 6 + assert "rank_local_heads" in draft.sources["cache_groups"] + assert '"dtype": "bfloat16"' in draft.sources["cache_groups"] + + +def test_inkling_convolution_caps_attention_tp_without_limiting_dep_width(tmp_path): + config = _synthetic_inkling(tmp_path) + with pytest.raises(ValueError, match="attention TP <= 2"): + derive_profile(config, _request("moe", tensor_parallel=4), _runtime()) + dep = derive_profile( + config, _request("moe", attention_data_parallel=8, moe_tensor_parallel=1, moe_expert_parallel=8), _runtime() + ) + assert dep.resolved["cache_layout"] == "grouped" + assert "cache_block_sizes" in dep.missing + + +@pytest.mark.parametrize("flag", [False, "absent"]) +def test_inkling_cannot_silently_omit_runtime_convolution_state(tmp_path, flag): + _synthetic_inkling(tmp_path) + path = tmp_path / "config.json" + document = json.loads(path.read_text()) + if flag == "absent": + document["text_config"].pop("use_sconv") + else: + document["text_config"]["use_sconv"] = flag + path.write_text(json.dumps(document)) + with pytest.raises(ValueError, match="allocates convolution state in every decoder layer"): + load_model_config(path) + + +@pytest.mark.parametrize("local_ids", [None, "absent"]) +def test_inkling_nominal_window_does_not_enable_absent_local_layers(tmp_path, local_ids): + _synthetic_inkling(tmp_path) + path = tmp_path / "config.json" + document = json.loads(path.read_text()) + if local_ids == "absent": + document["text_config"].pop("local_layer_ids") + else: + document["text_config"]["local_layer_ids"] = local_ids + path.write_text(json.dumps(document)) + draft = derive_profile(load_model_config(path), _request("moe"), _runtime(cache_block_sizes={"full_attention": 16})) + assert [group["name"] for group in draft.resolved["cache_groups"]] == ["full_attention", "full_convolution"] + assert all(group["num_layers"] == 8 for group in draft.resolved["cache_groups"]) + + +def test_group_override_preserves_runtime_padding_and_required_states(tmp_path): + config = _synthetic_inkling(tmp_path) + overrides = _runtime( + weights_bytes=4096, activations_bytes=1024, cache_block_sizes={"full_attention": 16, "sliding_attention": 8} + ) + initial = derive_profile(config, _request("moe", tensor_parallel=2), overrides) + groups = initial.resolved["cache_groups"] + groups[0]["page_size_bytes"] += 128 + overrides.pop("cache_block_sizes") + overrides["cache_groups"] = groups + edited = derive_profile(config, _request("moe", tensor_parallel=2), overrides) + assert edited.profile is not None + assert edited.resolved["cache_groups"][0]["page_size_bytes"] == 640 + assert "user override" in edited.sources["cache_groups"] + assert "cache_block_sizes" not in edited.missing + with pytest.raises(ValueError, match="layer counts and retention windows"): + derive_profile(config, _request("moe", tensor_parallel=2), {**overrides, "cache_groups": groups[:2]}) + + +@pytest.mark.parametrize( + "updates,match", + [ + ({"sliding_window": 64, "sliding_window_size": 128}, "conflicting sliding-window"), + ({"sliding_window": 64, "local_layer_ids": [0, 0]}, "unique indexes"), + ({"sliding_window": 64, "local_layer_ids": [2]}, "below num_hidden_layers"), + ( + {"sliding_window": 64, "local_layer_ids": [0], "layer_types": ["full_attention", "sliding_attention"]}, + "conflicting layer_types", + ), + ({"sliding_window": 64, "local_layer_ids": [0], "use_sliding_window": False}, "conflicts"), + ({"layer_types": ["full_attention", "sliding_attention"]}, "positive sliding_window"), + ({"use_sconv": True, "sconv_kernel_size": 4}, "unsupported use_sconv"), + ({"use_sliding_window": True}, "requires a positive sliding_window"), + ({"sliding_window": 128, "attn_type_list": [1, 1]}, "conflicts"), + ], +) +def test_cache_alias_conflicts_and_unknown_states_fail_at_intake(tmp_path, updates, match): + with pytest.raises(ValueError, match=match): + _config(tmp_path, **updates) + + +def test_grouped_sizing_requires_declared_kv_dtype_and_dynamic_scale_bounds(tmp_path): + config = _config(tmp_path, sliding_window=128) + unknown_dtype = derive_profile(config, _request(), {"cache_block_sizes": {"sliding_attention": 16}}) + assert {"kv_cache_dtype", "cache_groups"} <= unknown_dtype.missing.keys() + assert "cache_groups" not in unknown_dtype.resolved + dynamic = _config(tmp_path, sliding_window=128, quantization_config={"kv_cache_scheme": {"dynamic": True}}) + unknown_scales = derive_profile(dynamic, _request(), _runtime(cache_block_sizes={"sliding_attention": 16})) + assert "cache_groups" in unknown_scales.missing + assert "nonstandard KV" in unknown_scales.missing["cache_groups"] diff --git a/python/aisimulate/tests/unit/test_support_execution_evidence_compatibility.py b/python/aisimulate/tests/unit/test_support_execution_evidence_compatibility.py new file mode 100644 index 000000000..be6f0faa3 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_execution_evidence_compatibility.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Historical prepared campaigns retain their frozen observation semantics.""" + +import json + +import pytest +from collector.fpm_forward import repeatability + +from aisimulate import main as cli +from aisimulate.support import validation_workflow as workflow +from aisimulate.support.plan import create_plan +from aisimulate.support.schema import SupportRequest + +from .collector.test_fpm_measurement_evidence import _remove_execution_protocol +from .test_onboard_finalization_quality import assessed_collection # noqa: F401 + +pytestmark = pytest.mark.unit + + +def _legacy_execution(case): + case["execution_evidence"] = False + for path in case["root"].rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + _remove_execution_protocol(payload) + path.write_text(json.dumps(payload)) + + +def test_fresh_v2_public_quality_cannot_pass_without_execution_observations(assessed_collection): # noqa: F811 + case = assessed_collection + _legacy_execution(case) + assert cli.main([*case["validation_args"], "--execute"]) == 1 + report = json.loads(case["report"].read_text()) + selection = json.loads((case["report"].parent / "repeatability-selection.json").read_text()) + assert selection["observation_evidence_version"] == 2 + assert all( + "execution_evidence" not in cell["source_measurement"]["measurement_protocol"] for cell in selection["cells"] + ) + assert all(cell["source_measurement"]["status"] == "unestablished" for cell in selection["cells"]) + assert report["status"] == report["gates"]["execution"]["status"] == "incomplete" + + +@pytest.mark.parametrize("policy", ["runtime", "explicit"]) +@pytest.mark.parametrize("execute", [False, True]) +def test_ordinary_collection_reports_graph_advisory_before_launch( + assessed_collection, # noqa: F811 + tmp_path, + monkeypatch, + capsys, + policy, + execute, +): + from collector.fpm_forward import capabilities, entry + + case = assessed_collection + payload = case["request"].model_dump(mode="json") + payload["collection"].update( + prefill_cudagraph_policy=policy, max_prefill_cudagraph_size=2048 if policy == "explicit" else None + ) + request = SupportRequest.model_validate(payload) + root = tmp_path / "graph-advisory" + create_plan(request, root) + capsys.readouterr() + calls = [] + monkeypatch.setattr(capabilities, "load_model_config", lambda *args, **kwargs: case["plan"].capability.model_config) + + def launch(*args): + output = capsys.readouterr() + assert ("prefill only" in output.err) == (policy == "explicit") + calls.append(args) + return [] + + monkeypatch.setattr(entry, "run_resolved", launch) + args = ["onboard", "collect-fpm", "--config", str(root / "request.yaml"), "--output-dir", str(root)] + result = cli.main([*args, *(["--execute", "--smoke"] if execute else [])]) + if execute: + # The launch was reached; this synthetic launcher does not publish readiness artifacts. + assert calls + assert result == 1 + else: + output = capsys.readouterr() + assert result == 0 + assert ("prefill only" in output.err) == (policy == "explicit") + assert "collector.fpm_forward" in output.out + + +@pytest.mark.parametrize("mutation", ["running", "future_order"]) +def test_public_quality_cannot_pass_impossible_warmup_history(assessed_collection, monkeypatch, mutation): # noqa: F811 + case = assessed_collection + + def mutate(root): + for path in root.rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + record = payload["warmup_evidence"]["records"][0] + if mutation == "running": + record.update(status="running", forward_index_end=None) + else: + record["completed_points_before"] = 99 + path.write_text(json.dumps(payload)) + + mutate(case["root"]) + collect = repeatability.run_collection + + def mutated_repeat(plan, **kwargs): + from pathlib import Path + + result = collect(plan, **kwargs) + mutate(Path(kwargs["artifact_root"])) + return result + + monkeypatch.setattr(repeatability, "run_collection", mutated_repeat) + if mutation == "future_order": + with pytest.raises(SystemExit, match="2"): + cli.main([*case["validation_args"], "--execute"]) + assert not case["report"].exists() + else: + assert cli.main([*case["validation_args"], "--execute"]) == 1 + if case["report"].exists(): + report = json.loads(case["report"].read_text()) + assert report["status"] != "passed" + if mutation == "running": + assert report["status"] == report["gates"]["execution"]["status"] == "incomplete" + + +def test_legacy_preparation_can_resume_first_repeats_and_reuse_them(assessed_collection, tmp_path, monkeypatch): # noqa: F811 + case = assessed_collection + _legacy_execution(case) + current = workflow._repeat_plan + + def historical(*args, **kwargs): + return current(*args, **{**kwargs, "observation_evidence_version": 1}) + + monkeypatch.setattr(workflow, "_repeat_plan", historical) + assert cli.main(case["validation_args"]) == 0 + selection_path = case["report"].parent / "repeatability-selection.json" + selection = selection_path.read_bytes() + assert "observation_evidence_version" not in json.loads(selection) + monkeypatch.setattr(workflow, "_repeat_plan", current) + assert cli.main([*case["validation_args"], "--resume", "--execute"]) == 0 + assert json.loads(case["report"].read_text())["status"] == "passed" + assert selection_path.read_bytes() == selection + repeats = case["report"].parent / "repeatability" + assert "observation_evidence_version" not in json.loads((repeats / repeatability.PLAN_FILENAME).read_text()) + workflow.check_collection_report(case["report"]) + calls = len(case["calls"]) + new_output = tmp_path / "reused" + args = list(case["validation_args"]) + args[args.index("--validation-output-dir") + 1] = str(new_output) + assert cli.main([*args, "--repeatability-dir", str(repeats)]) == 0 + assert len(case["calls"]) == calls + assert json.loads((new_output / "collection-validation.json").read_text())["status"] == "passed" + workflow.check_collection_report(new_output / "collection-validation.json") diff --git a/python/aisimulate/tests/unit/test_support_interpolation_validation.py b/python/aisimulate/tests/unit/test_support_interpolation_validation.py new file mode 100644 index 000000000..a70b1dc7e --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_interpolation_validation.py @@ -0,0 +1,571 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Withheld timings use the native interpolator without an analytical graph.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from aisimulate.support.interpolation_validation import evaluate_interpolation_holdout +from aisimulate.support.schema import SupportRequest +from aisimulate_core.sdk import engine +from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY + +pytestmark = pytest.mark.unit + + +def _write_pair(path, rows, *, schema_version=7): + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": schema_version, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "collector_attempt_ids": ["synthetic-validation-fixture"], + } + ) + ) + + +def _execution_evidence(rows, directory, *, max_capture=512): + """Synthetic source-bound metadata; the CLI tests exercise its native reader.""" + directory.mkdir() + evidence = [] + for cell_id in sorted({row["cell_id"] for row in rows}): + cell_rows = [row for row in rows if row["cell_id"] == cell_id] + identity = { + "source_plan_sha256": "a" * 64, + "collector_attempt_id": "synthetic-source", + "runtime_run_id": f"run-{cell_id}", + "runtime_grid_digest": f"grid-{cell_id}", + } + points = [] + for row in cell_rows: + row.update(identity) + axis = row["total_prefill_tokens"] or row["batch_size"] + captured = axis <= max_capture + points.append( + { + **{ + key: row[key] + for key in ("workload_kind", "batch_size", "total_prefill_tokens", "total_kv_read_tokens") + }, + "expected_cudagraph_mode": ("PIECEWISE" if row["workload_kind"] == "prefill" else "FULL") + if captured + else "NONE", + "expected_capture_size": axis if captured else None, + } + ) + config = { + "mode": "FULL_AND_PIECEWISE", + "capture_sizes": sorted({p["expected_capture_size"] for p in points if p["expected_capture_size"]}), + } + graphs = [] + for rank in range(cell_rows[0]["dp"]): + path = directory / f"{cell_id}-{rank}.json" + path.write_text(json.dumps({"cudagraph": config, "points": points})) + graphs.append( + { + "dp_rank": rank, + "config": copy.deepcopy(config), + "source": { + "path": str(path), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "size_bytes": path.stat().st_size, + }, + } + ) + evidence.append({"cell_id": cell_id, "identity": identity, "points": points, "native_graph_config": graphs}) + return evidence + + +@pytest.fixture +def holdout_case(tmp_path, monkeypatch, request): + def reject_graph(*_args, **_kwargs): + raise AssertionError("holdout must not construct an analytical model or substitute op-level timings") + + monkeypatch.setattr(engine, "get_model", reject_graph) + monkeypatch.setattr(engine, "build_model_config", reject_graph) + topology = getattr(request, "param", "tp") + is_moe = topology != "tp" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["UnregisteredDecoderForCausalLM"]})) + profile = { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "synthetic-immutable-v1", + "architecture": "UnregisteredDecoderForCausalLM", + "context_length": 8192, + "num_experts": 8 if is_moe else 0, + "provenance": "Synthetic metadata for native path tests; no silicon qualification.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 1 if topology == "dep" else 2, + "dp": 2 if topology == "dep" else 1, + "moe_tp": 1, + "moe_ep": 2 if is_moe else 1, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 256, + "max_batch_size": 4, + "provenance": "Pending runtime memory; sufficient for standalone forward timing.", + }, + } + ], + } + deployment = profile["deployments"][0] + onboarding = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe" if is_moe else "dense", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": { + "tensor_parallel": deployment["tp"], + "attention_data_parallel": deployment["dp"], + "moe_tensor_parallel": deployment["moe_tp"], + "moe_expert_parallel": deployment["moe_ep"], + "context_length": 8192, + }, + "fpm_profile": profile, + } + ) + identity = onboarding.profile_deployment().model_dump(mode="json", exclude={"resources"}) + rows = [] + for phase in ("prefill", "decode"): + for batch in (1, 2, 4): + for prefill in (16, 32, 64, 128, 256) if phase == "prefill" else (0,): + for kv in (0, 64, 256, 1024, 2048) if phase == "prefill" else (4, 16, 64, 256, 2048): + rows.append( + { + **identity, + **dict(zip(EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, strict=True)), + "model_path": profile["model"], + "cell_id": f"synthetic-{phase}", + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": batch, + "total_prefill_tokens": prefill, + "total_kv_read_tokens": kv, + "latency_ms": 10 + batch / 4 + prefill / 100 + kv / 1000, + "kv_seed_regime": "real_kv", + } + ) + source = tmp_path / "source" + source.mkdir() + packaged = Path(engine.__file__).parents[1] / "systems/h200_sxm.yaml" + (source / "h200_sxm.yaml").write_bytes(packaged.read_bytes()) + parquet = source / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + parquet.parent.mkdir(parents=True) + _write_pair(parquet, rows) + return onboarding, source, parquet, rows + + +@pytest.mark.parametrize("holdout_case", ["tp", "dep", "tep"], indirect=True) +def test_native_holdout_is_graph_independent_and_removes_all_coordinates(holdout_case, tmp_path): + request, source, parquet, original_rows = holdout_case + before = {path: path.read_bytes() for path in source.rglob("*") if path.is_file()} + output = tmp_path / "assessment" + report = evaluate_interpolation_holdout(request, systems_root=source, output_dir=output) + assert not request.profile_deployment().resources.memory_ready + assert report["status"] == "passed" + assert report["serving_accuracy"] == "not_assessed" + assert report["source_artifacts_unchanged"] + assert before == {path: path.read_bytes() for path in before} + assert report["native_query_coverage"]["queries"] == { + "measured": 0, + "interpolated": 18, + "unsupported": 0, + } + for phase, expected in (("prefill", 15), ("decode", 3)): + phase_report = report["phases"][phase] + assert phase_report["selected_count"] == expected + assert phase_report["absolute_relative_error"]["p95"] == pytest.approx(0, abs=1e-14) + assert phase_report["absolute_relative_error"]["bias"] == pytest.approx(0, abs=1e-14) + config = report["native_diagnostics"]["provenance"]["config"] + assert config["estimation_mode"] == "fpm_interpolation" + assert config["fallback_policy"] == "deny" + assert config["systems_paths"] == [str(output / "systems")] + assert config["estimator_config"]["fpm_interpolation"]["method"] == "direct" + remaining = pq.read_table(report["artifacts"]["training_parquet"]["path"]).to_pylist() + assert len(remaining) == len(original_rows) - 18 + for point in report["predictions"]: + assert not any( + row["workload_kind"] == point["phase"] + and all(row[name] == value for name, value in point["coordinates"].items()) + for row in remaining + ) + plan = json.loads((output / "holdout-plan.json").read_text()) + assert plan["retained_envelope_anchors"] + assert plan["capture_boundaries"]["status"] == "not_assessed" + assert plan["inputs"]["parquet"]["sha256"] == hashlib.sha256(parquet.read_bytes()).hexdigest() + assert json.loads((output / "interpolation-validation.json").read_text()) == report + + +@pytest.mark.parametrize("role", ["prefill", "decode"]) +def test_role_holdout_qualifies_only_its_phase_without_counterpart_data(holdout_case, tmp_path, role): + request, source, parquet, rows = holdout_case + payload = request.model_dump(mode="json") + payload["worker_type"] = role + payload["fpm_profile"]["deployments"][0]["worker_type"] = role + request = SupportRequest.model_validate(payload) + selected_rows = [row for row in rows if row["workload_kind"] == role] + _write_pair(parquet, selected_rows) + output = tmp_path / f"{role}-assessment" + report = evaluate_interpolation_holdout(request, systems_root=source, output_dir=output) + assert report["status"] == "passed" + assert report["worker_type"] == role + assert report["collection_phases"] == [role] + assert set(report["phases"]) == {role} + assert report["phases"][role]["selected_count"] > 0 + assert all(point["phase"] == role for point in report["predictions"]) + config = json.loads((output / "perfmodel-config.json").read_text()) + assert config["worker_type"] == role + assert report["native_query_coverage"]["queries"]["measured"] == 0 + assert report["native_query_coverage"]["queries"]["unsupported"] == 0 + assert pq.read_table(parquet).num_rows == len(selected_rows) + + +def test_selection_is_deterministic_and_input_order_independent(holdout_case, tmp_path): + request, source, parquet, rows = holdout_case + reports = [] + for index in range(3): + if index == 1: + _write_pair(parquet, list(reversed(rows))) + reports.append( + evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / f"assessment-{index}", seed=42 if index < 2 else 43 + ) + ) + assert reports[0]["predictions"] == reports[1]["predictions"] + assert reports[1]["predictions"] != reports[2]["predictions"] + + +def test_holdout_retains_both_sides_of_graph_eager_transition(holdout_case, tmp_path): + request, source, parquet, rows = holdout_case + template = next(row for row in rows if row["workload_kind"] == "prefill") + # Measured campaign shape: the 513-token eager point is much slower than + # the adjacent 512-token PIECEWISE point. Withholding it must not erase + # the runtime transition that the full table already measured. + timings = {256: 38.0, 512: 75.982, 513: 678.985, 1024: 718.120, 2048: 796.54} + prefill = [ + {**template, "batch_size": 1, "total_prefill_tokens": tokens, "total_kv_read_tokens": 0, "latency_ms": timing} + for tokens, timing in timings.items() + ] + rows = prefill + [row for row in rows if row["workload_kind"] == "decode"] + evidence = _execution_evidence(rows, tmp_path / "native") + _write_pair(parquet, rows) + old = evaluate_interpolation_holdout(request, systems_root=source, output_dir=tmp_path / "old", seed=0) + assert old["status"] == "failed" + missed = next(point for point in old["predictions"] if point["phase"] == "prefill") + assert missed["coordinates"]["total_prefill_tokens"] == 513 + assert missed["predicted_ms"] == pytest.approx(77.2361758) + previous = {path: path.read_bytes() for path in (tmp_path / "old").rglob("*") if path.is_file()} + before = { + path: path.read_bytes() for root in (source, tmp_path / "native") for path in root.rglob("*") if path.is_file() + } + report = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "assessment", seed=0, execution_evidence=evidence + ) + remaining = pq.read_table(report["artifacts"]["training_parquet"]["path"]).to_pylist() + retained_tokens = {row["total_prefill_tokens"] for row in remaining if row["workload_kind"] == "prefill"} + assert {512, 513} <= retained_tokens + assert report["native_query_coverage"]["queries"]["measured"] == 0 + assert report["capture_boundaries"]["status"] == "assessed" + assert report["capture_boundaries"]["per_point_observed_dispatch"] == "unreported" + assert {point["total_prefill_tokens"] for point in report["capture_boundaries"]["retained_anchors"]} == {512, 513} + assert report["phases"]["prefill"]["selected_count"] == 1 + assert report["phases"]["prefill"]["absolute_relative_error"]["p95"] < 0.001 + assert before == {path: path.read_bytes() for path in before} + assert previous == {path: path.read_bytes() for path in previous} + + +def test_boundary_selection_is_deterministic_with_ragged_curves_and_fake_rows(holdout_case, tmp_path): + request, source, parquet, rows = holdout_case + rows = [ + row + for row in rows + if not ( + row["workload_kind"] == "prefill" + and row["total_kv_read_tokens"] == 64 + and row["total_prefill_tokens"] == 64 + ) + ] + for row in rows: + if ( + row["workload_kind"] == "prefill" + and row["total_kv_read_tokens"] == 256 + and row["total_prefill_tokens"] == 64 + ): + row["kv_seed_regime"] = "fake_fallback" + evidence = _execution_evidence(rows, tmp_path / "native", max_capture=32) + _write_pair(parquet, rows) + first = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "first", execution_evidence=evidence + ) + _write_pair(parquet, list(reversed(rows))) + reordered = copy.deepcopy(list(reversed(evidence))) + for cell in reordered: + cell["points"].reverse() + cell["native_graph_config"].reverse() + second = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "second", execution_evidence=reordered + ) + assert first["predictions"] == second["predictions"] + assert first["capture_boundaries"] == second["capture_boundaries"] + anchors = first["capture_boundaries"]["retained_anchors"] + assert any(point["total_prefill_tokens"] == 128 and point["total_kv_read_tokens"] == 64 for point in anchors) + assert not any(point["total_prefill_tokens"] == 64 and point["total_kv_read_tokens"] == 256 for point in anchors) + assert first["native_query_coverage"]["queries"]["measured"] == 0 + + +@pytest.mark.parametrize("missing", ["point_mode", "unknown_mode", "capture_size", "graph_config"]) +def test_incomplete_boundary_evidence_does_not_claim_qualification(holdout_case, tmp_path, missing): + request, source, parquet, rows = holdout_case + evidence = _execution_evidence(rows, tmp_path / "native") + cell = next(cell for cell in evidence if cell["cell_id"] == "synthetic-prefill") + if missing == "graph_config": + cell["native_graph_config"][0]["config"] = None + elif missing == "unknown_mode": + cell["points"][0]["expected_cudagraph_mode"] = "FUTURE_MODE" + elif missing == "capture_size": + cell["points"][0].pop("expected_capture_size") + else: + cell["points"][0].pop("expected_cudagraph_mode") + _write_pair(parquet, rows) + report = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "assessment", execution_evidence=evidence + ) + assert report["status"] == report["capture_boundaries"]["status"] == "incomplete" + assert all(phase["status"] == "passed" for phase in report["phases"].values()) + assert report["capture_boundaries"]["issues"] + assert report["capture_boundaries"]["retained_anchors"] == [] + + +@pytest.mark.parametrize( + "corruption", ["identity", "coordinates", "duplicate", "rank", "graph_rank", "capture", "source"] +) +@pytest.mark.parametrize("holdout_case", ["dep"], indirect=True) +def test_boundary_evidence_rejects_contradictions(holdout_case, tmp_path, corruption): + request, source, parquet, rows = holdout_case + evidence = _execution_evidence(rows, tmp_path / "native") + _write_pair(parquet, rows) + cell = next(cell for cell in evidence if cell["cell_id"] == "synthetic-prefill") + if corruption == "identity": + cell["identity"]["collector_attempt_id"] = "another-attempt" + elif corruption == "coordinates": + cell["points"].pop() + elif corruption == "duplicate": + cell["points"].append(cell["points"][0]) + elif corruption == "rank": + cell["native_graph_config"][1]["dp_rank"] = 0 + elif corruption == "graph_rank": + cell["native_graph_config"][1]["config"]["mode"] = "NONE" + elif corruption == "capture": + cell["points"][0]["expected_capture_size"] = 9999 + else: + Path(cell["native_graph_config"][0]["source"]["path"]).write_text("changed") + with pytest.raises(ValueError, match="holdout"): + evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "assessment", execution_evidence=evidence + ) + assert not (tmp_path / "assessment/holdout-plan.json").exists() + + +@pytest.mark.parametrize("with_evidence", [False, True]) +def test_errors_are_against_withheld_measurements_with_editable_gates(holdout_case, tmp_path, with_evidence): + request, source, parquet, rows = holdout_case + evidence = _execution_evidence(rows, tmp_path / "native", max_capture=32) if with_evidence else None + _write_pair(parquet, rows) + baseline = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "baseline", execution_evidence=evidence + ) + for row in rows: + if any( + row["workload_kind"] == point["phase"] + and all(row[name] == value for name, value in point["coordinates"].items()) + for point in baseline["predictions"] + ): + row["latency_ms"] *= 2 + _write_pair(parquet, rows) + strict = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "strict", execution_evidence=evidence + ) + assert strict["status"] == "failed" + for phase in ("prefill", "decode"): + distribution = strict["phases"][phase]["absolute_relative_error"] + assert distribution == pytest.approx({"p50": 0.5, "p95": 0.5, "max": 0.5, "bias": -0.5}) + assert strict["phases"][phase]["unsupported_count"] == 0 + relaxed = evaluate_interpolation_holdout( + request, + systems_root=source, + output_dir=tmp_path / "relaxed", + max_p95_relative_error=0.6, + execution_evidence=evidence, + ) + assert relaxed["status"] == "passed" + assert relaxed["predictions"] == strict["predictions"] + assert relaxed["policy"]["max_p95_relative_error"] == 0.6 + + +@pytest.mark.parametrize("with_evidence", [False, True]) +def test_ragged_unsupported_holdout_is_not_an_error_sample_or_a_pass(holdout_case, tmp_path, with_evidence): + request, source, parquet, rows = holdout_case + template = next(row for row in rows if row["workload_kind"] == "prefill") + sparse_prefill = [ + {**template, "batch_size": 1, "total_prefill_tokens": tokens, "total_kv_read_tokens": kv} + for tokens, kv in ((16, 0), (32, 0), (64, 128), (128, 256), (256, 256)) + ] + rows = sparse_prefill + [row for row in rows if row["workload_kind"] == "decode"] + evidence = _execution_evidence(rows, tmp_path / "native") if with_evidence else None + _write_pair(parquet, rows) + report = evaluate_interpolation_holdout( + request, systems_root=source, output_dir=tmp_path / "assessment", execution_evidence=evidence + ) + assert report["status"] == "failed" + prefill = report["phases"]["prefill"] + assert prefill["selected_count"] == prefill["unsupported_count"] == 1 + assert prefill["unsupported_fraction"] == 1 + assert prefill["absolute_relative_error"] is None + assert report["native_query_coverage"]["queries"]["unsupported"] == 1 + assert report["predictions"][0]["status"] == "unsupported" + assert "predicted_ms" not in report["predictions"][0] + permissive = evaluate_interpolation_holdout( + request, + systems_root=source, + output_dir=tmp_path / "permissive", + max_unsupported_fraction=1.0, + execution_evidence=evidence, + ) + assert permissive["status"] == "failed" # An entirely unsupported phase has no error assessment. + + +def test_sparse_tables_are_incomplete_and_fake_rows_are_not_holdouts(holdout_case, tmp_path): + request, source, parquet, rows = holdout_case + retained = [] + for phase in ("prefill", "decode"): + phase_rows = [row for row in rows if row["workload_kind"] == phase and row["batch_size"] == 1] + retained.extend([phase_rows[0], phase_rows[-1]]) + retained.append({**phase_rows[-1], "total_kv_read_tokens": 9000, "kv_seed_regime": "fake_fallback"}) + _write_pair(parquet, retained) + report = evaluate_interpolation_holdout(request, systems_root=source, output_dir=tmp_path / "assessment") + assert report["status"] == "incomplete" + assert report["predictions"] == [] + assert all(phase["status"] == "not_assessed" for phase in report["phases"].values()) + plan = json.loads(Path(report["artifacts"]["holdout_plan"]["path"]).read_text()) + assert plan["excluded_fake_fallback_count"] == 2 + assert plan["removed_row_count"] == 0 + + +@pytest.mark.parametrize("corruption", ["digest", "row_count", "identity", "duplicate", "coordinate_collision"]) +def test_source_validation_remains_native_and_rejects_leaking_duplicates(holdout_case, tmp_path, corruption): + request, source, parquet, rows = holdout_case + if corruption in {"duplicate", "coordinate_collision"}: + duplicate = copy.deepcopy(rows[0]) + if corruption == "coordinate_collision": + duplicate["cell_id"] += "-other-DP-evidence" + duplicate["weight_quantization"] = "other-source" + rows.append(duplicate) + _write_pair(parquet, rows) + else: + metadata_path = parquet.with_suffix(".metadata.json") + metadata = json.loads(metadata_path.read_text()) + metadata.update( + {"digest": {"parquet_sha256": "0" * 64}, "row_count": {"row_count": 1}, "identity": {"system": "b200_sxm"}}[ + corruption + ] + ) + metadata_path.write_text(json.dumps(metadata)) + before = parquet.read_bytes() + with pytest.raises((ValueError, PerfDataNotAvailableError)): + evaluate_interpolation_holdout(request, systems_root=source, output_dir=tmp_path / "assessment") + assert parquet.read_bytes() == before + assert not (tmp_path / "assessment/holdout-plan.json").exists() + + +def test_other_identity_and_legacy_metadata_are_preserved_without_becoming_holdouts(holdout_case, tmp_path): + request, source, parquet, rows = holdout_case + for row in rows: + for name in EXECUTION_COLUMNS: + row.pop(name) + foreign = {**rows[0], "model_path": "test/another-model", "latency_ms": 9999} + _write_pair(parquet, [*rows, foreign], schema_version=6) + report = evaluate_interpolation_holdout(request, systems_root=source, output_dir=tmp_path / "assessment") + assert report["status"] == "passed" + remaining = pq.read_table(report["artifacts"]["training_parquet"]["path"]).to_pylist() + assert foreign in remaining + metadata = json.loads(Path(report["artifacts"]["training_metadata"]["path"]).read_text()) + assert metadata["schema_version"] == 6 + assert metadata["collector_attempt_ids"] == ["synthetic-validation-fixture"] + + +@pytest.mark.parametrize( + "options", + [ + {"max_points_per_phase": 0}, + {"max_points_per_phase": True}, + {"seed": -1}, + {"seed": 0.1}, + {"max_p95_relative_error": float("nan")}, + {"max_p95_relative_error": -0.1}, + {"max_unsupported_fraction": 1.1}, + {"max_unsupported_fraction": True}, + ], +) +def test_invalid_policy_fails_before_writing(holdout_case, tmp_path, options): + request, source, _parquet, _rows = holdout_case + with pytest.raises(ValueError): + evaluate_interpolation_holdout(request, systems_root=source, output_dir=tmp_path / "assessment", **options) + assert not (tmp_path / "assessment").exists() + + +@pytest.mark.parametrize("destination", ["source", "nested", "parent", "existing", "symlink"]) +def test_assessment_never_overwrites_source_or_prior_outputs(holdout_case, tmp_path, destination): + request, source, _parquet, _rows = holdout_case + output = { + "source": source, + "nested": source / "assessment", + "parent": tmp_path, + "existing": tmp_path / "previous", + "symlink": tmp_path / "link", + }[destination] + if destination == "existing": + output.mkdir() + (output / "evidence.json").write_text("preserve me") + elif destination == "symlink": + output.symlink_to(source, target_is_directory=True) + before = {path: path.read_bytes() for path in source.rglob("*") if path.is_file()} + with pytest.raises(ValueError): + evaluate_interpolation_holdout(request, systems_root=source, output_dir=output) + assert before == {path: path.read_bytes() for path in before} diff --git a/python/aisimulate/tests/unit/test_support_modelopt_sidecar.py b/python/aisimulate/tests/unit/test_support_modelopt_sidecar.py new file mode 100644 index 000000000..23d1f343b --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_modelopt_sidecar.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Local ModelOpt metadata survives guided initialization and collector planning.""" + +from __future__ import annotations + +import hashlib +import json +import shutil +from pathlib import Path + +import pytest +from collector.fpm_forward.cli import _parser +from collector.fpm_forward.config import FPMCollectionOptions +from collector.fpm_forward.planner import build_collection_plan + +import aisimulate.main as cli +from aisimulate.support.config_profile import derive_profile, load_model_config +from aisimulate.support.fpm import fpm_cli_args +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +def _snapshot(tmp_path): + source = Path(__file__).parents[2] / "src/aisimulate_core/model_configs" + for filename in ("config.json", "hf_quant_config.json"): + shutil.copyfile(source / f"nvidia--Llama-3.1-70B-Instruct-FP8_{filename}", tmp_path / filename) + return tmp_path / "config.json" + + +@pytest.mark.parametrize( + "algorithm,kv,gemm,moe", [("FP8", "FP8", "fp8_static", "fp8"), ("NVFP4", "none", "nvfp4", "nvfp4")] +) +@pytest.mark.parametrize( + "metadata_layout", + ["flat", "nested-modelopt", "null-fields", "dynamic-conflict", "inline-conflict", "nested-inline-conflict"], +) +@pytest.mark.parametrize("duplicate_sidecar", [False, True]) +def test_public_init_with_modelopt_sidecar_plans_same_checkpoint( + tmp_path, algorithm, kv, gemm, moe, metadata_layout, duplicate_sidecar +): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + if metadata_layout == "nested-modelopt": + raw["quantization_config"] = {"quant_method": "modelopt"} + elif metadata_layout == "null-fields": + raw.update(quant_algo=None, quant_dynamic=None, kv_cache_quant_algo=None) + elif metadata_layout == "dynamic-conflict": + raw.update( + quant_dynamic=False, quantization_config={"quant_method": "modelopt", "activation_scheme": "dynamic"} + ) + elif metadata_layout in {"inline-conflict", "nested-inline-conflict"}: + raw["quantization_config"] = { + "quant_method": "modelopt", + "quant_algo": "NVFP4" if algorithm == "FP8" else "FP8", + } + sidecar = tmp_path / "hf_quant_config.json" + metadata = json.loads(sidecar.read_text()) + metadata["quantization"].update(quant_algo=algorithm, kv_cache_quant_algo=kv) + if duplicate_sidecar: + raw["hf_quant_config"] = metadata + if metadata_layout.startswith("nested"): + raw = {"architectures": ["ExampleMultimodalForConditionalGeneration"], "text_config": raw} + path.write_text(json.dumps(raw)) + sidecar.write_text(json.dumps(metadata)) + kv_dtype = "fp8" if kv == "FP8" else "bfloat16" + overrides = tmp_path / "overrides.json" + overrides.write_text(json.dumps({"fmha_quant_mode": kv_dtype, "kv_cache_dtype": kv_dtype})) + output = tmp_path / "request.yaml" + argv = [ + "onboard", + "init", + "--model-config", + str(path), + "--model", + "nvidia/Llama-3.1-70B-Instruct-FP8", + "--model-revision", + "a" * 40, + "--framework-version", + "0.27.0", + "--gpu", + "h200_sxm", + "--interconnect", + "nvswitch", + "--resource-overrides", + str(overrides), + "--context-length", + "16384", + "--tensor-parallel", + "2", + "--output", + str(output), + ] + if metadata_layout.endswith("conflict"): + with pytest.raises(SystemExit) as error: + cli.main(argv) + assert error.value.code == 2 + assert not output.exists() + from collector.fpm_forward.model_capability import load_model_config as collector_model_config + + conflict = "quant_dynamic" if metadata_layout == "dynamic-conflict" else "quant_algo" + with pytest.raises(ValueError, match=f"conflicting {conflict}"): + collector_model_config(str(path)) + return + assert cli.main(argv) == 0 + request = SupportRequest.from_yaml(output) + deployment = request.profile_deployment() + assert (deployment.gemm_quant_mode, deployment.moe_quant_mode) == (gemm, moe) + assert deployment.fmha_quant_mode == deployment.kv_cache_dtype == kv_dtype + provenance = json.loads(request.fpm_profile.provenance) + assert provenance["config_sha256"] == hashlib.sha256(path.read_bytes()).hexdigest() + assert provenance["config_source_files"] == { + "hf_quant_config.json": hashlib.sha256(sidecar.read_bytes()).hexdigest() + } + assert str(sidecar) in provenance["config_notes"]["quantization_sidecar"] + assert "weights_bytes" not in provenance["planning_estimates"] + assert "user override" in provenance["fields"]["kv_cache_dtype"]["source"] + args = _parser().parse_args(fpm_cli_args(request, output_dir=tmp_path / "collection", plan_only=True)[3:]) + plan = build_collection_plan( + backend="vllm", + model_path=request.identity.model, + model_architecture=args.model_architecture, + model_config_path=str(path), + system="h200_sxm", + selected_ops={"attention_context", "attention_generation"}, + options=FPMCollectionOptions.from_args(args), + fpm_profile=request.fpm_profile, + ) + assert {cell.weight_quantization for cell in plan.cells} == {gemm} + + +def test_sidecar_weights_do_not_supply_runtime_attention_or_unspecified_kv(tmp_path): + path = _snapshot(tmp_path) + (tmp_path / "hf_quant_config.json").write_text(json.dumps({"quantization": {"quant_algo": "NVFP4"}})) + draft = derive_profile(load_model_config(path)) + assert draft.resolved["gemm_quant_mode"] == draft.resolved["moe_quant_mode"] == "nvfp4" + assert {"fmha_quant_mode", "kv_cache_dtype"} <= draft.missing.keys() + assert "weights_bytes" not in draft.resolved + + +@pytest.mark.parametrize( + "inline", + [ + {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + {"quant_method": "fp8", "kv_cache_scheme": {"num_bits": 16, "type": "float"}}, + {"quant_method": "fp8", "kv_cache_scheme": {"num_bits": 8, "type": "int"}}, + ], +) +@pytest.mark.parametrize("duplicate_sidecar", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("consumer", ["onboarding", "collection"]) +def test_conflicting_inline_and_sidecar_metadata_is_rejected(tmp_path, inline, duplicate_sidecar, nested, consumer): + from collector.fpm_forward.model_capability import load_model_config as collector_model_config + + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + raw["quantization_config"] = inline + if duplicate_sidecar: + raw["hf_quant_config"] = json.loads((tmp_path / "hf_quant_config.json").read_text()) + if nested: + raw = {"architectures": ["ExampleMultimodalForConditionalGeneration"], "text_config": raw} + path.write_text(json.dumps(raw)) + with pytest.raises(ValueError, match="conflicting.*hf_quant_config"): + if consumer == "onboarding": + load_model_config(path) + else: + collector_model_config(str(path)) + + +@pytest.mark.parametrize("top_level", [None, "fp8"]) +def test_top_level_algorithm_does_not_hide_contradictory_inline_metadata(tmp_path, top_level): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + raw.update(quant_algo=top_level, quantization_config={"quant_method": "modelopt", "quant_algo": "NVFP4"}) + path.write_text(json.dumps(raw)) + with pytest.raises(ValueError, match="conflicting.*hf_quant_config"): + load_model_config(path) + + +def test_explicit_precision_override_remains_visible_with_sidecar(tmp_path): + config = load_model_config(_snapshot(tmp_path)) + draft = derive_profile(config, overrides={"gemm_quant_mode": "bfloat16", "kv_cache_dtype": "bfloat16"}) + assert draft.resolved["gemm_quant_mode"] == draft.resolved["kv_cache_dtype"] == "bfloat16" + assert "user override" in draft.sources["gemm_quant_mode"] + assert "user override" in draft.sources["kv_cache_dtype"] + + +@pytest.mark.parametrize( + "inline,gemm,moe", + [ + ({"quant_method": "modelopt"}, "fp8_static", "fp8"), + ({"quant_method": "fp8", "activation_scheme": "dynamic"}, "fp8", "fp8"), + ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8_block", "fp8_block"), + ], +) +@pytest.mark.parametrize("duplicate_sidecar", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_matching_inline_quantization_retains_activation_and_block_metadata( + tmp_path, inline, gemm, moe, duplicate_sidecar, nested +): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + raw["quantization_config"] = inline + if duplicate_sidecar: + raw["hf_quant_config"] = json.loads((tmp_path / "hf_quant_config.json").read_text()) + if nested: + raw = {"architectures": ["ExampleMultimodalForConditionalGeneration"], "text_config": raw} + path.write_text(json.dumps(raw)) + draft = derive_profile(load_model_config(path)) + assert (draft.resolved["gemm_quant_mode"], draft.resolved["moe_quant_mode"]) == (gemm, moe) + assert draft.resolved["kv_cache_dtype"] == "fp8" + assert "fmha_quant_mode" in draft.missing + + +def test_sidecar_conflict_with_string_kv_scheme_is_rejected(tmp_path): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + raw["quantization_config"] = {"kv_cache_scheme": "FP8"} + path.write_text(json.dumps(raw)) + (tmp_path / "hf_quant_config.json").write_text( + json.dumps({"quantization": {"quant_algo": "NVFP4", "kv_cache_quant_algo": "none"}}) + ) + with pytest.raises(ValueError, match="conflicting.*hf_quant_config"): + load_model_config(path) + + +@pytest.mark.parametrize("scheme,kv", [("FP8", "fp8"), ({"num_bits": 8, "type": "int"}, "int8")]) +def test_weight_only_sidecar_preserves_explicit_inline_kv_scheme(tmp_path, scheme, kv): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + raw["quantization_config"] = {"kv_cache_scheme": scheme} + path.write_text(json.dumps(raw)) + (tmp_path / "hf_quant_config.json").write_text(json.dumps({"quantization": {"quant_algo": "FP8"}})) + draft = derive_profile(load_model_config(path)) + assert draft.resolved["kv_cache_dtype"] == kv + assert draft.resolved["gemm_quant_mode"] == "fp8_static" + + +@pytest.mark.parametrize("sidecar", [{}, {"format": "unknown"}, {"quantization": {"quant_algo": "custom"}}]) +def test_unknown_sidecar_does_not_reintroduce_unquantized_defaults(tmp_path, sidecar): + path = _snapshot(tmp_path) + (tmp_path / "hf_quant_config.json").write_text(json.dumps(sidecar)) + draft = derive_profile(load_model_config(path)) + assert {"gemm_quant_mode", "moe_quant_mode"} <= draft.missing.keys() + assert "weights_bytes" not in draft.resolved + + +def test_nested_decoder_and_hf_snapshot_symlinks_keep_adjacent_sidecar(tmp_path): + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + path.write_text(json.dumps({"architectures": ["ExampleMultimodalForConditionalGeneration"], "text_config": raw})) + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + for filename in ("config.json", "hf_quant_config.json"): + (snapshot / filename).symlink_to(tmp_path / filename) + config = load_model_config(snapshot / "config.json") + draft = derive_profile(config) + assert draft.resolved["gemm_quant_mode"] == "fp8_static" + assert draft.resolved["kv_cache_dtype"] == "fp8" + assert str(snapshot / "hf_quant_config.json") in config.notes["quantization_sidecar"] + assert config.raw["hidden_size"] == raw["hidden_size"] + assert "text_config" not in config.raw + + +@pytest.mark.parametrize("layout", ["flat", "nested", "nested-inline"]) +def test_already_supported_sidecar_keeps_frozen_collector_source_identity(tmp_path, layout): + from collector.fpm_forward.model_capability import load_model_config as collector_model_config + + from aisimulate_core.sdk.utils import _attach_hf_quant_config + + path = _snapshot(tmp_path) + raw = json.loads(path.read_text()) + if layout == "nested-inline": + raw["quantization_config"] = {"quant_method": "fp8", "activation_scheme": "dynamic", "kv_cache_scheme": "FP8"} + if layout != "flat": + raw = {"architectures": ["ExampleMultimodalForConditionalGeneration"], "text_config": raw} + path.write_text(json.dumps(raw)) + sidecar = json.loads((tmp_path / "hf_quant_config.json").read_text()) + # This is the unchanged general SDK attachment used before this FPM fix. + baseline = _attach_hf_quant_config(raw, sidecar) + resolved = collector_model_config(str(path)) + assert resolved.payload == baseline + assert ( + resolved.sha256 + == hashlib.sha256( + json.dumps(baseline, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode() + ).hexdigest() + ) + assert resolved.effective_payload["quant_algo"] == "fp8" + + +@pytest.mark.parametrize( + "payload", + [ + "[]", + '{"quantization": []}', + '{"quantization": {"quant_algo": "FP8", "quant_algo": "NVFP4"}}', + '{"quantization": {"quant_algo": "FP8", "quantization_algo": "NVFP4"}}', + ], +) +def test_malformed_sidecar_fails_before_profile_creation(tmp_path, payload): + path = _snapshot(tmp_path) + (tmp_path / "hf_quant_config.json").write_text(payload) + with pytest.raises(ValueError): + load_model_config(path) diff --git a/python/aisimulate/tests/unit/test_support_plan.py b/python/aisimulate/tests/unit/test_support_plan.py index 8203a2071..57b527c35 100644 --- a/python/aisimulate/tests/unit/test_support_plan.py +++ b/python/aisimulate/tests/unit/test_support_plan.py @@ -19,6 +19,7 @@ from importlib.machinery import PathFinder from importlib.resources import files from pathlib import Path +from types import SimpleNamespace from typing import ClassVar import pytest @@ -26,7 +27,7 @@ from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig from aisimulate.support.fpm import run_fpm -from aisimulate.support.plan import create_plan, plan_lock, request_id +from aisimulate.support.plan import check_plan, create_plan, plan_lock, request_id, validation_id from aisimulate.support.schema import SupportRequest pytestmark = pytest.mark.unit @@ -114,9 +115,6 @@ def _request(**updates) -> SupportRequest: "model_kind": "moe", "framework_version": "0.25.1", "gpu": "h200_sxm", - "gpu_count": 12, - "node_count": 2, - "gpus_per_node": 6, "interconnect": "NVLink", }, "search": {"tensor_parallel": 4}, @@ -133,23 +131,47 @@ def _request(**updates) -> SupportRequest: {"identity": {"model_revision": "main"}}, {"identity": {"framework_version": " "}}, {"identity": {"gpu": "../h200_sxm"}}, - {"identity": {"gpu_count": 13}}, - {"search": {"tensor_parallel": 8}}, + {"search": {"moe_tensor_parallel": 2}}, {"search": {"tensor_parallel": True}}, - {"search": {"max_candidates": 3}}, {"search": {"context_length": 1151}}, {"workload": {"input_tokens": 0}}, {"workload": {"request_count": 0}}, {"workload": {"concurrency": 5}}, {"workload": {"slo": {"ttft_ms": float("inf"), "tpot_ms": 1.0}}}, {"workload": {"slo": {"ttft_ms": 1.0, "tpot_ms": True}}}, + {"collection": {"gpu_memory_utilization": True}}, + {"collection": {"gpu_memory_utilization": "0.8"}}, + {"collection": {"gpu_memory_utilization": float("nan")}}, + {"collection": {"gpu_memory_utilization": 0.0}}, + {"collection": {"gpu_memory_utilization": 1.1}}, + {"collection": {"prefill_cudagraph_policy": "runtime", "max_prefill_cudagraph_size": 512}}, ], ) -def test_request_rejects_invalid_allocation_workload_and_identity(updates): +def test_request_rejects_invalid_topology_workload_and_identity(updates): with pytest.raises(ValidationError): _request(**updates) +@pytest.mark.parametrize( + "updates,field", + [ + ({"identity": {"gpu_count": 4}}, "identity.gpu_count"), + ({"identity": {"node_count": 1}}, "identity.node_count"), + ({"identity": {"gpus_per_node": 4}}, "identity.gpus_per_node"), + ({"search": {"max_candidates": 1}}, "search.max_candidates"), + ({"search": {"max_candidates": 2}}, "search.max_candidates"), + ], +) +def test_legacy_allocation_fields_require_explicit_request_migration(updates, field): + with pytest.raises(ValidationError) as error: + _request(**updates) + message = str(error.value) + assert field in message + assert "Copy the request, remove these fields" in message + assert "new output directory" in message + assert "ordinary predict/recommend configs" in message + + def test_plan_does_not_resolve_unknown_model_and_reloads_public_configs(tmp_path, monkeypatch): real_import = builtins.__import__ @@ -187,7 +209,7 @@ def no_model_or_collector_import(name, *args, **kwargs): assert recommendation.optimizer.max_trials == 1 assert recommendation.optimizer.parallelism == 1 assert plan["search"]["candidate_count"] == 1 - assert "single" in plan["search"]["detail"] + assert "single-worker validation" in plan["search"]["detail"] assert {item["status"] for item in plan["prerequisites"]} == {"not_checked"} assert plan["accuracy"]["status"] == "not_assessed" assert (root / "systems/h200_sxm.yaml").is_file() @@ -195,6 +217,28 @@ def no_model_or_collector_import(name, *args, **kwargs): assert not (root / "evidence.yaml").exists() +@pytest.mark.parametrize("limit", [None, 512]) +def test_legacy_collection_policy_round_trip_retains_identity_and_launch_arguments(tmp_path, limit): + request = _request(collection={} if limit is None else {"max_prefill_cudagraph_size": limit}) + serialized = request.model_dump(mode="json", exclude_none=True) + assert "prefill_cudagraph_policy" not in serialized["collection"] + assert "gpu_memory_utilization" not in serialized["collection"] + assert request.collection.prefill_cudagraph_policy == "explicit" + assert request.collection_settings()["max_prefill_cudagraph_size"] == (limit or 2048) + plan = create_plan(request, tmp_path / "plan") + saved = SupportRequest.from_yaml(tmp_path / "plan/request.yaml") + assert request_id(saved) == request_id(request) + assert saved.model_dump(mode="json", exclude_none=True) == serialized + check_plan(saved, tmp_path / "plan") + command = plan["fpm"]["plan_command"] + assert "--fpm-prefill-cudagraph-policy" not in command + assert "--fpm-gpu-memory-utilization" not in command + if limit is None: + assert "--fpm-max-prefill-cudagraph-size" not in command + else: + assert command[command.index("--fpm-max-prefill-cudagraph-size") + 1] == str(limit) + + def test_onboarding_plan_and_preview_do_not_import_estimator_runtime(tmp_path): request = tmp_path / "request.json" request.write_text(_request().model_dump_json()) @@ -242,16 +286,20 @@ def find_spec(self, fullname, path=None, target=None): @pytest.mark.parametrize("model_kind,preset,moe_tensor", [("dense", "tp", 1), ("moe", "pure_tp", 4)]) -def test_plan_collects_one_chosen_worker_and_bounds_recommendation(tmp_path, model_kind, preset, moe_tensor): +def test_plan_collects_one_chosen_worker_and_generates_validation_configs(tmp_path, model_kind, preset, moe_tensor): from aisimulate.recommend import recommendation_to_sweeper - request = _request(identity={"model_kind": model_kind}, search={"max_candidates": 2}) + request = _request(identity={"model_kind": model_kind}) plan = create_plan(request, tmp_path) configs = [CoreRecommendationConfig.from_yaml(path) for path in plan["outputs"]["recommendation_configs"]] candidates = [candidate for config in configs for candidate in config.engine.workers.aggregated.parallelism.preset] - assert [(c.replicas, c.tensor, c.moe_tensor) for c in candidates] == [(1, 4, moe_tensor), (2, 4, moe_tensor)] + assert [(c.replicas, c.tensor, c.moe_tensor) for c in candidates] == [(1, 4, moe_tensor)] assert all(config.optimizer.max_trials == 1 for config in configs) - assert plan["search"]["candidate_count"] == 2 + assert plan["search"]["candidate_count"] == 1 + assert plan["fpm"]["collection_gpus_required"] == 4 + assert all(config.optimization.constraints.max_candidate_gpus == 4 for config in configs) + assert "max_candidates" not in plan["search"] + assert not (tmp_path / "recommend/replicas-2.yaml").exists() for config in configs: search = recommendation_to_sweeper(config).search_space assert search.agg_max_num_batched_tokens == [8192] @@ -261,26 +309,49 @@ def test_plan_collects_one_chosen_worker_and_bounds_recommendation(tmp_path, mod assert command[command.index("--fpm-gpu-counts") + 1] == "4" assert command[command.index("--fpm-max-gpus") + 1] == "4" assert command[command.index("--fpm-parallel-presets") + 1] == preset - assert command[command.index("--fpm-max-prefill-isl") + 1] == "1024" - assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "1" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "256" assert "--plan-only" in command assert "--execute" in commands["fpm_run_local"] -def test_one_worker_allocation_keeps_one_pilot_even_with_two_candidate_limit(tmp_path): - request = _request(identity={"node_count": 1, "gpus_per_node": 4, "gpu_count": 4}, search={"max_candidates": 2}) +@pytest.mark.parametrize( + "topology,preset", + [ + ({"tensor_parallel": 64}, "pure_tp"), + ( + {"tensor_parallel": 1, "attention_data_parallel": 64, "moe_tensor_parallel": 1, "moe_expert_parallel": 64}, + "dep", + ), + ({"tensor_parallel": 64, "moe_tensor_parallel": 1, "moe_expert_parallel": 64}, "tep"), + ], +) +def test_wide_worker_requirement_is_not_limited_by_node_allocation_or_runtime_default(tmp_path, topology, preset): + from aisimulate.recommend import recommendation_to_sweeper + + request = _request(search=topology) plan = create_plan(request, tmp_path) commands = json.loads((tmp_path / "commands.json").read_text()) + recommendation = CoreRecommendationConfig.from_yaml(tmp_path / "recommend/pilot.yaml") assert plan["search"]["candidate_count"] == 1 + assert plan["fpm"]["collection_gpus_required"] == 64 + assert recommendation_to_sweeper(recommendation).search_space.gpu_budget == 64 + assert recommendation.engine.workers.aggregated.parallelism.preset[0].replicas == 1 + command = plan["fpm"]["plan_command"] + assert command[command.index("--fpm-gpu-counts") + 1] == "64" + assert command[command.index("--fpm-parallel-presets") + 1] == preset assert plan["outputs"]["recommendation_configs"] == [str(tmp_path / "recommend/pilot.yaml")] assert len(commands["recommend"]) == 1 - assert "single" in plan["search"]["detail"] + assert "available allocation" in plan["search"]["detail"] @pytest.mark.parametrize("seed", [0, 7, 42]) @pytest.mark.parametrize("model_kind,moe_tensor", [("dense", 1), ("moe", 4)]) -def test_emitted_recommendations_evaluate_both_replica_choices(tmp_path, monkeypatch, seed, model_kind, moe_tensor): +@pytest.mark.parametrize("runtime_replicas", [1, 3]) +def test_validation_recommendation_and_user_runtime_replica_budget( + tmp_path, monkeypatch, seed, model_kind, moe_tensor, runtime_replicas +): from aisimulate.main import build_parser from aisimulate.recommend import run_recommendation from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig @@ -293,10 +364,10 @@ def test_emitted_recommendations_evaluate_both_replica_choices(tmp_path, monkeyp monkeypatch.setattr(RustForwardPassPerfModel, "best_available", _OfflinePerfModel.best_available) legal = [ ReplicaParallelConfig(shape=ParallelShape(tp=4, pp=1, dp=1, moe_tp=moe_tensor, moe_ep=1), replicas=count) - for count in (1, 2) + for count in (1, 2, 3) ] monkeypatch.setattr("aisimulate.sweeper.search_space.parallel_configs_for", lambda *args, **kwargs: legal) - request = _request(identity={"model_kind": model_kind}, search={"max_candidates": 2, "seed": seed}) + request = _request(identity={"model_kind": model_kind}, search={"seed": seed}) plan = create_plan(request, tmp_path) commands = json.loads((tmp_path / "commands.json").read_text())["recommend"] evaluated = [] @@ -306,57 +377,47 @@ def test_emitted_recommendations_evaluate_both_replica_choices(tmp_path, monkeyp for command in commands: args = build_parser().parse_args(command[1:]) config = CoreRecommendationConfig.from_yaml(args.config) + assert config.engine.workers.aggregated.parallelism.preset[0].replicas == 1 + if runtime_replicas > 1: + # Deployment choices remain ordinary runtime inputs; onboarding is unchanged. + payload = config.model_dump(mode="json", exclude_none=True) + payload["engine"]["workers"]["aggregated"]["parallelism"]["preset"][0]["replicas"] = runtime_replicas + payload["optimization"]["constraints"]["max_candidate_gpus"] = runtime_replicas * request.worker_gpus + config = CoreRecommendationConfig.model_validate(payload) + prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml").model_dump(mode="json") + prediction["engine"]["workers"]["aggregated"]["parallelism"]["replicas"] = runtime_replicas + assert ( + CorePredictionConfig.model_validate(prediction).engine.workers.aggregated.parallelism.replicas + == runtime_replicas + ) result = run_recommendation(config, stack="engine", runner_factory=_OfflineRunnerFactory(), show_progress=False) evaluated.extend(candidate.config["replicas"] for candidate in result.candidates) results.append(result) configs.append(args.config) outputs.append(args.output_dir) - assert evaluated == [1, 2] + assert evaluated == [runtime_replicas] assert configs == plan["outputs"]["recommendation_configs"] assert outputs == plan["outputs"]["recommendation_results"] - assert len(set(outputs)) == 2 + assert len(set(outputs)) == 1 assert all(result.counts.evaluated == 1 and result.counts.cache_hits == 0 for result in results) - assert len(_OfflinePerfModel.requests) == 2 + assert len(_OfflinePerfModel.requests) == 1 + assert SupportRequest.from_yaml(tmp_path / "request.yaml") == request _assert_estimator_requests(_OfflinePerfModel.requests, tmp_path / "systems", moe_tensor) -@pytest.mark.parametrize( - "max_candidates,missing_pilot_samples,statuses", - [(1, False, [0]), (2, False, [0, 0]), (2, True, [1, 0])], -) -@pytest.mark.parametrize("tampered_commands", [False, True]) -def test_documented_recommendation_loop_attempts_every_candidate( - tmp_path, max_candidates, missing_pilot_samples, statuses, tampered_commands -): +@pytest.mark.parametrize("missing_samples", [False, True]) +def test_documented_recommendation_command_evaluates_one_worker(tmp_path, missing_samples): guide = Path(__file__).resolve().parents[4] / "docs/fpm-self-service.md" - snippet = guide.read_text().split("python3 - <<'PY'\n", 1)[1].split("\nPY\n", 1)[0] + snippet = "aisimulate recommend " + guide.read_text().split("\naisimulate recommend ", 1)[1].split("\n```", 1)[0] workdir = tmp_path / "work with 'quotes'; $(literal)" workdir.mkdir() - request = _request( - identity={"model_kind": "dense", "gpu_count": 8, "node_count": 1, "gpus_per_node": 8}, - search={"max_candidates": max_candidates, "objective": "ttft"}, - ) + request = _request(identity={"model_kind": "dense"}, search={"objective": "ttft"}) plan = create_plan(request, workdir / "aisimulate-support") - commands = json.loads(Path(plan["outputs"]["commands"]).read_text())["recommend"] - injected_marker = tmp_path / "injected-command-ran" - if tampered_commands: - command_file = Path(plan["outputs"]["commands"]) - payload = json.loads(command_file.read_text()) - payload["recommend"].insert( - 0, - [ - sys.executable, - "-c", - "import sys; from pathlib import Path; Path(sys.argv[1]).write_text('unexpected execution')", - str(injected_marker), - ], - ) - command_file.write_text(json.dumps(payload)) evaluated = tmp_path / "evaluated.txt" estimator_requests = tmp_path / "estimator-requests.jsonl" wrapper = tmp_path / "bin/aisimulate" wrapper.parent.mkdir() - # Install the same Core double in each supervised child. The snippet executes + # Install the same Core double in each supervised child. The command executes # real subprocesses, CLI dispatch, estimator resolution, sampling and export. wrapper.write_text( f"#!{sys.executable}\n" @@ -394,7 +455,7 @@ def run(self, spec): replicas = spec.backend_deployment.num_workers with Path(os.environ["SUPPORT_TEST_EVALUATED"]).open("a") as stream: stream.write(str(replicas) + "\\n") - samples = 0.0 if {missing_pilot_samples!r} and replicas == 1 else 4.0 + samples = 0.0 if {missing_samples!r} else 4.0 return ReplayReport(metrics={{"mean_ttft_ms": 10.0, "num_ttft_samples": samples}}) def close(self): @@ -414,7 +475,7 @@ def create(self, worker_id): ) wrapper.chmod(0o755) result = subprocess.run( - [sys.executable, "-c", snippet], + shlex.split(snippet.replace("\\\n", "")), cwd=workdir, env={ **os.environ, @@ -428,53 +489,52 @@ def create(self, worker_id): check=False, ) - assert not injected_marker.exists() - assert result.returncode == (1 if any(statuses) else 0), result.stderr - assert evaluated.read_text().splitlines() == [str(index) for index in range(1, max_candidates + 1)] + assert result.returncode == (1 if missing_samples else 0), result.stderr + assert evaluated.read_text().splitlines() == ["1"] resolved_requests = [json.loads(line) for line in estimator_requests.read_text().splitlines()] - assert len(resolved_requests) == max_candidates + assert len(resolved_requests) == 1 _assert_estimator_requests(resolved_requests, workdir / "aisimulate-support/systems", 1) - for index, (output, status) in enumerate(zip(plan["outputs"]["recommendation_results"], statuses, strict=True), 1): + for output in plan["outputs"]["recommendation_results"]: root = Path(output) payload = json.loads((root / "recommendation.json").read_text()) assert payload["counts"]["evaluated"] == 1 - assert payload["counts"]["feasible"] == (0 if status else 1) - assert payload["counts"]["infeasible"] == (1 if status else 0) - assert payload["candidates"][0]["config"]["replicas"] == index + assert payload["counts"]["feasible"] == (0 if missing_samples else 1) + assert payload["counts"]["infeasible"] == (1 if missing_samples else 0) + assert payload["candidates"][0]["config"]["replicas"] == 1 exported = root / "recommendations/0001.yaml" - if status: + if missing_samples: assert "no qualifying samples" in payload["candidates"][0]["reason"] assert not exported.exists() else: prediction = CorePredictionConfig.from_yaml(exported) - assert prediction.engine.workers.aggregated.parallelism.replicas == index + assert prediction.engine.workers.aggregated.parallelism.replicas == 1 assert prediction.engine.systems_paths == [str(workdir / "aisimulate-support/systems")] timing = prediction.engine.workers.aggregated.timing assert timing.estimation_mode == "fpm_interpolation" assert timing.fallback_policy == "deny" assert timing.systems_paths == prediction.engine.systems_paths - assert [line for line in result.stdout.splitlines() if line.startswith("exit ")] == [ - f"exit {status}: {shlex.join(command)}" for command, status in zip(commands, statuses, strict=True) - ] -def test_prefill_bounds_follow_concurrent_workload_and_collector_minimum(tmp_path): +def test_collection_bounds_are_independent_of_synthetic_validation(tmp_path): request = _request(workload={"input_tokens": 10, "concurrency": 3, "request_count": 4}) command = create_plan(request, tmp_path)["fpm"]["plan_command"] - assert command[command.index("--fpm-max-prefill-isl") + 1] == "30" - assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "3" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "256" + assert command[command.index("--fpm-max-model-len") + 1] == "256000" small = _request(workload={"input_tokens": 1}) command = create_plan(small, tmp_path / "small")["fpm"]["plan_command"] - assert command[command.index("--fpm-max-prefill-isl") + 1] == "2" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" @pytest.mark.parametrize( "updates", [ {"search": {"tensor_parallel": 2}}, - {"search": {"max_candidates": 2}}, - {"search": {"seed": 7}}, - {"workload": {"input_tokens": 64}}, + {"search": {"context_length": 65536}}, + {"collection": {"max_num_tokens": 16384}}, + {"collection": {"max_prefill_cudagraph_size": 512}}, + {"collection": {"prefill_cudagraph_policy": "runtime"}}, + {"collection": {"gpu_memory_utilization": 0.7}}, {"identity": {"model_revision": "other-checkpoint"}}, {"identity": {"framework_version": "0.26.0"}}, {"identity": {"tokenizer_revision": "tokenizer-v2"}}, @@ -512,6 +572,81 @@ def test_same_request_overwrite_preserves_timings_and_rejects_modified_inputs(tm assert data.read_bytes() == b"timings" +@pytest.mark.parametrize( + "updates", + [ + {"workload": {"input_tokens": 4096, "output_tokens": 256, "concurrency": 8, "request_count": 16}}, + {"workload": {"slo": {"ttft_ms": 2500.0, "tpot_ms": 75.0}}}, + {"search": {"seed": 7, "objective": "goodput"}}, + ], +) +def test_validation_refresh_keeps_collection_artifacts_and_updates_examples(tmp_path, updates): + request = _request() + original = create_plan(request, tmp_path) + data = tmp_path / "systems/data/collected.parquet" + data.write_bytes(b"collected timings") + checkpoint = tmp_path / "fpm-checkpoint" + checkpoint.mkdir() + (checkpoint / "campaign.json").write_bytes(b"immutable checkpoint") + changed = _request(**updates) + assert request_id(changed) == request_id(request) + assert validation_id(changed) != validation_id(request) + + updated = create_plan(changed, tmp_path, overwrite=True) + + assert updated["request_id"] == original["request_id"] + assert updated["fpm"] == original["fpm"] + assert updated["validation_id"] != original["validation_id"] + assert data.read_bytes() == b"collected timings" + assert (checkpoint / "campaign.json").read_bytes() == b"immutable checkpoint" + assert SupportRequest.from_yaml(tmp_path / "request.yaml") == changed + prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(tmp_path / "recommend/pilot.yaml") + assert prediction.traffic.source.input_tokens == changed.workload.input_tokens + assert prediction.traffic.load.concurrency == changed.workload.concurrency + assert recommendation.optimizer.seed == changed.search.seed + check_plan(changed, tmp_path) + + +@pytest.mark.parametrize("relative", ["request.yaml", "predict/pilot.yaml", "commands.json", "support-plan.json"]) +def test_validation_refresh_does_not_bless_tampered_generated_files(tmp_path, relative): + create_plan(_request(), tmp_path) + target = tmp_path / relative + target.write_bytes(target.read_bytes() + b"\n") + before = _file_contents(tmp_path) + with pytest.raises(ValueError, match="modified|identity"): + create_plan(_request(workload={"input_tokens": 4096}), tmp_path, overwrite=True) + assert _file_contents(tmp_path) == before + + +@pytest.mark.parametrize( + "search,expected_context", + [(None, 16384), ({}, 16384), ({"tensor_parallel": 4, "context_length": 8192}, 8192)], +) +def test_v1_request_migration_preserves_declared_context_but_requires_a_new_plan(tmp_path, search, expected_context): + import yaml + + legacy = _request().model_dump(mode="json", exclude_none=True) + legacy["schema_version"] = "aisimulate-support-request/v1" + legacy.pop("collection") + if search is None: + legacy.pop("search") + else: + legacy["search"] = search + source = tmp_path / "old-request.yaml" + source.write_text(yaml.safe_dump(legacy)) + migrated = SupportRequest.from_yaml(source) + assert migrated.schema_version == "aisimulate-support-request/v2" + assert migrated.search.context_length == expected_context + assert migrated.scheduler_limits() == {"max_batched_tokens": 8192, "max_sequences": 256} + root = tmp_path / "old-plan" + root.mkdir() + (root / "support-plan.json").write_text(json.dumps({"schema_version": "aisimulate-support-plan/v1"})) + with pytest.raises(ValueError, match="new output directory"): + create_plan(migrated, root, overwrite=True) + assert create_plan(migrated, tmp_path / "new-plan")["fpm"]["runtime_limits"]["context_length"] == expected_context + + @pytest.mark.parametrize("modification", ["retired_command", "missing_execute", "formatting"]) @pytest.mark.parametrize("complete", [False, True]) def test_overwrite_rejects_modified_commands_before_repairing_missing_files(tmp_path, modification, complete): @@ -609,17 +744,29 @@ def no_collector_import(name, *args, **kwargs): assert not root.exists() -def test_execute_delegates_only_for_matching_plan(tmp_path, monkeypatch): - from collector.fpm_forward import cli +def _mock_collector_execution(monkeypatch, calls): + from collector.fpm_forward import cli, entry + + from aisimulate.support import collection_readiness, fpm + + def resolve(command): + calls.append(command[3:]) + return cli._parser().parse_args(command[3:]), (SimpleNamespace(cells=(), sha256="a" * 64), {}) + + monkeypatch.setattr(fpm, "_resolve_execution", resolve) + monkeypatch.setattr(entry, "run_resolved", lambda *_: [{"error_message": "synthetic collection failure"}]) + monkeypatch.setattr(collection_readiness, "assess_readiness", lambda *a, **k: {"ready_for_full_collection": True}) + +def test_execute_delegates_only_for_matching_plan(tmp_path, monkeypatch): calls = [] - monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + _mock_collector_execution(monkeypatch, calls) request = _request() with pytest.raises(ValueError, match="plan"): run_fpm(request, execute=True, output_dir=tmp_path) assert not calls create_plan(request, tmp_path) - assert run_fpm(request, execute=True, output_dir=tmp_path, smoke=True, limit=1) == 7 + assert run_fpm(request, execute=True, output_dir=tmp_path, smoke=True, limit=1) == 1 assert calls[0][calls[0].index("--limit") + 1] == "1" assert "--smoke" in calls[0] assert "--plan-only" not in calls[0] @@ -744,40 +891,36 @@ def test_resume_occupied_campaign_requires_usable_selected_checkpoint(tmp_path, @pytest.mark.parametrize("smoke", [False, True]) def test_matching_explicit_resume_delegates_and_preserves_campaign(tmp_path, monkeypatch, smoke): - from collector.fpm_forward import cli - request = _request() create_plan(request, tmp_path) checkpoint, _ = _seed_campaign(tmp_path, smoke=smoke, checkpoint_dir=tmp_path / "fpm-checkpoint/custom") before = _file_contents(tmp_path) calls = [] - monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + _mock_collector_execution(monkeypatch, calls) assert ( run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke, resume=True, checkpoint_dir=checkpoint.parent) - == 7 + == 1 ) assert "--resume" in calls[0] assert calls[0][calls[0].index("--checkpoint-dir") + 1] == str(checkpoint.parent) - assert _file_contents(tmp_path) == before + assert {key: value for key, value in _file_contents(tmp_path).items() if key != "fpm-readiness.json"} == before @pytest.mark.parametrize("smoke", [False, True]) def test_fresh_smoke_and_formal_campaigns_are_independent(tmp_path, monkeypatch, smoke): - from collector.fpm_forward import cli - request = _request() create_plan(request, tmp_path) _seed_campaign(tmp_path, smoke=not smoke) before = _file_contents(tmp_path) calls = [] - monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + _mock_collector_execution(monkeypatch, calls) - assert run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke) == 7 + assert run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke) == 1 assert len(calls) == 1 - assert _file_contents(tmp_path) == before + assert {key: value for key, value in _file_contents(tmp_path).items() if key != "fpm-readiness.json"} == before def test_execute_checks_for_campaign_outputs_while_holding_lock(tmp_path, monkeypatch): @@ -1037,7 +1180,7 @@ def test_overwrite_rejects_existing_invalid_metadata_before_repair(tmp_path, met (tmp_path / "predict/pilot.yaml").unlink() before_repair = _file_contents(tmp_path) - with pytest.raises(ValueError, match="readable onboarding plan|request identity"): + with pytest.raises(ValueError, match="readable onboarding plan|request identity|legacy or invalid onboarding plan"): create_plan(request, tmp_path, overwrite=True) assert _file_contents(tmp_path) == before_repair diff --git a/python/aisimulate/tests/unit/test_support_serving_validation.py b/python/aisimulate/tests/unit/test_support_serving_validation.py new file mode 100644 index 000000000..406b2aa71 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_serving_validation.py @@ -0,0 +1,611 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serving checks retain native benchmark records and fail closed on mismatched evidence.""" + +from __future__ import annotations + +import copy +import io +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from aisimulate import main as cli +from aisimulate.support import serving_validation as serving + +from .test_support_validation import validation_case as _replay_validation_case # noqa: F401 + +pytestmark = pytest.mark.unit + + +def _save(path, value): + path.write_text(json.dumps(value)) + + +@pytest.fixture +def materialized_workload(monkeypatch): + """Author only the isolated tokenizer boundary; native derived replay still runs.""" + + def materialize(recipe, output, _python): + play = json.loads(Path(recipe["benchmark_trace"]["path"]).read_text()) + inputs, payloads = [], [] + for index, source in enumerate(play["requests"]): + tokens = [block for block in source["hash_ids"] for _ in range(play["block_size"])][: source["in"]] + inputs.append({"turn_index": index, "input_token_ids": tokens, "input_length": len(tokens)}) + payloads.append( + { + "payload": { + "model": recipe["scope"]["model_projection"], + "messages": [{"role": "user", "content": "synthetic"}], + "add_generation_prompt": True, + "continue_final_message": False, + } + } + ) + _save(output / "tokenization.json", {"chat_template": "synthetic", "requests": inputs}) + (output / "requests.jsonl").write_text("".join(json.dumps(row) + "\n" for row in payloads)) + _save(output / "helper.json", {"fixture": "synthetic tokenizer boundary"}) + _save(output / "input.json", {"trace": recipe["trace"]}) + return { + "helper": serving._identity(output / "helper.json"), + "input": serving._identity(output / "input.json"), + "installation": {"source": {"vcs_info": {"commit_id": serving.AIPERF_REVISION}}}, + } + + monkeypatch.setattr(serving, "_materialize_payloads", materialize) + + +def write_server_tokenization(recipe_path, recipe): + frozen = json.loads(Path(recipe["matched_workload"]["tokenization"]["path"]).read_text()) + payloads = [ + json.loads(line) for line in Path(recipe["matched_workload"]["payloads"]["path"]).read_text().splitlines() + ] + rows = [] + for tokens, payload in zip(frozen["requests"], payloads, strict=True): + request = { + key: payload["payload"][key] + for key in ("model", "messages", "add_generation_prompt", "continue_final_message") + } + request["add_special_tokens"] = False + rows.append( + { + "turn_index": tokens["turn_index"], + "request": request, + "response": {"tokens": tokens["input_token_ids"], "count": tokens["input_length"]}, + } + ) + path = recipe_path.parent / "server-tokenization.json" + _save(path, {"recipe": serving._identity(recipe_path), "status": "matched", "requests": rows}) + return path + + +@pytest.fixture +def serving_case(request, tmp_path, materialized_workload): + args, request, _plan, trace, prediction_output = request.getfixturevalue("_replay_validation_case") + # This invokes the actual ordinary prediction/compiler/native FPM path; + # imported AIPerf-shaped observations below are synthetic contract fixtures. + assert cli.main(args) == 0 + config = yaml.safe_load((prediction_output / "predict.yaml").read_text()) + worker = config["engine"]["workers"]["aggregated"] + tokenizer = tmp_path / "tokenizer" + tokenizer.mkdir() + (tokenizer / "tokenizer_config.json").write_text('{"tokenizer_class":"synthetic"}') + expected = { + "model": request.identity.model, + "model_revision": request.identity.model_revision, + "backend": "vllm", + "backend_version": request.identity.framework_version, + "image_digest": "sha256:" + "a" * 64, + "hardware": request.identity.gpu, + "parallelism": request.parallelism(), + "precision": {"weights": "fp8", "fmha": "bfloat16", "kv": "fp8"}, + "attention_groups": [{"backend_name": "FLASH_ATTN", "layer_names": ["layer.0"]}], + "graph_config": {"cudagraph_mode": "FULL_AND_PIECEWISE", "cudagraph_capture_sizes": [1, 2]}, + "scheduler": worker["scheduler"], + "context_length": config["engine"]["context_length"], + "gpu_memory_utilization": 0.73, + "runtime_config": { + "model_config": { + "dtype": "torch.bfloat16", + "quantization": "fp8", + "revision": request.identity.model_revision, + }, + "cache_config": {"cache_dtype": "fp8"}, + "scheduler_config": {"max_num_batched_tokens": 64, "max_num_seqs": 4}, + }, + } + recipe_path = serving.prepare_serving_validation( + trace=trace, + prediction_config=prediction_output / "predict.yaml", + output=tmp_path / "serving", + endpoint="http://127.0.0.1:8000", + tokenizer=tokenizer, + expected_execution=expected, + aiperf_python=Path("/benchmark/bin/python"), + ) + recipe = json.loads(recipe_path.read_text()) + prediction_report = prediction_output / "prediction/prediction.json" + prediction = json.loads(Path(recipe["matched_workload"]["prediction_report"]["path"]).read_text()) + artifacts = Path(recipe["artifacts_dir"]) + artifacts.mkdir() + rows = [] + for index, row in enumerate(prediction["per_request"]): + rows.append( + { + "metadata": { + "conversation_id": recipe["play_id"], + "turn_index": index, + "x_request_id": f"http-{index}", + "was_cancelled": False, + "benchmark_phase": "profiling", + "request_start_ns": 1_000_000_000 + round(row["arrival_time_ms"] * 1_000_000), + "request_end_ns": 1_000_000_000 + round(row["last_token_ms"] * 1_000_000), + }, + "metrics": { + "input_sequence_length": {"value": row["input_length"], "unit": "tokens"}, + "output_sequence_length": {"value": row["output_length"], "unit": "tokens"}, + "time_to_first_token": {"value": row["ttft_ms"], "unit": "ms"}, + "inter_token_latency": {"value": row["itl_ms"], "unit": "ms"}, + }, + "error": None, + } + ) + records = artifacts / "profile_export.jsonl" + records.write_text("\n".join(json.dumps(row) for row in rows) + "\n") + summary = artifacts / "profile_export_aiperf.json" + _save(summary, {"was_cancelled": False, "is_complete": True, "request_count": {"avg": len(rows)}}) + runtime_dir = Path(recipe["serving_observer"]["output_dir"]) + runtime_dir.mkdir() + provenance_path = Path(recipe["serving_observer"]["provenance_path"]) + provenance = json.loads(provenance_path.read_text()) + runtime_sources = [] + native_config = copy.deepcopy(expected["runtime_config"]) + native_config["cache_config"]["enable_prefix_caching"] = True + for tp in range(expected["parallelism"]["tensor"]): + native_path = runtime_dir / f"fpm-execution-worker-dp0-tp{tp}-pp0.json" + _save( + native_path, + { + "schema_name": "aisimulate_fpm_runtime_execution", + "schema_version": 1, + "status": "observed", + "execution_provenance": provenance, + "provenance_source": serving._identity(provenance_path), + "dp_rank": 0, + "tp_rank": tp, + "pp_rank": 0, + **{key: expected[key] for key in ("attention_groups", "graph_config", "backend_version")}, + "resolved_config": native_config, + }, + ) + runtime_sources.append(serving._identity(native_path)) + receipt = { + "schema_version": serving.SCHEMA, + "recipe": serving._identity(recipe_path), + "status": "completed", + "exit_code": 0, + "command": ["python", *recipe["command_arguments"]], + "installation": {"source": {"vcs_info": {"commit_id": serving.AIPERF_REVISION}}}, + "artifacts": [serving._identity(records), serving._identity(summary)], + "runtime_observations": runtime_sources, + "serving_provenance": serving._identity(provenance_path), + "observer_modules": recipe["serving_observer"]["module_files"], + "server_tokenization": serving._identity(write_server_tokenization(recipe_path, recipe)), + } + _save(recipe_path.parent / "benchmark-execution.json", receipt) + observed = { + "schema_version": serving.SCHEMA, + "recipe": serving._identity(recipe_path), + "observed_execution": expected, + "scope": recipe["scope"], + "sources": [serving._identity(tokenizer / "tokenizer_config.json")], + } + observed_path = recipe_path.parent / "observed-execution.json" + _save(observed_path, observed) + return SimpleNamespace( + request=request, + recipe_path=recipe_path, + recipe=recipe, + prediction_report=prediction_report, + observed_path=observed_path, + observed=observed, + rows=rows, + records=records, + trace=trace, + prediction_config=prediction_output / "predict.yaml", + tokenizer=tokenizer, + ) + + +def _assess(case, **kwargs): + return serving.validate_serving_measurements( + case.recipe_path, + prediction_report=case.prediction_report, + execution_evidence_path=case.observed_path, + output_report=case.recipe_path.parent / "assessment.json", + **kwargs, + ) + + +def _rewrite_records(case): + case.records.write_text("\n".join(json.dumps(row) for row in case.rows) + "\n") + receipt_path = case.recipe_path.parent / "benchmark-execution.json" + receipt = json.loads(receipt_path.read_text()) + receipt["artifacts"] = [ + serving._identity(case.records), + serving._identity(case.records.parent / "profile_export_aiperf.json"), + ] + _save(receipt_path, receipt) + + +def test_serving_gate_compares_native_requests_and_keeps_accuracy_scope(serving_case): + report = _assess(serving_case) + assert report["status"] == "passed" + assert report["forward"]["status"] == "unavailable" + assert report["scope"]["qualification"] == "one_complete_single_stream_weka_play" + assert report["metrics"]["ttft"]["p95_relative_error"] == 0 + assert report["metrics"]["tpot"]["p95_relative_error"] == 0 + assert report["metrics"]["output_throughput"]["relative_error"] < 1e-6 + assert len(report["samples"]) == 2 + assert report["samples"][0]["serving_request_id"] == "http-0" + assert serving.check_serving_validation_report(serving_case.recipe_path.parent / "assessment.json") == report + + +def test_target_tokenized_replay_preserves_source_identity_and_dependencies(serving_case): + case = serving_case + matched = case.recipe["matched_workload"] + rows = [json.loads(row) for row in Path(matched["trace"]["path"]).read_text().splitlines()] + original = json.loads(case.prediction_report.read_text())["per_request"] + assert rows[0]["schema"] == "dynamo.agentic_mooncake" + assert [row["request_id"] for row in rows[1:]] == [row["request_id"] for row in original] + assert rows[2]["dependencies"] == [ + {"request_id": original[0]["request_id"], "trigger": "completion", "relation": "sequence", "delay_ms": 100.0} + ] + assert rows[2]["not_before_ms"] == 100.0 + config = yaml.safe_load(Path(matched["prediction_config"]["path"]).read_text()) + source_config = yaml.safe_load(case.prediction_config.read_text()) + assert config["engine"] == source_config["engine"] + assert config["traffic"]["source"]["format"] == "agentic_mooncake" + assert "mooncake_trace" in case.recipe["command_arguments"] + + +def test_native_prefix_observation_must_match_replay_even_after_protocol_normalization(serving_case): + case = serving_case + receipt_path = case.recipe_path.parent / "benchmark-execution.json" + receipt = json.loads(receipt_path.read_text()) + path = Path(receipt["runtime_observations"][0]["path"]) + observed = json.loads(path.read_text()) + observed["resolved_config"]["cache_config"]["enable_prefix_caching"] = False + _save(path, observed) + receipt["runtime_observations"][0] = serving._identity(path) + _save(receipt_path, receipt) + report = _assess(case) + assert report["status"] == "failed" + assert any("prefix cache setting" in issue for issue in report["failures"]) + + +@pytest.mark.parametrize("mismatch", [False, True]) +def test_actual_server_tokenization_preflight_checks_every_frozen_payload(serving_case, monkeypatch, mismatch): + case = serving_case + tokenization = json.loads(Path(case.recipe["matched_workload"]["tokenization"]["path"]).read_text()) + calls = [] + + def response(request, **_kwargs): + index = len(calls) + calls.append(json.loads(request.data)) + expected = tokenization["requests"][index] + tokens = list(expected["input_token_ids"]) + if mismatch and index == 1: + tokens[-1] += 1 + return io.BytesIO(json.dumps({"tokens": tokens, "count": len(tokens)}).encode()) + + monkeypatch.setattr(serving, "urlopen", response) + if mismatch: + with pytest.raises(ValueError, match="tokenizer/chat template differs"): + serving._tokenize_serving(case.recipe_path, case.recipe) + else: + path = serving._tokenize_serving(case.recipe_path, case.recipe) + assert json.loads(path.read_text())["status"] == "matched" + assert len(calls) == 2 + assert all(not row["add_special_tokens"] for row in calls) + + +def test_missing_server_tokenization_cannot_qualify_serving(serving_case): + case = serving_case + path = case.recipe_path.parent / "benchmark-execution.json" + receipt = json.loads(path.read_text()) + receipt.pop("server_tokenization") + _save(path, receipt) + report = _assess(case) + assert report["status"] == "incomplete" + assert any("/tokenize" in issue for issue in report["missing_evidence"]) + + +def test_serving_dispatch_tolerance_is_editable_without_remeasurement(serving_case): + case = serving_case + case.rows[1]["metadata"]["request_start_ns"] += 6_000_000 + case.rows[1]["metadata"]["request_end_ns"] += 6_000_000 + _rewrite_records(case) + first = _assess(case) + assert first["status"] == "failed" + second = serving.validate_serving_measurements( + case.recipe_path, + prediction_report=case.prediction_report, + execution_evidence_path=case.observed_path, + output_report=case.recipe_path.parent / "new-policy.json", + max_dispatch_delay_ms=7.0, + ) + assert second["status"] == "passed" + assert second["policy"]["max_dispatch_delay_ms"] == 7.0 + + +def test_reassessment_rejects_forged_status_and_changed_source(serving_case): + report = _assess(serving_case) + report["status"] = "failed" + path = serving_case.recipe_path.parent / "assessment.json" + _save(path, report) + with pytest.raises(ValueError, match="independent reassessment"): + serving.check_serving_validation_report(path) + report["status"] = "passed" + _save(path, report) + serving_case.records.write_text(serving_case.records.read_text() + "\n") + with pytest.raises(ValueError, match="independent reassessment"): + serving.check_serving_validation_report(path) + + +# Prefill emits token one; the two decode steps read 64 and 65 past-KV tokens. +@pytest.mark.parametrize( + "api_time,expected_arrival", [(0.0, 300 + 2.5 + (2 + 64 / 1024) + (2 + 65 / 1024)), (0.3, 300.0)] +) +def test_native_single_stream_matches_supported_timing_subsets(request, api_time, expected_arrival): + args, _request, _plan, trace, output = request.getfixturevalue("_replay_validation_case") + play = json.loads(trace.read_text()) + play["requests"][0].update(t=120.0, api_time=api_time) + play["requests"][1].update(t=120.3, api_time=0.3) + _save(trace, play) + assert cli.main(args) == 0 + report = json.loads((output / "prediction/prediction.json").read_text()) + rows = report["per_request"] + assert rows[0]["arrival_time_ms"] == 0 + assert rows[1]["arrival_time_ms"] == pytest.approx(expected_arrival) + assert serving._timing_subset(play) == ("zero_idle_gaps" if api_time else "zero_recorded_api_time") + + +@pytest.mark.parametrize("field", serving.EXECUTION_FIELDS) +def test_missing_effective_execution_never_qualifies_accuracy(serving_case, field): + serving_case.observed["observed_execution"][field] = None + _save(serving_case.observed_path, serving_case.observed) + report = _assess(serving_case) + assert report["status"] == "incomplete" + assert any(field in issue for issue in report["missing_evidence"]) + + +@pytest.mark.parametrize("field", ["model_revision", "precision", "attention_groups", "graph_config", "parallelism"]) +def test_different_serving_execution_fails_even_if_latencies_match(serving_case, field): + serving_case.observed["observed_execution"][field] = "different" + _save(serving_case.observed_path, serving_case.observed) + report = _assess(serving_case) + assert report["status"] == "failed" + assert any(field in issue for issue in report["failures"]) + + +@pytest.mark.parametrize("change", ["warmup", "cancelled", "wrong_shape", "duplicate", "wrong_units", "missing"]) +def test_bad_request_evidence_does_not_pass(serving_case, change): + row = serving_case.rows[0] + if change == "warmup": + row["metadata"]["benchmark_phase"] = "warmup" + elif change == "cancelled": + row["metadata"]["was_cancelled"] = True + elif change == "wrong_shape": + row["metrics"]["input_sequence_length"]["value"] += 1 + elif change == "duplicate": + serving_case.rows[1] = copy.deepcopy(row) + elif change == "wrong_units": + row["metrics"]["time_to_first_token"]["unit"] = "seconds" + else: + serving_case.rows.pop() + _rewrite_records(serving_case) + report = _assess(serving_case) + assert report["status"] != "passed" + assert report["metrics"]["output_throughput"]["status"] == "incomplete" + + +def test_latency_and_throughput_are_independent_gates_with_editable_limits(serving_case): + for row in serving_case.rows: + row["metrics"]["time_to_first_token"]["value"] *= 2 + _rewrite_records(serving_case) + report = _assess(serving_case) + assert report["status"] == "failed" + assert report["metrics"]["ttft"]["status"] == "failed" + assert report["metrics"]["tpot"]["status"] == "passed" + assert report["metrics"]["output_throughput"]["status"] == "passed" + second = serving.validate_serving_measurements( + serving_case.recipe_path, + prediction_report=serving_case.prediction_report, + execution_evidence_path=serving_case.observed_path, + output_report=serving_case.recipe_path.parent / "reviewed-policy-assessment.json", + latency_p95_error=0.6, + ) + assert second["status"] == "passed" + assert second["recipe"] == report["recipe"] + + +@pytest.mark.parametrize("argument", ["latency_p95_error", "throughput_error"]) +@pytest.mark.parametrize("limit", [False, True, -0.1, "0.2", float("nan"), float("inf"), float("-inf")]) +def test_serving_error_limits_reject_invalid_values_in_preparation_and_assessment( + serving_case, tmp_path, argument, limit +): + case = serving_case + with pytest.raises(ValueError, match=f"{argument} must"): + serving.prepare_serving_validation( + trace=case.trace, + prediction_config=case.prediction_config, + output=tmp_path / "invalid-limit", + endpoint="http://127.0.0.1:8000", + tokenizer=case.tokenizer, + expected_execution=case.recipe["expected_execution"], + aiperf_python=Path("/benchmark/bin/python"), + **{argument: limit}, + ) + with pytest.raises(ValueError, match=f"{argument} must"): + _assess(case, **{argument: limit}) + + +@pytest.mark.parametrize("change", ["nested", "overlap", "prefix_fork", "multiple_plays", "mixed_timing"]) +def test_unsupported_trace_preparation_remains_incomplete(serving_case, tmp_path, change): + case = serving_case + play = json.loads(case.trace.read_text()) + if change == "nested": + play["requests"].append({"type": "subagent", "requests": [play["requests"][0]]}) + elif change == "overlap": + play["requests"][0]["api_time"] = 1 + elif change == "prefix_fork": + play["requests"][1]["hash_ids"] = [8, 9] + elif change == "mixed_timing": + play["requests"][0]["api_time"] = 0.05 + case.trace.write_text(json.dumps(play) + "\n" + (json.dumps(play) + "\n" if change == "multiple_plays" else "")) + recipe_path = serving.prepare_serving_validation( + trace=case.trace, + prediction_config=case.prediction_config, + output=tmp_path / "unsupported", + endpoint="http://localhost:8000", + tokenizer=case.tokenizer, + expected_execution=case.recipe["expected_execution"], + ) + recipe = json.loads(recipe_path.read_text()) + assert recipe["status"] == "incomplete" + assert recipe["issues"] + assert recipe["command_arguments"] is None + with pytest.raises(ValueError, match="preparation is incomplete"): + serving.run_serving_benchmark(recipe_path, aiperf_python=Path("python")) + + +def test_real_producer_command_checks_pin_and_does_not_execute_saved_commands(serving_case, monkeypatch): + case = serving_case + (case.recipe_path.parent / "benchmark-execution.json").unlink() + case.records.unlink() + (case.records.parent / "profile_export_aiperf.json").unlink() + case.records.parent.rmdir() + calls = [] + + def run(command, **kwargs): + calls.append((command, kwargs)) + if command[1] == "-c": + return SimpleNamespace( + stdout=json.dumps( + { + "version": "pinned", + "source": { + "url": serving.AIPERF_SOURCE + ".git", + "vcs_info": {"commit_id": serving.AIPERF_REVISION}, + }, + } + ) + ) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(serving.subprocess, "run", run) + monkeypatch.setattr(serving, "_tokenize_serving", write_server_tokenization) + receipt_path = serving.run_serving_benchmark(case.recipe_path, aiperf_python=Path("/benchmark/bin/python")) + assert json.loads(receipt_path.read_text())["status"] == "completed" + assert calls[1][0][1:] == case.recipe["command_arguments"] + assert "--no-fixed-schedule" in calls[1][0] + assert "--scenario" not in calls[1][0] + assert "--warmup-request-count" not in calls[1][0] + case.recipe["command_arguments"] = ["-c", "raise AssertionError('must not execute')"] + _save(case.recipe_path, case.recipe) + with pytest.raises(ValueError, match="differs from the fixed"): + serving.run_serving_benchmark(case.recipe_path, aiperf_python=Path("/benchmark/bin/python")) + assert len(calls) == 2 + + +def test_changed_trace_is_rejected_before_execution(serving_case): + case = serving_case + case.trace.write_text(case.trace.read_text() + "\n") + with pytest.raises(ValueError, match="changed after preparation"): + _assess(case) + + +def test_matched_replay_preserves_prefix_inside_source_block(serving_case, tmp_path, monkeypatch): + case = serving_case + original_materialize = serving._materialize_payloads + # Explicit synthetic tokenizer output matching the real-tokenizer review case: + # a 96-token shared prefix ends inside the source's second 64-token block. + token_ids = [list(range(97)), [*range(96), *range(100, 167)]] + + def materialize(recipe, output, python): + result = original_materialize(recipe, output, python) + _save( + output / "tokenization.json", + { + "requests": [ + {"turn_index": i, "input_token_ids": tokens, "input_length": len(tokens)} + for i, tokens in enumerate(token_ids) + ] + }, + ) + return result + + monkeypatch.setattr(serving, "_materialize_payloads", materialize) + prediction = yaml.safe_load(case.prediction_config.read_text()) + prediction["engine"]["workers"]["aggregated"]["kv_cache"]["block_size"] = 16 + source = json.loads(Path(case.recipe["benchmark_trace"]["path"]).read_text()) + assert source["block_size"] == 64 + matched = serving._prepare_matched_workload(case.recipe, prediction, tmp_path / "exact-prefix", Path("python")) + replay = json.loads(Path(matched["prediction_report"]["path"]).read_text()) + graph = [json.loads(line) for line in Path(matched["trace"]["path"]).read_text().splitlines()] + original = json.loads(case.prediction_report.read_text()) + assert replay["completed_requests"] == 2 + assert replay["per_request"][1]["reused_input_tokens"] == 96 + assert graph[0]["block_size"] == 1 + assert [row["hash_ids"] for row in graph[1:]] == token_ids + for index, (row, old) in enumerate(zip(graph[1:], original["per_request"], strict=True)): + for key in ("request_id", "play_id", "session_id"): + assert row[key] == old[key] + assert row["not_before_ms"] == source["requests"][index]["t"] * 1000 + assert row["recorded_api_time_ms"] == 0 + assert row["output_length"] == source["requests"][index]["out"] + assert graph[1]["dependencies"] == [] + assert graph[2]["dependencies"] == [ + {"request_id": graph[1]["request_id"], "trigger": "completion", "relation": "sequence", "delay_ms": 100} + ] + coverage = json.loads(Path(matched["coverage_report"]["path"]).read_text()) + assert coverage["status"] == "covered" + assert coverage["queries"]["unsupported"] == 0 + + +def test_supplied_forward_samples_are_separate_from_request_latency(serving_case): + case = serving_case + forward = case.recipe_path.parent / "forward.json" + _save( + forward, + { + "recipe": serving._identity(case.recipe_path), + "sources": [serving._identity(case.records)], + "samples": { + "prefill": [ + { + "scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 64}, + "measured_ms": 0.01, + "unit": "ms", + } + ], + "decode": [ + { + "scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 65}, + "measured_ms": 2, + "unit": "ms", + } + ], + }, + }, + ) + report = _assess(case, forward_evidence_path=forward) + assert report["status"] == "failed" + assert report["forward"]["phases"]["prefill"]["status"] == "failed" + assert report["forward"]["phases"]["decode"]["status"] == "passed" + assert report["metrics"]["ttft"]["status"] == "passed" diff --git a/python/aisimulate/tests/unit/test_support_topology.py b/python/aisimulate/tests/unit/test_support_topology.py new file mode 100644 index 000000000..3120d5dd6 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_topology.py @@ -0,0 +1,626 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Adapted and modified from vLLM at commit 98dff2a81d747d1dba01a47f939f48c3526d4206: +# https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py + +"""Original synthetic geometries exercise onboarding's conservative byte precheck.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml + +from aisimulate.support import config_profile, topology +from aisimulate.support.config_profile import derive_profile, load_model_config, packaged_hardware_path +from aisimulate.support.schema import SupportRequest +from aisimulate.support.topology import suggest_topologies + +pytestmark = pytest.mark.unit + + +def _config(tmp_path, **updates): + raw = { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "vocab_size": 64, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", + } + raw.update(updates) + path = tmp_path / "config.json" + path.write_text(json.dumps(raw), encoding="utf-8") + return load_model_config(path) + + +def _large_config(tmp_path, **updates): + return _config( + tmp_path, + **{ + "hidden_size": 8192, + "intermediate_size": 28672, + "num_hidden_layers": 80, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 32768, + **updates, + }, + ) + + +def _request(kind="dense", gpu="h200_sxm", concurrency=2, interconnect="NVLink", model="example/synthetic", **search): + return SupportRequest.model_validate( + { + "identity": { + "model": model, + "model_revision": "checkpoint-123", + "model_kind": kind, + "framework_version": "0.25.1", + "gpu": gpu, + "interconnect": interconnect, + "tokenizer_revision": "tokenizer-456", + "chat_template_revision": "template-789", + }, + "workload": {"concurrency": concurrency, "request_count": max(4, concurrency)}, + "search": {"context_length": 4096, **search}, + } + ) + + +def _runtime(**updates): + return {"kv_cache_dtype": "bfloat16", "fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", **updates} + + +def _hardware(tmp_path, monkeypatch, **updates): + spec = { + "gpu": {"mem_capacity": 16 * 1024**3}, + "node": {"num_gpus_per_node": 8, "intra_node_bw": 1000, "inter_node_bw": 100}, + "misc": {"other_mem": 0, "nccl_mem": {1: 0, 2: 0, 4: 0, 8: 0}}, + **updates, + } + path = tmp_path / "hardware.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + monkeypatch.setattr(topology, "packaged_hardware_path", lambda system: path) + monkeypatch.setattr(config_profile, "packaged_hardware_path", lambda system: path) + return spec, path + + +def test_model_size_and_hardware_capacity_change_starting_widths(tmp_path): + small = suggest_topologies(_config(tmp_path), _request(), _runtime()) + assert [(item.family, item.required_gpus) for item in small.candidates] == [("tp", 1), ("tp", 2)] + large_config = _large_config(tmp_path) + large = suggest_topologies(large_config, _request(), _runtime()) + assert [item.required_gpus for item in large.candidates] == [2, 4] + assert large.default.required_gpus == 2 + assert large.rejected_candidates[0].required_gpus == 1 + assert "lower bound" in large.rejected_candidates[0].reasons[0] + bigger_gpu = suggest_topologies(large_config, _request(gpu="gb200"), _runtime()) + assert [item.required_gpus for item in bigger_gpu.candidates] == [1, 2] + assert bigger_gpu.default.required_gpus == 1 + assert bigger_gpu.hardware.per_gpu_bytes > large.hardware.per_gpu_bytes + assert all(item.status == "estimated_fit" for item in large.candidates + bigger_gpu.candidates) + + +def test_grouped_fit_uses_windowed_native_peak_instead_of_linear_context_rate(tmp_path, monkeypatch): + # The budget fits the short window plus a prefill chunk but not the full + # retained context at TP1. The existing linear route requires TP2. + non_kv = 13472 + 70 * 1024**2 + capacity = ((non_kv + 100_000) * 10 + 8) // 9 + _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": capacity}) + request = _request(context_length=2048, model=str(tmp_path)) + runtime = _runtime(max_num_tokens=256, max_batch_size=8) + linear = suggest_topologies(_config(tmp_path), request, runtime) + windowed = suggest_topologies( + _config(tmp_path, sliding_window=128), + request, + {**runtime, "cache_block_sizes": {"sliding_attention": 16}}, + ) + assert linear.default.required_gpus == 2 + assert windowed.default.required_gpus == 1 + assert windowed.default.estimated_required_bytes < windowed.hardware.memory_budget_bytes + assert windowed.default.draft.profile.deployments[0].resources.kv_bytes_per_token is None + assert any("Rust's window" in item for item in windowed.assumptions) + + +def test_grouped_missing_runtime_blocks_never_establish_fit(tmp_path): + report = suggest_topologies(_config(tmp_path, sliding_window=128), _request(), _runtime()) + assert report.default is None + assert report.candidates + assert all(candidate.status == "needs_inputs" for candidate in report.candidates) + assert all("cache_block_sizes" in candidate.missing for candidate in report.candidates) + + +def test_grouped_non_kv_budget_equality_does_not_admit_a_worker(tmp_path, monkeypatch): + non_kv = 13472 + 70 * 1024**2 + _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": (non_kv * 10 + 8) // 9}) + report = suggest_topologies( + _config(tmp_path, sliding_window=128), + _request(model=str(tmp_path)), + _runtime(cache_block_sizes={"sliding_attention": 16}), + ) + candidate = next(candidate for candidate in report.rejected_candidates if candidate.required_gpus == 1) + assert candidate.known_required_bytes == candidate.memory_budget_bytes + assert "no grouped cache budget" in candidate.reasons[0] + + +def test_inkling_topology_rejects_attention_tp_but_preserves_large_dep_candidates(tmp_path): + # Original synthetic counts with the pinned Inkling convolution semantics; + # this is not a checkpoint fixture or a declaration of runtime qualification. + # Modified metadata-only test of the Apache-2.0 vLLM packing constraint: + # https://github.com/vllm-project/vllm/blob/98dff2a81d747d1dba01a47f939f48c3526d4206/vllm/models/inkling/nvidia/sconv_swa_attn.py + text = { + "hidden_size": 64, + "num_hidden_layers": 4, + "num_attention_heads": 64, + "num_key_value_heads": 2, + "head_dim": 4, + "model_max_length": 4096, + "local_layer_ids": [0, 1, 2], + "sliding_window_size": 128, + "swa_num_key_value_heads": 4, + "use_sconv": True, + "sconv_kernel_size": 4, + "n_routed_experts": 64, + "torch_dtype": "bfloat16", + } + config = _config( + tmp_path, architectures=["InklingForConditionalGeneration"], model_type="inkling_mm_model", text_config=text + ) + report = suggest_topologies(config, _request("moe", gpu="gb200"), _runtime()) + assert report.default is None # Unknown decoder weight/activation bounds remain explicit. + assert not any(candidate.family == "dep" for candidate in report.rejected_candidates) + rejected = [candidate for candidate in report.rejected_candidates if candidate.sizes[0] > 2] + assert rejected and all("attention TP <= 2" in candidate.reasons[0] for candidate in rejected) + assert any(candidate.family == "dep" and candidate.required_gpus == 4 for candidate in report.candidates) + + +def test_grouped_topology_needs_current_extension_with_actionable_failure(tmp_path, monkeypatch): + from aisimulate_core.sdk.rust_engine_step import RustForwardPassPerfModel + + def missing(*args, **kwargs): + raise AttributeError("old extension") + + monkeypatch.setattr(RustForwardPassPerfModel, "estimate_cache_budget", missing) + with pytest.raises(ValueError, match="compiled AISimulate extension"): + suggest_topologies( + _config(tmp_path, sliding_window=128), + _request(), + _runtime(cache_block_sizes={"sliding_attention": 16}), + ) + + +def test_moe_families_have_exact_tuples_and_deduplicate_width_one(tmp_path): + config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) + report = suggest_topologies(config, _request("moe"), _runtime()) + assert [(item.family, item.sizes) for item in report.candidates] == [ + ("pure_tp", (1, 1, 1, 1)), + ("pure_tp", (2, 1, 2, 1)), + ("dep", (1, 2, 1, 2)), + ("dep", (1, 4, 1, 4)), + ("tep", (2, 1, 1, 2)), + ("tep", (4, 1, 1, 4)), + ] + assert len({item.sizes for item in report.candidates}) == 6 + assert report.default.family == "pure_tp" + for item in report.candidates: + assert item.apply(_request("moe")).worker_gpus == item.required_gpus + if item.family in {"dep", "tep"}: + assert item.status == "needs_inputs" + assert not item.missing + assert "comm_overhead_bytes" not in item.draft.resolved + assert item.draft.profile.deployments[0].resources.memory_source == "pending" + assert item.estimated_required_bytes is None + + +def test_moe_expert_and_attention_divisibility_rejects_only_incompatible_tuples(tmp_path): + config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=6) + report = suggest_topologies(config, _request("moe"), _runtime()) + assert [(item.family, item.required_gpus) for item in report.candidates if item.family != "pure_tp"] == [ + ("dep", 2), + ("tep", 2), + ] + rejected = {(item.family, item.required_gpus): item for item in report.rejected_candidates} + assert "num_experts" in rejected["dep", 4].reasons[0] + assert "num_experts" in rejected["tep", 4].reasons[0] + assert "num_attention_heads" in rejected["pure_tp", 8].reasons[0] + + +@pytest.mark.parametrize("gpu", ["gb200", "gb300"]) +def test_declared_high_speed_rack_domain_is_not_limited_to_physical_node(tmp_path, gpu): + config = _large_config( + tmp_path, + hidden_size=32768, + intermediate_size=131072, + num_hidden_layers=256, + num_attention_heads=256, + num_key_value_heads=32, + ) + report = suggest_topologies(config, _request(gpu=gpu), _runtime()) + assert report.hardware.gpus_per_node == 4 + assert report.hardware.gpus_per_rack == report.hardware.fast_domain_gpus == 72 + assert report.hardware.fast_domain == "rack" + assert "inter_node_bw >= intra_node_bw" in report.hardware.assumption + assert "not available GPU capacity" in report.hardware.assumption + assert report.candidates[-1].required_gpus == 64 + assert all(item.status == "needs_inputs" for item in report.candidates) + assert report.default is None + assert "activations_bytes" not in report.candidates[-1].draft.resolved + assert "comm_overhead_bytes" not in report.candidates[-1].draft.resolved + + +def test_slower_rack_connectivity_keeps_node_domain(tmp_path, monkeypatch): + _hardware( + tmp_path, + monkeypatch, + node={"num_gpus_per_node": 4, "num_gpus_per_rack": 72, "intra_node_bw": 1000, "inter_node_bw": 100}, + ) + report = suggest_topologies(_config(tmp_path), _request(), _runtime()) + assert report.hardware.fast_domain == "node" + assert report.hardware.fast_domain_gpus == 4 + assert report.hardware.gpus_per_rack == 72 + assert all(item.required_gpus <= 4 for item in report.candidates + report.rejected_candidates) + + +@pytest.mark.parametrize("interconnect", ["PCIe", "IB", "EFA", "custom-fabric"]) +def test_declared_fabric_does_not_silently_assume_nvlink_rack(tmp_path, interconnect): + report = suggest_topologies(_config(tmp_path), _request(gpu="gb200", interconnect=interconnect), _runtime()) + assert report.hardware.declared_interconnect == interconnect + assert report.hardware.fast_domain == "node" + assert report.hardware.fast_domain_gpus == 4 + assert "fabric compatibility is unverified" in report.hardware.assumption + + +def test_no_declared_interconnect_limits_automatic_choices_to_one_gpu(tmp_path): + report = suggest_topologies(_config(tmp_path), _request(gpu="gb200", interconnect="none"), _runtime()) + assert report.hardware.fast_domain == "single_gpu" + assert report.hardware.fast_domain_gpus == 1 + assert len(report.candidates) == 1 + assert report.default.required_gpus == 1 + large = suggest_topologies(_large_config(tmp_path), _request(interconnect="none"), _runtime()) + assert not large.candidates + assert large.default is None + assert [item.required_gpus for item in large.rejected_candidates] == [1] + + +@pytest.mark.parametrize( + "config_updates,overrides,missing", + [ + ({}, {}, {"fmha_quant_mode", "kv_cache_dtype", "activations_bytes", "kv_bytes_per_token"}), + ({"torch_dtype": None}, _runtime(), {"gemm_quant_mode", "moe_quant_mode", "weights_bytes"}), + ({"quantization_config": {"quant_method": "fp8"}}, _runtime(), {"weights_bytes"}), + ({"num_key_value_heads": None}, _runtime(), {"kv_bytes_per_token", "weights_bytes"}), + ], +) +def test_missing_precision_weight_or_cache_inputs_never_fabricate_fit(tmp_path, config_updates, overrides, missing): + report = suggest_topologies(_config(tmp_path, **config_updates), _request(), overrides) + assert report.default is None + assert report.candidates + for candidate in report.candidates: + assert candidate.status == "needs_inputs" + assert ( + missing - {"weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"} + <= candidate.missing.keys() + ) + assert ( + not {"weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"} + & candidate.missing.keys() + ) + assert candidate.estimated_required_bytes is None + assert candidate.known_required_bytes > 0 + + +def test_unknown_layout_preserves_unchecked_status(tmp_path): + config = _config(tmp_path, architectures=["CustomDecoderForCausalLM"], model_type="custom", num_experts=0) + report = suggest_topologies(config, _request(), _runtime()) + assert report.default is None + assert {"kv_bytes_per_token", "cache_layout"} <= report.candidates[0].missing.keys() + assert any("Unknown architecture constraints" in reason for reason in report.assumptions) + + +def test_collector_communication_identity_is_proposed_with_source(tmp_path): + report = suggest_topologies( + _config(tmp_path), _request(), {"kv_cache_dtype": "bfloat16", "fmha_quant_mode": "bfloat16"} + ) + assert report.default is not None + assert report.candidates[0].status == "estimated_fit" + assert report.candidates[0].estimated_required_bytes is not None + assert not report.candidates[0].missing + assert report.candidates[0].draft.resolved["comm_quant_mode"] == "half" + assert "collector FPM identity default" in report.candidates[0].draft.sources["comm_quant_mode"] + + +def test_known_weight_lower_bound_can_reject_without_precision_cache_or_overheads(tmp_path, monkeypatch): + _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": 1024}, misc={}) + report = suggest_topologies(_config(tmp_path), _request()) + assert report.default is None + assert not report.candidates + rejected = report.rejected_candidates[0] + assert rejected.status == "rejected" + assert rejected.known_required_bytes == 13472 + assert rejected.estimated_required_bytes is None + assert rejected.known_required_bytes > rejected.memory_budget_bytes + assert { + "kv_bytes_per_token", + } <= rejected.missing.keys() + + +def test_cache_precheck_requires_one_full_context_independent_of_scheduler_batch_limit(tmp_path): + config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) + request = _request("moe", concurrency=8) + report = suggest_topologies(config, request, _runtime(max_batch_size=3, max_num_tokens=64)) + dep4 = next(item for item in report.candidates if item.family == "dep" and item.required_gpus == 4) + known_non_kv = sum( + dep4.draft.resolved[field] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes") + if field in dep4.draft.resolved + ) + assert dep4.resident_sequences_per_rank == 1 + assert dep4.known_required_bytes - known_non_kv == 64 * (4096 + 1) + assert dep4.draft.resolved["max_num_tokens"] == 64 + assert report.context_length == 4096 + assert report.collection["max_sequences"] == 3 + assert report.collection["max_batched_tokens"] == 64 + assert "workload" not in report.to_dict() + assert any("not guarantee" in item for item in report.assumptions) + + +def test_validation_concurrency_does_not_change_topology_fit(tmp_path, monkeypatch): + _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": 85_000_000}) + config = _config(tmp_path) + light = suggest_topologies(config, _request(concurrency=1), _runtime()) + assert light.default.required_gpus == 1 + heavy = suggest_topologies(config, _request(concurrency=16), _runtime()) + assert heavy.to_dict() == light.to_dict() + + +def test_gpu_memory_fraction_changes_topology_admission(tmp_path, monkeypatch): + _hardware(tmp_path, monkeypatch) + config = _config(tmp_path) + request = _request() + first = suggest_topologies(config, request, _runtime()).default + _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": first.estimated_required_bytes * 2}) + high = suggest_topologies(config, request, _runtime()) + payload = request.model_dump() + payload["collection"]["gpu_memory_utilization"] = 0.4 + low = suggest_topologies(config, SupportRequest.model_validate(payload), _runtime()) + assert high.default.required_gpus == 1 + assert low.hardware.memory_budget_bytes == int(low.hardware.per_gpu_bytes * 0.4) + assert any(candidate.required_gpus == 1 and candidate.status == "rejected" for candidate in low.rejected_candidates) + assert all(candidate.required_gpus > 1 for candidate in low.candidates) + assert "40%" in low.rejected_candidates[0].reasons[0] + + +@pytest.mark.parametrize("extra_token", [0, 1]) +def test_full_context_boundary_matches_existing_plan_admission(tmp_path, monkeypatch, extra_token): + from aisimulate.support.plan import create_plan + from aisimulate_core.sdk import perf_database + + # Hand-counted TP1 tensors plus the existing 70 MiB minimum activation estimate. + total = 13472 + 70 * 1024**2 + 64 * (4096 + extra_token) + capacity = (total * 10 + 8) // 9 + spec, _ = _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": capacity}) + monkeypatch.setattr(perf_database, "load_system_spec", lambda *args, **kwargs: spec) + config, request = _config(tmp_path), _request(concurrency=1) + report = suggest_topologies(config, request, _runtime()) + payload = request.model_dump() + payload["fpm_profile"] = derive_profile( + config, + request, + { + **_runtime(), + "weights_bytes": 13472, + "activations_bytes": 70 * 1024**2, + "runtime_overhead_bytes": 0, + "comm_overhead_bytes": 0, + }, + ).profile + exact = SupportRequest.model_validate(payload) + if extra_token: + assert report.default.required_gpus == 1 + plan = create_plan(exact, tmp_path / "plan") + assert plan["resources"]["total_kv_size_tokens"] == 4097 + else: + assert report.default.required_gpus == 2 + assert report.rejected_candidates[0].required_gpus == 1 + with pytest.raises(ValueError, match="insufficient rank-local KV capacity"): + create_plan(exact, tmp_path / "plan") + + +@pytest.mark.parametrize( + "field", + ["weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes", "kv_bytes_per_token"], +) +def test_flat_per_rank_byte_overrides_require_explicit_topology(tmp_path, field): + with pytest.raises( + ValueError, match="cannot transfer rank-local byte overrides.*" + field + ".*select an exact topology" + ): + suggest_topologies(_config(tmp_path), _request(), _runtime(**{field: 64})) + + +def test_attached_exact_profile_is_not_silently_replaced(tmp_path): + config, request = _config(tmp_path), _request() + payload = request.model_dump() + payload["fpm_profile"] = derive_profile( + config, + request, + { + **_runtime(), + "weights_bytes": 13472, + "activations_bytes": 70 * 1024**2, + "runtime_overhead_bytes": 0, + "comm_overhead_bytes": 0, + }, + ).profile + request = SupportRequest.model_validate(payload) + with pytest.raises(ValueError, match="attached FPM profile.*exact topology"): + suggest_topologies(config, request, _runtime()) + + +@pytest.mark.parametrize( + "overrides", [{"max_batch_size": True}, {"max_num_tokens": "128"}, {"fmha_quant_mode": "auto"}] +) +def test_invalid_shared_overrides_fail_explicitly(tmp_path, overrides): + with pytest.raises(ValueError): + suggest_topologies(_config(tmp_path), _request(), overrides) + + +@pytest.mark.parametrize( + "updates,match", + [ + ({"gpu": {}}, "gpu.mem_capacity"), + ({"gpu": {"mem_capacity": True}}, "gpu.mem_capacity"), + ({"gpu": {"mem_capacity": "1024"}}, "gpu.mem_capacity"), + ({"node": None}, "node must be an object"), + ({"node": {}}, "num_gpus_per_node"), + ({"node": {"num_gpus_per_node": 4}}, "intra_node_bw"), + ({"node": {"num_gpus_per_node": 4, "intra_node_bw": float("nan"), "inter_node_bw": 1}}, "intra_node_bw"), + ({"node": {"num_gpus_per_node": 4, "intra_node_bw": 1, "inter_node_bw": True}}, "inter_node_bw"), + ( + {"node": {"num_gpus_per_node": 4, "intra_node_bw": 1, "inter_node_bw": 1, "num_gpus_per_rack": 6}}, + "num_gpus_per_rack", + ), + ({"misc": None}, "misc must be an object"), + ({"misc": {"other_mem": -1}}, "other_mem"), + ({"misc": {"nccl_mem": []}}, "nccl_mem"), + ({"misc": {"nccl_mem": {"1": 0}}}, "TP key"), + ({"misc": {"nccl_mem": {128: True}}}, "nccl_mem"), + ], +) +def test_partial_or_malformed_packaged_hardware_fails_explicitly(tmp_path, monkeypatch, updates, match): + _hardware(tmp_path, monkeypatch, **updates) + with pytest.raises(ValueError, match=match): + suggest_topologies(_config(tmp_path), _request(), _runtime()) + + +def test_unknown_packaged_hardware_fails_without_fallback(tmp_path): + with pytest.raises(ValueError, match="no packaged hardware metadata for nonexistent_gpu"): + suggest_topologies(_config(tmp_path), _request(gpu="nonexistent_gpu"), _runtime()) + + +@pytest.mark.parametrize("system", ["../h200_sxm", "/tmp/h200_sxm", "h200_sxm.yaml", "", None]) +def test_hardware_path_helper_rejects_paths(system): + with pytest.raises(ValueError, match="packaged system name"): + packaged_hardware_path(system) + + +@pytest.mark.parametrize( + "config_updates,kind,search,match", + [ + ({}, "moe", {}, "model_kind"), + ({}, "dense", {"context_length": 8192}, "context_length"), + ({"hidden_size": 17, "head_dim": None}, "dense", {}, "hidden_size"), + ], +) +def test_global_model_request_conflicts_are_not_hidden_as_candidate_rejections( + tmp_path, config_updates, kind, search, match +): + with pytest.raises(ValueError, match=match): + suggest_topologies(_config(tmp_path, **config_updates), _request(kind, **search), _runtime()) + + +def test_input_provenance_is_preserved_and_serialization_excludes_profile_objects(tmp_path): + config = _config(tmp_path, _name_or_path="source/different-id") + request = _request(tensor_parallel=4, objective="ttft", seed=17) + overrides = _runtime(max_batch_size=1, max_num_tokens=1024, provenance="confirmed runtime settings") + before = (copy.deepcopy(config), request.model_dump(), dict(overrides)) + report = suggest_topologies(config, request, overrides) + selected = report.default.apply(request) + serialized = json.loads(json.dumps(report.to_dict())) + assert config == before[0] + assert request.model_dump() == before[1] + assert overrides == before[2] + assert selected.identity == request.identity + assert selected.workload == request.workload + assert selected.search.seed == 17 + assert selected.search.objective == "ttft" + assert selected.search.tensor_parallel == 1 + assert selected.fpm_profile is None + assert serialized["identity"]["model"] == "example/synthetic" + assert serialized["identity"]["model_revision"] == "checkpoint-123" + assert serialized["overrides"] == overrides + assert serialized["config_sha256"] == hashlib.sha256((tmp_path / "config.json").read_bytes()).hexdigest() + hardware_sha = hashlib.sha256(packaged_hardware_path("h200_sxm").read_bytes()).hexdigest() + assert serialized["hardware"]["sha256"] == hardware_sha + assert serialized["default"]["status"] == "estimated_fit" + assert serialized["default"]["cli_flags"] == [ + "--tensor-parallel", + "1", + "--attention-data-parallel", + "1", + "--moe-tensor-parallel", + "1", + "--moe-expert-parallel", + "1", + ] + assert "draft" not in serialized["candidates"][0] + assert "profile" not in serialized["candidates"][0] + provenance = json.loads(report.default.draft.profile.provenance) + assert provenance["config_sha256"] == config.sha256 + assert provenance["deployment_identity"] == request.identity.model_dump() + assert hardware_sha in provenance["fields"]["runtime_overhead_bytes"]["source"] + assert "user override" in serialized["default"]["field_sources"]["provenance"] + + +@pytest.mark.parametrize("nested", [False, True]) +def test_cold_process_suggestions_do_not_import_native_runtime_models_or_timing_readers(tmp_path, nested): + _config(tmp_path) + if nested: + path = tmp_path / "config.json" + raw = json.loads(path.read_text()) + path.write_text(json.dumps({"torch_dtype": raw.pop("torch_dtype"), "text_config": raw})) + code = """ +import builtins, importlib.abc, json, sys +forbidden = ('aiconfigurator', 'aiconfigurator_core', 'aisimulate.sdk', 'aisimulate_core.sdk', + 'aisimulate_core._native', 'aisimulate._native', 'aisimulate._runtime', + 'aisimulate.runner', 'aisimulate.engine', 'transformers', 'huggingface_hub', + 'collector', 'numpy', 'pandas', 'pyarrow', 'torch') +class NoRuntime(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if any(fullname == name or fullname.startswith(name + '.') for name in forbidden): + raise AssertionError('unexpected dependency: ' + fullname) +sys.meta_path.insert(0, NoRuntime()) +from aisimulate.support.config_profile import load_model_config +from aisimulate.support.schema import SupportRequest +from aisimulate.support.topology import suggest_topologies +request = SupportRequest.model_validate(json.loads(sys.argv[2])) +report = suggest_topologies(load_model_config(sys.argv[1]), request, json.loads(sys.argv[3])) +assert report.default is not None +assert report.default.draft.resolved['runtime_overhead_bytes'] == 3758096384 +assert report.default.draft.resolved['comm_overhead_bytes'] == 0 +assert not any(name == prefix or name.startswith(prefix + '.') for name in sys.modules for prefix in forbidden) +core = {name for name in sys.modules if name == 'aisimulate_core' or name.startswith('aisimulate_core.')} +assert core <= {'aisimulate_core', 'aisimulate_core.fpm_profile', 'aisimulate_core.quantization'}, core +print(json.dumps(report.to_dict())) +""" + source = Path(__file__).parents[2] / "src" + result = subprocess.run( + [ + sys.executable, + "-c", + code, + str(tmp_path / "config.json"), + _request().model_dump_json(), + json.dumps(_runtime()), + ], + env={**os.environ, "PYTHONPATH": str(source)}, + text=True, + capture_output=True, + check=True, + ) + report = json.loads(result.stdout) + assert report["default"]["status"] == "estimated_fit" diff --git a/python/aisimulate/tests/unit/test_support_validation.py b/python/aisimulate/tests/unit/test_support_validation.py new file mode 100644 index 000000000..9cb844246 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_validation.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Onboarding validates actual replay lookups, including failed and partial runs.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +import yaml + +from aisimulate import main as cli +from aisimulate import supervision +from aisimulate.support.plan import create_plan +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def validation_case(tmp_path, monkeypatch, request): + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate_core.sdk import engine + + def reject_graph(*_args, **_kwargs): + raise AssertionError("FPM replay must not construct an analytical model") + + monkeypatch.setattr(engine, "get_model", reject_graph) + monkeypatch.setattr(engine, "build_model_config", reject_graph) + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["UnregisteredDecoderForCausalLM"]})) + profile = { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "synthetic-v1", + "architecture": "UnregisteredDecoderForCausalLM", + "context_length": 512, + "num_experts": 0, + "provenance": "Synthetic test metadata; not silicon qualification.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 2, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 128, + "max_batch_size": 2, + "provenance": "Declared synthetic rank-local bounds.", + }, + } + ], + } + topology = getattr(request, "param", "tp") + parallel = {"tensor_parallel": 2, "context_length": 512} + if topology in {"dep", "tep"}: + profile["num_experts"] = 4 + profile["deployments"][0].update(tp=1 if topology == "dep" else 2, dp=2 if topology == "dep" else 1, moe_ep=2) + parallel.update( + tensor_parallel=profile["deployments"][0]["tp"], + attention_data_parallel=profile["deployments"][0]["dp"], + moe_tensor_parallel=1, + moe_expert_parallel=2, + ) + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "dense" if topology == "tp" else "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": parallel, + "collection": {"max_num_tokens": 128, "max_batch_size": 2, "gpu_memory_utilization": 0.73}, + "workload": {"input_tokens": 64, "output_tokens": 3, "request_count": 1}, + "fpm_profile": profile, + } + ) + plan = tmp_path / "collection" + create_plan(request, plan) + deployment = request.profile_deployment().model_dump(mode="json", exclude={"resources"}) + rows = [] + for batch in (1, 2): + for phase, tokens, kv_values in ( + ("prefill", (1, 64, 128, 256), (0, 64, 128, 512)), + ("decode", (0,), (1, 2, 64, 128, 1024)), + ): + for tokens_in_batch in tokens: + for kv in kv_values: + rows.append( + { + **deployment, + "model_path": profile["model"], + "cell_id": f"synthetic-{phase}", + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": batch, + "total_prefill_tokens": tokens_in_batch, + "total_kv_read_tokens": kv, + "latency_ms": 1 + batch + tokens_in_batch / 128 + kv / 1024, + "kv_seed_regime": "real_kv", + } + ) + table = plan / "systems/data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + table.parent.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), table) + metadata = { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(table.read_bytes()).hexdigest(), + } + table.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + trace = tmp_path / "trace.jsonl" + play = { + "id": "synthetic-play", + "models": ["trace/source-model"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "model": "trace/source-model", "in": 64, "out": 3, "hash_ids": [1]}, + { + "t": 0.1, + "type": "s", + "model": "trace/source-model", + "in": 128, + "out": 3, + "hash_ids": [1, 2], + }, + ], + } + trace.write_text(json.dumps(play) + "\n") + output = tmp_path / "validation" + args = [ + "onboard", + "validate-fpm", + "--config", + str(plan / "request.yaml"), + "--output-dir", + str(plan), + "--trace", + str(trace), + "--validation-output-dir", + str(output), + ] + # Exercise the public CLI compiler/runner/native runtime in-process so the + # no-analytical-graph assertion above also covers estimator construction. + monkeypatch.setattr(supervision, "main", cli.main) + return args, request, plan, trace, output + + +@pytest.mark.parametrize("validation_case", ["tp", "dep", "tep"], indirect=True) +def test_onboard_validation_replays_profile_without_graph_and_preserves_timing(validation_case, tmp_path): + args, request, plan, _trace, output = validation_case + saved_inputs = {path: path.read_bytes() for path in plan.rglob("*") if path.is_file()} + assert cli.main(args) == 0 + validation = json.loads((output / "validation.json").read_text()) + covered = json.loads((output / "prediction/prediction.json").read_text()) + assert validation["status"] == "covered" + assert validation["accuracy"] == "not_assessed" + assert validation["fpm_query_coverage"]["queries"]["measured"] > 0 + assert validation["fpm_query_coverage"]["queries"]["interpolated"] > 0 + assert validation["fpm_query_coverage"]["queries"]["unsupported"] == 0 + assert covered["completed_requests"] == 2 + assert covered["agentic_model_projection"]["target_model"] == request.identity.model + assert covered["agentic_graph"]["source_models"] == ["trace/source-model"] + prediction = yaml.safe_load((output / "predict.yaml").read_text()) + engine = prediction["engine"] + assert engine["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert engine["workers"]["aggregated"]["scheduler"] == request.scheduler_limits() + assert engine["workers"]["aggregated"]["kv_cache"]["capacity"]["memory_fraction"] == 0.73 + assert saved_inputs == {path: path.read_bytes() for path in saved_inputs} + + # Coverage is diagnostic only: cold/warm prefix behavior and every native + # request output remain identical when collection is disabled. + ordinary = copy.deepcopy(prediction) + ordinary["engine"]["workers"]["aggregated"]["timing"]["estimator_config"]["fpm_interpolation"].pop( + "collect_coverage" + ) + ordinary_path = tmp_path / "ordinary.yaml" + ordinary_path.write_text(yaml.safe_dump(ordinary)) + ordinary_output = tmp_path / "ordinary" + assert ( + cli.main( + ["predict", "--config", str(ordinary_path), "--output-dir", str(ordinary_output), "--capture-per-request"] + ) + == 0 + ) + uncollected = json.loads((ordinary_output / "prediction.json").read_text()) + assert "fpm_query_coverage" not in uncollected + for key in ("per_request", "agentic_graph", "agentic_play_outcomes", "prefix_cache_reused_ratio"): + assert covered[key] == uncollected[key] + + +def test_onboard_validation_keeps_missing_query_evidence_on_native_failure(validation_case): + import pyarrow as pa + import pyarrow.parquet as pq + + args, _request, plan, trace, output = validation_case + play = json.loads(trace.read_text()) + play["requests"][0]["in"] = 384 + play["requests"][0]["hash_ids"] = list(range(6)) + trace.write_text(json.dumps(play) + "\n") + table = plan / "systems/data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + rows = [ + row + for row in pq.read_table(table).to_pylist() + if row["workload_kind"] != "prefill" or row["total_kv_read_tokens"] <= 128 + ] + pq.write_table(pa.Table.from_pylist(rows), table) + metadata_path = table.with_suffix(".metadata.json") + metadata = json.loads(metadata_path.read_text()) + metadata.update(row_count=len(rows), parquet_sha256=hashlib.sha256(table.read_bytes()).hexdigest()) + metadata_path.write_text(json.dumps(metadata)) + # The third prefill chunk needs a longer KV bracket than this table has. + assert cli.main(args) != 0 + validation = json.loads((output / "validation.json").read_text()) + assert validation["status"] == "incomplete" + assert validation["issues"] + assert validation["prediction_exit_code"] != 0 + coverage = validation["fpm_query_coverage"] + assert coverage["queries"]["unsupported"] > 0 + assert coverage["queries"]["measured"] > 0 + assert coverage["roles"][0]["coverage"]["gaps"][0]["coordinates"]["total_kv_read_tokens"] == 256 + + +def test_onboard_validation_rejects_changed_collection_identity_before_replay(validation_case, monkeypatch): + args, _request, plan, _trace, output = validation_case + config = yaml.safe_load((plan / "request.yaml").read_text()) + config["identity"]["model_revision"] = "different-revision" + supplied = plan.parent / "different.yaml" + supplied.write_text(yaml.safe_dump(config)) + args[args.index("--config") + 1] = str(supplied) + monkeypatch.setattr(supervision, "main", lambda _args: pytest.fail("replay must not start")) + with pytest.raises(SystemExit) as error: + cli.main(args) + assert error.value.code == 2 + assert not output.exists() + + +@pytest.mark.parametrize("flag", ["--trace", "--config"]) +def test_validation_overwrite_preserves_inputs_in_output_directory(validation_case, flag): + args, _request, _plan, _trace, output = validation_case + source = Path(args[args.index(flag) + 1]) + before = source.read_bytes() + output.mkdir() + input_path = output / "prediction.json" + input_path.write_bytes(before) + args[args.index(flag) + 1] = str(input_path) + with pytest.raises(SystemExit) as error: + cli.main([*args, "--overwrite"]) + assert error.value.code == 2 + assert input_path.read_bytes() == before + + +@pytest.mark.parametrize("failure", ["failed_request", "incomplete_play", "truncated_output", "zero_queries"]) +def test_validation_never_accepts_partial_replay(validation_case, monkeypatch, failure): + args, _request, _plan, _trace, output = validation_case + + def partial_prediction(command): + assert command[0] == "predict" + assert json.loads((output / "validation.json").read_text())["status"] == "incomplete" + report = { + "num_requests": 1, + "completed_requests": 1, + "agentic_graph": {"node_count": 1, "play_count": 1}, + "per_request": [{"terminal_status": "completed", "output_length": 2, "requested_output_length": 2}], + "agentic_play_outcomes": [{"status": "completed", "settled_at_ms": 5}], + } + coverage = {"status": "covered", "queries": {"measured": 1, "interpolated": 0, "unsupported": 0}} + if failure == "failed_request": + report["per_request"][0]["terminal_status"] = "failed" + elif failure == "incomplete_play": + report["agentic_play_outcomes"][0]["status"] = "incomplete" + elif failure == "truncated_output": + report["per_request"][0]["output_length"] = 1 + else: + coverage.update(status="incomplete", queries={"measured": 0, "interpolated": 0, "unsupported": 0}) + destination = Path(command[command.index("--output-dir") + 1]) + (destination / "prediction.json").write_text(json.dumps(report)) + (destination / "fpm-coverage.json").write_text(json.dumps(coverage)) + return 0 + + monkeypatch.setattr(supervision, "main", partial_prediction) + assert cli.main(args) == 1 + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == "incomplete" + assert report["issues"] + + +def test_validation_persists_incomplete_status_before_interrupted_prediction(validation_case, monkeypatch): + args, _request, _plan, _trace, output = validation_case + + def interrupt(_command): + assert json.loads((output / "validation.json").read_text())["status"] == "incomplete" + raise KeyboardInterrupt + + monkeypatch.setattr(supervision, "main", interrupt) + assert cli.main(args) == 130 + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == "incomplete" + assert "interrupted" in report["issues"][0] diff --git a/python/aisimulate/tests/unit/test_support_validation_workflow.py b/python/aisimulate/tests/unit/test_support_validation_workflow.py new file mode 100644 index 000000000..d7ba4bc88 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_validation_workflow.py @@ -0,0 +1,1165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public onboarding commands with synthetic native records, never silicon claims.""" + +from __future__ import annotations + +import hashlib +import json +import shutil +from dataclasses import replace +from pathlib import Path + +import pytest +import yaml +from collector.fpm_forward import cli as collector_cli +from collector.fpm_forward import planner, repeatability +from collector.fpm_forward.config import FPMCollectionOptions +from collector.fpm_forward.database import aggregate_cell, write_formal_database + +from aisimulate import main as cli +from aisimulate.support import serving_validation as serving +from aisimulate.support import validation_workflow as workflow +from aisimulate.support.fpm import fpm_cli_args +from aisimulate.support.plan import create_plan +from aisimulate.support.schema import SupportRequest + +from .collector.test_fpm_repeatability import _write_campaign +from .test_support_serving_validation import materialized_workload, write_server_tokenization # noqa: F401 +from .test_support_validation import validation_case # noqa: F401 + +pytestmark = pytest.mark.unit + + +def _write(path, value): + path.write_text(json.dumps(value)) + + +def _synthetic_campaign(root, checkpoint, plan, *, attempt_id, generator_overrides=None): + _write_campaign(root, checkpoint, plan, attempt_id=attempt_id, generator_overrides=generator_overrides) + for path in root.rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + for row in payload["results"]: + point = row["point"] + if point["point_type"] == "prefill" and point["total_kv_read_tokens"] == 0: + row["kv_seed_regime"] = "not_applicable" + point["sample_reasons"] = [] + for group in payload["iteration_groups"]: + point = group["point"] + if point["point_type"] == "prefill" and point["total_kv_read_tokens"] == 0: + point["sample_reasons"] = [] + _write(path, payload) + + +@pytest.fixture +def quality_case(validation_case, tmp_path, monkeypatch, materialized_workload, request): # noqa: F811 + selection = getattr(request, "param", False) + instrumented = selection if isinstance(selection, bool) else False + worker_type = selection if isinstance(selection, str) else None + replay_args, request, root, trace, replay_output = validation_case + payload = request.model_dump(mode="json") + payload["identity"]["framework_version"] = "0.27.0" + payload["fpm_profile"]["deployments"][0].update(backend_version="0.27.0", fmha_quant_mode="fp8") + payload["collection"]["prefill_cudagraph_policy"] = "runtime" + if worker_type is not None: + payload["worker_type"] = worker_type + payload["fpm_profile"]["deployments"][0]["worker_type"] = worker_type + request = SupportRequest.model_validate(payload) + old_root = root + root = tmp_path / "native-collection" + replay_args = [value.replace(str(old_root), str(root)) for value in replay_args] + create_plan(request, root) + model_config = tmp_path / "model-config.json" + _write( + model_config, + { + "architectures": [request.fpm_profile.architecture], + "model_type": "llama", + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 1024, + "max_position_embeddings": 512, + "torch_dtype": "bfloat16", + "num_local_experts": request.fpm_profile.num_experts, + "quantization_config": {"quant_method": "fp8", "activation_scheme": "dynamic"}, + }, + ) + parsed = collector_cli._parser().parse_args(fpm_cli_args(request, output_dir=root, plan_only=True)[3:]) + options = FPMCollectionOptions.from_args(parsed) + points = { + "schema_version": 3, + "prefill": [ + {"batch_size": batch, "total_prefill_tokens": tokens, "total_kv_read_tokens": batch * context} + for batch in (1, 2) + for tokens in (2, 16, 32, 64, 128) + for context in (0, 16, 64, 128, 256) + ], + "decode": [ + {"batch_size": batch, "total_kv_read_tokens": batch * context} + for batch in (1, 2) + for context in (1, 4, 16, 64, 128, 256) + ], + } + points_json = json.dumps(points, sort_keys=True, separators=(",", ":")) + options = replace( + options, benchmark_points_json=points_json, benchmark_points_sha256=planner._canonical_hash(points) + ) + overrides = {"K8sConfig": {"k8s_image": "example/runtime@sha256:" + "a" * 64}} + runtime = {} + if instrumented: + from collector.fpm_forward.bundled_instrumentation import bundled_instrumentation + + deployment = request.profile_deployment() + runtime = { + "runtime_instrumentation": bundled_instrumentation("0.27.0"), + "runtime_configuration": "worker", + "runtime_launch": { + "identity": request.identity.model_dump(mode="json"), + "topology": {name: getattr(deployment, name) for name in ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp")}, + "precision": { + **{ + name: getattr(deployment, name) + for name in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "comm_quant_mode", + "moe_backend", + "attention_backend", + "enable_wideep", + "enable_eplb", + ) + }, + "kvcache_quant_mode": deployment.kv_cache_dtype, + }, + "collection": { + "max_model_len": options.vllm_max_model_len, + "max_num_batched_tokens": options.max_num_batched_tokens, + "max_num_seqs": options.max_num_seqs, + "gpu_memory_utilization": request.collection.memory_fraction, + "prefill_cudagraph_policy": "runtime", + "max_prefill_cudagraph_size": None, + }, + "model_config": { + "path": str(model_config), + "sha256": hashlib.sha256(model_config.read_bytes()).hexdigest(), + }, + "deployment": {"executor": "kubernetes", "image": overrides["K8sConfig"]["k8s_image"]}, + }, + } + monkeypatch.setattr(planner, "_git_revision", lambda: "synthetic-workflow-revision") + plan = planner.build_collection_plan( + backend="vllm", + model_path=request.identity.model, + system=request.identity.gpu, + selected_ops=set(), + options=options, + model_architecture=request.fpm_profile.architecture, + model_config_path=str(model_config), + fpm_profile=request.fpm_profile, + generator_overrides=overrides, + **runtime, + ) + campaign = root / "fpm-artifacts" / plan.sha256[:16] + checkpoint = root / "fpm-checkpoint/fpm_forward.json" + _synthetic_campaign(campaign, checkpoint, plan, attempt_id="synthetic-source", generator_overrides=overrides) + rows = [] + for cell in plan.cells: + rows.extend( + aggregate_cell(plan, cell, campaign / "cells" / cell.cell_id, expected_attempt_id="synthetic-source") + ) + # Replace the other fixture's explicit synthetic table with this native publication. + for path in (root / "systems/data").rglob("fpm_forward_perf.*"): + path.unlink() + parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=root / "systems/data") + saved = json.loads(checkpoint.read_text()) + saved["database"] = { + "status": "passed", + "missing_cells": [], + "skipped_first_publisher_wins": list(skipped), + "plan_cells": len(plan.cells), + "published_cells": len(plan.cells), + "parquet": str(parquet), + "metadata": str(metadata), + } + _write(checkpoint, saved) + calls = [] + + def synthetic_repeat(sample_plan, **kwargs): + calls.append(kwargs) + raw = Path(kwargs["artifact_root"]) / sample_plan.sha256[:16] + cp = Path(kwargs["checkpoint_dir"]) / "fpm_forward.json" + _synthetic_campaign( + raw, + cp, + sample_plan, + attempt_id=f"synthetic-repeat-{len(calls)}", + generator_overrides=kwargs["generator_overrides"], + ) + return 0 + + monkeypatch.setattr(repeatability, "run_collection", synthetic_repeat) + output = tmp_path / "collection-quality" + args = [ + "onboard", + "validate-collection", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--validation-output-dir", + str(output), + ] + return { + "request": request, + "root": root, + "plan": plan, + "campaign": campaign, + "checkpoint": checkpoint, + "args": args, + "output": output, + "calls": calls, + "replay_args": replay_args, + "replay_output": replay_output, + "trace": trace, + } + + +@pytest.mark.parametrize("quality_case", [False, True], indirect=True, ids=["ordinary", "instrumented"]) +def test_public_prepare_then_repeat_and_holdout_preserve_originals(quality_case): + case = quality_case + originals = {p: p.read_bytes() for p in case["root"].rglob("*") if p.is_file()} + assert cli.main(case["args"]) == 0 + report = json.loads((case["output"] / "collection-validation.json").read_text()) + assert report["status"] == "incomplete" + assert report["gates"]["interpolation"]["status"] == "passed" + assert not case["calls"] + assert cli.main([*case["args"], "--resume", "--execute"]) == 0 + checked, _, _, _ = workflow.check_collection_report( + case["output"] / "collection-validation.json", workflow.ValidationPolicy() + ) + assert checked["status"] == "passed" + assert len(case["calls"]) == 10 + assert originals == {p: p.read_bytes() for p in originals} + if case["plan"].runtime_instrumentation is not None: + from collector.fpm_forward.runtime_memory import validate_saved_plan + + for path in case["output"].rglob("collection-plan.json"): + saved = json.loads(path.read_text()) + validate_saved_plan(saved) + assert saved["runtime_observation"] == case["plan"].to_dict()["runtime_observation"] + + +@pytest.mark.parametrize("quality_case", ["prefill", "decode"], indirect=True) +def test_public_role_quality_prepares_and_qualifies_without_the_other_phase(quality_case): + case = quality_case + role = case["request"].worker_type + assert {cell.workload_kind for cell in case["plan"].cells} == {role} + assert cli.main(case["args"]) == 0 + holdout = json.loads((case["output"] / "holdout/interpolation-validation.json").read_text()) + assert set(holdout["phases"]) == {role} + assert holdout["status"] == "passed" + assert cli.main([*case["args"], "--resume", "--execute"]) == 0 + report, _, _, _ = workflow.check_collection_report(case["output"] / "collection-validation.json") + assert report["status"] == "passed" + assert len(case["calls"]) == 5 + + +def test_holdout_uses_full_native_grid_not_repeat_subset(quality_case, tmp_path): + case = quality_case + policy = tmp_path / "bounded-policy.json" + _write(policy, {"repeatability": {"comparison_mode": "bounded"}}) + before = {path: path.read_bytes() for path in case["root"].rglob("*") if path.is_file()} + assert cli.main([*case["args"], "--policy", str(policy)]) == 0 + selection = json.loads((case["output"] / "repeatability-selection.json").read_text()) + assert all(len(cell["points"]) <= 12 for cell in selection["cells"]) + holdout = json.loads((case["output"] / "holdout/interpolation-validation.json").read_text()) + boundaries = holdout["capture_boundaries"] + assert boundaries["status"] == "assessed" + assert len(boundaries["retained_anchors"]) == 20 # Both sides, for every batch/KV curve. + assert {point["total_prefill_tokens"] for point in boundaries["retained_anchors"]} == {16, 32} + assert holdout["native_query_coverage"]["queries"]["measured"] == 0 + assert all(cell["known_mode_point_count"] == cell["eligible_point_count"] for cell in boundaries["cells"]) + assert not case["calls"] + assert before == {path: path.read_bytes() for path in before} + + +@pytest.mark.parametrize("missing", ["point_mode", "graph_config"]) +def test_native_missing_boundary_evidence_keeps_assessment_incomplete(quality_case, missing): + case = quality_case + for path in case["campaign"].rglob("benchmark-dp*.json"): + payload = json.loads(path.read_text()) + if missing == "graph_config": + payload.pop("cudagraph") + else: + for row in payload["results"]: + row["point"].pop("expected_cudagraph_mode") + for row in payload["iteration_groups"]: + row["point"].pop("expected_cudagraph_mode") + _write(path, payload) + assert cli.main(case["args"]) == 0 + holdout = json.loads((case["output"] / "holdout/interpolation-validation.json").read_text()) + assert holdout["status"] == holdout["capture_boundaries"]["status"] == "incomplete" + assert holdout["capture_boundaries"]["retained_anchors"] == [] + checked, _, _, _ = workflow.check_collection_report( + case["output"] / "collection-validation.json", workflow.ValidationPolicy() + ) + assert checked["status"] == "incomplete" + + +@pytest.mark.parametrize("validation_case", ["dep"], indirect=True) +@pytest.mark.parametrize("conflict", ["point_mode", "graph_config"]) +def test_native_rank_boundary_conflicts_are_rejected(quality_case, conflict): + case = quality_case + path = next(case["campaign"].rglob("benchmark-dp1.json")) + payload = json.loads(path.read_text()) + if conflict == "graph_config": + payload["cudagraph"]["capture_sizes"].append(512) + else: + payload["results"][0]["point"]["expected_cudagraph_mode"] = "NONE" + payload["iteration_groups"][0]["point"]["expected_cudagraph_mode"] = "NONE" + _write(path, payload) + with pytest.raises(SystemExit, match="2"): + cli.main(case["args"]) + assert not case["calls"] + assert not (case["output"] / "holdout/interpolation-validation.json").exists() + + +@pytest.mark.parametrize("changed", ["legacy_selection", "anchors"]) +def test_changed_holdout_selection_requires_fresh_output_and_preserves_prior_report(quality_case, changed): + case = quality_case + assert cli.main(case["args"]) == 0 + holdout_path = case["output"] / "holdout/interpolation-validation.json" + holdout = json.loads(holdout_path.read_text()) + if changed == "legacy_selection": + holdout["policy"]["selection"] = "retained_envelope_seeded_farthest_shape_v1" + else: + holdout["capture_boundaries"]["retained_anchors"] = [] + _write(holdout_path, holdout) + report_path = case["output"] / "collection-validation.json" + report = json.loads(report_path.read_text()) + report["gates"]["interpolation"]["report"] = workflow._identity(holdout_path) + _write(report_path, report) + before = {path: path.read_bytes() for path in case["output"].rglob("*") if path.is_file()} + with pytest.raises(ValueError, match="fresh assessment directory"): + workflow.check_collection_report(report_path, workflow.ValidationPolicy()) + with pytest.raises(SystemExit, match="2"): + cli.main([*case["args"], "--resume"]) + assert before == {path: path.read_bytes() for path in before} + assert not case["calls"] + + +def test_changed_threshold_reuses_raw_samples_but_changed_count_does_not(quality_case, tmp_path): + case = quality_case + assert cli.main([*case["args"], "--execute"]) == 0 + repeats = case["output"] / "repeatability" + saved = {p: p.read_bytes() for p in repeats.rglob("*") if p.is_file()} + policy = tmp_path / "policy.json" + _write(policy, {"repeatability": {"max_cv": 0.10}, "interpolation": {"max_p95_relative_error": 0.1}}) + args = [ + *case["args"][:-1], + str(tmp_path / "reassessed"), + "--repeatability-dir", + str(repeats), + "--policy", + str(policy), + ] + assert cli.main(args) == 0 + assert len(case["calls"]) == 10 + assert saved == {p: p.read_bytes() for p in saved} + _write(policy, {"repeatability": {"samples": 6}}) + args[args.index("--validation-output-dir") + 1] = str(tmp_path / "wrong-count") + with pytest.raises(SystemExit, match="2"): + cli.main(args) + assert len(case["calls"]) == 10 + + +@pytest.mark.parametrize("condition", ["missing", "failed", "failed_and_missing"]) +def test_source_execution_status_survives_report_recheck(quality_case, condition): + case = quality_case + paths = sorted(case["campaign"].rglob("fpm-execution-worker-*.json")) + if condition in {"missing", "failed_and_missing"}: + paths.pop().unlink() + if condition in {"failed", "failed_and_missing"}: + observed = json.loads(paths[0].read_text()) + observed["resolved_config"]["model_config"]["dtype"] = "torch.float32" + _write(paths[0], observed) + expected = "incomplete" if condition == "missing" else "failed" + assert cli.main(case["args"]) == (0 if expected == "incomplete" else 1) + report_path = case["output"] / "collection-validation.json" + report, _, _, _ = workflow.check_collection_report(report_path, workflow.ValidationPolicy()) + assert report["status"] == report["gates"]["execution"]["status"] == expected + cells = report["gates"]["execution"]["cells"] + assert any(cell["missing_evidence" if expected == "incomplete" else "failures"] for cell in cells) + report["gates"]["execution"]["status"] = "passed" + _write(report_path, report) + with pytest.raises(ValueError, match="execution gate differs"): + workflow.check_collection_report(report_path, workflow.ValidationPolicy()) + + +@pytest.mark.parametrize("relation", ["ancestor", "same", "descendant", "symlink_descendant"]) +def test_collection_reassessment_preserves_reused_attempt_tree(quality_case, tmp_path, relation): + case = quality_case + assert cli.main([*case["args"], "--execute"]) == 0 + repeats = case["output"] / "repeatability" + attempt = next((repeats / "samples").glob("*/sample-*/attempt-*")) + output = {"ancestor": repeats.parent, "same": repeats, "descendant": attempt / "assessment"}.get(relation) + if relation == "symlink_descendant": + link = tmp_path / "repeat-alias" + link.symlink_to(attempt, target_is_directory=True) + output = link / "assessment" + before = {p: p.read_bytes() for p in case["output"].rglob("*") if p.is_file()} + args = [*case["args"]] + args[args.index("--validation-output-dir") + 1] = str(output) + with pytest.raises(SystemExit, match="2"): + cli.main([*args, "--repeatability-dir", str(repeats)]) + assert before == {p: p.read_bytes() for p in case["output"].rglob("*") if p.is_file()} + checked, _, _, _ = workflow.check_collection_report( + case["output"] / "collection-validation.json", workflow.ValidationPolicy() + ) + assert checked["status"] == "passed" + + +@pytest.mark.parametrize("which", ["status", "formal", "policy", "sample"]) +def test_requalification_rejects_changed_or_forged_evidence(quality_case, which): + case = quality_case + assert cli.main([*case["args"], "--execute"]) == 0 + report_path = case["output"] / "collection-validation.json" + report = json.loads(report_path.read_text()) + if which == "status": + report["gates"]["interpolation"]["status"] = "failed" + _write(report_path, report) + elif which == "policy": + _write(case["output"] / "policy.json", {"repeatability": {"max_cv": 1.0}}) + elif which == "formal": + Path(report["inputs"]["formal_data"][0]["path"]).write_bytes(b"changed") + else: + next((case["output"] / "repeatability").rglob("benchmark-dp0.json")).write_text("{}") + with pytest.raises((ValueError, OSError)): + workflow.check_collection_report(report_path, workflow.ValidationPolicy()) + + +@pytest.mark.parametrize("policy_format", ["json", "yaml"]) +@pytest.mark.parametrize( + "values", + [ + {"repeatability": {"samples": True}}, + {"repeatability": {"comparison_mode": "automatic"}}, + {"repeatability": {"max_attempts_per_sample": 0}}, + {"repeatability": {"max_attempts_per_sample": True}}, + {"repeatability": {"max_source_relative_difference": float("inf")}}, + {"serving": {"max_latency_p95_relative_error": float("nan")}}, + {"serving": {"max_latency_p95_relative_error": -0.1}}, + {"serving": {"max_latency_p95_relative_error": "1e-05"}}, + {"serving": {"max_throughput_relative_error": True}}, + {"serving": {"max_throughput_relative_error": float("inf")}}, + {"serving": {"max_throughput_relative_error": float("-inf")}}, + {"interpolation": {"max_unsupported_fraction": 1.1}}, + {"typo": 1}, + [], + ], +) +def test_validation_policy_rejects_ambiguous_settings(values, tmp_path, policy_format): + policy = tmp_path / f"policy.{policy_format}" + policy.write_text(json.dumps(values) if policy_format == "json" else yaml.safe_dump(values)) + with pytest.raises(ValueError): + workflow._policy(policy) + + +def test_repeatability_policy_defaults_and_legacy_semantics(tmp_path): + policy = workflow.ValidationPolicy() + assert policy.schema_version == "aisimulate-onboarding-validation-policy/v2" + assert policy.repeatability.comparison_mode == "full_grid" + assert policy.repeatability.max_attempts_per_sample == 2 + assert policy.repeatability.max_cv == policy.repeatability.max_source_relative_difference == 0.05 + path = tmp_path / "legacy-policy.json" + legacy = {"schema_version": "aisimulate-onboarding-validation-policy/v1", "repeatability": {"samples": 3}} + _write(path, legacy) + parsed = workflow._policy(path) + assert parsed.schema_version == legacy["schema_version"] + assert parsed.repeatability.comparison_mode == "bounded" + assert parsed.repeatability.samples == 3 + assert json.loads(path.read_text()) == legacy + with pytest.raises(ValueError, match="legacy v1 policy uses bounded"): + workflow.ValidationPolicy.model_validate({**legacy, "repeatability": {"comparison_mode": "full_grid"}}) + + +@pytest.mark.parametrize("reason", ["legacy", "bounded", "source_missing", "source_failed", "population_unstable"]) +def test_stable_fresh_measurements_do_not_qualify_other_populations(reason): + policy = workflow.ValidationPolicy() + assessment = { + "schema_version": 2, + "policy": {"comparison_mode": "full_grid"}, + "status": "passed", + "source_qualification": {"status": "qualified"}, + "new_population": {"status": "qualified"}, + } + assert workflow._repeatability_gate_status(assessment, policy) == "passed" + if reason == "legacy": + assessment["schema_version"] = 1 + elif reason == "bounded": + policy = workflow.ValidationPolicy(repeatability={"comparison_mode": "bounded"}) + assessment["policy"]["comparison_mode"] = "bounded" + elif reason == "source_missing": + assessment.pop("source_qualification") + elif reason == "source_failed": + assessment["source_qualification"]["status"] = "failed" + else: + assessment["new_population"]["status"] = "unstable" + expected = "failed" if reason in {"source_failed", "population_unstable"} else "incomplete" + assert workflow._repeatability_gate_status(assessment, policy) == expected + assessment["status"] = "failed" + assert workflow._repeatability_gate_status(assessment, policy) == "failed" + + +@pytest.mark.parametrize( + "change", [{"comparison_mode": "bounded"}, {"max_attempts_per_sample": 3}, {"max_source_relative_difference": 0.1}] +) +def test_repeatability_policy_resume_rejects_drift_without_replacing_evidence(quality_case, tmp_path, change): + case = quality_case + policy = tmp_path / "repeat-policy.json" + _write(policy, workflow.ValidationPolicy().model_dump(mode="json")) + args = [*case["args"], "--policy", str(policy)] + assert cli.main(args) == 0 + preserved = {path: path.read_bytes() for path in case["output"].rglob("*") if path.is_file()} + _write(policy, {"repeatability": change}) + with pytest.raises(SystemExit, match="2"): + cli.main([*args, "--resume", "--execute"]) + assert not case["calls"] + assert preserved == {path: path.read_bytes() for path in preserved} + + +def test_collector_cli_passes_frozen_repeatability_settings(quality_case, capsys): + case = quality_case + assert ( + collector_cli.main( + [ + "--repeatability-source-campaign", + str(case["campaign"]), + "--repeatability-source-checkpoint", + str(case["checkpoint"]), + "--repeatability-output-dir", + str(case["output"]), + "--repeatability-comparison-mode", + "bounded", + "--repeatability-max-attempts", + "3", + "--repeatability-source-agreement-threshold", + "0.04", + "--plan-only", + ] + ) + == 0 + ) + frozen = json.loads(capsys.readouterr().out) + assert frozen["policy"]["comparison_mode"] == "bounded" + assert frozen["policy"]["max_attempts_per_sample"] == 3 + assert frozen["policy"]["source_agreement_threshold"] == 0.04 + assert not case["calls"] + + +def test_repeatability_settings_survive_checkpoint_resume_without_execution(quality_case, tmp_path, capsys): + from aisimulate.support.checkpoint import save_checkpoint + + case = quality_case + policy = workflow.ValidationPolicy().model_dump(mode="json") + checkpoint = tmp_path / "onboarding-checkpoint.json" + state, _ = save_checkpoint( + checkpoint, + patch={ + "configurations": { + "worker": { + "draft_request": case["request"].model_dump(mode="json"), + "validation_inputs": {"policy": policy}, + } + } + }, + expected_revision=None, + accept=["worker"], + ) + assert state.configurations["worker"].acceptance is not None + assert cli.main(["onboard", "resume", "--checkpoint", str(checkpoint)]) == 0 + resumed = json.loads(capsys.readouterr().out) + assert resumed["state"]["configurations"]["worker"]["validation_inputs"]["policy"] == policy + assert resumed["configurations"]["worker"]["profile_accepted"] + assert not case["calls"] + + +def _prepare_matched(case, tmp_path, *, policy=None): + assert cli.main([*case["args"], "--execute"]) == 0 + assert cli.main(case["replay_args"]) == 0 + tokenizer = tmp_path / "tokenizer" + tokenizer.mkdir() + _write(tokenizer / "tokenizer_config.json", {"synthetic": True}) + prepared = tmp_path / "prepared" + args = [ + "onboard", + "validate-serving", + "--collection-report", + str(case["output"] / "collection-validation.json"), + "--replay-report", + str(case["replay_output"] / "validation.json"), + ] + if policy is not None: + args.extend(["--policy", str(policy)]) + assert ( + cli.main( + [ + *args, + "--action", + "prepare", + "--validation-output-dir", + str(prepared), + "--endpoint", + "http://127.0.0.1:8000", + "--tokenizer", + str(tokenizer), + "--aiperf-python", + "/synthetic/pinned-aiperf/python", + ] + ) + == 0 + ) + recipe_path = prepared / "serving/recipe.json" + recipe = json.loads(recipe_path.read_text()) + return args, recipe_path, recipe, tokenizer + + +def _synthetic_serving(recipe_path, recipe, tokenizer): + """AIPerf-shaped fixtures generated from native prediction; no HTTP/GPU execution.""" + prediction_ref = recipe.get("matched_workload", {}).get("prediction_report", recipe["prediction_report"]) + prediction = json.loads(Path(prediction_ref["path"]).read_text()) + artifacts = Path(recipe["artifacts_dir"]) + artifacts.mkdir() + rows = [] + for index, row in enumerate(prediction["per_request"]): + rows.append( + { + "metadata": { + "conversation_id": recipe["play_id"], + "turn_index": index, + "x_request_id": f"synthetic-{index}", + "was_cancelled": False, + "benchmark_phase": "profiling", + "request_start_ns": 1_000_000_000 + round(row["arrival_time_ms"] * 1_000_000), + "request_end_ns": 1_000_000_000 + round(row["last_token_ms"] * 1_000_000), + }, + "metrics": { + "input_sequence_length": {"value": row["input_length"], "unit": "tokens"}, + "output_sequence_length": {"value": row["output_length"], "unit": "tokens"}, + "time_to_first_token": {"value": row["ttft_ms"], "unit": "ms"}, + "inter_token_latency": {"value": row["itl_ms"], "unit": "ms"}, + }, + "error": None, + } + ) + records = artifacts / "profile_export.jsonl" + records.write_text("\n".join(json.dumps(row) for row in rows) + "\n") + summary = artifacts / "profile_export_aiperf.json" + _write(summary, {"was_cancelled": False, "is_complete": True, "request_count": {"avg": len(rows)}}) + runtime_dir = Path(recipe["serving_observer"]["output_dir"]) + runtime_dir.mkdir() + provenance = Path(recipe["serving_observer"]["provenance_path"]) + expected = recipe["expected_execution"] + native_runtime = json.loads(json.dumps(expected["runtime_config"])) + native_runtime.setdefault("cache_config", {})["enable_prefix_caching"] = recipe["scope"]["prefix_caching"] + runtime_sources = [] + for dp in range(expected["parallelism"]["attention_data"]): + for tp in range(expected["parallelism"]["tensor"]): + path = runtime_dir / f"fpm-execution-worker-dp{dp}-tp{tp}-pp0.json" + _write( + path, + { + "schema_name": "aisimulate_fpm_runtime_execution", + "schema_version": 1, + "status": "observed", + "execution_provenance": json.loads(provenance.read_text()), + "provenance_source": workflow._identity(provenance), + "dp_rank": dp, + "tp_rank": tp, + "pp_rank": 0, + "resolved_config": native_runtime, + **{key: expected[key] for key in ("attention_groups", "graph_config", "backend_version")}, + }, + ) + runtime_sources.append(workflow._identity(path)) + _write( + recipe_path.parent / "benchmark-execution.json", + { + "schema_version": serving.SCHEMA, + "recipe": workflow._identity(recipe_path), + "status": "completed", + "exit_code": 0, + "command": ["python", *recipe["command_arguments"]], + "installation": {"source": {"vcs_info": {"commit_id": serving.AIPERF_REVISION}}}, + "artifacts": [workflow._identity(records), workflow._identity(summary)], + "runtime_observations": runtime_sources, + "serving_provenance": workflow._identity(provenance), + "observer_modules": recipe["serving_observer"]["module_files"], + "server_tokenization": workflow._identity(write_server_tokenization(recipe_path, recipe)), + }, + ) + observed = recipe_path.parent / "observed-execution.json" + _write( + observed, + { + "schema_version": serving.SCHEMA, + "recipe": workflow._identity(recipe_path), + "observed_execution": expected, + "scope": recipe["scope"], + "sources": [workflow._identity(tokenizer / "tokenizer_config.json")], + }, + ) + return observed + + +@pytest.mark.parametrize("validation_case", ["tp", "dep", "tep"], indirect=True) +def test_public_four_steps_qualify_only_matched_synthetic_scope(quality_case, tmp_path): + args, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path) + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + raw = {p: p.read_bytes() for p in recipe_path.parent.rglob("*") if p.is_file()} + output = tmp_path / "qualified" + assert ( + cli.main( + [ + *args, + "--action", + "assess", + "--recipe", + str(recipe_path), + "--execution-evidence", + str(observed), + "--validation-output-dir", + str(output), + ] + ) + == 0 + ) + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == "passed" + assert report["accuracy"] == "qualified_for_evaluated_scope" + assert report["gates"]["serving"]["forward"]["status"] == "unavailable" + assert set(report["gates"]["serving"]["metrics"]) == {"ttft", "tpot", "output_throughput"} + assert raw == {p: p.read_bytes() for p in raw} + + +@pytest.mark.parametrize("exit_code", [0, 7]) +def test_serving_run_returns_producer_status_and_preserves_logs(quality_case, tmp_path, monkeypatch, capsys, exit_code): + _, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path) + _synthetic_serving(recipe_path, recipe, tokenizer) + receipt_path = recipe_path.parent / "benchmark-execution.json" + receipt_path.unlink() + shutil.rmtree(recipe["artifacts_dir"]) + producer = tmp_path / "synthetic-aiperf-python" + installation = { + "version": "synthetic-process-boundary", + "source": {"url": serving.AIPERF_SOURCE, "vcs_info": {"commit_id": serving.AIPERF_REVISION}}, + } + producer.write_text( + "#!/usr/bin/env python3\nimport sys\nif '-c' in sys.argv:\n" + f" print({json.dumps(installation)!r})\nelse:\n" + f" print('producer stdout')\n print('producer stderr', file=sys.stderr)\n sys.exit({exit_code})\n" + ) + producer.chmod(0o755) + monkeypatch.setattr(serving, "_tokenize_serving", write_server_tokenization) + result = cli.main( + [ + "onboard", + "validate-serving", + "--action", + "run", + "--recipe", + str(recipe_path), + "--aiperf-python", + str(producer), + ] + ) + receipt = json.loads(receipt_path.read_text()) + assert receipt["exit_code"] == exit_code + assert receipt["status"] == ("completed" if exit_code == 0 else "failed") + assert result == (0 if exit_code == 0 else 1) + stdout = capsys.readouterr().out + assert str(receipt_path) in stdout + for name, content in (("stdout", "producer stdout"), ("stderr", "producer stderr")): + path = recipe_path.parent / f"aiperf.{name}.log" + assert path.read_text().strip() == content + assert str(path) in stdout + + +@pytest.mark.parametrize("missing", ["receipt", "execution"]) +def test_missing_serving_evidence_retains_incomplete_diagnostic(quality_case, tmp_path, missing): + args, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path) + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + missing_path = recipe_path.parent / "benchmark-execution.json" if missing == "receipt" else observed + missing_path.unlink() + output = tmp_path / "missing-assessment" + assert ( + cli.main( + [ + *args, + "--action", + "assess", + "--recipe", + str(recipe_path), + "--execution-evidence", + str(observed), + "--validation-output-dir", + str(output), + ] + ) + == 1 + ) + combined = json.loads((output / "validation.json").read_text()) + assessed = json.loads((output / "serving-validation.json").read_text()) + gate = combined["gates"]["serving"] + assert combined["status"] == gate["status"] == assessed["status"] == "incomplete" + assert gate["missing_evidence"] == assessed["missing_evidence"] + assert any(str(missing_path) in issue for issue in gate["missing_evidence"]) + assert gate["failures"] == [] + + +@pytest.mark.parametrize("action", ["prepare", "assess"]) +@pytest.mark.parametrize("relation", ["ancestor", "same", "descendant"]) +@pytest.mark.parametrize( + "source_kind", ["original_raw", "reused_repeat", "collection_report", "replay_report", "recipe", "tokenizer"] +) +def test_serving_rejects_nested_outputs_before_writing(quality_case, tmp_path, action, source_kind, relation): + case = quality_case + args, recipe_path, recipe, tokenizer = _prepare_matched(case, tmp_path) + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + repeats = case["output"] / "repeatability" + reassessed = tmp_path / "reassessed-collection" + collection_args = [*case["args"]] + collection_args[collection_args.index("--validation-output-dir") + 1] = str(reassessed) + assert cli.main([*collection_args, "--repeatability-dir", str(repeats)]) == 0 + report_path = reassessed / "collection-validation.json" + args[args.index("--collection-report") + 1] = str(report_path) + source = { + "original_raw": next((case["campaign"] / "cells").glob("*/raw")), + "reused_repeat": next((repeats / "samples").glob("*/sample-*/attempt-*")), + "collection_report": report_path.parent, + "replay_report": case["replay_output"], + "recipe": recipe_path.parent, + "tokenizer": tokenizer, + }[source_kind] + output = {"ancestor": source.parent, "same": source, "descendant": source / "must-not-exist"}[relation] + before = {p: p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + with pytest.raises(SystemExit, match="2"): + cli.main( + [ + *args, + "--action", + action, + "--recipe", + str(recipe_path), + "--execution-evidence", + str(observed), + "--validation-output-dir", + str(output), + "--endpoint", + "http://localhost:8000", + "--tokenizer", + str(tokenizer), + "--aiperf-python", + "/not/executed", + ] + ) + if relation == "descendant": + assert not output.exists() + assert before == {p: p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + checked, original, plan, _ = workflow.check_collection_report(report_path, workflow.ValidationPolicy()) + assert checked["status"] == "passed" + assert ( + workflow._check_replay(case["replay_output"] / "validation.json", original, checked, plan)["status"] == "passed" + ) + assert serving._checked_recipe(recipe_path)["status"] == "ready" + + +@pytest.mark.parametrize( + ("changed", "expected_status"), + [ + ("missing_collection", "incomplete"), + ("prediction", "failed"), + ("coverage", "failed"), + ("legacy_replay", "incomplete"), + ("missing_serving", "incomplete"), + ], +) +def test_combined_gate_rejects_missing_and_stale_evidence(quality_case, tmp_path, changed, expected_status): + args, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path) + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + if changed == "missing_collection": + Path(args[args.index("--collection-report") + 1]).unlink() + elif changed in {"prediction", "coverage"}: + path = ( + quality_case["replay_output"] + / "prediction" + / ("prediction.json" if changed == "prediction" else "fpm-coverage.json") + ) + path.write_text(path.read_text() + "\n") + elif changed == "legacy_replay": + path = quality_case["replay_output"] / "validation.json" + payload = json.loads(path.read_text()) + del payload["prediction_evidence"] + _write(path, payload) + else: + (recipe_path.parent / "benchmark-execution.json").unlink() + output = tmp_path / "rejected" + assert ( + cli.main( + [ + *args, + "--action", + "assess", + "--recipe", + str(recipe_path), + "--execution-evidence", + str(observed), + "--validation-output-dir", + str(output), + ] + ) + == 1 + ) + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == expected_status + assert report["accuracy"] == "not_qualified" + if changed != "missing_collection": + assert report["gates"]["collection"]["status"] == "passed" + + +@pytest.mark.parametrize("policy_format", ["json", "yaml"]) +@pytest.mark.parametrize("limit", [0.0, 0.00001, 1.1]) +def test_serving_error_limits_survive_preparation_assessment_and_recheck(quality_case, tmp_path, limit, policy_format): + policy = tmp_path / f"policy.{policy_format}" + values = {"serving": {"max_latency_p95_relative_error": limit, "max_throughput_relative_error": limit}} + policy.write_text(json.dumps(values) if policy_format == "json" else yaml.safe_dump(values)) + args, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path, policy=policy) + assert recipe["policy"]["latency_p95_relative_error"] == limit + assert recipe["policy"]["throughput_relative_error"] == limit + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + raw = {path: path.read_bytes() for path in recipe_path.parent.rglob("*") if path.is_file()} + output = tmp_path / "assessed" + assert ( + cli.main( + [ + *args, + "--action", + "assess", + "--recipe", + str(recipe_path), + "--execution-evidence", + str(observed), + "--validation-output-dir", + str(output), + ] + ) + == 0 + ) + report = serving.check_serving_validation_report(output / "serving-validation.json") + assert report["status"] == "passed" + assert report["policy"] == recipe["policy"] + assert report["metrics"]["ttft"]["p95_relative_error"] == 0 + assert report["metrics"]["tpot"]["p95_relative_error"] == 0 + assert report["metrics"]["output_throughput"]["relative_error"] == 0 + assert raw == {path: path.read_bytes() for path in raw} + + +@pytest.mark.parametrize( + "duration_scale,strict_limit,relaxed_limit,failed_metric", + [ + (1.5, 0.2, 0.5, "ttft"), + (1.000001, 0.0, 0.00001, "ttft"), + (0.48, 1.0, 1.1, "ttft"), + # Two 30 ms requests plus 100 ms idle give 160 ms baseline; at + # 3.8x request duration, (2 * 30 * 3.8 + 100) / 160 - 1 = 1.05. + (3.8, 1.0, 1.1, "output_throughput"), + ], +) +def test_serving_threshold_reassessment_preserves_measurements( + quality_case, tmp_path, duration_scale, strict_limit, relaxed_limit, failed_metric +): + args, recipe_path, recipe, tokenizer = _prepare_matched(quality_case, tmp_path) + observed = _synthetic_serving(recipe_path, recipe, tokenizer) + records = Path(recipe["artifacts_dir"]) / "profile_export.jsonl" + rows = [json.loads(line) for line in records.read_text().splitlines()] + previous = None + for row in rows: + metadata = row["metadata"] + duration = metadata["request_end_ns"] - metadata["request_start_ns"] + if previous is not None: + metadata["request_start_ns"] = previous + 100_000_000 + metadata["request_end_ns"] = metadata["request_start_ns"] + round(duration * duration_scale) + previous = metadata["request_end_ns"] + for key in ("time_to_first_token", "inter_token_latency"): + row["metrics"][key]["value"] *= duration_scale + records.write_text("\n".join(json.dumps(row) for row in rows) + "\n") + receipt_path = recipe_path.parent / "benchmark-execution.json" + receipt = json.loads(receipt_path.read_text()) + receipt["artifacts"] = [workflow._identity(Path(ref["path"])) for ref in receipt["artifacts"]] + _write(receipt_path, receipt) + raw = {path: path.read_bytes() for path in recipe_path.parent.rglob("*") if path.is_file()} + assess = [*args, "--action", "assess", "--recipe", str(recipe_path), "--execution-evidence", str(observed)] + strict_policy = tmp_path / "strict.json" + _write( + strict_policy, + {"serving": {"max_latency_p95_relative_error": strict_limit, "max_throughput_relative_error": strict_limit}}, + ) + assert cli.main([*assess, "--validation-output-dir", str(tmp_path / "strict"), "--policy", str(strict_policy)]) == 1 + report = json.loads((tmp_path / "strict/validation.json").read_text()) + assert report["gates"]["serving"]["metrics"][failed_metric]["status"] == "failed" + assert serving.check_serving_validation_report(tmp_path / "strict/serving-validation.json")["status"] == "failed" + policy = tmp_path / "relaxed.json" + _write( + policy, + {"serving": {"max_latency_p95_relative_error": relaxed_limit, "max_throughput_relative_error": relaxed_limit}}, + ) + assert cli.main([*assess, "--validation-output-dir", str(tmp_path / "relaxed"), "--policy", str(policy)]) == 0 + assert serving.check_serving_validation_report(tmp_path / "relaxed/serving-validation.json")["status"] == "passed" + assert raw == {path: path.read_bytes() for path in raw} + + +@pytest.mark.parametrize("legacy_provenance", [False, True]) +def test_replay_accepts_verified_memory_finalization_and_rejects_changed_capacity( + tmp_path, monkeypatch, legacy_provenance +): + from collector.fpm_forward.repeatability import load_repeatability_source + + from aisimulate import supervision + from aisimulate.support.finalization import _merge_resources, _verify_collection, finalize + + from .test_onboard_finalization import build_completed_collection + + original, root = build_completed_collection(tmp_path) + target = tmp_path / "resolved" + finalize(original, root, target) + request = SupportRequest.from_yaml(target / "request.yaml") + if legacy_provenance: + # Reproduce the previous embedded observation format, keeping its exact + # profile/request identity through subsequent replay reassessment. + observations, manifest, formal_data, _ = _verify_collection(original, root) + payload = request.model_dump(mode="json") + payload["fpm_profile"]["deployments"][0]["resources"] = _merge_resources(observations, manifest) + request = SupportRequest.model_validate(payload) + target = tmp_path / "legacy-resolved" + create_plan(request, target) + for relative, content in formal_data.items(): + destination = target / relative + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(content) + saved_request = (target / "request.yaml").read_bytes() + trace = tmp_path / "trace.json" + _write( + trace, + { + "id": "synthetic-memory", + "models": ["source"], + "block_size": 16, + "hash_id_scope": "local", + "requests": [{"t": 0, "type": "s", "model": "source", "in": 16, "out": 2, "hash_ids": [1]}], + }, + ) + monkeypatch.setattr(supervision, "main", cli.main) + replay = tmp_path / "replay-resolved" + args = [ + "onboard", + "validate-fpm", + "--config", + str(target / "request.yaml"), + "--output-dir", + str(target), + "--trace", + str(trace), + "--validation-output-dir", + str(replay), + ] + assert cli.main(args) == 0 + plan = load_repeatability_source(next((root / "fpm-artifacts").iterdir())) + collection = { + "inputs": { + "collection_directory": str(root), + "formal_data": [ + workflow._identity(path) for path in sorted((root / "systems/data").rglob("fpm_forward_perf.*")) + ], + } + } + assert workflow._check_replay(replay / "validation.json", original, collection, plan)["status"] == "passed" + assert (target / "request.yaml").read_bytes() == saved_request + + changed = request.model_dump(mode="json") + changed["fpm_profile"]["deployments"][0]["resources"]["runtime_memory"]["kv_cache_bytes"] += 1024 + altered = SupportRequest.model_validate(changed) + bad_root = tmp_path / "altered-resolved" + create_plan(altered, bad_root) + for source in (target / "systems/data").rglob("*"): + if source.is_file(): + destination = bad_root / source.relative_to(target) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(source.read_bytes()) + args = [ + str(value).replace(str(target), str(bad_root)).replace(str(replay), str(tmp_path / "altered-replay")) + for value in args + ] + assert cli.main(args) == 0 + with pytest.raises(ValueError, match="resources differ"): + workflow._check_replay(tmp_path / "altered-replay/validation.json", original, collection, plan) + + +def test_replay_revalidates_capacity_revision_against_original_timing_collection(tmp_path, capsys, monkeypatch): + from collector.fpm_forward.repeatability import load_repeatability_source + + from aisimulate import supervision + from aisimulate.support.finalization import finalize + + from .test_onboard_runtime import _reviewed_capacity_revision + + _, original, _, root, memory_config = _reviewed_capacity_revision(tmp_path, capsys) + target = tmp_path / "resolved" + finalize(original, root, target, memory_config=memory_config) + trace = tmp_path / "trace.json" + _write( + trace, + { + "id": "synthetic-capacity-revision", + "models": ["source"], + "block_size": 16, + "hash_id_scope": "local", + "requests": [{"t": 0, "type": "s", "model": "source", "in": 16, "out": 2, "hash_ids": [1]}], + }, + ) + monkeypatch.setattr(supervision, "main", cli.main) + replay = tmp_path / "replay" + assert ( + cli.main( + [ + "onboard", + "validate-fpm", + "--config", + str(target / "request.yaml"), + "--output-dir", + str(target), + "--trace", + str(trace), + "--validation-output-dir", + str(replay), + ] + ) + == 0 + ) + plan = load_repeatability_source(next((root / "fpm-artifacts").iterdir())) + collection = { + "inputs": { + "collection_directory": str(root), + "formal_data": [ + workflow._identity(path) for path in sorted((root / "systems/data").rglob("fpm_forward_perf.*")) + ], + } + } + assert workflow._check_replay(replay / "validation.json", original, collection, plan)["status"] == "passed" + memory_config.write_text(memory_config.read_text() + "\n") + with pytest.raises(ValueError, match="memory revision.*changed"): + workflow._check_replay(replay / "validation.json", original, collection, plan) diff --git a/tests/sweeper/test_deploy.py b/tests/sweeper/test_deploy.py index 25e9f5bb6..eb2dc9223 100644 --- a/tests/sweeper/test_deploy.py +++ b/tests/sweeper/test_deploy.py @@ -269,6 +269,104 @@ def test_disagg_auto_transfer_geometry_uses_prefill_source_shape(monkeypatch): assert deployment.decode_engine_args["kv_transfer_bytes_per_token"] == 20_001 +def _role_geometry_profile(roles): + return { + "schema_version": 1, + "model": "example/role-geometry", + "model_revision": "synthetic-v1", + "architecture": "SyntheticDecoder", + "context_length": 4096, + "num_experts": 0, + "provenance": "Synthetic role geometry; not measured.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 1, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "bfloat16", + "moe_quant_mode": "bfloat16", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "bfloat16", + **({"worker_type": role} if role is not None else {}), + "resources": { + "weights_bytes": 1024, + "activations_bytes": 128, + "runtime_overhead_bytes": 128, + "comm_overhead_bytes": 0, + "kv_bytes_per_token": {None: 512, "prefill": 512, "decode": 1024, "aggregated": 2048}[role], + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Synthetic role geometry; not measured.", + }, + } + for role in roles + ], + } + + +@pytest.mark.parametrize("roles", [(None,), ("prefill", "decode"), ("prefill", "decode", "aggregated")]) +@pytest.mark.parametrize("feature", ["transfer", "host_offload"]) +def test_disagg_auto_geometry_selects_matching_profile_role(roles, feature): + profile = _role_geometry_profile(roles) + host_offload = {"num_host_blocks": 4096, "d2h_bandwidth_gbps": 7.0, "h2d_bandwidth_gbps": 38.0} + settings = ( + {"kv_transfer_bytes_per_token": "auto", "kv_transfer_bandwidth": 400.0} + if feature == "transfer" + else { + "prefill_native_host_offload": host_offload, + "decode_native_host_offload": host_offload, + "prefill_kv_bytes_per_token": "auto", + "decode_kv_bytes_per_token": "auto", + } + ) + parallel = ReplicaParallelConfig(ParallelShape(tp=1, dp=1, moe_tp=1, moe_ep=1), 1) + sample = unroll_sample( + search_space=_space(model_name=profile["model"], hardware_sku="h200_sxm", fpm_profile=profile, **settings), + selection=_agg_selection( + deployment_mode="disagg", + backend="vllm", + prefill_max_num_batched_tokens=8192, + prefill_max_num_seqs=1, + decode_max_num_batched_tokens=8192, + decode_max_num_seqs=256, + ), + parallel_config=DisaggParallelConfig(prefill=parallel, decode=parallel), + ) + + deployment = build_backend_deployment(sample, backend_version="0.25.1") + + field = "kv_transfer_bytes_per_token" if feature == "transfer" else "kv_cache_bytes_per_token" + assert deployment.prefill_engine_args[field] == 512 + assert deployment.decode_engine_args[field] == (512 if feature == "transfer" or roles == (None,) else 1024) + + +@pytest.mark.parametrize("roles,expected", [((None,), 512), (("prefill", "decode", "aggregated"), 2048)]) +def test_agg_auto_host_offload_geometry_selects_aggregated_profile_role(roles, expected): + profile = _role_geometry_profile(roles) + sample = unroll_sample( + search_space=_space( + model_name=profile["model"], + hardware_sku="h200_sxm", + fpm_profile=profile, + agg_native_host_offload={"num_host_blocks": 4096, "d2h_bandwidth_gbps": 7.0, "h2d_bandwidth_gbps": 38.0}, + agg_kv_bytes_per_token="auto", + ), + selection=_agg_selection(backend="vllm"), + parallel_config=ReplicaParallelConfig(ParallelShape(tp=1, dp=1, moe_tp=1, moe_ep=1), 1), + ) + + deployment = build_backend_deployment(sample, backend_version="0.25.1") + + assert deployment.agg_engine_args["worker_type"] == "aggregated" + assert deployment.agg_engine_args["kv_cache_bytes_per_token"] == expected + + def test_serialized_search_space_transfer_geometry_enables_pd_transfer(): legacy = _space( kv_transfer_bytes_per_token=333, diff --git a/tests/sweeper/test_search_space.py b/tests/sweeper/test_search_space.py index a9310e618..cb035a0ce 100644 --- a/tests/sweeper/test_search_space.py +++ b/tests/sweeper/test_search_space.py @@ -268,13 +268,14 @@ def fake_parallel_configs( gpu_budget, deployment_mode, backend, + worker_type, backend_version=None, min_gpu_budget=None, max_seq_len=None, role_runtime=None, systems_paths=None, ): - calls.append((hardware, deployment_mode, backend_version, role_runtime)) + calls.append((hardware, deployment_mode, backend_version, role_runtime, worker_type)) return [_AGG_CFG if hardware == "h200_sxm" else _DP8_CFG] monkeypatch.setattr("aisimulate.sweeper.search_space.parallel_configs_for", fake_parallel_configs) @@ -293,6 +294,7 @@ def fake_parallel_configs( ("gb200", "agg", "1.0"), ] assert all(set(call[3]) == {"agg"} for call in calls) + assert [call[4] for call in calls] == ["prefill", "decode"] expected = DisaggParallelConfig(prefill=_AGG_CFG, decode=_DP8_CFG) assert branch.parallel_configs == (expected,) assert branch.supported_backends[expected] == frozenset({"trtllm"}) diff --git a/tests/test_fpm_grouped_workflow.py b/tests/test_fpm_grouped_workflow.py new file mode 100644 index 000000000..46c42c1fc --- /dev/null +++ b/tests/test_fpm_grouped_workflow.py @@ -0,0 +1,905 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public grouped-cache workflows with fictional resources and synthetic timings.""" + +from __future__ import annotations + +import hashlib +import json +from copy import deepcopy +from pathlib import Path + +import pytest +import yaml +from pydantic import ValidationError + +from aisimulate import main as cli +from aisimulate.capacity import materialize_aic_num_gpu_blocks +from aisimulate.compiler import prediction_to_replay_spec +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.plan import create_plan +from aisimulate.support.schema import SupportRequest +from aisimulate.sweeper.config import Workload +from aisimulate.sweeper.kv_load import resolve_kv_load +from aisimulate.sweeper.model_hw import parallel_configs_for + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def grouped_case(tmp_path, monkeypatch): + import aisimulate_core + from aisimulate_core.sdk import engine, memory + + def no_graph(*_args, **_kwargs): + pytest.fail("grouped FPM workflow constructed an analytical model") + + monkeypatch.setattr(engine, "get_model", no_graph) + monkeypatch.setattr(engine, "build_model_config", no_graph) + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", no_graph) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", no_graph) + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + resources = { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "cache_layout": "grouped", + "cache_groups": [ + { + "name": "full", + "kind": "attention", + "num_layers": 1, + "block_size_tokens": 64, + "page_size_bytes": 64, + }, + { + "name": "window", + "kind": "attention", + "num_layers": 2, + "block_size_tokens": 16, + "page_size_bytes": 16, + "sliding_window": 512, + }, + { + "name": "conv", + "kind": "convolution", + "num_layers": 3, + "block_size_tokens": 4, + "page_size_bytes": 8, + "sliding_window": 4, + }, + ], + "max_num_tokens": 128, + "max_batch_size": 8, + "provenance": "Hand-sized fake cache pages including aggregate group layers; not measured GPU resources.", + } + deployment = { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": resources, + } + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["UnregisteredWindowMoeForCausalLM"]})) + profile = { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "synthetic-v1", + "architecture": "UnregisteredWindowMoeForCausalLM", + "context_length": 2048, + "num_experts": 4, + "provenance": "Fictional model and timings; functional coverage only, no silicon accuracy claim.", + "deployments": [ + {**deepcopy(deployment), "tp": 2, "dp": 1, "moe_tp": 2, "moe_ep": 1}, + {**deepcopy(deployment), "tp": 1, "dp": 2, "moe_tp": 1, "moe_ep": 2}, + ], + } + root = tmp_path / "systems" + root.mkdir() + system = yaml.safe_load((Path(aisimulate_core.__file__).parent / "systems/h200_sxm.yaml").read_text()) + # A deliberately tiny fake GPU creates meaningful byte pressure in CPU tests. + system["gpu"]["mem_capacity"] = 4096 + (root / "h200_sxm.yaml").write_text(yaml.safe_dump(system)) + _write_timings(root, profile) + return profile, root + + +def _write_timings(root, profile, *, decode_ceiling=8192): + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + rows = [] + for index, deployment in enumerate(load_fpm_profile(profile).deployments): + identity = deployment.model_dump(mode="json", exclude={"resources"}) + for batch in (1, 2, 4, 8): + for phase, tokens, kvs in ( + ( + "prefill", + (batch, 128, 256, 512, 1024), + (0, 1, 512, 1024, 2048, 8192), + ), + ("decode", (0,), (0, 1, 512, decode_ceiling)), + ): + for new_tokens in sorted(set(tokens)): + for kv in sorted(set(kvs)): + rows.append( + { + **identity, + "cell_id": f"synthetic-{index}-{phase}-{batch}-{new_tokens}-{kv}", + "model_path": profile["model"], + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": batch, + "total_prefill_tokens": new_tokens, + "total_kv_read_tokens": kv, + "latency_ms": 1 + batch + new_tokens / 128 + kv / 1024, + "kv_seed_regime": "real_kv", + } + ) + path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + ) + + +def _prediction(profile, root, *, dep=False): + return { + "engine": { + "mode": "aggregated", + "model": profile["model"], + "fpm_profile": profile, + "hardware": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "context_length": 2048, + "systems_paths": [str(root)], + "workers": { + "aggregated": { + "parallelism": { + "tensor": 1 if dep else 2, + "attention_data": 2 if dep else 1, + "moe_tensor": 1 if dep else 2, + "moe_expert": 2 if dep else 1, + }, + "scheduler": {"max_batched_tokens": 128, "max_sequences": 4}, + "kv_cache": {"prefix_caching": False}, + "timing": { + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": { + "fpm_interpolation": { + "method": "direct", + "collect_coverage": True, + } + }, + }, + } + }, + }, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1152, "output_tokens": 4}, + "load": {"type": "concurrency", "concurrency": 2}, + "stop": {"requests": 2}, + }, + } + + +def _mixed_profile(profile): + profile = deepcopy(profile) + resources = profile["deployments"][0]["resources"] + resources["cache_layout"] = "linear" + resources["kv_bytes_per_token"] = 1 + resources.pop("cache_groups") + # A full-only grouped deployment is valid beside a linear deployment. + profile["deployments"][1]["resources"]["cache_groups"] = profile["deployments"][1]["resources"]["cache_groups"][:1] + return profile + + +@pytest.mark.parametrize("unrelated", ["topology", "hardware", "version"]) +def test_mixed_profile_predict_preserves_selected_linear_prefix_cache(grouped_case, tmp_path, unrelated): + profile, root = grouped_case + profile = _mixed_profile(profile) + if unrelated != "topology": + profile["deployments"][1].update(tp=2, dp=1, moe_tp=2, moe_ep=1) + profile["deployments"][1].update( + {"system": "h100_sxm"} if unrelated == "hardware" else {"backend_version": "0.25.2"} + ) + raw = _prediction(profile, root) + raw["engine"]["workers"]["aggregated"]["kv_cache"]["prefix_caching"] = True + CorePredictionConfig.model_validate(raw) + path = tmp_path / "linear.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "linear-prediction" + assert cli.main(["predict", "-c", str(path), "--output-dir", str(output), "--detail", "memory"]) == 0 + report = json.loads((output / "prediction.json").read_text()) + assert report["completed_requests"] == 2 + assert report["details"]["sections"]["memory"]["roles"]["aggregated"]["total_kv_size_tokens"] > 0 + + +@pytest.mark.parametrize("preset", [False, True]) +def test_mixed_profile_recommendation_preserves_pinned_linear_prefix_cache(grouped_case, tmp_path, monkeypatch, preset): + from aisimulate.recommend import _run_recommendation, recommendation_to_sweeper + from aisimulate.sweeper.search_space import enumerate_branches + + profile, root = grouped_case + raw = _prediction(_mixed_profile(profile), root) + worker = raw["engine"]["workers"]["aggregated"] + worker["kv_cache"]["prefix_caching"] = True + worker["parallelism"]["preset"] = False + if preset: + worker["parallelism"] = {"preset": [{**worker["parallelism"], "replicas": 1, "pipeline": 1}]} + worker["parallelism"]["preset"][0].pop("preset") + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + smart = recommendation_to_sweeper(CoreRecommendationConfig.model_validate(raw)) + branches = enumerate_branches(smart) + assert len(branches) == 1 + assert {candidate.shape.tp for candidate in branches[0].parallel_configs} == {2} + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + path = tmp_path / "linear-recommend.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "linear-recommendation" + assert cli.main(["recommend", "-c", str(path), "--output-dir", str(output)]) == 0 + exported = CorePredictionConfig.from_yaml(output / "recommendations/0001.yaml") + assert exported.engine.workers.aggregated.parallelism.tensor == 2 + assert exported.engine.workers.aggregated.kv_cache.prefix_caching is True + + +def _mixed_search(profile, root): + return { + "deployment_mode": ["agg"], + "backend": ["vllm"], + "backend_version": "0.25.1", + "model_name": profile["model"], + "hardware_sku": "h200_sxm", + "fpm_profile": _mixed_profile(profile), + "systems_paths": [str(root)], + "gpu_budget": 2, + "agg_max_num_batched_tokens": [128], + "agg_max_num_seqs": [4], + "agg_enable_prefix_caching": True, + } + + +@pytest.mark.parametrize("filter_by", ["hardware", "version", "budget", "pin", "mode_pin", "choices", "log_range"]) +def test_mixed_search_ignores_unreachable_grouped_deployments(grouped_case, filter_by): + from aisimulate.sweeper.config import SmartSearchConfig + from aisimulate.sweeper.search_space import enumerate_branches + + profile, root = grouped_case + raw = _mixed_search(profile, root) + grouped = raw["fpm_profile"]["deployments"][1] + if filter_by == "hardware": + grouped["system"] = "h100_sxm" + elif filter_by == "version": + grouped["backend_version"] = "0.25.2" + raw["backend_version"] = {"vllm": " 0.25.1 "} + elif filter_by == "budget": + grouped.update(tp=4, dp=1, moe_tp=1, moe_ep=4) + elif filter_by in {"pin", "mode_pin"}: + pinned = [{"tp": 2, "moe_tp": 2}] + raw["parallel_configs" if filter_by == "pin" else "parallel_configs_by_mode"] = ( + pinned if filter_by == "pin" else {"agg": pinned} + ) + elif filter_by == "choices": + raw["parallel_independent_by_mode"] = {"agg": {"tp": [2]}} + else: + raw["parallel_independent_by_mode"] = {"agg": {"tp": None}} + raw["parallel_independent_log_ranges_by_mode"] = {"agg": {"tp": [2, 4]}} + config = SmartSearchConfig(search_space=raw, workload=Workload(isl=1152, osl=4, concurrency=2, request_count=2)) + branches = enumerate_branches(config) + assert len(branches) == 1 + assert {candidate.shape.tp for candidate in branches[0].parallel_configs} == {2} + + +@pytest.mark.parametrize("placement", ["aggregated", "disaggregated", "heterogeneous"]) +def test_mixed_recommendation_linear_fixed_capacity_and_disaggregation(grouped_case, placement): + from aisimulate.recommend import recommendation_to_sweeper + from aisimulate.sweeper.search_space import enumerate_branches + + profile, root = grouped_case + raw = _prediction(_mixed_profile(profile), root) + # Unselected deployments must not reach resource-envelope validation either. + raw["engine"]["fpm_profile"]["deployments"][1]["resources"]["max_num_tokens"] = 64 + worker = raw["engine"]["workers"]["aggregated"] + worker["kv_cache"].update(prefix_caching=True, capacity={"type": "fixed", "blocks": 64}) + disaggregated = placement != "aggregated" + if disaggregated: + raw["engine"]["mode"] = "disaggregated" + raw["engine"]["workers"] = {"prefill": deepcopy(worker), "decode": deepcopy(worker)} + if placement == "heterogeneous": + linear = deepcopy(raw["engine"]["fpm_profile"]["deployments"][0]) + linear["system"] = "h100_sxm" + raw["engine"]["fpm_profile"]["deployments"].append(linear) + raw["engine"]["workers"]["decode"]["hardware"] = "h100_sxm" + (root / "h100_sxm.yaml").write_bytes((root / "h200_sxm.yaml").read_bytes()) + CorePredictionConfig.model_validate(raw) + for worker in raw["engine"]["workers"].values(): + worker["parallelism"]["preset"] = False + worker["parallelism"]["replicas"] = 1 + if placement == "heterogeneous": + parallel = raw["engine"]["workers"]["prefill"]["parallelism"] + parallel.pop("preset") + raw["engine"]["workers"]["prefill"]["parallelism"] = {"preset": [{**parallel, "pipeline": 1}]} + raw["optimization"] = {"constraints": {"max_candidate_gpus": 4 if disaggregated else 2}} + branches = enumerate_branches(recommendation_to_sweeper(CoreRecommendationConfig.model_validate(raw))) + assert len(branches) == 1 + assert len(branches[0].parallel_configs) == 1 + candidate = branches[0].parallel_configs[0] + assert (candidate.prefill.shape.tp if disaggregated else candidate.shape.tp) == 2 + + +def test_mixed_search_keeps_both_layouts_when_grouped_settings_are_supported(grouped_case): + from aisimulate.sweeper.config import SmartSearchConfig + from aisimulate.sweeper.search_space import enumerate_branches + + profile, root = grouped_case + raw = _mixed_search(profile, root) + raw["agg_enable_prefix_caching"] = False + config = SmartSearchConfig(search_space=raw, workload=Workload(isl=1152, osl=4, concurrency=2, request_count=2)) + assert {candidate.shape.tp for candidate in enumerate_branches(config)[0].parallel_configs} == {1, 2} + + +@pytest.mark.parametrize( + "update,error", + [ + ({}, "prefix_caching=false"), + ({"agg_enable_prefix_caching": False, "agg_native_host_offload": {"num_host_blocks": 64}}, "only HBM"), + ({"agg_enable_prefix_caching": False, "agg_num_gpu_blocks": 64}, "not scalar capacity"), + ({"agg_enable_prefix_caching": False, "agg_kv_bytes_per_token": 1}, "not scalar capacity"), + ({"deployment_mode": ["disagg"], "gpu_budget": 4}, "without speculative decoding"), + ], +) +def test_mixed_search_rejects_reachable_grouped_unsupported_settings(grouped_case, update, error): + from aisimulate.sweeper.config import SearchSpace + + profile, root = grouped_case + raw = _mixed_search(profile, root) + raw.update(update) + with pytest.raises(ValidationError, match=error): + SearchSpace.model_validate(raw) + + +def test_public_recommendation_rejects_reachable_grouped_before_estimating(grouped_case, tmp_path, monkeypatch, capsys): + from aisimulate.recommend import _run_recommendation + from aisimulate_core.sdk import RustForwardPassPerfModel + + def no_estimator(*_args, **_kwargs): + pytest.fail("unsupported grouped recommendation reached estimator construction") + + profile, root = grouped_case + raw = _prediction(_mixed_profile(profile), root) + raw["engine"]["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw["engine"]["workers"]["aggregated"]["kv_cache"]["prefix_caching"] = True + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + CoreRecommendationConfig.model_validate(raw) + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", no_estimator) + path = tmp_path / "unsupported-recommend.yaml" + path.write_text(yaml.safe_dump(raw)) + with pytest.raises(SystemExit) as error: + cli.main(["recommend", "-c", str(path), "--output-dir", str(tmp_path / "unsupported")]) + assert error.value.code == 2 + output = capsys.readouterr() + assert "grouped FPM cache requires agg_enable_prefix_caching=false" in output.out + output.err + + +@pytest.mark.parametrize("dep", [False, True]) +@pytest.mark.parametrize("layout", ["mixed", "window_only", "full_only"]) +def test_public_predict_uses_group_bytes_with_tp_and_dep(grouped_case, tmp_path, dep, layout): + profile, root = grouped_case + if layout != "mixed": + for deployment in profile["deployments"]: + groups = deployment["resources"]["cache_groups"] + deployment["resources"]["cache_groups"] = groups[:1] if layout == "full_only" else groups[1:] + raw = _prediction(profile, root, dep=dep) + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) + diagnostics = {} + lowered = materialize_aic_num_gpu_blocks(spec.backend_deployment.agg_engine_args, memory_diagnostics=diagnostics) + assert "num_gpu_blocks" not in lowered + assert lowered["kv_cache_capacity_bytes"] == 3486 # floor(4096 * .9 - 200) + assert lowered["kv_cache_groups"] == diagnostics["cache_groups"] + assert diagnostics["total_kv_size_tokens"] is None + path = tmp_path / "predict.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "prediction" + assert ( + cli.main( + [ + "predict", + "-c", + str(path), + "--output-dir", + str(output), + "--detail", + "memory", + "--capture-per-request", + ] + ) + == 0 + ) + report = json.loads((output / "prediction.json").read_text()) + assert report["completed_requests"] == 2 + assert all(row["output_length"] == 4 for row in report["per_request"]) + assert report["fpm_query_coverage"]["status"] == "covered" + memory = report["details"]["sections"]["memory"]["roles"]["aggregated"] + assert memory["total_kv_size_bytes"] == 3486 + assert memory["total_kv_size_tokens"] is None + assert "estimated_num_gpu_blocks" not in memory + + +def test_public_predict_does_not_clamp_logical_context_to_window(grouped_case, tmp_path): + profile, root = grouped_case + # Window-only allocation is bounded, but the first decode after the + # prefill-emitted token must read all 1152 prompt KV tokens. A table ending + # at 512 must fail, not shrink the logical context or physical byte budget. + for deployment in profile["deployments"]: + deployment["resources"]["cache_groups"] = deployment["resources"]["cache_groups"][1:] + _write_timings(root, profile, decode_ceiling=512) + raw = _prediction(profile, root) + raw["traffic"]["load"]["concurrency"] = 1 + raw["traffic"]["stop"]["requests"] = 1 + path = tmp_path / "uncovered.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "uncovered" + assert cli.main(["predict", "-c", str(path), "--output-dir", str(output)]) != 0 + coverage = json.loads((output / "fpm-coverage.json").read_text()) + gaps = coverage["roles"][0]["coverage"]["gaps"] + assert any(gap["coordinates"]["total_kv_read_tokens"] == 1152 for gap in gaps) + + +@pytest.mark.parametrize("relative_load", [False, True]) +@pytest.mark.parametrize("observed_memory", [False, True]) +def test_grouped_recommendation_exports_a_runnable_prediction( + grouped_case, tmp_path, monkeypatch, relative_load, observed_memory +): + from aisimulate.recommend import _run_recommendation + + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + profile, root = grouped_case + if observed_memory: + for deployment in profile["deployments"]: + resources = deployment["resources"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(name) + resources["max_batch_size"] = 4 + resources["runtime_memory"] = { + "kv_cache_bytes": 3486, + "max_model_len": 2048, + "gpu_memory_utilization": 0.9, + "provenance": "Synthetic initialized cache pool; no GPU qualification.", + } + raw = _prediction(profile, root) + raw["engine"]["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 2, "parallelism": 1} + if relative_load: + raw["traffic"]["load"] = {"type": "kv_capacity_fraction", "fraction": 1.0} + CoreRecommendationConfig.model_validate(raw) + path = tmp_path / "recommend.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "recommendation" + assert cli.main(["recommend", "-c", str(path), "--output-dir", str(output)]) == 0 + exported_path = output / "recommendations/0001.yaml" + exported = CorePredictionConfig.from_yaml(exported_path) + assert exported.engine.fpm_profile.deployments[0].resources.cache_layout == "grouped" + assert exported.engine.workers.aggregated.kv_cache.prefix_caching is False + assert ( + cli.main( + [ + "predict", + "-c", + str(exported_path), + "--output-dir", + str(tmp_path / "exported-prediction"), + ] + ) + == 0 + ) + + +@pytest.mark.parametrize( + "field,value,error", + [ + ("kv_cache_capacity_bytes", 9999, "byte capacity conflicts"), + ("kv_cache_groups", None, "cache groups conflict"), + ("tensor_parallel_size", 4, "tensor_parallel_size|topology"), + ("dp_size", 2, "dp_size|topology"), + ], +) +def test_native_grouped_rematerialization_checks_identity_and_budget(grouped_case, field, value, error): + from aisimulate import _runtime + from aisimulate.runner import EngineReplayRunnerFactory + + profile, root = grouped_case + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(_prediction(profile, root))) + + class ChangedNativeInput: + def run_replay_json(self, serialized): + payload = json.loads(serialized) + engine = payload.get("spec", payload)["engine"] + if field in {"tensor_parallel_size", "dp_size"}: + engine[field] = value + elif field == "kv_cache_groups": + # Wrong but internally valid physical page geometry must not + # replace the canonical profile after Python materialization. + engine["rank"][field][0]["page_size_bytes"] += 1 + else: + engine["rank"][field] = value + return _runtime.run_replay_json(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + EngineReplayRunnerFactory(runtime=ChangedNativeInput()).create(0).run(spec) + + +def test_grouped_capacity_preserves_graph_reservation_and_rejects_fixed_blocks(grouped_case): + profile, root = grouped_case + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(_prediction(profile, root))) + args = deepcopy(spec.backend_deployment.agg_engine_args) + args["cuda_graph_reserved_bytes"] = 1000 + args["timing_model"]["config"]["cuda_graph_reserved_bytes"] = 1000 + budget = materialize_aic_num_gpu_blocks(args) + assert budget["kv_cache_capacity_bytes"] == 2486 + args["num_gpu_blocks"] = 1 + with pytest.raises(ValueError, match="fixed num_gpu_blocks"): + materialize_aic_num_gpu_blocks(args) + + +def test_grouped_topology_admission_uses_each_rank_byte_budget_before_collection(grouped_case, monkeypatch): + from aisimulate_core.sdk import RustForwardPassPerfModel + + profile, root = grouped_case + # DEP has an explicitly larger rank-local weight reservation. It leaves + # 1086 bytes, insufficient for even the 2048-byte full-history group. + profile["deployments"][1]["resources"]["weights_bytes"] = 2500 + + def no_timing(*_args, **_kwargs): + pytest.fail("candidate fit constructed a timing estimator before collection") + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", no_timing) + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + max_num_tokens=128, + max_batch_size=4, + fpm_profile=profile, + systems_paths=[str(root)], + ) + assert [candidate.shape.strategy for candidate in candidates] == ["tp"] + + +def test_grouped_kv_load_uses_bytes_without_inventing_token_capacity(grouped_case): + profile, root = grouped_case + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + max_num_tokens=128, + max_batch_size=4, + fpm_profile=profile, + systems_paths=[str(root)], + ) + assert len(candidates) == 2 + for candidate in candidates: + sample = { + "model_name": profile["model"], + "hardware_sku": "h200_sxm", + "backend": "vllm", + "fpm_profile": profile, + "systems_paths": [str(root)], + "agg_block_size": 64, + "agg_max_num_batched_tokens": 128, + "agg_max_num_seqs": 4, + "agg_gpu_memory_utilization": 0.9, + } + result = resolve_kv_load( + sample, + workload=Workload(isl=1152, osl=4, kv_load_ratio=1.0, request_count=2), + parallel_config=candidate, + ratio=1.0, + backend_version="0.25.1", + ) + # 1216 full-history + 528 sliding-window + 16 convolution bytes. + assert result.role_request_cache_bytes == {"agg": 1760} + assert result.role_capacity_bytes == {"agg": 3486 * candidate.shape.dp} + assert result.concurrency_capacity == candidate.shape.dp + assert result.role_capacity_tokens == {} + + +@pytest.mark.parametrize( + "update,error", + [ + ({"prefix_caching": True}, "prefix_caching=false"), + ({"capacity": {"type": "fixed", "blocks": 1024}}, "byte budget"), + ({"bytes_per_token": 8}, "byte budget"), + ({"host_offload": {"num_host_blocks": 64}}, "only HBM"), + ], +) +def test_grouped_incompatible_cache_settings_fail_before_replay(grouped_case, update, error): + profile, root = grouped_case + raw = _prediction(profile, root) + raw["engine"]["workers"]["aggregated"]["kv_cache"].update(update) + with pytest.raises(ValidationError, match=error): + CorePredictionConfig.model_validate(raw) + + +@pytest.mark.parametrize( + "changes", + [ + {"mode": "disaggregated"}, + {"nextn": 1, "nextn_accepted": 1.0}, + {"speculation": {"kind": "ngram", "num_speculative_tokens": 1, "acceptance_rates": [0.5]}}, + ], +) +def test_grouped_unsupported_execution_modes_fail_before_replay(grouped_case, changes): + profile, root = grouped_case + raw = _prediction(profile, root) + raw["engine"].update(changes) + if changes.get("mode") == "disaggregated": + worker = raw["engine"]["workers"].pop("aggregated") + raw["engine"]["workers"] = {"prefill": deepcopy(worker), "decode": worker} + with pytest.raises(ValidationError, match="grouped FPM cache supports only aggregated vLLM"): + CorePredictionConfig.model_validate(raw) + + +def test_grouped_onboarding_profile_reaches_real_collector_preview_and_render( + grouped_case, tmp_path, monkeypatch, capsys +): + from collector.fpm_forward import cli as collector_cli + from collector.fpm_forward import runner as collector_runner + + profile, _root = grouped_case + profile = deepcopy(profile) + for deployment in profile["deployments"]: + deployment["gemm_quant_mode"] = "fp8_static" + deployment["fmha_quant_mode"] = "fp8" + local_config = tmp_path / "config.json" + local_config.write_text( + json.dumps( + { + "architectures": [profile["architecture"]], + "model_type": "example_moe", + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 3, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 1024, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 128, + "max_position_embeddings": 2048, + "torch_dtype": "bfloat16", + "quantization_config": {"quant_method": "fp8", "kv_cache_scheme": "FP8"}, + "sliding_window": 512, + "layer_types": ["full_attention", "sliding_attention", "sliding_attention"], + } + ) + ) + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": { + "tensor_parallel": 2, + "attention_data_parallel": 1, + "moe_tensor_parallel": 2, + "moe_expert_parallel": 1, + "context_length": 2048, + }, + "collection": {"max_num_tokens": 128, "max_batch_size": 4}, + "fpm_profile": profile, + } + ) + plan_dir = tmp_path / "collection" + create_plan(request, plan_dir) + command = json.loads((plan_dir / "commands.json").read_text())["fpm_plan_local"][3:] + command += ["--fpm-model-config", str(local_config)] + monkeypatch.delenv("COLLECTOR_MODEL_PATH", raising=False) + assert collector_cli.main(command) == 0 + preview = json.loads(capsys.readouterr().out) + assert preview["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert preview["point_generation"]["method"] == "native_self_benchmark" + assert preview["point_generation"]["planned_point_count"] is None + assert {cell["parallel_strategy"] for cell in preview["cells"]} == {"pure_tp"} + assert ( + preview["topology_memory_admission"][0]["estimates"][0]["provenance"] + == (profile["deployments"][0]["resources"]["provenance"]) + ) + + def render_without_launch(_args, resolved): + plan, overrides = resolved + assert plan.to_dict()["fpm_profile"] == preview["fpm_profile"] + for cell in plan.cells: + target = tmp_path / f"render-{cell.workload_kind}" + target.mkdir() + collector_runner._render_cell(plan, cell, target, overrides) + script = (target / "run.sh").read_text() + assert f"--model {profile['model']}" in script + assert "--max-num-batched-tokens 128" in script + assert "--max-num-seqs 4" in script + return [] + + monkeypatch.setattr(collector_cli, "run_resolved", render_without_launch) + assert collector_cli.main([arg for arg in command if arg != "--plan-only"]) == 0 + + +@pytest.mark.parametrize("dep", [False, True]) +def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case, tmp_path, monkeypatch, dep): + from aisimulate import supervision + + profile, _root = grouped_case + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": { + "tensor_parallel": 1 if dep else 2, + "attention_data_parallel": 2 if dep else 1, + "moe_tensor_parallel": 1 if dep else 2, + "moe_expert_parallel": 2 if dep else 1, + "context_length": 2048, + }, + "collection": {"max_num_tokens": 128, "max_batch_size": 4}, + "fpm_profile": profile, + } + ) + plan = tmp_path / "collection" + create_plan(request, plan) + _write_timings(plan / "systems", profile) + trace = tmp_path / "trace.jsonl" + trace.write_text( + json.dumps( + { + "id": "window-boundary-play", + "models": ["trace/source-model"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "s", + "model": "trace/source-model", + "in": 512, + "out": 3, + "hash_ids": list(range(8)), + }, + { + "t": 0.1, + "type": "s", + "model": "trace/source-model", + "in": 1152, + "out": 4, + "hash_ids": list(range(18)), + }, + ], + } + ) + + "\n" + ) + monkeypatch.setattr(supervision, "main", cli.main) + output = tmp_path / "validation" + assert ( + cli.main( + [ + "onboard", + "validate-fpm", + "--config", + str(plan / "request.yaml"), + "--output-dir", + str(plan), + "--trace", + str(trace), + "--validation-output-dir", + str(output), + ] + ) + == 0 + ) + validation = json.loads((output / "validation.json").read_text()) + assert validation["status"] == "covered" + assert validation["accuracy"] == "not_assessed" + report = json.loads((output / "prediction/prediction.json").read_text()) + assert report["completed_requests"] == 2 + assert report["agentic_model_projection"]["target_model"] == profile["model"] + saved = CorePredictionConfig.from_yaml(output / "predict.yaml") + assert saved.engine.workers.aggregated.kv_cache.prefix_caching is False + assert saved.engine.fpm_profile.deployments[0].resources.cache_layout == "grouped" + + +@pytest.mark.parametrize("profile_field", ["aic_fpm_profile", "fpm_profile"]) +def test_grouped_runner_capability_rejects_unsupported_stack(grouped_case, profile_field): + from aisimulate.runner import EngineReplayRunnerFactory + from aisimulate.sweeper.replay import RunnerCapabilities + + profile, root = grouped_case + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(_prediction(profile, root))) + with pytest.raises(ValueError, match="grouped FPM caches; use --stack engine"): + RunnerCapabilities(supported_backend_topologies=(("vllm", "agg"),)).require_compatible(spec) + EngineReplayRunnerFactory().capabilities().require_compatible(spec) + linear_spec = prediction_to_replay_spec( + CorePredictionConfig.model_validate(_prediction(_mixed_profile(profile), root)) + ) + RunnerCapabilities(supported_backend_topologies=(("vllm", "agg"),)).require_compatible(linear_spec) + + from dataclasses import replace + + for profile_variant, grouped in ((profile, True), (_mixed_profile(profile), False)): + selected = profile_variant["deployments"][0] + flat = { + "aic_model_path": profile_variant["model"], + "aic_system": selected["system"], + profile_field: profile_variant, + "aic_tp_size": selected["tp"], + "aic_attention_dp_size": selected["dp"], + "aic_moe_tp_size": selected["moe_tp"], + "aic_moe_ep_size": selected["moe_ep"], + } + flat_spec = replace(spec, backend_deployment=replace(spec.backend_deployment, agg_engine_args=flat)) + capabilities = RunnerCapabilities(supported_backend_topologies=(("vllm", "agg"),)) + if grouped: + with pytest.raises(ValueError, match="grouped FPM caches"): + capabilities.require_compatible(flat_spec) + else: + capabilities.require_compatible(flat_spec) + + +def test_aggregated_only_grouped_profile_skips_unavailable_disagg(grouped_case): + from aisimulate.sweeper.config import SmartSearchConfig + from aisimulate.sweeper.search_space import enumerate_branches + + profile, root = grouped_case + for deployment in profile["deployments"]: + deployment["worker_type"] = "aggregated" + raw = _mixed_search(profile, root) + raw.update(fpm_profile=profile, agg_enable_prefix_caching=False) + raw.pop("deployment_mode") + config = SmartSearchConfig(search_space=raw, workload=Workload(isl=1152, osl=4, concurrency=1, request_count=2)) + with pytest.warns(UserWarning, match="disagg.*skipped"): + branches = enumerate_branches(config, max_seq_len=2048) + assert [branch.deployment_mode for branch in branches] == ["agg"] diff --git a/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py index 9786fb4f3..dccb7bbba 100644 --- a/tests/test_fpm_profile_workflow.py +++ b/tests/test_fpm_profile_workflow.py @@ -206,6 +206,174 @@ def test_prediction_materializes_rank_local_capacity_without_a_model(profile, ti assert reserved["num_gpu_blocks"] < lowered["num_gpu_blocks"] +@pytest.mark.parametrize("memory_state", ["pending", "runtime"]) +def test_profile_memory_cannot_be_bypassed_by_explicit_blocks(profile, timing_systems, memory_state): + config = CorePredictionConfig.model_validate({"engine": {**_engine(profile), "systems_paths": [timing_systems]}}) + args = deepcopy(prediction_to_replay_spec(config).backend_deployment.agg_engine_args) + resources = args["timing_model"]["config"]["fpm_profile"]["deployments"][0]["resources"] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(field) + if memory_state == "runtime": + resources["runtime_memory"] = { + "kv_cache_bytes": 1024**3, + "gpu_memory_utilization": 0.9, + "max_model_len": 4096, + "provenance": "Synthetic initialized worker; no silicon qualification.", + } + args["num_gpu_blocks"] = 1 + with pytest.raises(ValueError, match="pending runtime profiling|runtime FPM memory cannot use fixed"): + materialize_aic_num_gpu_blocks(args) + + +@pytest.mark.parametrize("memory_state,missing_timings", [("pending", False), ("runtime", False), ("pending", True)]) +def test_native_profile_memory_rejects_explicit_capacity_bypass(profile, timing_systems, memory_state, missing_timings): + from aisimulate import _runtime + from aisimulate.runner import EngineReplayRunnerFactory + + config = CorePredictionConfig.model_validate( + { + "engine": {**_engine(profile), "systems_paths": [timing_systems]}, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 1}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + }, + } + ) + spec = prediction_to_replay_spec(config) + + class ChangedNativeInput: + def run_replay_json(self, serialized): + payload = json.loads(serialized) + engine = payload.get("spec", payload)["engine"] + engine["num_gpu_blocks_is_explicit"] = True + rank = engine["rank"] + if missing_timings: + rank["timing_model"]["config"]["systems_paths"] = [str(Path(timing_systems) / "missing")] + resources = rank["timing_model"]["config"]["fpm_profile"]["deployments"][0]["resources"] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(field) + if memory_state == "runtime": + resources["runtime_memory"] = { + "kv_cache_bytes": 1024**3, + "gpu_memory_utilization": 0.9, + "max_model_len": 4096, + "provenance": "Synthetic initialized worker; no silicon qualification.", + } + return _runtime.run_replay_json(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="pending runtime profiling|runtime FPM memory cannot use fixed"): + EngineReplayRunnerFactory(runtime=ChangedNativeInput()).create(0).run(spec) + + +@pytest.mark.parametrize("memory_state", ["pending", "runtime"]) +def test_recommendation_cannot_bypass_profile_memory_with_fixed_tokens(profile, memory_state): + for deployment in profile["deployments"]: + resources = deployment["resources"] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(field) + if memory_state == "runtime": + resources["runtime_memory"] = { + "kv_cache_bytes": 1024**3, + "gpu_memory_utilization": 0.9, + "max_model_len": 4096, + "provenance": "Synthetic initialized worker; no silicon qualification.", + } + with pytest.raises(ValueError, match="pending runtime profiling|runtime FPM memory cannot use a fixed"): + parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + fpm_profile=profile, + role_runtime={"agg": (8192, 8, 0.9, 1_000_000)}, + ) + + +def test_runtime_recommendation_propagates_selected_context(profile): + for deployment in profile["deployments"]: + resources = deployment["resources"] + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(field) + resources["runtime_memory"] = { + "kv_cache_bytes": 1024**3, + "gpu_memory_utilization": 0.9, + "max_model_len": 2048, + "provenance": "Synthetic initialized worker; no silicon qualification.", + } + assert parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + fpm_profile=profile, + max_seq_len=1024, + max_batch_size=8, + ) + with pytest.raises(ValueError, match="max_model_len"): + parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + fpm_profile=profile, + max_seq_len=2049, + max_batch_size=8, + ) + + +@pytest.mark.parametrize("command", ["predict", "recommend"]) +def test_public_commands_report_pending_memory_before_missing_timings(profile, tmp_path, monkeypatch, capsys, command): + import aisimulate.main as cli + import aisimulate_core + from aisimulate.recommend import _run_recommendation + + for deployment in profile["deployments"]: + for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + deployment["resources"].pop(field) + root = tmp_path / "empty-systems" + root.mkdir() + packaged = Path(aisimulate_core.__file__).parent / "systems/h200_sxm.yaml" + (root / packaged.name).write_bytes(packaged.read_bytes()) + engine = {**_engine(profile), "mode": "aggregated", "systems_paths": [str(root)]} + raw = { + "engine": engine, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 1}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + }, + } + if command == "recommend": + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + CoreRecommendationConfig.model_validate(raw) # Pending templates remain schema-valid. + else: + CorePredictionConfig.model_validate(raw) + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + path = tmp_path / f"{command}.json" + path.write_text(json.dumps(raw)) + try: + status = cli.main( + [command, "--config", str(path), "--output-dir", str(tmp_path / "output"), "--format", "json"] + ) + except SystemExit as exc: + status = exc.code + assert status in {1, 2} + message = capsys.readouterr().err + assert "pending runtime profiling" in message + assert "finalize runtime memory" in message + assert "perf data directory not found" not in message + + def test_disaggregated_transfer_uses_profile_cache_geometry(profile): engine = _engine(profile) engine.update( @@ -218,6 +386,112 @@ def test_disaggregated_transfer_uses_profile_cache_geometry(profile): assert spec.backend_deployment.decode_engine_args["kv_transfer_bytes_per_token"] == 1024 +def test_disaggregated_profile_selects_same_topology_resources_by_role(profile, timing_systems): + prefill = deepcopy(profile["deployments"][0]) + decode = deepcopy(prefill) + prefill["worker_type"] = "prefill" + prefill["resources"]["weights_bytes"] += 10 * 1024**3 + decode["worker_type"] = "decode" + profile["deployments"] = [prefill, decode] + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(profile), + "mode": "disaggregated", + "systems_paths": [timing_systems], + "workers": {"prefill": _worker(), "decode": _worker()}, + } + } + ) + reloaded = CorePredictionConfig.model_validate_json(config.model_dump_json()) + deployment = prediction_to_replay_spec(reloaded).backend_deployment + blocks = {} + for role in ("prefill", "decode"): + args = getattr(deployment, f"{role}_engine_args") + assert args["timing_model"]["config"]["worker_type"] == role + diagnostics = {} + lowered = materialize_aic_num_gpu_blocks(args, memory_diagnostics=diagnostics) + blocks[role] = lowered["num_gpu_blocks"] + expected = prefill if role == "prefill" else decode + assert diagnostics["memory_breakdown"]["weights_bytes"] == expected["resources"]["weights_bytes"] + assert blocks["prefill"] < blocks["decode"] + conflicting = deepcopy(deployment.decode_engine_args) + conflicting["worker_type"] = "prefill" + with pytest.raises(ValueError, match="worker_type conflicts"): + materialize_aic_num_gpu_blocks(conflicting) + + +def test_disaggregated_profile_candidates_do_not_swap_roles(profile): + profile["deployments"][0]["worker_type"] = "decode" + profile["deployments"][1]["worker_type"] = "prefill" + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="disagg", + gpu_budget=4, + max_num_tokens=1024, + max_batch_size=4, + fpm_profile=profile, + ) + assert len(candidates) == 1 + assert candidates[0].prefill.shape.strategy == "dep" + assert candidates[0].decode.shape.strategy == "tp" + with pytest.raises(ValueError, match="no parallel config"): + parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=4, + max_num_tokens=1024, + max_batch_size=4, + fpm_profile=profile, + ) + + +def test_native_disaggregated_capacity_keeps_canonical_worker_role(profile, timing_systems, monkeypatch): + from aisimulate.runner import EngineReplayRunnerFactory + from aisimulate_core.sdk import memory + + prefill = deepcopy(profile["deployments"][0]) + decode = deepcopy(prefill) + prefill["worker_type"] = "prefill" + decode["worker_type"] = "decode" + profile["deployments"] = [prefill, decode] + worker = _worker() + worker["kv_cache"] = {"prefix_caching": False} + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(profile), + "mode": "disaggregated", + "systems_paths": [timing_systems], + "workers": {"prefill": worker, "decode": worker}, + }, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 2}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + }, + } + ) + roles = [] + original = memory.estimate_num_gpu_blocks + + def observe(*args, **kwargs): + roles.append(kwargs.get("worker_type")) + return original(*args, **kwargs) + + monkeypatch.setattr(memory, "estimate_num_gpu_blocks", observe) + result = EngineReplayRunnerFactory().create(0).run(prediction_to_replay_spec(config)) + assert result.metrics["completed_requests"] == 1 + assert "prefill" in roles and "decode" in roles + assert set(roles) == {"prefill", "decode"} + + @pytest.mark.parametrize("mode,budget,expected", [("agg", 2, 2), ("disagg", 4, 4)]) def test_profile_candidates_use_only_declared_topologies(profile, mode, budget, expected, monkeypatch): from aisimulate_core.sdk import RustForwardPassPerfModel @@ -582,6 +856,134 @@ def test_profile_recommendation_cli_accepts_version_mapping_and_exports_literal( assert CorePredictionConfig.model_validate(exported).engine.fpm_profile.model == profile["model"] +def _observed_profile(profile, *, context=2048, fraction=0.9): + for deployment in profile["deployments"]: + resources = deployment["resources"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(name) + resources.update(max_num_tokens=1024, max_batch_size=4) + resources["runtime_memory"] = { + "kv_cache_bytes": 1024**3, + "max_model_len": context, + "gpu_memory_utilization": fraction, + "provenance": "Synthetic initialized cache pool; no GPU qualification.", + } + + +def test_public_linear_runtime_kv_load_preserves_selected_context(profile, timing_systems, tmp_path, monkeypatch): + import aisimulate.main as cli + from aisimulate.recommend import _run_recommendation + + _observed_profile(profile) + for deployment in profile["deployments"]: + deployment["resources"]["runtime_memory"]["kv_cache_bytes"] = 1024**2 + engine = _engine(profile) + engine.update(mode="aggregated", systems_paths=[timing_systems], context_length=1024) + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + engine["workers"]["aggregated"]["kv_cache"] = { + "prefix_caching": False, + "block_size": 64, + "capacity": {"type": "default", "memory_fraction": 0.9}, + } + path = tmp_path / "recommend.json" + path.write_text( + json.dumps( + { + "engine": engine, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 1}, + "load": {"type": "kv_capacity_fraction", "fraction": 0.0001}, + "stop": {"requests": 1}, + }, + "optimization": {"constraints": {"max_candidate_gpus": 2}}, + "optimizer": {"algorithm": "random", "max_trials": 1, "parallelism": 1}, + } + ) + ) + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + output = tmp_path / "recommendation" + assert cli.main(["recommend", "--config", str(path), "--output-dir", str(output)]) == 0 + exported = CorePredictionConfig.from_yaml(output / "recommendations/0001.yaml") + assert exported.engine.context_length == 1024 + assert all(item.resources.runtime_memory.max_model_len == 2048 for item in exported.engine.fpm_profile.deployments) + + +def test_observed_linear_pool_exposes_missing_batch_timing(profile, timing_systems): + from aisimulate.runner import EngineReplayRunnerFactory + + _observed_profile(profile, context=4096) + engine = {**_engine(profile), "systems_paths": [timing_systems]} + engine["workers"]["aggregated"]["kv_cache"] = {"prefix_caching": False} + config = CorePredictionConfig.model_validate( + { + "engine": engine, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 1}, + "load": {"type": "concurrency", "concurrency": 2}, + "stop": {"requests": 2}, + }, + } + ) + spec = prediction_to_replay_spec(config) + # 1 GiB / 512 bytes per token / 64 tokens per block = 32768 blocks. + assert materialize_aic_num_gpu_blocks(spec.backend_deployment.agg_engine_args)["num_gpu_blocks"] == 32768 + # The physical pool admits both requests. The table only measured batch=1, + # so strict lookup must fail instead of serializing through a clamped pool. + with pytest.raises(RuntimeError, match="direct.*lookup|uncovered|not.*found|unsupported"): + EngineReplayRunnerFactory().create(0).run(spec) + + +@pytest.mark.parametrize("override", [(), ("--context-length", "256"), ("--gpu-memory-utilization", "0.9")]) +def test_runtime_profile_reinitialization_reuses_observed_defaults(profile, tmp_path, override): + import aisimulate.main as cli + + _observed_profile(profile, context=128, fraction=0.8) + source, target = tmp_path / "profile.json", tmp_path / "request.yaml" + source.write_text(json.dumps(profile)) + command = [ + "onboard", + "init", + "--model", + profile["model"], + "--model-revision", + profile["model_revision"], + "--model-kind", + "moe", + "--framework-version", + "0.25.1", + "--gpu", + "h200_sxm", + "--interconnect", + "NVLink", + "--fpm-profile", + str(source), + "--tensor-parallel", + "2", + "--input-tokens", + "1", + "--output-tokens", + "1", + "--output", + str(target), + *override, + ] + if override: + with pytest.raises(SystemExit) as error: + cli.main(command) + assert error.value.code == 2 + assert not target.exists() + else: + assert cli.main(command) == 0 + request = SupportRequest.from_yaml(target) + assert request.search.context_length == 128 + assert request.collection.gpu_memory_utilization == 0.8 + assert "observed runtime" in request.collection_settings()["sources"]["context_length"] + assert ( + "recorded runtime memory fraction 0.8" in request.collection_settings()["sources"]["gpu_memory_utilization"] + ) + assert cli.main(["onboard", "plan", "--config", str(target), "--output-dir", str(tmp_path / "plan")]) == 0 + + @pytest.mark.parametrize( "update,error", [ @@ -612,8 +1014,6 @@ def no_timing(*args, **kwargs): "model_kind": "moe", "framework_version": "0.25.1", "gpu": "h200_sxm", - "gpu_count": 2, - "gpus_per_node": 2, "interconnect": "NVLink", }, "search": { @@ -627,8 +1027,8 @@ def no_timing(*args, **kwargs): } ) plan = create_plan(request, tmp_path) - assert plan["fpm"]["worker_gpus"] == 2 - assert plan["search"]["candidates"][0]["total_gpus"] == 2 + assert plan["fpm"]["collection_gpus_required"] == 2 + assert plan["search"]["candidates"][0]["required_gpus"] == 2 assert plan["resources"]["source"] == "profile" assert plan["resources"]["total_kv_size_tokens"] > 4096 loaded = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") @@ -645,8 +1045,10 @@ def no_timing(*args, **kwargs): tampered = json.loads(profile_path.read_text()) tampered["deployments"][0]["fmha_quant_mode"] = "fp8" profile_path.write_text(json.dumps(tampered)) - with pytest.raises(ValueError, match="differs from the requested identity"): + before = {p.relative_to(tmp_path): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + with pytest.raises(ValueError, match="fpm-model-profile[.]json was modified or missing"): check_plan(request, tmp_path) + assert {p.relative_to(tmp_path): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} == before @pytest.mark.parametrize( @@ -694,8 +1096,6 @@ def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( "model_kind": "moe", "framework_version": "0.25.1", "gpu": "b200_sxm", - "gpu_count": 8, - "gpus_per_node": 8, "interconnect": "NVLink", }, "search": { @@ -717,10 +1117,14 @@ def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( assert collected["counts"]["cells"] == 2 assert {cell["parallel_strategy"] for cell in collected["cells"]} == {preset} assert collected["options"]["moe_backend"] == moe_backend + assert collected["options"]["vllm_max_model_len"] == 8192 + assert collected["options"]["max_num_batched_tokens"] == max_tokens + assert collected["options"]["max_num_seqs"] == max_sequences sampling = collected["options"]["prefill_sampling"] - expected_tokens = min(max_tokens, 1024 * max_sequences) - assert sampling["max_total_prefill_tokens"] == expected_tokens + assert sampling["max_total_prefill_tokens"] == max_tokens assert sampling["max_batch_size"] == max_sequences + assert collected["point_generation"]["method"] == "native_self_benchmark" + assert collected["point_generation"]["planned_point_count"] is None assert collected["topology_memory_admission"][0]["activation_envelope"]["scope"] == "rank_local" prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") scheduler = prediction.engine.workers.aggregated.scheduler @@ -728,11 +1132,28 @@ def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( assert scheduler.max_sequences == max_sequences assert prediction.traffic.load.concurrency == 8 assert prediction.traffic.stop.requests == 8 - assert "timing coverage" in plan["fpm"]["sampling"] + + updated_input = request.model_dump(mode="json") + updated_input["workload"].update(input_tokens=64, output_tokens=16, concurrency=2, request_count=4) + updated_request = SupportRequest.model_validate(updated_input) + updated_plan = create_plan(updated_request, tmp_path, overwrite=True) + assert updated_plan["request_id"] == plan["request_id"] + assert updated_plan["validation_id"] != plan["validation_id"] + assert updated_plan["fpm"] == plan["fpm"] + assert updated_plan["resources"] == plan["resources"] + updated_command = json.loads((tmp_path / "commands.json").read_text())["fpm_plan_local"] + assert updated_command == command + assert collector_cli.main(updated_command[3:]) == 0 + assert json.loads(capsys.readouterr().out) == collected + updated_prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + assert updated_prediction.traffic.source.input_tokens == 64 + assert updated_prediction.traffic.source.output_tokens == 16 + assert updated_prediction.traffic.load.concurrency == 2 + assert updated_prediction.traffic.stop.requests == 4 def test_onboard_rejects_profile_below_collector_minimum_token_axis(profile, tmp_path): - profile["deployments"][0]["resources"]["max_num_tokens"] = 1 + profile["deployments"][0]["resources"].update(max_num_tokens=1, max_batch_size=1) request = SupportRequest.model_validate( { "identity": { @@ -741,16 +1162,16 @@ def test_onboard_rejects_profile_below_collector_minimum_token_axis(profile, tmp "model_kind": "moe", "framework_version": "0.25.1", "gpu": "h200_sxm", - "gpu_count": 2, - "gpus_per_node": 2, "interconnect": "NVLink", }, "search": {"tensor_parallel": 2, "context_length": 4096}, "fpm_profile": profile, } ) + output = tmp_path / "plan" with pytest.raises(ValueError, match="FPM collection requires.*at least 2"): - create_plan(request, tmp_path) + create_plan(request, output) + assert not output.exists() @pytest.mark.parametrize( @@ -778,8 +1199,6 @@ def test_onboard_cli_embeds_profile_and_preserves_topology(profile, tmp_path, tp "0.25.1", "--gpu", "h200_sxm", - "--gpu-count", - "2", "--interconnect", "NVLink", "--context-length", diff --git a/tests/test_fpm_query_evidence.py b/tests/test_fpm_query_evidence.py index 5977dfc63..f13b06475 100644 --- a/tests/test_fpm_query_evidence.py +++ b/tests/test_fpm_query_evidence.py @@ -101,6 +101,7 @@ def decode(batch=3, kv=12): return {"scheduled_requests": {"num_decode_requests": batch, "sum_decode_kv_tokens": kv}} +@pytest.mark.parametrize("collect_coverage", [False, True]) @pytest.mark.parametrize( "metrics,latency,resolution,weights,measurements", [ @@ -125,9 +126,11 @@ def decode(batch=3, kv=12): ], ) def test_canonical_detailed_query_support( - profile, evidence_systems, metrics, latency, resolution, weights, measurements + profile, evidence_systems, metrics, latency, resolution, weights, measurements, collect_coverage ): - model = RustForwardPassPerfModel.best_available(canonical(profile, evidence_systems)) + config = canonical(profile, evidence_systems) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = collect_coverage + model = RustForwardPassPerfModel.best_available(config) detailed = model.estimate_forward_pass_detailed(metrics) assert detailed["latency_ms"] == model.estimate_forward_pass_time_ms(metrics) == latency assert detailed["native_latency_ms"] == latency @@ -148,6 +151,40 @@ def test_canonical_detailed_query_support( for point in evidence["support"] ] == measurements assert json.loads(json.dumps(detailed)) == detailed + coverage = model.fpm_query_coverage() + if collect_coverage: + assert coverage["queries"] == { + "measured": 2 if resolution == "exact_lookup" else 0, + "interpolated": 0 if resolution == "exact_lookup" else 2, + "unsupported": 0, + } + else: + assert coverage is None + + +def test_detailed_coverage_retains_mixed_baseline_and_unsupported_queries(profile, evidence_systems): + config = canonical(profile, evidence_systems) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + model = RustForwardPassPerfModel.best_available(config) + mixed = {"scheduled_requests": {**prefill()["scheduled_requests"], **decode()["scheduled_requests"]}} + # The measured 10 ms prefill plus interpolated decode 23 ms minus its + # 15 ms interpolated floor is 18 ms. All three native lookups count once. + detailed = model.estimate_forward_pass_detailed(mixed) + assert detailed["latency_ms"] == 18.0 + assert len(detailed["ranks"][0]["queries"]) == 3 + coverage = model.fpm_query_coverage() + assert coverage["queries"] == {"measured": 1, "interpolated": 2, "unsupported": 0} + assert coverage["mixed_decode_baseline"] == {"measured": 0, "interpolated": 1, "unsupported": 0} + with pytest.raises(Exception, match="unsupported|outside"): + model.estimate_forward_pass_detailed(prefill(20)) + assert model.estimate_forward_pass_detailed({})["latency_ms"] == 0.0 + coverage = model.fpm_query_coverage() + assert coverage["queries"] == {"measured": 1, "interpolated": 2, "unsupported": 1} + assert coverage["gaps"][0]["coordinates"] == { + "batch_size": 1.0, + "total_prefill_tokens": 20.0, + "total_kv_read_tokens": 0.0, + } def test_mixed_baseline_rank_max_correction_and_clamp(profile, evidence_systems): @@ -230,20 +267,34 @@ def records(evidence): ] -def test_native_prediction_replay_and_source_detail_roundtrip(profile, evidence_systems): +@pytest.mark.parametrize("collect_coverage", [False, True]) +def test_native_prediction_replay_and_source_detail_roundtrip(profile, evidence_systems, collect_coverage): import jsonschema - config = CorePredictionConfig.model_validate(prediction(profile, evidence_systems)) + raw = prediction(profile, evidence_systems) + raw["engine"]["workers"]["aggregated"]["timing"]["estimator_config"]["fpm_interpolation"]["collect_coverage"] = ( + collect_coverage + ) + config = CorePredictionConfig.model_validate(raw) # Round-trip the saved public config before lowering it to the real runner. config = CorePredictionConfig.model_validate_json(config.model_dump_json()) spec = prediction_to_replay_spec(config) runner = EngineReplayRunnerFactory().create(0) try: - plain = runner.run(spec) + plain = runner.run(spec, output_requirements=ReplayOutputRequirements(include_raw_report=True)) report = runner.run(spec, output_requirements=ReplayOutputRequirements(capture_performance_diagnostics=True)) finally: runner.close() assert "fpm_query_evidence" not in plain.metadata + for replay in (plain, report): + native = replay.metadata["native_report"] + if collect_coverage: + assert native["fpm_query_coverage"]["status"] == "covered" + # Each request has one measured prefill, one measured decode and one + # interpolated decode; its first output token is produced by prefill. + assert native["fpm_query_coverage"]["queries"] == {"measured": 6, "interpolated": 3, "unsupported": 0} + else: + assert "fpm_query_coverage" not in native wall_clock_metrics = {"wall_time_ms", "processed_tokens_per_s", "processed_output_tokens_per_s"} assert {k: v for k, v in report.metrics.items() if k not in wall_clock_metrics} == { k: v for k, v in plain.metrics.items() if k not in wall_clock_metrics