Refactoring VMware Integration and Implementing AWS Support (#44)

* Initailize aws support * Add README for the VM server * Refactor OSWorld for supporting more cloud services. * Initialize vmware and aws implementation v1, waiting for verification * Initlize files for azure, gcp and virtualbox support * Debug on the VMware provider * Fix on aws interface mapping * Fix instance type * Refactor * Clean * hk region; debug * Fix lock * Remove print * Remove key_name requirements when allocating aws vm * Clean README --------- Co-authored-by: XinyuanWangCS <xywang626@gmail.com>
2024-06-15 20:52:29 +08:00
parent c121869219
commit fffa8f8da6
31 changed files with 847 additions and 302 deletions
--- a/desktop_env/providers/aws/provider.py
+++ b/desktop_env/providers/aws/provider.py
@@ -0,0 +1,112 @@
+import boto3
+from botocore.exceptions import ClientError
+
+import logging
+
+from .manager import INSTANCE_TYPE
+from desktop_env.providers.base import Provider
+
+logger = logging.getLogger("desktopenv.providers.aws.AWSProvider")
+logger.setLevel(logging.INFO)
+
+WAIT_DELAY = 15
+MAX_ATTEMPTS = 10
+
+
+class AWSProvider(Provider):
+
+    def start_emulator(self, path_to_vm: str, headless: bool):
+        logger.info("Starting AWS VM...")
+        ec2_client = boto3.client('ec2', region_name=self.region)
+
+        try:
+            # Start the instance
+            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 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', '')
+                    return private_ip_address
+            return ''  # Return an empty string if no IP address is found
+        except ClientError as e:
+            logger.error(f"Failed to retrieve private 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, ImageId=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:
+            # Step 1: Retrieve the original instance details
+            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']
+            iam_instance_profile = instance.get('IamInstanceProfile', {}).get('Arn', '')
+
+            # Step 2: Launch a new instance from the snapshot
+            logger.info(f"Launching a new instance from snapshot {snapshot_name}...")
+            new_instance = ec2_client.run_instances(
+                ImageId=snapshot_name,
+                InstanceType=instance_type,
+                SecurityGroupIds=security_groups,
+                SubnetId=subnet_id,
+                IamInstanceProfile={'Arn': iam_instance_profile} if iam_instance_profile else {},
+                MinCount=1,
+                MaxCount=1
+            )
+            new_instance_id = new_instance['Instances'][0]['InstanceId']
+            logger.info(f"New instance {new_instance_id} launched from snapshot {snapshot_name}.")
+
+            # Step 3: Terminate the old instance
+            ec2_client.terminate_instances(InstanceIds=[path_to_vm])
+            logger.info(f"Old instance {path_to_vm} has been terminated.")
+
+            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