Repository: skypilot-org/skypilot
Pull request: [SKY-1304] AWS throws `InvalidGroup.Duplicate` when user submits a lot of parallel launch
<!-- Describe the changes in this PR -->
Fix SKY 1304
Fixes #4578
Root cause:
- [Get_security_groups_from_vpc_ids](get_security_groups_from_vpc_ids) can't query security groups when security groups is under creating.

Solution:
- When `InvalidGroup.Duplicate` exception occurs during security groups creating, query security again and return the created else raise created but not found exception.

Test:
```bash
D 01-17 17:49:33 config.py:88] Creating or updating security groups...
W 01-17 17:49:34 config.py:593] sky-sg-hong-4292 already exists when creating.
I 01-17 17:49:37 config.py:598] Found existing security group sky-sg-hong-4292 [id=sg-0f255f9fecda37b1f]
I 01-17 17:49:37 config.py:103] Security groups created or updated in 3.44550 seconds.
```


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`


Problem statement:
[AWS] AWS throws `InvalidGroup.Duplicate` when user submits a lot of parallel launch
To reproduce:
`seq 10 | xargs -P 16 -n 1 bash -c 'sky launch -c test-$0 --cpus 2 --use-spot --down -i0 -y &'`

Observed this with error
```
ot error(s) in AWS:[sky.exceptions.NoClusterLaunchedError] Failed to create security group. Error: botocore.exceptions.ClientError: An error occurred (InvalidGroup.Duplicate) when calling the CreateSecurityGroup operation: The security group 'sky-sg-root-xxx' already exists for VPC 'vpc-xxxxxxx'
```

Hints from issue discussion:
Does this need to be a new VPC that you haven't used skypilot on before?
I don't think it is a new VPC. The issue happens for only some of the launches. This issue happens for some of the tests with `pytest tests/test_smoke.py --aws`.

Commit message: wip

Patch to review:
diff --git a/sky/provision/aws/config.py b/sky/provision/aws/config.py
index ffa87c3a011..d07d364c56d 100644
--- a/sky/provision/aws/config.py
+++ b/sky/provision/aws/config.py
@@ -543,19 +543,41 @@ def _configure_security_group(ec2, vpc_id: str, expected_sg_name: str,
     return sg_ids
 
 
-def _get_or_create_vpc_security_group(ec2, vpc_id: str,
-                                      expected_sg_name: str) -> Any:
-    # Figure out which security groups with this name exist for each VPC...
+def _get_security_groups(ec2, vpc_id: str, group_name: str) -> Optional[Any]:
+    """Get security group by VPC ID and group name."""
     vpc_to_existing_sg = {
         sg.vpc_id: sg for sg in _get_security_groups_from_vpc_ids(
             ec2,
             [vpc_id],
-            [expected_sg_name],
+            [group_name],
         )
     }
+    return vpc_to_existing_sg.get(vpc_id)
+
+
+def _get_or_create_vpc_security_group(ec2, vpc_id: str,
+                                      expected_sg_name: str) -> Any:
+    """Find or create a security group in the specified VPC.
+
+    Args:
+        ec2: The initialized EC2 client object.
+        vpc_id: The ID of the VPC where the security group should be queried
+            or created.
+        expected_sg_name: The expected name of the security group.
 
-    if vpc_id in vpc_to_existing_sg:
-        return vpc_to_existing_sg[vpc_id]
+    Returns:
+        The security group object containing the details of the security group.
+
+    Raises:
+        exceptions.NoClusterLaunchedError: If the security group creation fails
+            and is not due to an existing duplicate.
+        botocore.exceptions.ClientError: If the security group creation fails
+            due to AWS service issues.
+    """
+    # Figure out which security groups with this name exist for each VPC...
+    security_group = _get_security_groups(ec2, vpc_id, expected_sg_name)
+    if security_group is not None:
+        return security_group
 
     try:
         # create a new security group
@@ -565,17 +587,26 @@ def _get_or_create_vpc_security_group(ec2, vpc_id: str,
             VpcId=vpc_id,
         )
     except ec2.meta.client.exceptions.ClientError as e:
+        if e.response['Error']['Code'] == 'InvalidGroup.Duplicate':
+            # The security group already exists, but we didn't see it
+            # because of eventual consistency.
+            logger.warning(f'{expected_sg_name} already exists when creating.')
+            security_group = _get_security_groups(ec2, vpc_id, expected_sg_name)
+            assert (security_group is not None and
+                    security_group.group_name == expected_sg_name), \
+                f'Expected {expected_sg_name} but got {security_group}'
+            logger.info(
+                f'Found existing security group {colorama.Style.BRIGHT}'
+                f'{security_group.group_name}{colorama.Style.RESET_ALL} '
+                f'[id={security_group.id}]')
+            return security_group
         message = ('Failed to create security group. Error: '
                    f'{common_utils.format_exception(e)}')
         logger.warning(message)
         raise exceptions.NoClusterLaunchedError(message) from e
 
-    security_group = _get_security_groups_from_vpc_ids(ec2, [vpc_id],
-                                                       [expected_sg_name])
-
-    assert security_group, 'Failed to create security group'
-    security_group = security_group[0]
-
+    security_group = _get_security_groups(ec2, vpc_id, expected_sg_name)
+    assert security_group is not None, 'Failed to create security group'
     logger.info(f'Created new security group {colorama.Style.BRIGHT}'
                 f'{security_group.group_name}{colorama.Style.RESET_ALL} '
                 f'[id={security_group.id}]')


Code review comments:
diff_hunk: @@ -543,19 +543,41 @@ def _configure_security_group(ec2, vpc_id: str, expected_sg_name: str,
     return sg_ids
 
 
-def _get_or_create_vpc_security_group(ec2, vpc_id: str,
-                                      expected_sg_name: str) -> Any:
-    # Figure out which security groups with this name exist for each VPC...
+def _get_security_groups(ec2, vpc_id: str, group_name: str) -> Optional[Any]:
path: sky/provision/aws/config.py
text: @user1:
```suggestion
def _get_security_group_from_vpc_id(ec2, vpc_id: str, group_name: str) -> Optional[Any]:
```