#[cfg(feature = "python")]
use std::path::{Path, PathBuf};
use std::sync::Arc;
use serde::{Deserialize, Serialize};
#[cfg(feature = "python")]
use crate::perfmodel::EngineConfig;
use crate::perfmodel::engine::Engine;
use crate::{AicError, ForwardPassMetrics};
use super::correction::CorrectionBuckets;
use super::metrics::validate_forward_pass_metrics;
use super::options::{ForwardPassPerfOptions, validate_options};
use super::regression::BucketedRegression;
use super::samples::{AxisRange, StoreStats, WithOptions};
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ForwardPassPerfDiagnostics {
pub source: ForwardPassPerfSource,
pub readiness: ForwardPassPerfReadiness,
pub retained_observations: usize,
pub correction_ready_buckets: usize,
pub last_warning: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ForwardPassPerfSource {
Aic,
FallbackRegression,
AicWithCorrection,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ForwardPassPerfReadiness {
Ready,
InsufficientData,
UnsupportedConfig,
InvalidInput,
}
#[derive(Clone, Debug)]
pub struct ForwardPassPerfModel {
mode: ForwardPassPerfMode,
options: ForwardPassPerfOptions,
last_warning: Option<String>,
}
#[derive(Clone, Debug)]
enum ForwardPassPerfMode {
Native {
engine: Arc<Engine>,
corrections: WorkloadStores<CorrectionBuckets>,
},
Regression {
regressions: WorkloadStores<BucketedRegression>,
},
}
impl ForwardPassPerfModel {
#[cfg(feature = "python")]
pub fn from_native(
config: EngineConfig,
options: ForwardPassPerfOptions,
) -> Result<Self, AicError> {
validate_options(&options)?;
let engine = build_engine_via_python(&config, None)?;
Ok(Self::from_engine(Arc::new(engine), options))
}
#[cfg(feature = "python")]
pub fn from_native_with_roots(
config: EngineConfig,
options: ForwardPassPerfOptions,
systems_root: impl AsRef<Path>,
) -> Result<Self, AicError> {
validate_options(&options)?;
let engine = build_engine_via_python(&config, Some(systems_root.as_ref()))?;
Ok(Self::from_engine(Arc::new(engine), options))
}
pub(crate) fn from_engine(engine: Arc<Engine>, options: ForwardPassPerfOptions) -> Self {
Self {
mode: ForwardPassPerfMode::Native {
engine,
corrections: WorkloadStores::with_options(&options),
},
options,
last_warning: None,
}
}
pub fn from_regression(options: ForwardPassPerfOptions) -> Result<Self, AicError> {
validate_options(&options)?;
Ok(Self {
mode: ForwardPassPerfMode::Regression {
regressions: WorkloadStores::with_options(&options),
},
options,
last_warning: None,
})
}
#[cfg(feature = "python")]
pub fn best_available(
config: EngineConfig,
options: ForwardPassPerfOptions,
) -> Result<Self, AicError> {
match Self::from_native(config, options.clone()) {
Ok(model) => Ok(model),
Err(err) if can_fallback_to_regression(&err) => {
Self::regression_with_warning(options, err)
}
Err(err) => Err(err),
}
}
#[cfg(feature = "python")]
pub fn best_available_with_roots(
config: EngineConfig,
options: ForwardPassPerfOptions,
systems_root: impl AsRef<Path>,
) -> Result<Self, AicError> {
match Self::from_native_with_roots(config, options.clone(), systems_root) {
Ok(model) => Ok(model),
Err(err) if can_fallback_to_regression(&err) => {
Self::regression_with_warning(options, err)
}
Err(err) => Err(err),
}
}
#[cfg(feature = "python")]
fn regression_with_warning(
options: ForwardPassPerfOptions,
err: AicError,
) -> Result<Self, AicError> {
let mut model = Self::from_regression(options)?;
model.last_warning = Some(format!(
"native forward-pass estimator unavailable; using fallback regression: {err}"
));
Ok(model)
}
pub fn estimate_forward_pass_time_ms(
&self,
metrics_by_rank: &[ForwardPassMetrics],
) -> Result<Option<f64>, AicError> {
let feature = IterationFeatures::from_metrics(metrics_by_rank)?;
let Some(feature) = feature else {
return Ok(Some(0.0));
};
match &self.mode {
ForwardPassPerfMode::Native {
engine,
corrections,
} => {
let native = engine.forward_pass_time_ms(metrics_by_rank)?;
let corrected = native
* corrections
.store(feature.workload_kind)
.correction_factor_for(&feature.x);
Ok(Some(corrected))
}
ForwardPassPerfMode::Regression { regressions } => {
Ok(regressions.store(feature.workload_kind).predict(&feature.x))
}
}
}
pub fn tune_with_fpms(
&mut self,
iterations: &[Vec<ForwardPassMetrics>],
) -> Result<(), AicError> {
for metrics_by_rank in iterations {
let observation = IterationObservation::from_metrics(metrics_by_rank)?;
let Some(observation) = observation else {
continue;
};
match &mut self.mode {
ForwardPassPerfMode::Native {
engine,
corrections,
} => {
let native = engine.forward_pass_time_ms(metrics_by_rank)?;
corrections
.store_mut(observation.feature.workload_kind)
.add_observation(observation.feature.x, observation.wall_time_ms, native);
}
ForwardPassPerfMode::Regression { regressions } => {
regressions
.store_mut(observation.feature.workload_kind)
.add_observation(observation.feature.x, observation.wall_time_ms);
}
}
}
Ok(())
}
pub fn diagnostics(&self) -> ForwardPassPerfDiagnostics {
match &self.mode {
ForwardPassPerfMode::Native { corrections, .. } => {
let ready_buckets = corrections.ready_bucket_count();
ForwardPassPerfDiagnostics {
source: if ready_buckets > 0 {
ForwardPassPerfSource::AicWithCorrection
} else {
ForwardPassPerfSource::Aic
},
readiness: ForwardPassPerfReadiness::Ready,
retained_observations: corrections.observation_count(),
correction_ready_buckets: ready_buckets,
last_warning: self.last_warning.clone(),
}
}
ForwardPassPerfMode::Regression { regressions } => {
let ready = regressions.any_ready();
ForwardPassPerfDiagnostics {
source: ForwardPassPerfSource::FallbackRegression,
readiness: if ready {
ForwardPassPerfReadiness::Ready
} else if self.last_warning.is_some() {
ForwardPassPerfReadiness::UnsupportedConfig
} else {
ForwardPassPerfReadiness::InsufficientData
},
retained_observations: regressions.observation_count(),
correction_ready_buckets: 0,
last_warning: self.last_warning.clone(),
}
}
}
}
pub fn min_correction_factor(&self) -> Option<f64> {
self.correction_factors()
.into_iter()
.reduce(|a, b| a.min(b))
}
pub fn max_correction_factor(&self) -> Option<f64> {
self.correction_factors()
.into_iter()
.reduce(|a, b| a.max(b))
}
pub fn avg_correction_factor(&self) -> Option<f64> {
let factors = self.correction_factors();
if factors.is_empty() {
None
} else {
Some(factors.iter().sum::<f64>() / factors.len() as f64)
}
}
pub fn options(&self) -> &ForwardPassPerfOptions {
&self.options
}
fn correction_factors(&self) -> Vec<f64> {
match &self.mode {
ForwardPassPerfMode::Native { corrections, .. } => corrections.correction_factors(),
ForwardPassPerfMode::Regression { .. } => Vec::new(),
}
}
}
#[cfg(feature = "python")]
fn build_engine_via_python(
config: &EngineConfig,
systems_root: Option<&Path>,
) -> Result<Engine, AicError> {
let systems_path: Option<PathBuf> = systems_root
.map(PathBuf::from)
.or_else(|| config.systems_path.clone());
let systems_path_str = match systems_path.as_ref() {
Some(p) => Some(p.to_str().ok_or_else(|| {
AicError::InvalidEngineConfig(format!(
"systems_path is not valid UTF-8: {}",
p.display()
))
})?),
None => None,
};
crate::py::compile_engine_to_engine(config, systems_path_str)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum WorkloadKind {
Prefill,
Decode,
Mixed,
}
#[derive(Clone, Debug)]
pub(crate) struct IterationFeatures {
pub(crate) workload_kind: WorkloadKind,
pub(crate) x: Vec<f64>,
}
impl IterationFeatures {
pub(crate) fn from_metrics(
metrics_by_rank: &[ForwardPassMetrics],
) -> Result<Option<Self>, AicError> {
if metrics_by_rank.is_empty() {
return Err(AicError::InvalidForwardPassMetrics(
"at least one attention-DP rank metric is required".to_string(),
));
}
for metrics in metrics_by_rank {
validate_forward_pass_metrics(metrics)?;
}
Ok(metrics_by_rank
.iter()
.filter_map(Self::from_single_rank)
.max_by(|left, right| {
left.load_score()
.partial_cmp(&right.load_score())
.unwrap_or(std::cmp::Ordering::Equal)
}))
}
fn from_single_rank(metrics: &ForwardPassMetrics) -> Option<Self> {
let scheduled = &metrics.scheduled_requests;
let has_prefill = scheduled.sum_prefill_tokens > 0;
let has_decode = scheduled.num_decode_requests > 0 || scheduled.sum_decode_kv_tokens > 0;
let feature = match (has_prefill, has_decode) {
(false, false) => return None,
(true, false) => Self {
workload_kind: WorkloadKind::Prefill,
x: vec![f64::from(scheduled.sum_prefill_tokens)],
},
(false, true) => Self {
workload_kind: WorkloadKind::Decode,
x: vec![
f64::from(scheduled.num_decode_requests),
f64::from(scheduled.sum_decode_kv_tokens),
],
},
(true, true) => Self {
workload_kind: WorkloadKind::Mixed,
x: vec![
f64::from(scheduled.sum_prefill_tokens),
f64::from(scheduled.sum_decode_kv_tokens),
],
},
};
Some(feature)
}
fn load_score(&self) -> f64 {
self.x.iter().sum()
}
}
#[derive(Clone, Debug)]
pub(crate) struct IterationObservation {
pub(crate) feature: IterationFeatures,
pub(crate) wall_time_ms: f64,
}
impl IterationObservation {
pub(crate) fn from_metrics(
metrics_by_rank: &[ForwardPassMetrics],
) -> Result<Option<Self>, AicError> {
let Some(feature) = IterationFeatures::from_metrics(metrics_by_rank)? else {
return Ok(None);
};
let wall_time = metrics_by_rank
.iter()
.map(|metrics| metrics.wall_time)
.filter(|wall_time| wall_time.is_finite() && *wall_time > 0.0)
.fold(0.0_f64, f64::max);
if wall_time <= 0.0 {
return Ok(None);
}
Ok(Some(Self {
feature,
wall_time_ms: wall_time * 1000.0,
}))
}
}
#[derive(Clone, Debug)]
pub(crate) struct WorkloadStores<T> {
prefill: T,
decode: T,
mixed: T,
}
impl<T: WithOptions> WorkloadStores<T> {
fn with_options(options: &ForwardPassPerfOptions) -> Self {
Self {
prefill: T::with_options(options, &[AxisRange::from_zero_to(options.max_num_tokens)]),
decode: T::with_options(
options,
&[
AxisRange::from_zero_to(options.max_batch_size),
AxisRange::from_zero_to(options.max_kv_tokens),
],
),
mixed: T::with_options(
options,
&[
AxisRange::from_zero_to(options.max_num_tokens),
AxisRange::from_zero_to(options.max_kv_tokens),
],
),
}
}
}
impl<T: StoreStats> WorkloadStores<T> {
fn observation_count(&self) -> usize {
self.prefill.observation_count()
+ self.decode.observation_count()
+ self.mixed.observation_count()
}
fn any_ready(&self) -> bool {
self.prefill.is_ready() || self.decode.is_ready() || self.mixed.is_ready()
}
}
impl WorkloadStores<CorrectionBuckets> {
fn ready_bucket_count(&self) -> usize {
self.prefill.ready_bucket_count()
+ self.decode.ready_bucket_count()
+ self.mixed.ready_bucket_count()
}
fn correction_factors(&self) -> Vec<f64> {
let mut factors = self.prefill.correction_factors();
factors.extend(self.decode.correction_factors());
factors.extend(self.mixed.correction_factors());
factors
}
}
impl<T> WorkloadStores<T> {
fn store(&self, workload_kind: WorkloadKind) -> &T {
match workload_kind {
WorkloadKind::Prefill => &self.prefill,
WorkloadKind::Decode => &self.decode,
WorkloadKind::Mixed => &self.mixed,
}
}
fn store_mut(&mut self, workload_kind: WorkloadKind) -> &mut T {
match workload_kind {
WorkloadKind::Prefill => &mut self.prefill,
WorkloadKind::Decode => &mut self.decode,
WorkloadKind::Mixed => &mut self.mixed,
}
}
}
#[cfg(feature = "python")]
fn can_fallback_to_regression(err: &AicError) -> bool {
matches!(
err,
AicError::UnsupportedModel(_)
| AicError::DataRoot(_)
| AicError::ModelConfig(_)
| AicError::PerfDatabase(_)
| AicError::Io { .. }
| AicError::Parquet { .. }
)
}