Skip to main content

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}