Improve code logic for password & resolution (#252)
* fix chrome * fix: fix proxy setup * feat&fix: add proxy support in setup and remove hardcoded proxy from example * fix tasks * fix chrome finished * fix * clean chrome_fix code * clean chrome_fix code * fix chrome 2888b4e6-5b47-4b57-8bf5-c73827890774 * fix multiapps * fix chrome 2888b4e6-5b47-4b57-8bf5-c73827890774 * fix some multi_apps tasks * fix some multi_apps tasks * fix password&resolution * fix password&resolution * Improve code logic for password & resolution * edit * Merge branch 'main' into fix_chrome * fix chrome tasks --------- Co-authored-by: adlsdztony <zzl0712@connect.hku.hk>
This commit is contained in:
@@ -27,12 +27,6 @@ import dotenv
|
||||
# Load environment variables from .env file
|
||||
dotenv.load_dotenv()
|
||||
|
||||
if os.environ.get("PROVIDER_NAME") == "aws":
|
||||
os.environ["CLIENT_PASSWORD"] = os.environ.get("CLIENT_PASSWORD_AWS", "osworld-public-evaluation")
|
||||
else:
|
||||
os.environ["CLIENT_PASSWORD"] = os.environ.get("CLIENT_PASSWORD", "password")
|
||||
|
||||
CLIENT_PASSWORD = os.environ["CLIENT_PASSWORD"]
|
||||
|
||||
PROXY_CONFIG_FILE = os.getenv("PROXY_CONFIG_FILE", "evaluation_examples/settings/proxy/dataimpulse.json") # Default proxy config file
|
||||
|
||||
@@ -45,7 +39,7 @@ init_proxy_pool(PROXY_CONFIG_FILE) # initialize the global proxy pool
|
||||
MAX_RETRIES = 20
|
||||
|
||||
class SetupController:
|
||||
def __init__(self, vm_ip: str, server_port: int = 5000, chromium_port: int = 9222, vlc_port: int = 8080, cache_dir: str = "cache"):
|
||||
def __init__(self, vm_ip: str, server_port: int = 5000, chromium_port: int = 9222, vlc_port: int = 8080, cache_dir: str = "cache", client_password: str = "", screen_width: int = 1920, screen_height: int = 1080):
|
||||
self.vm_ip: str = vm_ip
|
||||
self.server_port: int = server_port
|
||||
self.chromium_port: int = chromium_port
|
||||
@@ -54,6 +48,9 @@ class SetupController:
|
||||
self.http_server_setup_root: str = f"http://{vm_ip}:{server_port}/setup"
|
||||
self.cache_dir: str = cache_dir
|
||||
self.use_proxy: bool = False
|
||||
self.client_password: str = client_password
|
||||
self.screen_width: int = screen_width
|
||||
self.screen_height: int = screen_height
|
||||
|
||||
def reset_cache_dir(self, cache_dir: str):
|
||||
self.cache_dir = cache_dir
|
||||
@@ -304,22 +301,31 @@ class SetupController:
|
||||
terminates: bool = False
|
||||
nb_failings = 0
|
||||
|
||||
def replace_screen_env_in_command(command_list):
|
||||
width = int(os.environ.get("SCREEN_WIDTH", 1920))
|
||||
height = int(os.environ.get("SCREEN_HEIGHT", 1080))
|
||||
def replace_screen_env_in_command(command):
|
||||
password = self.client_password
|
||||
width = self.screen_width
|
||||
height = self.screen_height
|
||||
width_half = str(width // 2)
|
||||
height_half = str(height // 2)
|
||||
new_command_list = []
|
||||
for item in command_list:
|
||||
if isinstance(item, str):
|
||||
new_command = ""
|
||||
if isinstance(command, str):
|
||||
new_command = command.replace("{CLIENT_PASSWORD}", password)
|
||||
new_command = new_command.replace("{SCREEN_WIDTH_HALF}", width_half)
|
||||
new_command = new_command.replace("{SCREEN_HEIGHT_HALF}", height_half)
|
||||
new_command = new_command.replace("{SCREEN_WIDTH}", str(width))
|
||||
new_command = new_command.replace("{SCREEN_HEIGHT}", str(height))
|
||||
return new_command
|
||||
else:
|
||||
for item in command:
|
||||
item = item.replace("{CLIENT_PASSWORD}", password)
|
||||
item = item.replace("{SCREEN_WIDTH_HALF}", width_half)
|
||||
item = item.replace("{SCREEN_HEIGHT_HALF}", height_half)
|
||||
item = item.replace("{SCREEN_WIDTH}", str(width))
|
||||
item = item.replace("{SCREEN_HEIGHT}", str(height))
|
||||
new_command_list.append(item)
|
||||
return new_command_list
|
||||
if isinstance(command, list):
|
||||
command = replace_screen_env_in_command(command)
|
||||
new_command_list.append(item)
|
||||
return new_command_list
|
||||
command = replace_screen_env_in_command(command)
|
||||
payload = json.dumps({"command": command, "shell": shell})
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
@@ -467,7 +473,7 @@ class SetupController:
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error("An error occurred while trying to send the request: %s", e)
|
||||
|
||||
def _proxy_setup(self, client_password: str = CLIENT_PASSWORD):
|
||||
def _proxy_setup(self, client_password: str = ""):
|
||||
"""Setup system-wide proxy configuration using proxy pool
|
||||
|
||||
Args:
|
||||
|
||||
@@ -26,7 +26,7 @@ class DesktopEnv(gym.Env):
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
provider_name: str = "vmware",
|
||||
provider_name: str = "aws",
|
||||
region: str = None,
|
||||
path_to_vm: str = None,
|
||||
snapshot_name: str = "init_state",
|
||||
@@ -38,6 +38,7 @@ class DesktopEnv(gym.Env):
|
||||
require_terminal: bool = False,
|
||||
os_type: str = "Ubuntu",
|
||||
enable_proxy: bool = False,
|
||||
client_password: str = "",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -59,6 +60,16 @@ class DesktopEnv(gym.Env):
|
||||
self.region = region
|
||||
self.provider_name = provider_name
|
||||
self.enable_proxy = enable_proxy # Store proxy enablement setting
|
||||
if client_password == "":
|
||||
if self.provider_name == "aws":
|
||||
self.client_password = "osworld-public-evaluation"
|
||||
else:
|
||||
self.client_password = "password"
|
||||
else:
|
||||
self.client_password = client_password
|
||||
|
||||
self.screen_width = screen_size[0]
|
||||
self.screen_height = screen_size[1]
|
||||
|
||||
# Default
|
||||
self.server_port = 5000
|
||||
@@ -88,7 +99,7 @@ class DesktopEnv(gym.Env):
|
||||
if provider_name in {"vmware", "virtualbox"} else path_to_vm
|
||||
else:
|
||||
|
||||
self.path_to_vm = self.manager.get_vm_path(os_type=self.os_type, region=region)
|
||||
self.path_to_vm = self.manager.get_vm_path(os_type=self.os_type, region=region, screen_size=(self.screen_width, self.screen_height))
|
||||
try:
|
||||
self.snapshot_name = snapshot_name
|
||||
self.cache_dir_base: str = cache_dir
|
||||
@@ -136,7 +147,7 @@ class DesktopEnv(gym.Env):
|
||||
self.vnc_port = int(vm_ip_ports[3])
|
||||
self.vlc_port = int(vm_ip_ports[4])
|
||||
self.controller = PythonController(vm_ip=self.vm_ip, server_port=self.server_port)
|
||||
self.setup_controller = SetupController(vm_ip=self.vm_ip, server_port=self.server_port, chromium_port=self.chromium_port, vlc_port=self.vlc_port, cache_dir=self.cache_dir_base)
|
||||
self.setup_controller = SetupController(vm_ip=self.vm_ip, server_port=self.server_port, chromium_port=self.chromium_port, vlc_port=self.vlc_port, cache_dir=self.cache_dir_base, client_password=self.client_password, screen_width=self.screen_width, screen_height=self.screen_height)
|
||||
|
||||
def _revert_to_snapshot(self):
|
||||
# Revert to certain snapshot of the virtual machine, and refresh the path to vm and ip of vm
|
||||
@@ -197,7 +208,7 @@ class DesktopEnv(gym.Env):
|
||||
if task_config is not None:
|
||||
if task_config.get("proxy", False) and self.enable_proxy:
|
||||
# If using proxy and proxy is enabled, set up the proxy configuration
|
||||
self.setup_controller._proxy_setup()
|
||||
self.setup_controller._proxy_setup(self.client_password)
|
||||
self._set_task_info(task_config)
|
||||
self.setup_controller.reset_cache_dir(self.cache_dir)
|
||||
logger.info("Setting up environment...")
|
||||
|
||||
@@ -164,11 +164,11 @@ def _allocate_vm(region=DEFAULT_REGION, screen_size=(1920, 1080)):
|
||||
return instance_id
|
||||
|
||||
|
||||
def _allocate_vm_with_proxy(region=DEFAULT_REGION, proxy_config_file=None):
|
||||
def _allocate_vm_with_proxy(region=DEFAULT_REGION, proxy_config_file=None, screen_size=(1920, 1080)):
|
||||
"""Allocate a VM with proxy configuration"""
|
||||
if not PROXY_SUPPORT_AVAILABLE:
|
||||
logger.warning("Proxy support not available, falling back to regular VM allocation")
|
||||
return _allocate_vm(region)
|
||||
return _allocate_vm(region, screen_size=screen_size)
|
||||
|
||||
from desktop_env.providers.aws.provider_with_proxy import AWSProviderWithProxy
|
||||
|
||||
@@ -268,11 +268,11 @@ class AWSVMManager(VMManager):
|
||||
def _list_free_vms(self, region=DEFAULT_REGION):
|
||||
pass
|
||||
|
||||
def get_vm_path(self, region=DEFAULT_REGION, **kwargs):
|
||||
def get_vm_path(self, region=DEFAULT_REGION, screen_size=(1920, 1080), **kwargs):
|
||||
if self.proxy_config_file:
|
||||
logger.info("Allocating a new VM with proxy configuration in region: {}".format(region))
|
||||
new_vm_path = _allocate_vm_with_proxy(region, self.proxy_config_file)
|
||||
new_vm_path = _allocate_vm_with_proxy(region, self.proxy_config_file, screen_size=screen_size)
|
||||
else:
|
||||
logger.info("Allocating a new VM in region: {}".format(region))
|
||||
new_vm_path = _allocate_vm(region)
|
||||
new_vm_path = _allocate_vm(region, screen_size=screen_size)
|
||||
return new_vm_path
|
||||
Reference in New Issue
Block a user