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):
|
def reset_cache_dir(self, cache_dir: str):
|
||||||
self.cache_dir = cache_dir
|
self.cache_dir = cache_dir
|
||||||
|
|
||||||
def setup(self, config: List[Dict[str, Any]]):
|
def setup(self, config: List[Dict[str, Any]])-> bool:
|
||||||
"""
|
"""
|
||||||
Args:
|
Args:
|
||||||
config (List[Dict[str, Any]]): list of dict like {str: Any}. each
|
config (List[Dict[str, Any]]): list of dict like {str: Any}. each
|
||||||
@@ -64,13 +64,18 @@ class SetupController:
|
|||||||
# make sure connection can be established
|
# make sure connection can be established
|
||||||
logger.info(f"try to connect {self.http_server}")
|
logger.info(f"try to connect {self.http_server}")
|
||||||
retry = 0
|
retry = 0
|
||||||
while retry < 30:
|
while retry < 50:
|
||||||
try:
|
try:
|
||||||
_ = requests.get(self.http_server + "/terminal")
|
_ = requests.get(self.http_server + "/terminal")
|
||||||
break
|
break
|
||||||
except:
|
except:
|
||||||
time.sleep(5)
|
time.sleep(5)
|
||||||
retry += 1
|
retry += 1
|
||||||
|
logger.info(f"retry: {retry}/50")
|
||||||
|
|
||||||
|
if retry == 50:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
for cfg in config:
|
for cfg in config:
|
||||||
config_type: str = cfg["type"]
|
config_type: str = cfg["type"]
|
||||||
@@ -85,6 +90,8 @@ class SetupController:
|
|||||||
|
|
||||||
logger.info("SETUP: %s(%s)", setup_function, str(parameters))
|
logger.info("SETUP: %s(%s)", setup_function, str(parameters))
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
def _download_setup(self, files: List[Dict[str, str]]):
|
def _download_setup(self, files: List[Dict[str, str]]):
|
||||||
"""
|
"""
|
||||||
Args:
|
Args:
|
||||||
|
|||||||
@@ -145,6 +145,7 @@ class DesktopEnv(gym.Env):
|
|||||||
self.provider.stop_emulator(self.path_to_vm)
|
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]:
|
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
|
# Reset to certain task in OSWorld
|
||||||
logger.info("Resetting environment...")
|
logger.info("Resetting environment...")
|
||||||
logger.info("Switching task...")
|
logger.info("Switching task...")
|
||||||
@@ -152,6 +153,7 @@ class DesktopEnv(gym.Env):
|
|||||||
self._traj_no += 1
|
self._traj_no += 1
|
||||||
self._step_no = 0
|
self._step_no = 0
|
||||||
self.action_history.clear()
|
self.action_history.clear()
|
||||||
|
while retry < 5:
|
||||||
|
|
||||||
# Check and handle proxy requirement changes BEFORE starting emulator
|
# Check and handle proxy requirement changes BEFORE starting emulator
|
||||||
if task_config is not None:
|
if task_config is not None:
|
||||||
@@ -179,6 +181,7 @@ class DesktopEnv(gym.Env):
|
|||||||
else:
|
else:
|
||||||
logger.info("Using regular AWS provider.")
|
logger.info("Using regular AWS provider.")
|
||||||
|
|
||||||
|
|
||||||
logger.info("Reverting to snapshot to {}...".format(self.snapshot_name))
|
logger.info("Reverting to snapshot to {}...".format(self.snapshot_name))
|
||||||
self._revert_to_snapshot()
|
self._revert_to_snapshot()
|
||||||
logger.info("Starting emulator...")
|
logger.info("Starting emulator...")
|
||||||
@@ -189,7 +192,14 @@ class DesktopEnv(gym.Env):
|
|||||||
self._set_task_info(task_config)
|
self._set_task_info(task_config)
|
||||||
self.setup_controller.reset_cache_dir(self.cache_dir)
|
self.setup_controller.reset_cache_dir(self.cache_dir)
|
||||||
logger.info("Setting up environment...")
|
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.")
|
logger.info("Environment setup complete.")
|
||||||
|
|
||||||
if task_config.get("proxy", False):
|
if task_config.get("proxy", False):
|
||||||
|
|||||||
Reference in New Issue
Block a user