aisimulate_core/engine/
runtime.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Factory for one logical mock engine.
5
6use std::num::NonZeroU32;
7use std::sync::Arc;
8
9use anyhow::Result;
10
11use crate::engine::generalized::{EngineIdentity, GeneralizedMockerEngine, RankIdentity};
12use crate::engine::scheduler::{SchedulerRank, engine_seed_offset};
13use crate::engine::{EngineConfig, TimingModel};
14
15/// A single-rank or attention-DP engine.
16pub type Engine = GeneralizedMockerEngine<SchedulerRank>;
17
18/// Constructs engines with one process-local timing provider.
19///
20/// The serializable engine configuration contains only a provider descriptor.
21/// A Runner resolves external timing providers before constructing this
22/// factory, which then shares the provider across all attention-DP ranks.
23#[derive(Clone)]
24pub struct EngineFactory {
25    config: EngineConfig,
26    timing: Arc<dyn TimingModel>,
27}
28
29impl EngineFactory {
30    /// Construct a factory for a built-in timing model.
31    pub fn new(config: EngineConfig) -> Result<Self> {
32        config.validate()?;
33        let timing = config.built_in_timing_model()?;
34        Ok(Self { config, timing })
35    }
36
37    /// Construct a factory with a process-local timing provider.
38    pub fn with_timing_model(config: EngineConfig, timing: Arc<dyn TimingModel>) -> Result<Self> {
39        config.validate()?;
40        Ok(Self { config, timing })
41    }
42
43    /// Build one scheduler/KV/timing rank with an explicit identity.
44    pub fn build_rank(&self, identity: RankIdentity) -> Result<SchedulerRank> {
45        let seed_offset = engine_seed_offset(identity)?;
46        SchedulerRank::new_with_timing_model(
47            identity,
48            &self.config,
49            Arc::clone(&self.timing),
50            seed_offset,
51        )
52    }
53
54    /// Build a single-rank or attention-DP logical engine.
55    pub fn build(&self, identity: EngineIdentity, dp_size: NonZeroU32) -> Result<Engine> {
56        GeneralizedMockerEngine::new_with_rank_factory(identity, dp_size, |rank_identity| {
57            self.build_rank(rank_identity)
58        })
59    }
60}
aisimulate_core/engine/runtime.rs

aisimulate_core/engine/
runtime.rs
