4fa59ebba2
- Added a new function to ensure URLs have a scheme, defaulting to 'http://' if missing. - Integrated tldextract to normalize URLs by extracting domain parts and handling 'www' subdomains. - Updated the compare_urls function to include logging for better traceability during URL comparisons. - Added tldextract to requirements.txt to support the new functionality. - Updated the AWS manager with a new AMI ID for the specified resolution. - Modified Chrome desktop launcher to include --remote-debugging-port=1337 for GUI debugging support. These changes improve the robustness of URL handling and enable consistent Chrome debugging capabilities without altering existing logic.
277 lines
11 KiB
Python
277 lines
11 KiB
Python
import os
|
|
from filelock import FileLock
|
|
import boto3
|
|
import psutil
|
|
import logging
|
|
import dotenv
|
|
import signal
|
|
|
|
|
|
INSTANCE_TYPE = "t3.xlarge"
|
|
|
|
# Load environment variables from .env file
|
|
dotenv.load_dotenv()
|
|
|
|
# Ensure the AWS region is set in the environment
|
|
if not os.getenv('AWS_REGION'):
|
|
raise EnvironmentError("AWS_REGION must be set in the environment variables.")
|
|
|
|
# Ensure the AWS subnet and security group IDs are set in the environment
|
|
if not os.getenv('AWS_SUBNET_ID') or not os.getenv('AWS_SECURITY_GROUP_ID'):
|
|
raise EnvironmentError("AWS_SUBNET_ID and AWS_SECURITY_GROUP_ID must be set in the environment variables.")
|
|
|
|
from desktop_env.providers.base import VMManager
|
|
|
|
# Import proxy-related modules only when needed
|
|
try:
|
|
from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool
|
|
PROXY_SUPPORT_AVAILABLE = True
|
|
except ImportError:
|
|
PROXY_SUPPORT_AVAILABLE = False
|
|
|
|
logger = logging.getLogger("desktopenv.providers.aws.AWSVMManager")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
DEFAULT_REGION = "us-east-1"
|
|
# todo: Add doc for the configuration of image, security group and network interface
|
|
# todo: public the AMI images
|
|
IMAGE_ID_MAP = {
|
|
"us-east-1": {
|
|
(1920, 1080): "ami-0d23263edb96951d8"
|
|
},
|
|
"ap-east-1": {
|
|
(1920, 1080): "ami-0c092a5b8be4116f5"
|
|
}
|
|
}
|
|
|
|
|
|
def _allocate_vm(region=DEFAULT_REGION, screen_size=(1920, 1080)):
|
|
|
|
if region not in IMAGE_ID_MAP:
|
|
raise ValueError(f"Region {region} is not supported. Supported regions are: {list(IMAGE_ID_MAP.keys())}")
|
|
if screen_size not in IMAGE_ID_MAP[region]:
|
|
raise ValueError(f"Screen size {screen_size} not supported for region {region}. Supported: {list(IMAGE_ID_MAP[region].keys())}")
|
|
ami_id = IMAGE_ID_MAP[region][screen_size]
|
|
|
|
ec2_client = boto3.client('ec2', region_name=region)
|
|
instance_id = None
|
|
original_sigint_handler = signal.getsignal(signal.SIGINT)
|
|
original_sigterm_handler = signal.getsignal(signal.SIGTERM)
|
|
|
|
def signal_handler(sig, frame):
|
|
if instance_id:
|
|
signal_name = "SIGINT" if sig == signal.SIGINT else "SIGTERM"
|
|
logger.warning(f"Received {signal_name} signal, terminating instance {instance_id}...")
|
|
try:
|
|
ec2_client.terminate_instances(InstanceIds=[instance_id])
|
|
logger.info(f"Successfully terminated instance {instance_id} after {signal_name}.")
|
|
except Exception as cleanup_error:
|
|
logger.error(f"Failed to terminate instance {instance_id} after {signal_name}: {str(cleanup_error)}")
|
|
|
|
# Restore original signal handlers
|
|
signal.signal(signal.SIGINT, original_sigint_handler)
|
|
signal.signal(signal.SIGTERM, original_sigterm_handler)
|
|
|
|
# Raise appropriate exception based on signal type
|
|
if sig == signal.SIGINT:
|
|
raise KeyboardInterrupt
|
|
else:
|
|
# For SIGTERM, exit gracefully
|
|
import sys
|
|
sys.exit(0)
|
|
|
|
try:
|
|
# Set up signal handlers for both SIGINT and SIGTERM
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
|
|
if not os.getenv('AWS_SECURITY_GROUP_ID'):
|
|
raise ValueError("AWS_SECURITY_GROUP_ID is not set in the environment variables.")
|
|
if not os.getenv('AWS_SUBNET_ID'):
|
|
raise ValueError("AWS_SUBNET_ID is not set in the environment variables.")
|
|
|
|
run_instances_params = {
|
|
"MaxCount": 1,
|
|
"MinCount": 1,
|
|
"ImageId": ami_id,
|
|
"InstanceType": INSTANCE_TYPE,
|
|
"EbsOptimized": True,
|
|
"NetworkInterfaces": [
|
|
{
|
|
"SubnetId": os.getenv('AWS_SUBNET_ID'),
|
|
"AssociatePublicIpAddress": True,
|
|
"DeviceIndex": 0,
|
|
"Groups": [
|
|
os.getenv('AWS_SECURITY_GROUP_ID')
|
|
]
|
|
}
|
|
],
|
|
"BlockDeviceMappings": [
|
|
{
|
|
"DeviceName": "/dev/sda1",
|
|
"Ebs": {
|
|
# "VolumeInitializationRate": 300
|
|
"VolumeSize": 30, # Size in GB
|
|
"VolumeType": "gp3", # General Purpose SSD
|
|
"Throughput": 1000,
|
|
"Iops": 4000 # Adjust IOPS as needed
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
response = ec2_client.run_instances(**run_instances_params)
|
|
instance_id = response['Instances'][0]['InstanceId']
|
|
|
|
waiter = ec2_client.get_waiter('instance_running')
|
|
logger.info(f"Waiting for instance {instance_id} to be running...")
|
|
waiter.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}")
|
|
logger.info("="*80)
|
|
print(f"\n🌐 VNC Web Access URL: {vnc_url}")
|
|
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}")
|
|
except KeyboardInterrupt:
|
|
logger.warning("VM allocation interrupted by user (SIGINT).")
|
|
if instance_id:
|
|
logger.info(f"Terminating instance {instance_id} due to interruption.")
|
|
ec2_client.terminate_instances(InstanceIds=[instance_id])
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Failed to allocate VM: {e}", exc_info=True)
|
|
if instance_id:
|
|
logger.info(f"Terminating instance {instance_id} due to an error.")
|
|
ec2_client.terminate_instances(InstanceIds=[instance_id])
|
|
raise
|
|
finally:
|
|
# Restore original signal handlers
|
|
signal.signal(signal.SIGINT, original_sigint_handler)
|
|
signal.signal(signal.SIGTERM, original_sigterm_handler)
|
|
|
|
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.
|
|
|
|
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
|
|
# 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
|
|
|
|
def add_vm(self, vm_path, region=DEFAULT_REGION, lock_needed=True, **kwargs):
|
|
pass
|
|
|
|
def _add_vm(self, vm_path, region=DEFAULT_REGION):
|
|
pass
|
|
|
|
def delete_vm(self, vm_path, region=DEFAULT_REGION, lock_needed=True, **kwargs):
|
|
pass
|
|
|
|
def _delete_vm(self, vm_path, region=DEFAULT_REGION):
|
|
pass
|
|
|
|
def occupy_vm(self, vm_path, pid, region=DEFAULT_REGION, lock_needed=True, **kwargs):
|
|
pass
|
|
|
|
def _occupy_vm(self, vm_path, pid, region=DEFAULT_REGION):
|
|
pass
|
|
|
|
def check_and_clean(self, lock_needed=True, **kwargs):
|
|
pass
|
|
|
|
def _check_and_clean(self):
|
|
pass
|
|
|
|
def list_free_vms(self, region=DEFAULT_REGION, lock_needed=True, **kwargs):
|
|
pass
|
|
|
|
def _list_free_vms(self, region=DEFAULT_REGION):
|
|
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)
|
|
return new_vm_path |