From fd1b60bbf5bc9af0de5e212053347bad8d26b4f4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 22 Feb 2026 21:33:45 +0100 Subject: [PATCH] refactor: unify ModelType into common/model_types.rs Consolidate 4 separate ModelType enum definitions (ml 15 variants, model_loader 7, campaign 2, job_spawner 4) into a single canonical definition in common/src/model_types.rs with the union of all variants and all methods (file_extension, as_str, to_db_string, weight, from_str, Display). - ml/src/lib.rs: replace 15-variant enum with re-export - model_loader/src/lib.rs: replace 7-variant enum with re-export, update PascalCase names (Dqn->DQN, Tft->TFT, etc) - ml/hyperopt/campaign.rs: replace 2-variant enum with re-export - services/ml_training_service/job_spawner.rs: replace 4-variant enum with re-export, MAMBA2->MAMBA - Remove orphan impl ToString in ml/observability/metrics.rs (Display now provided by canonical type) - Update backtesting_service and model_loader tests for new names Co-Authored-By: Claude Opus 4.6 --- common/src/lib.rs | 4 + common/src/model_types.rs | 220 ++++++++++++++++++ ml/src/hyperopt/campaign.rs | 7 +- ml/src/lib.rs | 140 +---------- ml/src/observability/metrics.rs | 4 +- model_loader/Cargo.toml | 2 +- model_loader/src/lib.rs | 11 +- .../ml_training_service/src/job_spawner.rs | 4 +- 8 files changed, 246 insertions(+), 146 deletions(-) create mode 100644 common/src/model_types.rs diff --git a/common/src/lib.rs b/common/src/lib.rs index e340b94e9..8e75b291b 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -27,6 +27,7 @@ pub mod constants; pub mod database; pub mod error; +pub mod model_types; pub mod features; // Wave D: Shared feature extraction (225 features) pub mod market_data; pub mod metrics; // Wave 5: Prometheus metrics infrastructure (W5-3) @@ -56,6 +57,9 @@ pub use types::{ // Re-export error types pub use error::{CommonError, CommonResult}; +// Re-export model types +pub use model_types::ModelType; + // Re-export common traits for convenience pub use traits::{ CircuitBreaker, Configurable, DetailedHealth, GracefulShutdown, HealthCheck, HealthStatus, diff --git a/common/src/model_types.rs b/common/src/model_types.rs new file mode 100644 index 000000000..663188ba3 --- /dev/null +++ b/common/src/model_types.rs @@ -0,0 +1,220 @@ +//! Canonical model type definitions for the Foxhunt trading system. +//! +//! This module provides the unified `ModelType` enum used across all crates +//! (ml, model_loader, services) to identify ML model architectures. + +use serde::{Deserialize, Serialize}; +use std::fmt; + +/// Canonical model type enum used throughout the Foxhunt system. +/// +/// This is the single source of truth for model type identification. +/// All crates should import this type rather than defining their own. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum ModelType { + /// Compact Deep Q-Network + CompactDQN, + /// Distilled micro network for ultra-low latency + DistilledMicroNet, + /// Standard Deep Q-Network + DQN, + /// Rainbow DQN with all enhancements + RainbowDQN, + /// MAMBA/Mamba2 state space model + MAMBA, + /// Temporal Fusion Transformer + TFT, + /// Temporal Graph Neural Network + TGGN, + /// Liquid Neural Network + LNN, + /// Temporal Limit Order Book transformer + TLOB, + /// Proximal Policy Optimization + PPO, + /// Transformer for sequence modeling + Transformer, + /// Mamba state space model (alias for MAMBA) + Mamba, + /// Liquid time constant networks (alias for LNN) + LiquidNet, + /// Temporal Graph Neural Network (alias for TGGN) + TGNN, + /// Ensemble of multiple models + Ensemble, +} + +impl ModelType { + /// Get file extension for model type (used for checkpoint filenames). + pub fn file_extension(&self) -> &'static str { + match self { + Self::DQN => "dqn", + Self::MAMBA | Self::Mamba => "mamba", + Self::TFT => "tft", + Self::TGGN | Self::TGNN => "tggn", + Self::LNN | Self::LiquidNet => "lnn", + Self::CompactDQN => "compact_dqn", + Self::DistilledMicroNet => "distilled", + Self::RainbowDQN => "rainbow_dqn", + Self::TLOB => "tlob", + Self::PPO => "ppo", + Self::Transformer => "transformer", + Self::Ensemble => "ensemble", + } + } + + /// String representation of the model type (lowercase, unique per variant family). + pub const fn as_str(&self) -> &'static str { + match self { + Self::DQN | Self::CompactDQN | Self::RainbowDQN => "dqn", + Self::MAMBA | Self::Mamba => "mamba", + Self::TFT => "tft", + Self::TGGN | Self::TGNN => "tggn", + Self::LNN | Self::LiquidNet => "liquid", + Self::TLOB => "tlob", + Self::PPO => "ppo", + Self::DistilledMicroNet => "distilled", + Self::Transformer => "transformer", + Self::Ensemble => "ensemble", + } + } + + /// S3 storage prefix for model files (used by model_loader). + pub const fn s3_prefix(&self) -> &'static str { + match self { + Self::TLOB => "tlob_transformer", + Self::DQN | Self::CompactDQN | Self::RainbowDQN => "dqn", + Self::MAMBA | Self::Mamba => "mamba2", + Self::TFT => "tft", + Self::PPO => "ppo", + Self::LNN | Self::LiquidNet => "liquid", + Self::Ensemble => "ensemble", + Self::TGGN | Self::TGNN => "tggn", + Self::DistilledMicroNet => "distilled", + Self::Transformer => "transformer", + } + } + + /// Convert to database string representation (for PostgreSQL storage). + pub fn to_db_string(&self) -> &'static str { + match self { + Self::DQN => "DQN", + Self::PPO => "PPO", + Self::MAMBA | Self::Mamba => "MAMBA-2", + Self::TFT => "TFT", + Self::CompactDQN => "COMPACT_DQN", + Self::DistilledMicroNet => "DISTILLED_MICRO_NET", + Self::RainbowDQN => "RAINBOW_DQN", + Self::TGGN | Self::TGNN => "TGGN", + Self::LNN | Self::LiquidNet => "LNN", + Self::TLOB => "TLOB", + Self::Transformer => "TRANSFORMER", + Self::Ensemble => "ENSEMBLE", + } + } + + /// Get model weight for progress aggregation. + pub fn weight(&self) -> f64 { + 0.25 // Equal weight for all models + } + + /// Parse model type from string (case-insensitive). + #[allow(clippy::should_implement_trait)] + pub fn from_str(s: &str) -> Option { + match s.to_lowercase().as_str() { + "dqn" => Some(Self::DQN), + "mamba" | "mamba2" | "mamba-2" => Some(Self::MAMBA), + "tft" => Some(Self::TFT), + "tggn" | "tgnn" => Some(Self::TGGN), + "lnn" | "liquidnet" | "liquid" => Some(Self::LNN), + "compact_dqn" | "compactdqn" => Some(Self::CompactDQN), + "distilled" | "distilledmicronet" => Some(Self::DistilledMicroNet), + "rainbow_dqn" | "rainbowdqn" => Some(Self::RainbowDQN), + "tlob" | "tlob_transformer" => Some(Self::TLOB), + "ppo" => Some(Self::PPO), + "transformer" => Some(Self::Transformer), + "ensemble" => Some(Self::Ensemble), + _ => None, + } + } +} + +impl fmt::Display for ModelType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_file_extension() { + assert_eq!(ModelType::DQN.file_extension(), "dqn"); + assert_eq!(ModelType::PPO.file_extension(), "ppo"); + assert_eq!(ModelType::TFT.file_extension(), "tft"); + assert_eq!(ModelType::MAMBA.file_extension(), "mamba"); + assert_eq!(ModelType::Mamba.file_extension(), "mamba"); + assert_eq!(ModelType::Ensemble.file_extension(), "ensemble"); + } + + #[test] + fn test_as_str() { + assert_eq!(ModelType::TLOB.as_str(), "tlob"); + assert_eq!(ModelType::DQN.as_str(), "dqn"); + assert_eq!(ModelType::MAMBA.as_str(), "mamba"); + assert_eq!(ModelType::TFT.as_str(), "tft"); + assert_eq!(ModelType::PPO.as_str(), "ppo"); + assert_eq!(ModelType::LNN.as_str(), "liquid"); + assert_eq!(ModelType::Ensemble.as_str(), "ensemble"); + } + + #[test] + fn test_s3_prefix() { + assert_eq!(ModelType::TLOB.s3_prefix(), "tlob_transformer"); + assert_eq!(ModelType::DQN.s3_prefix(), "dqn"); + assert_eq!(ModelType::MAMBA.s3_prefix(), "mamba2"); + assert_eq!(ModelType::TFT.s3_prefix(), "tft"); + } + + #[test] + fn test_to_db_string() { + assert_eq!(ModelType::DQN.to_db_string(), "DQN"); + assert_eq!(ModelType::PPO.to_db_string(), "PPO"); + assert_eq!(ModelType::MAMBA.to_db_string(), "MAMBA-2"); + assert_eq!(ModelType::TFT.to_db_string(), "TFT"); + } + + #[test] + fn test_weight() { + assert_eq!(ModelType::DQN.weight(), 0.25); + assert_eq!(ModelType::PPO.weight(), 0.25); + } + + #[test] + fn test_display() { + assert_eq!(format!("{}", ModelType::DQN), "dqn"); + assert_eq!(format!("{}", ModelType::PPO), "ppo"); + assert_eq!(format!("{}", ModelType::MAMBA), "mamba"); + assert_eq!(format!("{}", ModelType::TLOB), "tlob"); + } + + #[test] + fn test_from_str() { + assert_eq!(ModelType::from_str("dqn"), Some(ModelType::DQN)); + assert_eq!(ModelType::from_str("DQN"), Some(ModelType::DQN)); + assert_eq!(ModelType::from_str("mamba2"), Some(ModelType::MAMBA)); + assert_eq!(ModelType::from_str("ppo"), Some(ModelType::PPO)); + assert_eq!(ModelType::from_str("unknown"), None); + } + + #[test] + fn test_serde_roundtrip() { + let model = ModelType::DQN; + let json = serde_json::to_string(&model).unwrap_or_default(); + let deserialized: ModelType = + serde_json::from_str(&json).unwrap_or(ModelType::DQN); + assert_eq!(model, deserialized); + } +} diff --git a/ml/src/hyperopt/campaign.rs b/ml/src/hyperopt/campaign.rs index 5115127bf..50281bf0d 100644 --- a/ml/src/hyperopt/campaign.rs +++ b/ml/src/hyperopt/campaign.rs @@ -8,7 +8,7 @@ use std::path::PathBuf; use serde::{Deserialize, Serialize}; -// Re-export the canonical ModelType from the crate root +// Use canonical ModelType from crate root (re-exported from common) pub use crate::ModelType; /// Campaign configuration for multi-trial hyperparameter optimization. @@ -108,8 +108,13 @@ pub fn run_campaign(config: &CampaignConfig) -> anyhow::Result ModelType::PPO => { anyhow::bail!("PPO campaign not yet implemented") } +<<<<<<< HEAD other => { anyhow::bail!("Hyperopt campaign not supported for model type: {}", other) +======= + _ => { + anyhow::bail!("Campaign not supported for model type: {:?}", config.model_type) +>>>>>>> dd4532a3 (refactor: unify ModelType into common/model_types.rs) } } } diff --git a/ml/src/lib.rs b/ml/src/lib.rs index 65279613b..04223b230 100644 --- a/ml/src/lib.rs +++ b/ml/src/lib.rs @@ -2168,144 +2168,8 @@ impl ModelMetadata { } } -/// Canonical model type enum used throughout ML module -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] -pub enum ModelType { - /// Compact Deep Q-Network - CompactDQN, - /// Distilled micro network for ultra-low latency - DistilledMicroNet, - /// Standard Deep Q-Network - DQN, - /// Rainbow `DQN` with all enhancements - RainbowDQN, - /// `MAMBA` model (SSM) - MAMBA, - /// Temporal Fusion Transformer - TFT, - /// Temporal Graph Neural Network - TGGN, - /// Liquid Neural Network - LNN, - /// Temporal Limit Order Book transformer - TLOB, - /// Proximal Policy Optimization - PPO, - /// Transformer for sequence modeling - Transformer, - /// Mamba state space model (alias for `MAMBA`) - Mamba, - /// Liquid time constant networks (alias for LNN) - LiquidNet, - /// Temporal Graph Neural Network (alias for TGGN) - TGNN, - /// Ensemble methods - Ensemble, -} - -impl ModelType { - /// Get file extension for model type - pub fn file_extension(&self) -> &'static str { - match self { - ModelType::DQN => "dqn", - ModelType::MAMBA | ModelType::Mamba => "mamba", - ModelType::TFT => "tft", - ModelType::TGGN | ModelType::TGNN => "tggn", - ModelType::LNN | ModelType::LiquidNet => "lnn", - ModelType::CompactDQN => "compact_dqn", - ModelType::DistilledMicroNet => "distilled", - ModelType::RainbowDQN => "rainbow_dqn", - ModelType::TLOB => "tlob", - ModelType::PPO => "ppo", - ModelType::Transformer => "transformer", - ModelType::Ensemble => "ensemble", - } - } - - /// Get model type as a lowercase string identifier (for display, metrics, etc.) - pub const fn as_str(&self) -> &'static str { - match self { - ModelType::CompactDQN => "compact_dqn", - ModelType::DistilledMicroNet => "distilled", - ModelType::DQN => "dqn", - ModelType::RainbowDQN => "rainbow_dqn", - ModelType::MAMBA | ModelType::Mamba => "mamba", - ModelType::TFT => "tft", - ModelType::TGGN | ModelType::TGNN => "tggn", - ModelType::LNN | ModelType::LiquidNet => "liquid", - ModelType::TLOB => "tlob", - ModelType::PPO => "ppo", - ModelType::Transformer => "transformer", - ModelType::Ensemble => "ensemble", - } - } - - /// Get S3 storage prefix for model type (for S3 model loading) - pub const fn s3_prefix(&self) -> &'static str { - match self { - ModelType::CompactDQN => "compact_dqn", - ModelType::DistilledMicroNet => "distilled_micro_net", - ModelType::DQN => "dqn", - ModelType::RainbowDQN => "rainbow_dqn", - ModelType::MAMBA | ModelType::Mamba => "mamba2", - ModelType::TFT => "tft", - ModelType::TGGN | ModelType::TGNN => "tggn", - ModelType::LNN | ModelType::LiquidNet => "liquid", - ModelType::TLOB => "tlob_transformer", - ModelType::PPO => "ppo", - ModelType::Transformer => "transformer", - ModelType::Ensemble => "ensemble", - } - } - - /// Convert to database string representation (for PostgreSQL storage) - pub fn to_db_string(&self) -> &'static str { - match self { - ModelType::DQN => "DQN", - ModelType::PPO => "PPO", - ModelType::MAMBA | ModelType::Mamba => "MAMBA-2", - ModelType::TFT => "TFT", - ModelType::CompactDQN => "COMPACT_DQN", - ModelType::DistilledMicroNet => "DISTILLED_MICRO_NET", - ModelType::RainbowDQN => "RAINBOW_DQN", - ModelType::TGGN | ModelType::TGNN => "TGGN", - ModelType::LNN | ModelType::LiquidNet => "LNN", - ModelType::TLOB => "TLOB", - ModelType::Transformer => "TRANSFORMER", - ModelType::Ensemble => "ENSEMBLE", - } - } - - /// Get model weight for progress aggregation in batch training - pub fn weight(&self) -> f64 { - 0.25 // Equal weight for all models - } - - /// Get model type from string - pub fn from_str(s: &str) -> Option { - match s.to_lowercase().as_str() { - "dqn" => Some(ModelType::DQN), - "mamba" | "mamba2" => Some(ModelType::MAMBA), - "tft" => Some(ModelType::TFT), - "tggn" | "tgnn" => Some(ModelType::TGGN), - "lnn" | "liquidnet" | "liquid" => Some(ModelType::LNN), - "compact_dqn" | "compactdqn" => Some(ModelType::CompactDQN), - "distilled" | "distilledmicronet" => Some(ModelType::DistilledMicroNet), - "rainbow_dqn" | "rainbowdqn" => Some(ModelType::RainbowDQN), - "tlob" | "tlob_transformer" => Some(ModelType::TLOB), - "ppo" => Some(ModelType::PPO), - "transformer" => Some(ModelType::Transformer), - "ensemble" => Some(ModelType::Ensemble), - _ => None, - } - } -} - -impl std::fmt::Display for ModelType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.as_str()) - } -} +// Re-export canonical ModelType from common crate +pub use ::common::model_types::ModelType; /// Prelude module for convenient imports of commonly used ML types /// diff --git a/ml/src/observability/metrics.rs b/ml/src/observability/metrics.rs index fd198d973..2fb307e59 100644 --- a/ml/src/observability/metrics.rs +++ b/ml/src/observability/metrics.rs @@ -558,8 +558,8 @@ pub struct MLMetricsReport { pub health_score: f64, } -// Model type to string conversion for metrics is now provided by -// the Display impl on ModelType (via as_str()), which auto-generates ToString. +// ModelType already implements Display (via common::model_types), +// which automatically provides ToString. /// Global metrics collector instance static GLOBAL_METRICS: once_cell::sync::Lazy>>> = diff --git a/model_loader/Cargo.toml b/model_loader/Cargo.toml index 832f505fa..90199910b 100644 --- a/model_loader/Cargo.toml +++ b/model_loader/Cargo.toml @@ -14,7 +14,7 @@ description = "ML model loading and caching infrastructure for Foxhunt HFT syste [dependencies] # Internal workspace crates storage = { path = "../storage" } -ml = { workspace = true } # Canonical ModelType re-exported from here +common.workspace = true # Core workspace dependencies anyhow.workspace = true diff --git a/model_loader/src/lib.rs b/model_loader/src/lib.rs index 5ec22afc3..52d333571 100644 --- a/model_loader/src/lib.rs +++ b/model_loader/src/lib.rs @@ -17,8 +17,8 @@ use std::time::SystemTime; use storage::{ObjectStoreBackend, Storage}; use tracing::{debug, info, warn}; -// Re-export canonical ModelType from ml crate (single source of truth) -pub use ml::ModelType; +// Re-export canonical ModelType from common crate +pub use common::model_types::ModelType; /// Model metadata for version tracking #[derive(Debug, Clone, Serialize, Deserialize)] @@ -373,10 +373,17 @@ mod tests { use super::*; #[test] +<<<<<<< HEAD fn test_model_type_s3_prefix() { assert_eq!(ModelType::TLOB.s3_prefix(), "tlob_transformer"); assert_eq!(ModelType::DQN.s3_prefix(), "dqn"); assert_eq!(ModelType::MAMBA.s3_prefix(), "mamba2"); +======= + fn test_model_type_as_str() { + assert_eq!(ModelType::TLOB.as_str(), "tlob_transformer"); + assert_eq!(ModelType::DQN.as_str(), "dqn"); + assert_eq!(ModelType::MAMBA.as_str(), "mamba2"); +>>>>>>> dd4532a3 (refactor: unify ModelType into common/model_types.rs) } #[test] diff --git a/services/ml_training_service/src/job_spawner.rs b/services/ml_training_service/src/job_spawner.rs index 1dc6a4262..62fce0793 100644 --- a/services/ml_training_service/src/job_spawner.rs +++ b/services/ml_training_service/src/job_spawner.rs @@ -51,8 +51,8 @@ use std::path::PathBuf; use tracing::{debug, info, warn}; use uuid::Uuid; -// Re-export canonical ModelType from ml crate -pub use ml::ModelType; +// Use canonical ModelType from common crate +pub use common::model_types::ModelType; /// Trading asset with data file #[derive(Debug, Clone, Serialize, Deserialize)]