rustytorch/crates/training/rtx-auto/src/lib.rs

//! # RTX-Auto: Autonomous Platform Agents for RustyTorch++
//!
//! This crate provides autonomous optimization agents that continuously monitor,
//! analyze, and optimize the RustyTorch++ platform for maximum performance while
//! maintaining safety and reliability.
//!
//! ## Features
//!
//! - **Data Engineering Agent**: Optimizes data layouts and access patterns
//! - **Parallel Planner Agent**: Automatically selects optimal parallelization strategies
//! - **Quantization Guardian Agent**: Monitors accuracy while applying quantization
//! - **Kernel Synthesizer Agent**: Generates and optimizes GPU kernels
//! - **Proposal System**: Safe optimization proposal generation and validation
//! - **Rollback Management**: Automatic rollback on performance degradation
//!
//! ## Quick Start
//!
//! ```rust,no_run
//! use rtx_auto::{
//!     agents::{DataEngineeringAgent, ParallelPlannerAgent},
//!     proposal::ProposalValidator,
//!     rollback::RollbackManager,
//! };
//! use rtx_runtime::Runtime;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let runtime = Runtime::new()?;
//!
//!     // Create autonomous agents
//!     let data_agent = DataEngineeringAgent::new(&runtime)?;
//!     let parallel_agent = ParallelPlannerAgent::new(&runtime)?;
//!
//!     // Set up proposal validation and rollback
//!     let validator = ProposalValidator::new(&runtime)?;
//!     let rollback_manager = RollbackManager::new(&runtime)?;
//!
//!     println!("Autonomous optimization agents initialized");
//!     Ok(())
//! }
//! ```

pub mod agents;
pub mod error;
pub mod proposal;
pub mod rollback;

use std::sync::Arc;

// Re-export commonly used types
pub use agents::{
    AutonomousAgent, DataEngineeringAgent, KernelSynthesizerAgent, ParallelPlannerAgent,
    QuantGuardianAgent,
};
pub use error::{AutoError, AutoResult};
pub use proposal::{Proposal, ProposalStatus, ProposalType, ProposalValidator};
pub use rollback::{Checkpoint, CheckpointType, RollbackManager};

/// Version information for the autonomous optimization system.
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

/// Initialize the autonomous optimization system with default configuration.
pub async fn initialize_autonomous_system(
    runtime: Arc<rtx_runtime::Runtime>,
) -> AutoResult<AutonomousOptimizer> {
    tracing::info!("Initializing autonomous optimization system v{}", VERSION);

    let data_agent = DataEngineeringAgent::new(runtime.clone())?;
    let parallel_agent = ParallelPlannerAgent::new(runtime.clone())?;
    let quant_agent = QuantGuardianAgent::new(runtime.clone())?;
    let kernel_agent = KernelSynthesizerAgent::new(runtime.clone())?;

    let validator = ProposalValidator::new(runtime.clone())?;
    let rollback_manager = RollbackManager::new(runtime)?;

    Ok(AutonomousOptimizer {
        data_agent,
        parallel_agent,
        quant_agent,
        kernel_agent,
        validator,
        rollback_manager,
    })
}

/// Main autonomous optimizer that coordinates all agents.
pub struct AutonomousOptimizer {
    pub data_agent: DataEngineeringAgent,
    pub parallel_agent: ParallelPlannerAgent,
    pub quant_agent: QuantGuardianAgent,
    pub kernel_agent: KernelSynthesizerAgent,
    pub validator: ProposalValidator,
    pub rollback_manager: RollbackManager,
}

impl AutonomousOptimizer {
    /// Run a complete optimization cycle across all agents.
    pub async fn run_optimization_cycle(&mut self) -> AutoResult<Vec<Proposal>> {
        tracing::info!("Starting autonomous optimization cycle");

        let all_proposals = Vec::new();

        // Collect proposals from all agents
        // In practice, this would analyze current workloads and generate relevant proposals

        // Validate and rank all proposals
        let ranked_proposals = self.validator.rank_proposals(&all_proposals).await?;

        // Return the ranked proposals for potential application
        Ok(ranked_proposals
            .into_iter()
            .map(|(proposal, _score)| proposal)
            .collect())
    }

    /// Check the health status of all agents.
    pub fn check_agent_health(&self) -> std::collections::HashMap<String, bool> {
        let mut health_status = std::collections::HashMap::new();

        health_status.insert(
            self.data_agent.name().to_string(),
            self.data_agent.is_healthy(),
        );
        health_status.insert(
            self.parallel_agent.name().to_string(),
            self.parallel_agent.is_healthy(),
        );
        health_status.insert(
            self.quant_agent.name().to_string(),
            self.quant_agent.is_healthy(),
        );
        health_status.insert(
            self.kernel_agent.name().to_string(),
            self.kernel_agent.is_healthy(),
        );

        health_status
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_autonomous_system_initialization() {
        let runtime = Arc::new(rtx_runtime::Runtime::new().unwrap());
        let optimizer = initialize_autonomous_system(runtime).await;
        assert!(optimizer.is_ok(), "Failed to initialize autonomous system");
    }

    #[tokio::test]
    async fn test_agent_health_check() {
        let runtime = Arc::new(rtx_runtime::Runtime::new().unwrap());
        let optimizer = initialize_autonomous_system(runtime).await.unwrap();

        let health_status = optimizer.check_agent_health();
        assert_eq!(
            health_status.len(),
            4,
            "Should check health of all 4 agents"
        );

        for (agent_name, is_healthy) in health_status {
            assert!(is_healthy, "Agent {} should be healthy", agent_name);
        }
    }
}