add more delete functionalities

Author: Jyothirmaikottu

test/test_utils/ec2.py

@@ -1829,8 +1829,13 @@ def get_subnet_id_by_vpc(ec2_client, vpc_id):
             },
         ],
     )
-    print(response)
-    return response
+
+    subnet_ids = []
+    for subnet in response["Subnets"]:
+        if subnet["SubnetId"] is not None:
+            subnet_ids.append(subnet["SubnetId"])
+
+    return subnet_ids
 
 
 def get_vpc_id_by_name(ec2_client, vpc_name):

vllm/infra/ec2.py

@@ -222,19 +222,23 @@ def setup():
     ec2_cli = ec2_client(DEFAULT_REGION)
 
     vpc_id = get_default_vpc_id(ec2_cli)
-    subnet_id = get_subnet_id_by_vpc(ec2_cli, vpc_id)
+    subnet_ids = get_subnet_id_by_vpc(ec2_cli, vpc_id)
 
     # create fsx
     sg_fsx = fsx.create_security_group(vpc_id, "vllm-ec2-fsx-sg", "SG for Fsx Mounting")
     ingress_rules = {"protocol": "tcp", "port": "988-1023"}
     fsx.add_security_group_ingress_rules(sg_fsx, ingress_rules)
 
     fsx_config = fsx.create_fsx_filesystem(
-        subnet_id, sg_fsx, 1200, "SCRATCH_2", {"Name": "vllm-fsx-storage"}
+        subnet_ids[0], sg_fsx, 1200, "SCRATCH_2", {"Name": "vllm-fsx-storage"}
     )
 
     print(fsx_config)
 
+    fsx.delete_fsx_filesystem(fsx_config["filesystem_id"])
+
+    fsx.delete_security_group(sg_fsx)
+
 
 if __name__ == "__main__":
     setup()

vllm/infra/utils/fsx_utils.py

@@ -58,14 +58,30 @@ def create_fsx_filesystem(
             logger.error(f"Failed to create FSx filesystem: {e}")
             raise
 
+    def delete_fsx_filesystem(self, fsx_id: str):
+
+        try:
+            fsx_id = run(
+                f"aws fsx delete-file-system"
+                f" --file-system-id {fsx_id}"
+                f' --query "FileSystem.FileSystemId"'
+                f" --output text"
+            ).stdout.strip()
+
+            print(f"Deleted FSx filesystem: {fsx_id}")
+
+        except Exception as e:
+            logger.error(f"Failed to create FSx filesystem: {e}")
+            raise
+
     def wait_for_filesystem(self, filesystem_id: str):
         """
         Wait for FSx filesystem to become available and return its details
         : param filesystem_id: FSx filesystem ID
         : return: dictionary containing filesystem details (filesystem_id, dns_name, mount_name)
         : raises: Exception if filesystem enters FAILED, DELETING, or DELETED state
         """
-        logger.info(f"Waiting for FSx filesystem {filesystem_id} to be available...")
+        print(f"Waiting for FSx filesystem {filesystem_id} to be available...")
         while True:
             status = run(
                 f"aws fsx describe-file-systems --file-system-id {filesystem_id} "
@@ -77,7 +93,7 @@ def wait_for_filesystem(self, filesystem_id: str):
             elif status in ["FAILED", "DELETING", "DELETED"]:
                 raise Exception(f"FSx filesystem entered {status} state")
 
-            logger.info(f"FSx status: {status}, waiting...")
+            print(f"FSx status: {status}, waiting...")
             time.sleep(30)
 
         # get fs DNS and mount name
@@ -111,13 +127,30 @@ def create_security_group(self, vpc_id: str, name: str, description: str):
                 f' --query "GroupId"'
                 f" --output text"
             ).stdout.strip()
-            logger.info(f"Created security group: {sg_id}")
+            print(f"Created security group: {sg_id}")
             return sg_id
 
         except Exception as e:
             logger.error(f"Failed to create security group: {e}")
             raise
 
+    def delete_security_group(self, group_id: str):
+        """
+        Create a security group in the specified VPC
+        : param vpc_id: VPC ID where the security group will be created
+        : param name: name of the security group
+        : param description: description of the security group
+        : return: created security group ID
+        : raises: Exception if security group creation fails
+        """
+        try:
+            sg_id = run(f"aws ec2 delete-security-group" f" --group-id {group_id}").stdout.strip()
+            print(f"Deleted security group: {sg_id}")
+
+        except Exception as e:
+            logger.error(f"Failed to create security group: {e}")
+            raise
+
     def add_security_group_ingress_rules(
         self, security_group_id: str, ingress_rules: List[Dict[str, Any]]
     ):