Update ann-benchmarks.py to connect to running instance (#500)

Author: jparismorgan

src/benchmarks/ann-benchmarks.py

@@ -64,6 +64,13 @@
         f"Key file not found at {key_path}. Please set the correct path before running."
     )
 
+# If you want to connect to a running instance, set these. Note that we could fetch instance_id from
+# the public_dns (via ec2.describe_instances()), but I had IAM permissions issues, so we manually set.
+# Example: "ec2-9-38-120-343.eu-central-1.compute.amazonaws.com"
+connect_to_running_instance_public_dns = ""
+# Example: "i-0790c23df32093234"
+connect_to_running_instance_id = ""
+
 # You do not need to change these.
 security_group_ids = ["sg-04258b401ce76d246"]
 # 64 vCPU, 512 GiB, EBS-Only.
@@ -137,54 +144,62 @@ def execute_commands(ssh, commands):
 
 
 try:
-    # Launch an EC2 instance
-    logger.info("Launching EC2 instance...")
-    response = ec2.run_instances(
-        ImageId=ami_id,
-        InstanceType=instance_type,
-        KeyName=key_name,
-        SecurityGroupIds=security_group_ids,
-        MinCount=1,
-        MaxCount=1,
-        BlockDeviceMappings=[
-            {
-                "DeviceName": "/dev/xvda",
-                "Ebs": {
-                    # Size in GiB.
-                    "VolumeSize": 30,
-                    # General Purpose SSD (gp3).
-                    "VolumeType": "gp3",
+    if connect_to_running_instance_public_dns and connect_to_running_instance_id:
+        # Connect to an existng EC2 instance.
+        public_dns = connect_to_running_instance_public_dns
+        instance_id = connect_to_running_instance_id
+        logger.info(
+            f"Will connect to running instance at public_dns: {public_dns} and instance_id: {instance_id}"
+        )
+    else:
+        # Launch an EC2 instance.
+        logger.info("Launching EC2 instance...")
+        response = ec2.run_instances(
+            ImageId=ami_id,
+            InstanceType=instance_type,
+            KeyName=key_name,
+            SecurityGroupIds=security_group_ids,
+            MinCount=1,
+            MaxCount=1,
+            BlockDeviceMappings=[
+                {
+                    "DeviceName": "/dev/xvda",
+                    "Ebs": {
+                        # Size in GiB.
+                        "VolumeSize": 30,
+                        # General Purpose SSD (gp3).
+                        "VolumeType": "gp3",
+                    },
+                }
+            ],
+        )
+        instance_id = response["Instances"][0]["InstanceId"]
+        logger.info(f"Launched EC2 instance with ID: {instance_id}")
+
+        # Wait for the instance to be in a running state.
+        logger.info("Waiting for instance to enter running state...")
+        waiter = ec2.get_waiter("instance_running")
+        waiter.wait(InstanceIds=[instance_id])
+
+        # Get the public DNS name of the instance.
+        instance_description = ec2.describe_instances(InstanceIds=[instance_id])
+        public_dns = instance_description["Reservations"][0]["Instances"][0][
+            "PublicDnsName"
+        ]
+        logger.info(f"Public DNS of the instance: {public_dns}")
+
+        # Tag the instance.
+        instance_name = f"vector-search-ann-benchmarks-{socket.gethostname()}"
+        logger.info(f"Will name the instance: {instance_name}")
+        ec2.create_tags(
+            Resources=[instance_id],
+            Tags=[
+                {
+                    "Key": "Name",
+                    "Value": instance_name,
                 },
-            }
-        ],
-    )
-    instance_id = response["Instances"][0]["InstanceId"]
-    logger.info(f"Launched EC2 instance with ID: {instance_id}")
-
-    # Wait for the instance to be in a running state.
-    logger.info("Waiting for instance to enter running state...")
-    waiter = ec2.get_waiter("instance_running")
-    waiter.wait(InstanceIds=[instance_id])
-
-    # Get the public DNS name of the instance.
-    instance_description = ec2.describe_instances(InstanceIds=[instance_id])
-    public_dns = instance_description["Reservations"][0]["Instances"][0][
-        "PublicDnsName"
-    ]
-    logger.info(f"Public DNS of the instance: {public_dns}")
-
-    # Tag the instance.
-    instance_name = f"vector-search-ann-benchmarks-{socket.gethostname()}"
-    logger.info(f"Will name the instance: {instance_name}")
-    ec2.create_tags(
-        Resources=[instance_id],
-        Tags=[
-            {
-                "Key": "Name",
-                "Value": instance_name,
-            },
-        ],
-    )
+            ],
+        )
 
     # Wait for SSH to be ready
     if not check_ssh_ready(public_dns=public_dns, key_filename=key_path):
@@ -233,7 +248,8 @@ def execute_commands(ssh, commands):
             )
         for algorithm in algorithms:
             post_reconnect_commands += [
-                f"cd {ann_benchmarks_dir} && python3 run.py --dataset sift-128-euclidean --algorithm {algorithm} --force --batch",
+                # NOTE: If you want to force re-running a benchmark even if the results exist, add --force.
+                f"cd {ann_benchmarks_dir} && python3 run.py --dataset sift-128-euclidean --algorithm {algorithm} --batch",
                 f"cd {ann_benchmarks_dir} && sudo chmod -R 777 results/sift-128-euclidean/10/{algorithm}-batch",
             ]
         post_reconnect_commands.append(