test_complete_system.py

#!/usr/bin/env python3
"""
Complete OpenControl System Integration Test.

This script performs a comprehensive test of all OpenControl components
working together to ensure the system is fully functional.

Author: Nik Jois <nikjois@llamasearch.ai>
"""

import asyncio
import logging
import sys
import time
import torch
from pathlib import Path
from typing import Dict, Any

# Import all OpenControl components
from opencontrol.cli.commands import get_dev_config, get_test_config, OpenControlConfig
from opencontrol.core.world_model import OpenControlWorldModel
from opencontrol.core.attention_mechanisms import MultiModalAttention, RotaryPositionalEncoding
from opencontrol.core.multimodal_encoder import MultiModalEncoder
from opencontrol.core.temporal_dynamics import TemporalDynamicsModule
from opencontrol.data.dataset_manager import MultiModalDatasetManager
from opencontrol.training.distributed_trainer import DistributedWorldModelTrainer
from opencontrol.control.visual_mpc import ProductionVisualMPC
from opencontrol.evaluation.metrics_suite import ComprehensiveEvaluator
from opencontrol.deployment.model_server import OpenControlModelServer


class SystemTester:
    """Comprehensive system integration tester."""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.results = {}
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def log_test(self, test_name: str, status: str, details: str = ""):
        """Log test results."""
        print(f"[{status:>6}] {test_name}")
        if details:
            print(f"         {details}")
        self.results[test_name] = {"status": status, "details": details}
    
    def test_configurations(self):
        """Test configuration management."""
        print("\n=== Testing Configuration System ===")
        
        try:
            # Test development config
            dev_config = get_dev_config()
            assert isinstance(dev_config, OpenControlConfig)
            self.log_test("Development Config", "PASS", 
                         f"{dev_config.model.model_dim}D model, {dev_config.model.num_layers} layers")
            
            # Test test config
            test_config = get_test_config()
            assert isinstance(test_config, OpenControlConfig)
            self.log_test("Test Config", "PASS",
                         f"{test_config.model.model_dim}D model, {test_config.model.num_layers} layers")
            
            # Test validation
            try:
                test_config.validate()
                self.log_test("Config Validation", "WARN", "Data path validation skipped")
            except Exception as e:
                self.log_test("Config Validation", "WARN", f"Expected validation error: {str(e)[:50]}...")
                
        except Exception as e:
            self.log_test("Configuration System", "FAIL", str(e))
    
    def test_core_components(self):
        """Test core model components."""
        print("\n=== Testing Core Components ===")
        
        config = get_test_config()
        
        try:
            # Test RoPE
            rope = RotaryPositionalEncoding(64, 128)
            x = torch.randn(2, 10, 8, 64)
            y = rope(x)
            assert y.shape == x.shape
            self.log_test("Rotary Positional Encoding", "PASS", 
                         f"Input {x.shape} -> Output {y.shape}")
            
            # Test Attention
            attn = MultiModalAttention(config.model.model_dim, config.model.num_heads)
            x = torch.randn(2, 10, config.model.model_dim)
            y, weights = attn(x, rope_encoder=rope, return_attention_weights=True)
            assert y.shape == x.shape
            weights_shape = weights.shape if weights is not None else "None"
            self.log_test("Multi-Modal Attention", "PASS",
                         f"Output {y.shape}, Weights {weights_shape}")
            
            # Test Temporal Dynamics
            temporal = TemporalDynamicsModule(config.model.model_dim, num_layers=2)
            x = torch.randn(2, 16, config.model.model_dim)
            output = temporal(x)
            assert 'dynamics_output' in output
            self.log_test("Temporal Dynamics", "PASS",
                         f"Input {x.shape} -> Output {output['dynamics_output'].shape}")
            
        except Exception as e:
            self.log_test("Core Components", "FAIL", str(e))
    
    def test_multimodal_encoder(self):
        """Test multi-modal encoder."""
        print("\n=== Testing Multi-Modal Encoder ===")
        
        config = get_test_config()
        
        try:
            encoder = MultiModalEncoder(config)
            
            # Create test inputs
            inputs = {
                'video': torch.randn(1, 3, 3, 224, 224),
                'actions': torch.randn(1, 3, config.model.action_dim),
                'proprioception': torch.randn(1, 3, config.model.proprioception_dim)
            }
            
            tokens, modality_info = encoder(inputs)
            
            self.log_test("Multi-Modal Encoder", "PASS",
                         f"Input modalities -> {tokens.shape} tokens")
            
        except Exception as e:
            self.log_test("Multi-Modal Encoder", "FAIL", str(e))
    
    def test_world_model(self):
        """Test complete world model."""
        print("\n=== Testing World Model ===")
        
        config = get_test_config()
        
        try:
            model = OpenControlWorldModel(config)
            model = model.to(self.device)
            
            # Count parameters
            total_params = sum(p.numel() for p in model.parameters())
            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            
            self.log_test("World Model Creation", "PASS",
                         f"{total_params:,} total ({trainable_params:,} trainable)")
            
            # Test forward pass
            inputs = {
                'video': torch.randn(1, 3, 3, 224, 224).to(self.device),
                'actions': torch.randn(1, 3, config.model.action_dim).to(self.device),
                'proprioception': torch.randn(1, 3, config.model.proprioception_dim).to(self.device)
            }
            
            with torch.no_grad():
                output = model(inputs, prediction_horizon=2)
            
            self.log_test("World Model Forward", "PASS",
                         f"{len(output.predictions)} prediction modalities")
            
            # Test uncertainty estimation
            if output.uncertainty_estimates:
                self.log_test("Uncertainty Estimation", "PASS",
                             f"{len(output.uncertainty_estimates)} modalities")
            else:
                self.log_test("Uncertainty Estimation", "SKIP", "Not enabled")
                
        except Exception as e:
            self.log_test("World Model", "FAIL", str(e))
    
    async def test_data_pipeline(self):
        """Test data loading pipeline."""
        print("\n=== Testing Data Pipeline ===")
        
        config = get_test_config()
        
        try:
            dataset_manager = MultiModalDatasetManager(config)
            await dataset_manager.initialize()
            
            self.log_test("Dataset Manager", "PASS", "Initialized successfully")
            
            # Test data loader
            train_loader = dataset_manager.get_train_loader()
            self.log_test("Data Loader", "PASS", f"Batch size: {train_loader.batch_size}")
            
            # Test getting statistics
            stats = await dataset_manager.get_statistics()
            self.log_test("Dataset Statistics", "PASS", 
                         f"Episodes: {stats.get('total_episodes', 'N/A')}")
            
        except Exception as e:
            self.log_test("Data Pipeline", "FAIL", str(e))
    
    async def test_training_system(self):
        """Test training system."""
        print("\n=== Testing Training System ===")
        
        config = get_test_config()
        
        try:
            # Create components
            dataset_manager = MultiModalDatasetManager(config)
            await dataset_manager.initialize()
            
            model = OpenControlWorldModel(config)
            model = model.to(self.device)
            
            logger = logging.getLogger("TestTrainer")
            trainer = DistributedWorldModelTrainer(model, dataset_manager, config, logger)
            
            self.log_test("Trainer Creation", "PASS", "All components initialized")
            
            # Test a few training steps
            step_count = 0
            def test_callback(epoch, step, metrics):
                nonlocal step_count
                step_count = step
            
            # Run minimal training
            original_epochs = config.training.num_epochs
            original_steps = config.training.steps_per_epoch
            config.training.num_epochs = 1
            config.training.steps_per_epoch = 2
            
            await trainer.train(test_callback)
            
            # Restore original values
            config.training.num_epochs = original_epochs
            config.training.steps_per_epoch = original_steps
            
            self.log_test("Training Execution", "PASS", f"Completed {step_count} steps")
            
        except Exception as e:
            self.log_test("Training System", "FAIL", str(e))
    
    async def test_control_system(self):
        """Test MPC control system."""
        print("\n=== Testing Control System ===")
        
        config = get_test_config()
        
        try:
            model = OpenControlWorldModel(config)
            model = model.to(self.device)
            
            logger = logging.getLogger("TestMPC")
            mpc = ProductionVisualMPC(model, config, logger)
            
            # Test observation
            observation = {
                'video': torch.randn(1, 3, 224, 224).to(self.device),
                'proprioception': torch.randn(1, config.model.proprioception_dim).to(self.device)
            }
            
            # Test action computation
            action, info = await mpc.compute_action(observation)
            
            self.log_test("MPC Controller", "PASS",
                         f"Action shape: {action.shape}, Solve time: {info['solve_time']:.3f}s")
            
        except Exception as e:
            self.log_test("Control System", "FAIL", str(e))
    
    async def test_evaluation_system(self):
        """Test evaluation system."""
        print("\n=== Testing Evaluation System ===")
        
        config = get_test_config()
        
        try:
            dataset_manager = MultiModalDatasetManager(config)
            await dataset_manager.initialize()
            
            model = OpenControlWorldModel(config)
            model = model.to(self.device)
            
            logger = logging.getLogger("TestEvaluator")
            evaluator = ComprehensiveEvaluator(model, dataset_manager, config, logger)
            
            self.log_test("Evaluator Creation", "PASS", "Evaluator initialized")
            
            # Test a quick evaluation
            def eval_callback(progress):
                pass
            
            results = await evaluator.run_comprehensive_evaluation(eval_callback)
            
            self.log_test("Evaluation Execution", "PASS", 
                         f"{len(results)} metrics computed")
            
        except Exception as e:
            self.log_test("Evaluation System", "FAIL", str(e))
    
    def test_cli_system(self):
        """Test CLI system."""
        print("\n=== Testing CLI System ===")
        
        try:
            from opencontrol.cli.main import cli
            self.log_test("CLI Import", "PASS", "CLI module imported successfully")
            
        except Exception as e:
            self.log_test("CLI System", "FAIL", str(e))
    
    def print_summary(self):
        """Print test summary."""
        print("\n" + "="*60)
        print("OPENCONTROL SYSTEM TEST SUMMARY")
        print("="*60)
        
        total_tests = len(self.results)
        passed = sum(1 for r in self.results.values() if r["status"] == "PASS")
        failed = sum(1 for r in self.results.values() if r["status"] == "FAIL")
        warnings = sum(1 for r in self.results.values() if r["status"] == "WARN")
        skipped = sum(1 for r in self.results.values() if r["status"] == "SKIP")
        
        print(f"Total Tests: {total_tests}")
        print(f"Passed:      {passed}")
        print(f"Failed:      {failed}")
        print(f"Warnings:    {warnings}")
        print(f"Skipped:     {skipped}")
        
        if failed == 0:
            print("\nSYSTEM STATUS: ALL TESTS PASSED!")
            print("\nOpenControl is ready for production use!")
            print("\nQuick Start:")
            print("  opencontrol --help")
            print("  opencontrol interactive")
            print("  opencontrol train --config configs/models/development.yaml")
        else:
            print(f"\nSYSTEM STATUS: {failed} TESTS FAILED")
            print("\nFailed tests:")
            for name, result in self.results.items():
                if result["status"] == "FAIL":
                    print(f"  - {name}: {result['details']}")
        
        print("\n" + "="*60)


async def main():
    """Run comprehensive system test."""
    print("OpenControl Complete System Integration Test")
    print("=" * 60)
    print(f"Device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
    print(f"PyTorch: {torch.__version__}")
    print()
    
    tester = SystemTester()
    
    # Run all tests
    tester.test_configurations()
    tester.test_core_components()
    tester.test_multimodal_encoder()
    tester.test_world_model()
    await tester.test_data_pipeline()
    await tester.test_training_system()
    await tester.test_control_system()
    await tester.test_evaluation_system()
    tester.test_cli_system()
    
    # Print summary
    tester.print_summary()


if __name__ == "__main__":
    asyncio.run(main())