118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import os
|
|
from filelock import FileLock
|
|
import boto3
|
|
import psutil
|
|
import logging
|
|
import dotenv
|
|
# 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)
|
|
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.")
|
|
|
|
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
|