395 lines
11 KiB
Rust
395 lines
11 KiB
Rust
//! HPC Channels integration for distributed training.
|
|
//!
|
|
//! This module bridges distributed training events to the hpc-channels message bus,
|
|
//! enabling real-time monitoring and coordination of multi-GPU/multi-node training.
|
|
//!
|
|
//! # Channels Used
|
|
//!
|
|
//! - `hpc.training.start` - Training job started
|
|
//! - `hpc.training.progress` - Training progress updates
|
|
//! - `hpc.training.metrics` - Training metrics (loss, accuracy, throughput)
|
|
//! - `hpc.training.status` - Training status changes
|
|
//! - `hpc.training.control` - Training control signals (pause, resume, stop)
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_distributed::channels::TrainingChannelBridge;
|
|
//!
|
|
//! let bridge = TrainingChannelBridge::new();
|
|
//!
|
|
//! // Publish training started
|
|
//! bridge.publish_training_start(&config);
|
|
//!
|
|
//! // Publish progress update
|
|
//! bridge.publish_progress(epoch, step, total_steps);
|
|
//!
|
|
//! // Publish metrics
|
|
//! bridge.publish_metrics(&metrics);
|
|
//! ```
|
|
|
|
use std::sync::Arc;
|
|
|
|
use tokio::sync::broadcast;
|
|
|
|
/// Training job started event.
|
|
#[derive(Clone, Debug)]
|
|
pub struct TrainingStartEvent {
|
|
/// Job identifier.
|
|
pub job_id: String,
|
|
/// Model name.
|
|
pub model_name: String,
|
|
/// Total number of epochs.
|
|
pub total_epochs: u32,
|
|
/// Batch size per GPU.
|
|
pub batch_size: u32,
|
|
/// Number of GPUs.
|
|
pub num_gpus: u32,
|
|
/// World size (total processes).
|
|
pub world_size: u32,
|
|
/// Timestamp (epoch ms).
|
|
pub timestamp_ms: u64,
|
|
}
|
|
|
|
/// Training progress update.
|
|
#[derive(Clone, Debug)]
|
|
pub struct TrainingProgressEvent {
|
|
/// Job identifier.
|
|
pub job_id: String,
|
|
/// Current epoch.
|
|
pub epoch: u32,
|
|
/// Current step within epoch.
|
|
pub step: u64,
|
|
/// Total steps in epoch.
|
|
pub total_steps: u64,
|
|
/// Estimated time remaining (seconds).
|
|
pub eta_seconds: Option<u64>,
|
|
/// Timestamp (epoch ms).
|
|
pub timestamp_ms: u64,
|
|
}
|
|
|
|
/// Training metrics snapshot.
|
|
#[derive(Clone, Debug)]
|
|
pub struct TrainingMetricsEvent {
|
|
/// Job identifier.
|
|
pub job_id: String,
|
|
/// Current loss.
|
|
pub loss: f64,
|
|
/// Learning rate.
|
|
pub learning_rate: f64,
|
|
/// Throughput (samples/second).
|
|
pub throughput: f64,
|
|
/// GPU memory used (bytes).
|
|
pub gpu_memory_used: u64,
|
|
/// GPU utilization (0-100).
|
|
pub gpu_utilization: f32,
|
|
/// Gradient norm.
|
|
pub grad_norm: Option<f64>,
|
|
/// Timestamp (epoch ms).
|
|
pub timestamp_ms: u64,
|
|
}
|
|
|
|
/// Training status.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum TrainingStatus {
|
|
/// Training is initializing.
|
|
Initializing,
|
|
/// Training is running.
|
|
Running,
|
|
/// Training is paused.
|
|
Paused,
|
|
/// Training completed successfully.
|
|
Completed,
|
|
/// Training failed.
|
|
Failed(String),
|
|
/// Training was cancelled.
|
|
Cancelled,
|
|
}
|
|
|
|
/// Training status change event.
|
|
#[derive(Clone, Debug)]
|
|
pub struct TrainingStatusEvent {
|
|
/// Job identifier.
|
|
pub job_id: String,
|
|
/// New status.
|
|
pub status: TrainingStatus,
|
|
/// Timestamp (epoch ms).
|
|
pub timestamp_ms: u64,
|
|
}
|
|
|
|
/// Training control commands.
|
|
#[derive(Clone, Debug)]
|
|
pub enum TrainingControlCommand {
|
|
/// Pause training.
|
|
Pause,
|
|
/// Resume training.
|
|
Resume,
|
|
/// Stop training gracefully.
|
|
Stop,
|
|
/// Checkpoint now.
|
|
CheckpointNow,
|
|
/// Adjust learning rate.
|
|
AdjustLr(f64),
|
|
}
|
|
|
|
/// Bridge between distributed training and hpc-channels.
|
|
pub struct TrainingChannelBridge {
|
|
/// Broadcast sender for training start events.
|
|
start_tx: broadcast::Sender<TrainingStartEvent>,
|
|
/// Broadcast sender for progress events.
|
|
progress_tx: broadcast::Sender<TrainingProgressEvent>,
|
|
/// Broadcast sender for metrics events.
|
|
metrics_tx: broadcast::Sender<TrainingMetricsEvent>,
|
|
/// Broadcast sender for status events.
|
|
status_tx: broadcast::Sender<TrainingStatusEvent>,
|
|
/// Broadcast sender for control commands.
|
|
control_tx: broadcast::Sender<TrainingControlCommand>,
|
|
}
|
|
|
|
impl TrainingChannelBridge {
|
|
/// Create a new training channel bridge.
|
|
///
|
|
/// Registers channels with the hpc-channels global registry.
|
|
pub fn new() -> Self {
|
|
let start_tx = hpc_channels::broadcast::<TrainingStartEvent>(
|
|
hpc_channels::channels::TRAINING_START,
|
|
256,
|
|
);
|
|
let progress_tx = hpc_channels::broadcast::<TrainingProgressEvent>(
|
|
hpc_channels::channels::TRAINING_PROGRESS,
|
|
1024,
|
|
);
|
|
let metrics_tx = hpc_channels::broadcast::<TrainingMetricsEvent>(
|
|
hpc_channels::channels::TRAINING_METRICS,
|
|
1024,
|
|
);
|
|
let status_tx = hpc_channels::broadcast::<TrainingStatusEvent>(
|
|
hpc_channels::channels::TRAINING_STATUS,
|
|
256,
|
|
);
|
|
let control_tx = hpc_channels::broadcast::<TrainingControlCommand>(
|
|
hpc_channels::channels::TRAINING_CONTROL,
|
|
64,
|
|
);
|
|
|
|
Self {
|
|
start_tx,
|
|
progress_tx,
|
|
metrics_tx,
|
|
status_tx,
|
|
control_tx,
|
|
}
|
|
}
|
|
|
|
fn now_ms() -> u64 {
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.map(|d| d.as_millis() as u64)
|
|
.unwrap_or(0)
|
|
}
|
|
|
|
/// Publish training started event.
|
|
pub fn publish_training_start(
|
|
&self,
|
|
job_id: &str,
|
|
model_name: &str,
|
|
total_epochs: u32,
|
|
batch_size: u32,
|
|
num_gpus: u32,
|
|
world_size: u32,
|
|
) {
|
|
let _ = self.start_tx.send(TrainingStartEvent {
|
|
job_id: job_id.to_string(),
|
|
model_name: model_name.to_string(),
|
|
total_epochs,
|
|
batch_size,
|
|
num_gpus,
|
|
world_size,
|
|
timestamp_ms: Self::now_ms(),
|
|
});
|
|
}
|
|
|
|
/// Publish training progress.
|
|
pub fn publish_progress(
|
|
&self,
|
|
job_id: &str,
|
|
epoch: u32,
|
|
step: u64,
|
|
total_steps: u64,
|
|
eta_seconds: Option<u64>,
|
|
) {
|
|
let _ = self.progress_tx.send(TrainingProgressEvent {
|
|
job_id: job_id.to_string(),
|
|
epoch,
|
|
step,
|
|
total_steps,
|
|
eta_seconds,
|
|
timestamp_ms: Self::now_ms(),
|
|
});
|
|
}
|
|
|
|
/// Publish training metrics.
|
|
pub fn publish_metrics(
|
|
&self,
|
|
job_id: &str,
|
|
loss: f64,
|
|
learning_rate: f64,
|
|
throughput: f64,
|
|
gpu_memory_used: u64,
|
|
gpu_utilization: f32,
|
|
grad_norm: Option<f64>,
|
|
) {
|
|
let _ = self.metrics_tx.send(TrainingMetricsEvent {
|
|
job_id: job_id.to_string(),
|
|
loss,
|
|
learning_rate,
|
|
throughput,
|
|
gpu_memory_used,
|
|
gpu_utilization,
|
|
grad_norm,
|
|
timestamp_ms: Self::now_ms(),
|
|
});
|
|
}
|
|
|
|
/// Publish training status change.
|
|
pub fn publish_status(&self, job_id: &str, status: TrainingStatus) {
|
|
let _ = self.status_tx.send(TrainingStatusEvent {
|
|
job_id: job_id.to_string(),
|
|
status,
|
|
timestamp_ms: Self::now_ms(),
|
|
});
|
|
}
|
|
|
|
/// Send a control command.
|
|
pub fn send_control(&self, command: TrainingControlCommand) {
|
|
let _ = self.control_tx.send(command);
|
|
}
|
|
|
|
/// Subscribe to training start events.
|
|
pub fn subscribe_start(&self) -> broadcast::Receiver<TrainingStartEvent> {
|
|
self.start_tx.subscribe()
|
|
}
|
|
|
|
/// Subscribe to progress events.
|
|
pub fn subscribe_progress(&self) -> broadcast::Receiver<TrainingProgressEvent> {
|
|
self.progress_tx.subscribe()
|
|
}
|
|
|
|
/// Subscribe to metrics events.
|
|
pub fn subscribe_metrics(&self) -> broadcast::Receiver<TrainingMetricsEvent> {
|
|
self.metrics_tx.subscribe()
|
|
}
|
|
|
|
/// Subscribe to status events.
|
|
pub fn subscribe_status(&self) -> broadcast::Receiver<TrainingStatusEvent> {
|
|
self.status_tx.subscribe()
|
|
}
|
|
|
|
/// Subscribe to control commands.
|
|
pub fn subscribe_control(&self) -> broadcast::Receiver<TrainingControlCommand> {
|
|
self.control_tx.subscribe()
|
|
}
|
|
}
|
|
|
|
impl Default for TrainingChannelBridge {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Shared channel bridge type.
|
|
pub type SharedTrainingChannelBridge = Arc<TrainingChannelBridge>;
|
|
|
|
/// Create a new shared channel bridge.
|
|
#[must_use]
|
|
pub fn shared_channel_bridge() -> SharedTrainingChannelBridge {
|
|
Arc::new(TrainingChannelBridge::new())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_bridge_creation() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
assert!(hpc_channels::exists(hpc_channels::channels::TRAINING_START));
|
|
assert!(hpc_channels::exists(
|
|
hpc_channels::channels::TRAINING_PROGRESS
|
|
));
|
|
assert!(hpc_channels::exists(
|
|
hpc_channels::channels::TRAINING_METRICS
|
|
));
|
|
let _ = bridge;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_training_start_publishing() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
let mut rx = bridge.subscribe_start();
|
|
|
|
bridge.publish_training_start("job-123", "resnet50", 100, 32, 8, 8);
|
|
|
|
let event = rx.recv().await.expect("Should receive event");
|
|
assert_eq!(event.job_id, "job-123");
|
|
assert_eq!(event.model_name, "resnet50");
|
|
assert_eq!(event.total_epochs, 100);
|
|
assert_eq!(event.batch_size, 32);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_progress_publishing() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
let mut rx = bridge.subscribe_progress();
|
|
|
|
bridge.publish_progress("job-123", 5, 1000, 5000, Some(3600));
|
|
|
|
let event = rx.recv().await.expect("Should receive event");
|
|
assert_eq!(event.epoch, 5);
|
|
assert_eq!(event.step, 1000);
|
|
assert_eq!(event.total_steps, 5000);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_metrics_publishing() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
let mut rx = bridge.subscribe_metrics();
|
|
|
|
bridge.publish_metrics(
|
|
"job-123",
|
|
0.05, // loss
|
|
0.001, // learning rate
|
|
1500.0, // throughput
|
|
8_000_000, // gpu memory
|
|
95.0, // gpu utilization
|
|
Some(1.5), // grad norm
|
|
);
|
|
|
|
let event = rx.recv().await.expect("Should receive event");
|
|
assert!((event.loss - 0.05).abs() < f64::EPSILON);
|
|
assert!((event.throughput - 1500.0).abs() < f64::EPSILON);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_status_publishing() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
let mut rx = bridge.subscribe_status();
|
|
|
|
bridge.publish_status("job-123", TrainingStatus::Running);
|
|
|
|
let event = rx.recv().await.expect("Should receive event");
|
|
assert_eq!(event.status, TrainingStatus::Running);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_control_commands() {
|
|
let bridge = TrainingChannelBridge::new();
|
|
let mut rx = bridge.subscribe_control();
|
|
|
|
bridge.send_control(TrainingControlCommand::Pause);
|
|
|
|
let cmd = rx.recv().await.expect("Should receive command");
|
|
assert!(matches!(cmd, TrainingControlCommand::Pause));
|
|
}
|
|
}
|