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

desktop_env/desktop_env.py

@@ -71,29 +71,39 @@ class DesktopEnv(gym.Env):
         else:
 
             self.path_to_vm = self.manager.get_vm_path(os_type=self.os_type, region=region)
+        try:
+            self.snapshot_name = snapshot_name
+            self.cache_dir_base: str = cache_dir
+            # todo: add the logic to get the screen size from the VM
+            self.headless = headless
+            self.require_a11y_tree = require_a11y_tree
+            self.require_terminal = require_terminal
 
-        self.snapshot_name = snapshot_name
-        self.cache_dir_base: str = cache_dir
-        # todo: add the logic to get the screen size from the VM
-        self.headless = headless
-        self.require_a11y_tree = require_a11y_tree
-        self.require_terminal = require_terminal
+            # Initialize emulator and controller
+            if provider_name != "docker": # Check if this is applicable to other VM providers
+                logger.info("Initializing...")
+                self._start_emulator()
 
-        # Initialize emulator and controller
-        if provider_name != "docker": # Check if this is applicable to other VM providers
-            logger.info("Initializing...")
-            self._start_emulator()
+            # mode: human or machine
+            self.instruction = None
+            assert action_space in ["computer_13", "pyautogui"]
+            self.action_space = action_space  # todo: refactor it to the ActType
 
-        # mode: human or machine
-        self.instruction = None
-        assert action_space in ["computer_13", "pyautogui"]
-        self.action_space = action_space  # todo: refactor it to the ActType
-
-        # episodic stuffs, like counters, will be updated or reset
-        # when calling self.reset()
-        self._traj_no: int = -1
-        self._step_no: int = 0
-        self.action_history: List[Dict[str, any]] = []
+            # episodic stuffs, like counters, will be updated or reset
+            # when calling self.reset()
+            self._traj_no: int = -1
+            self._step_no: int = 0
+            self.action_history: List[Dict[str, any]] = []
+        except Exception as e:
+            logger.error(f"Failed to initialize DesktopEnv: {e}")
+            # If initialization fails, we should clean up the VM
+            try:
+                self.close()
+                self.manager.delete_vm(self.path_to_vm, self.region)
+                logger.info(f"Cleaned up VM {self.path_to_vm}.")
+            except Exception as cleanup_error:
+                logger.error(f"Failed to clean up VM {self.path_to_vm}: {cleanup_error}")
+            raise
 
     def _start_emulator(self):
         # Power on the virtual machine