rustytorch/crates/production/rtx-inference/src/lib.rs

//! `RustyTorch`++ Inference Runtime
//!
//! A production-grade inference runtime for transformer models with:
//! - Continuous batching with SLA lanes
//! - Paged KV cache with GPU/CPU/NVMe tiering
//! - Speculative and assisted decoding
//! - Quantization support (INT8/INT4/FP8)
//! - Integration with auto-kernel synthesis
//!
//! This crate provides vLLM-class performance with strict safety guarantees
//! and comprehensive error handling throughout the inference pipeline.

#![deny(clippy::unwrap_used)]
#![cfg_attr(test, allow(clippy::unwrap_used))]
#![allow(clippy::module_name_repetitions)]
#![allow(clippy::too_many_lines)]

pub mod beam_search;
pub use beam_search::{Beam, BeamSearchConfig, BeamSearchDecoder, DiverseBeamSearchDecoder};

pub mod logit_processors;
pub use logit_processors::{
    EtaSamplingProcessor, FrequencyPenaltyProcessor, LogitProcessor, LogitProcessorList,
    MinPProcessor, PresencePenaltyProcessor, RepetitionPenaltyProcessor, TemperatureProcessor,
    TopKProcessor, TopPProcessor, argmax, log_softmax, sample_token, softmax,
};

pub mod medusa;
pub use medusa::{
    MedusaConfig as MedusaHeadsConfig, MedusaHead, MedusaHeads, MedusaLossResult, MedusaTree,
    MedusaVerifyResult,
};
pub mod eagle;
pub use eagle::{
    EagleDraftHead, EagleDraftStep, EagleHeadOutput, EagleHeads, EagleHeadsConfig, EagleLossResult,
};
pub mod batch_processor;
pub mod cache;
pub mod gqa;
pub use gqa::{GqaConfig, GqaError, expand_kv_heads, gqa_attention_cpu, kv_head_for_q};
pub mod chunked_prefill;
pub use chunked_prefill::{
    ChunkedPrefillConfig, ChunkedPrefillScheduler, ChunkedStep, PrefillChunkState,
};
pub mod inference_graph;
pub use inference_graph::{InferenceGraphCapture, StepMode};
pub mod engine;
pub mod error;
pub mod model_loader;
pub mod monitoring;
pub mod quantization;
pub mod request;
pub mod scheduler;
pub mod speculative;
pub mod tokenizer;
pub use tokenizer::ServingTokenizer;

// Re-export key types for convenience
pub use cache::{
    AttentionScoreEviction, AttentionSinkEviction, CacheKey, CachePage, CacheStats,
    CpuOffloadConfig, EvictionPolicy, KvCacheConfig, KvCpuOffloadManager, MemoryTier, OffloadStats,
    PageId, PagedKvCache, PagedKvCacheManager, PrefixIndex,
};
pub use engine::{
    HealthStatus, InferenceEngine, InferenceEngineConfig, MemoryStats, ModelConfig, ModelHealth,
    ModelInfo, ModelMetrics, OptimizationStats, PerformanceMetrics, StreamingToken,
};
pub use error::{ErrorSeverity, InferenceError, InferenceResult};
pub use model_loader::{
    LayerConfig as ModelLayerConfig, LoadedModel, LoadingProgress, LoadingStage, ModelFormat,
    ModelLoader, ModelLoaderConfig, OptimizationStats as LoaderOptimizationStats, ShapeConstraint,
    TokenizerConfig, ValidationResults,
};
pub use quantization::{
    BatchQuantizer, CalibrationMethod, DataType, DynamicQuantizer, GroupSize, LayerConfig,
    MixedPrecisionQuantizer, OptimizationLevel, OutlierHandling, QuantizationConfig,
    QuantizationGranularity, QuantizationParams, QuantizationScheme, QuantizationValidator,
    Quantizer, ScalingMethod, TestTensor, ValidationMetrics,
};
pub use request::{
    FinishReason, InferenceRequest, OverflowStrategy, QueueStats, RequestId, RequestManager,
    RequestManagerConfig, RequestMetrics, RequestPriority, RequestResult, RequestStatus,
};
pub use scheduler::{
    BatchScheduler, BatchSchedulerConfig, InferenceBatch, PreemptionDecision,
    QueueStats as SchedulerQueueStats, SlaLane, SlaViolation,
};
pub use speculative::{
    AcceptanceDecision,
    AdaptiveConfig,
    AdvancedMetrics,
    AdvancedSpeculativeDecoder,
    CandidateTree,
    DecodingResult,
    DraftModel,
    DraftModelType,
    EagleConfig,
    EagleDraftModel,
    FusionMethod,
    LookaheadConfig as SpeculativeLookaheadConfig,
    MedusaConfig,
    MedusaDraftModel,
    NgramPool,
    PerformanceMetrics as SpeculativePerformanceMetrics,
    SelfSpeculativeConfig,
    SelfSpeculativeModel,
    SpeculativeConfig,
    SpeculativeDecoder,
    SpeculativeError,
    SpeculativeStreamConfig,
    SpeculativeStreamer,
    StreamStats,
    StreamedToken,
    TargetModel,
    Token,
    // Advanced speculative decoding
    TreeNode,
    // Speculative streaming
    collect_stream,
};

pub mod lookahead;
pub use lookahead::{LookaheadConfig, LookaheadDecoder, LookaheadStats, NGram, NGramCache};

// Re-export ONNX Runtime types when feature is enabled
#[cfg(feature = "onnx-runtime")]
pub use model_loader::OnnxRuntimeModel;
#[cfg(feature = "onnx-runtime")]
pub use rtx_onnx::{
    ExecutionProviderType, OnnxSessionConfig, OptimizationLevel as OnnxOptimizationLevel,
};

// Re-export Burn types when feature is enabled
#[cfg(feature = "burn")]
pub use model_loader::BurnRuntimeModel;
#[cfg(feature = "burn")]
pub use rtx_burn::{BurnBackend, BurnConfig, BurnModel, SessionStats as BurnSessionStats};

// Re-export Candle types when feature is enabled
pub use batch_processor::{
    BatchProcessor, BatchProcessorConfig, BatchStats, BatchedRequest, ProcessedBatch, SlaType,
};
#[cfg(feature = "candle")]
pub use model_loader::CandleRuntimeModel;
pub use monitoring::{
    Alert, HealthLevel, HealthStatus as MonitoringHealthStatus, MetricsCollector,
    ModelSpecificMetrics, MonitoringConfig, PerformanceMetrics as MonitoringPerformanceMetrics,
    RequestTrace, ResourceMetrics,
};
#[cfg(feature = "candle")]
pub use rtx_candle::{
    CandleBackend, CandleConfig, CandleModel, SessionStats as CandleSessionStats,
};

/// Version information
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

/// Build metadata
pub const BUILD_INFO: BuildInfo = BuildInfo {
    version: VERSION,
    git_commit: match option_env!("RUSTYTORCH_GIT_COMMIT") {
        Some(commit) => commit,
        None => "development",
    },
    build_time: env!("RUSTYTORCH_BUILD_TIMESTAMP"),
    rust_version: env!("RUSTYTORCH_RUST_VERSION"),
    target_arch: env!("RUSTYTORCH_TARGET_ARCH"),
};

/// Build information structure
#[derive(Debug, Clone)]
pub struct BuildInfo {
    /// Crate version
    pub version: &'static str,
    /// Git commit hash
    pub git_commit: &'static str,
    /// Build timestamp
    pub build_time: &'static str,
    /// Rust compiler version
    pub rust_version: &'static str,
    /// Target architecture
    pub target_arch: &'static str,
}

impl std::fmt::Display for BuildInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "rtx-inference {} (commit: {}, built: {}, rust: {}, arch: {})",
            self.version, self.git_commit, self.build_time, self.rust_version, self.target_arch
        )
    }
}

/// Initialize the inference runtime with logging
///
/// This should be called once at application startup to configure
/// logging and any global state required by the inference system.
pub fn init() -> InferenceResult<()> {
    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

    tracing_subscriber::registry()
        .with(tracing_subscriber::fmt::layer())
        .try_init()
        .map_err(|e| InferenceError::internal_error("logging init", e.to_string()))?;

    tracing::info!("Initialized rtx-inference {}", BUILD_INFO);
    Ok(())
}

/// Shutdown the inference runtime cleanly
///
/// This performs any necessary cleanup when the application is shutting down.
pub fn shutdown() {
    use std::io::Write;

    tracing::info!("Shutting down rtx-inference runtime");

    // Flush standard output/error streams
    let _ = std::io::stderr().flush();
    let _ = std::io::stdout().flush();

    // Force flush tracing subscribers
    tracing::debug!("Runtime cleanup completed");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_version_info() {
        assert!(!VERSION.is_empty());
        assert!(!BUILD_INFO.version.is_empty());
        assert!(!BUILD_INFO.target_arch.is_empty());
    }

    #[test]
    fn test_build_info_display() {
        let info_str = BUILD_INFO.to_string();
        assert!(info_str.contains("rtx-inference"));
        assert!(info_str.contains(VERSION));
    }

    #[tokio::test]
    async fn test_init_shutdown() {
        // Note: This test might fail if init() is called multiple times
        // In a real application, init() should only be called once
        let result = init();
        // Don't assert success as logging might already be initialized

        shutdown();
        // Shutdown should always succeed
    }
}