feat&fix: enhance error handling during environment initialization and VM allocation

run_multienv_openaicua.py

@@ -2,12 +2,15 @@
 Utils and basic architecture credit to https://github.com/web-arena-x/webarena/blob/main/run.py.
 """
 
+from __future__ import annotations
 import argparse
 import datetime
 import json
 import logging
 import os
 import sys
+import signal
+import time
 from typing import List, Dict
 import math
 from tqdm import tqdm
@@ -16,6 +19,11 @@ import lib_run_single
 from desktop_env.desktop_env import DesktopEnv
 from mm_agents.openai_cua_agent import OpenAICUAAgent
 
+# Global variables for signal handling
+active_environments = []
+processes = []
+is_terminating = False
+
 # import wandb
 
 # load the environment variables from .env file
@@ -147,8 +155,39 @@ def distribute_tasks(test_all_meta: dict, num_envs: int) -> List[Dict]:
     return distributed_tasks
 
 
+def process_signal_handler(signum, frame, env_idx):
+    """Signal handler for child processes to gracefully shut down their environments."""
+    logger.info(f"Process {env_idx + 1} received signal {signum}. Shutting down...")
+    
+    # Get the active_environments from the caller's frame
+    local_vars = frame.f_locals
+    print(f"Local variables in process {env_idx + 1}: {local_vars}")
+    active_environments = local_vars.get('active_environments', [])
+    
+    # Close environment in the current process context
+    for env in active_environments:
+        if env is not None:
+            try:
+                logger.info(f"Process {env_idx + 1} closing environment...")
+                env.close()
+                logger.info(f"Process {env_idx + 1} environment closed successfully")
+            except Exception as e:
+                logger.error(f"Process {env_idx + 1} error closing environment: {e}")
+    
+    logger.info(f"Process {env_idx + 1} shutdown complete. Exiting.")
+    sys.exit(0)
+
+
 def run_env_tasks(env_idx: int, env_tasks: dict, args: argparse.Namespace, shared_scores: list):
     """Run tasks for a single environment."""
+    # Each process has its own list of active environments
+    active_environments = []
+    env = None
+    
+    # Setup signal handlers for this process too
+    signal.signal(signal.SIGINT, lambda signum, frame: process_signal_handler(signum, frame, env_idx))
+    signal.signal(signal.SIGTERM, lambda signum, frame: process_signal_handler(signum, frame, env_idx))
+    
     # ami-05e7d7bd279ea4f14
     env = DesktopEnv(
         path_to_vm=args.path_to_vm,
@@ -163,6 +202,7 @@ def run_env_tasks(env_idx: int, env_tasks: dict, args: argparse.Namespace, share
         os_type="Ubuntu",
         require_a11y_tree=args.observation_type in ["a11y_tree", "screenshot_a11y_tree", "som"],
     )
+    active_environments.append(env)
     agent = OpenAICUAAgent(
         env=env,
         model=args.model,
@@ -175,56 +215,114 @@ def run_env_tasks(env_idx: int, env_tasks: dict, args: argparse.Namespace, share
     )
     logger.info(f"Executing tasks in environment {env_idx + 1}/{args.num_envs}")
     
-    for domain in tqdm(env_tasks, desc=f"Env{env_idx+1}-Domain"):
-         for example_id in tqdm(env_tasks[domain], desc="Example", leave=False):
-            config_file = os.path.join(
-                args.test_config_base_dir, f"examples/{domain}/{example_id}.json"
-            )
-            with open(config_file, "r", encoding="utf-8") as f:
-                example = json.load(f)
-
-            logger.info(f"[Env {env_idx+1}][Domain]: {domain}")
-            logger.info(f"[Env {env_idx+1}][Example ID]: {example_id}")
-            logger.info(f"[Env {env_idx+1}][Instruction]: {example['instruction']}")
-            
-            example_result_dir = os.path.join(
-                args.result_dir,
-                args.action_space,
-                args.observation_type,
-                args.model,
-                domain,
-                example_id,
-            )
-            os.makedirs(example_result_dir, exist_ok=True)
-
-            try:
-                lib_run_single.run_single_example_openaicua(
-                    agent,
-                    env,
-                    example,
-                    args.max_steps,
-                    example["instruction"],
-                    args,
-                    example_result_dir,
-                    shared_scores,
+    try:
+        for domain in tqdm(env_tasks, desc=f"Env{env_idx+1}-Domain"):
+             for example_id in tqdm(env_tasks[domain], desc="Example", leave=False):
+                config_file = os.path.join(
+                    args.test_config_base_dir, f"examples/{domain}/{example_id}.json"
                 )
-            except Exception as e:
-                logger.error(f"Exception in Env{env_idx+1} {domain}/{example_id}: {e}")
-                env.controller.end_recording(
-                    os.path.join(example_result_dir, "recording.mp4")
+                with open(config_file, "r", encoding="utf-8") as f:
+                    example = json.load(f)
+
+                logger.info(f"[Env {env_idx+1}][Domain]: {domain}")
+                logger.info(f"[Env {env_idx+1}][Example ID]: {example_id}")
+                logger.info(f"[Env {env_idx+1}][Instruction]: {example['instruction']}")
+                
+                example_result_dir = os.path.join(
+                    args.result_dir,
+                    args.action_space,
+                    args.observation_type,
+                    args.model,
+                    domain,
+                    example_id,
                 )
-                with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
-                    f.write(
-                        json.dumps(
-                            {"Error": f"Time limit exceeded in {domain}/{example_id}"}
-                        )
+                os.makedirs(example_result_dir, exist_ok=True)
+
+                try:
+                    lib_run_single.run_single_example_openaicua(
+                        agent,
+                        env,
+                        example,
+                        args.max_steps,
+                        example["instruction"],
+                        args,
+                        example_result_dir,
+                        shared_scores,
                     )
-                    f.write("\n")
+                except Exception as e:
+                    logger.error(f"Exception in Env{env_idx+1} {domain}/{example_id}: {e}")
+                    try:
+                        env.controller.end_recording(
+                            os.path.join(example_result_dir, "recording.mp4")
+                        )
+                    except Exception as rec_e:
+                        logger.error(f"Failed to end recording: {rec_e}")
+                    
+                    with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
+                        f.write(
+                            json.dumps(
+                                {"Error": f"Time limit exceeded in {domain}/{example_id}"}
+                            )
+                        )
+                        f.write("\n")
+    finally:
+        # This ensures the environment is closed even if there's an exception
+        logger.info(f"Process {env_idx + 1} cleaning up environment...")
+        try:
+            env.close()
+            logger.info(f"Process {env_idx + 1} environment closed successfully")
+        except Exception as e:
+            logger.error(f"Process {env_idx + 1} error during environment cleanup: {e}")
+
+
+def signal_handler(signum, frame):
+    """Handle termination signals (SIGINT, SIGTERM) to gracefully shutdown environments."""
+    global is_terminating, active_environments, processes
     
-    env.close()
+    # Avoid duplicate handling
+    if is_terminating:
+        return
+    
+    is_terminating = True
+    logger.info(f"Received signal {signum}. Gracefully shutting down...")
+    
+    # Close all registered environments in the main process
+    for env in active_environments:
+        try:
+            logger.info(f"Closing environment...")
+            env.close()
+            logger.info(f"Environment closed successfully")
+        except Exception as e:
+            logger.error(f"Error closing environment: {e}")
+    
+    # Send termination signal to all child processes first
+    for p in processes:
+        if p.is_alive():
+            try:
+                logger.info(f"Sending termination signal to process {p.name}...")
+                p.terminate()
+            except Exception as e:
+                logger.error(f"Error sending termination signal to process: {e}")
+    
+    # Allow a short time for processes to handle their own cleanup
+    time.sleep(1)
+    
+    # Forcefully terminate any processes that didn't exit
+    for p in processes:
+        if p.is_alive():
+            try:
+                logger.info(f"Forcefully terminating process {p.name}...")
+                import signal
+                os.kill(p.pid, signal.SIGKILL)
+            except Exception as e:
+                logger.error(f"Error forcefully terminating process: {e}")
+    
+    logger.info("Shutdown complete. Exiting.")
+    sys.exit(0)
 
 
 def test(args: argparse.Namespace, test_all_meta: dict) -> None:
+    global processes
     logger.info("Args: %s", args)
     
     distributed_tasks = distribute_tasks(test_all_meta, args.num_envs)
@@ -244,10 +342,28 @@ def test(args: argparse.Namespace, test_all_meta: dict) -> None:
             )
             processes.append(p)
             p.start()
+            logger.info(f"Started process {p.name} with PID {p.pid}")
         
-        # Wait for all processes to complete
-        for p in processes:
-            p.join()
+        try:
+            # Wait for all processes to complete
+            for p in processes:
+                p.join()
+                logger.info(f"Process {p.name} completed")
+        except KeyboardInterrupt:
+            logger.info("Main process received KeyboardInterrupt. Initiating graceful shutdown...")
+            # Let the signal handler do the cleanup
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error while waiting for processes: {e}", exc_info=True)
+            # Ensure cleanup happens
+            for p in processes:
+                if p.is_alive():
+                    try:
+                        logger.info(f"Terminating process {p.name} due to error...")
+                        p.terminate()
+                    except Exception as term_e:
+                        logger.error(f"Error terminating process {p.name}: {term_e}")
+            raise
         
         # Convert shared list to regular list
         scores = list(shared_scores)
@@ -331,31 +447,76 @@ if __name__ == "__main__":
     ####### The complete version of the list of examples #######
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
    
-    args = config()
+    # Register signal handlers for graceful termination
+    signal.signal(signal.SIGINT, signal_handler)  # Handle Ctrl+C
+    signal.signal(signal.SIGTERM, signal_handler)  # Handle termination signal
+    
+    try:
+        args = config()
 
-    with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
-        test_all_meta = json.load(f)
+        with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
+            test_all_meta = json.load(f)
 
-    if args.domain != "all":
-        test_all_meta = {args.domain: test_all_meta[args.domain]}
+        if args.domain != "all":
+            test_all_meta = {args.domain: test_all_meta[args.domain]}
 
-    test_file_list = get_unfinished(
-        args.action_space,
-        args.model,
-        args.observation_type,
-        args.result_dir,
-        test_all_meta,
-    )
-    left_info = ""
-    for domain in test_file_list:
-        left_info += f"{domain}: {len(test_file_list[domain])}\n"
-    logger.info(f"Left tasks:\n{left_info}")
+        test_file_list = get_unfinished(
+            args.action_space,
+            args.model,
+            args.observation_type,
+            args.result_dir,
+            test_all_meta,
+        )
+        left_info = ""
+        for domain in test_file_list:
+            left_info += f"{domain}: {len(test_file_list[domain])}\n"
+        logger.info(f"Left tasks:\n{left_info}")
 
-    get_result(
-        args.action_space,
-        args.model,
-        args.observation_type,
-        args.result_dir,
-        test_all_meta,
-    )
-    test(args, test_file_list)
+        get_result(
+            args.action_space,
+            args.model,
+            args.observation_type,
+            args.result_dir,
+            test_all_meta,
+        )
+        test(args, test_file_list)
+    except KeyboardInterrupt:
+        logger.info("Main process received KeyboardInterrupt.")
+        # Signal handler will take care of cleanup
+    except Exception as e:
+        logger.error(f"Unexpected error in main process: {e}", exc_info=True)
+        # Also trigger cleanup for unhandled exceptions
+        signal_handler(signal.SIGTERM, None)
+    finally:
+        # Final cleanup in case any environments or processes remain
+        logger.info("Main process final cleanup...")
+        for env in active_environments:
+            if env is not None:
+                try:
+                    logger.info(f"Closing environment in final cleanup...")
+                    env.close()
+                    logger.info(f"Environment closed successfully in final cleanup")
+                except Exception as e:
+                    logger.error(f"Error during final environment cleanup: {e}")
+        
+        # First try gentle termination
+        for p in processes:
+            if p is not None and p.is_alive():
+                try:
+                    logger.info(f"Terminating process {p.name}...")
+                    p.terminate()
+                except Exception as e:
+                    logger.error(f"Error terminating process: {e}")
+        
+        # Wait a moment for processes to terminate
+        time.sleep(1)
+        
+        # Then force kill if needed
+        for p in processes:
+            if p is not None and p.is_alive():
+                try:
+                    logger.info(f"Force killing process {p.name}...")
+                    os.kill(p.pid, signal.SIGKILL)
+                    logger.info(f"Process {p.name} force killed")
+                except Exception as e:
+                    logger.error(f"Error force killing process: {e}")