#!/usr/bin/env python3 """ Multi-Agent Training Deployment Script This script provides comprehensive deployment capabilities for the multi-agent training system, including Docker container management, environment setup, and training execution. """ import os import sys import json import yaml import argparse import subprocess import logging from pathlib import Path from typing import Dict, List, Optional, Any from dataclasses import dataclass # Add src to path for imports sys.path.append('src') from training.multi_agent_trainer import MultiAgentTrainingConfig, MultiAgentTrainingPipeline from datasets.multi_agent_loader import MultiAgentDatasetConfig @dataclass class DeploymentConfig: """Configuration for deployment""" project_root: str dataset_path: str model_repo_id: str dataset_repo_id: Optional[str] = None agents_file: Optional[str] = None config_file: Optional[str] = None docker_image_name: str = "phi35moe-cpu:latest" output_dir: str = "./outputs" logs_dir: str = "./logs" max_steps: int = 50 balance_agents: bool = True push_to_hub: bool = True dry_run: bool = False class MultiAgentTrainingDeployment: """ Comprehensive deployment manager for multi-agent training """ def __init__(self, config: DeploymentConfig): self.config = config self.setup_logging() def setup_logging(self): """Setup logging configuration""" logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(), logging.FileHandler(os.path.join(self.config.logs_dir, 'deployment.log')) ] ) self.logger = logging.getLogger(__name__) # Create logs directory os.makedirs(self.config.logs_dir, exist_ok=True) def validate_environment(self) -> bool: """Validate deployment environment""" self.logger.info("Validating deployment environment") # Check required environment variables required_env_vars = ["HF_TOKEN"] missing_vars = [] for var in required_env_vars: if not os.getenv(var): missing_vars.append(var) if missing_vars: self.logger.error(f"Missing required environment variables: {missing_vars}") return False # Check Docker availability try: result = subprocess.run(["docker", "--version"], capture_output=True, text=True) if result.returncode != 0: self.logger.error("Docker not available") return False except FileNotFoundError: self.logger.error("Docker not installed") return False # Check dataset path if not os.path.exists(self.config.dataset_path): self.logger.error(f"Dataset path not found: {self.config.dataset_path}") return False # Check agents file if specified if self.config.agents_file and not os.path.exists(self.config.agents_file): self.logger.error(f"Agents file not found: {self.config.agents_file}") return False self.logger.info("Environment validation passed") return True def build_docker_image(self) -> bool: """Build Docker image for training""" self.logger.info("Building Docker image") dockerfile_path = "docker/multi_agent_training/Dockerfile.cpu" if not os.path.exists(dockerfile_path): self.logger.error(f"Dockerfile not found: {dockerfile_path}") return False try: cmd = [ "docker", "build", "-f", dockerfile_path, "-t", self.config.docker_image_name, "docker/multi_agent_training/" ] self.logger.info(f"Running command: {' '.join(cmd)}") if not self.config.dry_run: result = subprocess.run(cmd, check=True, capture_output=True, text=True) self.logger.info("Docker image built successfully") else: self.logger.info("Dry run: Would build Docker image") return True except subprocess.CalledProcessError as e: self.logger.error(f"Docker build failed: {e}") self.logger.error(f"Error output: {e.stderr}") return False def create_training_command(self) -> List[str]: """Create training command for Docker execution""" cmd = [ "python", "/app/train_lora_cpu_multiagent.py", "--dataset_path", self.config.dataset_path, "--hub_repo_id", self.config.model_repo_id, "--output_dir", "/app/outputs", "--max_steps", str(self.config.max_steps), "--logging_steps", "5", "--save_steps", "50", "--eval_steps", "25" ] if self.config.balance_agents: cmd.append("--balance_agents") if self.config.push_to_hub: cmd.append("--push_to_hub") if self.config.agents_file: cmd.extend(["--agents_file", self.config.agents_file]) if self.config.dataset_repo_id: cmd.extend(["--push_dataset_repo", self.config.dataset_repo_id]) return cmd def run_training(self) -> bool: """Run training in Docker container""" self.logger.info("Starting training in Docker container") # Create training command training_cmd = self.create_training_command() # Prepare Docker run command docker_cmd = [ "docker", "run", "--rm", "-it", "-e", f"HF_TOKEN={os.getenv('HF_TOKEN')}", "-v", f"{os.path.abspath(self.config.dataset_path)}:{self.config.dataset_path}:ro", "-v", f"{os.path.abspath(self.config.output_dir)}:/app/outputs", "-v", f"{os.path.abspath(self.config.logs_dir)}:/app/logs", self.config.docker_image_name, "bash", "-lc" ] # Create full command with training script full_cmd = " ".join(training_cmd) docker_cmd.append(full_cmd) self.logger.info(f"Running command: {' '.join(docker_cmd[:-1])} '{full_cmd}'") try: if not self.config.dry_run: result = subprocess.run(docker_cmd, check=True) self.logger.info("Training completed successfully") else: self.logger.info("Dry run: Would execute training") return True except subprocess.CalledProcessError as e: self.logger.error(f"Training failed: {e}") return False def generate_deployment_report(self) -> Dict[str, Any]: """Generate deployment report""" report = { "deployment_config": { "project_root": self.config.project_root, "dataset_path": self.config.dataset_path, "model_repo_id": self.config.model_repo_id, "dataset_repo_id": self.config.dataset_repo_id, "docker_image_name": self.config.docker_image_name, "max_steps": self.config.max_steps, "balance_agents": self.config.balance_agents, "push_to_hub": self.config.push_to_hub }, "environment": { "hf_token_set": bool(os.getenv("HF_TOKEN")), "docker_available": self._check_docker_available(), "dataset_exists": os.path.exists(self.config.dataset_path), "agents_file_exists": os.path.exists(self.config.agents_file) if self.config.agents_file else None }, "deployment_status": "completed" if not self.config.dry_run else "dry_run" } return report def _check_docker_available(self) -> bool: """Check if Docker is available""" try: result = subprocess.run(["docker", "--version"], capture_output=True, text=True) return result.returncode == 0 except FileNotFoundError: return False def save_deployment_report(self, report: Dict[str, Any]): """Save deployment report to file""" report_file = os.path.join(self.config.logs_dir, "deployment_report.json") with open(report_file, 'w') as f: json.dump(report, f, indent=2) self.logger.info(f"Deployment report saved to {report_file}") def deploy(self) -> bool: """Execute complete deployment process""" self.logger.info("Starting multi-agent training deployment") try: # Validate environment if not self.validate_environment(): return False # Build Docker image if not self.build_docker_image(): return False # Run training if not self.run_training(): return False # Generate and save report report = self.generate_deployment_report() self.save_deployment_report(report) self.logger.info("Deployment completed successfully") return True except Exception as e: self.logger.error(f"Deployment failed: {e}") return False def load_config_from_file(config_file: str) -> Dict[str, Any]: """Load configuration from YAML file""" with open(config_file, 'r') as f: return yaml.safe_load(f) def create_deployment_config_from_yaml(config_data: Dict[str, Any], dataset_path: str, model_repo_id: str) -> DeploymentConfig: """Create deployment configuration from YAML data""" training_config = config_data.get("training", {}) multi_agent_config = config_data.get("multi_agent", {}) huggingface_config = config_data.get("huggingface", {}) return DeploymentConfig( project_root=os.getcwd(), dataset_path=dataset_path, model_repo_id=model_repo_id, dataset_repo_id=huggingface_config.get("push_dataset_repo"), agents_file=config_data.get("dataset", {}).get("agents_file"), config_file=None, # Will be set if provided docker_image_name="phi35moe-cpu:latest", output_dir=training_config.get("output_dir", "./outputs"), logs_dir=training_config.get("logging", {}).get("logging_dir", "./logs"), max_steps=training_config.get("max_steps", 50), balance_agents=multi_agent_config.get("balance_agents", True), push_to_hub=huggingface_config.get("push_to_hub", True) ) def main(): """Main deployment function""" parser = argparse.ArgumentParser(description="Deploy Multi-Agent Training System") # Required arguments parser.add_argument("--dataset_path", required=True, help="Path to multi-agent dataset") parser.add_argument("--model_repo_id", required=True, help="Hugging Face model repository ID") # Optional arguments parser.add_argument("--dataset_repo_id", default="", help="Optional dataset repository ID") parser.add_argument("--agents_file", default="", help="Optional agents YAML file") parser.add_argument("--config_file", default="", help="Optional configuration YAML file") parser.add_argument("--docker_image_name", default="phi35moe-cpu:latest", help="Docker image name") parser.add_argument("--output_dir", default="./outputs", help="Output directory") parser.add_argument("--logs_dir", default="./logs", help="Logs directory") parser.add_argument("--max_steps", type=int, default=50, help="Maximum training steps") parser.add_argument("--balance_agents", action="store_true", help="Balance dataset across agents") parser.add_argument("--push_to_hub", action="store_true", help="Push model to Hugging Face Hub") parser.add_argument("--dry_run", action="store_true", help="Perform dry run without actual execution") parser.add_argument("--log_level", default="INFO", help="Logging level") args = parser.parse_args() # Setup logging logging.basicConfig(level=getattr(logging, args.log_level.upper())) try: # Load configuration from file if provided if args.config_file and os.path.exists(args.config_file): config_data = load_config_from_file(args.config_file) deployment_config = create_deployment_config_from_yaml( config_data, args.dataset_path, args.model_repo_id ) else: # Create configuration from command line arguments deployment_config = DeploymentConfig( project_root=os.getcwd(), dataset_path=args.dataset_path, model_repo_id=args.model_repo_id, dataset_repo_id=args.dataset_repo_id if args.dataset_repo_id else None, agents_file=args.agents_file if args.agents_file else None, config_file=args.config_file if args.config_file else None, docker_image_name=args.docker_image_name, output_dir=args.output_dir, logs_dir=args.logs_dir, max_steps=args.max_steps, balance_agents=args.balance_agents, push_to_hub=args.push_to_hub, dry_run=args.dry_run ) # Create deployment manager deployment = MultiAgentTrainingDeployment(deployment_config) # Execute deployment success = deployment.deploy() if success: print("✅ Deployment completed successfully!") print(f"📁 Outputs: {deployment_config.output_dir}") print(f"📊 Logs: {deployment_config.logs_dir}") if deployment_config.push_to_hub: print(f"🚀 Model: https://huggingface.co/{deployment_config.model_repo_id}") else: print("❌ Deployment failed!") sys.exit(1) except Exception as e: print(f"❌ Deployment error: {e}") sys.exit(1) if __name__ == "__main__": main()