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

logger = logging.getLogger("desktopenv.providers.aws.AWSVMManager")
logger.setLevel(logging.INFO)

REGISTRY_PATH = '.aws_vms'

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


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 remains the interface of VMManager for compatibility with other components.
    """
    def __init__(self, registry_path=REGISTRY_PATH):
        self.registry_path = registry_path
        # self.lock = FileLock(".aws_lck", timeout=60)
        self.initialize_registry()

    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):
        logger.info("Allocating a new VM in region: {}".format(region))
        new_vm_path = _allocate_vm(region)
        return new_vm_path