Merge remote-tracking branch 'upstream/fix/aws-proxy'
This commit is contained in:
@@ -47,11 +47,12 @@ class SetupController:
|
||||
self.http_server: str = f"http://{vm_ip}:{server_port}"
|
||||
self.http_server_setup_root: str = f"http://{vm_ip}:{server_port}/setup"
|
||||
self.cache_dir: str = cache_dir
|
||||
self.use_proxy: bool = False
|
||||
|
||||
def reset_cache_dir(self, cache_dir: str):
|
||||
self.cache_dir = cache_dir
|
||||
|
||||
def setup(self, config: List[Dict[str, Any]])-> bool:
|
||||
def setup(self, config: List[Dict[str, Any]], use_proxy: bool = False)-> bool:
|
||||
"""
|
||||
Args:
|
||||
config (List[Dict[str, Any]]): list of dict like {str: Any}. each
|
||||
@@ -59,10 +60,11 @@ class SetupController:
|
||||
{
|
||||
"type": str, corresponding to the `_{:}_setup` methods of
|
||||
this class
|
||||
"parameters": dick like {str, Any} providing the keyword
|
||||
"parameters": dict like {str, Any} providing the keyword
|
||||
parameters
|
||||
}
|
||||
"""
|
||||
self.use_proxy = use_proxy
|
||||
# make sure connection can be established
|
||||
logger.info(f"try to connect {self.http_server}")
|
||||
retry = 0
|
||||
@@ -264,6 +266,9 @@ class SetupController:
|
||||
if not shell and isinstance(command, str) and len(command.split()) > 1:
|
||||
logger.warning("Command should be a list of strings. Now it is a string. Will split it by space.")
|
||||
command = command.split()
|
||||
|
||||
if command[0] == "google-chrome" and self.use_proxy:
|
||||
command.append("--proxy-server=http://127.0.0.1:18888") # Use the proxy server set up by _proxy_setup
|
||||
|
||||
payload = json.dumps({"command": command, "shell": shell})
|
||||
headers = {"Content-Type": "application/json"}
|
||||
@@ -446,6 +451,19 @@ class SetupController:
|
||||
Args:
|
||||
client_password (str): Password for sudo operations, defaults to "password"
|
||||
"""
|
||||
retry = 0
|
||||
while retry < MAX_RETRIES:
|
||||
try:
|
||||
_ = requests.get(self.http_server + "/terminal")
|
||||
break
|
||||
except:
|
||||
time.sleep(5)
|
||||
retry += 1
|
||||
logger.info(f"retry: {retry}/{MAX_RETRIES}")
|
||||
|
||||
if retry == MAX_RETRIES:
|
||||
return False
|
||||
|
||||
# Get proxy from global proxy pool
|
||||
proxy_pool = get_global_proxy_pool()
|
||||
current_proxy = proxy_pool.get_next_proxy()
|
||||
@@ -458,14 +476,21 @@ class SetupController:
|
||||
proxy_url = proxy_pool._format_proxy_url(current_proxy)
|
||||
logger.info(f"Setting up proxy: {current_proxy.host}:{current_proxy.port}")
|
||||
|
||||
# Configure system proxy environment variables
|
||||
# Configure system proxy environment variables
|
||||
proxy_commands = [
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'export http_proxy={proxy_url}' >> /etc/environment\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'export https_proxy={proxy_url}' >> /etc/environment\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'export HTTP_PROXY={proxy_url}' >> /etc/environment\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'export HTTPS_PROXY={proxy_url}' >> /etc/environment\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"apt-get update\"", ## TODO: remove this line if ami is already updated
|
||||
f"echo '{client_password}' | sudo -S bash -c \"apt-get install -y tinyproxy\"", ## TODO: remove this line if tinyproxy is already installed
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'Port 18888' > /tmp/tinyproxy.conf\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'Allow 127.0.0.1' >> /tmp/tinyproxy.conf\"",
|
||||
f"echo '{client_password}' | sudo -S bash -c \"echo 'Upstream http {current_proxy.username}:{current_proxy.password}@{current_proxy.host}:{current_proxy.port}' >> /tmp/tinyproxy.conf\"",
|
||||
|
||||
# CML commands to set environment variables for proxy
|
||||
f"echo 'export http_proxy={proxy_url}' >> ~/.bashrc",
|
||||
f"echo 'export https_proxy={proxy_url}' >> ~/.bashrc",
|
||||
f"echo 'export HTTP_PROXY={proxy_url}' >> ~/.bashrc",
|
||||
f"echo 'export HTTPS_PROXY={proxy_url}' >> ~/.bashrc",
|
||||
]
|
||||
|
||||
|
||||
# Execute all proxy configuration commands
|
||||
for cmd in proxy_commands:
|
||||
try:
|
||||
@@ -475,6 +500,8 @@ class SetupController:
|
||||
proxy_pool.mark_proxy_failed(current_proxy)
|
||||
raise
|
||||
|
||||
self._launch_setup(["tinyproxy -c /tmp/tinyproxy.conf -d"], shell=True)
|
||||
|
||||
# Reload environment variables
|
||||
reload_cmd = "source /etc/environment"
|
||||
try:
|
||||
|
||||
@@ -215,10 +215,13 @@ class DesktopEnv(gym.Env):
|
||||
logger.info("Environment is clean, skipping snapshot revert (provider: {}).".format(self.provider_name))
|
||||
|
||||
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._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)
|
||||
success = self.setup_controller.setup(self.config, task_config.get("proxy", False) and self.enable_proxy)
|
||||
if success:
|
||||
# Mark environment as used when setup is successfully executed
|
||||
if self.config: # Only mark as used if there were actual setup operations
|
||||
@@ -235,10 +238,6 @@ class DesktopEnv(gym.Env):
|
||||
break
|
||||
|
||||
logger.info("Environment setup complete.")
|
||||
|
||||
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()
|
||||
|
||||
observation = self._get_obs()
|
||||
return observation
|
||||
|
||||
@@ -21,14 +21,8 @@ def create_vm_manager_and_provider(provider_name: str, region: str, use_proxy: b
|
||||
return VirtualBoxVMManager(), VirtualBoxProvider(region)
|
||||
elif provider_name in ["aws", "amazon web services"]:
|
||||
from desktop_env.providers.aws.manager import AWSVMManager
|
||||
if use_proxy:
|
||||
# Use proxy-enabled AWS provider
|
||||
from desktop_env.providers.aws.provider_with_proxy import AWSProviderWithProxy
|
||||
return AWSVMManager(proxy_config_file="dataimpulse_proxy_config.json"), AWSProviderWithProxy(region, proxy_config_file="dataimpulse_proxy_config.json")
|
||||
else:
|
||||
# Use regular AWS provider
|
||||
from desktop_env.providers.aws.provider import AWSProvider
|
||||
return AWSVMManager(), AWSProvider(region)
|
||||
from desktop_env.providers.aws.provider import AWSProvider
|
||||
return AWSVMManager(), AWSProvider(region)
|
||||
elif provider_name == "azure":
|
||||
from desktop_env.providers.azure.manager import AzureVMManager
|
||||
from desktop_env.providers.azure.provider import AzureProvider
|
||||
|
||||
Reference in New Issue
Block a user