4e11eafd1d
* Refactor evaluator structure in LibreOffice Writer example JSON to support multiple expected and result files, enhancing evaluation flexibility. * Update instance type to t3.large and add VNC access URL logging for allocated VMs, enhancing remote access capabilities. * Update instance type to t3.large and add VNC access URL logging for allocated VMs, enhancing remote access capabilities. * Update time format in get_vm_file function to include hours, minutes, and seconds for more precise file naming with time suffix. * More delay for 936321ce-5236-426a-9a20-e0e3c5dc536f; support one more potential solutions. * Enhance SetupController with configurable retry limit and improved error handling for file opening requests. Introduce new function to compare unique training records, and update logging for better debugging. Adjust JSON examples for evaluation to support multiple expected and result files. * Clean debug code --------- Co-authored-by: yuanmengqi <yuanmengqi@mail.ustc.edu.cn>
315 lines
13 KiB
Python
315 lines
13 KiB
Python
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
|
|
|
|
# 初始化代理池
|
|
if proxy_config_file:
|
|
init_proxy_pool(proxy_config_file)
|
|
logger.info(f"Initialized proxy pool from {proxy_config_file}")
|
|
|
|
# 获取下一个可用代理
|
|
self._rotate_proxy()
|
|
|
|
def _rotate_proxy(self):
|
|
"""轮换到下一个可用代理"""
|
|
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:
|
|
"""生成包含代理配置的user data脚本"""
|
|
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:
|
|
"""格式化代理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:
|
|
# 如果实例已经存在，直接启动
|
|
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:
|
|
"""创建带有代理配置的新实例"""
|
|
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]:
|
|
"""获取当前代理信息"""
|
|
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):
|
|
"""强制轮换代理"""
|
|
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:
|
|
"""获取代理池统计信息"""
|
|
proxy_pool = get_global_proxy_pool()
|
|
return proxy_pool.get_stats() |