Merge pull request #198 from yuanmengqi/main

aws_communication_success_log

.gitignore

@@ -194,3 +194,5 @@ vmware_vm_data
 
 # result
 **/result*/**/*
+
+.vscode

aws/README.md

@@ -0,0 +1,54 @@
+# AWS CLI v2
+
+This bundle contains a built executable of the AWS CLI v2.
+
+## Installation
+
+To install the AWS CLI v2, run the `install` script:
+```
+$ sudo ./install 
+You can now run: /usr/local/bin/aws --version
+```
+This will install the AWS CLI v2 at `/usr/local/bin/aws`.  Assuming
+`/usr/local/bin` is on your `PATH`, you can now run:
+```
+$ aws --version
+```
+
+
+### Installing without sudo
+
+If you don't have ``sudo`` permissions or want to install the AWS
+CLI v2 only for the current user, run the `install` script with the `-b`
+and `-i` options:
+```
+$ ./install -i ~/.local/aws-cli -b ~/.local/bin
+``` 
+This will install the AWS CLI v2 in `~/.local/aws-cli` and create
+symlinks for `aws` and `aws_completer` in `~/.local/bin`. For more
+information about these options, run the `install` script with `-h`:
+```
+$ ./install -h
+```
+
+### Updating
+
+If you run the `install` script and there is a previously installed version
+of the AWS CLI v2, the script will error out. To update to the version included
+in this bundle, run the `install` script with `--update`:
+```
+$ sudo ./install --update
+```
+
+
+### Removing the installation
+
+To remove the AWS CLI v2, delete the its installation and symlinks:
+```
+$ sudo rm -rf /usr/local/aws-cli
+$ sudo rm /usr/local/bin/aws
+$ sudo rm /usr/local/bin/aws_completer
+```
+Note if you installed the AWS CLI v2 using the `-b` or `-i` options, you will
+need to remove the installation and the symlinks in the directories you
+specified.

aws/install

@@ -0,0 +1,155 @@
+#!/bin/sh
+# Copyright 2012-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+usage() {
+  cat 1>&2 <<EOF
+Installs the AWS CLI v2
+
+USAGE:
+    install [FLAGS] [OPTIONS]
+
+FLAGS:
+    -u, --update              Updates the AWS CLI v2 if a different version
+                              is previously installed. By default, this script
+                              will not update the AWS CLI if a previous
+                              installation is detected.
+
+    -h, --help                Prints help information
+
+OPTIONS:
+    -i, --install-dir <path>  The directory to install the AWS CLI v2. By
+                              default, this directory is: /usr/local/aws-cli
+
+    -b, --bin-dir <path>      The directory to store symlinks to executables
+                              for the AWS CLI v2. By default, the directory
+                              used is: /usr/local/bin
+EOF
+}
+
+parse_commandline() {
+  while test $# -gt 0
+  do
+    key="$1"
+	case "$key" in
+	  -i|--install-dir)
+	    PARSED_INSTALL_DIR="$2"
+		shift
+	   ;;
+	  -b|--bin-dir)
+	    PARSED_BIN_DIR="$2"
+		shift
+	   ;;
+	  -u|--update)
+	    PARSED_UPGRADE="yes"
+	  ;;
+	  -h|--help)
+	    usage
+        exit 0
+	  ;;
+	  *)
+	   die "Got an unexpected argument: $1"
+	  ;;
+    esac
+	shift
+  done
+}
+
+set_global_vars() {
+  ROOT_INSTALL_DIR=${PARSED_INSTALL_DIR:-/usr/local/aws-cli}
+  BIN_DIR=${PARSED_BIN_DIR:-/usr/local/bin}
+  UPGRADE=${PARSED_UPGRADE:-no}
+
+  EXE_NAME="aws"
+  COMPLETER_EXE_NAME="aws_completer"
+  INSTALLER_DIR="$( cd "$( dirname "$0" )" >/dev/null 2>&1 && pwd )"
+  INSTALLER_DIST_DIR="$INSTALLER_DIR/dist"
+  INSTALLER_EXE="$INSTALLER_DIST_DIR/$EXE_NAME"
+  AWS_EXE_VERSION=$($INSTALLER_EXE --version | cut -d ' ' -f 1 | cut -d '/' -f 2)
+
+  INSTALL_DIR="$ROOT_INSTALL_DIR/v2/$AWS_EXE_VERSION"
+  INSTALL_DIR="$INSTALL_DIR"
+  INSTALL_DIST_DIR="$INSTALL_DIR/dist"
+  INSTALL_BIN_DIR="$INSTALL_DIR/bin"
+  INSTALL_AWS_EXE="$INSTALL_BIN_DIR/$EXE_NAME"
+  INSTALL_AWS_COMPLETER_EXE="$INSTALL_BIN_DIR/$COMPLETER_EXE_NAME"
+
+  CURRENT_INSTALL_DIR="$ROOT_INSTALL_DIR/v2/current"
+  CURRENT_AWS_EXE="$CURRENT_INSTALL_DIR/bin/$EXE_NAME"
+  CURRENT_AWS_COMPLETER_EXE="$CURRENT_INSTALL_DIR/bin/$COMPLETER_EXE_NAME"
+
+  BIN_AWS_EXE="$BIN_DIR/$EXE_NAME"
+  BIN_AWS_COMPLETER_EXE="$BIN_DIR/$COMPLETER_EXE_NAME"
+}
+
+create_install_dir() {
+  mkdir -p "$INSTALL_DIR" || exit 1
+  {
+    setup_install_dist &&
+    setup_install_bin &&
+    create_current_symlink
+  } || {
+    rm -rf "$INSTALL_DIR"
+    exit 1
+  }
+}
+
+check_preexisting_install() {
+  if [ -L "$CURRENT_INSTALL_DIR" ] && [ "$UPGRADE" = "no" ]
+  then
+    die "Found preexisting AWS CLI installation: $CURRENT_INSTALL_DIR. Please rerun install script with --update flag."
+  fi
+  if [ -d "$INSTALL_DIR" ]
+  then
+    echo "Found same AWS CLI version: $INSTALL_DIR. Skipping install."
+    exit 0
+  fi
+}
+
+setup_install_dist() {
+  cp -r "$INSTALLER_DIST_DIR" "$INSTALL_DIST_DIR"
+}
+
+setup_install_bin() {
+  mkdir -p "$INSTALL_BIN_DIR"
+  ln -s "../dist/$EXE_NAME" "$INSTALL_AWS_EXE"
+  ln -s "../dist/$COMPLETER_EXE_NAME" "$INSTALL_AWS_COMPLETER_EXE"
+}
+
+create_current_symlink() {
+  ln -snf "$INSTALL_DIR" "$CURRENT_INSTALL_DIR"
+}
+
+create_bin_symlinks() {
+  mkdir -p "$BIN_DIR"
+  ln -sf "$CURRENT_AWS_EXE" "$BIN_AWS_EXE"
+  ln -sf "$CURRENT_AWS_COMPLETER_EXE" "$BIN_AWS_COMPLETER_EXE"
+}
+
+die() {
+	err_msg="$1"
+	echo "$err_msg" >&2
+	exit 1
+}
+
+main() {
+  parse_commandline "$@"
+  set_global_vars
+  check_preexisting_install
+  create_install_dir
+  create_bin_symlinks
+  echo "You can now run: $BIN_AWS_EXE --version"
+  exit 0
+}
+
+main "$@" || exit 1

desktop_env/desktop_env.py

@@ -23,10 +23,10 @@ class DesktopEnv(gym.Env):
     """
     DesktopEnv with OpenAI Gym interface. It provides a desktop environment for setting and evaluating desktop automation tasks.
     """
-
+    #TODO:provider_name: str = "vmware",
     def __init__(
             self,
-            provider_name: str = "vmware",
+            provider_name: str = "aws",
             region: str = None,
             path_to_vm: str = None,
             snapshot_name: str = "init_state",
@@ -55,7 +55,7 @@ class DesktopEnv(gym.Env):
         # Initialize VM manager and vitualization provider
         self.region = region
 
-        # Default
+        # Default TODO:
         self.server_port = 5000
         self.chromium_port = 9222
         self.vnc_port = 8006
@@ -69,6 +69,7 @@ class DesktopEnv(gym.Env):
             self.path_to_vm = os.path.abspath(os.path.expandvars(os.path.expanduser(path_to_vm))) \
                 if provider_name in {"vmware", "virtualbox"} else path_to_vm
         else:
+
             self.path_to_vm = self.manager.get_vm_path(os_type=self.os_type, region=region)
 
         self.snapshot_name = snapshot_name
@@ -114,7 +115,8 @@ class DesktopEnv(gym.Env):
         # due to the fact it could be changed when implemented by cloud services
         path_to_vm = self.provider.revert_to_snapshot(self.path_to_vm, self.snapshot_name)
         if path_to_vm and not path_to_vm == self.path_to_vm:
-            # path_to_vm has to be a new path
+            # path_to_vm has to be a new path 
+            
             self.manager.delete_vm(self.path_to_vm, self.region)
             self.manager.add_vm(path_to_vm, self.region)
             self.manager.occupy_vm(path_to_vm, os.getpid(), self.region)

desktop_env/providers/aws/manager.py

@@ -21,14 +21,15 @@ IMAGE_ID_MAP = {
 
 INSTANCE_TYPE = "t3.medium"
 
+# sg-0342574803206ee9c   subnet-037edfff66c2eb894
 NETWORK_INTERFACE_MAP = {
     "us-east-1": [
         {
-            "SubnetId": "subnet-037edfff66c2eb894",
+            "SubnetId": "subnet-0a4b0c5b8f6066712",
             "AssociatePublicIpAddress": True,
             "DeviceIndex": 0,
             "Groups": [
-                "sg-0342574803206ee9c"
+                "sg-08a53433e9b4abde6"
             ]
         }
     ],
@@ -240,19 +241,20 @@ class AWSVMManager(VMManager):
                 AWSVMManager.checked_and_cleaned = True
                 self._check_and_clean()
 
-        allocation_needed = False
+        # allocation_needed = False
-        with self.lock:
+        # with self.lock:
-            free_vms_paths = self._list_free_vms(region)
+        #     free_vms_paths = self._list_free_vms(region)
 
-            if len(free_vms_paths) == 0:
+        #     if len(free_vms_paths) == 0:
-                # No free virtual machine available, generate a new one
+        #         # No free virtual machine available, generate a new one
-                allocation_needed = True
+        #         allocation_needed = True
-            else:
+        #     else:
-                # Choose the first free virtual machine
+        #         # Choose the first free virtual machine
-                chosen_vm_path = free_vms_paths[0][0]
+        #         chosen_vm_path = free_vms_paths[0][0]
-                self._occupy_vm(chosen_vm_path, os.getpid(), region)
+        #         self._occupy_vm(chosen_vm_path, os.getpid(), region)
-                return chosen_vm_path
+        #         return chosen_vm_path
         
+        allocation_needed = True
         if allocation_needed:
             logger.info("No free virtual machine available. Generating a new one, which would take a while...☕")
             new_vm_path = _allocate_vm(region)

desktop_env/providers/aws/provider.py

@@ -4,6 +4,9 @@ from botocore.exceptions import ClientError
 import logging
 
 from desktop_env.providers.base import Provider
+from datetime import datetime
+import time
+
 
 logger = logging.getLogger("desktopenv.providers.aws.AWSProvider")
 logger.setLevel(logging.INFO)
@@ -14,24 +17,43 @@ MAX_ATTEMPTS = 10
 
 class AWSProvider(Provider):
 
+
     def start_emulator(self, path_to_vm: str, headless: bool, *args, **kwargs):
         logger.info("Starting AWS VM...")
         ec2_client = boto3.client('ec2', region_name=self.region)
 
         try:
-            # Start the instance
+            # Check the current state of the instance
-            ec2_client.start_instances(InstanceIds=[path_to_vm])
+            response = ec2_client.describe_instances(InstanceIds=[path_to_vm])
-            logger.info(f"Instance {path_to_vm} is starting...")
+            state = response['Reservations'][0]['Instances'][0]['State']['Name']
+            logger.info(f"Instance {path_to_vm} current state: {state}")
 
-            # Wait for the instance to be in the 'running' state
+            if state == 'running':
-            waiter = ec2_client.get_waiter('instance_running')
+                # If the instance is already running, skip starting it
-            waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS})
+                logger.info(f"Instance {path_to_vm} is already running. Skipping start.")
-            logger.info(f"Instance {path_to_vm} is now running.")
+                return
+
+            if state == 'stopped':
+                # Start the instance if it's currently stopped
+                ec2_client.start_instances(InstanceIds=[path_to_vm])
+                logger.info(f"Instance {path_to_vm} is starting...")
+
+                # Wait until the instance reaches '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.")
+            else:
+                # For all other states (terminated, pending, etc.), log a warning
+                logger.warning(f"Instance {path_to_vm} is in state '{state}' and cannot be started.")
 
         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)
@@ -71,21 +93,46 @@ class AWSProvider(Provider):
             security_groups = [sg['GroupId'] for sg in instance['SecurityGroups']]
             subnet_id = instance['SubnetId']
             instance_type = instance['InstanceType']
+            instance_snapshot = instance_details['Reservations'][0]['Instances'][0]['ImageId']
 
-            # Step 2: Terminate the old instance
+            # TODO:Step 2: Terminate the old instance
             ec2_client.terminate_instances(InstanceIds=[path_to_vm])
             logger.info(f"Old instance {path_to_vm} has been terminated.")
 
             # Step 3: Launch a new instance from the snapshot
-            logger.info(f"Launching a new instance from snapshot {snapshot_name}...")
+            logger.info(f"Launching a new instance from snapshot {instance_snapshot}...")
 
-            run_instances_params = {
+            # run_instances_params = {
-                "MaxCount": 1,
+            #     "MaxCount": 1,
-                "MinCount": 1,
+            #     "MinCount": 1,
-                "ImageId": snapshot_name,
+            #     "ImageId": instance_snapshot, 
-                "InstanceType": instance_type,
+            #     "InstanceType": instance_type,
-                "EbsOptimized": True,
+            #     "EbsOptimized": True,
-                "NetworkInterfaces": [
+            #     "NetworkInterfaces": [
+            #         {
+            #             "SubnetId": subnet_id,
+            #             "AssociatePublicIpAddress": True,
+            #             "DeviceIndex": 0,
+            #             "Groups": security_groups
+            #         }
+            #     ],
+            #     "BlockDeviceMappings":[
+            #         {
+            #             "Ebs": {
+            #                 "VolumeSize": 30,
+            #                 "VolumeType": "gp3"
+            #             },
+            #         },
+            #     ],
+            # }
+
+            new_instance = ec2_client.run_instances(
+                MaxCount = 1,
+                MinCount = 1,
+                ImageId = instance_snapshot, 
+                InstanceType = instance_type,
+                EbsOptimized = True,
+                NetworkInterfaces = [
                     {
                         "SubnetId": subnet_id,
                         "AssociatePublicIpAddress": True,
@@ -93,13 +140,13 @@ class AWSProvider(Provider):
                         "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])
+            # wait 60 seconds for the instance to be ready
+            time.sleep(60)
             logger.info(f"Instance {new_instance_id} is ready.")
 
             return new_instance_id
@@ -108,6 +155,121 @@ class AWSProvider(Provider):
             logger.error(f"Failed to revert to snapshot {snapshot_name} for the instance {path_to_vm}: {str(e)}")
             raise
 
+
+        # # 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']
+        # #TODO:instance_type = instance['InstanceType']
+        # instance_type = 't3.large'
+        # instance_snapshot = instance_details['Reservations'][0]['Instances'][0]['ImageId']
+
+        # # TODO:Step 2: Terminate the old instance
+        # if not path_to_vm == 'i-00017dfb534d22011':
+        #     ec2_client.terminate_instances(InstanceIds=[path_to_vm])
+        # logger.info(f"Old instance {path_to_vm} has been terminated.")
+
+        # # Step 3: Launch a new instance from the snapshot
+        # logger.info(f"Launching a new instance from snapshot {instance_snapshot}...")
+
+        # timestamp_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        # instance_name = "/dev/sda1"
+
+        # new_instance = ec2_client.run_instances(
+        #     BlockDeviceMappings = [
+        #         { 
+        #             "Ebs": {
+        #                 "VolumeSize": 30,
+        #                 "VolumeType": "gp3"
+        #             },
+        #             'DeviceName':instance_name,
+        #         },
+        #     ],
+        #     MaxCount = 1,
+        #     MinCount = 1,
+        #     ImageId = instance_snapshot, 
+        #     InstanceType = instance_type,
+        #     EbsOptimized = True,
+        #     NetworkInterfaces = [
+        #         {
+        #             "AssociatePublicIpAddress": True,
+        #             "DeviceIndex": 0,
+        #             "Groups": security_groups
+        #         }
+        #     ]
+        # )
+        # '''NetworkInterfaces = [
+        #         {
+        #             "SubnetId": subnet_id,
+        #             "AssociatePublicIpAddress": True,
+        #             "DeviceIndex": 0,
+        #             "Groups": security_groups
+        #         }
+        #     ]'''
+        # 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"Instance {new_instance_id} is ready.")
+
+        # # # Step 4: set inbound rules
+        # # # TODO: get host sg automatically
+        # # host = ec2_client.describe_instances(InstanceIds=['i-027eab0d007b62793'])
+        # # host_sg_id = host['Reservations'][0]['Instances'][0]['SecurityGroups'][0]['GroupId']
+        # # vm_sg_id = new_instance['Instances'][0]['SecurityGroups'][0]['GroupId']
+
+        # # # add inbound rules to the host security group
+        # # try:
+        # #     host.authorize_security_group_ingress(
+        # #         GroupId= host_sg_id,
+        # #         IpPermissions=[
+        # #             {
+        # #                 "IpProtocol": "tcp",
+        # #                 "FromPort": 5000,
+        # #                 "ToPort": 5000,
+        # #                 "UserIdGroupPairs": [
+        # #                     {
+        # #                         "GroupId": vm_sg_id
+        # #                     }
+        # #                 ]
+        # #             }
+        # #         ]
+        # #     )
+        # #     print(f"Port 5000 opened on {host_sg_id} for {vm_sg_id}")
+        # # except ClientError as e:
+        # #     if "InvalidPermission.Duplicate" in str(e):
+        # #         print(f"Rule already exists on {host_sg_id}")
+        # #     else:
+        # #         print(f"Error updating {host_sg_id}: {e}")
+
+        # # # add inbound rules to the new instance security group
+        # # try:
+        # #     new_instance.authorize_security_group_ingress(
+        # #         GroupId= new_instance_id,
+        # #         IpPermissions=[
+        # #             {
+        # #                 "IpProtocol": "tcp",
+        # #                 "FromPort": 6000,
+        # #                 "ToPort": 6000,
+        # #                 "UserIdGroupPairs": [
+        # #                     {
+        # #                         "GroupId": host_sg_id
+        # #                     }
+        # #                 ]
+        # #             }
+        # #         ]
+        # #     )
+        # #     print(f"Port 6000 opened on {new_instance_id} for {host_sg_id}")
+        # # except ClientError as e:
+        # #     if "InvalidPermission.Duplicate" in str(e):
+        # #         print(f"Rule already exists on {new_instance_id}")
+        # #     else:
+        # #         print(f"Error updating {new_instance_id}: {e}")
+
+        # return new_instance_id
+
+
     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)

run.sh

@@ -0,0 +1,7 @@
+python run_multienv.py \
+--headless \
+--observation_type screenshot \
+--model gpt-4-turbo \
+--result_dir ./results_aws \
+--test_all_meta_path evaluation_examples/test_small.json \
+--region us-east-1

run_multienv.py

@@ -107,6 +107,11 @@ def config() -> argparse.Namespace:
     # logging related
     parser.add_argument("--result_dir", type=str, default="./results")
     parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to run in parallel")
+
+    # aws config
+    parser.add_argument(
+        "--region", type=str, default="us-east-1", help="AWS region for the VM"
+    )
     
     args = parser.parse_args()
     return args
@@ -335,6 +340,7 @@ def get_result(action_space, use_model, observation_type, result_dir, total_file
 if __name__ == "__main__":
     ####### The complete version of the list of examples #######
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
     args = config()
 
     with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
@@ -363,3 +369,6 @@ if __name__ == "__main__":
         test_all_meta,
     )
     test(args, test_file_list)
+
+
+# path_to_vm can be a list["xxx","xxx"]