diff --git a/services/trading_service/src/services/enhanced_ml.rs b/services/trading_service/src/services/enhanced_ml.rs index 9b77d8e3b..3cfe22dfa 100644 --- a/services/trading_service/src/services/enhanced_ml.rs +++ b/services/trading_service/src/services/enhanced_ml.rs @@ -23,8 +23,9 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; use std::time::{Instant, SystemTime, UNIX_EPOCH}; -use tokio::sync::{broadcast, RwLock}; +use tokio::sync::{broadcast, OnceCell, RwLock}; use tokio_stream::{wrappers::BroadcastStream, wrappers::ReceiverStream, Stream, StreamExt}; +use tonic::transport::Channel; use tonic::{Request, Response, Status}; use tracing::{debug, info, warn, instrument}; @@ -189,6 +190,8 @@ pub struct EnhancedMLServiceImpl { feature_preprocessor: Arc, // System metrics tracking system: Arc>, + // gRPC client for forwarding retrain requests to ml_training_service + ml_training_client: Arc>>, } impl EnhancedMLServiceImpl { @@ -210,9 +213,29 @@ impl EnhancedMLServiceImpl { ml_fallback_manager, feature_preprocessor: Arc::new(FeaturePreprocessor::new()), system: Arc::new(RwLock::new(System::new_all())), + ml_training_client: Arc::new(OnceCell::new()), } } + /// Lazily connect to ml_training_service. Endpoint from ML_TRAINING_SERVICE_URL env var. + async fn get_training_client( + &self, + ) -> Result, Status> { + self.ml_training_client + .get_or_try_init(|| async { + let url = std::env::var("ML_TRAINING_SERVICE_URL") + .unwrap_or_else(|_| "http://ml-training-service:50052".to_string()); + let channel = Channel::from_shared(url) + .map_err(|e| Status::internal(format!("invalid training service URL: {e}")))? + .connect() + .await + .map_err(|e| Status::unavailable(format!("ml_training_service unreachable: {e}")))?; + Ok(crate::proto::ml_training::ml_training_service_client::MlTrainingServiceClient::new(channel)) + }) + .await + .cloned() + } + /// Load model from checkpoint file (production implementation) pub async fn load_model_from_file( &self, @@ -1020,16 +1043,51 @@ impl MlService for EnhancedMLServiceImpl { request: Request, ) -> Result, Status> { let req = request.into_inner(); + info!(model_name = %req.model_name, "retrain_model: forwarding to ml_training_service"); - warn!( - model_name = %req.model_name, - "retrain_model called but training service integration not yet implemented" - ); + // Look up model to validate it exists and get its ModelType + let models = self.models.read().await; + let model_info = models.get(&req.model_name).ok_or_else(|| { + Status::not_found(format!("model '{}' not registered", req.model_name)) + })?; + let model_type_str = model_info.model_type.as_str().to_string(); + drop(models); - Err(Status::unavailable(format!( - "Model retraining for '{}' not yet wired to ml_training_service", - req.model_name - ))) + // Connect to ml_training_service + let mut client = self.get_training_client().await?; + + // Build fine-tune request + use crate::proto::ml_training::{StartTrainingRequest as MlStartReq, TrainingMode}; + let training_req = MlStartReq { + model_type: model_type_str, + data_source: None, + hyperparameters: None, + use_gpu: true, + description: format!("Fine-tune {} via retrain_model RPC", req.model_name), + tags: req.parameters.clone(), + mode: TrainingMode::FineTune.into(), + resume_checkpoint_path: String::new(), + max_epochs: 10, + }; + + let resp = client + .start_training(tonic::Request::new(training_req)) + .await + .map_err(|e| { + warn!(error = %e, "ml_training_service StartTraining failed"); + Status::unavailable(format!("training service error: {e}")) + })? + .into_inner(); + + Ok(Response::new(RetrainModelResponse { + success: true, + message: resp.message, + job_id: Some(resp.job_id), + started_at: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0), + })) } #[instrument(skip_all)] @@ -2769,24 +2827,24 @@ mod enhanced_ml_tests { } // ----------------------------------------------------------------------- - // 12. retrain_model returns Unimplemented (H7 fix) + // 12. retrain_model returns NotFound for unknown models // ----------------------------------------------------------------------- #[test] - fn test_retrain_model_status_code_is_unimplemented() { - // Verify the Status::unimplemented message format matches implementation - let model_name = "test-ppo-v2"; - let status = Status::unimplemented(format!( - "Model retraining for '{}' not yet connected to training service", + fn test_retrain_model_not_found_for_unknown_model() { + // retrain_model now validates the model exists before forwarding + let model_name = "nonexistent-model"; + let status = Status::not_found(format!( + "model '{}' not registered", model_name )); - assert_eq!(status.code(), tonic::Code::Unimplemented); + assert_eq!(status.code(), tonic::Code::NotFound); assert!( status.message().contains(model_name), "Error message should contain the model name" ); assert!( - status.message().contains("not yet connected"), - "Error message should indicate the feature is not connected" + status.message().contains("not registered"), + "Error message should indicate the model is not registered" ); }