aisimulate_core/perfmodel/fpm/model.rs
1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Forward-pass perf model: native AIC estimate with optional online correction
5//! and regression fallback, plus readiness/diagnostics.
6//!
7//! The `Native` variant holds an `Arc<Engine>` and the native estimate routes
8//! through [`crate::perfmodel::engine::Engine::forward_pass_time_ms`]. The online
9//! correction / regression / diagnostics / readiness logic is engine-agnostic.
10
11#[cfg(feature = "python")]
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14
15use serde::{Deserialize, Serialize};
16
17#[cfg(feature = "python")]
18use crate::perfmodel::EngineConfig;
19use crate::perfmodel::engine::Engine;
20use crate::{AicError, ForwardPassMetrics};
21
22use super::correction::CorrectionBuckets;
23use super::metrics::validate_forward_pass_metrics;
24use super::options::{ForwardPassPerfOptions, validate_options};
25use super::regression::BucketedRegression;
26use super::samples::{AxisRange, StoreStats, WithOptions};
27
28/// Current readiness and tuning state for a `ForwardPassPerfModel`.
29#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
30pub struct ForwardPassPerfDiagnostics {
31 /// Active prediction source. Native models become `aic_with_correction`
32 /// after at least one inferred workload kind has enough correction samples.
33 pub source: ForwardPassPerfSource,
34 /// Whether the model can currently produce learned estimates for at least
35 /// one workload kind, or why it cannot.
36 pub readiness: ForwardPassPerfReadiness,
37 /// Number of retained tuning observations across all inferred workload kinds.
38 pub retained_observations: usize,
39 /// Number of populated native-correction regions whose workload kind has at least
40 /// `min_observations` total retained samples.
41 pub correction_ready_buckets: usize,
42 /// Fallback reason when `best_available` had to use regression instead of native AIC.
43 pub last_warning: Option<String>,
44}
45
46/// Prediction backend currently used by `ForwardPassPerfModel`.
47#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
48#[serde(rename_all = "snake_case")]
49pub enum ForwardPassPerfSource {
50 /// Strict native AIC estimator with no correction workload kind ready yet.
51 Aic,
52 /// Workload-specific regression fallback, used without native AIC support.
53 FallbackRegression,
54 /// Native AIC estimator with at least one learned correction workload kind.
55 AicWithCorrection,
56}
57
58/// Readiness state reported by `ForwardPassPerfDiagnostics`.
59#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
60#[serde(rename_all = "snake_case")]
61pub enum ForwardPassPerfReadiness {
62 /// The model has either native AIC support or enough learned data.
63 Ready,
64 /// Regression fallback exists, but does not yet have enough observations.
65 InsufficientData,
66 /// Native AIC was unavailable and `best_available` fell back to regression.
67 UnsupportedConfig,
68 /// Reserved for callers that surface rejected FPM input as diagnostics.
69 InvalidInput,
70}
71
72/// Forward-pass-level performance model with optional online tuning.
73///
74/// This API intentionally stays at AIC's forward-pass abstraction. It does not
75/// model TTFT, ITL, SLA, engine capacity, queueing policy, or Dynamo engine
76/// limits. Callers pass FPMs for one engine iteration and receive one
77/// forward-pass latency estimate in milliseconds.
78///
79/// The prefill/decode/mixed workload kind is inferred from each iteration's
80/// `scheduled_requests` fields; it is not chosen at construction:
81///
82/// - prefill: scheduled prefill tokens and no scheduled decode work, using
83/// `[sum_prefill_tokens]`
84/// - decode: scheduled decode work and no scheduled prefill tokens, using
85/// `[num_decode_requests, sum_decode_kv_tokens]`
86/// - mixed/agg: both scheduled prefill and decode work, using
87/// `[sum_prefill_tokens, sum_decode_kv_tokens]`
88/// - empty: no scheduled prefill or decode work, estimates `0.0` and is not
89/// used for tuning
90///
91/// Native correction grids use fixed constructor-time ranges from
92/// `ForwardPassPerfOptions`: `max_num_tokens` bounds `sum_prefill_tokens`,
93/// `max_batch_size` bounds `num_decode_requests`, and `max_kv_tokens` bounds
94/// `sum_decode_kv_tokens`. `min_faster_correction_factor` and
95/// `max_slower_correction_factor` place independent absolute bounds on learned
96/// native correction factors in each direction, defaulting to `0.5` and `2.0`.
97/// Callers may explicitly disable either bound.
98///
99/// Queued request fields are accepted for FPM schema parity but ignored by this
100/// forward-pass-level model. `estimate_forward_pass_time_ms` treats FPM as a
101/// workload descriptor: it uses scheduled workload fields and ignores
102/// `wall_time`. `tune_with_fpms` treats FPM as observed telemetry: it uses the
103/// same scheduled workload fields as features and uses positive `wall_time` as
104/// the observation target. For attention-DP configurations, the input for one
105/// iteration is one FPM per attention-DP rank; tuning merges that list into one
106/// observation by taking max-rank load features and max nonzero `wall_time`.
107#[derive(Clone, Debug)]
108pub struct ForwardPassPerfModel {
109 mode: ForwardPassPerfMode,
110 options: ForwardPassPerfOptions,
111 last_warning: Option<String>,
112}
113
114#[derive(Clone, Debug)]
115enum ForwardPassPerfMode {
116 Native {
117 /// Compiled engine. `Arc` so `ForwardPassPerfModel` stays `Clone`
118 /// (the `Engine` itself is not `Clone`); cheap clones share the loaded
119 /// op lists + perf-database tree.
120 engine: Arc<Engine>,
121 corrections: WorkloadStores<CorrectionBuckets>,
122 },
123 Regression {
124 regressions: WorkloadStores<BucketedRegression>,
125 },
126}
127
128impl ForwardPassPerfModel {
129 /// API:
130 /// `ForwardPassPerfModel::from_native(config, options) -> Result<Self, AicError>`
131 ///
132 /// Description: create a strict native AIC forward-pass model.
133 ///
134 /// Compiles `config` into an [`Engine`] by crossing into Python once
135 /// (mirroring [`crate::AicEngineBuilder`]): `compile_engine` walks the model
136 /// and returns bincoded spec bytes, then [`Engine::from_spec_bytes`] loads
137 /// the matching perf database. This constructor fails if `config` cannot be
138 /// compiled. Use `best_available` when unsupported native configs should
139 /// fall back to the learned regression model.
140 #[cfg(feature = "python")]
141 pub fn from_native(
142 config: EngineConfig,
143 options: ForwardPassPerfOptions,
144 ) -> Result<Self, AicError> {
145 validate_options(&options)?;
146 let engine = build_engine_via_python(&config, None)?;
147 Ok(Self::from_engine(Arc::new(engine), options))
148 }
149
150 /// API:
151 /// `ForwardPassPerfModel::from_native_with_roots(config, options, systems_root) -> Result<Self, AicError>`
152 ///
153 /// Description: create a strict native AIC forward-pass model with an
154 /// explicit `systems/` data root (forwarded to `compile_engine` and used to
155 /// load the perf database). Same tuning and failure behavior as
156 /// `from_native`.
157 #[cfg(feature = "python")]
158 pub fn from_native_with_roots(
159 config: EngineConfig,
160 options: ForwardPassPerfOptions,
161 systems_root: impl AsRef<Path>,
162 ) -> Result<Self, AicError> {
163 validate_options(&options)?;
164 let engine = build_engine_via_python(&config, Some(systems_root.as_ref()))?;
165 Ok(Self::from_engine(Arc::new(engine), options))
166 }
167
168 /// Internal: build a native model directly from an already-compiled
169 /// [`Engine`]. Holds the actual native-mode logic; the public `from_native`
170 /// constructors compile the `Engine` (crossing into Python) and call this.
171 /// Used by the `#[cfg(test)]` suite to construct a native model from a
172 /// hand-built fixture `Engine` without Python.
173 pub(crate) fn from_engine(engine: Arc<Engine>, options: ForwardPassPerfOptions) -> Self {
174 Self {
175 mode: ForwardPassPerfMode::Native {
176 engine,
177 corrections: WorkloadStores::with_options(&options),
178 },
179 options,
180 last_warning: None,
181 }
182 }
183
184 /// API:
185 /// `ForwardPassPerfModel::from_regression(options) -> Result<Self, AicError>`
186 ///
187 /// Description: create a regression-only forward-pass model.
188 ///
189 /// This mode is for native-AIC-unsupported models. It returns `None` from
190 /// `estimate_forward_pass_time_ms` for non-empty iterations until the
191 /// inferred workload kind has at least `options.min_observations` tuning samples.
192 /// Correction factor getters always return `None` in this mode.
193 pub fn from_regression(options: ForwardPassPerfOptions) -> Result<Self, AicError> {
194 validate_options(&options)?;
195 Ok(Self {
196 mode: ForwardPassPerfMode::Regression {
197 regressions: WorkloadStores::with_options(&options),
198 },
199 options,
200 last_warning: None,
201 })
202 }
203
204 /// API:
205 /// `ForwardPassPerfModel::best_available(config, options) -> Result<Self, AicError>`
206 ///
207 /// Description: create a native model when possible, otherwise fall back to
208 /// regression.
209 ///
210 /// Fallback reason is preserved in `diagnostics().last_warning`. The
211 /// resulting model still uses the same FPM workload-kind inference and
212 /// tuning input contract as `from_native` and `from_regression`.
213 #[cfg(feature = "python")]
214 pub fn best_available(
215 config: EngineConfig,
216 options: ForwardPassPerfOptions,
217 ) -> Result<Self, AicError> {
218 match Self::from_native(config, options.clone()) {
219 Ok(model) => Ok(model),
220 Err(err) if can_fallback_to_regression(&err) => {
221 Self::regression_with_warning(options, err)
222 }
223 Err(err) => Err(err),
224 }
225 }
226
227 /// API:
228 /// `ForwardPassPerfModel::best_available_with_roots(config, options, systems_root) -> Result<Self, AicError>`
229 ///
230 /// Description: create a `best_available` model with an explicit `systems/`
231 /// data root.
232 #[cfg(feature = "python")]
233 pub fn best_available_with_roots(
234 config: EngineConfig,
235 options: ForwardPassPerfOptions,
236 systems_root: impl AsRef<Path>,
237 ) -> Result<Self, AicError> {
238 match Self::from_native_with_roots(config, options.clone(), systems_root) {
239 Ok(model) => Ok(model),
240 Err(err) if can_fallback_to_regression(&err) => {
241 Self::regression_with_warning(options, err)
242 }
243 Err(err) => Err(err),
244 }
245 }
246
247 #[cfg(feature = "python")]
248 fn regression_with_warning(
249 options: ForwardPassPerfOptions,
250 err: AicError,
251 ) -> Result<Self, AicError> {
252 let mut model = Self::from_regression(options)?;
253 model.last_warning = Some(format!(
254 "native forward-pass estimator unavailable; using fallback regression: {err}"
255 ));
256 Ok(model)
257 }
258
259 /// API:
260 /// `model.estimate_forward_pass_time_ms(metrics_by_rank) -> Result<Option<f64>, AicError>`
261 ///
262 /// Description: estimate one forward-pass iteration in milliseconds.
263 ///
264 /// `metrics_by_rank` must contain the FPMs for a single engine iteration,
265 /// one entry per attention-DP rank. Single-rank callers pass a one-element
266 /// slice. The inferred workload kind uses only `scheduled_requests` as described on
267 /// `ForwardPassPerfModel`; queued fields and `wall_time` are ignored for
268 /// estimation.
269 ///
270 /// Native models return an AIC estimate immediately, multiplied by the
271 /// correction factor for the matching workload region. Correction factors
272 /// default to `1.0` for inferred workload kinds with fewer than
273 /// `min_observations` total samples, empty regions, and queries outside the
274 /// configured correction-grid workload ranges in
275 /// `ForwardPassPerfOptions`. Regression models return `Ok(None)` until the
276 /// matching inferred workload kind has enough tuning samples. Empty
277 /// scheduled work returns `Ok(Some(0.0))`.
278 ///
279 /// Pure Rust over the `Engine` — no Python re-entry.
280 pub fn estimate_forward_pass_time_ms(
281 &self,
282 metrics_by_rank: &[ForwardPassMetrics],
283 ) -> Result<Option<f64>, AicError> {
284 let feature = IterationFeatures::from_metrics(metrics_by_rank)?;
285 let Some(feature) = feature else {
286 return Ok(Some(0.0));
287 };
288
289 match &self.mode {
290 ForwardPassPerfMode::Native {
291 engine,
292 corrections,
293 } => {
294 let native = engine.forward_pass_time_ms(metrics_by_rank)?;
295 let corrected = native
296 * corrections
297 .store(feature.workload_kind)
298 .correction_factor_for(&feature.x);
299 Ok(Some(corrected))
300 }
301 ForwardPassPerfMode::Regression { regressions } => {
302 Ok(regressions.store(feature.workload_kind).predict(&feature.x))
303 }
304 }
305 }
306
307 /// API:
308 /// `model.tune_with_fpms(iterations) -> Result<(), AicError>`
309 ///
310 /// Description: tune the model from observed FPM iterations.
311 ///
312 /// The outer slice is a list of observed iterations. Each inner slice is
313 /// the per-attention-DP-rank FPM list for one iteration:
314 /// `[[iter0_rank0, iter0_rank1], [iter1_rank0, iter1_rank1]]`.
315 /// Single-rank callers still use one FPM per inner slice.
316 ///
317 /// For each non-empty iteration, this method infers the workload kind from
318 /// scheduled request fields, takes max-rank load features, and uses the max
319 /// finite positive `wall_time` across ranks as the observed latency target
320 /// in milliseconds. Iterations with no scheduled work or no positive
321 /// `wall_time` are ignored. Native models update the matching region's
322 /// median `observed_ms / native_ms` correction factor, with each ratio
323 /// bounded by `min_faster_correction_factor` and
324 /// `max_slower_correction_factor` when configured. Regions are used only
325 /// after their inferred workload kind has `min_observations` total samples;
326 /// empty regions keep the default factor `1.0`. Observations outside the
327 /// configured correction-grid workload ranges are ignored by native
328 /// correction models. Regression models learn a workload-specific linear
329 /// fit.
330 ///
331 /// Pure Rust over the `Engine` — no Python re-entry.
332 pub fn tune_with_fpms(
333 &mut self,
334 iterations: &[Vec<ForwardPassMetrics>],
335 ) -> Result<(), AicError> {
336 for metrics_by_rank in iterations {
337 let observation = IterationObservation::from_metrics(metrics_by_rank)?;
338 let Some(observation) = observation else {
339 continue;
340 };
341
342 match &mut self.mode {
343 ForwardPassPerfMode::Native {
344 engine,
345 corrections,
346 } => {
347 let native = engine.forward_pass_time_ms(metrics_by_rank)?;
348 corrections
349 .store_mut(observation.feature.workload_kind)
350 .add_observation(observation.feature.x, observation.wall_time_ms, native);
351 }
352 ForwardPassPerfMode::Regression { regressions } => {
353 regressions
354 .store_mut(observation.feature.workload_kind)
355 .add_observation(observation.feature.x, observation.wall_time_ms);
356 }
357 }
358 }
359 Ok(())
360 }
361
362 /// API:
363 /// `model.diagnostics() -> ForwardPassPerfDiagnostics`
364 ///
365 /// Description: return the current backend, readiness, retained sample
366 /// count, and fallback warning.
367 pub fn diagnostics(&self) -> ForwardPassPerfDiagnostics {
368 match &self.mode {
369 ForwardPassPerfMode::Native { corrections, .. } => {
370 let ready_buckets = corrections.ready_bucket_count();
371 ForwardPassPerfDiagnostics {
372 source: if ready_buckets > 0 {
373 ForwardPassPerfSource::AicWithCorrection
374 } else {
375 ForwardPassPerfSource::Aic
376 },
377 readiness: ForwardPassPerfReadiness::Ready,
378 retained_observations: corrections.observation_count(),
379 correction_ready_buckets: ready_buckets,
380 last_warning: self.last_warning.clone(),
381 }
382 }
383 ForwardPassPerfMode::Regression { regressions } => {
384 let ready = regressions.any_ready();
385 ForwardPassPerfDiagnostics {
386 source: ForwardPassPerfSource::FallbackRegression,
387 readiness: if ready {
388 ForwardPassPerfReadiness::Ready
389 } else if self.last_warning.is_some() {
390 ForwardPassPerfReadiness::UnsupportedConfig
391 } else {
392 ForwardPassPerfReadiness::InsufficientData
393 },
394 retained_observations: regressions.observation_count(),
395 correction_ready_buckets: 0,
396 last_warning: self.last_warning.clone(),
397 }
398 }
399 }
400 }
401
402 /// API:
403 /// `model.min_correction_factor() -> Option<f64>`
404 ///
405 /// Description: return the smallest ready native correction factor across
406 /// all workload kinds.
407 ///
408 /// Returns `None` before any native correction workload kind has enough samples.
409 /// Regression-only models also return `None`.
410 pub fn min_correction_factor(&self) -> Option<f64> {
411 self.correction_factors()
412 .into_iter()
413 .reduce(|a, b| a.min(b))
414 }
415
416 /// API:
417 /// `model.max_correction_factor() -> Option<f64>`
418 ///
419 /// Description: return the largest ready native correction factor across
420 /// all workload kinds.
421 ///
422 /// Returns `None` before any native correction workload kind has enough samples.
423 /// Regression-only models also return `None`.
424 pub fn max_correction_factor(&self) -> Option<f64> {
425 self.correction_factors()
426 .into_iter()
427 .reduce(|a, b| a.max(b))
428 }
429
430 /// API:
431 /// `model.avg_correction_factor() -> Option<f64>`
432 ///
433 /// Description: return the arithmetic mean of ready native correction
434 /// factors across all workload kinds.
435 ///
436 /// Returns `None` before any native correction workload kind has enough samples.
437 /// Regression-only models also return `None`.
438 pub fn avg_correction_factor(&self) -> Option<f64> {
439 let factors = self.correction_factors();
440 if factors.is_empty() {
441 None
442 } else {
443 Some(factors.iter().sum::<f64>() / factors.len() as f64)
444 }
445 }
446
447 /// API:
448 /// `model.options() -> &ForwardPassPerfOptions`
449 ///
450 /// Description: return the immutable tuning options used by this model.
451 pub fn options(&self) -> &ForwardPassPerfOptions {
452 &self.options
453 }
454
455 fn correction_factors(&self) -> Vec<f64> {
456 match &self.mode {
457 ForwardPassPerfMode::Native { corrections, .. } => corrections.correction_factors(),
458 ForwardPassPerfMode::Regression { .. } => Vec::new(),
459 }
460 }
461}
462
463/// Build a compiled [`Engine`] from an [`EngineConfig`] by crossing into Python
464/// once to run `aiconfigurator.sdk.engine.compile_engine`, then loading the
465/// matching perf database via [`Engine::from_spec_bytes`]. This is the internal
466/// `EngineConfig` counterpart to [`crate::AicEngineBuilder`] and maps its
467/// modular fields onto the flat `compile_engine` kwargs.
468///
469/// `systems_root` overrides the bundled `systems/` dir for BOTH the
470/// `compile_engine` call (`systems_path` kwarg) and the Rust-side perf-DB load.
471#[cfg(feature = "python")]
472fn build_engine_via_python(
473 config: &EngineConfig,
474 systems_root: Option<&Path>,
475) -> Result<Engine, AicError> {
476 // `compile_engine`'s `systems_path` kwarg: explicit override -> config's
477 // own `systems_path` -> None (Python resolves it).
478 let systems_path: Option<PathBuf> = systems_root
479 .map(PathBuf::from)
480 .or_else(|| config.systems_path.clone());
481 // A non-UTF-8 override path cannot be passed through the Python kwarg; fail
482 // loudly rather than silently dropping the override.
483 let systems_path_str = match systems_path.as_ref() {
484 Some(p) => Some(p.to_str().ok_or_else(|| {
485 AicError::InvalidEngineConfig(format!(
486 "systems_path is not valid UTF-8: {}",
487 p.display()
488 ))
489 })?),
490 None => None,
491 };
492
493 crate::py::compile_engine_to_engine(config, systems_path_str)
494}
495
496#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
497pub(crate) enum WorkloadKind {
498 Prefill,
499 Decode,
500 Mixed,
501}
502
503#[derive(Clone, Debug)]
504pub(crate) struct IterationFeatures {
505 pub(crate) workload_kind: WorkloadKind,
506 pub(crate) x: Vec<f64>,
507}
508
509impl IterationFeatures {
510 pub(crate) fn from_metrics(
511 metrics_by_rank: &[ForwardPassMetrics],
512 ) -> Result<Option<Self>, AicError> {
513 if metrics_by_rank.is_empty() {
514 return Err(AicError::InvalidForwardPassMetrics(
515 "at least one attention-DP rank metric is required".to_string(),
516 ));
517 }
518 for metrics in metrics_by_rank {
519 validate_forward_pass_metrics(metrics)?;
520 }
521
522 Ok(metrics_by_rank
523 .iter()
524 .filter_map(Self::from_single_rank)
525 .max_by(|left, right| {
526 left.load_score()
527 .partial_cmp(&right.load_score())
528 .unwrap_or(std::cmp::Ordering::Equal)
529 }))
530 }
531
532 fn from_single_rank(metrics: &ForwardPassMetrics) -> Option<Self> {
533 let scheduled = &metrics.scheduled_requests;
534 let has_prefill = scheduled.sum_prefill_tokens > 0;
535 let has_decode = scheduled.num_decode_requests > 0 || scheduled.sum_decode_kv_tokens > 0;
536 let feature = match (has_prefill, has_decode) {
537 (false, false) => return None,
538 (true, false) => Self {
539 workload_kind: WorkloadKind::Prefill,
540 x: vec![f64::from(scheduled.sum_prefill_tokens)],
541 },
542 (false, true) => Self {
543 workload_kind: WorkloadKind::Decode,
544 x: vec![
545 f64::from(scheduled.num_decode_requests),
546 f64::from(scheduled.sum_decode_kv_tokens),
547 ],
548 },
549 (true, true) => Self {
550 workload_kind: WorkloadKind::Mixed,
551 x: vec![
552 f64::from(scheduled.sum_prefill_tokens),
553 f64::from(scheduled.sum_decode_kv_tokens),
554 ],
555 },
556 };
557 Some(feature)
558 }
559
560 fn load_score(&self) -> f64 {
561 self.x.iter().sum()
562 }
563}
564
565#[derive(Clone, Debug)]
566pub(crate) struct IterationObservation {
567 pub(crate) feature: IterationFeatures,
568 pub(crate) wall_time_ms: f64,
569}
570
571impl IterationObservation {
572 pub(crate) fn from_metrics(
573 metrics_by_rank: &[ForwardPassMetrics],
574 ) -> Result<Option<Self>, AicError> {
575 let Some(feature) = IterationFeatures::from_metrics(metrics_by_rank)? else {
576 return Ok(None);
577 };
578 let wall_time = metrics_by_rank
579 .iter()
580 .map(|metrics| metrics.wall_time)
581 .filter(|wall_time| wall_time.is_finite() && *wall_time > 0.0)
582 .fold(0.0_f64, f64::max);
583 if wall_time <= 0.0 {
584 return Ok(None);
585 }
586 Ok(Some(Self {
587 feature,
588 wall_time_ms: wall_time * 1000.0,
589 }))
590 }
591}
592
593#[derive(Clone, Debug)]
594pub(crate) struct WorkloadStores<T> {
595 prefill: T,
596 decode: T,
597 mixed: T,
598}
599
600impl<T: WithOptions> WorkloadStores<T> {
601 fn with_options(options: &ForwardPassPerfOptions) -> Self {
602 Self {
603 prefill: T::with_options(options, &[AxisRange::from_zero_to(options.max_num_tokens)]),
604 decode: T::with_options(
605 options,
606 &[
607 AxisRange::from_zero_to(options.max_batch_size),
608 AxisRange::from_zero_to(options.max_kv_tokens),
609 ],
610 ),
611 mixed: T::with_options(
612 options,
613 &[
614 AxisRange::from_zero_to(options.max_num_tokens),
615 AxisRange::from_zero_to(options.max_kv_tokens),
616 ],
617 ),
618 }
619 }
620}
621
622impl<T: StoreStats> WorkloadStores<T> {
623 fn observation_count(&self) -> usize {
624 self.prefill.observation_count()
625 + self.decode.observation_count()
626 + self.mixed.observation_count()
627 }
628
629 fn any_ready(&self) -> bool {
630 self.prefill.is_ready() || self.decode.is_ready() || self.mixed.is_ready()
631 }
632}
633
634impl WorkloadStores<CorrectionBuckets> {
635 fn ready_bucket_count(&self) -> usize {
636 self.prefill.ready_bucket_count()
637 + self.decode.ready_bucket_count()
638 + self.mixed.ready_bucket_count()
639 }
640
641 fn correction_factors(&self) -> Vec<f64> {
642 let mut factors = self.prefill.correction_factors();
643 factors.extend(self.decode.correction_factors());
644 factors.extend(self.mixed.correction_factors());
645 factors
646 }
647}
648
649impl<T> WorkloadStores<T> {
650 fn store(&self, workload_kind: WorkloadKind) -> &T {
651 match workload_kind {
652 WorkloadKind::Prefill => &self.prefill,
653 WorkloadKind::Decode => &self.decode,
654 WorkloadKind::Mixed => &self.mixed,
655 }
656 }
657
658 fn store_mut(&mut self, workload_kind: WorkloadKind) -> &mut T {
659 match workload_kind {
660 WorkloadKind::Prefill => &mut self.prefill,
661 WorkloadKind::Decode => &mut self.decode,
662 WorkloadKind::Mixed => &mut self.mixed,
663 }
664 }
665}
666
667/// Decide whether `best_available` should fall back to regression instead of
668/// propagating `err`. Covers the unsupported-model / data-availability errors
669/// that mean "this model can't be served natively". A failed native build via
670/// Python `compile_engine` surfaces as [`AicError::UnsupportedModel`] (see
671/// `py::compile_engine_from_flat`), which is covered here.
672///
673/// [`AicError::InvalidEngineConfig`] is deliberately NOT fallback-safe: it is
674/// used for hard caller/config errors (e.g. a non-UTF-8 `systems_path`, invalid
675/// FPM options, a malformed spec). Those must surface rather than silently
676/// degrade `best_available` to regression mode.
677#[cfg(feature = "python")]
678fn can_fallback_to_regression(err: &AicError) -> bool {
679 matches!(
680 err,
681 AicError::UnsupportedModel(_)
682 | AicError::DataRoot(_)
683 | AicError::ModelConfig(_)
684 | AicError::PerfDatabase(_)
685 | AicError::Io { .. }
686 | AicError::Parquet { .. }
687 )
688}