Skip to main content

eredu_runtime/
execution_plan.rs

1//! Portable execution-plan policy derivation before backend resource realization.
2
3use eredu_core::{
4    residency::OffloadConfig, ExecutionPlan, ResidencyPlan, WeightTransformationPlan,
5};
6
7use crate::{
8    DenseDiskStreamLoadOptions, LayerwiseLoadOptions, NormalizedLoadRequest,
9    NormalizedLoadRequestError, OrdinaryWeightResidency, ParallelLoadRequest,
10    ParameterBankLoadOptions, WeightResidency,
11};
12
13/// Optional observations requested from the backend's residency mechanisms.
14///
15/// Deriving a request never samples memory or creates a native resource.
16#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
17pub struct ResidencyDiagnostics {
18    backend_memory: bool,
19    process_memory: bool,
20}
21
22impl ResidencyDiagnostics {
23    /// Selects backend allocator and process-memory observations independently.
24    pub const fn new(backend_memory: bool, process_memory: bool) -> Self {
25        Self {
26            backend_memory,
27            process_memory,
28        }
29    }
30}
31
32/// Portable failure deriving a load request; backend support is checked separately.
33#[derive(Debug, thiserror::Error)]
34#[non_exhaustive]
35pub enum ExecutionPlanLoadError {
36    /// The serialized or caller-built plan violates its portable contract.
37    #[error(transparent)]
38    Plan(#[from] eredu_core::ExecutionPlanError),
39    /// Complete rank, wire, invocation and completion policy is required.
40    #[error("single-device automatic plans require a 1x1x1 parallel topology; distributed plans require an explicit parallel load request")]
41    MissingParallelRequest,
42    /// The supplied parallel request does not belong to the plan's topology.
43    #[error("parallel load request topology differs from the execution plan")]
44    ParallelTopologyMismatch,
45    /// The plan selected policy not understood by this runtime.
46    #[error("unsupported execution-plan {0}")]
47    Unsupported(&'static str),
48    /// An invalid offload policy cannot become a load request.
49    #[error(transparent)]
50    Offload(#[from] eredu_core::residency::OffloadError),
51    /// Invalid ordinary or parameter-bank residency controls.
52    #[error(transparent)]
53    Residency(#[from] crate::WeightResidencyPolicyError),
54    /// Invalid quantization, drafting, or completion controls.
55    #[error(transparent)]
56    Request(#[from] NormalizedLoadRequestError),
57}
58
59/// Derives the exact load-time transformation without backend capability policy.
60pub fn execution_plan_quantization(
61    transformation: WeightTransformationPlan,
62) -> Result<Option<eredu_core::QuantizationRequest>, ExecutionPlanLoadError> {
63    use eredu_core::QuantizationRequest;
64    let request = match transformation {
65        WeightTransformationPlan::PreserveCheckpoint => return Ok(None),
66        WeightTransformationPlan::Affine { bits, group_size } => QuantizationRequest::Affine {
67            group_size: u32::try_from(group_size).map_err(|_| {
68                NormalizedLoadRequestError::Quantization(format!(
69                    "group_size must be non-negative, got {group_size}"
70                ))
71            })?,
72            bits: u8::try_from(bits).map_err(|_| {
73                NormalizedLoadRequestError::Quantization(format!("bits must fit in u8, got {bits}"))
74            })?,
75        },
76        WeightTransformationPlan::MxFp4 => QuantizationRequest::MxFp4,
77        _ => return Err(ExecutionPlanLoadError::Unsupported("weight transformation")),
78    };
79    NormalizedLoadRequest::with_quantization(request).weight_quantization()?;
80    Ok(Some(request))
81}
82
83impl NormalizedLoadRequest {
84    /// Derives all portable load policy from a plan and optional complete rank request.
85    ///
86    /// Distributed plans supply their exact rank, wire, invocation limits and completion
87    /// policy explicitly because an execution plan does not contain those values. Native
88    /// device identity and realized resources never enter this conversion.
89    pub fn from_execution_plan(
90        plan: &ExecutionPlan,
91        diagnostics: ResidencyDiagnostics,
92        parallel: Option<ParallelLoadRequest>,
93    ) -> Result<Self, ExecutionPlanLoadError> {
94        plan.validate_structure()?;
95        match parallel {
96            Some(parallel) if parallel.rank().topology() != *plan.topology() => {
97                return Err(ExecutionPlanLoadError::ParallelTopologyMismatch);
98            }
99            None if !plan.topology().is_replicated() => {
100                return Err(ExecutionPlanLoadError::MissingParallelRequest);
101            }
102            _ => {}
103        }
104        let ordinary = match plan.residency() {
105            ResidencyPlan::FullyResident => OrdinaryWeightResidency::FullyResident,
106            ResidencyPlan::LayerwiseHost {
107                device_layer_window,
108                device_budget_bytes,
109                host_budget_bytes,
110            } => OrdinaryWeightResidency::LayerwiseHost(
111                LayerwiseLoadOptions::new(OffloadConfig::new(
112                    *device_budget_bytes,
113                    *host_budget_bytes,
114                    *device_layer_window,
115                )?)
116                .with_max_cached_shards(plan.max_cached_shards())
117                .with_memory_sampling(diagnostics.backend_memory, diagnostics.process_memory),
118            ),
119            ResidencyPlan::DenseDiskStream {
120                device_budget_bytes,
121                host_budget_bytes,
122                host_lookahead,
123                background_queue,
124            } => OrdinaryWeightResidency::DenseDiskStream(
125                DenseDiskStreamLoadOptions::new(
126                    *device_budget_bytes,
127                    *host_budget_bytes,
128                    *host_lookahead,
129                    *background_queue,
130                )?
131                .with_max_cached_shards(plan.max_cached_shards())
132                .with_memory_sampling(diagnostics.backend_memory, diagnostics.process_memory),
133            ),
134            _ => return Err(ExecutionPlanLoadError::Unsupported("residency")),
135        };
136        let residency = match plan.expert_cache() {
137            Some(bank) => WeightResidency::with_independent_parameter_banks(
138                ordinary,
139                ParameterBankLoadOptions::new(
140                    OffloadConfig::new(bank.device_budget_bytes(), bank.host_budget_bytes(), 1)?
141                        .with_eviction_policy(bank.eviction_policy()),
142                    bank.scratch_bytes(),
143                    bank.prefill_bank_bytes(),
144                )?,
145            ),
146            None => WeightResidency::with_layers(ordinary.layers()),
147        };
148        let mut request = execution_plan_quantization(plan.weight_transformation())?
149            .map_or_else(Self::default, Self::with_quantization)
150            .with_weight_residency(residency)
151            .with_max_cached_shards(
152                std::num::NonZeroUsize::new(plan.max_cached_shards())
153                    .expect("plan structure validates a positive reader limit"),
154            )
155            .with_required_session_capabilities(*plan.required_session_capabilities())
156            .with_prompt_cache_persistence(plan.prompt_cache_persistence())
157            .with_drafting_plan(plan.drafting())?;
158        if let Some(parallel) = parallel {
159            request = request.with_parallel_execution(parallel)?;
160        }
161        request.validate_model_preparation()?;
162        Ok(request)
163    }
164}
165
166#[cfg(test)]
167mod tests;