""" Main entry point for Mamba Swarm 100 units of 70M parameter Mamba encoders for distributed language modeling """ import os import sys import argparse import logging import asyncio from pathlib import Path from typing import Dict, Any, Optional # Add project root to path project_root = Path(__file__).parent sys.path.insert(0, str(project_root)) # Import core components from core.config import MambaSwarmConfig from system.mambaSwarm import SwarmEngine from system.inference import InferenceEngine from api.api_server import run_server from api.load_balancer import run_load_balancer, LoadBalancingStrategy from training.trainer import DistributedTrainer from monitoring.metrics import MambaSwarmMetrics from monitoring.profiler import MambaSwarmProfiler from monitoring.evaluator import MambaSwarmEvaluator from checkpoints.checkpoint_manager import CheckpointManager from training.trainer import setup_logging, get_device_info def setup_argument_parser(): """Setup command line argument parser""" parser = argparse.ArgumentParser(description="Mamba Swarm - Distributed Language Model") # Main mode selection parser.add_argument("mode", choices=["train", "serve", "evaluate", "load_balance"], help="Operation mode") # Configuration parser.add_argument("--config", type=str, default="config/default.yaml", help="Configuration file path") parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint to load") # Training arguments parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs") parser.add_argument("--batch-size", type=int, default=8, help="Training batch size") parser.add_argument("--learning-rate", type=float, default=1e-4, help="Learning rate") parser.add_argument("--data-path", type=str, default="data/", help="Training data path") # Serving arguments parser.add_argument("--host", type=str, default="0.0.0.0", help="Server host") parser.add_argument("--port", type=int, default=8000, help="Server port") parser.add_argument("--workers", type=int, default=1, help="Number of worker processes") # Load balancer arguments parser.add_argument("--servers", type=str, nargs="+", help="Backend server addresses (host:port)") parser.add_argument("--strategy", type=str, default="resource_aware", choices=["round_robin", "least_connections", "weighted_round_robin", "least_response_time", "hash_based", "resource_aware"], help="Load balancing strategy") # Evaluation arguments parser.add_argument("--eval-data", type=str, default="data/eval/", help="Evaluation data path") parser.add_argument("--output-report", type=str, default=None, help="Evaluation report output path") # System arguments parser.add_argument("--num-encoders", type=int, default=100, help="Number of Mamba encoders") parser.add_argument("--encoder-params", type=int, default=70000000, help="Parameters per encoder (70M)") parser.add_argument("--device", type=str, default="auto", help="Device to use (cuda, cpu, auto)") parser.add_argument("--distributed", action="store_true", help="Enable distributed training") # Monitoring arguments parser.add_argument("--enable-metrics", action="store_true", help="Enable metrics collection") parser.add_argument("--enable-profiling", action="store_true", help="Enable performance profiling") parser.add_argument("--metrics-port", type=int, default=9090, help="Metrics server port") # Logging parser.add_argument("--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Logging level") parser.add_argument("--log-file", type=str, default=None, help="Log file path") return parser async def train_mode(args, config: MambaSwarmConfig): """Training mode""" logging.info("Starting Mamba Swarm training...") # Initialize components metrics = MambaSwarmMetrics() if args.enable_metrics else None profiler = MambaSwarmProfiler() if args.enable_profiling else None # Initialize swarm engine swarm_engine = SwarmEngine(config) swarm_engine.initialize() # Initialize checkpoint manager checkpoint_manager = CheckpointManager( checkpoint_dir=config.checkpoint_dir, max_checkpoints=config.max_checkpoints, save_interval=config.save_interval ) # Load checkpoint if specified if args.checkpoint: checkpoint_data = checkpoint_manager.load_checkpoint(args.checkpoint) if checkpoint_data: swarm_engine.load_state_dict(checkpoint_data["model_state"]) logging.info(f"Loaded checkpoint: {args.checkpoint}") # Initialize trainer trainer = DistributedTrainer( swarm_engine=swarm_engine, config=config, checkpoint_manager=checkpoint_manager, metrics=metrics, profiler=profiler ) try: # Start monitoring if metrics: metrics.start_monitoring() if profiler: profiler.start_profiling() # Train model await trainer.train( data_path=args.data_path, epochs=args.epochs, batch_size=args.batch_size, learning_rate=args.learning_rate ) finally: # Cleanup if metrics: metrics.stop_monitoring() if profiler: profiler.cleanup() swarm_engine.shutdown() def serve_mode(args, config: MambaSwarmConfig): """API serving mode""" logging.info("Starting Mamba Swarm API server...") # Run API server run_server( host=args.host, port=args.port, workers=args.workers ) def load_balance_mode(args, config: MambaSwarmConfig): """Load balancer mode""" logging.info("Starting Mamba Swarm load balancer...") # Parse server addresses servers = [] for server_addr in args.servers or []: if ":" in server_addr: host, port = server_addr.split(":", 1) servers.append((host, int(port))) else: servers.append((server_addr, 8000)) # Default port if not servers: logging.error("No backend servers specified") return # Map strategy name to enum strategy_map = { "round_robin": LoadBalancingStrategy.ROUND_ROBIN, "least_connections": LoadBalancingStrategy.LEAST_CONNECTIONS, "weighted_round_robin": LoadBalancingStrategy.WEIGHTED_ROUND_ROBIN, "least_response_time": LoadBalancingStrategy.LEAST_RESPONSE_TIME, "hash_based": LoadBalancingStrategy.HASH_BASED, "resource_aware": LoadBalancingStrategy.RESOURCE_AWARE } strategy = strategy_map.get(args.strategy, LoadBalancingStrategy.RESOURCE_AWARE) # Run load balancer run_load_balancer( servers=servers, host=args.host, port=args.port, strategy=strategy ) async def evaluate_mode(args, config: MambaSwarmConfig): """Evaluation mode""" logging.info("Starting Mamba Swarm evaluation...") # Initialize swarm engine swarm_engine = SwarmEngine(config) swarm_engine.initialize() # Load checkpoint if specified if args.checkpoint: checkpoint_manager = CheckpointManager(config.checkpoint_dir) checkpoint_data = checkpoint_manager.load_checkpoint(args.checkpoint) if checkpoint_data: swarm_engine.load_state_dict(checkpoint_data["model_state"]) logging.info(f"Loaded checkpoint: {args.checkpoint}") # Initialize evaluator evaluator = MambaSwarmEvaluator(swarm_engine, config.__dict__) try: # Run comprehensive evaluation result = evaluator.run_comprehensive_evaluation() # Print results print(f"\nEvaluation Results:") print(f"Overall Score: {result.overall_score:.3f}") print(f"Execution Time: {result.execution_time:.2f}s") print(f"Total Metrics: {len(result.individual_metrics)}") # Print top metrics print(f"\nTop Metrics:") for metric in result.individual_metrics[:10]: print(f" {metric.metric_name}: {metric.score:.3f}") # Export report output_path = args.output_report or f"evaluation_report_{int(result.timestamp)}.json" report_file = evaluator.export_evaluation_report(result, output_path) print(f"\nDetailed report saved to: {report_file}") finally: swarm_engine.shutdown() def validate_config(args) -> MambaSwarmConfig: """Validate and create configuration""" # Load base configuration if os.path.exists(args.config): config = MambaSwarmConfig.from_file(args.config) else: logging.warning(f"Config file {args.config} not found, using defaults") config = MambaSwarmConfig() # Override with command line arguments if args.num_encoders: config.num_encoders = args.num_encoders if args.encoder_params: config.encoder_params = args.encoder_params # Device configuration if args.device == "auto": device_info = get_device_info() config.device = "cuda" if device_info["cuda_available"] else "cpu" else: config.device = args.device # Validate configuration total_params = config.num_encoders * config.encoder_params logging.info(f"Configuration: {config.num_encoders} encoders × {config.encoder_params/1e6:.0f}M params = {total_params/1e9:.1f}B total parameters") return config def main(): """Main entry point""" parser = setup_argument_parser() args = parser.parse_args() # Setup logging setup_logging( level=getattr(logging, args.log_level), log_file=args.log_file ) # Print banner print("=" * 60) print("🐍 Mamba Swarm - Distributed Language Model") print("100 × 70M Parameter Mamba Encoders") print("=" * 60) # Validate configuration try: config = validate_config(args) except Exception as e: logging.error(f"Configuration validation failed: {e}") sys.exit(1) # Print system information device_info = get_device_info() logging.info(f"System: {device_info['cpu_count']} CPUs, {device_info['memory_gb']:.1f}GB RAM") if device_info["cuda_available"]: logging.info(f"GPU: {device_info['gpu_count']} devices, {device_info['gpu_memory_gb']:.1f}GB VRAM") # Run mode-specific logic try: if args.mode == "train": asyncio.run(train_mode(args, config)) elif args.mode == "serve": serve_mode(args, config) elif args.mode == "load_balance": load_balance_mode(args, config) elif args.mode == "evaluate": asyncio.run(evaluate_mode(args, config)) else: logging.error(f"Unknown mode: {args.mode}") sys.exit(1) except KeyboardInterrupt: logging.info("Received interrupt signal, shutting down...") except Exception as e: logging.error(f"Application error: {e}", exc_info=True) sys.exit(1) logging.info("Mamba Swarm shutdown complete") def print_usage_examples(): """Print usage examples""" examples = """ Usage Examples: 1. Training: python main.py train --data-path ./data/train --epochs 10 --batch-size 8 --enable-metrics 2. Serving: python main.py serve --host 0.0.0.0 --port 8000 --checkpoint best_model.pt 3. Load Balancing: python main.py load_balance --servers localhost:8000 localhost:8001 localhost:8002 --strategy resource_aware 4. Evaluation: python main.py evaluate --checkpoint best_model.pt --eval-data ./data/eval --output-report eval_results.json 5. Distributed Training: python main.py train --distributed --num-encoders 100 --batch-size 4 --enable-profiling Configuration File Example (config.yaml): --- num_encoders: 100 encoder_params: 70000000 hidden_size: 2048 num_layers: 32 vocab_size: 50000 max_sequence_length: 2048 device: "auto" checkpoint_dir: "./checkpoints" max_checkpoints: 10 save_interval: 1000 learning_rate: 1e-4 warmup_steps: 1000 weight_decay: 0.01 gradient_clip_norm: 1.0 mixed_precision: true gradient_accumulation_steps: 8 """ print(examples) if __name__ == "__main__": # Check for help with examples if len(sys.argv) == 2 and sys.argv[1] in ["--help-examples", "-he"]: print_usage_examples() sys.exit(0) main()