use crate::registry::{ComponentConfig, ComponentRegistry};
use ferrum_interfaces::engine::{InferenceEngine, LlmInferenceEngine};
use ferrum_interfaces::{
KvCacheManager, ModelExecutor, RecurrentStateManager, Sampler, SchedulerInterface as Scheduler,
TensorFactory, Tokenizer,
};
use ferrum_models::vnext::{PreparedProductionModel, ProductionModelSourceBundle};
use ferrum_types::{EngineConfig, FerrumError, Result};
use std::sync::Arc;
use tracing::{debug, info};
pub struct EngineBuilder {
registry: Arc<ComponentRegistry>,
config: EngineConfig,
model_sources: Option<Arc<ProductionModelSourceBundle>>,
prepared_model: Option<Arc<PreparedProductionModel>>,
tokenizer_name: Option<String>,
sampler_name: Option<String>,
scheduler_name: Option<String>,
kv_cache_name: Option<String>,
executor_name: Option<String>,
custom_tokenizer: Option<Arc<dyn Tokenizer + Send + Sync>>,
custom_sampler: Option<Arc<dyn Sampler + Send + Sync>>,
custom_scheduler: Option<Arc<dyn Scheduler + Send + Sync>>,
custom_kv_cache: Option<Arc<dyn KvCacheManager + Send + Sync>>,
custom_recurrent_state_manager: Option<Arc<dyn RecurrentStateManager + Send + Sync>>,
custom_executor: Option<Arc<dyn ModelExecutor + Send + Sync>>,
}
impl EngineBuilder {
pub fn new(config: EngineConfig) -> Self {
Self::with_registry(config, crate::registry::global_registry())
}
pub fn with_registry(config: EngineConfig, registry: Arc<ComponentRegistry>) -> Self {
Self {
registry,
config,
model_sources: None,
prepared_model: None,
tokenizer_name: None,
sampler_name: None,
scheduler_name: None,
kv_cache_name: None,
executor_name: None,
custom_tokenizer: None,
custom_sampler: None,
custom_scheduler: None,
custom_kv_cache: None,
custom_recurrent_state_manager: None,
custom_executor: None,
}
}
pub fn with_model_sources(mut self, sources: Arc<ProductionModelSourceBundle>) -> Self {
self.model_sources = Some(sources);
self.prepared_model = None;
self
}
pub fn with_prepared_model(mut self, prepared: Arc<PreparedProductionModel>) -> Self {
self.model_sources = Some(Arc::clone(prepared.sources()));
self.prepared_model = Some(prepared);
self
}
pub fn with_tokenizer(mut self, name: impl Into<String>) -> Self {
self.tokenizer_name = Some(name.into());
self
}
pub fn with_custom_tokenizer(mut self, tokenizer: Arc<dyn Tokenizer + Send + Sync>) -> Self {
self.custom_tokenizer = Some(tokenizer);
self
}
pub fn with_sampler(mut self, name: impl Into<String>) -> Self {
self.sampler_name = Some(name.into());
self
}
pub fn with_custom_sampler(mut self, sampler: Arc<dyn Sampler + Send + Sync>) -> Self {
self.custom_sampler = Some(sampler);
self
}
pub fn with_scheduler(mut self, name: impl Into<String>) -> Self {
self.scheduler_name = Some(name.into());
self
}
pub fn with_custom_scheduler(mut self, scheduler: Arc<dyn Scheduler + Send + Sync>) -> Self {
self.custom_scheduler = Some(scheduler);
self
}
pub fn with_kv_cache(mut self, name: impl Into<String>) -> Self {
self.kv_cache_name = Some(name.into());
self
}
pub fn with_custom_kv_cache(mut self, kv_cache: Arc<dyn KvCacheManager + Send + Sync>) -> Self {
self.custom_kv_cache = Some(kv_cache);
self
}
pub fn with_custom_recurrent_state_manager(
mut self,
manager: Arc<dyn RecurrentStateManager + Send + Sync>,
) -> Self {
self.custom_recurrent_state_manager = Some(manager);
self
}
pub fn with_executor(mut self, name: impl Into<String>) -> Self {
self.executor_name = Some(name.into());
self
}
pub fn with_custom_executor(mut self, executor: Arc<dyn ModelExecutor + Send + Sync>) -> Self {
self.custom_executor = Some(executor);
self
}
fn resolve_tokenizer_name(&self) -> String {
if let Some(ref name) = self.tokenizer_name {
return name.clone();
}
if self.has_typed_model_path() || self.config.runtime.model_path.is_some() {
return "huggingface".to_string();
}
"stub".to_string()
}
fn resolve_sampler_name(&self) -> String {
if let Some(ref name) = self.sampler_name {
return name.clone();
}
"multinomial".to_string()
}
fn resolve_kv_cache_name(&self) -> String {
if let Some(ref name) = self.kv_cache_name {
return name.clone();
}
match self.config.kv_cache.cache_type {
ferrum_types::KvCacheType::Contiguous => "default".to_string(),
ferrum_types::KvCacheType::Paged => "paged".to_string(),
ferrum_types::KvCacheType::Tree => "default".to_string(), }
}
fn resolve_executor_name(&self) -> String {
if let Some(ref name) = self.executor_name {
return name.clone();
}
if self.has_typed_model_path() || self.config.runtime.model_path.is_some() {
return "llm".to_string();
}
"stub".to_string()
}
fn has_typed_model_path(&self) -> bool {
self.model_sources.is_some()
|| self
.config
.backend
.backend_options
.get("model_path")
.and_then(|value| value.as_str())
.is_some()
}
pub async fn build(self) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
info!(
"Building inference engine for model: {}",
self.config.model.model_id
);
if self.scheduler_name.is_some() || self.custom_scheduler.is_some() {
return Err(FerrumError::config(
"EngineBuilder scheduler component overrides are no longer accepted; configure the typed EngineConfig.scheduler used by ContinuousBatchScheduler",
));
}
let tokenizer_name = self.resolve_tokenizer_name();
let sampler_name = self.resolve_sampler_name();
let kv_cache_name = self.resolve_kv_cache_name();
let executor_name = self.resolve_executor_name();
let explicit_kv_cache_override = self.kv_cache_name.is_some();
let component_config = ComponentConfig::from_engine_config_and_product_model(
&self.config,
self.model_sources.clone(),
self.prepared_model.clone(),
);
validate_layer_split_plan(&component_config)?;
let typed_model_path = component_config.get_string_option("model_path");
let has_model_path = typed_model_path.is_some() || self.config.runtime.model_path.is_some();
let registry = self.registry.clone();
let config = self.config;
let custom_tokenizer = self.custom_tokenizer;
let custom_sampler = self.custom_sampler;
let custom_kv_cache = self.custom_kv_cache;
let custom_recurrent_state_manager = self.custom_recurrent_state_manager;
let custom_executor = self.custom_executor;
let tokenizer = if let Some(tokenizer) = custom_tokenizer {
debug!("Using custom tokenizer");
tokenizer
} else {
debug!("Creating tokenizer: {}", tokenizer_name);
match registry
.create_tokenizer(&tokenizer_name, &component_config)
.await
{
Ok(t) => t,
Err(e) => {
if has_model_path {
return Err(FerrumError::config(format!(
"Failed to create tokenizer '{}' in model mode: {}",
tokenizer_name, e
)));
}
tracing::warn!(
"Failed to create tokenizer '{}': {}, falling back to stub",
tokenizer_name,
e
);
registry.create_tokenizer("stub", &component_config).await?
}
}
};
let sampler = if let Some(sampler) = custom_sampler {
debug!("Using custom sampler");
sampler
} else {
debug!("Creating sampler: {}", sampler_name);
registry
.create_sampler(&sampler_name, &component_config)
.await?
};
let executor = if let Some(executor) = custom_executor {
debug!("Using custom executor");
executor
} else {
debug!("Creating executor: {}", executor_name);
match registry
.create_executor(&executor_name, &component_config)
.await
{
Ok(e) => e,
Err(err) => {
if has_model_path {
return Err(FerrumError::config(format!(
"Failed to create executor '{}' in model mode: {}",
executor_name, err
)));
}
tracing::warn!(
"Failed to create executor '{}': {}, falling back to stub",
executor_name,
err
);
registry.create_executor("stub", &component_config).await?
}
}
};
let execution_resource_authority = executor.execution_resource_authority();
let (kv_cache, recurrent_state_manager) = match execution_resource_authority {
ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime => {
if explicit_kv_cache_override || custom_kv_cache.is_some() {
return Err(FerrumError::config(
"plan runtime cannot be combined with a legacy engine KV-cache override",
));
}
if custom_recurrent_state_manager.is_some() {
return Err(FerrumError::config(
"plan runtime cannot be combined with a legacy engine recurrent-state manager",
));
}
if executor.resolved_model_plan().is_none() {
return Err(FerrumError::config(
"plan-runtime executor did not expose its authoritative ResolvedModelPlan",
));
}
(None, None)
}
ferrum_interfaces::model_executor::ExecutionResourceAuthority::LegacyEngine => {
let kv_cache = if let Some(kv_cache) = custom_kv_cache {
debug!("Using custom KV cache");
kv_cache
} else {
debug!("Creating KV cache: {}", kv_cache_name);
registry
.create_kv_cache(&kv_cache_name, &component_config)
.await?
};
let recurrent_state_manager = custom_recurrent_state_manager
.or_else(|| default_recurrent_state_manager(&config));
(Some(kv_cache), recurrent_state_manager)
}
};
info!("All components created, building ContinuousBatchEngine");
let cb_scheduler = Arc::new(
ferrum_scheduler::implementations::ContinuousBatchScheduler::new(
config.scheduler.clone(),
),
);
let tensor_factory: Arc<dyn TensorFactory> = Arc::new(
crate::tensor_factory::candle::CandleTensorFactory::new(config.backend.device.clone()),
);
let spec_draft = component_config
.get_string_option("spec_draft")
.or_else(|| config.runtime.spec_draft.clone());
if execution_resource_authority
== ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime
&& spec_draft.is_some()
{
return Err(FerrumError::unsupported(
"speculative decoding is not yet part of the plan-runtime contract",
));
}
let spec_n = component_config
.get_option::<usize>("spec_n")
.unwrap_or(config.runtime.spec_n.unwrap_or(4));
let (draft_executor, spec_config) = match spec_draft.as_ref() {
Some(draft_path) => {
info!("Speculative decoding: loading draft model from {draft_path}");
let mut draft_cfg = component_config.clone();
draft_cfg.component_options.insert(
"model_path".to_string(),
serde_json::Value::String(draft_path.to_string()),
);
let draft = registry
.create_executor(&executor_name, &draft_cfg)
.await
.map_err(|error| {
FerrumError::config(format!(
"requested speculative draft executor failed to load: {error}"
))
})?;
if draft.execution_resource_authority() != execution_resource_authority {
return Err(FerrumError::config(
"target and speculative draft executors declare different resource authority",
));
}
(
Some(draft),
Some(crate::speculative::SpeculativeDecodingConfig {
num_speculative_tokens: spec_n,
temperature: 1.0,
}),
)
}
_ => (None, None),
};
let engine = match execution_resource_authority {
ferrum_interfaces::model_executor::ExecutionResourceAuthority::PlanRuntime => {
crate::ContinuousBatchEngine::new_plan_runtime(
config,
cb_scheduler,
tokenizer,
sampler,
Arc::clone(&executor),
tensor_factory,
)?
}
ferrum_interfaces::model_executor::ExecutionResourceAuthority::LegacyEngine => {
let kv_cache = kv_cache.ok_or_else(|| {
FerrumError::internal("legacy-engine composition lost its KV-cache manager")
})?;
crate::ContinuousBatchEngine::new_with_speculation_and_recurrent_state_manager(
config,
cb_scheduler,
tokenizer,
sampler,
kv_cache,
Arc::clone(&executor),
tensor_factory,
draft_executor.clone(),
spec_config,
recurrent_state_manager,
)?
}
};
let startup_result = async {
executor.prepare_startup().await?;
if let Some(draft) = draft_executor.as_ref() {
draft.prepare_startup().await?;
}
Ok(())
}
.await;
if let Err(startup_error) = startup_result {
if let Err(shutdown_error) = engine.shutdown().await {
tracing::warn!(
"Failed to close engine resources after startup rejection: {shutdown_error}"
);
}
return Err(startup_error);
}
Ok(Box::new(engine))
}
}
fn default_recurrent_state_manager(
config: &EngineConfig,
) -> Option<Arc<dyn RecurrentStateManager + Send + Sync>> {
let total_batch_slots = config
.runtime
.recurrent_state_max_slots
.unwrap_or(usize::MAX);
Some(
Arc::new(crate::recurrent_state::InMemoryRecurrentStateManager::new(
crate::recurrent_state::InMemoryRecurrentStateConfig {
total_memory_bytes: usize::MAX,
total_batch_slots,
},
)) as Arc<dyn RecurrentStateManager + Send + Sync>,
)
}
fn validate_layer_split_plan(component_config: &ComponentConfig) -> Result<()> {
if component_config
.get_string_option("selected_distributed_strategy")
.as_deref()
!= Some("layer_split")
{
return Ok(());
}
let requested = component_config
.get_option::<Vec<usize>>("requested_gpu_devices")
.unwrap_or_default();
let selected = component_config
.get_option::<Vec<usize>>("selected_gpu_devices")
.unwrap_or_default();
let plan_raw = component_config.get_string_option("selected_layer_split_plan");
let parsed_plan = if let Some(stages) = component_config
.component_options
.get("selected_layer_split_stages")
{
crate::layer_split::parse_layer_split_stage_documents(stages)?
} else {
let plan_raw = plan_raw.as_deref().ok_or_else(|| {
FerrumError::config(
"selected_distributed_strategy=layer_split requires selected_layer_split_plan",
)
})?;
crate::layer_split::parse_layer_split_plan(plan_raw)?
};
crate::layer_split::validate_layer_split_plan_for_devices(&parsed_plan, &selected)?;
let execution_plan = parsed_plan.to_execution_plan();
let stage_ranges = execution_plan
.layer_distribution
.stage_layers
.iter()
.map(|range| format!("{}-{}", range.start, range.end.saturating_sub(1)))
.collect::<Vec<_>>()
.join(",");
let plan_label = plan_raw.unwrap_or_else(|| format!("{:?}", parsed_plan.stages));
tracing::info!(
"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={}",
parsed_plan.total_layers(),
execution_plan.parallel_config.pipeline_parallel_size,
execution_plan.parallel_config.communication_backend,
);
Ok(())
}
pub async fn create_engine(
config: EngineConfig,
) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
EngineBuilder::new(config).build().await
}
pub async fn create_product_engine(
config: EngineConfig,
sources: Arc<ProductionModelSourceBundle>,
) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
EngineBuilder::new(config)
.with_model_sources(sources)
.build()
.await
}
pub async fn create_prepared_product_engine(
config: EngineConfig,
prepared: Arc<PreparedProductionModel>,
) -> Result<Box<dyn LlmInferenceEngine + Send + Sync>> {
EngineBuilder::new(config)
.with_prepared_model(prepared)
.build()
.await
}
#[cfg(test)]
mod tests {
use super::*;
use ferrum_interfaces::{
model_executor::{
DecodeInput, DecodeOutput, ExecutionResourceAuthority, ExecutorCapabilities,
ExecutorStatus, PlanRuntimeResourceSnapshot, PrefillInput, PrefillOutput,
},
vnext::ExecutionEventSink,
RecurrentStateHandle, RecurrentStateManager, RecurrentStateManagerStats,
RecurrentStateSpec, RecurrentStateTensorSpec,
};
use ferrum_types::{DataType, Device, RequestId};
use std::sync::{
atomic::{AtomicBool, AtomicUsize, Ordering},
Mutex,
};
#[derive(Debug)]
struct NoopRecurrentStateManager;
struct PlanRuntimeBuilderExecutor {
inner: ferrum_testkit::MockModelExecutor,
}
struct StartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor,
calls: Arc<AtomicUsize>,
fail: bool,
}
struct ProfileStartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor,
event_sink: Mutex<Option<Arc<dyn ExecutionEventSink>>>,
saw_sink_during_startup: Arc<AtomicBool>,
}
#[async_trait::async_trait]
impl ModelExecutor for StartupProbeExecutor {
fn info(&self) -> &ferrum_types::ModelInfo {
self.inner.info()
}
async fn prepare_startup(&self) -> Result<()> {
self.calls.fetch_add(1, Ordering::Relaxed);
if self.fail {
return Err(FerrumError::backend("startup preparation rejected"));
}
Ok(())
}
async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
self.inner.prefill(input).await
}
async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
self.inner.decode(input).await
}
fn capabilities(&self) -> ExecutorCapabilities {
self.inner.capabilities()
}
fn status(&self) -> ExecutorStatus {
self.inner.status()
}
}
#[async_trait::async_trait]
impl ModelExecutor for ProfileStartupProbeExecutor {
fn info(&self) -> &ferrum_types::ModelInfo {
self.inner.info()
}
async fn prepare_startup(&self) -> Result<()> {
use ferrum_interfaces::vnext::{ExecutionEventEmitter, TrustedExecutionEventContext};
let sink = self
.event_sink
.lock()
.expect("profile startup probe sink lock")
.clone();
self.saw_sink_during_startup
.store(sink.is_some(), Ordering::Release);
let Some(sink) = sink else {
return Ok(());
};
let (run_id, request_id, event) = startup_profile_test_event();
ExecutionEventEmitter::from_shared(sink, run_id.clone(), request_id.clone())
.emit(
event,
&TrustedExecutionEventContext::pre_plan(&run_id, &request_id),
)
.map_err(|error| {
FerrumError::internal(format!("emit startup profile probe: {error}"))
})
}
fn attach_execution_event_sink(&self, sink: Arc<dyn ExecutionEventSink>) {
*self
.event_sink
.lock()
.expect("profile startup probe sink lock") = Some(sink);
}
async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
self.inner.prefill(input).await
}
async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
self.inner.decode(input).await
}
fn capabilities(&self) -> ExecutorCapabilities {
self.inner.capabilities()
}
fn status(&self) -> ExecutorStatus {
self.inner.status()
}
}
fn startup_profile_test_event() -> (
ferrum_interfaces::vnext::RunId,
ferrum_interfaces::vnext::RequestIdentity,
ferrum_interfaces::vnext::ExecutionEvent,
) {
use ferrum_interfaces::vnext::{
ExecutionEvent, ExecutionEventDetail, ExecutionEventKind, ExecutionIdentityEnvelope,
ExecutionIdentityParts, ExecutionPhase, MonotonicTimestamp, RequestIdentity, RunId,
SpanId, EXECUTION_IDENTITY_VERSION,
};
let run_id = RunId::new("run.vnext.builder-startup-profile").unwrap();
let request_id = RequestIdentity::new("request.vnext.builder-startup-profile").unwrap();
let event = ExecutionEvent::new(
MonotonicTimestamp {
nanos_since_run_start: 1,
},
ExecutionPhase::Resolution,
ExecutionEventKind::RequestAccepted,
ExecutionIdentityEnvelope::new(ExecutionIdentityParts {
version: EXECUTION_IDENTITY_VERSION,
run_id: run_id.clone(),
request_id: request_id.clone(),
sequence: 1,
plan_id: None,
plan_hash: None,
frame_id: None,
node_invocation_id: None,
node_id: None,
operation_id: None,
provider_id: None,
device_id: None,
resource_pool_id: None,
resource_pool_identity_fingerprint: None,
provisioning_run_id: None,
provisioning_request_id: None,
transaction_id: None,
active_sequence_slot: None,
admission_generation: None,
activation_epoch: None,
runtime_implementation_fingerprint: None,
active_sequence_fingerprint: None,
completed_sequence_fingerprint: None,
aborted_sequence_fingerprint: None,
resource_id: None,
resource_generation: None,
resource_batch_fingerprint: None,
span_id: SpanId::new("vnext/request/builder-startup-profile").unwrap(),
parent_span_id: None,
async_links: Vec::new(),
})
.unwrap(),
ExecutionEventDetail::None,
)
.unwrap();
(run_id, request_id, event)
}
#[async_trait::async_trait]
impl ModelExecutor for PlanRuntimeBuilderExecutor {
fn info(&self) -> &ferrum_types::ModelInfo {
self.inner.info()
}
fn execution_resource_authority(&self) -> ExecutionResourceAuthority {
ExecutionResourceAuthority::PlanRuntime
}
fn plan_runtime_resource_snapshot(&self) -> Result<Option<PlanRuntimeResourceSnapshot>> {
PlanRuntimeResourceSnapshot::new(1_000, 900, 700, 700, 400, 300, 200, 0, 0).map(Some)
}
async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput> {
self.inner.prefill(input).await
}
async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput> {
self.inner.decode(input).await
}
fn capabilities(&self) -> ExecutorCapabilities {
self.inner.capabilities()
}
fn status(&self) -> ExecutorStatus {
self.inner.status()
}
}
struct CountingKvFactory {
calls: Arc<AtomicUsize>,
}
#[async_trait::async_trait]
impl crate::registry::ComponentFactory<Arc<dyn KvCacheManager + Send + Sync>>
for CountingKvFactory
{
async fn create(
&self,
_config: &ComponentConfig,
) -> Result<Arc<dyn KvCacheManager + Send + Sync>> {
self.calls.fetch_add(1, Ordering::Relaxed);
Ok(Arc::new(ferrum_testkit::MockKvCacheManager::new(8)))
}
fn metadata(&self) -> crate::registry::ComponentMetadata {
crate::registry::ComponentMetadata::default()
}
}
#[async_trait::async_trait]
impl RecurrentStateManager for NoopRecurrentStateManager {
async fn allocate(
&self,
_spec: &RecurrentStateSpec,
) -> Result<Arc<dyn RecurrentStateHandle>> {
Err(FerrumError::unsupported(
"noop recurrent-state manager does not allocate",
))
}
async fn deallocate(&self, _request_id: RequestId) -> Result<()> {
Ok(())
}
fn can_allocate(&self, _spec: &RecurrentStateSpec) -> bool {
false
}
fn get_handle(&self, _request_id: RequestId) -> Option<Arc<dyn RecurrentStateHandle>> {
None
}
fn list_handles(&self) -> Vec<(RequestId, Arc<dyn RecurrentStateHandle>)> {
Vec::new()
}
fn stats(&self) -> RecurrentStateManagerStats {
RecurrentStateManagerStats {
total_memory_bytes: 0,
used_memory_bytes: 0,
active_states: 0,
active_state_tensors: 0,
total_batch_slots: 0,
used_batch_slots: 0,
allocation_count: 0,
allocation_failures: 0,
eviction_count: 0,
}
}
async fn reset(&self) -> Result<()> {
Ok(())
}
}
#[test]
fn test_builder_creation() {
let config = EngineConfig::default();
let builder = EngineBuilder::new(config);
assert!(builder.tokenizer_name.is_none());
assert!(builder.custom_recurrent_state_manager.is_none());
}
#[test]
fn test_builder_with_overrides() {
let config = EngineConfig::default();
let builder = EngineBuilder::new(config)
.with_tokenizer("custom_tokenizer")
.with_sampler("greedy")
.with_scheduler("priority")
.with_kv_cache("paged")
.with_executor("custom_executor");
assert_eq!(builder.tokenizer_name, Some("custom_tokenizer".to_string()));
assert_eq!(builder.sampler_name, Some("greedy".to_string()));
assert_eq!(builder.scheduler_name, Some("priority".to_string()));
assert_eq!(builder.kv_cache_name, Some("paged".to_string()));
assert_eq!(builder.executor_name, Some("custom_executor".to_string()));
}
#[test]
fn test_builder_with_custom_recurrent_state_manager() {
let config = EngineConfig::default();
let manager = Arc::new(NoopRecurrentStateManager);
let builder = EngineBuilder::new(config).with_custom_recurrent_state_manager(manager);
assert!(builder.custom_recurrent_state_manager.is_some());
}
#[test]
fn test_builder_typed_model_path_selects_model_components() {
let mut config = EngineConfig::default();
config.backend.backend_options.insert(
"model_path".to_string(),
serde_json::Value::String("/models/target".to_string()),
);
let builder = EngineBuilder::new(config);
assert!(builder.has_typed_model_path());
assert_eq!(builder.resolve_tokenizer_name(), "huggingface");
assert_eq!(builder.resolve_executor_name(), "llm");
}
#[test]
fn test_builder_retains_one_typed_source_bundle_for_components() {
let root = std::env::temp_dir().join(format!(
"ferrum-builder-source-bundle-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&root).unwrap();
std::fs::write(
root.join("config.json"),
br#"{"architectures":["Fixture"]}"#,
)
.unwrap();
std::fs::write(root.join("tokenizer.json"), br#"{"version":"1.0"}"#).unwrap();
std::fs::write(root.join("model.safetensors"), b"fixture").unwrap();
let original = ferrum_interfaces::vnext::OriginalModelSource {
kind: ferrum_interfaces::vnext::ModelSourceKind::LocalDirectory,
location: root.display().to_string(),
requested_revision: None,
};
let sources = Arc::new(
ProductionModelSourceBundle::open(
&root,
&root,
ferrum_models::vnext::ProductionWeightArtifact::safetensors_directory(&root),
ferrum_interfaces::vnext::OriginalModelSources {
semantic: original.clone(),
tokenizer: original.clone(),
weights: original,
},
)
.unwrap(),
);
let builder =
EngineBuilder::new(EngineConfig::default()).with_model_sources(Arc::clone(&sources));
assert!(builder.has_typed_model_path());
assert!(Arc::ptr_eq(
builder.model_sources.as_ref().unwrap(),
&sources
));
assert_eq!(builder.resolve_tokenizer_name(), "huggingface");
assert_eq!(builder.resolve_executor_name(), "llm");
std::fs::remove_dir_all(root).unwrap();
}
#[test]
fn test_builder_typed_spec_options_parse_from_component_config() {
let mut config = EngineConfig::default();
config.backend.backend_options.insert(
"model_path".to_string(),
serde_json::Value::String("/models/target".to_string()),
);
config.backend.backend_options.insert(
"spec_draft".to_string(),
serde_json::Value::String("/models/draft".to_string()),
);
config.backend.backend_options.insert(
"spec_n".to_string(),
serde_json::Value::Number(serde_json::Number::from(6)),
);
let component_config = ComponentConfig::from_engine_config(&config);
assert_eq!(
component_config.get_string_option("spec_draft").as_deref(),
Some("/models/draft")
);
assert_eq!(component_config.get_option::<usize>("spec_n"), Some(6));
}
#[test]
fn test_builder_cuda_recurrent_state_manager_uses_recurrent_state_slot_cap() {
let mut config = EngineConfig::default();
config.backend.device = Device::CUDA(0);
config.runtime.recurrent_state_max_slots = Some(2);
let manager = default_recurrent_state_manager(&config)
.expect("cuda product path should install admission recurrent-state manager");
let spec = |request_id| RecurrentStateSpec {
request_id,
num_layers: 1,
tensors: vec![RecurrentStateTensorSpec::new(
0,
"delta_state",
vec![1, 1, 1],
DataType::FP32,
)],
device: Device::CUDA(0),
max_batch_slots: 1,
};
tokio_test::block_on(manager.allocate(&spec(RequestId::new()))).unwrap();
tokio_test::block_on(manager.allocate(&spec(RequestId::new()))).unwrap();
let err = tokio_test::block_on(manager.allocate(&spec(RequestId::new())))
.expect_err("third recurrent allocation should exceed the two-slot cap");
assert!(matches!(err, FerrumError::ResourceExhausted { .. }));
let stats = manager.stats();
assert_eq!(stats.total_batch_slots, 2);
assert_eq!(stats.used_batch_slots, 2);
assert_eq!(stats.allocation_failures, 1);
}
#[test]
fn test_builder_validates_layer_split_plan_without_executor_reject() {
let mut config = EngineConfig::default();
config.backend.backend_options.insert(
"model_path".to_string(),
serde_json::Value::String("/models/target".to_string()),
);
config.backend.backend_options.insert(
"selected_distributed_strategy".to_string(),
serde_json::Value::String("layer_split".to_string()),
);
config.backend.backend_options.insert(
"requested_gpu_devices".to_string(),
serde_json::json!([0, 1]),
);
config.backend.backend_options.insert(
"selected_gpu_devices".to_string(),
serde_json::json!([0, 1]),
);
config.backend.backend_options.insert(
"selected_layer_split_plan".to_string(),
serde_json::Value::String(
"stage0:cuda:0:layers=0-39;stage1:cuda:1:layers=40-79".to_string(),
),
);
config.backend.backend_options.insert(
"selected_layer_split_stages".to_string(),
serde_json::json!([
{"stage": 0, "device": 0, "layer_start": 0, "layer_end": 39},
{"stage": 1, "device": 1, "layer_start": 40, "layer_end": 79}
]),
);
let component_config = ComponentConfig::from_engine_config(&config);
validate_layer_split_plan(&component_config).unwrap();
}
#[tokio::test]
async fn test_builder_rejects_invalid_layer_split_plan_before_executor_build() {
let mut config = EngineConfig::default();
config.backend.backend_options.insert(
"model_path".to_string(),
serde_json::Value::String("/models/target".to_string()),
);
config.backend.backend_options.insert(
"selected_distributed_strategy".to_string(),
serde_json::Value::String("layer_split".to_string()),
);
config.backend.backend_options.insert(
"requested_gpu_devices".to_string(),
serde_json::json!([0, 1]),
);
config.backend.backend_options.insert(
"selected_gpu_devices".to_string(),
serde_json::json!([0, 1]),
);
config.backend.backend_options.insert(
"selected_layer_split_plan".to_string(),
serde_json::Value::String(
"stage0:cuda:0:layers=auto;stage1:cuda:1:layers=auto".to_string(),
),
);
let err = match EngineBuilder::new(config).build().await {
Ok(_) => panic!("layer_split build unexpectedly succeeded"),
Err(err) => err,
};
assert!(err.to_string().contains("expected START-END"));
}
#[test]
fn test_resolve_defaults() {
let config = EngineConfig::default();
let builder = EngineBuilder::new(config);
assert_eq!(builder.resolve_sampler_name(), "multinomial");
assert_eq!(builder.resolve_kv_cache_name(), "default");
}
#[tokio::test]
async fn test_build_with_defaults() {
let config = EngineConfig::default();
let result = EngineBuilder::new(config).build().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn startup_preparation_runs_once_and_blocks_engine_publication_on_failure() {
let success_calls = Arc::new(AtomicUsize::new(0));
let success: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(StartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
calls: Arc::clone(&success_calls),
fail: false,
});
EngineBuilder::new(EngineConfig::default())
.with_custom_executor(success)
.build()
.await
.expect("successful startup preparation builds the engine");
assert_eq!(success_calls.load(Ordering::Relaxed), 1);
let failure_calls = Arc::new(AtomicUsize::new(0));
let failure: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(StartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
calls: Arc::clone(&failure_calls),
fail: true,
});
let error = EngineBuilder::new(EngineConfig::default())
.with_custom_executor(failure)
.build()
.await
.err()
.expect("failed startup preparation must stop engine construction");
assert!(error.to_string().contains("startup preparation rejected"));
assert_eq!(failure_calls.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn product_profile_captures_startup_events_before_engine_readiness() {
let trace_path = std::env::temp_dir().join(format!(
"ferrum-builder-startup-profile-{}-{}.jsonl",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
let _ = std::fs::remove_file(&trace_path);
let saw_sink_during_startup = Arc::new(AtomicBool::new(false));
let executor: Arc<dyn ModelExecutor + Send + Sync> =
Arc::new(ProfileStartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
event_sink: Mutex::new(None),
saw_sink_during_startup: Arc::clone(&saw_sink_during_startup),
});
let mut config = EngineConfig::default();
config.runtime.profile_jsonl = Some(trace_path.clone());
config.runtime.profile_entrypoint = Some(ferrum_types::ProfileEntrypoint::Run);
let engine = EngineBuilder::new(config)
.with_custom_executor(executor)
.build()
.await
.expect("profile-enabled startup builds the engine");
assert!(saw_sink_during_startup.load(Ordering::Acquire));
engine.shutdown().await.unwrap();
let startup_events = std::fs::read_to_string(&trace_path)
.unwrap()
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.filter(|event| event["phase"] == "vnext.request_accepted")
.count();
assert_eq!(startup_events, 1);
let _ = std::fs::remove_file(trace_path);
}
#[tokio::test]
async fn product_without_profile_does_not_attach_execution_event_sink() {
let saw_sink_during_startup = Arc::new(AtomicBool::new(false));
let executor: Arc<dyn ModelExecutor + Send + Sync> =
Arc::new(ProfileStartupProbeExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
event_sink: Mutex::new(None),
saw_sink_during_startup: Arc::clone(&saw_sink_during_startup),
});
let engine = EngineBuilder::new(EngineConfig::default())
.with_custom_executor(executor)
.build()
.await
.expect("profile-disabled engine builds");
assert!(!saw_sink_during_startup.load(Ordering::Acquire));
engine.shutdown().await.unwrap();
}
#[tokio::test]
async fn plan_runtime_without_resolved_plan_rejects_before_legacy_kv_factory() {
let calls = Arc::new(AtomicUsize::new(0));
let registry = Arc::new(ComponentRegistry::with_defaults());
registry.register_kv_cache_factory(
"default",
Arc::new(CountingKvFactory {
calls: Arc::clone(&calls),
}),
);
let executor: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(PlanRuntimeBuilderExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
});
let result = EngineBuilder::with_registry(EngineConfig::default(), registry)
.with_custom_executor(executor)
.build()
.await;
let error = result
.err()
.expect("plan runtime without a resolved plan must fail closed");
assert!(error
.to_string()
.contains("authoritative ResolvedModelPlan"));
assert_eq!(calls.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn plan_runtime_build_rejects_legacy_resource_override() {
let executor: Arc<dyn ModelExecutor + Send + Sync> = Arc::new(PlanRuntimeBuilderExecutor {
inner: ferrum_testkit::MockModelExecutor::instant(128),
});
let kv_cache: Arc<dyn KvCacheManager + Send + Sync> =
Arc::new(ferrum_testkit::MockKvCacheManager::new(8));
let error = EngineBuilder::new(EngineConfig::default())
.with_custom_executor(executor)
.with_custom_kv_cache(kv_cache)
.build()
.await
.err()
.expect("plan runtime must reject a legacy KV manager override");
assert!(error
.to_string()
.contains("cannot be combined with a legacy engine KV-cache override"));
}
}