1use 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
22pub struct EngineBuilder {
30 registry: Arc<ComponentRegistry>,
32 config: EngineConfig,
34 model_sources: Option<Arc<ProductionModelSourceBundle>>,
36 prepared_model: Option<Arc<PreparedProductionModel>>,
39 tokenizer_name: Option<String>,
41 sampler_name: Option<String>,
43 scheduler_name: Option<String>,
45 kv_cache_name: Option<String>,
47 executor_name: Option<String>,
49 custom_tokenizer: Option<Arc<dyn Tokenizer + Send + Sync>>,
51 custom_sampler: Option<Arc<dyn Sampler + Send + Sync>>,
53 custom_scheduler: Option<Arc<dyn Scheduler + Send + Sync>>,
55 custom_kv_cache: Option<Arc<dyn KvCacheManager + Send + Sync>>,
57 custom_recurrent_state_manager: Option<Arc<dyn RecurrentStateManager + Send + Sync>>,
59 custom_executor: Option<Arc<dyn ModelExecutor + Send + Sync>>,
61}
62
63impl EngineBuilder {
64 pub fn new(config: EngineConfig) -> Self {
66 Self::with_registry(config, crate::registry::global_registry())
67 }
68
69 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 pub fn with_tokenizer(mut self, name: impl Into<String>) -> Self {
104 self.tokenizer_name = Some(name.into());
105 self
106 }
107
108 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 pub fn with_sampler(mut self, name: impl Into<String>) -> Self {
116 self.sampler_name = Some(name.into());
117 self
118 }
119
120 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 pub fn with_scheduler(mut self, name: impl Into<String>) -> Self {
128 self.scheduler_name = Some(name.into());
129 self
130 }
131
132 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 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 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 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 pub fn with_executor(mut self, name: impl Into<String>) -> Self {
161 self.executor_name = Some(name.into());
162 self
163 }
164
165 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 fn resolve_tokenizer_name(&self) -> String {
173 if let Some(ref name) = self.tokenizer_name {
174 return name.clone();
175 }
176
177 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 fn resolve_sampler_name(&self) -> String {
187 if let Some(ref name) = self.sampler_name {
188 return name.clone();
189 }
190
191 "multinomial".to_string()
193 }
194
195 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(), }
206 }
207
208 fn resolve_executor_name(&self) -> String {
210 if let Some(ref name) = self.executor_name {
211 return name.clone();
212 }
213
214 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 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 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 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 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 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 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 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 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 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 let tensor_factory: Arc<dyn TensorFactory> = Arc::new(
397 crate::tensor_factory::candle::CandleTensorFactory::new(config.backend.device.clone()),
398 );
399
400 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 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 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
570pub async fn create_engine(
572 config: EngineConfig,
573) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
574 EngineBuilder::new(config).build().await
575}
576
577pub 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
588pub 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#[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 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}