sci-gui-agent-benchmark/desktop_env/providers/aws/manager.py

import os
from filelock import FileLock
import boto3
import psutil
import logging
import dotenv
import signal
# 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
# ami-05e7d7bd279ea4f14
IMAGE_ID_MAP = {
    "us-east-1": "ami-02fea2e5b77c79c17",
    "ap-east-1": "ami-0c092a5b8be4116f5",
}

INSTANCE_TYPE = "t3.medium"

def _allocate_vm(region=DEFAULT_REGION):
    
    if region not in IMAGE_ID_MAP:
        raise ValueError(f"Region {region} is not supported. Supported regions are: {list(IMAGE_ID_MAP.keys())}")

    run_instances_params = {
        "MaxCount": 1,
        "MinCount": 1,
        "ImageId": IMAGE_ID_MAP[region],
        "InstanceType": INSTANCE_TYPE,
        "EbsOptimized": True,
        "NetworkInterfaces": [
            {
                "SubnetId": os.getenv('AWS_SUBNET_ID'),
                "AssociatePublicIpAddress": True,
                "DeviceIndex": 0,
                "Groups": [
                    os.getenv('AWS_SECURITY_GROUP_ID')
                ]
            }
        ]
    }

    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)
        
        response = ec2_client.run_instances(**run_instances_params)
        instance_id = response['Instances'][0]['InstanceId']
        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.")
    except KeyboardInterrupt:
        logger.warning("VM allocation interrupted by user (SIGINT).")
        raise
    except SystemExit:
        logger.warning("VM allocation terminated by parent process (SIGTERM).")
        raise
    except Exception as e:
        logger.error(f"Failed to allocate VM in region {region}: {str(e)}")
        # try to clean up any resources that were created
        try:
            if instance_id:
                ec2_client.terminate_instances(InstanceIds=[instance_id])
                logger.info(f"Terminated instance {instance_id} due to allocation failure.")
        except Exception as cleanup_error:
            logger.error(f"May fail to clean up instance {instance_id}: {str(cleanup_error)}")
        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):
    """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)
    
    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')
    )
    
    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, **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)
        else:
            logger.info("Allocating a new VM in region: {}".format(region))
            new_vm_path = _allocate_vm(region)
        return new_vm_path