From 30050d4178dd76b86657aa3bce3609e227038f96 Mon Sep 17 00:00:00 2001
From: Tianbao Xie <47296835+Timothyxxx@users.noreply.github.com>
Date: Sat, 15 Jun 2024 22:21:13 +0800
Subject: [PATCH]  Refactoring VMware Integration and Implementing AWS Support
 (#45)

* 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

* Fix reset

* Fix bugs

---------

Co-authored-by: XinyuanWangCS <xywang626@gmail.com>
---
 desktop_env/providers/aws/provider.py | 33 +++++++++++++++++++--------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/desktop_env/providers/aws/provider.py b/desktop_env/providers/aws/provider.py
index 17f6bff..2f0c5cb 100644
--- a/desktop_env/providers/aws/provider.py
+++ b/desktop_env/providers/aws/provider.py
@@ -72,21 +72,34 @@ class AWSProvider(Provider):
             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
-            )
+
+            run_instances_params = {
+                "MaxCount": 1,
+                "MinCount": 1,
+                "ImageId": snapshot_name,
+                "InstanceType": instance_type,
+                "EbsOptimized": True,
+                "NetworkInterfaces": [
+                    {
+                        "SubnetId": subnet_id,
+                        "AssociatePublicIpAddress": True,
+                        "DeviceIndex": 0,
+                        "Groups": security_groups
+                    }
+                ]
+            }
+
+            new_instance = ec2_client.run_instances(**run_instances_params)
             new_instance_id = new_instance['Instances'][0]['InstanceId']
             logger.info(f"New instance {new_instance_id} launched from snapshot {snapshot_name}.")
+            logger.info(f"Waiting for instance {new_instance_id} to be running...")
+            ec2_client.get_waiter('instance_running').wait(InstanceIds=[new_instance_id])
+            logger.info(f"Waiting for instance {new_instance_id} status checks to pass...")
+            ec2_client.get_waiter('instance_status_ok').wait(InstanceIds=[new_instance_id])
+            logger.info(f"Instance {new_instance_id} is ready.")
 
             # Step 3: Terminate the old instance
             ec2_client.terminate_instances(InstanceIds=[path_to_vm])