Initial commit

crates/training/rtx-distributed/src/elastic_enhancements/factory.rs

@@ -0,0 +1,63 @@
+//! Shared Types and Factory Functions
+//!
+//! Provides shared type aliases and factory functions for elastic training components.
+
+use std::sync::Arc;
+
+use super::agent::{ElasticAgent, ElasticAgentConfig};
+use super::batch_size::{BatchSizeAdapter, BatchSizeConfig};
+use super::events::ElasticEventManager;
+use super::loss_scaling::{LossScalingRecovery, LossScalingRecoveryConfig};
+use super::lr_adaptation::{LRAdaptationConfig, LRAdapter};
+
+/// Shared elastic event manager
+pub type SharedElasticEventManager = Arc<ElasticEventManager>;
+
+/// Shared elastic agent
+pub type SharedElasticAgent = Arc<ElasticAgent>;
+
+/// Shared batch size adapter
+pub type SharedBatchSizeAdapter = Arc<BatchSizeAdapter>;
+
+/// Shared LR adapter
+pub type SharedLRAdapter = Arc<LRAdapter>;
+
+/// Shared loss scaling recovery
+pub type SharedLossScalingRecovery = Arc<LossScalingRecovery>;
+
+/// Create a shared elastic event manager
+pub fn shared_event_manager(max_history: usize) -> SharedElasticEventManager {
+    Arc::new(ElasticEventManager::new(max_history))
+}
+
+/// Create a shared elastic agent
+pub fn shared_elastic_agent(
+    config: ElasticAgentConfig,
+    event_manager: SharedElasticEventManager,
+) -> SharedElasticAgent {
+    Arc::new(ElasticAgent::new(config, event_manager))
+}
+
+/// Create a shared batch size adapter
+pub fn shared_batch_adapter(
+    config: BatchSizeConfig,
+    event_manager: SharedElasticEventManager,
+) -> SharedBatchSizeAdapter {
+    Arc::new(BatchSizeAdapter::new(config, event_manager))
+}
+
+/// Create a shared LR adapter
+pub fn shared_lr_adapter(
+    config: LRAdaptationConfig,
+    event_manager: SharedElasticEventManager,
+) -> SharedLRAdapter {
+    Arc::new(LRAdapter::new(config, event_manager))
+}
+
+/// Create a shared loss scaling recovery
+pub fn shared_loss_scaling_recovery(
+    config: LossScalingRecoveryConfig,
+    event_manager: SharedElasticEventManager,
+) -> SharedLossScalingRecovery {
+    Arc::new(LossScalingRecovery::new(config, event_manager))
+}