fix: enhance setup method with retry logic and return status
This commit is contained in:
@@ -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]]):
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user