Handle new exception (#213)

NixOS · Jan 27, 2025 · 80b9d65 · 80b9d65
1 parent 15d02f1
commit 80b9d65
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 33 deletions.
diff --git a/upload-ami/src/upload_ami/delete_deprecated_images.py b/upload-ami/src/upload_ami/delete_deprecated_images.py
@@ -33,7 +33,9 @@ def delete_deprecated_images(ec2: EC2Client, dry_run: bool) -> None:
                 if current_time >= image["DeprecationTime"]:
                     assert "ImageId" in image
                     assert "Name" in image
-                    logger.info(f"Deleting image {image['Name']} : {image['ImageId']}. DeprecationTime: {image['DeprecationTime']}")
+                    logger.info(
+                        f"Deleting image {image['Name']} : {image['ImageId']}. DeprecationTime: {image['DeprecationTime']}"
+                    )
                     try:
                         ec2.deregister_image(ImageId=image["ImageId"], DryRun=dry_run)
                     except botocore.exceptions.ClientError as e:

diff --git a/upload-ami/src/upload_ami/smoke_test.py b/upload-ami/src/upload_ami/smoke_test.py
@@ -1,4 +1,5 @@
 import boto3
+import botocore.exceptions
 import time
 import argparse
 import logging
@@ -22,52 +23,54 @@ def smoke_test(image_id: str, run_id: str, cancel: bool) -> None:
         instance_type = "t4g.nano"
     else:
         raise Exception("Unknown architecture: " + architecture)
-
     logging.info("Starting instance")
-    run_instances = ec2.run_instances(
-        ImageId=image_id,
-        InstanceType=instance_type,
-        MinCount=1,
-        MaxCount=1,
-        ClientToken=image_id + run_id if run_id else image_id,
-        InstanceMarketOptions={"MarketType": "spot"},
-    )
+    try:
+        run_instances = ec2.run_instances(
+            ImageId=image_id,
+            InstanceType=instance_type,
+            MinCount=1,
+            MaxCount=1,
+            ClientToken=image_id + run_id if run_id else image_id,
+            InstanceMarketOptions={"MarketType": "spot"},
+        )
+    except botocore.exceptions.ClientError as error:
+        if error.response["Error"]["Code"] == "IdempotentInstanceTerminated":
+            logging.warn(error)
+            return
+        else:
+            raise error
 
     instance = run_instances["Instances"][0]
     assert "InstanceId" in instance
     instance_id = instance["InstanceId"]
 
     try:
-        if not cancel and instance["State"]["Name"] != "terminated":
-            # This basically waits for DHCP to have finished; as it uses ARP to check if the instance is healthy
-            logging.info(f"Waiting for instance {instance_id} to be running")
-            ec2.get_waiter("instance_running").wait(InstanceIds=[instance_id])
-            logging.info(f"Waiting for instance {instance_id} to be healthy")
-            ec2.get_waiter("instance_status_ok").wait(InstanceIds=[instance_id])
-            tries = 5
+        if cancel:
+            return
+        # This basically waits for DHCP to have finished; as it uses ARP to check if the instance is healthy
+        logging.info(f"Waiting for instance {instance_id} to be running")
+        ec2.get_waiter("instance_running").wait(InstanceIds=[instance_id])
+        logging.info(f"Waiting for instance {instance_id} to be healthy")
+        ec2.get_waiter("instance_status_ok").wait(InstanceIds=[instance_id])
+        tries = 5
+        console_output = ec2.get_console_output(InstanceId=instance_id, Latest=True)
+        output = console_output.get("Output")
+        while not output and tries > 0:
+            time.sleep(10)
+            logging.info(
+                f"Waiting for console output to become available ({tries} tries left)"
+            )
             console_output = ec2.get_console_output(InstanceId=instance_id, Latest=True)
             output = console_output.get("Output")
-            while not output and tries > 0:
-                time.sleep(10)
-                logging.info(
-                    f"Waiting for console output to become available ({tries} tries left)"
-                )
-                console_output = ec2.get_console_output(
-                    InstanceId=instance_id, Latest=True
-                )
-                output = console_output.get("Output")
-                tries -= 1
-            logging.info(f"Console output: {output}")
+            tries -= 1
+        logging.info(f"Console output: {output}")
     except Exception as e:
         logging.error(f"Error: {e}")
         raise
     finally:
         logging.info(f"Terminating instance {instance_id}")
-        assert "State" in instance
-        assert "Name" in instance["State"]
-        if instance["State"]["Name"] != "terminated":
-            ec2.terminate_instances(InstanceIds=[instance_id])
-            ec2.get_waiter("instance_terminated").wait(InstanceIds=[instance_id])
+        ec2.terminate_instances(InstanceIds=[instance_id])
+        ec2.get_waiter("instance_terminated").wait(InstanceIds=[instance_id])
 
 
 def main() -> None: