e4870b17b9
Reduce log noise for non-critical operational paths: connection retries, expected fallbacks, graceful degradation, and optional feature absence. Keeps warn/error for genuine failures requiring attention. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1292 lines
47 KiB
Rust
1292 lines
47 KiB
Rust
//! Training Orchestrator
|
|
//!
|
|
//! This module coordinates training jobs using the existing ML training infrastructure
|
|
//! from the ml crate. It manages job queues, resource allocation, and progress tracking.
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use anyhow::Result;
|
|
use chrono::{DateTime, Utc};
|
|
use tokio::sync::{broadcast, mpsc, Mutex, RwLock};
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::{debug, error, info, warn};
|
|
use uuid::Uuid;
|
|
|
|
// Import from existing ML infrastructure
|
|
use ml::training_pipeline::{
|
|
FinancialFeatures, ProductionMLTrainingSystem, ProductionTrainingConfig, TrainingResult,
|
|
};
|
|
|
|
use crate::data_config::TrainingDataSourceConfig;
|
|
use crate::database::DatabaseManager;
|
|
use crate::storage::ModelStorageManager;
|
|
use config::MLConfig;
|
|
|
|
/// Training job status
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum JobStatus {
|
|
Pending,
|
|
Running,
|
|
Completed,
|
|
Failed,
|
|
Stopped,
|
|
Paused,
|
|
}
|
|
|
|
/// Training job metadata
|
|
#[derive(Debug, Clone)]
|
|
pub struct TrainingJob {
|
|
pub id: Uuid,
|
|
pub model_type: String,
|
|
pub status: JobStatus,
|
|
pub config: ProductionTrainingConfig,
|
|
pub created_at: DateTime<Utc>,
|
|
pub started_at: Option<DateTime<Utc>>,
|
|
pub completed_at: Option<DateTime<Utc>>,
|
|
pub description: String,
|
|
pub tags: HashMap<String, String>,
|
|
pub progress_percentage: f32,
|
|
pub current_epoch: u32,
|
|
pub total_epochs: u32,
|
|
pub metrics: HashMap<String, f64>,
|
|
pub error_message: Option<String>,
|
|
pub model_artifact_path: Option<String>,
|
|
}
|
|
|
|
impl TrainingJob {
|
|
pub fn new(
|
|
model_type: String,
|
|
config: ProductionTrainingConfig,
|
|
description: String,
|
|
tags: HashMap<String, String>,
|
|
) -> Self {
|
|
Self {
|
|
id: Uuid::new_v4(),
|
|
model_type,
|
|
status: JobStatus::Pending,
|
|
config,
|
|
created_at: Utc::now(),
|
|
started_at: None,
|
|
completed_at: None,
|
|
description,
|
|
tags,
|
|
progress_percentage: 0.0,
|
|
current_epoch: 0,
|
|
total_epochs: 0,
|
|
metrics: HashMap::new(),
|
|
error_message: None,
|
|
model_artifact_path: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Resource allocation for training jobs
|
|
#[derive(Debug, Clone)]
|
|
pub struct ResourceAllocation {
|
|
pub gpu_id: Option<u32>,
|
|
pub cpu_cores: u32,
|
|
pub memory_gb: f64,
|
|
pub worker_id: u32,
|
|
}
|
|
|
|
/// Training orchestrator that manages job lifecycle and resources
|
|
pub struct TrainingOrchestrator {
|
|
config: MLConfig,
|
|
|
|
// Job management
|
|
jobs: Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
|
|
job_queue: Arc<Mutex<mpsc::Sender<Uuid>>>,
|
|
job_receiver: Arc<Mutex<mpsc::Receiver<Uuid>>>,
|
|
|
|
// Resource management
|
|
available_resources: Arc<Mutex<Vec<ResourceAllocation>>>,
|
|
resource_assignments: Arc<RwLock<HashMap<Uuid, ResourceAllocation>>>,
|
|
|
|
// Progress tracking
|
|
status_broadcasters: Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
|
|
// External dependencies
|
|
database: Arc<crate::database::DatabaseManager>,
|
|
storage: Arc<ModelStorageManager>,
|
|
|
|
// Worker pool
|
|
worker_handles: Vec<tokio::task::JoinHandle<()>>,
|
|
|
|
// Graceful shutdown
|
|
cancellation_token: CancellationToken,
|
|
}
|
|
|
|
/// Status update message for broadcasting
|
|
#[derive(Debug, Clone)]
|
|
pub struct TrainingStatusUpdate {
|
|
pub job_id: Uuid,
|
|
pub status: JobStatus,
|
|
pub progress_percentage: f32,
|
|
pub current_epoch: u32,
|
|
pub total_epochs: u32,
|
|
pub metrics: HashMap<String, f64>,
|
|
pub message: String,
|
|
pub timestamp: DateTime<Utc>,
|
|
pub financial_metrics: Option<ml::training_pipeline::FinancialPerformanceMetrics>,
|
|
pub resource_usage: ResourceUsage,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ResourceUsage {
|
|
pub cpu_usage_percent: f32,
|
|
pub memory_usage_gb: f32,
|
|
pub gpu_usage_percent: Option<f32>,
|
|
pub gpu_memory_usage_gb: Option<f32>,
|
|
pub active_workers: u32,
|
|
}
|
|
|
|
impl TrainingOrchestrator {
|
|
/// Create a new training orchestrator
|
|
pub async fn new(
|
|
config: MLConfig,
|
|
database: Arc<crate::database::DatabaseManager>,
|
|
storage: Arc<ModelStorageManager>,
|
|
) -> Result<Self> {
|
|
let (job_sender, job_receiver) = mpsc::channel(1000); // Default queue capacity
|
|
|
|
// Initialize available resources based on system capabilities
|
|
let available_resources = Self::detect_system_resources(&config).await?;
|
|
|
|
let orchestrator = Self {
|
|
config: config.clone(),
|
|
jobs: Arc::new(RwLock::new(HashMap::new())),
|
|
job_queue: Arc::new(Mutex::new(job_sender)),
|
|
job_receiver: Arc::new(Mutex::new(job_receiver)),
|
|
available_resources: Arc::new(Mutex::new(available_resources)),
|
|
resource_assignments: Arc::new(RwLock::new(HashMap::new())),
|
|
status_broadcasters: Arc::new(RwLock::new(HashMap::new())),
|
|
database,
|
|
storage,
|
|
worker_handles: Vec::new(),
|
|
cancellation_token: CancellationToken::new(),
|
|
};
|
|
|
|
info!("Training orchestrator initialized with default configuration");
|
|
|
|
Ok(orchestrator)
|
|
}
|
|
|
|
/// Start the orchestrator worker pool
|
|
pub async fn start(&mut self) -> Result<()> {
|
|
info!("Starting training orchestrator with default worker configuration");
|
|
|
|
// Start worker tasks with default thread count
|
|
let worker_count = 4; // Default worker threads
|
|
for worker_id in 0..worker_count {
|
|
let worker_handle = self.spawn_worker(worker_id).await?;
|
|
self.worker_handles.push(worker_handle);
|
|
}
|
|
|
|
// Start resource monitor
|
|
let resource_monitor_handle = self.spawn_resource_monitor().await?;
|
|
self.worker_handles.push(resource_monitor_handle);
|
|
|
|
// Start snapshot broadcaster for fallback status updates
|
|
let snapshot_broadcaster_handle = self.spawn_snapshot_broadcaster().await?;
|
|
self.worker_handles.push(snapshot_broadcaster_handle);
|
|
|
|
info!("Training orchestrator started successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Gracefully shutdown the orchestrator
|
|
pub async fn shutdown(&mut self) -> Result<()> {
|
|
info!("Initiating graceful shutdown of training orchestrator");
|
|
|
|
// Signal all workers to stop
|
|
self.cancellation_token.cancel();
|
|
|
|
// Wait for all workers to finish with timeout
|
|
let shutdown_timeout = Duration::from_secs(30);
|
|
let mut remaining_handles = Vec::new();
|
|
|
|
for handle in self.worker_handles.drain(..) {
|
|
remaining_handles.push(handle);
|
|
}
|
|
|
|
// Wait for workers with timeout
|
|
if let Err(_) = tokio::time::timeout(shutdown_timeout, async {
|
|
for handle in remaining_handles {
|
|
let _ = handle.await;
|
|
}
|
|
})
|
|
.await
|
|
{
|
|
warn!("Some workers did not shutdown gracefully within timeout");
|
|
}
|
|
|
|
// Clean up any remaining broadcasters for proper resource cleanup
|
|
{
|
|
let mut broadcasters = self.status_broadcasters.write().await;
|
|
broadcasters.clear();
|
|
}
|
|
|
|
info!("Training orchestrator shutdown completed");
|
|
Ok(())
|
|
}
|
|
|
|
/// Clean up disconnected broadcasters to prevent memory leaks
|
|
pub async fn cleanup_disconnected_broadcasters(&self) {
|
|
let mut broadcasters = self.status_broadcasters.write().await;
|
|
let mut to_remove = Vec::new();
|
|
|
|
for (job_id, broadcaster) in broadcasters.iter() {
|
|
if broadcaster.receiver_count() == 0 {
|
|
to_remove.push(*job_id);
|
|
}
|
|
}
|
|
|
|
for job_id in to_remove {
|
|
broadcasters.remove(&job_id);
|
|
debug!("Cleaned up broadcaster for job {}", job_id);
|
|
}
|
|
}
|
|
|
|
/// Submit a new training job
|
|
pub async fn submit_job(
|
|
&self,
|
|
model_type: String,
|
|
config: ProductionTrainingConfig,
|
|
description: String,
|
|
tags: HashMap<String, String>,
|
|
) -> Result<Uuid> {
|
|
let job = TrainingJob::new(model_type, config, description, tags);
|
|
let job_id = job.id;
|
|
|
|
// Store job in database
|
|
let job_record = crate::database::TrainingJobRecord::from_training_job(&job);
|
|
if let Err(e) = self.database.insert_training_job(&job_record).await {
|
|
error!("Failed to store job {} in database: {}", job_id, e);
|
|
// Continue - in-memory storage still works
|
|
} else {
|
|
info!("Job {} successfully stored in database", job_id);
|
|
}
|
|
|
|
// Add to in-memory jobs
|
|
{
|
|
let mut jobs = self.jobs.write().await;
|
|
jobs.insert(job_id, job);
|
|
}
|
|
|
|
// Create status broadcaster for this job
|
|
let (tx, _) = broadcast::channel(100); // Default broadcast capacity
|
|
{
|
|
let mut broadcasters = self.status_broadcasters.write().await;
|
|
broadcasters.insert(job_id, tx);
|
|
}
|
|
|
|
// Queue the job for execution
|
|
{
|
|
let queue = self.job_queue.lock().await;
|
|
match queue.try_send(job_id) {
|
|
Ok(()) => {},
|
|
Err(mpsc::error::TrySendError::Full(_)) => {
|
|
return Err(anyhow::anyhow!(
|
|
"Job queue is full. System is under high load. Please try again later."
|
|
));
|
|
},
|
|
Err(mpsc::error::TrySendError::Closed(_)) => {
|
|
return Err(anyhow::anyhow!(
|
|
"Job queue is closed. System is shutting down."
|
|
));
|
|
},
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"Submitted training job {} for model type {}",
|
|
job_id,
|
|
self.get_job_model_type(job_id)
|
|
.await
|
|
.unwrap_or("unknown".to_string())
|
|
);
|
|
|
|
Ok(job_id)
|
|
}
|
|
|
|
/// Stop a running training job
|
|
pub async fn stop_job(&self, job_id: Uuid, reason: String) -> Result<bool> {
|
|
let mut job_updated = false;
|
|
|
|
// Update job status
|
|
{
|
|
let mut jobs = self.jobs.write().await;
|
|
if let Some(job) = jobs.get_mut(&job_id) {
|
|
if matches!(job.status, JobStatus::Running | JobStatus::Pending) {
|
|
job.status = JobStatus::Stopped;
|
|
job.completed_at = Some(Utc::now());
|
|
job.error_message = Some(format!("Stopped: {}", reason));
|
|
job_updated = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
if job_updated {
|
|
// Update database
|
|
let jobs_read = self.jobs.read().await;
|
|
if let Some(job) = jobs_read.get(&job_id) {
|
|
let job_record = crate::database::TrainingJobRecord::from_training_job(job);
|
|
drop(jobs_read); // Release lock before async call
|
|
|
|
if let Err(e) = self.database.update_training_job(&job_record).await {
|
|
error!("Failed to update job {} in database: {}", job_id, e);
|
|
} else {
|
|
info!("Job {} successfully updated in database", job_id);
|
|
}
|
|
}
|
|
|
|
// Broadcast status update
|
|
self.broadcast_status_update(job_id, "Job stopped by user request".to_string())
|
|
.await;
|
|
|
|
info!("Stopped training job {}: {}", job_id, reason);
|
|
}
|
|
|
|
Ok(job_updated)
|
|
}
|
|
|
|
/// Get job status
|
|
pub async fn get_job(&self, job_id: Uuid) -> Result<TrainingJob> {
|
|
let jobs = self.jobs.read().await;
|
|
jobs.get(&job_id)
|
|
.cloned()
|
|
.ok_or_else(|| anyhow::anyhow!("Job {} not found", job_id))
|
|
}
|
|
|
|
/// List all jobs with optional filtering
|
|
pub async fn list_jobs(
|
|
&self,
|
|
status_filter: Option<JobStatus>,
|
|
model_type_filter: Option<String>,
|
|
limit: Option<usize>,
|
|
offset: Option<usize>,
|
|
) -> Result<Vec<TrainingJob>> {
|
|
let jobs = self.jobs.read().await;
|
|
let mut filtered_jobs: Vec<_> = jobs
|
|
.values()
|
|
.filter(|job| {
|
|
if let Some(ref status) = status_filter {
|
|
if job.status != *status {
|
|
return false;
|
|
}
|
|
}
|
|
if let Some(ref model_type) = model_type_filter {
|
|
if job.model_type != *model_type {
|
|
return false;
|
|
}
|
|
}
|
|
true
|
|
})
|
|
.cloned()
|
|
.collect();
|
|
|
|
// Sort by creation time (newest first)
|
|
filtered_jobs.sort_by(|a, b| b.created_at.cmp(&a.created_at));
|
|
|
|
// Apply pagination
|
|
if let Some(offset) = offset {
|
|
if offset < filtered_jobs.len() {
|
|
filtered_jobs = filtered_jobs.into_iter().skip(offset).collect();
|
|
} else {
|
|
filtered_jobs.clear();
|
|
}
|
|
}
|
|
|
|
if let Some(limit) = limit {
|
|
filtered_jobs.truncate(limit);
|
|
}
|
|
|
|
Ok(filtered_jobs)
|
|
}
|
|
|
|
/// Subscribe to status updates for a job
|
|
pub async fn subscribe_to_job_status(
|
|
&self,
|
|
job_id: Uuid,
|
|
) -> Result<broadcast::Receiver<TrainingStatusUpdate>> {
|
|
let broadcasters = self.status_broadcasters.read().await;
|
|
if let Some(broadcaster) = broadcasters.get(&job_id) {
|
|
Ok(broadcaster.subscribe())
|
|
} else {
|
|
Err(anyhow::anyhow!("Job {} not found", job_id))
|
|
}
|
|
}
|
|
|
|
/// Detect available system resources
|
|
async fn detect_system_resources(_config: &MLConfig) -> Result<Vec<ResourceAllocation>> {
|
|
let mut resources = Vec::new();
|
|
|
|
// Default configuration for resource allocation
|
|
let worker_threads = 4u32;
|
|
let default_device = "cpu"; // Default to CPU for safety
|
|
let max_memory_gb = 8.0; // Default memory allocation
|
|
|
|
// For now, create one resource allocation per worker thread
|
|
// In a real implementation, this would detect actual GPU devices
|
|
for worker_id in 0..worker_threads {
|
|
let allocation = ResourceAllocation {
|
|
gpu_id: if default_device == "cuda" {
|
|
Some(0)
|
|
} else {
|
|
None
|
|
},
|
|
cpu_cores: num_cpus::get() as u32 / worker_threads,
|
|
memory_gb: max_memory_gb / worker_threads as f64,
|
|
worker_id,
|
|
};
|
|
resources.push(allocation);
|
|
}
|
|
|
|
info!("Detected {} resource allocations", resources.len());
|
|
Ok(resources)
|
|
}
|
|
|
|
/// Spawn a worker task
|
|
async fn spawn_worker(&self, worker_id: usize) -> Result<tokio::task::JoinHandle<()>> {
|
|
let jobs = Arc::clone(&self.jobs);
|
|
let job_receiver = Arc::clone(&self.job_receiver);
|
|
let available_resources = Arc::clone(&self.available_resources);
|
|
let resource_assignments = Arc::clone(&self.resource_assignments);
|
|
let status_broadcasters = Arc::clone(&self.status_broadcasters);
|
|
let database = Arc::clone(&self.database);
|
|
let storage = Arc::clone(&self.storage);
|
|
let config = self.config.clone();
|
|
let cancellation_token = self.cancellation_token.child_token();
|
|
|
|
let handle = tokio::spawn(async move {
|
|
info!("Training worker {} started", worker_id);
|
|
|
|
loop {
|
|
// Wait for a job to process
|
|
let job_id = {
|
|
let mut receiver = job_receiver.lock().await;
|
|
tokio::select! {
|
|
job_result = receiver.recv() => {
|
|
match job_result {
|
|
Some(id) => id,
|
|
None => {
|
|
warn!("Worker {} job receiver closed", worker_id);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
_ = cancellation_token.cancelled() => {
|
|
info!("Worker {} received cancellation signal", worker_id);
|
|
break;
|
|
}
|
|
}
|
|
};
|
|
|
|
// Process the job
|
|
if let Err(e) = Self::process_job(
|
|
job_id,
|
|
worker_id,
|
|
&jobs,
|
|
&available_resources,
|
|
&resource_assignments,
|
|
&status_broadcasters,
|
|
&database,
|
|
&storage,
|
|
&config,
|
|
)
|
|
.await
|
|
{
|
|
error!(
|
|
"Worker {} failed to process job {}: {}",
|
|
worker_id, job_id, e
|
|
);
|
|
}
|
|
}
|
|
|
|
info!("Training worker {} stopped", worker_id);
|
|
});
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
/// Process a single training job
|
|
async fn process_job(
|
|
job_id: Uuid,
|
|
worker_id: usize,
|
|
jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
|
|
available_resources: &Arc<Mutex<Vec<ResourceAllocation>>>,
|
|
resource_assignments: &Arc<RwLock<HashMap<Uuid, ResourceAllocation>>>,
|
|
status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
database: &Arc<DatabaseManager>,
|
|
storage: &Arc<ModelStorageManager>,
|
|
_config: &MLConfig,
|
|
) -> Result<()> {
|
|
info!("Worker {} processing job {}", worker_id, job_id);
|
|
|
|
// Get job details
|
|
let (job_config, model_type) = {
|
|
let jobs_read = jobs.read().await;
|
|
if let Some(job) = jobs_read.get(&job_id) {
|
|
(job.config.clone(), job.model_type.clone())
|
|
} else {
|
|
return Err(anyhow::anyhow!("Job {} not found", job_id));
|
|
}
|
|
};
|
|
|
|
// Allocate resources
|
|
let resource_allocation = {
|
|
let mut resources = available_resources.lock().await;
|
|
if let Some(resource) = resources.pop() {
|
|
resource
|
|
} else {
|
|
error!("No available resources for job {}", job_id);
|
|
return Self::mark_job_failed(
|
|
job_id,
|
|
"No available resources".to_string(),
|
|
jobs,
|
|
database,
|
|
)
|
|
.await;
|
|
}
|
|
};
|
|
|
|
// Assign resource to job
|
|
{
|
|
let mut assignments = resource_assignments.write().await;
|
|
assignments.insert(job_id, resource_allocation.clone());
|
|
}
|
|
|
|
// Update job status to running
|
|
{
|
|
let mut jobs_write = jobs.write().await;
|
|
if let Some(job) = jobs_write.get_mut(&job_id) {
|
|
job.status = JobStatus::Running;
|
|
job.started_at = Some(Utc::now());
|
|
}
|
|
}
|
|
|
|
// Broadcast status update
|
|
Self::send_status_update(
|
|
job_id,
|
|
JobStatus::Running,
|
|
0.0,
|
|
0,
|
|
job_config.training_params.max_epochs as u32,
|
|
HashMap::new(),
|
|
"Training started".to_string(),
|
|
None,
|
|
&resource_allocation,
|
|
status_broadcasters,
|
|
)
|
|
.await;
|
|
|
|
// Execute training using existing ML infrastructure
|
|
let training_result = Self::execute_training(
|
|
job_id,
|
|
job_config,
|
|
model_type,
|
|
&resource_allocation,
|
|
status_broadcasters,
|
|
)
|
|
.await;
|
|
|
|
// Handle training result
|
|
match training_result {
|
|
Ok(result) => {
|
|
Self::handle_training_success(job_id, result, jobs, database, storage).await?;
|
|
},
|
|
Err(e) => {
|
|
Self::mark_job_failed(job_id, e.to_string(), jobs, database).await?;
|
|
},
|
|
}
|
|
|
|
// Release resources
|
|
{
|
|
let mut resources = available_resources.lock().await;
|
|
resources.push(resource_allocation);
|
|
}
|
|
{
|
|
let mut assignments = resource_assignments.write().await;
|
|
assignments.remove(&job_id);
|
|
}
|
|
|
|
info!("Worker {} completed job {}", worker_id, job_id);
|
|
Ok(())
|
|
}
|
|
|
|
/// Execute training using the existing ML training infrastructure
|
|
async fn execute_training(
|
|
job_id: Uuid,
|
|
config: ProductionTrainingConfig,
|
|
model_type: String,
|
|
_resource_allocation: &ResourceAllocation,
|
|
_status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
) -> Result<TrainingResult> {
|
|
info!(
|
|
"Starting training execution for job {} with model type {}",
|
|
job_id, model_type
|
|
);
|
|
|
|
// Create the production training system
|
|
let training_system = ProductionMLTrainingSystem::new(config.clone())
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("Failed to create training system: {:?}", e))?;
|
|
|
|
// Load training data from configured source
|
|
// Phase 1: Configuration established, Phase 2-3 will implement actual loading
|
|
let (training_data, validation_data) = Self::load_training_data().await?;
|
|
|
|
// Execute training with progress callbacks
|
|
let result = training_system
|
|
.train_model(training_data, Some(validation_data))
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("Training failed: {:?}", e))?;
|
|
|
|
info!(
|
|
"Training completed for job {} with final loss: {:.6}",
|
|
job_id, result.final_train_loss
|
|
);
|
|
Ok(result)
|
|
}
|
|
|
|
/// Load training data from configured source with 225-feature extraction
|
|
///
|
|
/// Loads real market data from DBN files and applies full 225-feature extraction
|
|
/// (Wave C 201 features + Wave D 24 features) using FeatureExtractor.
|
|
pub async fn load_training_data() -> Result<(
|
|
Vec<(FinancialFeatures, Vec<f64>)>,
|
|
Vec<(FinancialFeatures, Vec<f64>)>,
|
|
)> {
|
|
|
|
|
|
// Primary: Try to load real DBN market data
|
|
let dbn_file_path = std::env::var("DBN_DATA_FILE").unwrap_or_else(|_| {
|
|
"test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn".to_string()
|
|
});
|
|
|
|
if std::path::Path::new(&dbn_file_path).exists() {
|
|
info!(
|
|
"📊 Loading REAL market data from DBN file with 225-feature extraction: {}",
|
|
dbn_file_path
|
|
);
|
|
|
|
match Self::load_training_data_with_225_features(&dbn_file_path, 0.8).await {
|
|
Ok((training_data, validation_data)) => {
|
|
info!(
|
|
"✅ Loaded {} training samples, {} validation samples with 225 features (Wave C + Wave D)",
|
|
training_data.len(),
|
|
validation_data.len()
|
|
);
|
|
return Ok((training_data, validation_data));
|
|
},
|
|
Err(e) => {
|
|
warn!("Failed to load DBN data with 225-feature extraction: {}, falling back to database", e);
|
|
},
|
|
}
|
|
} else {
|
|
debug!("DBN file not found at {}, trying database", dbn_file_path);
|
|
}
|
|
|
|
// Fallback: Try database loading or mock data
|
|
#[cfg(feature = "mock-data")]
|
|
{
|
|
error!("Using MOCK training data - NOT FOR PRODUCTION USE!");
|
|
error!("Rebuild without --features mock-data for production");
|
|
let training_data = Self::generate_mock_training_data()?;
|
|
let validation_data = Self::generate_mock_validation_data()?;
|
|
return Ok((training_data, validation_data));
|
|
}
|
|
|
|
#[cfg(not(feature = "mock-data"))]
|
|
{
|
|
use crate::data_config::DataSourceType;
|
|
use crate::data_loader::HistoricalDataLoader;
|
|
|
|
// Load data source configuration
|
|
let data_config = TrainingDataSourceConfig::from_env()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load data source configuration: {}", e))?;
|
|
|
|
data_config
|
|
.validate()
|
|
.map_err(|e| anyhow::anyhow!("Invalid data source configuration: {}", e))?;
|
|
|
|
info!(
|
|
"📊 Training data configuration loaded: {:?}",
|
|
data_config.summary()
|
|
);
|
|
|
|
// Load data based on source type
|
|
match data_config.source_type {
|
|
DataSourceType::Historical | DataSourceType::Hybrid => {
|
|
info!("Loading historical training data from PostgreSQL");
|
|
|
|
let mut loader = HistoricalDataLoader::new(data_config)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("Failed to create data loader: {}", e))?;
|
|
|
|
let (training_data, validation_data) = loader
|
|
.load_training_data()
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("Failed to load training data: {}", e))?;
|
|
|
|
info!(
|
|
"✅ Loaded {} training samples, {} validation samples from database",
|
|
training_data.len(),
|
|
validation_data.len()
|
|
);
|
|
|
|
Ok((training_data, validation_data))
|
|
},
|
|
DataSourceType::RealTime => Err(anyhow::anyhow!(
|
|
"❌ RealTime data source not yet implemented (Phase 3)\n\
|
|
\n\
|
|
📋 Supported data sources:\n\
|
|
- DBN: Real market data files (Phase 2 ✅)\n\
|
|
- Historical: PostgreSQL database (Phase 2 ✅)\n\
|
|
- RealTime: Live streaming data (Phase 3 pending)\n\
|
|
- Hybrid: Historical + RealTime (Phase 3 pending)\n\
|
|
- Parquet: S3 parquet files (Phase 4 pending)\n\
|
|
\n\
|
|
🔧 Set DBN_DATA_FILE=/path/to/file.dbn for real market data\n\
|
|
Set DATA_SOURCE_TYPE=historical to use database loading\n\
|
|
Set DATABASE_URL to your PostgreSQL instance"
|
|
)),
|
|
DataSourceType::Parquet => Err(anyhow::anyhow!(
|
|
"❌ Parquet data source not yet implemented (Phase 4)\n\
|
|
\n\
|
|
📋 Supported data sources:\n\
|
|
- DBN: Real market data files (Phase 2 ✅)\n\
|
|
- Historical: PostgreSQL database (Phase 2 ✅)\n\
|
|
- Parquet: S3 parquet files (Phase 4 pending)\n\
|
|
\n\
|
|
🔧 Set DBN_DATA_FILE=/path/to/file.dbn for real market data\n\
|
|
Set DATA_SOURCE_TYPE=historical to use database loading\n\
|
|
Set DATABASE_URL to your PostgreSQL instance"
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Load training data from DBN file with 225-feature extraction (Wave C + Wave D)
|
|
///
|
|
/// Uses FeatureExtractor to produce all 225 features per bar:
|
|
/// - Wave C (indices 0-200): 201 features
|
|
/// - Wave D (indices 201-224): 24 regime detection features
|
|
async fn load_training_data_with_225_features(
|
|
dbn_file_path: &str,
|
|
train_split: f64,
|
|
) -> Result<(
|
|
Vec<(FinancialFeatures, Vec<f64>)>,
|
|
Vec<(FinancialFeatures, Vec<f64>)>,
|
|
)> {
|
|
use dbn::decode::{DbnDecoder, DecodeRecordRef};
|
|
use dbn::OhlcvMsg;
|
|
use ml::features::extraction::{FeatureExtractor, OHLCVBar};
|
|
use anyhow::Context;
|
|
use chrono::TimeZone;
|
|
|
|
// Load OHLCV bars from DBN file (reuse existing loader logic)
|
|
debug!("Loading DBN file for 225-feature extraction: {}", dbn_file_path);
|
|
|
|
let mut decoder = DbnDecoder::from_file(dbn_file_path)
|
|
.context(format!("Failed to create DBN decoder for file: {}", dbn_file_path))?;
|
|
|
|
let mut ohlcv_bars = Vec::new();
|
|
|
|
// Decode all OHLCV records from DBN file
|
|
while let Some(record_ref) = decoder.decode_record_ref()
|
|
.context("Failed to decode DBN record")?
|
|
{
|
|
if let Some(ohlcv_msg) = record_ref.get::<OhlcvMsg>() {
|
|
// Convert timestamp from nanoseconds
|
|
let ts_nanos = ohlcv_msg.hd.ts_event as i64;
|
|
let secs = ts_nanos / 1_000_000_000;
|
|
let nanos = (ts_nanos % 1_000_000_000) as u32;
|
|
let timestamp = chrono::Utc
|
|
.timestamp_opt(secs, nanos)
|
|
.single()
|
|
.ok_or_else(|| anyhow::anyhow!("Invalid timestamp: {}", ts_nanos))?;
|
|
|
|
// Convert DBN fixed-point prices to f64 (9 decimal places)
|
|
let dbn_to_f64 = |price: i64| price as f64 / 1_000_000_000.0;
|
|
|
|
ohlcv_bars.push(OHLCVBar {
|
|
timestamp,
|
|
open: dbn_to_f64(ohlcv_msg.open),
|
|
high: dbn_to_f64(ohlcv_msg.high),
|
|
low: dbn_to_f64(ohlcv_msg.low),
|
|
close: dbn_to_f64(ohlcv_msg.close),
|
|
volume: ohlcv_msg.volume as f64,
|
|
});
|
|
}
|
|
}
|
|
|
|
info!("Loaded {} OHLCV bars from DBN file", ohlcv_bars.len());
|
|
|
|
// Extract 225 features using FeatureExtractor (Wave C + Wave D)
|
|
const WARMUP_PERIOD: usize = 50;
|
|
|
|
if ohlcv_bars.len() < WARMUP_PERIOD {
|
|
return Err(anyhow::anyhow!(
|
|
"Insufficient data: {} bars provided, {} required for warmup",
|
|
ohlcv_bars.len(),
|
|
WARMUP_PERIOD
|
|
));
|
|
}
|
|
|
|
let mut extractor = FeatureExtractor::new();
|
|
let mut feature_vectors_225 = Vec::with_capacity(ohlcv_bars.len() - WARMUP_PERIOD);
|
|
|
|
// Feed bars sequentially to build rolling windows
|
|
for (i, bar) in ohlcv_bars.iter().enumerate() {
|
|
extractor.update(bar).context(format!("Feature extraction failed at bar {}", i))?;
|
|
|
|
// Start extracting features after warmup
|
|
if i >= WARMUP_PERIOD {
|
|
let features_225 = extractor.extract_current_features()
|
|
.context(format!("Failed to extract features at bar {}", i))?;
|
|
feature_vectors_225.push(features_225);
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"Extracted {} feature vectors (225 dimensions each, Wave C + Wave D)",
|
|
feature_vectors_225.len()
|
|
);
|
|
|
|
// Convert [f64; 225] to (FinancialFeatures, Vec<f64>) format for compatibility.
|
|
//
|
|
// This conversion exists because UnifiedTrainable::train() accepts
|
|
// Vec<(FinancialFeatures, Vec<f64>)>. The FinancialFeatures fields below are
|
|
// stubs — only the Vec<f64> component is consumed by the underlying models.
|
|
//
|
|
// To remove this layer: change UnifiedTrainable::train() to accept &[Vec<f64>]
|
|
// directly and update all 4 model adapters (DQN, PPO, TFT, Mamba2).
|
|
let training_samples: Vec<(FinancialFeatures, Vec<f64>)> = feature_vectors_225
|
|
.iter()
|
|
.zip(ohlcv_bars.iter().skip(WARMUP_PERIOD))
|
|
.map(|(features_225, bar)| {
|
|
// Create stub FinancialFeatures for compatibility
|
|
let close_price = common::Price::new(bar.close).unwrap_or_default();
|
|
let financial_features = FinancialFeatures {
|
|
timestamp: bar.timestamp,
|
|
prices: vec![close_price],
|
|
volumes: vec![bar.volume as i64],
|
|
technical_indicators: std::collections::HashMap::new(),
|
|
microstructure: ml::training_pipeline::MicrostructureFeatures {
|
|
spread_bps: 0,
|
|
imbalance: 0.0,
|
|
trade_intensity: 0.0,
|
|
vwap: close_price,
|
|
},
|
|
risk_metrics: ml::training_pipeline::RiskFeatures {
|
|
var_5pct: 0.0,
|
|
expected_shortfall: 0.0,
|
|
max_drawdown: 0.0,
|
|
sharpe_ratio: 0.0,
|
|
},
|
|
};
|
|
|
|
// Convert [f64; 225] to Vec<f64>
|
|
(financial_features, features_225.to_vec())
|
|
})
|
|
.collect();
|
|
|
|
// Split into train/validation (preserve time-series ordering)
|
|
let split_idx = (training_samples.len() as f64 * train_split) as usize;
|
|
let train_data = training_samples[..split_idx].to_vec();
|
|
let val_data = training_samples[split_idx..].to_vec();
|
|
|
|
info!(
|
|
"✅ Created {} training samples, {} validation samples (225 features each)",
|
|
train_data.len(),
|
|
val_data.len()
|
|
);
|
|
|
|
Ok((train_data, val_data))
|
|
}
|
|
|
|
/// Generate mock training data for demonstration
|
|
#[cfg(feature = "mock-data")]
|
|
fn generate_mock_training_data() -> Result<Vec<(FinancialFeatures, Vec<f64>)>> {
|
|
// This is a simplified mock implementation
|
|
// In production, this would load real market data
|
|
let mut data = Vec::new();
|
|
|
|
for i in 0..1000 {
|
|
let price = 100.0 + (i as f64 * 0.01);
|
|
let features = FinancialFeatures {
|
|
prices: vec![
|
|
common::Price::from_f64(price).unwrap_or_default()
|
|
],
|
|
volumes: vec![1000 + i as i64],
|
|
technical_indicators: [("rsi".to_string(), 0.5 + 0.3 * (i as f64 / 1000.0).sin())]
|
|
.iter()
|
|
.cloned()
|
|
.collect(),
|
|
microstructure: ml::training_pipeline::MicrostructureFeatures {
|
|
spread_bps: 10,
|
|
imbalance: 0.1 * (i as f64 / 100.0).sin(),
|
|
trade_intensity: 2.5,
|
|
vwap: common::Price::from_f64(price * 0.9995).unwrap_or_default(),
|
|
},
|
|
risk_metrics: ml::training_pipeline::RiskFeatures {
|
|
var_5pct: -0.02,
|
|
expected_shortfall: -0.03,
|
|
max_drawdown: -0.05,
|
|
sharpe_ratio: 1.2,
|
|
},
|
|
timestamp: chrono::Utc::now(),
|
|
};
|
|
|
|
let targets = vec![price * 1.001]; // Predict small price increase
|
|
data.push((features, targets));
|
|
}
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
/// Generate mock validation data
|
|
#[cfg(feature = "mock-data")]
|
|
fn generate_mock_validation_data() -> Result<Vec<(FinancialFeatures, Vec<f64>)>> {
|
|
// Similar to training data but smaller
|
|
Self::generate_mock_training_data().map(|mut data| {
|
|
data.truncate(200);
|
|
data
|
|
})
|
|
}
|
|
|
|
/// Handle successful training completion
|
|
async fn handle_training_success(
|
|
job_id: Uuid,
|
|
result: TrainingResult,
|
|
jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
|
|
_database: &Arc<DatabaseManager>,
|
|
storage: &Arc<ModelStorageManager>,
|
|
) -> Result<()> {
|
|
// Serialize training result to model data (simplified)
|
|
let model_data = format!(
|
|
"{{\"final_train_loss\":{},\"final_val_loss\":{}}}",
|
|
result.final_train_loss, result.final_val_loss
|
|
)
|
|
.into_bytes();
|
|
|
|
// Attempt to store model artifact with S3 upload
|
|
let model_path = match storage.store_model(job_id, &model_data).await {
|
|
Ok(path) => {
|
|
info!(
|
|
"Successfully uploaded model for job {} to storage: {}",
|
|
job_id, path
|
|
);
|
|
path
|
|
},
|
|
Err(e) => {
|
|
error!("Failed to upload model for job {}: {}", job_id, e);
|
|
// Fallback to local path for tracking
|
|
let fallback_path = format!("models/{}.bin", job_id);
|
|
warn!("Using fallback path for job {}: {}", job_id, fallback_path);
|
|
fallback_path
|
|
},
|
|
};
|
|
|
|
// Extract accuracy from training result metrics history
|
|
let accuracy = result
|
|
.metrics_history
|
|
.last()
|
|
.map(|m| m.prediction_accuracy)
|
|
.unwrap_or(0.0);
|
|
|
|
// Create model metadata for tracking and versioning
|
|
let model_metadata = {
|
|
let job_guard = jobs.read().await;
|
|
if let Some(job) = job_guard.get(&job_id) {
|
|
config::ModelRegistryEntry {
|
|
id: job_id,
|
|
name: job.model_type.clone(),
|
|
version: format!("v{}", chrono::Utc::now().format("%Y%m%d_%H%M%S")),
|
|
created_at: chrono::Utc::now(),
|
|
updated_at: chrono::Utc::now(),
|
|
training_metrics: config::TrainingMetrics {
|
|
accuracy,
|
|
loss: result.final_train_loss,
|
|
validation_accuracy: accuracy, // Same as training accuracy for now
|
|
validation_loss: result.final_val_loss,
|
|
epochs: result.epochs_trained as u32,
|
|
training_time_seconds: result.training_duration.num_seconds() as f64,
|
|
},
|
|
architecture: config::ModelArchitecture {
|
|
model_type: job.model_type.clone(),
|
|
input_dim: job.config.model_config.input_dim,
|
|
output_dim: job.config.model_config.output_dim,
|
|
hidden_layers: job.config.model_config.hidden_dims.clone(),
|
|
activation: job.config.model_config.activation.clone(),
|
|
optimizer: "adamw".to_string(), // Standard optimizer for production ML
|
|
learning_rate: job.config.training_params.learning_rate,
|
|
},
|
|
}
|
|
} else {
|
|
return Err(anyhow::anyhow!(
|
|
"Job {} not found for metadata creation",
|
|
job_id
|
|
));
|
|
}
|
|
};
|
|
|
|
info!(
|
|
"Created model metadata for job {} with version {}",
|
|
job_id, model_metadata.version
|
|
);
|
|
|
|
// Update job status
|
|
{
|
|
let mut jobs_write = jobs.write().await;
|
|
if let Some(job) = jobs_write.get_mut(&job_id) {
|
|
job.status = JobStatus::Completed;
|
|
job.completed_at = Some(Utc::now());
|
|
job.progress_percentage = 100.0;
|
|
job.model_artifact_path = Some(model_path.clone());
|
|
job.metrics
|
|
.insert("final_train_loss".to_string(), result.final_train_loss);
|
|
job.metrics
|
|
.insert("final_val_loss".to_string(), result.final_val_loss);
|
|
job.metrics.insert(
|
|
"model_version".to_string(),
|
|
model_metadata.version.parse().unwrap_or(0.0),
|
|
);
|
|
job.metrics
|
|
.insert("model_size_bytes".to_string(), model_data.len() as f64);
|
|
}
|
|
}
|
|
|
|
info!("Training job {} completed successfully", job_id);
|
|
Ok(())
|
|
}
|
|
|
|
/// Mark job as failed
|
|
async fn mark_job_failed(
|
|
job_id: Uuid,
|
|
error_message: String,
|
|
jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
|
|
_database: &Arc<DatabaseManager>,
|
|
) -> Result<()> {
|
|
{
|
|
let mut jobs_write = jobs.write().await;
|
|
if let Some(job) = jobs_write.get_mut(&job_id) {
|
|
job.status = JobStatus::Failed;
|
|
job.completed_at = Some(Utc::now());
|
|
job.error_message = Some(error_message.clone());
|
|
}
|
|
}
|
|
|
|
error!("Training job {} failed: {}", job_id, error_message);
|
|
Ok(())
|
|
}
|
|
|
|
/// Send status update
|
|
async fn send_status_update(
|
|
job_id: Uuid,
|
|
status: JobStatus,
|
|
progress: f32,
|
|
current_epoch: u32,
|
|
total_epochs: u32,
|
|
metrics: HashMap<String, f64>,
|
|
message: String,
|
|
financial_metrics: Option<ml::training_pipeline::FinancialPerformanceMetrics>,
|
|
resource_allocation: &ResourceAllocation,
|
|
status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
) {
|
|
let update = TrainingStatusUpdate {
|
|
job_id,
|
|
status,
|
|
progress_percentage: progress,
|
|
current_epoch,
|
|
total_epochs,
|
|
metrics,
|
|
message,
|
|
timestamp: Utc::now(),
|
|
financial_metrics,
|
|
resource_usage: ResourceUsage {
|
|
cpu_usage_percent: 75.0, // Mock values
|
|
memory_usage_gb: resource_allocation.memory_gb as f32,
|
|
gpu_usage_percent: resource_allocation.gpu_id.map(|_| 80.0),
|
|
gpu_memory_usage_gb: resource_allocation.gpu_id.map(|_| 4.0),
|
|
active_workers: 1,
|
|
},
|
|
};
|
|
|
|
let broadcasters = status_broadcasters.read().await;
|
|
if let Some(broadcaster) = broadcasters.get(&job_id) {
|
|
match broadcaster.send(update.clone()) {
|
|
Ok(_) => {
|
|
debug!("Status update sent successfully for job {}", job_id);
|
|
},
|
|
Err(broadcast::error::SendError(_)) => {
|
|
warn!(
|
|
"Failed to send status update for job {} - all receivers dropped. \
|
|
This indicates clients have disconnected.",
|
|
job_id
|
|
);
|
|
// Note: The broadcast channel automatically handles the case where
|
|
// receivers are lagging behind. Slow receivers will miss old messages
|
|
// but will still receive new ones, providing natural back-pressure.
|
|
},
|
|
}
|
|
} else {
|
|
debug!("No broadcaster found for job {}", job_id);
|
|
}
|
|
}
|
|
|
|
/// Broadcast status update for a job
|
|
async fn broadcast_status_update(&self, job_id: Uuid, message: String) {
|
|
let jobs = self.jobs.read().await;
|
|
if let Some(job) = jobs.get(&job_id) {
|
|
Self::send_status_update(
|
|
job_id,
|
|
job.status.clone(),
|
|
job.progress_percentage,
|
|
job.current_epoch,
|
|
job.total_epochs,
|
|
job.metrics.clone(),
|
|
message,
|
|
None,
|
|
&ResourceAllocation {
|
|
gpu_id: None,
|
|
cpu_cores: 1,
|
|
memory_gb: 1.0,
|
|
worker_id: 0,
|
|
},
|
|
&self.status_broadcasters,
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
/// Get job model type
|
|
async fn get_job_model_type(&self, job_id: Uuid) -> Option<String> {
|
|
let jobs = self.jobs.read().await;
|
|
jobs.get(&job_id).map(|job| job.model_type.clone())
|
|
}
|
|
|
|
/// Spawn resource monitor
|
|
async fn spawn_resource_monitor(&self) -> Result<tokio::task::JoinHandle<()>> {
|
|
let interval = Duration::from_secs(30);
|
|
|
|
let handle = tokio::spawn(async move {
|
|
let mut interval_timer = tokio::time::interval(interval);
|
|
|
|
loop {
|
|
interval_timer.tick().await;
|
|
|
|
// Monitor system resources
|
|
// This is a simplified implementation
|
|
debug!("Resource monitor tick - system resources OK");
|
|
}
|
|
});
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
/// Periodically send snapshot status updates for all active jobs
|
|
///
|
|
/// This provides a fallback mechanism when streaming updates are overwhelmed
|
|
async fn spawn_snapshot_broadcaster(&self) -> Result<tokio::task::JoinHandle<()>> {
|
|
let jobs = Arc::clone(&self.jobs);
|
|
let status_broadcasters = Arc::clone(&self.status_broadcasters);
|
|
let interval = Duration::from_secs(30); // Default snapshot interval
|
|
let cancellation_token = self.cancellation_token.child_token();
|
|
|
|
let handle = tokio::spawn(async move {
|
|
let mut interval_timer = tokio::time::interval(interval);
|
|
|
|
loop {
|
|
tokio::select! {
|
|
_ = interval_timer.tick() => {
|
|
// Send snapshot updates for all active jobs
|
|
let jobs_read = jobs.read().await;
|
|
for job in jobs_read.values() {
|
|
if matches!(job.status, JobStatus::Running | JobStatus::Pending) {
|
|
Self::send_snapshot_update(job, &status_broadcasters).await;
|
|
}
|
|
}
|
|
|
|
// Clean up disconnected broadcasters periodically
|
|
Self::cleanup_disconnected_broadcasters_static(&status_broadcasters).await;
|
|
}
|
|
_ = cancellation_token.cancelled() => {
|
|
info!("Snapshot broadcaster received cancellation signal");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
/// Send a snapshot status update for a job
|
|
async fn send_snapshot_update(
|
|
job: &TrainingJob,
|
|
status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
) {
|
|
let snapshot_update = TrainingStatusUpdate {
|
|
job_id: job.id,
|
|
status: job.status.clone(),
|
|
progress_percentage: job.progress_percentage,
|
|
current_epoch: job.current_epoch,
|
|
total_epochs: job.total_epochs,
|
|
metrics: job.metrics.clone(),
|
|
message: "Snapshot status update".to_string(),
|
|
timestamp: Utc::now(),
|
|
financial_metrics: None,
|
|
resource_usage: ResourceUsage {
|
|
cpu_usage_percent: 0.0, // Snapshot doesn't have real-time resource data
|
|
memory_usage_gb: 0.0,
|
|
gpu_usage_percent: None,
|
|
gpu_memory_usage_gb: None,
|
|
active_workers: 0,
|
|
},
|
|
};
|
|
|
|
let broadcasters = status_broadcasters.read().await;
|
|
if let Some(broadcaster) = broadcasters.get(&job.id) {
|
|
// For snapshots, we use send (broadcast channels don't have try_send)
|
|
if let Err(e) = broadcaster.send(snapshot_update) {
|
|
tracing::warn!("Channel send failed (receiver dropped?): {:?}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Static version of cleanup for use in spawned tasks
|
|
async fn cleanup_disconnected_broadcasters_static(
|
|
status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
|
|
) {
|
|
let mut broadcasters = status_broadcasters.write().await;
|
|
let mut to_remove = Vec::new();
|
|
|
|
for (job_id, broadcaster) in broadcasters.iter() {
|
|
if broadcaster.receiver_count() == 0 {
|
|
to_remove.push(*job_id);
|
|
}
|
|
}
|
|
|
|
for job_id in to_remove {
|
|
broadcasters.remove(&job_id);
|
|
debug!("Cleaned up broadcaster for job {}", job_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Shutdown the orchestrator
|
|
impl Drop for TrainingOrchestrator {
|
|
fn drop(&mut self) {
|
|
info!("Shutting down training orchestrator");
|
|
|
|
// Cancel all worker tasks
|
|
for handle in self.worker_handles.drain(..) {
|
|
handle.abort();
|
|
}
|
|
}
|
|
}
|