From cbe650d0bbcfa87b88e93a8439baeeca2680ccfc Mon Sep 17 00:00:00 2001 From: Zilong Zhou Date: Thu, 24 Jul 2025 16:27:18 +0800 Subject: [PATCH] refactor&delete: simplify AWS VM allocation and remove proxy support (#284) --- desktop_env/providers/aws/manager.py | 71 +--- .../providers/aws/provider_with_proxy.py | 315 ------------------ 2 files changed, 3 insertions(+), 383 deletions(-) delete mode 100644 desktop_env/providers/aws/provider_with_proxy.py diff --git a/desktop_env/providers/aws/manager.py b/desktop_env/providers/aws/manager.py index 687b745..2c32249 100644 --- a/desktop_env/providers/aws/manager.py +++ b/desktop_env/providers/aws/manager.py @@ -1,7 +1,6 @@ import os from filelock import FileLock import boto3 -import psutil import logging import dotenv import signal @@ -163,60 +162,6 @@ 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, 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, screen_size=screen_size) - - from desktop_env.providers.aws.provider_with_proxy import AWSProviderWithProxy - - # Initialize proxy pool if needed - if proxy_config_file: - init_proxy_pool(proxy_config_file) - - # Get current proxy - proxy_pool = get_global_proxy_pool() - current_proxy = proxy_pool.get_next_proxy() - - if current_proxy: - logger.info(f"Allocating VM with proxy: {current_proxy.host}:{current_proxy.port}") - - # Create provider instance - provider = AWSProviderWithProxy(region=region, proxy_config_file=proxy_config_file) - - # Create new instance - instance_id = provider.create_instance_with_proxy( - image_id=IMAGE_ID_MAP[region], - instance_type=INSTANCE_TYPE, - security_groups=[os.getenv('AWS_SECURITY_GROUP_ID')], - subnet_id=os.getenv('AWS_SUBNET_ID') - ) - - try: - ec2_client = boto3.client('ec2', region_name=region) - instance_details = ec2_client.describe_instances(InstanceIds=[instance_id]) - instance = instance_details['Reservations'][0]['Instances'][0] - public_ip = instance.get('PublicIpAddress', '') - if public_ip: - vnc_url = f"http://{public_ip}:5910/vnc.html" - logger.info("="*80) - logger.info(f"šŸ–„ļø VNC Web Access URL: {vnc_url}") - logger.info(f"šŸ“” Public IP: {public_ip}") - logger.info(f"šŸ†” Instance ID: {instance_id}") - if current_proxy: - logger.info(f"šŸŒ Proxy: {current_proxy.host}:{current_proxy.port}") - logger.info("="*80) - print(f"\nšŸŒ VNC Web Access URL: {vnc_url}") - if current_proxy: - print(f"šŸ”„ Current Proxy: {current_proxy.host}:{current_proxy.port}") - print(f"šŸ“ Please open the above address in the browser for remote desktop access\n") - except Exception as e: - logger.warning(f"Failed to get VNC address for proxy instance {instance_id}: {e}") - - return instance_id - - class AWSVMManager(VMManager): """ AWS VM Manager for managing virtual machines on AWS. @@ -224,15 +169,9 @@ class AWSVMManager(VMManager): AWS does not need to maintain a registry of VMs, as it can dynamically allocate and deallocate VMs. This class supports both regular VM allocation and proxy-enabled VM allocation. """ - def __init__(self, proxy_config_file=None, **kwargs): - self.proxy_config_file = proxy_config_file + def __init__(self, **kwargs): # self.lock = FileLock(".aws_lck", timeout=60) self.initialize_registry() - - # Initialize proxy pool if proxy configuration is provided - if proxy_config_file and PROXY_SUPPORT_AVAILABLE: - init_proxy_pool(proxy_config_file) - logger.info(f"Proxy pool initialized with config: {proxy_config_file}") def initialize_registry(self, **kwargs): pass @@ -268,10 +207,6 @@ class AWSVMManager(VMManager): pass 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, screen_size=screen_size) - else: - logger.info("Allocating a new VM in region: {}".format(region)) - new_vm_path = _allocate_vm(region, screen_size=screen_size) + logger.info("Allocating a new VM in region: {}".format(region)) + new_vm_path = _allocate_vm(region, screen_size=screen_size) return new_vm_path \ No newline at end of file diff --git a/desktop_env/providers/aws/provider_with_proxy.py b/desktop_env/providers/aws/provider_with_proxy.py deleted file mode 100644 index 230cc4b..0000000 --- a/desktop_env/providers/aws/provider_with_proxy.py +++ /dev/null @@ -1,315 +0,0 @@ -import boto3 -from botocore.exceptions import ClientError -import base64 -import logging -import json -from typing import Optional - -from desktop_env.providers.base import Provider -from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool, ProxyInfo - -logger = logging.getLogger("desktopenv.providers.aws.AWSProviderWithProxy") -logger.setLevel(logging.INFO) - -WAIT_DELAY = 15 -MAX_ATTEMPTS = 10 - - -class AWSProviderWithProxy(Provider): - - def __init__(self, region: str = None, proxy_config_file: str = None): - super().__init__(region) - self.current_proxy: Optional[ProxyInfo] = None - - # Initialize proxy pool - if proxy_config_file: - init_proxy_pool(proxy_config_file) - logger.info(f"Initialized proxy pool from {proxy_config_file}") - - # Get next available proxy - self._rotate_proxy() - - def _rotate_proxy(self): - """Rotate to next available proxy""" - proxy_pool = get_global_proxy_pool() - self.current_proxy = proxy_pool.get_next_proxy() - - if self.current_proxy: - logger.info(f"Switched to proxy: {self.current_proxy.host}:{self.current_proxy.port}") - else: - logger.warning("No proxy available, using direct connection") - - def _generate_proxy_user_data(self) -> str: - """Generate user data script with proxy configuration""" - if not self.current_proxy: - return "" - - proxy_url = self._format_proxy_url(self.current_proxy) - - user_data_script = f"""#!/bin/bash -# Configure system proxy -echo 'export http_proxy={proxy_url}' >> /etc/environment -echo 'export https_proxy={proxy_url}' >> /etc/environment -echo 'export HTTP_PROXY={proxy_url}' >> /etc/environment -echo 'export HTTPS_PROXY={proxy_url}' >> /etc/environment - -# Configure apt proxy -cat > /etc/apt/apt.conf.d/95proxy << EOF -Acquire::http::Proxy "{proxy_url}"; -Acquire::https::Proxy "{proxy_url}"; -EOF - -# Configure chrome/chromium proxy -mkdir -p /etc/opt/chrome/policies/managed -cat > /etc/opt/chrome/policies/managed/proxy.json << EOF -{{ - "ProxyMode": "fixed_servers", - "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}" -}} -EOF - -# Configure chromium proxy (Ubuntu default) -mkdir -p /etc/chromium/policies/managed -cat > /etc/chromium/policies/managed/proxy.json << EOF -{{ - "ProxyMode": "fixed_servers", - "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}" -}} -EOF - -# Configure firefox proxy - support multiple possible paths -for firefox_dir in /etc/firefox/policies /usr/lib/firefox/distribution/policies /etc/firefox-esr/policies; do - if [ -d "$(dirname "$firefox_dir")" ]; then - mkdir -p "$firefox_dir" - cat > "$firefox_dir/policies.json" << EOF -{{ - "policies": {{ - "Proxy": {{ - "Mode": "manual", - "HTTPProxy": "{self.current_proxy.host}:{self.current_proxy.port}", - "HTTPSProxy": "{self.current_proxy.host}:{self.current_proxy.port}", - "UseHTTPProxyForAllProtocols": true - }} - }} -}} -EOF - break - fi -done - -# Reload environment variables -source /etc/environment - -# Log proxy configuration -echo "$(date): Configured proxy {self.current_proxy.host}:{self.current_proxy.port}" >> /var/log/proxy-setup.log -""" - - return base64.b64encode(user_data_script.encode()).decode() - - def _format_proxy_url(self, proxy: ProxyInfo) -> str: - """Format proxy URL""" - if proxy.username and proxy.password: - return f"{proxy.protocol}://{proxy.username}:{proxy.password}@{proxy.host}:{proxy.port}" - else: - return f"{proxy.protocol}://{proxy.host}:{proxy.port}" - - def start_emulator(self, path_to_vm: str, headless: bool, *args, **kwargs): - logger.info("Starting AWS VM with proxy configuration...") - ec2_client = boto3.client('ec2', region_name=self.region) - - try: - # If instance already exists, start it directly - ec2_client.start_instances(InstanceIds=[path_to_vm]) - logger.info(f"Instance {path_to_vm} is starting...") - - # Wait for the instance to be in the 'running' state - waiter = ec2_client.get_waiter('instance_running') - waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS}) - logger.info(f"Instance {path_to_vm} is now running.") - - except ClientError as e: - logger.error(f"Failed to start the AWS VM {path_to_vm}: {str(e)}") - raise - - def create_instance_with_proxy(self, image_id: str, instance_type: str, - security_groups: list, subnet_id: str) -> str: - """Create new instance with proxy configuration""" - ec2_client = boto3.client('ec2', region_name=self.region) - - user_data = self._generate_proxy_user_data() - - run_instances_params = { - "MaxCount": 1, - "MinCount": 1, - "ImageId": image_id, - "InstanceType": instance_type, - "EbsOptimized": True, - "NetworkInterfaces": [ - { - "SubnetId": subnet_id, - "AssociatePublicIpAddress": True, - "DeviceIndex": 0, - "Groups": security_groups - } - ] - } - - if user_data: - run_instances_params["UserData"] = user_data - - try: - response = ec2_client.run_instances(**run_instances_params) - instance_id = response['Instances'][0]['InstanceId'] - - logger.info(f"Created new instance {instance_id} with proxy configuration") - - logger.info(f"Waiting for instance {instance_id} to be running...") - ec2_client.get_waiter('instance_running').wait(InstanceIds=[instance_id]) - logger.info(f"Instance {instance_id} is ready.") - - try: - instance_details = ec2_client.describe_instances(InstanceIds=[instance_id]) - instance = instance_details['Reservations'][0]['Instances'][0] - public_ip = instance.get('PublicIpAddress', '') - if public_ip: - vnc_url = f"http://{public_ip}:5910/vnc.html" - logger.info("="*80) - logger.info(f"šŸ–„ļø VNC Web Access URL: {vnc_url}") - logger.info(f"šŸ“” Public IP: {public_ip}") - logger.info(f"šŸ†” Instance ID: {instance_id}") - if self.current_proxy: - logger.info(f"šŸŒ Proxy: {self.current_proxy.host}:{self.current_proxy.port}") - logger.info("="*80) - print(f"\nšŸŒ VNC Web Access URL: {vnc_url}") - if self.current_proxy: - print(f"šŸ”„ Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}") - print(f"šŸ“ Please open the above address in the browser for remote desktop access\n") - except Exception as e: - logger.warning(f"Failed to get VNC address for instance {instance_id}: {e}") - - return instance_id - - except ClientError as e: - logger.error(f"Failed to create instance with proxy: {str(e)}") - if self.current_proxy: - proxy_pool = get_global_proxy_pool() - proxy_pool.mark_proxy_failed(self.current_proxy) - self._rotate_proxy() - raise - - def get_ip_address(self, path_to_vm: str) -> str: - logger.info("Getting AWS VM IP address...") - ec2_client = boto3.client('ec2', region_name=self.region) - - try: - response = ec2_client.describe_instances(InstanceIds=[path_to_vm]) - for reservation in response['Reservations']: - for instance in reservation['Instances']: - private_ip_address = instance.get('PrivateIpAddress', '') - public_ip_address = instance.get('PublicIpAddress', '') - - if public_ip_address: - vnc_url = f"http://{public_ip_address}:5910/vnc.html" - logger.info("="*80) - logger.info(f"šŸ–„ļø VNC Web Access URL: {vnc_url}") - logger.info(f"šŸ“” Public IP: {public_ip_address}") - logger.info(f"šŸ  Private IP: {private_ip_address}") - if self.current_proxy: - logger.info(f"šŸŒ Proxy: {self.current_proxy.host}:{self.current_proxy.port}") - logger.info("="*80) - print(f"\nšŸŒ VNC Web Access URL: {vnc_url}") - if self.current_proxy: - print(f"šŸ”„ Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}") - print(f"šŸ“ Please open the above address in the browser for remote desktop access\n") - else: - logger.warning("No public IP address available for VNC access") - - return private_ip_address - return '' - except ClientError as e: - logger.error(f"Failed to retrieve IP address for the instance {path_to_vm}: {str(e)}") - raise - - def save_state(self, path_to_vm: str, snapshot_name: str): - logger.info("Saving AWS VM state...") - ec2_client = boto3.client('ec2', region_name=self.region) - - try: - image_response = ec2_client.create_image(InstanceId=path_to_vm, Name=snapshot_name) - image_id = image_response['ImageId'] - logger.info(f"AMI {image_id} created successfully from instance {path_to_vm}.") - return image_id - except ClientError as e: - logger.error(f"Failed to create AMI from the instance {path_to_vm}: {str(e)}") - raise - - def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str): - logger.info(f"Reverting AWS VM to snapshot: {snapshot_name}...") - ec2_client = boto3.client('ec2', region_name=self.region) - - try: - # Get original instance details for config. - instance_details = ec2_client.describe_instances(InstanceIds=[path_to_vm]) - instance = instance_details['Reservations'][0]['Instances'][0] - security_groups = [sg['GroupId'] for sg in instance['SecurityGroups']] - subnet_id = instance['SubnetId'] - instance_type = instance['InstanceType'] - - # Terminate the old instance. This is a non-blocking call. - logger.info(f"Initiating termination for old instance {path_to_vm}...") - ec2_client.terminate_instances(InstanceIds=[path_to_vm]) - logger.info(f"Old instance {path_to_vm} termination initiated.") - - # Rotate to a new proxy - self._rotate_proxy() - - # Create a new instance - new_instance_id = self.create_instance_with_proxy( - snapshot_name, instance_type, security_groups, subnet_id - ) - - # Note: VNC address is displayed within create_instance_with_proxy - logger.info(f"Successfully launched new instance {new_instance_id} for revert.") - - return new_instance_id - - except ClientError as e: - logger.error(f"Failed to revert to snapshot {snapshot_name} for the instance {path_to_vm}: {str(e)}") - raise - - def stop_emulator(self, path_to_vm, region=None): - logger.info(f"Stopping AWS VM {path_to_vm}...") - ec2_client = boto3.client('ec2', region_name=self.region) - - try: - ec2_client.stop_instances(InstanceIds=[path_to_vm]) - waiter = ec2_client.get_waiter('instance_stopped') - waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS}) - logger.info(f"Instance {path_to_vm} has been stopped.") - except ClientError as e: - logger.error(f"Failed to stop the AWS VM {path_to_vm}: {str(e)}") - raise - - def get_current_proxy_info(self) -> Optional[dict]: - """Get current proxy information""" - if self.current_proxy: - return { - 'host': self.current_proxy.host, - 'port': self.current_proxy.port, - 'protocol': self.current_proxy.protocol, - 'failed_count': self.current_proxy.failed_count - } - return None - - def force_rotate_proxy(self): - """Force rotate proxy""" - logger.info("Force rotating proxy...") - if self.current_proxy: - proxy_pool = get_global_proxy_pool() - proxy_pool.mark_proxy_failed(self.current_proxy) - self._rotate_proxy() - - def get_proxy_stats(self) -> dict: - """Get proxy pool statistics""" - proxy_pool = get_global_proxy_pool() - return proxy_pool.get_stats() \ No newline at end of file