fix: enhance setup method with retry logic and return status

desktop_env/controllers/setup.py

@@ -49,7 +49,7 @@ class SetupController:
     def reset_cache_dir(self, cache_dir: str):
         self.cache_dir = cache_dir
 
-    def setup(self, config: List[Dict[str, Any]]):
+    def setup(self, config: List[Dict[str, Any]])-> bool:
         """
         Args:
             config (List[Dict[str, Any]]): list of dict like {str: Any}. each
@@ -64,13 +64,18 @@ class SetupController:
         # make sure connection can be established
         logger.info(f"try to connect {self.http_server}")
         retry = 0
-        while retry < 30:
+        while retry < 50:
             try:
                 _ = requests.get(self.http_server + "/terminal")
                 break
             except:
                 time.sleep(5)
                 retry += 1
+                logger.info(f"retry: {retry}/50")
+            
+            if retry == 50:
+                return False
+                
 
         for cfg in config:
             config_type: str = cfg["type"]
@@ -85,6 +90,8 @@ class SetupController:
 
             logger.info("SETUP: %s(%s)", setup_function, str(parameters))
         
+        return True
+
     def _download_setup(self, files: List[Dict[str, str]]):
         """
         Args:

desktop_env/desktop_env.py

@@ -145,6 +145,7 @@ class DesktopEnv(gym.Env):
         self.provider.stop_emulator(self.path_to_vm)
 
     def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None) -> Dict[str, Any]:
+        retry = 0
         # Reset to certain task in OSWorld
         logger.info("Resetting environment...")
         logger.info("Switching task...")
@@ -152,6 +153,7 @@ class DesktopEnv(gym.Env):
         self._traj_no += 1
         self._step_no = 0
         self.action_history.clear()
+        while retry < 5:
 
             # Check and handle proxy requirement changes BEFORE starting emulator
             if task_config is not None:
@@ -179,6 +181,7 @@ class DesktopEnv(gym.Env):
                     else:
                         logger.info("Using regular AWS provider.")
 
+        
             logger.info("Reverting to snapshot to {}...".format(self.snapshot_name))
             self._revert_to_snapshot()
             logger.info("Starting emulator...")
@@ -189,7 +192,14 @@ class DesktopEnv(gym.Env):
                 self._set_task_info(task_config)
                 self.setup_controller.reset_cache_dir(self.cache_dir)
                 logger.info("Setting up environment...")
-            self.setup_controller.setup(self.config)
+                success = self.setup_controller.setup(self.config)
+                if success:
+                    break
+                else:
+                    logger.error("Environment setup failed, retrying...")
+                    retry += 1
+                    time.sleep(5)
+                    
         logger.info("Environment setup complete.")
     
         if task_config.get("proxy", False):