Skip to main content

ferrum_engine/
builder.rs

1//! Engine builder with registry-based component creation
2//!
3//! This module provides a fluent builder API for creating inference engines
4//! using the component registry pattern. The builder supports:
5//!
6//! - Configuration-driven component selection
7//! - Custom component overrides
8//! - Automatic fallback to defaults
9//! - Validation before engine creation
10
11use crate::registry::{ComponentConfig, ComponentRegistry};
12use ferrum_interfaces::engine::{InferenceEngine, LlmInferenceEngine};
13use ferrum_interfaces::{
14    KvCacheManager, ModelExecutor, RecurrentStateManager, Sampler, SchedulerInterface as Scheduler,
15    TensorFactory, Tokenizer,
16};
17use ferrum_models::vnext::{PreparedProductionModel, ProductionModelSourceBundle};
18use ferrum_types::{EngineConfig, FerrumError, Result};
19use std::sync::Arc;
20use tracing::{debug, info};
21
22// Engine-build composition knobs (FERRUM_MODEL_PATH / FERRUM_SPEC_DRAFT /
23// FERRUM_SPEC_N) are no longer read from the environment here. The CLI
24// composition root captures them via `RuntimeConfigSnapshot::capture_current()`
25// and lands them in `EngineConfig.runtime` through
26// `apply_runtime_config_snapshot`; the builder reads `self.config.runtime`.
27
28/// Engine builder for creating inference engines with registry-based components
29pub struct EngineBuilder {
30    /// Component registry to use
31    registry: Arc<ComponentRegistry>,
32    /// Engine configuration
33    config: EngineConfig,
34    /// Product-resolved semantic, tokenizer, and weight sources.
35    model_sources: Option<Arc<ProductionModelSourceBundle>>,
36    /// Immutable typed model package prepared once by the product composition
37    /// root and reused by startup policy and executor construction.
38    prepared_model: Option<Arc<PreparedProductionModel>>,
39    /// Override: custom tokenizer name
40    tokenizer_name: Option<String>,
41    /// Override: custom sampler name
42    sampler_name: Option<String>,
43    /// Override: custom scheduler name
44    scheduler_name: Option<String>,
45    /// Override: custom KV cache name
46    kv_cache_name: Option<String>,
47    /// Override: custom executor name
48    executor_name: Option<String>,
49    /// Pre-created tokenizer (skip factory)
50    custom_tokenizer: Option<Arc<dyn Tokenizer + Send + Sync>>,
51    /// Pre-created sampler (skip factory)
52    custom_sampler: Option<Arc<dyn Sampler + Send + Sync>>,
53    /// Pre-created scheduler (skip factory)
54    custom_scheduler: Option<Arc<dyn Scheduler + Send + Sync>>,
55    /// Pre-created KV cache (skip factory)
56    custom_kv_cache: Option<Arc<dyn KvCacheManager + Send + Sync>>,
57    /// Pre-created recurrent-state manager (skip factory)
58    custom_recurrent_state_manager: Option<Arc<dyn RecurrentStateManager + Send + Sync>>,
59    /// Pre-created executor (skip factory)
60    custom_executor: Option<Arc<dyn ModelExecutor + Send + Sync>>,
61}
62
63impl EngineBuilder {
64    /// Create a new engine builder with default registry
65    pub fn new(config: EngineConfig) -> Self {
66        Self::with_registry(config, crate::registry::global_registry())
67    }
68
69    /// Create a new engine builder with a custom registry
70    pub fn with_registry(config: EngineConfig, registry: Arc<ComponentRegistry>) -> Self {
71        Self {
72            registry,
73            config,
74            model_sources: None,
75            prepared_model: None,
76            tokenizer_name: None,
77            sampler_name: None,
78            scheduler_name: None,
79            kv_cache_name: None,
80            executor_name: None,
81            custom_tokenizer: None,
82            custom_sampler: None,
83            custom_scheduler: None,
84            custom_kv_cache: None,
85            custom_recurrent_state_manager: None,
86            custom_executor: None,
87        }
88    }
89
90    pub fn with_model_sources(mut self, sources: Arc<ProductionModelSourceBundle>) -> Self {
91        self.model_sources = Some(sources);
92        self.prepared_model = None;
93        self
94    }
95
96    pub fn with_prepared_model(mut self, prepared: Arc<PreparedProductionModel>) -> Self {
97        self.model_sources = Some(Arc::clone(prepared.sources()));
98        self.prepared_model = Some(prepared);
99        self
100    }
101
102    /// Set the tokenizer to use by name
103    pub fn with_tokenizer(mut self, name: impl Into<String>) -> Self {
104        self.tokenizer_name = Some(name.into());
105        self
106    }
107
108    /// Set a pre-created tokenizer
109    pub fn with_custom_tokenizer(mut self, tokenizer: Arc<dyn Tokenizer + Send + Sync>) -> Self {
110        self.custom_tokenizer = Some(tokenizer);
111        self
112    }
113
114    /// Set the sampler to use by name
115    pub fn with_sampler(mut self, name: impl Into<String>) -> Self {
116        self.sampler_name = Some(name.into());
117        self
118    }
119
120    /// Set a pre-created sampler
121    pub fn with_custom_sampler(mut self, sampler: Arc<dyn Sampler + Send + Sync>) -> Self {
122        self.custom_sampler = Some(sampler);
123        self
124    }
125
126    /// Set the scheduler to use by name
127    pub fn with_scheduler(mut self, name: impl Into<String>) -> Self {
128        self.scheduler_name = Some(name.into());
129        self
130    }
131
132    /// Set a pre-created scheduler
133    pub fn with_custom_scheduler(mut self, scheduler: Arc<dyn Scheduler + Send + Sync>) -> Self {
134        self.custom_scheduler = Some(scheduler);
135        self
136    }
137
138    /// Set the KV cache to use by name
139    pub fn with_kv_cache(mut self, name: impl Into<String>) -> Self {
140        self.kv_cache_name = Some(name.into());
141        self
142    }
143
144    /// Set a pre-created KV cache manager
145    pub fn with_custom_kv_cache(mut self, kv_cache: Arc<dyn KvCacheManager + Send + Sync>) -> Self {
146        self.custom_kv_cache = Some(kv_cache);
147        self
148    }
149
150    /// Set a pre-created recurrent-state manager.
151    pub fn with_custom_recurrent_state_manager(
152        mut self,
153        manager: Arc<dyn RecurrentStateManager + Send + Sync>,
154    ) -> Self {
155        self.custom_recurrent_state_manager = Some(manager);
156        self
157    }
158
159    /// Set the executor to use by name
160    pub fn with_executor(mut self, name: impl Into<String>) -> Self {
161        self.executor_name = Some(name.into());
162        self
163    }
164
165    /// Set a pre-created model executor
166    pub fn with_custom_executor(mut self, executor: Arc<dyn ModelExecutor + Send + Sync>) -> Self {
167        self.custom_executor = Some(executor);
168        self
169    }
170
171    /// Determine which tokenizer to use based on config and overrides
172    fn resolve_tokenizer_name(&self) -> String {
173        if let Some(ref name) = self.tokenizer_name {
174            return name.clone();
175        }
176
177        // If model path is set, try huggingface first
178        if self.has_typed_model_path() || self.config.runtime.model_path.is_some() {
179            return "huggingface".to_string();
180        }
181
182        "stub".to_string()
183    }
184
185    /// Determine which sampler to use based on config and overrides
186    fn resolve_sampler_name(&self) -> String {
187        if let Some(ref name) = self.sampler_name {
188            return name.clone();
189        }
190
191        // Could derive from sampling params in the future
192        "multinomial".to_string()
193    }
194
195    /// Determine which KV cache to use based on config and overrides
196    fn resolve_kv_cache_name(&self) -> String {
197        if let Some(ref name) = self.kv_cache_name {
198            return name.clone();
199        }
200
201        match self.config.kv_cache.cache_type {
202            ferrum_types::KvCacheType::Contiguous => "default".to_string(),
203            ferrum_types::KvCacheType::Paged => "paged".to_string(),
204            ferrum_types::KvCacheType::Tree => "default".to_string(), // Fallback
205        }
206    }
207
208    /// Determine which executor to use based on config and overrides
209    fn resolve_executor_name(&self) -> String {
210        if let Some(ref name) = self.executor_name {
211            return name.clone();
212        }
213
214        // If model path is set, try candle executor
215        if self.has_typed_model_path() || self.config.runtime.model_path.is_some() {
216            return "llm".to_string();
217        }
218
219        "stub".to_string()
220    }
221
222    fn has_typed_model_path(&self) -> bool {
223        self.model_sources.is_some()
224            || self
225                .config
226                .backend
227                .backend_options
228                .get("model_path")
229                .and_then(|value| value.as_str())
230                .is_some()
231    }
232
233    /// Build the inference engine
234    pub async fn build(self) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
235        info!(
236            "Building inference engine for model: {}",
237            self.config.model.model_id
238        );
239
240        if self.scheduler_name.is_some() || self.custom_scheduler.is_some() {
241            return Err(FerrumError::config(
242                "EngineBuilder scheduler component overrides are no longer accepted; configure the typed EngineConfig.scheduler used by ContinuousBatchScheduler",
243            ));
244        }
245
246        // Pre-compute all component names before consuming self
247        let tokenizer_name = self.resolve_tokenizer_name();
248        let sampler_name = self.resolve_sampler_name();
249        let kv_cache_name = self.resolve_kv_cache_name();
250        let executor_name = self.resolve_executor_name();
251        let explicit_kv_cache_override = self.kv_cache_name.is_some();
252
253        let component_config = ComponentConfig::from_engine_config_and_product_model(
254            &self.config,
255            self.model_sources.clone(),
256            self.prepared_model.clone(),
257        );
258        validate_layer_split_plan(&component_config)?;
259        let typed_model_path = component_config.get_string_option("model_path");
260        let has_model_path = typed_model_path.is_some() || self.config.runtime.model_path.is_some();
261        let registry = self.registry.clone();
262        let config = self.config;
263
264        // Extract custom components. Phase 3e+ deleted the legacy
265        // `ComputeBackend` trait, so there's no "backend" component to
266        // build here — real GPU dispatch goes through `Backend<B>` in
267        // `ferrum-kernels`. The stub executor now wires
268        // `CandleTensorFactory` directly from the registry.
269        let custom_tokenizer = self.custom_tokenizer;
270        let custom_sampler = self.custom_sampler;
271        let custom_kv_cache = self.custom_kv_cache;
272        let custom_recurrent_state_manager = self.custom_recurrent_state_manager;
273        let custom_executor = self.custom_executor;
274
275        // 2. Create or use provided tokenizer
276        let tokenizer = if let Some(tokenizer) = custom_tokenizer {
277            debug!("Using custom tokenizer");
278            tokenizer
279        } else {
280            debug!("Creating tokenizer: {}", tokenizer_name);
281
282            // Try primary tokenizer, fallback to stub
283            match registry
284                .create_tokenizer(&tokenizer_name, &component_config)
285                .await
286            {
287                Ok(t) => t,
288                Err(e) => {
289                    if has_model_path {
290                        return Err(FerrumError::config(format!(
291                            "Failed to create tokenizer '{}' in model mode: {}",
292                            tokenizer_name, e
293                        )));
294                    }
295
296                    tracing::warn!(
297                        "Failed to create tokenizer '{}': {}, falling back to stub",
298                        tokenizer_name,
299                        e
300                    );
301                    registry.create_tokenizer("stub", &component_config).await?
302                }
303            }
304        };
305
306        // 3. Create or use provided sampler
307        let sampler = if let Some(sampler) = custom_sampler {
308            debug!("Using custom sampler");
309            sampler
310        } else {
311            debug!("Creating sampler: {}", sampler_name);
312            registry
313                .create_sampler(&sampler_name, &component_config)
314                .await?
315        };
316
317        // 4. Resolve the executor before resource managers. Its typed
318        // authority decides whether engine-side KV/recurrent managers may
319        // exist at all.
320        let executor = if let Some(executor) = custom_executor {
321            debug!("Using custom executor");
322            executor
323        } else {
324            debug!("Creating executor: {}", executor_name);
325
326            // Try primary executor, fallback to stub
327            match registry
328                .create_executor(&executor_name, &component_config)
329                .await
330            {
331                Ok(e) => e,
332                Err(err) => {
333                    if has_model_path {
334                        return Err(FerrumError::config(format!(
335                            "Failed to create executor '{}' in model mode: {}",
336                            executor_name, err
337                        )));
338                    }
339
340                    tracing::warn!(
341                        "Failed to create executor '{}': {}, falling back to stub",
342                        executor_name,
343                        err
344                    );
345                    registry.create_executor("stub", &component_config).await?
346                }
347            }
348        };
349        let execution_resource_authority = executor.execution_resource_authority();
350
351        let (kv_cache, recurrent_state_manager) = match execution_resource_authority {
352            ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime => {
353                if explicit_kv_cache_override || custom_kv_cache.is_some() {
354                    return Err(FerrumError::config(
355                        "plan runtime cannot be combined with a legacy engine KV-cache override",
356                    ));
357                }
358                if custom_recurrent_state_manager.is_some() {
359                    return Err(FerrumError::config(
360                        "plan runtime cannot be combined with a legacy engine recurrent-state manager",
361                    ));
362                }
363                if executor.resolved_model_plan().is_none() {
364                    return Err(FerrumError::config(
365                        "plan-runtime executor did not expose its authoritative ResolvedModelPlan",
366                    ));
367                }
368                (None, None)
369            }
370            ferrum_interfaces::model_executor::ExecutionResourceAuthority::LegacyEngine => {
371                let kv_cache = if let Some(kv_cache) = custom_kv_cache {
372                    debug!("Using custom KV cache");
373                    kv_cache
374                } else {
375                    debug!("Creating KV cache: {}", kv_cache_name);
376                    registry
377                        .create_kv_cache(&kv_cache_name, &component_config)
378                        .await?
379                };
380                let recurrent_state_manager = custom_recurrent_state_manager
381                    .or_else(|| default_recurrent_state_manager(&config));
382                (Some(kv_cache), recurrent_state_manager)
383            }
384        };
385
386        // 5. Create the engine — always ContinuousBatchEngine.
387        info!("All components created, building ContinuousBatchEngine");
388
389        let cb_scheduler = Arc::new(
390            ferrum_scheduler::implementations::ContinuousBatchScheduler::new(
391                config.scheduler.clone(),
392            ),
393        );
394
395        // Create TensorFactory for the configured device
396        let tensor_factory: Arc<dyn TensorFactory> = Arc::new(
397            crate::tensor_factory::candle::CandleTensorFactory::new(config.backend.device.clone()),
398        );
399
400        // Opt-in speculative decoding: provide an absolute HF snapshot path
401        // for a second smaller model. The draft must use the same tokenizer
402        // + vocab as the target (same family e.g. Qwen3). Backend options are
403        // the typed startup path; the legacy speculative env names remain
404        // compatibility aliases.
405        let spec_draft = component_config
406            .get_string_option("spec_draft")
407            .or_else(|| config.runtime.spec_draft.clone());
408        if execution_resource_authority
409            == ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime
410            && spec_draft.is_some()
411        {
412            return Err(FerrumError::unsupported(
413                "speculative decoding is not yet part of the plan-runtime contract",
414            ));
415        }
416        let spec_n = component_config
417            .get_option::<usize>("spec_n")
418            .unwrap_or(config.runtime.spec_n.unwrap_or(4));
419        let (draft_executor, spec_config) = match spec_draft.as_ref() {
420            Some(draft_path) => {
421                info!("Speculative decoding: loading draft model from {draft_path}");
422                let mut draft_cfg = component_config.clone();
423                draft_cfg.component_options.insert(
424                    "model_path".to_string(),
425                    serde_json::Value::String(draft_path.to_string()),
426                );
427                let draft = registry
428                    .create_executor(&executor_name, &draft_cfg)
429                    .await
430                    .map_err(|error| {
431                        FerrumError::config(format!(
432                            "requested speculative draft executor failed to load: {error}"
433                        ))
434                    })?;
435                if draft.execution_resource_authority() != execution_resource_authority {
436                    return Err(FerrumError::config(
437                        "target and speculative draft executors declare different resource authority",
438                    ));
439                }
440                (
441                    Some(draft),
442                    Some(crate::speculative::SpeculativeDecodingConfig {
443                        num_speculative_tokens: spec_n,
444                        temperature: 1.0,
445                    }),
446                )
447            }
448            _ => (None, None),
449        };
450
451        // Construct the unpublished engine shell first so its typed profile
452        // sink is attached before executor startup emits warmup/capture events.
453        let engine = match execution_resource_authority {
454            ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime => {
455                crate::ContinuousBatchEngine::new_plan_runtime(
456                    config,
457                    cb_scheduler,
458                    tokenizer,
459                    sampler,
460                    Arc::clone(&executor),
461                    tensor_factory,
462                )?
463            }
464            ferrum_interfaces::model_executor::ExecutionResourceAuthority::LegacyEngine => {
465                let kv_cache = kv_cache.ok_or_else(|| {
466                    FerrumError::internal("legacy-engine composition lost its KV-cache manager")
467                })?;
468                crate::ContinuousBatchEngine::new_with_speculation_and_recurrent_state_manager(
469                    config,
470                    cb_scheduler,
471                    tokenizer,
472                    sampler,
473                    kv_cache,
474                    Arc::clone(&executor),
475                    tensor_factory,
476                    draft_executor.clone(),
477                    spec_config,
478                    recurrent_state_manager,
479                )?
480            }
481        };
482
483        // This is the single product readiness boundary shared by `run` and
484        // `serve`. The engine is not exposed until executor-owned compilation
485        // and warmup complete, but those events now share the product trace.
486        let startup_result = async {
487            executor.prepare_startup().await?;
488            if let Some(draft) = draft_executor.as_ref() {
489                draft.prepare_startup().await?;
490            }
491            Ok(())
492        }
493        .await;
494        if let Err(startup_error) = startup_result {
495            if let Err(shutdown_error) = engine.shutdown().await {
496                tracing::warn!(
497                    "Failed to close engine resources after startup rejection: {shutdown_error}"
498                );
499            }
500            return Err(startup_error);
501        }
502        Ok(Box::new(engine))
503    }
504}
505
506fn default_recurrent_state_manager(
507    config: &EngineConfig,
508) -> Option<Arc<dyn RecurrentStateManager + Send + Sync>> {
509    let total_batch_slots = config
510        .runtime
511        .recurrent_state_max_slots
512        .unwrap_or(usize::MAX);
513    Some(
514        Arc::new(crate::recurrent_state::InMemoryRecurrentStateManager::new(
515            crate::recurrent_state::InMemoryRecurrentStateConfig {
516                total_memory_bytes: usize::MAX,
517                total_batch_slots,
518            },
519        )) as Arc<dyn RecurrentStateManager + Send + Sync>,
520    )
521}
522
523fn validate_layer_split_plan(component_config: &ComponentConfig) -> Result<()> {
524    if component_config
525        .get_string_option("selected_distributed_strategy")
526        .as_deref()
527        != Some("layer_split")
528    {
529        return Ok(());
530    }
531    let requested = component_config
532        .get_option::<Vec<usize>>("requested_gpu_devices")
533        .unwrap_or_default();
534    let selected = component_config
535        .get_option::<Vec<usize>>("selected_gpu_devices")
536        .unwrap_or_default();
537    let plan_raw = component_config.get_string_option("selected_layer_split_plan");
538    let parsed_plan = if let Some(stages) = component_config
539        .component_options
540        .get("selected_layer_split_stages")
541    {
542        crate::layer_split::parse_layer_split_stage_documents(stages)?
543    } else {
544        let plan_raw = plan_raw.as_deref().ok_or_else(|| {
545            FerrumError::config(
546                "selected_distributed_strategy=layer_split requires selected_layer_split_plan",
547            )
548        })?;
549        crate::layer_split::parse_layer_split_plan(plan_raw)?
550    };
551    crate::layer_split::validate_layer_split_plan_for_devices(&parsed_plan, &selected)?;
552    let execution_plan = parsed_plan.to_execution_plan();
553    let stage_ranges = execution_plan
554        .layer_distribution
555        .stage_layers
556        .iter()
557        .map(|range| format!("{}-{}", range.start, range.end.saturating_sub(1)))
558        .collect::<Vec<_>>()
559        .join(",");
560    let plan_label = plan_raw.unwrap_or_else(|| format!("{:?}", parsed_plan.stages));
561    tracing::info!(
562        "validated CUDA layer_split plan: requested_gpu_devices={requested:?} selected_gpu_devices={selected:?} selected_layer_split_plan={plan_label} total_layers={} pipeline_stages={} stage_ranges={stage_ranges} communication_backend={}",
563        parsed_plan.total_layers(),
564        execution_plan.parallel_config.pipeline_parallel_size,
565        execution_plan.parallel_config.communication_backend,
566    );
567    Ok(())
568}
569
570/// Create an engine with the default configuration and registry
571pub async fn create_engine(
572    config: EngineConfig,
573) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
574    EngineBuilder::new(config).build().await
575}
576
577/// Create a product engine from one immutable role-specific source bundle.
578pub async fn create_product_engine(
579    config: EngineConfig,
580    sources: Arc<ProductionModelSourceBundle>,
581) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
582    EngineBuilder::new(config)
583        .with_model_sources(sources)
584        .build()
585        .await
586}
587
588/// Create a product engine from the exact typed model package already used by
589/// startup capability and resource-policy resolution.
590pub async fn create_prepared_product_engine(
591    config: EngineConfig,
592    prepared: Arc<PreparedProductionModel>,
593) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
594    EngineBuilder::new(config)
595        .with_prepared_model(prepared)
596        .build()
597        .await
598}
599
600// ============================================================================
601// Tests
602// ============================================================================
603
604#[cfg(test)]
605mod tests {
606    use super::*;
607    use ferrum_interfaces::{
608        model_executor::{
609            DecodeInput, DecodeOutput, ExecutionResourceAuthority, ExecutorCapabilities,
610            ExecutorStatus, PlanRuntimeResourceSnapshot, PrefillInput, PrefillOutput,
611        },
612        vnext::ExecutionEventSink,
613        RecurrentStateHandle, RecurrentStateManager, RecurrentStateManagerStats,
614        RecurrentStateSpec, RecurrentStateTensorSpec,
615    };
616    use ferrum_types::{DataType, Device, RequestId};
617    use std::sync::{
618        atomic::{AtomicBool, AtomicUsize, Ordering},
619        Mutex,
620    };
621
622    #[derive(Debug)]
623    struct NoopRecurrentStateManager;
624
625    struct PlanRuntimeBuilderExecutor {
626        inner: ferrum_testkit::MockModelExecutor,
627    }
628
629    struct StartupProbeExecutor {
630        inner: ferrum_testkit::MockModelExecutor,
631        calls: Arc<AtomicUsize>,
632        fail: bool,
633    }
634
635    struct ProfileStartupProbeExecutor {
636        inner: ferrum_testkit::MockModelExecutor,
637        event_sink: Mutex<Option<Arc<dyn ExecutionEventSink>>>,
638        saw_sink_during_startup: Arc<AtomicBool>,
639    }
640
641    #[async_trait::async_trait]
642    impl ModelExecutor for StartupProbeExecutor {
643        fn info(&self) -> &ferrum_types::ModelInfo {
644            self.inner.info()
645        }
646
647        async fn prepare_startup(&self) -> Result<()> {
648            self.calls.fetch_add(1, Ordering::Relaxed);
649            if self.fail {
650                return Err(FerrumError::backend("startup preparation rejected"));
651            }
652            Ok(())
653        }
654
655        async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
656            self.inner.prefill(input).await
657        }
658
659        async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
660            self.inner.decode(input).await
661        }
662
663        fn capabilities(&self) -> ExecutorCapabilities {
664            self.inner.capabilities()
665        }
666
667        fn status(&self) -> ExecutorStatus {
668            self.inner.status()
669        }
670    }
671
672    #[async_trait::async_trait]
673    impl ModelExecutor for ProfileStartupProbeExecutor {
674        fn info(&self) -> &ferrum_types::ModelInfo {
675            self.inner.info()
676        }
677
678        async fn prepare_startup(&self) -> Result<()> {
679            use ferrum_interfaces::vnext::{ExecutionEventEmitter, TrustedExecutionEventContext};
680
681            let sink = self
682                .event_sink
683                .lock()
684                .expect("profile startup probe sink lock")
685                .clone();
686            self.saw_sink_during_startup
687                .store(sink.is_some(), Ordering::Release);
688            let Some(sink) = sink else {
689                return Ok(());
690            };
691            let (run_id, request_id, event) = startup_profile_test_event();
692            ExecutionEventEmitter::from_shared(sink, run_id.clone(), request_id.clone())
693                .emit(
694                    event,
695                    &TrustedExecutionEventContext::pre_plan(&run_id, &request_id),
696                )
697                .map_err(|error| {
698                    FerrumError::internal(format!("emit startup profile probe: {error}"))
699                })
700        }
701
702        fn attach_execution_event_sink(&self, sink: Arc<dyn ExecutionEventSink>) {
703            *self
704                .event_sink
705                .lock()
706                .expect("profile startup probe sink lock") = Some(sink);
707        }
708
709        async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
710            self.inner.prefill(input).await
711        }
712
713        async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
714            self.inner.decode(input).await
715        }
716
717        fn capabilities(&self) -> ExecutorCapabilities {
718            self.inner.capabilities()
719        }
720
721        fn status(&self) -> ExecutorStatus {
722            self.inner.status()
723        }
724    }
725
726    fn startup_profile_test_event() -> (
727        ferrum_interfaces::vnext::RunId,
728        ferrum_interfaces::vnext::RequestIdentity,
729        ferrum_interfaces::vnext::ExecutionEvent,
730    ) {
731        use ferrum_interfaces::vnext::{
732            ExecutionEvent, ExecutionEventDetail, ExecutionEventKind, ExecutionIdentityEnvelope,
733            ExecutionIdentityParts, ExecutionPhase, MonotonicTimestamp, RequestIdentity, RunId,
734            SpanId, EXECUTION_IDENTITY_VERSION,
735        };
736
737        let run_id = RunId::new("run.vnext.builder-startup-profile").unwrap();
738        let request_id = RequestIdentity::new("request.vnext.builder-startup-profile").unwrap();
739        let event = ExecutionEvent::new(
740            MonotonicTimestamp {
741                nanos_since_run_start: 1,
742            },
743            ExecutionPhase::Resolution,
744            ExecutionEventKind::RequestAccepted,
745            ExecutionIdentityEnvelope::new(ExecutionIdentityParts {
746                version: EXECUTION_IDENTITY_VERSION,
747                run_id: run_id.clone(),
748                request_id: request_id.clone(),
749                sequence: 1,
750                plan_id: None,
751                plan_hash: None,
752                frame_id: None,
753                node_invocation_id: None,
754                node_id: None,
755                operation_id: None,
756                provider_id: None,
757                device_id: None,
758                resource_pool_id: None,
759                resource_pool_identity_fingerprint: None,
760                provisioning_run_id: None,
761                provisioning_request_id: None,
762                transaction_id: None,
763                active_sequence_slot: None,
764                admission_generation: None,
765                activation_epoch: None,
766                runtime_implementation_fingerprint: None,
767                active_sequence_fingerprint: None,
768                completed_sequence_fingerprint: None,
769                aborted_sequence_fingerprint: None,
770                resource_id: None,
771                resource_generation: None,
772                resource_batch_fingerprint: None,
773                span_id: SpanId::new("vnext/request/builder-startup-profile").unwrap(),
774                parent_span_id: None,
775                async_links: Vec::new(),
776            })
777            .unwrap(),
778            ExecutionEventDetail::None,
779        )
780        .unwrap();
781        (run_id, request_id, event)
782    }
783
784    #[async_trait::async_trait]
785    impl ModelExecutor for PlanRuntimeBuilderExecutor {
786        fn info(&self) -> &ferrum_types::ModelInfo {
787            self.inner.info()
788        }
789
790        fn execution_resource_authority(&self) -> ExecutionResourceAuthority {
791            ExecutionResourceAuthority::PlanRuntime
792        }
793
794        fn plan_runtime_resource_snapshot(&self) -> Result<Option<PlanRuntimeResourceSnapshot>> {
795            PlanRuntimeResourceSnapshot::new(1_000, 900, 700, 700, 400, 300, 200, 0, 0).map(Some)
796        }
797
798        async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
799            self.inner.prefill(input).await
800        }
801
802        async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
803            self.inner.decode(input).await
804        }
805
806        fn capabilities(&self) -> ExecutorCapabilities {
807            self.inner.capabilities()
808        }
809
810        fn status(&self) -> ExecutorStatus {
811            self.inner.status()
812        }
813    }
814
815    struct CountingKvFactory {
816        calls: Arc<AtomicUsize>,
817    }
818
819    #[async_trait::async_trait]
820    impl crate::registry::ComponentFactory<Arc<dyn KvCacheManager + Send + Sync>>
821        for CountingKvFactory
822    {
823        async fn create(
824            &self,
825            _config: &ComponentConfig,
826        ) -> Result<Arc<dyn KvCacheManager + Send + Sync>> {
827            self.calls.fetch_add(1, Ordering::Relaxed);
828            Ok(Arc::new(ferrum_testkit::MockKvCacheManager::new(8)))
829        }
830
831        fn metadata(&self) -> crate::registry::ComponentMetadata {
832            crate::registry::ComponentMetadata::default()
833        }
834    }
835
836    #[async_trait::async_trait]
837    impl RecurrentStateManager for NoopRecurrentStateManager {
838        async fn allocate(
839            &self,
840            _spec: &RecurrentStateSpec,
841        ) -> Result<Arc<dyn RecurrentStateHandle>> {
842            Err(FerrumError::unsupported(
843                "noop recurrent-state manager does not allocate",
844            ))
845        }
846
847        async fn deallocate(&self, _request_id: RequestId) -> Result<()> {
848            Ok(())
849        }
850
851        fn can_allocate(&self, _spec: &RecurrentStateSpec) -> bool {
852            false
853        }
854
855        fn get_handle(&self, _request_id: RequestId) -> Option<Arc<dyn RecurrentStateHandle>> {
856            None
857        }
858
859        fn list_handles(&self) -> Vec<(RequestId, Arc<dyn RecurrentStateHandle>)> {
860            Vec::new()
861        }
862
863        fn stats(&self) -> RecurrentStateManagerStats {
864            RecurrentStateManagerStats {
865                total_memory_bytes: 0,
866                used_memory_bytes: 0,
867                active_states: 0,
868                active_state_tensors: 0,
869                total_batch_slots: 0,
870                used_batch_slots: 0,
871                allocation_count: 0,
872                allocation_failures: 0,
873                eviction_count: 0,
874            }
875        }
876
877        async fn reset(&self) -> Result<()> {
878            Ok(())
879        }
880    }
881
882    #[test]
883    fn test_builder_creation() {
884        let config = EngineConfig::default();
885        let builder = EngineBuilder::new(config);
886
887        assert!(builder.tokenizer_name.is_none());
888        assert!(builder.custom_recurrent_state_manager.is_none());
889    }
890
891    #[test]
892    fn test_builder_with_overrides() {
893        let config = EngineConfig::default();
894        let builder = EngineBuilder::new(config)
895            .with_tokenizer("custom_tokenizer")
896            .with_sampler("greedy")
897            .with_scheduler("priority")
898            .with_kv_cache("paged")
899            .with_executor("custom_executor");
900
901        assert_eq!(builder.tokenizer_name, Some("custom_tokenizer".to_string()));
902        assert_eq!(builder.sampler_name, Some("greedy".to_string()));
903        assert_eq!(builder.scheduler_name, Some("priority".to_string()));
904        assert_eq!(builder.kv_cache_name, Some("paged".to_string()));
905        assert_eq!(builder.executor_name, Some("custom_executor".to_string()));
906    }
907
908    #[test]
909    fn test_builder_with_custom_recurrent_state_manager() {
910        let config = EngineConfig::default();
911        let manager = Arc::new(NoopRecurrentStateManager);
912        let builder = EngineBuilder::new(config).with_custom_recurrent_state_manager(manager);
913
914        assert!(builder.custom_recurrent_state_manager.is_some());
915    }
916
917    #[test]
918    fn test_builder_typed_model_path_selects_model_components() {
919        let mut config = EngineConfig::default();
920        config.backend.backend_options.insert(
921            "model_path".to_string(),
922            serde_json::Value::String("/models/target".to_string()),
923        );
924        let builder = EngineBuilder::new(config);
925
926        assert!(builder.has_typed_model_path());
927        assert_eq!(builder.resolve_tokenizer_name(), "huggingface");
928        assert_eq!(builder.resolve_executor_name(), "llm");
929    }
930
931    #[test]
932    fn test_builder_retains_one_typed_source_bundle_for_components() {
933        let root = std::env::temp_dir().join(format!(
934            "ferrum-builder-source-bundle-{}-{}",
935            std::process::id(),
936            std::time::SystemTime::now()
937                .duration_since(std::time::UNIX_EPOCH)
938                .unwrap()
939                .as_nanos()
940        ));
941        std::fs::create_dir_all(&root).unwrap();
942        std::fs::write(
943            root.join("config.json"),
944            br#"{"architectures":["Fixture"]}"#,
945        )
946        .unwrap();
947        std::fs::write(root.join("tokenizer.json"), br#"{"version":"1.0"}"#).unwrap();
948        std::fs::write(root.join("model.safetensors"), b"fixture").unwrap();
949        let original = ferrum_interfaces::vnext::OriginalModelSource {
950            kind: ferrum_interfaces::vnext::ModelSourceKind::LocalDirectory,
951            location: root.display().to_string(),
952            requested_revision: None,
953        };
954        let sources = Arc::new(
955            ProductionModelSourceBundle::open(
956                &root,
957                &root,
958                ferrum_models::vnext::ProductionWeightArtifact::safetensors_directory(&root),
959                ferrum_interfaces::vnext::OriginalModelSources {
960                    semantic: original.clone(),
961                    tokenizer: original.clone(),
962                    weights: original,
963                },
964            )
965            .unwrap(),
966        );
967
968        let builder =
969            EngineBuilder::new(EngineConfig::default()).with_model_sources(Arc::clone(&sources));
970        assert!(builder.has_typed_model_path());
971        assert!(Arc::ptr_eq(
972            builder.model_sources.as_ref().unwrap(),
973            &sources
974        ));
975        assert_eq!(builder.resolve_tokenizer_name(), "huggingface");
976        assert_eq!(builder.resolve_executor_name(), "llm");
977        std::fs::remove_dir_all(root).unwrap();
978    }
979
980    #[test]
981    fn test_builder_typed_spec_options_parse_from_component_config() {
982        let mut config = EngineConfig::default();
983        config.backend.backend_options.insert(
984            "model_path".to_string(),
985            serde_json::Value::String("/models/target".to_string()),
986        );
987        config.backend.backend_options.insert(
988            "spec_draft".to_string(),
989            serde_json::Value::String("/models/draft".to_string()),
990        );
991        config.backend.backend_options.insert(
992            "spec_n".to_string(),
993            serde_json::Value::Number(serde_json::Number::from(6)),
994        );
995        let component_config = ComponentConfig::from_engine_config(&config);
996
997        assert_eq!(
998            component_config.get_string_option("spec_draft").as_deref(),
999            Some("/models/draft")
1000        );
1001        assert_eq!(component_config.get_option::<usize>("spec_n"), Some(6));
1002    }
1003
1004    #[test]
1005    fn test_builder_cuda_recurrent_state_manager_uses_recurrent_state_slot_cap() {
1006        let mut config = EngineConfig::default();
1007        config.backend.device = Device::CUDA(0);
1008        config.runtime.recurrent_state_max_slots = Some(2);
1009        let manager = default_recurrent_state_manager(&config)
1010            .expect("cuda product path should install admission recurrent-state manager");
1011        let spec = |request_id| RecurrentStateSpec {
1012            request_id,
1013            num_layers: 1,
1014            tensors: vec![RecurrentStateTensorSpec::new(
1015                0,
1016                "delta_state",
1017                vec![1, 1, 1],
1018                DataType::FP32,
1019            )],
1020            device: Device::CUDA(0),
1021            max_batch_slots: 1,
1022        };
1023
1024        tokio_test::block_on(manager.allocate(&spec(RequestId::new()))).unwrap();
1025        tokio_test::block_on(manager.allocate(&spec(RequestId::new()))).unwrap();
1026        let err = tokio_test::block_on(manager.allocate(&spec(RequestId::new())))
1027            .expect_err("third recurrent allocation should exceed the two-slot cap");
1028
1029        assert!(matches!(err, FerrumError::ResourceExhausted { .. }));
1030        let stats = manager.stats();
1031        assert_eq!(stats.total_batch_slots, 2);
1032        assert_eq!(stats.used_batch_slots, 2);
1033        assert_eq!(stats.allocation_failures, 1);
1034    }
1035
1036    #[test]
1037    fn test_builder_validates_layer_split_plan_without_executor_reject() {
1038        let mut config = EngineConfig::default();
1039        config.backend.backend_options.insert(
1040            "model_path".to_string(),
1041            serde_json::Value::String("/models/target".to_string()),
1042        );
1043        config.backend.backend_options.insert(
1044            "selected_distributed_strategy".to_string(),
1045            serde_json::Value::String("layer_split".to_string()),
1046        );
1047        config.backend.backend_options.insert(
1048            "requested_gpu_devices".to_string(),
1049            serde_json::json!([0, 1]),
1050        );
1051        config.backend.backend_options.insert(
1052            "selected_gpu_devices".to_string(),
1053            serde_json::json!([0, 1]),
1054        );
1055        config.backend.backend_options.insert(
1056            "selected_layer_split_plan".to_string(),
1057            serde_json::Value::String(
1058                "stage0:cuda:0:layers=0-39;stage1:cuda:1:layers=40-79".to_string(),
1059            ),
1060        );
1061        config.backend.backend_options.insert(
1062            "selected_layer_split_stages".to_string(),
1063            serde_json::json!([
1064                {"stage": 0, "device": 0, "layer_start": 0, "layer_end": 39},
1065                {"stage": 1, "device": 1, "layer_start": 40, "layer_end": 79}
1066            ]),
1067        );
1068        let component_config = ComponentConfig::from_engine_config(&config);
1069
1070        validate_layer_split_plan(&component_config).unwrap();
1071    }
1072
1073    #[tokio::test]
1074    async fn test_builder_rejects_invalid_layer_split_plan_before_executor_build() {
1075        let mut config = EngineConfig::default();
1076        config.backend.backend_options.insert(
1077            "model_path".to_string(),
1078            serde_json::Value::String("/models/target".to_string()),
1079        );
1080        config.backend.backend_options.insert(
1081            "selected_distributed_strategy".to_string(),
1082            serde_json::Value::String("layer_split".to_string()),
1083        );
1084        config.backend.backend_options.insert(
1085            "requested_gpu_devices".to_string(),
1086            serde_json::json!([0, 1]),
1087        );
1088        config.backend.backend_options.insert(
1089            "selected_gpu_devices".to_string(),
1090            serde_json::json!([0, 1]),
1091        );
1092        config.backend.backend_options.insert(
1093            "selected_layer_split_plan".to_string(),
1094            serde_json::Value::String(
1095                "stage0:cuda:0:layers=auto;stage1:cuda:1:layers=auto".to_string(),
1096            ),
1097        );
1098
1099        let err = match EngineBuilder::new(config).build().await {
1100            Ok(_) => panic!("layer_split build unexpectedly succeeded"),
1101            Err(err) => err,
1102        };
1103        assert!(err.to_string().contains("expected START-END"));
1104    }
1105
1106    #[test]
1107    fn test_resolve_defaults() {
1108        let config = EngineConfig::default();
1109        let builder = EngineBuilder::new(config);
1110
1111        assert_eq!(builder.resolve_sampler_name(), "multinomial");
1112        assert_eq!(builder.resolve_kv_cache_name(), "default");
1113    }
1114
1115    #[tokio::test]
1116    async fn test_build_with_defaults() {
1117        let config = EngineConfig::default();
1118        let result = EngineBuilder::new(config).build().await;
1119
1120        // Should succeed with stub components
1121        assert!(result.is_ok());
1122    }
1123
1124    #[tokio::test]
1125    async fn startup_preparation_runs_once_and_blocks_engine_publication_on_failure() {
1126        let success_calls = Arc::new(AtomicUsize::new(0));
1127        let success: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(StartupProbeExecutor {
1128            inner: ferrum_testkit::MockModelExecutor::instant(128),
1129            calls: Arc::clone(&success_calls),
1130            fail: false,
1131        });
1132        EngineBuilder::new(EngineConfig::default())
1133            .with_custom_executor(success)
1134            .build()
1135            .await
1136            .expect("successful startup preparation builds the engine");
1137        assert_eq!(success_calls.load(Ordering::Relaxed), 1);
1138
1139        let failure_calls = Arc::new(AtomicUsize::new(0));
1140        let failure: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(StartupProbeExecutor {
1141            inner: ferrum_testkit::MockModelExecutor::instant(128),
1142            calls: Arc::clone(&failure_calls),
1143            fail: true,
1144        });
1145        let error = EngineBuilder::new(EngineConfig::default())
1146            .with_custom_executor(failure)
1147            .build()
1148            .await
1149            .err()
1150            .expect("failed startup preparation must stop engine construction");
1151        assert!(error.to_string().contains("startup preparation rejected"));
1152        assert_eq!(failure_calls.load(Ordering::Relaxed), 1);
1153    }
1154
1155    #[tokio::test]
1156    async fn product_profile_captures_startup_events_before_engine_readiness() {
1157        let trace_path = std::env::temp_dir().join(format!(
1158            "ferrum-builder-startup-profile-{}-{}.jsonl",
1159            std::process::id(),
1160            std::time::SystemTime::now()
1161                .duration_since(std::time::UNIX_EPOCH)
1162                .unwrap()
1163                .as_nanos()
1164        ));
1165        let _ = std::fs::remove_file(&trace_path);
1166        let saw_sink_during_startup = Arc::new(AtomicBool::new(false));
1167        let executor: Arc<dyn ModelExecutor + Send + Sync> =
1168            Arc::new(ProfileStartupProbeExecutor {
1169                inner: ferrum_testkit::MockModelExecutor::instant(128),
1170                event_sink: Mutex::new(None),
1171                saw_sink_during_startup: Arc::clone(&saw_sink_during_startup),
1172            });
1173        let mut config = EngineConfig::default();
1174        config.runtime.profile_jsonl = Some(trace_path.clone());
1175        config.runtime.profile_entrypoint = Some(ferrum_types::ProfileEntrypoint::Run);
1176
1177        let engine = EngineBuilder::new(config)
1178            .with_custom_executor(executor)
1179            .build()
1180            .await
1181            .expect("profile-enabled startup builds the engine");
1182        assert!(saw_sink_during_startup.load(Ordering::Acquire));
1183        engine.shutdown().await.unwrap();
1184
1185        let startup_events = std::fs::read_to_string(&trace_path)
1186            .unwrap()
1187            .lines()
1188            .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
1189            .filter(|event| event["phase"] == "vnext.request_accepted")
1190            .count();
1191        assert_eq!(startup_events, 1);
1192        let _ = std::fs::remove_file(trace_path);
1193    }
1194
1195    #[tokio::test]
1196    async fn product_without_profile_does_not_attach_execution_event_sink() {
1197        let saw_sink_during_startup = Arc::new(AtomicBool::new(false));
1198        let executor: Arc<dyn ModelExecutor + Send + Sync> =
1199            Arc::new(ProfileStartupProbeExecutor {
1200                inner: ferrum_testkit::MockModelExecutor::instant(128),
1201                event_sink: Mutex::new(None),
1202                saw_sink_during_startup: Arc::clone(&saw_sink_during_startup),
1203            });
1204
1205        let engine = EngineBuilder::new(EngineConfig::default())
1206            .with_custom_executor(executor)
1207            .build()
1208            .await
1209            .expect("profile-disabled engine builds");
1210        assert!(!saw_sink_during_startup.load(Ordering::Acquire));
1211        engine.shutdown().await.unwrap();
1212    }
1213
1214    #[tokio::test]
1215    async fn plan_runtime_without_resolved_plan_rejects_before_legacy_kv_factory() {
1216        let calls = Arc::new(AtomicUsize::new(0));
1217        let registry = Arc::new(ComponentRegistry::with_defaults());
1218        registry.register_kv_cache_factory(
1219            "default",
1220            Arc::new(CountingKvFactory {
1221                calls: Arc::clone(&calls),
1222            }),
1223        );
1224        let executor: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(PlanRuntimeBuilderExecutor {
1225            inner: ferrum_testkit::MockModelExecutor::instant(128),
1226        });
1227
1228        let result = EngineBuilder::with_registry(EngineConfig::default(), registry)
1229            .with_custom_executor(executor)
1230            .build()
1231            .await;
1232
1233        let error = result
1234            .err()
1235            .expect("plan runtime without a resolved plan must fail closed");
1236        assert!(error
1237            .to_string()
1238            .contains("authoritative ResolvedModelPlan"));
1239        assert_eq!(calls.load(Ordering::Relaxed), 0);
1240    }
1241
1242    #[tokio::test]
1243    async fn plan_runtime_build_rejects_legacy_resource_override() {
1244        let executor: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(PlanRuntimeBuilderExecutor {
1245            inner: ferrum_testkit::MockModelExecutor::instant(128),
1246        });
1247        let kv_cache: Arc<dyn KvCacheManager + Send + Sync> =
1248            Arc::new(ferrum_testkit::MockKvCacheManager::new(8));
1249
1250        let error = EngineBuilder::new(EngineConfig::default())
1251            .with_custom_executor(executor)
1252            .with_custom_kv_cache(kv_cache)
1253            .build()
1254            .await
1255            .err()
1256            .expect("plan runtime must reject a legacy KV manager override");
1257
1258        assert!(error
1259            .to_string()
1260            .contains("cannot be combined with a legacy engine KV-cache override"));
1261    }
1262}