diff --git a/desktop_env/controllers/setup.py b/desktop_env/controllers/setup.py index 096d8e2..16985ea 100644 --- a/desktop_env/controllers/setup.py +++ b/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"] @@ -84,6 +89,8 @@ class SetupController: getattr(self, setup_function)(**parameters) logger.info("SETUP: %s(%s)", setup_function, str(parameters)) + + return True def _download_setup(self, files: List[Dict[str, str]]): """ diff --git a/desktop_env/desktop_env.py b/desktop_env/desktop_env.py index 43114fe..8aff6a8 100644 --- a/desktop_env/desktop_env.py +++ b/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,49 +153,58 @@ 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: - task_use_proxy = task_config.get("proxy", False) - if task_use_proxy != self.current_use_proxy: - logger.info(f"Task proxy requirement changed: {self.current_use_proxy} -> {task_use_proxy}") - - # Close current provider if it exists - if hasattr(self, 'provider') and self.provider: - try: - self.provider.stop_emulator(self.path_to_vm) - except Exception as e: - logger.warning(f"Failed to stop current provider: {e}") - - # Create new provider with appropriate proxy setting - self.current_use_proxy = task_use_proxy - self.manager, self.provider = create_vm_manager_and_provider( - self.provider_name, - self.region, - use_proxy=task_use_proxy - ) - - if task_use_proxy: - logger.info("Using proxy-enabled AWS provider.") - else: - logger.info("Using regular AWS provider.") + # Check and handle proxy requirement changes BEFORE starting emulator + if task_config is not None: + task_use_proxy = task_config.get("proxy", False) + if task_use_proxy != self.current_use_proxy: + logger.info(f"Task proxy requirement changed: {self.current_use_proxy} -> {task_use_proxy}") + + # Close current provider if it exists + if hasattr(self, 'provider') and self.provider: + try: + self.provider.stop_emulator(self.path_to_vm) + except Exception as e: + logger.warning(f"Failed to stop current provider: {e}") + + # Create new provider with appropriate proxy setting + self.current_use_proxy = task_use_proxy + self.manager, self.provider = create_vm_manager_and_provider( + self.provider_name, + self.region, + use_proxy=task_use_proxy + ) + + if task_use_proxy: + logger.info("Using proxy-enabled AWS provider.") + 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...") - self._start_emulator() - logger.info("Emulator started.") - - if task_config is not None: - 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) - logger.info("Environment setup complete.") - if task_config.get("proxy", False): - # If using proxy, set up the proxy configuration - self.setup_controller._proxy_setup() + logger.info("Reverting to snapshot to {}...".format(self.snapshot_name)) + self._revert_to_snapshot() + logger.info("Starting emulator...") + self._start_emulator() + logger.info("Emulator started.") + + if task_config is not None: + self._set_task_info(task_config) + self.setup_controller.reset_cache_dir(self.cache_dir) + logger.info("Setting up environment...") + 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): + # If using proxy, set up the proxy configuration + self.setup_controller._proxy_setup() observation = self._get_obs() return observation