eredu_runtime/
load_request.rs

1//! Normalized portable policy for cold model preparation.
2
3use std::num::NonZeroUsize;
4
5use eredu_checkpoint::{AffineQuantization, WeightQuantization};
6use eredu_core::{
7    CompletionCancellationMode, DraftingPlan, ParallelRankTopology, PreparationPolicy,
8    QuantizationRequest, ResidencyRequest, SessionCapabilities,
9};
10
11use crate::{
12    CacheResidencyPolicy, CommunicationCompletionPolicy, LayerWeightResidency,
13    PipelineWireContract, WeightResidency,
14};
15
16/// Portable drafting intent resolved before checkpoint payload selection.
17#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
18pub enum DraftingLoadRequest {
19    /// Preserve direct-load behavior by installing an admitted embedded extension.
20    #[default]
21    ArchitectureDefault,
22    /// Materialize only the ordinary target.
23    Disabled,
24    /// Require and install an embedded extension with this proposal bound.
25    Embedded {
26        /// Maximum number of draft tokens proposed per target step.
27        max_draft_tokens: NonZeroUsize,
28    },
29    /// Prepare only the ordinary target for later external-assistant pairing.
30    ExternalTarget,
31}
32
33impl DraftingLoadRequest {
34    /// Creates an embedded drafting request with a positive proposal capacity.
35    pub fn embedded(max_draft_tokens: usize) -> Result<Self, NormalizedLoadRequestError> {
36        let max_draft_tokens = NonZeroUsize::new(max_draft_tokens)
37            .ok_or(NormalizedLoadRequestError::ZeroEmbeddedDraftCapacity)?;
38        Ok(Self::Embedded { max_draft_tokens })
39    }
40
41    /// Returns the selected embedded proposal capacity, when requested explicitly.
42    pub const fn embedded_capacity(self) -> Option<NonZeroUsize> {
43        match self {
44            Self::Embedded { max_draft_tokens } => Some(max_draft_tokens),
45            Self::ArchitectureDefault | Self::Disabled | Self::ExternalTarget => None,
46        }
47    }
48}
49
50/// Portable topology, wire, and invocation geometry selected as one atomic request.
51#[derive(Debug, Clone, Copy, Eq, PartialEq)]
52pub struct ParallelLoadRequest {
53    rank: ParallelRankTopology,
54    wire: PipelineWireContract,
55    maximum_batch_size: i32,
56    maximum_sequence_length: i32,
57    completion: CommunicationCompletionPolicy,
58}
59
60impl ParallelLoadRequest {
61    /// Binds a rank, activation wire, and maximum invocation geometry.
62    pub fn new(
63        rank: ParallelRankTopology,
64        wire: PipelineWireContract,
65        maximum_batch_size: i32,
66        maximum_sequence_length: i32,
67        completion: CommunicationCompletionPolicy,
68    ) -> Result<Self, NormalizedLoadRequestError> {
69        if rank.is_replicated() {
70            return Err(NormalizedLoadRequestError::ReplicatedParallelTopology);
71        }
72        if maximum_batch_size <= 0 || maximum_sequence_length <= 0 {
73            return Err(NormalizedLoadRequestError::InvalidInvocationLimits {
74                maximum_batch_size,
75                maximum_sequence_length,
76            });
77        }
78        Ok(Self {
79            rank,
80            wire,
81            maximum_batch_size,
82            maximum_sequence_length,
83            completion,
84        })
85    }
86
87    /// Returns the selected portable rank.
88    pub const fn rank(self) -> ParallelRankTopology {
89        self.rank
90    }
91
92    /// Returns the exact pipeline activation wire contract.
93    pub const fn wire(self) -> PipelineWireContract {
94        self.wire
95    }
96
97    /// Returns the maximum admitted batch and sequence geometry.
98    pub const fn invocation_limits(self) -> (i32, i32) {
99        (self.maximum_batch_size, self.maximum_sequence_length)
100    }
101
102    /// Returns the bounded completion policy selected for this parallel request.
103    pub const fn completion(self) -> CommunicationCompletionPolicy {
104        self.completion
105    }
106
107    const fn with_completion(mut self, completion: CommunicationCompletionPolicy) -> Self {
108        self.completion = completion;
109        self
110    }
111}
112
113/// Failure while validating or lowering a normalized cold-load request.
114#[derive(Debug, thiserror::Error)]
115#[non_exhaustive]
116pub enum NormalizedLoadRequestError {
117    /// A prepared source must have a positive reader-cache bound.
118    #[error("source reader-cache limit must be positive")]
119    ZeroCachedShards,
120    /// Parallel invocation limits must both be positive.
121    #[error(
122        "partitioned invocation limits must be positive, got batch {maximum_batch_size} and sequence {maximum_sequence_length}"
123    )]
124    InvalidInvocationLimits {
125        /// Maximum admitted batch size.
126        maximum_batch_size: i32,
127        /// Maximum admitted sequence length.
128        maximum_sequence_length: i32,
129    },
130    /// A replicated topology was incorrectly attached as parallel execution.
131    #[error("explicit parallel execution requires a non-replicated topology")]
132    ReplicatedParallelTopology,
133    /// A local completion setting was attached before a complete parallel request.
134    #[error("parallel execution cannot replace an existing local completion policy")]
135    ConflictingCompletionPolicy,
136    /// An ordinary model request retained a communication-only policy.
137    #[error("a communication completion policy requires a parallel topology")]
138    OrphanedModelCompletion,
139    /// Embedded drafting requires positive capacity.
140    #[error("embedded draft capacity must be positive")]
141    ZeroEmbeddedDraftCapacity,
142    /// The portable planner supplied a drafting mode this runtime does not understand.
143    #[error("unsupported speculative drafting plan")]
144    UnsupportedDraftingPlan,
145    /// The default realtime completion contract could not be constructed.
146    #[error("{0}")]
147    Completion(String),
148    /// A quantization request could not be represented by the checkpoint contract.
149    #[error("{0}")]
150    Quantization(String),
151}
152
153/// One normalized backend-neutral request for cold model preparation.
154///
155/// Parallel rank, wire, and invocation fields are represented atomically. A concrete backend
156/// may pair this value with a separate native device or resource token, but must not add native
157/// identity to this request.
158#[derive(Debug, Clone, Default, Eq, PartialEq)]
159pub struct NormalizedLoadRequest {
160    quantization: Option<QuantizationRequest>,
161    max_cached_shards: Option<NonZeroUsize>,
162    parallel: Option<ParallelLoadRequest>,
163    communication_completion: Option<CommunicationCompletionPolicy>,
164    weight_residency: WeightResidency,
165    state_residency: CacheResidencyPolicy,
166    required_session_capabilities: SessionCapabilities,
167    prompt_cache_persistence: bool,
168    drafting: DraftingLoadRequest,
169}
170
171/// A normalized model request whose cross-field invariants were checked once.
172#[derive(Debug, Clone, Copy)]
173pub struct ValidatedModelLoadRequest<'a> {
174    request: &'a NormalizedLoadRequest,
175    policy: PreparationPolicy,
176}
177
178impl<'a> ValidatedModelLoadRequest<'a> {
179    /// Exact portable preparation policy projected during validation.
180    pub const fn preparation_policy(self) -> PreparationPolicy {
181        self.policy
182    }
183
184    /// Original normalized request covered by this validation proof.
185    pub const fn request(self) -> &'a NormalizedLoadRequest {
186        self.request
187    }
188}
189
190impl NormalizedLoadRequest {
191    /// Requires persisted prompt-prefix import/export independently of state residency.
192    pub const fn with_prompt_cache_persistence(mut self, required: bool) -> Self {
193        self.prompt_cache_persistence = required;
194        self
195    }
196
197    /// Returns explicit persisted prompt-prefix import/export intent.
198    pub const fn prompt_cache_persistence(&self) -> bool {
199        self.prompt_cache_persistence
200    }
201
202    /// Selects the exact reader-cache limit, including fully resident execution.
203    pub const fn with_max_cached_shards(mut self, maximum: NonZeroUsize) -> Self {
204        self.max_cached_shards = Some(maximum);
205        self
206    }
207
208    /// Returns the source reader-cache bound retained by cold selection.
209    pub const fn max_cached_shards(&self) -> usize {
210        match self.max_cached_shards {
211            Some(maximum) => maximum.get(),
212            None => self.weight_residency.max_cached_shards(),
213        }
214    }
215
216    /// Creates a request that quantizes eligible dense weights on load.
217    pub fn with_quantization(quantization: QuantizationRequest) -> Self {
218        Self {
219            quantization: Some(quantization),
220            ..Self::default()
221        }
222    }
223
224    /// Attaches one complete portable parallel-execution request.
225    pub fn with_parallel_execution(
226        mut self,
227        parallel: ParallelLoadRequest,
228    ) -> Result<Self, NormalizedLoadRequestError> {
229        if self.communication_completion.is_some() {
230            return Err(NormalizedLoadRequestError::ConflictingCompletionPolicy);
231        }
232        self.parallel = Some(parallel);
233        Ok(self)
234    }
235
236    /// Selects the bounded completion policy used by communication or local realtime work.
237    pub const fn with_communication_completion_policy(
238        mut self,
239        policy: CommunicationCompletionPolicy,
240    ) -> Self {
241        self.set_communication_completion_policy(policy);
242        self
243    }
244
245    /// Replaces the bounded completion policy without changing another request field.
246    pub const fn set_communication_completion_policy(
247        &mut self,
248        policy: CommunicationCompletionPolicy,
249    ) {
250        match self.parallel {
251            Some(parallel) => self.parallel = Some(parallel.with_completion(policy)),
252            None => self.communication_completion = Some(policy),
253        }
254    }
255
256    /// Selects fully resident or bounded checkpoint-weight execution.
257    pub fn with_weight_residency(mut self, residency: WeightResidency) -> Self {
258        self.weight_residency = residency;
259        self
260    }
261
262    /// Selects mutable-state residency and paging controls.
263    pub fn with_state_residency(mut self, residency: CacheResidencyPolicy) -> Self {
264        self.state_residency = residency;
265        self
266    }
267
268    /// Requires capabilities from the exact inspected and realized session.
269    pub const fn with_required_session_capabilities(
270        mut self,
271        capabilities: SessionCapabilities,
272    ) -> Self {
273        self.required_session_capabilities = capabilities;
274        self
275    }
276
277    /// Selects portable drafting intent.
278    pub const fn with_drafting(mut self, drafting: DraftingLoadRequest) -> Self {
279        self.drafting = drafting;
280        self
281    }
282
283    /// Applies a portable execution plan's drafting mode before payload selection.
284    pub fn with_drafting_plan(
285        self,
286        plan: &DraftingPlan,
287    ) -> Result<Self, NormalizedLoadRequestError> {
288        let drafting = match plan {
289            DraftingPlan::Disabled => DraftingLoadRequest::Disabled,
290            DraftingPlan::Embedded {
291                max_draft_tokens, ..
292            } => DraftingLoadRequest::embedded(*max_draft_tokens)?,
293            DraftingPlan::External { .. } => DraftingLoadRequest::ExternalTarget,
294            _ => return Err(NormalizedLoadRequestError::UnsupportedDraftingPlan),
295        };
296        Ok(self.with_drafting(drafting))
297    }
298
299    /// Returns the requested dense-weight transformation.
300    pub const fn quantization(&self) -> Option<QuantizationRequest> {
301        self.quantization
302    }
303
304    /// Returns the complete portable parallel request.
305    pub const fn parallel_execution(&self) -> Option<ParallelLoadRequest> {
306        self.parallel
307    }
308
309    /// Returns the selected portable rank, when parallel execution was attached.
310    pub const fn parallel_topology(&self) -> Option<ParallelRankTopology> {
311        match self.parallel {
312            Some(parallel) => Some(parallel.rank()),
313            None => None,
314        }
315    }
316
317    /// Returns the activation wire contract for parallel execution.
318    pub const fn pipeline_wire_contract(&self) -> Option<PipelineWireContract> {
319        match self.parallel {
320            Some(parallel) => Some(parallel.wire()),
321            None => None,
322        }
323    }
324
325    /// Reports whether a portable parallel execution request is attached.
326    pub const fn has_parallel_execution(&self) -> bool {
327        self.parallel.is_some()
328    }
329
330    /// Returns invocation limits already validated by parallel construction.
331    pub fn partitioned_invocation_limits(&self) -> Option<(i32, i32)> {
332        self.parallel.map(ParallelLoadRequest::invocation_limits)
333    }
334
335    /// Returns bounded communication completion for model preparation.
336    pub fn communication_completion_policy(
337        &self,
338    ) -> Result<Option<CommunicationCompletionPolicy>, NormalizedLoadRequestError> {
339        match (self.parallel, self.communication_completion) {
340            (Some(parallel), None) => Ok(Some(parallel.completion())),
341            (None, None) => Ok(None),
342            (None, Some(_)) => Err(NormalizedLoadRequestError::OrphanedModelCompletion),
343            (Some(_), Some(_)) => Err(NormalizedLoadRequestError::ConflictingCompletionPolicy),
344        }
345    }
346
347    /// Returns completion policy for realtime work, including local async evaluation.
348    pub fn realtime_completion_policy(
349        &self,
350    ) -> Result<CommunicationCompletionPolicy, NormalizedLoadRequestError> {
351        if let Some(parallel) = self.parallel {
352            return Ok(parallel.completion());
353        }
354        self.communication_completion.map_or_else(
355            || {
356                CommunicationCompletionPolicy::new(
357                    std::time::Duration::from_secs(30),
358                    CompletionCancellationMode::QuarantineUntilComplete,
359                )
360                .map_err(|error| NormalizedLoadRequestError::Completion(error.to_string()))
361            },
362            Ok,
363        )
364    }
365
366    /// Returns selected immutable-weight residency.
367    pub const fn weight_residency(&self) -> WeightResidency {
368        self.weight_residency
369    }
370
371    /// Returns selected mutable-state residency.
372    pub const fn state_residency(&self) -> &CacheResidencyPolicy {
373        &self.state_residency
374    }
375
376    /// Returns capabilities required from the realized session.
377    pub const fn required_session_capabilities(&self) -> SessionCapabilities {
378        self.required_session_capabilities
379    }
380
381    /// Returns the selected pre-payload drafting intent.
382    pub const fn drafting(&self) -> DraftingLoadRequest {
383        self.drafting
384    }
385
386    /// Converts the requested transformation to the checkpoint lowering vocabulary.
387    pub fn weight_quantization(
388        &self,
389    ) -> Result<Option<WeightQuantization>, NormalizedLoadRequestError> {
390        self.quantization
391            .map(|request| match request {
392                QuantizationRequest::Affine { group_size, bits } => {
393                    let group_size = i32::try_from(group_size).map_err(|_| {
394                        NormalizedLoadRequestError::Quantization(format!(
395                            "group_size must fit in i32, got {group_size}"
396                        ))
397                    })?;
398                    AffineQuantization::new(group_size, i32::from(bits))
399                        .map(WeightQuantization::Affine)
400                        .map_err(|error| {
401                            NormalizedLoadRequestError::Quantization(error.to_string())
402                        })
403                }
404                QuantizationRequest::MxFp4 => Ok(WeightQuantization::MxFp4),
405                _ => Err(NormalizedLoadRequestError::Quantization(
406                    "unknown load-time transformation request".into(),
407                )),
408            })
409            .transpose()
410    }
411
412    /// Validates the complete ordinary model-preparation request.
413    pub fn validate_model_preparation(
414        &self,
415    ) -> Result<ValidatedModelLoadRequest<'_>, NormalizedLoadRequestError> {
416        if self.max_cached_shards() == 0 {
417            return Err(NormalizedLoadRequestError::ZeroCachedShards);
418        }
419        self.weight_quantization()?;
420        self.communication_completion_policy()?;
421        Ok(ValidatedModelLoadRequest {
422            request: self,
423            policy: self.project_preparation_policy(),
424        })
425    }
426
427    /// Converts this normalized request into core's portable preparation policy.
428    pub fn preparation_policy(&self) -> Result<PreparationPolicy, NormalizedLoadRequestError> {
429        self.validate_model_preparation()
430            .map(ValidatedModelLoadRequest::preparation_policy)
431    }
432
433    fn project_preparation_policy(&self) -> PreparationPolicy {
434        let residency = if self.weight_residency.parameter_bank_cache().is_some() {
435            ResidencyRequest::AddressableParameterBanks
436        } else {
437            match self.weight_residency.layers() {
438                LayerWeightResidency::FullyResident => ResidencyRequest::FullyResident,
439                LayerWeightResidency::LayerwiseHost(_) => ResidencyRequest::LayerwiseHost,
440                LayerWeightResidency::DenseDiskStream(_) => ResidencyRequest::DenseDiskStream,
441            }
442        };
443        let mut policy = PreparationPolicy::new(self.quantization, residency)
444            .with_required_session_capabilities(self.required_session_capabilities);
445        if let Some(topology) = self.parallel_topology() {
446            policy = policy.with_topology(topology.topology());
447        }
448        policy
449    }
450}
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455    use eredu_core::{ParallelTopology, QuantizationRequest};
456
457    fn completion() -> CommunicationCompletionPolicy {
458        CommunicationCompletionPolicy::new(
459            std::time::Duration::from_secs(1),
460            CompletionCancellationMode::QuarantineUntilComplete,
461        )
462        .unwrap()
463    }
464
465    #[test]
466    fn parallel_policy_is_atomic_and_exact() {
467        let rank =
468            ParallelRankTopology::new(ParallelTopology::new(2, 1, 1, 1).unwrap(), 1).unwrap();
469        let parallel = ParallelLoadRequest::new(
470            rank,
471            PipelineWireContract::new(crate::PipelineActivationDtype::Float32),
472            2,
473            128,
474            completion(),
475        )
476        .unwrap();
477        let request = NormalizedLoadRequest::with_quantization(QuantizationRequest::MxFp4)
478            .with_parallel_execution(parallel)
479            .unwrap()
480            .with_required_session_capabilities(SessionCapabilities::new(true, false, true));
481
482        request.validate_model_preparation().unwrap();
483        assert_eq!(request.parallel_execution(), Some(parallel));
484        assert_eq!(request.parallel_topology(), Some(rank));
485        assert_eq!(request.partitioned_invocation_limits(), Some((2, 128)));
486        assert_eq!(
487            request.communication_completion_policy().unwrap(),
488            Some(completion())
489        );
490        assert_eq!(
491            request.preparation_policy().unwrap().topology(),
492            Some(rank.topology())
493        );
494    }
495
496    #[test]
497    fn invalid_parallel_geometry_fails_at_parallel_construction() {
498        let rank =
499            ParallelRankTopology::new(ParallelTopology::new(2, 1, 1, 1).unwrap(), 0).unwrap();
500        let request = ParallelLoadRequest::new(
501            rank,
502            PipelineWireContract::new(crate::PipelineActivationDtype::Float32),
503            0,
504            128,
505            completion(),
506        );
507        assert!(matches!(
508            request,
509            Err(NormalizedLoadRequestError::InvalidInvocationLimits { .. })
510        ));
511
512        let request = ParallelLoadRequest::new(
513            rank,
514            PipelineWireContract::new(crate::PipelineActivationDtype::Float32),
515            1,
516            -1,
517            completion(),
518        );
519        assert!(matches!(
520            request,
521            Err(NormalizedLoadRequestError::InvalidInvocationLimits { .. })
522        ));
523    }
524
525    #[test]
526    fn replicated_topology_is_rejected_at_parallel_construction() {
527        let rank =
528            ParallelRankTopology::new(ParallelTopology::new(1, 1, 1, 1).unwrap(), 0).unwrap();
529        assert!(matches!(
530            ParallelLoadRequest::new(
531                rank,
532                PipelineWireContract::new(crate::PipelineActivationDtype::Float32),
533                1,
534                128,
535                completion(),
536            ),
537            Err(NormalizedLoadRequestError::ReplicatedParallelTopology)
538        ));
539    }
540
541    #[test]
542    fn local_completion_cannot_be_reinterpreted_as_parallel_completion() {
543        let rank =
544            ParallelRankTopology::new(ParallelTopology::new(2, 1, 1, 1).unwrap(), 0).unwrap();
545        let parallel = ParallelLoadRequest::new(
546            rank,
547            PipelineWireContract::new(crate::PipelineActivationDtype::Float32),
548            1,
549            128,
550            completion(),
551        )
552        .unwrap();
553        let request =
554            NormalizedLoadRequest::default().with_communication_completion_policy(completion());
555        assert!(matches!(
556            request.with_parallel_execution(parallel),
557            Err(NormalizedLoadRequestError::ConflictingCompletionPolicy)
558        ));
559    }
560
561    #[test]
562    fn embedded_drafting_capacity_is_positive_by_construction() {
563        assert!(matches!(
564            DraftingLoadRequest::embedded(0),
565            Err(NormalizedLoadRequestError::ZeroEmbeddedDraftCapacity)
566        ));
567        assert_eq!(
568            DraftingLoadRequest::embedded(4)
569                .unwrap()
570                .embedded_capacity()
571                .unwrap()
572                .get(),
573            4
574        );
575    }
576
577    #[test]
578    fn local_realtime_completion_does_not_become_valid_model_communication() {
579        let request =
580            NormalizedLoadRequest::default().with_communication_completion_policy(completion());
581        assert_eq!(request.realtime_completion_policy().unwrap(), completion());
582        assert!(matches!(
583            request.validate_model_preparation(),
584            Err(NormalizedLoadRequestError::OrphanedModelCompletion)
585        ));
586    }
587}
eredu_runtime/load_request.rs

eredu_runtime/
load_request.rs
