Feat(aws): Start batch monitor after jobs already submitted

Author: klieret

aws/run_job.py

@@ -61,8 +61,8 @@ def get_latest_job_definition_arn(self) -> str:
         logger.debug(f"Latest job definition:\n{json.dumps(latest_job_def, indent=2, default=str)}")
         return latest_job_def["jobDefinitionArn"]
 
-    def submit_job(self, command: list[str], job_name: str | None = None) -> str:
-        """Submit a job to AWS Batch. Returns jobs ID."""
+    def submit_job(self, command: list[str], job_name: str | None = None) -> tuple[str, str]:
+        """Submit a job to AWS Batch. Returns (job_id, job_name)."""
         if job_name is None:
             human_name = command[0]
             # If have config take that instead
@@ -97,7 +97,7 @@ def submit_job(self, command: list[str], job_name: str | None = None) -> str:
         logger.info(f"Full Command (passed to container entrypoint/bootstrap.sh): {shlex.join(full_command)}")
         logger.info(f"To retrieve logs later, run: python get_job_log.py {job_id}")
 
-        return job_id
+        return job_id, job_name
 
     def get_job_status(self, job_id: str) -> dict[str, Any]:
         """Get the current status of a job."""
@@ -186,7 +186,7 @@ def main():
     )
 
     # Submit the job
-    job_id = launcher.submit_job(command, args.job_name)
+    job_id, job_name = launcher.submit_job(command, args.job_name)
 
     # Wait for completion if requested
     if args.wait or args.show_logs:

aws/utils/job_monitor.py

@@ -55,6 +55,7 @@ def create_status_table(self, jobs_info: dict[str, dict]) -> Table:
         table.add_column("Status")
         table.add_column("Runtime", justify="right", style="magenta")
         table.add_column("Job ID", style="yellow")
+        table.add_column("Job Name", style="blue")
         table.add_column("Config", style="green")
 
         for job_id, info in jobs_info.items():
@@ -65,43 +66,61 @@ def create_status_table(self, jobs_info: dict[str, dict]) -> Table:
             table.add_row(
                 status_colored,
                 info["runtime"],
-                job_id,
+                job_id[:13] + "...",
+                info["job_name"],
                 info["config"],
             )
 
         return table
 
-    def monitor(self, job_id_to_config: dict[str, str]) -> None:
-        """Monitor submitted jobs and display status in real-time"""
+    def monitor(self, job_info_dict: dict[str, dict[str, str]]) -> None:
+        """Monitor submitted jobs and display status in real-time.
+
+        Args:
+            job_info_dict: Dict mapping job_id -> {job_name, config}
+        """
         jobs_info: dict[str, dict] = {}
-        start_times: dict[str, float] = {job_id: time.time() for job_id in job_id_to_config}
+        start_times: dict[str, float] = {job_id: time.time() for job_id in job_info_dict}
+        completion_times: dict[str, float] = {}
 
         self.console.print("\n[bold green]Note:[/] You can ^C this script without your jobs dying\n")
 
         try:
             with Live(self.create_status_table({}), refresh_per_second=1, console=self.console) as live:
                 while True:
                     all_done = True
-                    for job_id, config in job_id_to_config.items():
+                    for job_id, info in job_info_dict.items():
                         try:
                             job_status = self.get_job_status(job_id)
                             status = job_status["status"]
-                            elapsed = time.time() - start_times[job_id]
+
+                            # Freeze runtime when job reaches terminal state
+                            if status in ["SUCCEEDED", "FAILED"] and job_id not in completion_times:
+                                completion_times[job_id] = time.time()
+
+                            if job_id in completion_times:
+                                elapsed = completion_times[job_id] - start_times[job_id]
+                            else:
+                                elapsed = time.time() - start_times[job_id]
 
                             jobs_info[job_id] = {
                                 "status": status,
                                 "runtime": self.format_duration(elapsed),
-                                "config": config,
+                                "job_name": info["job_name"],
+                                "config": info["config"],
                             }
 
                             if status not in ["SUCCEEDED", "FAILED"]:
                                 all_done = False
                         except Exception as e:
                             logger.error(f"Error getting status for {job_id}: {e}", exc_info=True)
+                            if job_id not in completion_times:
+                                completion_times[job_id] = time.time()
                             jobs_info[job_id] = {
                                 "status": "ERROR",
-                                "runtime": self.format_duration(time.time() - start_times[job_id]),
-                                "config": config,
+                                "runtime": self.format_duration(completion_times[job_id] - start_times[job_id]),
+                                "job_name": info["job_name"],
+                                "config": info["config"],
                             }
 
                     live.update(self.create_status_table(jobs_info))

batch_monitor.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+
+import argparse
+import json
+from pathlib import Path
+
+from aws.utils.job_monitor import JobMonitor
+from codeclash.utils.log import get_logger
+
+logger = get_logger("batch_monitor", emoji="👁️")
+
+
+def load_job_files(job_files: list[Path]) -> dict[str, dict[str, str]]:
+    """Load job information from multiple JSON files.
+
+    Args:
+        job_files: List of paths to batch_submit_{timestamp}.json files
+
+    Returns:
+        Combined dict mapping job_id -> {job_name, config}
+    """
+    all_jobs: dict[str, dict[str, str]] = {}
+
+    for job_file in job_files:
+        logger.info(f"Loading jobs from {job_file}")
+        jobs = json.loads(job_file.read_text())
+
+        # Handle both old format (job_id -> config) and new format (job_id -> {job_name, config})
+        for job_id, value in jobs.items():
+            if isinstance(value, dict):
+                all_jobs[job_id] = value
+            else:
+                # Old format, convert to new format
+                all_jobs[job_id] = {"job_name": "N/A", "config": value}
+
+    logger.info(f"Loaded {len(all_jobs)} jobs total")
+    return all_jobs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Monitor AWS Batch jobs from multiple batch_submit JSON files")
+    parser.add_argument(
+        "job_files",
+        type=Path,
+        nargs="+",
+        help="One or more batch_submit_{timestamp}.json files to monitor",
+    )
+    parser.add_argument(
+        "--region",
+        default="us-east-1",
+        help="AWS region (default: us-east-1)",
+    )
+    args = parser.parse_args()
+
+    all_jobs = load_job_files(args.job_files)
+
+    monitor = JobMonitor(region=args.region)
+    monitor.monitor(all_jobs)
+
+
+if __name__ == "__main__":
+    main()

batch_submit.py

@@ -29,32 +29,32 @@ def __init__(
         )
         self.monitor = JobMonitor(region=region)
 
-    def submit_configs(self, configs: list[str]) -> dict[str, str]:
-        """Submit jobs for a list of config files. Returns job_id -> config mapping."""
+    def submit_configs(self, configs: list[str]) -> dict[str, dict[str, str]]:
+        """Submit jobs for a list of config files. Returns job_id -> {job_name, config} mapping."""
         logger.info(f"Launching {len(configs)} configs")
-        job_id_to_config: dict[str, str] = {}
+        job_info: dict[str, dict[str, str]] = {}
 
         for config in configs:
             config_path = self.config_dir / config
             command = ["aws/docker_and_sync.sh", "python", "main.py", str(config_path)]
-            job_id = self.launcher.submit_job(command)
-            job_id_to_config[job_id] = config
+            job_id, job_name = self.launcher.submit_job(command)
+            job_info[job_id] = {"job_name": job_name, "config": config}
 
-        return job_id_to_config
+        return job_info
 
-    def save_job_ids(self, job_id_to_config: dict[str, str]) -> Path:
-        """Save job IDs to a timestamped JSON file."""
+    def save_job_ids(self, job_info: dict[str, dict[str, str]]) -> Path:
+        """Save job IDs and names to a timestamped JSON file."""
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         output_file = Path(f"batch_submit_{timestamp}.json")
-        output_file.write_text(json.dumps(job_id_to_config, indent=2))
+        output_file.write_text(json.dumps(job_info, indent=2))
         logger.info(f"Job IDs saved to {output_file}")
         return output_file
 
     def run(self, configs: list[str]) -> None:
         """Submit jobs, save IDs, and monitor."""
-        job_id_to_config = self.submit_configs(configs)
-        self.save_job_ids(job_id_to_config)
-        self.monitor.monitor(job_id_to_config)
+        job_info = self.submit_configs(configs)
+        self.save_job_ids(job_info)
+        self.monitor.monitor(job_info)
 
 
 def main() -> None: