Retry AWS SSI across zones

docs/understand/scenarios/onboarding.md

@@ -375,7 +375,9 @@ All system-tests assertions and utilities are based on python and pytests. You n
 
 Before execute the "onboarding" tests you must configure some environment variables:
 
-- **ONBOARDING_AWS_INFRA_SUBNET_ID:** AWS subnet id.
+- **ONBOARDING_AWS_INFRA_SUBNET_ID:** One or more comma-separated AWS subnet IDs. Configure subnets from
+  different availability zones so EC2 launches can move to another zone when the selected zone has insufficient
+  instance capacity.
 - **ONBOARDING_AWS_INFRA_SECURITY_GROUPS_ID:** AWS security groups id.
 - **DD_API_KEY_ONBOARDING:** Datadog API key.
 - **DD_APP_KEY_ONBOARDING:** Datadog APP key.
@@ -1311,4 +1313,3 @@ A failure in the GitLab runners can propagate across all repositories that rely
 1. If the problem persists, report it in the #ci-infra-support channel so the CI infrastructure team can investigate and assist.
 2. The runners used for SSI tests in system-tests can be found at: https://github.com/DataDog/libdatadog-build/tree/main/docker
 
-

tests/test_the_test/test_aws_provider.py

@@ -0,0 +1,101 @@
+# ruff: noqa: SLF001
+
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+
+from utils import features, scenarios
+from utils.virtual_machine import aws_provider
+from utils.virtual_machine.virtual_machines import AWSInfraConfig
+
+
+class _CommandError(Exception):
+    pass
+
+
+class _Stack:
+    def __init__(self, outcomes: list[Exception | None]) -> None:
+        self.outcomes = outcomes
+        self.up_calls = 0
+
+    def up(self, *, on_output: Callable[[str], Any]) -> None:
+        del on_output
+        outcome = self.outcomes[self.up_calls]
+        self.up_calls += 1
+        if outcome is not None:
+            raise outcome
+
+
+@features.not_reported
+@scenarios.test_the_test
+class Test_AWSProvider:
+    def test_infra_config_normalizes_comma_separated_ids(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("ONBOARDING_AWS_INFRA_SUBNET_ID", " subnet-a,subnet-b, ,subnet-c ")
+        monkeypatch.setenv("ONBOARDING_AWS_INFRA_SECURITY_GROUPS_ID", " sg-a, sg-b ")
+
+        config = AWSInfraConfig()
+
+        assert config.subnet_id == ["subnet-a", "subnet-b", "subnet-c"]
+        assert config.vpc_security_group_ids == ["sg-a", "sg-b"]
+
+    def test_capacity_failure_rotates_through_subnets(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(aws_provider.pulumi.automation.errors, "CommandError", _CommandError)
+        provider = aws_provider.AWSPulumiProvider()
+        provider._subnet_ids = ["subnet-a", "subnet-b", "subnet-c"]
+        provider.stack = _Stack(
+            [
+                _CommandError("InsufficientInstanceCapacity"),
+                _CommandError("InsufficientInstanceCapacity"),
+                None,
+            ]
+        )
+        destroyed_subnets: list[str] = []
+        monkeypatch.setattr(
+            provider,
+            "stack_destroy",
+            lambda: destroyed_subnets.append(provider._subnet_ids[provider._subnet_index]),
+        )
+
+        provider._stack_up_with_transient_retry()
+
+        assert provider.stack.up_calls == 3
+        assert provider._subnet_index == 2
+        assert destroyed_subnets == ["subnet-b", "subnet-c"]
+
+    def test_capacity_failure_is_raised_after_all_subnets(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(aws_provider.pulumi.automation.errors, "CommandError", _CommandError)
+        provider = aws_provider.AWSPulumiProvider()
+        provider._subnet_ids = ["subnet-a", "subnet-b"]
+        provider.stack = _Stack(
+            [
+                _CommandError("InsufficientInstanceCapacity"),
+                _CommandError("InsufficientInstanceCapacity"),
+            ]
+        )
+        monkeypatch.setattr(provider, "stack_destroy", lambda: None)
+
+        with pytest.raises(_CommandError, match="InsufficientInstanceCapacity"):
+            provider._stack_up_with_transient_retry()
+
+        assert provider.stack.up_calls == 2
+        assert provider._subnet_index == 1
+
+    def test_idempotency_failure_retries_without_rotating_subnet(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(aws_provider.pulumi.automation.errors, "CommandError", _CommandError)
+        provider = aws_provider.AWSPulumiProvider()
+        provider._subnet_ids = ["subnet-a", "subnet-b"]
+        provider.stack = _Stack([_CommandError("IdempotentParameterMismatch"), None])
+        destroy_calls = 0
+
+        def record_destroy() -> None:
+            nonlocal destroy_calls
+            destroy_calls += 1
+
+        monkeypatch.setattr(provider, "stack_destroy", record_destroy)
+
+        provider._stack_up_with_transient_retry()
+
+        assert provider.stack.up_calls == 2
+        assert provider._subnet_index == 0
+        assert destroy_calls == 1

utils/virtual_machine/aws_infra_exceptions.json

@@ -1,5 +1,6 @@
 {
     "IdempotentParameterMismatch": "creating EC2 Instance: IdempotentParameterMismatch",
+    "InsufficientInstanceCapacity": "creating EC2 Instance: InsufficientInstanceCapacity",
     "InternalError": "last error: Server.InternalError: Internal error on launch",
     "error_ssh_connection": "error: after 30 failed attempts:",
     "error_pulumi_plugin": "error: could not read plugin",

utils/virtual_machine/aws_provider.py

@@ -45,9 +45,14 @@ def __init__(self):
         self.pulumi_ssh = None
         self.datadog_event_sender = DatadogEventSender()
         self.stack_name = "system-tests_dev_onboarding"
+        self._subnet_ids: list[str] = []
+        self._subnet_index = 0
 
     def configure(self, virtual_machine: _VirtualMachine) -> None:
         super().configure(virtual_machine)
+        self._subnet_ids = virtual_machine.aws_config.aws_infra_config.subnet_id.copy()
+        random.shuffle(self._subnet_ids)
+        self._subnet_index = 0
         # Configure the ssh connection for the VMs
         self.pulumi_ssh = PulumiSSH()
         self.pulumi_ssh.load(virtual_machine)
@@ -82,7 +87,7 @@ def pulumi_start_program() -> None:
             )
             if os.getenv("ONBOARDING_LOCAL_TEST") is None:
                 self.stack.set_config("aws:SkipMetadataApiCheck", auto.ConfigValue("false"))
-            self._stack_up_with_idempotency_retry()
+            self._stack_up_with_transient_retry()
             self.datadog_event_sender.send_event_to_datadog(
                 f"[E2E] Stack {self.stack_name}  : success on Pulumi stack up",
                 "",
@@ -117,21 +122,32 @@ def pulumi_start_program() -> None:
             )
             self._handle_provision_error(pulumi_exception)
 
-    def _stack_up_with_idempotency_retry(self, attempts: int = 3) -> None:
-        """Retry the stack up in-process on IdempotentParameterMismatch.
+    def _stack_up_with_transient_retry(self, idempotency_attempts: int = 3) -> None:
+        """Retry transient EC2 launch failures without restarting the CI job.
 
-        This is a transient AWS-side token collision (see _start_vm's ec2_resource_id comment).
-        A fresh attempt gets a new idempotency token, so it's worth retrying here rather than
-        immediately failing the whole CI job (which is a much more expensive way to retry).
+        IdempotentParameterMismatch gets a fresh idempotency token. InsufficientInstanceCapacity
+        moves the launch to the next configured subnet, allowing AWS to use another availability
+        zone. Each subnet is attempted at most once for a capacity failure.
         """
-        for attempt in range(1, attempts + 1):
+        idempotency_retries = idempotency_attempts - 1
+        subnet_retries = len(self._subnet_ids) - 1
+
+        while True:
             try:
                 self.stack.up(on_output=logger.info)
                 return
             except pulumi.automation.errors.CommandError as pulumi_command_exception:
-                if "IdempotentParameterMismatch" not in str(pulumi_command_exception) or attempt == attempts:
+                exception_message = str(pulumi_command_exception)
+                if "IdempotentParameterMismatch" in exception_message and idempotency_retries > 0:
+                    idempotency_retries -= 1
+                    retry_reason = "IdempotentParameterMismatch"
+                elif "InsufficientInstanceCapacity" in exception_message and subnet_retries > 0:
+                    subnet_retries -= 1
+                    self._subnet_index += 1
+                    retry_reason = "InsufficientInstanceCapacity in the selected availability zone"
+                else:
                     raise
-                logger.stdout(f"⚠️ IdempotentParameterMismatch on attempt {attempt}/{attempts}, retrying stack up ⚠️")
+                logger.stdout(f"⚠️ {retry_reason}, retrying stack up ⚠️")
                 self.stack_destroy()
 
     def get_windows_user_data(self) -> str:
@@ -198,6 +214,8 @@ def _start_vm(self, vm: _VirtualMachine):
         logger.info(
             f"Starting VM: {vm.name} with iam_instance_profile: {vm.aws_config.aws_infra_config.iam_instance_profile}"
         )
+        if not self._subnet_ids:
+            raise ValueError("ONBOARDING_AWS_INFRA_SUBNET_ID must contain at least one subnet ID")
         # Startup VM and prepare connection
         # The resource name (not the "Name" tag) must be unique per CI job: several parallel jobs
         # (one per weblog) can provision the same vm.name at the same time, and a shared resource
@@ -208,7 +226,7 @@ def _start_vm(self, vm: _VirtualMachine):
             ec2_resource_id,
             instance_type=vm.aws_config.ami_instance_type,
             vpc_security_group_ids=vm.aws_config.aws_infra_config.vpc_security_group_ids,
-            subnet_id=random.choice(vm.aws_config.aws_infra_config.subnet_id),
+            subnet_id=self._subnet_ids[self._subnet_index],
             key_name=self.pulumi_ssh.keypair_name,
             ami=vm.aws_config.ami_id,
             tags=self._get_ec2_tags(vm),

utils/virtual_machine/virtual_machines.py

@@ -14,8 +14,16 @@
 class AWSInfraConfig:
     def __init__(self) -> None:
         # Mandatory parameters
-        self.subnet_id = os.getenv("ONBOARDING_AWS_INFRA_SUBNET_ID", "").split(",")
-        self.vpc_security_group_ids = os.getenv("ONBOARDING_AWS_INFRA_SECURITY_GROUPS_ID", "").split(",")
+        self.subnet_id = [
+            subnet_id.strip()
+            for subnet_id in os.getenv("ONBOARDING_AWS_INFRA_SUBNET_ID", "").split(",")
+            if subnet_id.strip()
+        ]
+        self.vpc_security_group_ids = [
+            security_group_id.strip()
+            for security_group_id in os.getenv("ONBOARDING_AWS_INFRA_SECURITY_GROUPS_ID", "").split(",")
+            if security_group_id.strip()
+        ]
         self.iam_instance_profile = os.getenv("ONBOARDING_AWS_INFRA_IAM_INSTANCE_PROFILE")