Support multi-region custom security groups/NSGs

security_group_id / network_security_group / network_security_group_id were
flat single-value fields, but security groups and NSGs are scoped per
VPC/region, so they only worked for single-region backends. Fixed to mirror
the existing vpc_ids/subnet_ids/subnet_ids_per_region conventions:

- AWS: security_group_name (a name that must exist in every region's VPC —
  AWS allows reusing the same group name across regions) and security_group_ids
  (an explicit region -> ID map for when names differ).
- Azure: network_security_group_ids (location -> NSG name map). Azure NSG
  names are unique per resource group regardless of location, so a single
  name can never cover more than one region.
- OCI: network_security_group_ids (region -> NSG OCID map).

Regions/locations not covered by the mapping fall back to dstack's
auto-created security group, so partial custom-SG adoption across regions
works too. GCP is unaffected (its firewall rules are VPC-wide, not
per-region).

Author: james-boydell

mkdocs/docs/concepts/backends.md

@@ -245,7 +245,8 @@ There are two ways to configure AWS: using an access key or using the default cr
 ??? info "Custom security group"
     By default, `dstack` creates and manages its own security group per project (opening SSH to `0.0.0.0/0`
     and allowing all traffic within the group so multi-node clusters work out of the box).
-    To use a security group you manage yourself instead, set `security_group_id`:
+    To use a security group you manage yourself instead, set `security_group_name` if you create a
+    security group with the same name in every configured region's VPC:
 
     ```yaml
     projects:
@@ -255,12 +256,30 @@ There are two ways to configure AWS: using an access key or using the default cr
             creds:
               type: default
 
-            security_group_id: sg-0a1b2c3d4e5f6g7h8
+            security_group_name: my-security-group
     ```
 
-    When `security_group_id` is set, `dstack` attaches it to instances as-is and never adds, removes, or modifies
-    its rules. You're responsible for SSH reachability (from wherever the `dstack` server and users connect from)
-    and, for multi-node clusters, for allowing traffic between instances in the group.
+    If your security groups have different names (or IDs are more convenient) per region, use
+    `security_group_ids` instead:
+
+    ```yaml
+    projects:
+      - name: main
+        backends:
+          - type: aws
+            creds:
+              type: default
+
+            security_group_ids:
+              us-east-1: sg-0a1b2c3d4e5f6g7h8
+              us-west-2: sg-1b2c3d4e5f6g7h8i9
+    ```
+
+    Regions not covered by `security_group_ids` fall back to `security_group_name` if set, or to
+    dstack's auto-created security group otherwise. Either way, `dstack` attaches the security group
+    to instances as-is and never adds, removes, or modifies its rules. You're responsible for SSH
+    reachability (from wherever the `dstack` server and users connect from) and, for multi-node
+    clusters, for allowing traffic between instances in the group.
 
     You can also override this per fleet or run using the `security_group` profile property.
 
@@ -469,7 +488,8 @@ There are two ways to configure Azure: using a client secret or using the defaul
 ??? info "Custom network security group"
     By default, `dstack` creates and manages its own network security group (opening SSH to the internet
     and allowing all traffic within the group so multi-node clusters work out of the box).
-    To use a network security group you manage yourself instead, set `network_security_group`:
+    Azure NSG names must be unique within a resource group regardless of region, so a custom NSG is
+    configured per location via `network_security_group_ids`:
 
     ```yaml
     projects:
@@ -478,12 +498,16 @@ There are two ways to configure Azure: using a client secret or using the defaul
           - type: azure
             creds:
               type: default
-            network_security_group: my-network-security-group
+            regions: [westeurope, eastus]
+            network_security_group_ids:
+              westeurope: my-network-security-group-we
+              eastus: my-network-security-group-eus
     ```
 
-    When `network_security_group` is set, `dstack` attaches it to instances as-is and never adds, removes, or
-    modifies its rules. You're responsible for SSH reachability and, for multi-node clusters, for allowing
-    traffic between instances in the group.
+    Locations not covered by `network_security_group_ids` fall back to dstack's auto-created network
+    security group. Either way, `dstack` attaches the network security group to instances as-is and
+    never adds, removes, or modifies its rules. You're responsible for SSH reachability and, for
+    multi-node clusters, for allowing traffic between instances in the group.
 
     You can also override this per fleet or run using the `security_group` profile property.
 
@@ -1136,7 +1160,8 @@ There are two ways to configure OCI: using client credentials or using the defau
 ??? info "Custom network security group"
     By default, `dstack` creates and manages its own network security group per project (opening SSH to
     `0.0.0.0/0` and allowing all traffic within the VCN so multi-node clusters work out of the box).
-    To use a network security group you manage yourself instead, set `network_security_group_id`:
+    OCI network security groups are region-scoped, so a custom NSG is configured per region via
+    `network_security_group_ids`:
 
     ```yaml
     projects:
@@ -1145,12 +1170,15 @@ There are two ways to configure OCI: using client credentials or using the defau
       - type: oci
         creds:
           type: default
-        network_security_group_id: ocid1.networksecuritygroup.oc1..aaaaaaaa
+        network_security_group_ids:
+          eu-frankfurt-1: ocid1.networksecuritygroup.oc1..aaaaaaaa
+          us-ashburn-1: ocid1.networksecuritygroup.oc1..bbbbbbbb
     ```
 
-    When `network_security_group_id` is set, `dstack` attaches it to instances as-is and never adds, removes,
-    or modifies its rules. You're responsible for SSH reachability and, for multi-node clusters, for allowing
-    traffic between instances in the group.
+    Regions not covered by `network_security_group_ids` fall back to dstack's auto-created network
+    security group. Either way, `dstack` attaches the network security group to instances as-is and
+    never adds, removes, or modifies its rules. You're responsible for SSH reachability and, for
+    multi-node clusters, for allowing traffic between instances in the group.
 
     You can also override this per fleet or run using the `security_group` profile property.
 

src/dstack/_internal/core/backends/aws/compute.py

@@ -331,7 +331,20 @@ def create_instance(
                 instance_type=instance_offer.instance.name,
                 image_config=self.config.os_images,
             )
-            security_group_id = instance_config.security_group or self.config.security_group_id
+            security_group_id = instance_config.security_group
+            if security_group_id is None and self.config.security_group_ids is not None:
+                security_group_id = self.config.security_group_ids.get(instance_offer.region)
+            if security_group_id is None and self.config.security_group_name is not None:
+                security_group_id = aws_resources.get_security_group_id_by_name(
+                    ec2_client=ec2_client,
+                    name=self.config.security_group_name,
+                    vpc_id=vpc_id,
+                )
+                if security_group_id is None:
+                    raise ComputeError(
+                        f"Security group '{self.config.security_group_name}' not found in"
+                        f" VPC {vpc_id} (region {instance_offer.region})"
+                    )
             if security_group_id is None:
                 security_group_id = self._create_security_group(
                     ec2_client=ec2_client,

src/dstack/_internal/core/backends/aws/configurator.py

@@ -77,6 +77,7 @@ def validate_config(self, config: AWSBackendConfigWithCreds, default_creds_enabl
                 raise_invalid_credentials_error(fields=[["creds"]])
         self._check_config_tags(config)
         self._check_config_iam_instance_profile(session, config)
+        self._check_config_security_group(config)
         self._check_config_vpc(session, config)
 
     def create_backend(
@@ -146,6 +147,12 @@ def _check_config_iam_instance_profile(
                 f"Failed to check IAM instance profile {config.iam_instance_profile}"
             )
 
+    def _check_config_security_group(self, config: AWSBackendConfigWithCreds):
+        if config.security_group_name is not None and config.security_group_ids is not None:
+            raise ServerClientError(
+                msg="Only one of `security_group_name` and `security_group_ids` can be specified"
+            )
+
     def _check_config_vpc(self, session: Session, config: AWSBackendConfigWithCreds):
         allocate_public_ip = config.public_ips if config.public_ips is not None else True
         use_default_vpcs = config.default_vpcs if config.default_vpcs is not None else True

src/dstack/_internal/core/backends/aws/models.py

@@ -96,18 +96,37 @@ class AWSBackendConfig(CoreModel):
             )
         ),
     ] = None
-    security_group_id: Annotated[
+    security_group_name: Annotated[
         Optional[str],
         Field(
             description=(
-                "The ID of an existing security group to use for instances instead of the one"
+                "The name of an existing security group to use for instances instead of the one"
                 " `dstack` creates and manages automatically (`dstack_security_group_<project>`)."
+                " The security group must exist in every VPC `dstack` provisions into."
+                " If your custom security groups don't have names or have different names in"
+                " different regions, use `security_group_ids` instead."
                 " When set, `dstack` does not add, remove, or modify any rules on this security group"
                 " — you are responsible for SSH reachability and, for multi-node clusters,"
                 " for allowing traffic between instances in the group"
             )
         ),
     ] = None
+    security_group_ids: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The mapping from AWS regions to the IDs of existing security groups to use for"
+                " instances instead of the one `dstack` creates and manages automatically."
+                " Use this instead of `security_group_name` when your security groups don't have"
+                " names or have different names in different regions."
+                " Regions not present in this mapping fall back to `security_group_name` if set,"
+                " or to dstack's auto-created security group otherwise."
+                " When set, `dstack` does not add, remove, or modify any rules on these security"
+                " groups — you are responsible for SSH reachability and, for multi-node clusters,"
+                " for allowing traffic between instances in the group"
+            )
+        ),
+    ] = None
     tags: Annotated[
         Optional[Dict[str, str]],
         Field(description="The tags that will be assigned to resources created by `dstack`"),

src/dstack/_internal/core/backends/aws/resources.py

@@ -136,6 +136,21 @@ def create_security_group(
     return security_group_id
 
 
+def get_security_group_id_by_name(
+    ec2_client: botocore.client.BaseClient,
+    name: str,
+    vpc_id: Optional[str],
+) -> Optional[str]:
+    filters = [{"Name": "group-name", "Values": [name]}]
+    if vpc_id is not None:
+        filters.append({"Name": "vpc-id", "Values": [vpc_id]})
+    response = ec2_client.describe_security_groups(Filters=filters)
+    groups = response.get("SecurityGroups")
+    if not groups:
+        return None
+    return groups[0]["GroupId"]
+
+
 def create_instances_struct(
     disk_size: int,
     image_id: str,

src/dstack/_internal/core/backends/azure/compute.py

@@ -148,14 +148,14 @@ def create_instance(
             location=location,
             allocate_public_ip=allocate_public_ip,
         )
-        network_security_group = (
-            instance_config.security_group
-            or self.config.network_security_group
-            or azure_utils.get_default_network_security_group_name(
+        network_security_group = instance_config.security_group
+        if network_security_group is None and self.config.network_security_group_ids is not None:
+            network_security_group = self.config.network_security_group_ids.get(location)
+        if network_security_group is None:
+            network_security_group = azure_utils.get_default_network_security_group_name(
                 resource_group=self.config.resource_group,
                 location=location,
             )
-        )
 
         managed_identity_resource_group, managed_identity_name = parse_vm_managed_identity(
             self.config.vm_managed_identity

src/dstack/_internal/core/backends/azure/configurator.py

@@ -1,6 +1,6 @@
 import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import azure.core.exceptions
 from azure.core.credentials import TokenCredential
@@ -126,7 +126,7 @@ def create_backend(
             resource_group=config.resource_group,
             locations=config.regions,
             create_default_network=config.vpc_ids is None and config.subnet_ids is None,
-            create_instance_network_security_group=config.network_security_group is None,
+            network_security_group_ids=config.network_security_group_ids,
         )
         return BackendRecord(
             config=AzureStoredConfig(
@@ -341,7 +341,7 @@ def _create_network_resources(
         resource_group: str,
         locations: List[str],
         create_default_network: bool,
-        create_instance_network_security_group: bool = True,
+        network_security_group_ids: Optional[Dict[str, str]] = None,
     ):
         def func(location: str):
             network_manager = NetworkManager(
@@ -354,9 +354,10 @@ def func(location: str):
                     name=azure_utils.get_default_network_name(resource_group, location),
                     subnet_name=azure_utils.get_default_subnet_name(resource_group, location),
                 )
-            if create_instance_network_security_group:
-                # Skipped when the user supplies their own network security group via
-                # `network_security_group` - dstack does not create or manage it in that case.
+            if location not in (network_security_group_ids or {}):
+                # Skipped when the user supplies their own network security group for this
+                # location via `network_security_group_ids` - dstack does not create or manage
+                # it in that case.
                 network_manager.create_network_security_group(
                     resource_group=resource_group,
                     location=location,

src/dstack/_internal/core/backends/azure/models.py

@@ -81,14 +81,17 @@ class AzureBackendConfig(CoreModel):
             )
         ),
     ] = None
-    network_security_group: Annotated[
-        Optional[str],
+    network_security_group_ids: Annotated[
+        Optional[Dict[str, str]],
         Field(
             description=(
-                "The name of an existing network security group (in the configured resource group)"
-                " to use for instances instead of the one `dstack` creates and manages automatically."
-                " When set, `dstack` does not add, remove, or modify any rules on this network"
-                " security group — you are responsible for SSH reachability and, for multi-node"
+                "The mapping from Azure locations to the names of existing network security groups"
+                " (in the configured resource group) to use for instances instead of the one `dstack`"
+                " creates and manages automatically."
+                " Locations not present in this mapping fall back to dstack's auto-created"
+                " network security group."
+                " When set, `dstack` does not add, remove, or modify any rules on these network"
+                " security groups — you are responsible for SSH reachability and, for multi-node"
                 " clusters, for allowing traffic between instances in the group"
             )
         ),

src/dstack/_internal/core/backends/oci/compute.py

@@ -138,9 +138,11 @@ def create_instance(
         subnet: oci.core.models.Subnet = region.virtual_network_client.get_subnet(
             self.config.subnet_ids_per_region[instance_offer.region]
         ).data
-        security_group_id = (
-            instance_config.security_group or self.config.network_security_group_id
-        )
+        security_group_id = instance_config.security_group
+        if security_group_id is None and self.config.network_security_group_ids is not None:
+            security_group_id = self.config.network_security_group_ids.get(
+                instance_offer.region
+            )
         if security_group_id is None:
             security_group = resources.get_or_create_security_group(
                 f"dstack-{instance_config.project_name}-default-security-group",

src/dstack/_internal/core/backends/oci/models.py

@@ -69,14 +69,16 @@ class OCIBackendConfig(CoreModel):
             )
         ),
     ] = None
-    network_security_group_id: Annotated[
-        Optional[str],
+    network_security_group_ids: Annotated[
+        Optional[Dict[str, str]],
         Field(
             description=(
-                "The OCID of an existing network security group to use for instances instead of the"
-                " one `dstack` creates and manages automatically."
-                " When set, `dstack` does not add, remove, or modify any rules on this network"
-                " security group — you are responsible for SSH reachability and, for multi-node"
+                "The mapping from OCI regions to the OCIDs of existing network security groups to"
+                " use for instances instead of the one `dstack` creates and manages automatically."
+                " Regions not present in this mapping fall back to dstack's auto-created network"
+                " security group."
+                " When set, `dstack` does not add, remove, or modify any rules on these network"
+                " security groups — you are responsible for SSH reachability and, for multi-node"
                 " clusters, for allowing traffic between instances in the group"
             )
         ),