#!/usr/bin/env python3 """ Chain submission script for DPO training jobs. Automatically extracts checkpoint paths from completed jobs and updates subsequent scripts. Supports both YAML configuration files and command-line arguments. """ import os import re import sys import time import json import yaml import shutil import logging import argparse import subprocess from datetime import datetime, timedelta from pathlib import Path from typing import List, Dict, Optional, Tuple, Union class JobConfig: """Configuration for a single job in the chain.""" def __init__(self, name: str, script: str, nodes: int, **kwargs): self.name = name self.script = script self.nodes = nodes self.description = kwargs.get("description", "") self.time = kwargs.get("time") self.partition = kwargs.get("partition") # New optional parameters for overriding script values self.model_cache_loss_name = kwargs.get("model_cache_loss_name") self.wandb_run_name = kwargs.get("wandb_run_name") self.sft_loss_scale = kwargs.get("sft_loss_scale") # Add this line self.extra_params = kwargs class ChainConfig: """Configuration for the entire chain.""" def __init__(self, config_dict: Dict): self.chain_name = config_dict.get("chain_name", "DPO Training Chain") self.description = config_dict.get("description", "") # Settings settings = config_dict.get("settings", {}) # Smart default: use script directory for base_dir default_base_dir = Path(__file__).parent.absolute() self.base_dir = Path(settings.get("base_dir", default_base_dir)) # Smart default: find neon* root and use for log directory default_log_dir = "/app/suno/slurm/logs" neon_root = self._find_neon_root() if neon_root and not settings.get("log_dir"): # If we found a neon root, suggest using it for logs if path doesn't exist default_log_dir = f"/app/suno/slurm/logs" # Keep default for now self.log_dir = Path(settings.get("log_dir", default_log_dir)) self.state_file = settings.get("state_file", ".chain_state.json") # SLURM defaults self.slurm_defaults = settings.get("slurm_defaults", {}) # Jobs self.jobs = [] for job_dict in config_dict.get("jobs", []): # Merge with defaults job_config = {**self.slurm_defaults, **job_dict} self.jobs.append(JobConfig(**job_config)) # Patterns self.checkpoint_patterns = config_dict.get( "checkpoint_patterns", [ r"logging checkpoint here:\s*(/[^\s]+)", r"saving checkpoint to\s*(/[^\s]+)", r"(/app2?/suno/checkpoints/\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})", ], ) self.validation_patterns = config_dict.get( "validation_patterns", [ r"--preload_checkpoint=", r"--out_dir=", r"train_dpo\.py", r"--wandb_run_name=", ], ) @staticmethod def _find_neon_root() -> Optional[Path]: """Find the neon* parent directory (e.g., neon_sweep, neon_2, neon).""" current = Path(__file__).parent.absolute() # Walk up the directory tree for parent in [current] + list(current.parents): if parent.name.startswith("neon"): return parent return None class DPOChainSubmitter: def __init__( self, base_dir: Union[str, Path] = None, log_dir: Union[str, Path] = "/app/suno/slurm/logs", config: Optional[ChainConfig] = None, sft_loss_scale: Optional[float] = None, # Add this parameter ): self.config = config self.sft_loss_scale = sft_loss_scale # Store the sft_loss_scale # Use config values if provided, otherwise defaults if self.config: self.base_dir = self.config.base_dir self.log_dir = self.config.log_dir self.state_file = self.base_dir / self.config.state_file else: # Default to script directory if base_dir not provided self.base_dir = Path(base_dir) if base_dir else Path(__file__).parent.absolute() self.log_dir = Path(log_dir) self.state_file = self.base_dir / ".chain_state.json" # Setup logging self.setup_logging() def setup_logging(self): """Setup logging configuration.""" log_file = self.base_dir / f"chain_submit_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", handlers=[logging.FileHandler(log_file), logging.StreamHandler(sys.stdout)], ) self.logger = logging.getLogger(__name__) def save_state(self, state: Dict): """Save current chain state for potential resume.""" with open(self.state_file, "w") as f: json.dump(state, f, indent=2) def load_state(self) -> Optional[Dict]: """Load saved chain state if exists.""" if self.state_file.exists(): with open(self.state_file, "r") as f: return json.load(f) return None def validate_script(self, script_path: Path, patterns: Optional[List[str]] = None) -> bool: """Validate that a script has the expected structure.""" try: with open(script_path, "r") as f: content = f.read() # Use provided patterns or defaults required_patterns = patterns or ( self.config.validation_patterns if self.config else [ r"--preload_checkpoint=", r"--out_dir=", r"train_dpo\.py", r"--wandb_run_name=", ] ) for pattern in required_patterns: if not re.search(pattern, content): self.logger.warning(f"Script {script_path.name} missing pattern: {pattern}") return False return True except Exception as e: self.logger.error(f"Error validating script {script_path}: {e}") return False def submit_job( self, script_path: Path, job_config: Optional[JobConfig] = None, node_count: Optional[int] = None, job_name: Optional[str] = None, ) -> Optional[str]: """Submit a SLURM job and return the job ID.""" script_name = script_path.name # Use job config if provided if job_config: nodes = job_config.nodes job_name = job_name or "dpo" # Just "dpo" for all jobs time_limit = job_config.time or "48:00:00" else: nodes = node_count or 8 job_name = job_name or "dpo" # Just "dpo" for all jobs time_limit = "48:00:00" # Create a temporary sbatch script sbatch_content = f"""#!/bin/bash #SBATCH --job-name="{job_name}" #SBATCH --nodes={nodes} #SBATCH --ntasks-per-node=8 #SBATCH --cpus-per-task=4 #SBATCH --gres=gpu:8 #SBATCH --output={self.log_dir}/run_%x_%j.txt #SBATCH --error={self.log_dir}/run_%x_%j_err.txt #SBATCH --time={time_limit} srun -K1 {script_path.absolute()} """ temp_sbatch = f"/tmp/sbatch_{script_name}_{int(time.time())}" try: with open(temp_sbatch, "w") as f: f.write(sbatch_content) # Submit the job result = subprocess.run(["sbatch", temp_sbatch], capture_output=True, text=True, check=True) # Extract job ID from output match = re.search(r"Submitted batch job (\d+)", result.stdout) if match: job_id = match.group(1) self.logger.info(f"Submitted {script_name} as job {job_id} with {nodes} nodes") return job_id else: raise ValueError(f"Could not extract job ID from: {result.stdout}") except subprocess.CalledProcessError as e: self.logger.error(f"Failed to submit job: {e.stderr}") return None except Exception as e: self.logger.error(f"Unexpected error submitting job: {e}") return None finally: # Clean up temp file if os.path.exists(temp_sbatch): os.remove(temp_sbatch) def wait_for_job(self, job_id: str, check_interval: int = 60) -> Tuple[bool, Optional[timedelta]]: """Wait for a SLURM job to complete. Returns (success, runtime).""" self.logger.info(f"Waiting for job {job_id} to complete...") start_time = time.time() last_state = None while True: try: # Check job status result = subprocess.run( ["squeue", "-j", job_id, "-h", "-o", "%T,%M"], capture_output=True, text=True, ) if result.returncode != 0 or not result.stdout.strip(): # Job is no longer in queue - check completion status sacct_result = subprocess.run( ["sacct", "-j", job_id, "-n", "-o", "State,Elapsed", "-X"], capture_output=True, text=True, ) if sacct_result.returncode == 0 and sacct_result.stdout.strip(): parts = sacct_result.stdout.strip().split() if len(parts) >= 2: state, elapsed = parts[0], parts[1] runtime = self._parse_elapsed_time(elapsed) if "COMPLETED" in state: self.logger.info(f"Job {job_id} completed successfully in {elapsed}") return True, runtime else: self.logger.error(f"Job {job_id} ended with state: {state}") return False, runtime self.logger.warning(f"Could not determine final state of job {job_id}") return False, None else: # Job is still running parts = result.stdout.strip().split(",") if len(parts) >= 2: state, runtime = parts[0], parts[1] if state != last_state: self.logger.info(f"Job {job_id} state: {state}, runtime: {runtime}") last_state = state except Exception as e: self.logger.error(f"Error checking job status: {e}") return False, None time.sleep(check_interval) def _parse_elapsed_time(self, elapsed: str) -> Optional[timedelta]: """Parse SLURM elapsed time format (DD-HH:MM:SS or HH:MM:SS).""" try: if "-" in elapsed: days, time_part = elapsed.split("-") hours, minutes, seconds = map(int, time_part.split(":")) return timedelta(days=int(days), hours=hours, minutes=minutes, seconds=seconds) else: parts = elapsed.split(":") if len(parts) == 3: hours, minutes, seconds = map(int, parts) return timedelta(hours=hours, minutes=minutes, seconds=seconds) except: return None def extract_checkpoint_path(self, job_id: str, max_wait: int = 30) -> Optional[str]: """Extract the checkpoint path from job output log with retry logic.""" log_file = self.log_dir / f"run_dpo_{job_id}.txt" self.logger.info(f"Looking for checkpoint in log file: {log_file}") # Wait for log file to appear and be written start_time = time.time() while time.time() - start_time < max_wait: if log_file.exists() and log_file.stat().st_size > 0: break time.sleep(5) if not log_file.exists(): self.logger.error(f"Log file not found after {max_wait}s: {log_file}") return None try: with open(log_file, "r") as f: content = f.read() # Use config patterns if available patterns = ( self.config.checkpoint_patterns if self.config else [ r"logging checkpoint here:\s*(/[^\s]+)", r"saving checkpoint to\s*(/[^\s]+)", r"(/app2?/suno/checkpoints/\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})", ] ) for pattern in patterns: match = re.search(pattern, content) if match: checkpoint_dir = match.group(1) checkpoint_path = os.path.join(checkpoint_dir, "last_ckpt_infer.pt") # Verify the checkpoint exists if Path(checkpoint_path).exists(): self.logger.info(f"Found and verified checkpoint: {checkpoint_path}") return checkpoint_path else: self.logger.warning( f"Checkpoint path found but file doesn't exist: {checkpoint_path}" ) self.logger.error("Could not find valid checkpoint path in log") return None except Exception as e: self.logger.error(f"Error reading log file: {e}") return None def update_script_parameters( self, script_path: Path, checkpoint_path: Optional[str] = None, model_cache_loss_name: Optional[str] = None, wandb_run_name: Optional[str] = None, sft_loss_scale: Optional[float] = None, # Add this parameter ) -> bool: """Update multiple parameters in a script.""" # Create backup with timestamp timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") backup_path = script_path.with_suffix(f".sh.bak.{timestamp}") shutil.copy2(script_path, backup_path) self.logger.info(f"Created backup: {backup_path}") try: with open(script_path, "r") as f: content = f.read() # Track if any changes were made original_content = content # Update checkpoint path if provided if checkpoint_path: content = re.sub( r'--preload_checkpoint="[^"]*"', f'--preload_checkpoint="{checkpoint_path}"', content, ) # Update model_cache_loss_name if provided if model_cache_loss_name: content = re.sub( r'--model_cache_loss_name="[^"]*"', f'--model_cache_loss_name="{model_cache_loss_name}"', content, ) # Update wandb_run_name if provided if wandb_run_name: content = re.sub( r'--wandb_run_name="[^"]*"', f'--wandb_run_name="{wandb_run_name}"', content, ) # Update sft_loss_scale if provided if sft_loss_scale is not None: # First try to match existing --sft_loss_scale parameter pattern = r"--sft_loss_scale=[\d.]+" if re.search(pattern, content): content = re.sub(pattern, f"--sft_loss_scale={sft_loss_scale}", content) else: # If not found, add it after train_dpo.py (or other appropriate location) # Look for the train command line train_pattern = r"(train_dpo\.py[^\n]*)" match = re.search(train_pattern, content) if match: # Add sft_loss_scale at the end of the train command content = re.sub( train_pattern, rf"\1 --sft_loss_scale={sft_loss_scale}", content, ) if content == original_content: self.logger.warning(f"No parameters were updated in {script_path.name}") return False # Write updated script with open(script_path, "w") as f: f.write(content) # Log what was updated updates = [] if checkpoint_path: updates.append(f"checkpoint: {checkpoint_path}") if model_cache_loss_name: updates.append(f"model_cache_loss_name: {model_cache_loss_name}") if wandb_run_name: updates.append(f"wandb_run_name: {wandb_run_name}") if sft_loss_scale is not None: updates.append(f"sft_loss_scale: {sft_loss_scale}") self.logger.info(f"Updated {script_path.name} - {', '.join(updates)}") return True except Exception as e: self.logger.error(f"Error updating script: {e}") # Restore backup on error shutil.copy2(backup_path, script_path) return False def update_script_checkpoint(self, script_path: Path, new_checkpoint_path: str) -> bool: """Update the preload_checkpoint path in a script (legacy method).""" return self.update_script_parameters(script_path, checkpoint_path=new_checkpoint_path) def restore_scripts(self, script_paths: List[Path]): """Restore scripts from their most recent backups.""" for script_path in script_paths: backup_pattern = f"{script_path.name}.bak.*" backups = sorted(script_path.parent.glob(backup_pattern)) if backups: latest_backup = backups[-1] shutil.copy2(latest_backup, script_path) self.logger.info(f"Restored {script_path.name} from {latest_backup.name}") else: self.logger.warning(f"No backup found for {script_path.name}") def run_chain_from_config(self, resume: bool = False) -> List[Dict]: """Run a chain of jobs from configuration.""" if not self.config: raise ValueError("No configuration loaded") self.logger.info(f"Starting chain: {self.config.chain_name}") self.logger.info(f"Description: {self.config.description}") self.logger.info(f"Total jobs: {len(self.config.jobs)}") # Check for resume state start_index = 0 completed_jobs = [] if resume and self.state_file.exists(): state = self.load_state() if state: start_index = state.get("next_index", 0) completed_jobs = state.get("completed_jobs", []) self.logger.info(f"Resuming from job {start_index + 1}") total_runtime = timedelta() for i in range(start_index, len(self.config.jobs)): job = self.config.jobs[i] script_path = self.base_dir / job.script self.logger.info(f"\nJob {i+1}/{len(self.config.jobs)}: {job.name}") if job.description: self.logger.info(f"Description: {job.description}") # Validate script if not script_path.exists(): self.logger.error(f"Script not found: {script_path}") break if not self.validate_script(script_path): self.logger.error(f"Script validation failed for {job.script}") break # Update script parameters from YAML config if provided # Only update for the first job (i == 0) # For i > 0, previous job should have already updated everything if i == 0 and ( job.model_cache_loss_name or job.wandb_run_name or job.sft_loss_scale is not None # Change from self.sft_loss_scale ): self.logger.info(f"Updating {job.script} with YAML config parameters...") if not self.update_script_parameters( script_path, model_cache_loss_name=job.model_cache_loss_name, wandb_run_name=job.wandb_run_name, sft_loss_scale=job.sft_loss_scale if job.sft_loss_scale is not None else self.sft_loss_scale, # Use job-specific or fall back to command-line override ): self.logger.error(f"Failed to update parameters in {job.script}") # but don't need to break! # For i > 0, verify that checkpoint was properly updated by previous job if i > 0: with open(script_path, "r") as f: content = f.read() # Get the checkpoint from the previous job previous_checkpoint = completed_jobs[-1]["checkpoint"] # Check if the script contains the expected checkpoint from previous job if previous_checkpoint not in content: self.logger.error( f"Checkpoint path not updated in {job.script} - expected checkpoint from previous job: {previous_checkpoint}" ) break # Submit the job job_id = self.submit_job(script_path, job_config=job) if not job_id: self.logger.error(f"Failed to submit {job.name}") break # Wait for completion success, runtime = self.wait_for_job(job_id) if not success: self.logger.error(f"Job {job_id} failed") break if runtime: total_runtime += runtime # Extract checkpoint path checkpoint_path = self.extract_checkpoint_path(job_id) if not checkpoint_path: self.logger.error(f"Could not extract checkpoint path from job {job_id}") break completed_jobs.append( { "name": job.name, "script": job.script, "job_id": job_id, "checkpoint": checkpoint_path, "runtime": str(runtime) if runtime else None, "nodes": job.nodes, } ) # Save state for potential resume self.save_state( { "next_index": i + 1, "completed_jobs": completed_jobs, "timestamp": datetime.now().isoformat(), } ) # Update next script if not the last one if i < len(self.config.jobs) - 1: next_job = self.config.jobs[i + 1] next_script = self.base_dir / next_job.script self.logger.info(f"Updating {next_job.name} with new checkpoint...") # Include YAML config parameters for next job if provided if not self.update_script_parameters( next_script, checkpoint_path=checkpoint_path, model_cache_loss_name=next_job.model_cache_loss_name, wandb_run_name=next_job.wandb_run_name, sft_loss_scale=next_job.sft_loss_scale if next_job.sft_loss_scale is not None else self.sft_loss_scale, # Use job-specific or fall back to command-line override ): self.logger.error(f"Failed to update {next_job.script}") break # Clean up state file on successful completion if len(completed_jobs) == len(self.config.jobs): if self.state_file.exists(): self.state_file.unlink() self.logger.info( f"\nChain submission completed. Ran {len(completed_jobs)} jobs in {total_runtime}" ) for job in completed_jobs: self.logger.info( f" - {job['name']} (job {job['job_id']}, {job['nodes']} nodes) - {job.get('runtime', 'N/A')}" ) return completed_jobs def submit_single_mega_job(self, partition: Optional[str] = None) -> str: """Submit all rounds as a single Slurm job.""" if not self.config: raise ValueError("No configuration loaded") if not self.config.jobs: raise ValueError("No jobs configured") # Verify all jobs use same node count nodes = self.config.jobs[0].nodes for job in self.config.jobs: if job.nodes != nodes: self.logger.warning(f"{job.name} configured for {job.nodes} nodes, using {nodes}") # Create master wrapper wrapper_path = self._create_mega_wrapper() # Submit job_id = self._submit_mega_job(wrapper_path, nodes, partition) if not job_id: raise RuntimeError("Failed to submit mega job") self.logger.info(f"āœ… Submitted mega job {job_id}") self.logger.info(f" Running {len(self.config.jobs)} rounds on {nodes} nodes") return job_id def _create_mega_wrapper(self) -> Path: """Generate master script that runs all rounds sequentially.""" timestamp = int(time.time()) wrapper_path = self.base_dir / f".mega_wrapper_{timestamp}.sh" log_dir = self.base_dir / "round_logs" lines = [ "#!/bin/bash", "set -e # Exit immediately on any error", "set -o pipefail # Catch errors in pipes", "", "echo '=========================================='", f"echo 'DPO Chain: {self.config.chain_name}'", "echo 'Job ID: '$SLURM_JOB_ID", "echo 'Nodes: '$SLURM_JOB_NUM_NODES", f"echo 'Rounds: {len(self.config.jobs)}'", "echo 'Start time: '$(date)", "echo '=========================================='", "", f"# Create log directory", f"mkdir -p {log_dir}", "", "START_TIME=$(date +%s)", "", ] for i, job in enumerate(self.config.jobs): script_path = self.base_dir / job.script round_log = log_dir / f"round{i+1}_{job.name}.log" checkpoint_file = f"/tmp/checkpoint_round{i+1}_${{SLURM_JOB_ID}}.txt" lines.extend( [ f"# ========== Round {i+1}/{len(self.config.jobs)}: {job.name} ==========", f"echo ''", f"echo 'Round {i+1}: {job.name}'", f"echo 'Time: '$(date)", f"ROUND_START=$(date +%s)", "", ] ) # Update script parameters from YAML config if i == 0: # First round: update from YAML config if job.model_cache_loss_name: lines.append( f'sed -i \'s|--model_cache_loss_name="[^"]*"|--model_cache_loss_name="{job.model_cache_loss_name}"|g\' "{script_path}"' ) else: # Subsequent rounds: get checkpoint from previous round prev_checkpoint_file = f"/tmp/checkpoint_round{i}_${{SLURM_JOB_ID}}.txt" lines.extend( [ "# Extract checkpoint from previous round", f'if [ ! -f "{prev_checkpoint_file}" ]; then', f' echo "ERROR: Previous checkpoint marker not found"', " exit 1", "fi", "", f'PREV_CKPT=$(cat "{prev_checkpoint_file}")', 'echo "Previous checkpoint: $PREV_CKPT"', "", 'if [ ! -f "$PREV_CKPT" ]; then', ' echo "ERROR: Checkpoint file does not exist: $PREV_CKPT"', " exit 1", "fi", "", "# Update script with checkpoint", f'sed -i "s|--preload_checkpoint=\\"[^\\"]*\\"|--preload_checkpoint=\\"$PREV_CKPT\\"|g" "{script_path}"', "", ] ) if job.model_cache_loss_name: lines.append( f'sed -i \'s|--model_cache_loss_name="[^"]*"|--model_cache_loss_name="{job.model_cache_loss_name}"|g\' "{script_path}"' ) # Update wandb_run_name if provided if job.wandb_run_name: lines.append( f'sed -i \'s|--wandb_run_name="[^"]*"|--wandb_run_name="{job.wandb_run_name}"|g\' "{script_path}"' ) # Update sft_loss_scale if provided if job.sft_loss_scale is not None: lines.append( f"sed -i 's|--sft_loss_scale=[0-9.]*|--sft_loss_scale={job.sft_loss_scale}|g' \"{script_path}\"" ) lines.extend( [ "", "# Run training and capture output", f'echo "Executing {script_path.name}..."', f'srun -K1 {script_path} 2>&1 | tee "{round_log}"', "", "# Extract checkpoint path from output", f'CKPT_DIR=$(grep -oP "logging checkpoint here: \\K.*" "{round_log}" | tail -1)', "", 'if [ -z "$CKPT_DIR" ]; then', f' echo "ERROR: Could not find checkpoint in {round_log}"', " exit 1", "fi", "", 'CKPT_PATH="${CKPT_DIR}/last_ckpt_infer.pt"', "", 'if [ ! -f "$CKPT_PATH" ]; then', ' echo "ERROR: Checkpoint not found: $CKPT_PATH"', " exit 1", "fi", "", f"# Save checkpoint for next round", f'echo "$CKPT_PATH" > "{checkpoint_file}"', f'echo "Saved checkpoint: $CKPT_PATH"', "", "ROUND_END=$(date +%s)", "ROUND_ELAPSED=$((ROUND_END - ROUND_START))", "ROUND_MINS=$((ROUND_ELAPSED / 60))", f'echo "āœ“ Round {i+1} completed in ${{ROUND_MINS}}m"', "", ] ) # Footer lines.extend( [ "END_TIME=$(date +%s)", "TOTAL_ELAPSED=$((END_TIME - START_TIME))", "TOTAL_HOURS=$((TOTAL_ELAPSED / 3600))", "TOTAL_MINS=$(((TOTAL_ELAPSED % 3600) / 60))", "", "echo '=========================================='", f"echo 'SUCCESS: All {len(self.config.jobs)} rounds completed!'", 'echo "Total time: ${TOTAL_HOURS}h ${TOTAL_MINS}m"', "echo 'End time: '$(date)", "echo '=========================================='", "", "# Cleanup temp files", "rm -f /tmp/checkpoint_round*_${SLURM_JOB_ID}.txt", ] ) with open(wrapper_path, "w") as f: f.write("\n".join(lines)) wrapper_path.chmod(0o755) self.logger.info(f"Created mega wrapper: {wrapper_path}") return wrapper_path def _submit_mega_job( self, wrapper_path: Path, nodes: int, partition: Optional[str] ) -> Optional[str]: """Submit mega wrapper as Slurm job.""" # Estimate time: 3h per round + 20% buffer, cap at 48h est_hours = min(len(self.config.jobs) * 3 * 1.2, 48) time_limit = f"{int(est_hours):02d}:00:00" sbatch_script = f"""#!/bin/bash #SBATCH --job-name=dpo_chain #SBATCH --nodes={nodes} #SBATCH --ntasks-per-node=8 #SBATCH --cpus-per-task=4 #SBATCH --gres=gpu:8 #SBATCH --time={time_limit} #SBATCH --output={self.log_dir}/dpo_chain_%j.txt #SBATCH --error={self.log_dir}/dpo_chain_%j_err.txt #SBATCH --exclusive """ if partition: sbatch_script += f"#SBATCH --partition={partition}\n" sbatch_script += f"\nbash {wrapper_path}\n" temp_file = f"/tmp/sbatch_mega_{int(time.time())}.sh" try: with open(temp_file, "w") as f: f.write(sbatch_script) result = subprocess.run( ["sbatch", "--parsable", temp_file], capture_output=True, text=True, check=True ) return result.stdout.strip() except subprocess.CalledProcessError as e: self.logger.error(f"Submission failed: {e.stderr}") return None finally: if os.path.exists(temp_file): os.remove(temp_file) def run_chain( self, script_sequence: List[Path], node_counts: Optional[List[int]] = None, resume: bool = False, ) -> List[Dict]: """Run a sequence of scripts, updating checkpoint paths between runs. Legacy method for command-line compatibility. """ self.logger.info(f"Starting chain submission of {len(script_sequence)} jobs") self.logger.info(f"Scripts: {' -> '.join([s.name for s in script_sequence])}") # Check for resume state start_index = 0 completed_jobs = [] if resume and self.state_file.exists(): state = self.load_state() if state: start_index = state.get("next_index", 0) completed_jobs = state.get("completed_jobs", []) self.logger.info(f"Resuming from step {start_index + 1}") # Validate all scripts first for script in script_sequence: if not self.validate_script(script): self.logger.error(f"Script validation failed for {script.name}") return completed_jobs total_runtime = timedelta() for i in range(start_index, len(script_sequence)): script_path = script_sequence[i] script_name = script_path.name self.logger.info(f"\nStep {i+1}/{len(script_sequence)}: {script_name}") # Get node count for this script node_count = node_counts[i] if node_counts and i < len(node_counts) else 8 # Submit the job job_id = self.submit_job(script_path, node_count=node_count) if not job_id: self.logger.error(f"Failed to submit {script_name}") break # Wait for completion success, runtime = self.wait_for_job(job_id) if not success: self.logger.error(f"Job {job_id} failed") break if runtime: total_runtime += runtime # Extract checkpoint path checkpoint_path = self.extract_checkpoint_path(job_id) if not checkpoint_path: self.logger.error(f"Could not extract checkpoint path from job {job_id}") break completed_jobs.append( { "script": script_name, "job_id": job_id, "checkpoint": checkpoint_path, "runtime": str(runtime) if runtime else None, } ) # Save state for potential resume self.save_state( { "next_index": i + 1, "completed_jobs": completed_jobs, "timestamp": datetime.now().isoformat(), } ) # Update next script if not the last one if i < len(script_sequence) - 1: next_script = script_sequence[i + 1] self.logger.info(f"Updating {next_script.name} with new checkpoint...") if not self.update_script_parameters( next_script, checkpoint_path=checkpoint_path, sft_loss_scale=self.sft_loss_scale, # Add this line ): self.logger.error(f"Failed to update {next_script.name}") break # Clean up state file on successful completion if len(completed_jobs) == len(script_sequence): if self.state_file.exists(): self.state_file.unlink() self.logger.info( f"\nChain submission completed. Ran {len(completed_jobs)} jobs in {total_runtime}" ) for job in completed_jobs: self.logger.info(f" - {job['script']} (job {job['job_id']}) - {job.get('runtime', 'N/A')}") return completed_jobs def load_config(config_path: Path) -> ChainConfig: """Load configuration from YAML file.""" with open(config_path, "r") as f: config_dict = yaml.safe_load(f) return ChainConfig(config_dict) def main(): parser = argparse.ArgumentParser( description="Chain submit DPO training jobs", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Run with config file %(prog)s --config chain_config.yaml # Run with config file and resume %(prog)s --config chain_config.yaml --resume # Run with config and override sft_loss_scale %(prog)s --config chain_config.yaml --sft_loss_scale 0.5 # Dry run with config %(prog)s --config chain_config.yaml --dry-run # Legacy: Run a simple chain %(prog)s r3 r4 r5 # Legacy: Run with custom prefix %(prog)s r1 r2 r3 --prefix run_ipo_dodo_ """, ) # Config-based arguments parser.add_argument("--config", type=Path, help="Path to YAML configuration file") # Parameter override arguments parser.add_argument("--sft_loss_scale", type=float, help="Override SFT loss scale value in scripts") # Legacy arguments parser.add_argument("scripts", nargs="*", help="Script names to run in sequence (legacy mode)") parser.add_argument( "--base-dir", default=None, help="Base directory for scripts (legacy mode, defaults to script directory)", ) parser.add_argument( "--prefix", default="run_ipo_dodo_", help="Script filename prefix (legacy mode)", ) # Common arguments parser.add_argument( "--dry-run", action="store_true", help="Show what would be done without executing", ) parser.add_argument("--resume", action="store_true", help="Resume from last successful job") parser.add_argument("--restore", action="store_true", help="Restore scripts from backups and exit") parser.add_argument( "--mega-job", action="store_true", help="Submit all rounds as single job (guarantees node reservation)", ) parser.add_argument("--partition", type=str, help="Partition name for job submission") parser.add_argument( "--log-level", choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="INFO", help="Set logging level", ) args = parser.parse_args() # Set logging level logging.getLogger().setLevel(getattr(logging, args.log_level)) # Process config-based run if args.config: # Load configuration try: config = load_config(args.config) except Exception as e: print(f"āŒ Error loading configuration: {e}") sys.exit(1) # Initialize submitter with config submitter = DPOChainSubmitter( config=config, sft_loss_scale=args.sft_loss_scale, # Pass the override value ) # Handle dry run if args.dry_run: print(f"šŸ” DRY RUN - {config.chain_name}") print(f"šŸ“ {config.description}") print(f"\nWould execute {len(config.jobs)} jobs:") for i, job in enumerate(config.jobs): print(f" {i+1}. {job.name} - {job.script} ({job.nodes} nodes)") if job.description: print(f" {job.description}") # Show overrides if present if job.model_cache_loss_name: print(f" model_cache_loss_name: {job.model_cache_loss_name}") if job.wandb_run_name: print(f" wandb_run_name: {job.wandb_run_name}") sys.exit(0) # Handle restore if args.restore: script_paths = [config.base_dir / job.script for job in config.jobs] submitter.logger.info("Restoring scripts from backups...") submitter.restore_scripts(script_paths) sys.exit(0) # Run the chain if args.mega_job: job_id = submitter.submit_single_mega_job(partition=args.partition) print(f"\nāœ… Mega job submitted: {job_id}") print(f"Log: tail -f {config.log_dir}/dpo_chain_{job_id}.txt") print(f"Status: squeue -j {job_id}") else: submitter.run_chain_from_config(resume=args.resume) else: # Legacy mode if not args.scripts: parser.error("Either --config or script names must be provided") # Convert script names to full paths script_sequence = [] for script in args.scripts: script_path = Path(args.base_dir) / f"{args.prefix}{script}.sh" if not script_path.exists(): print(f"āŒ Script not found: {script_path}") sys.exit(1) script_sequence.append(script_path) # Initialize submitter with sft_loss_scale override submitter = DPOChainSubmitter( base_dir=args.base_dir, sft_loss_scale=args.sft_loss_scale, # Pass the override value ) # Handle restore operation if args.restore: submitter.logger.info("Restoring scripts from backups...") submitter.restore_scripts(script_sequence) sys.exit(0) # Handle dry run if args.dry_run: print("šŸ” DRY RUN - Would execute:") for i, script in enumerate(script_sequence): print(f" {i+1}. {script.name} (8 nodes)") # Default nodes in legacy mode sys.exit(0) # Run the chain submitter.run_chain(script_sequence, resume=args.resume) if __name__ == "__main__": main()