Feat: custom security groups

mkdocs/docs/concepts/backends.md

@@ -242,6 +242,47 @@ There are two ways to configure AWS: using an access key or using the default cr
     Using private subnets assumes that both the `dstack` server and users can access the configured VPC's private subnets.
     Additionally, private subnets must have outbound internet connectivity provided by NAT Gateway, Transit Gateway, or other mechanism.
 
+??? 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_name` if you create a
+    security group with the same name in every configured region's VPC:
+
+    ```yaml
+    projects:
+      - name: main
+        backends:
+          - type: aws
+            creds:
+              type: default
+
+            security_group_name: my-security-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.
+
 ??? info "OS images"
     By default, `dstack` uses its own [AMI](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AMIs.html)
     optimized for `dstack`.
@@ -444,6 +485,34 @@ There are two ways to configure Azure: using a client secret or using the defaul
     Using private subnets assumes that both the `dstack` server and users can access the configured VPC's private subnets.
     Additionally, private subnets must have outbound internet connectivity provided by [NAT Gateway or other mechanism](https://learn.microsoft.com/en-us/azure/nat-gateway/nat-overview).
 
+??? 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).
+    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_names`. The values are plain NSG names
+    within the configured `resource_group` (not full Azure resource IDs), so NSGs in a different
+    resource group cannot be referenced:
+
+    ```yaml
+    projects:
+      - name: main
+        backends:
+          - type: azure
+            creds:
+              type: default
+            regions: [westeurope, eastus]
+            network_security_group_names:
+              westeurope: my-network-security-group-we
+              eastus: my-network-security-group-eus
+    ```
+
+    Locations not covered by `network_security_group_names` 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.
+
 ### GCP
 
 There are two ways to configure GCP: using a service account or using the default credentials.
@@ -662,6 +731,30 @@ gcloud projects list --format="json(projectId)"
     Using private subnets assumes that both the `dstack` server and users can access the configured VPC's private subnets.
     Additionally, [Cloud NAT](https://cloud.google.com/nat/docs/overview) must be configured to provide access to external resources for provisioned instances.
 
+??? info "Custom firewall rules"
+    By default, `dstack` creates a VPC firewall rule allowing inbound SSH from the internet to instances,
+    scoped to the `dstack-runner-instance` target tag.
+    Unlike AWS/Azure/OCI, GCP firewall rules apply to the whole VPC rather than to a single attachable resource,
+    so there's no per-fleet override — if you manage your own firewall rules and don't want `dstack` creating
+    a rule that opens port 22 to `0.0.0.0/0`, disable this at the project level with `create_firewall_rules: false`:
+
+    ```yaml
+    projects:
+    - name: main
+      backends:
+        - type: gcp
+          project_id: gcp-project-id
+          creds:
+            type: default
+
+          create_firewall_rules: false
+    ```
+
+    You're then responsible for ensuring your VPC's own firewall rules allow whatever SSH and cluster traffic
+    `dstack` needs. This setting only affects the instance SSH rule — the separate firewall rule `dstack`
+    creates for gateways (allowing HTTP/HTTPS from the internet, scoped to the `dstack-gateway-instance`
+    target tag) is always auto-managed, since gateways are meant to be internet-reachable.
+
 ### Lambda
 
 Log into your [Lambda Cloud](https://lambdalabs.com/service/gpu-cloud) account, click API keys in the sidebar, and then click the `Generate API key`
@@ -1068,6 +1161,48 @@ There are two ways to configure OCI: using client credentials or using the defau
         compartment_id: ocid1.compartment.oc1..aaaaaaaa
     ```
 
+??? info "Custom network security group"
+    By default, `dstack` places instances in a shared subnet that has **no** OCI security list
+    attached. On top of that, `dstack` creates and manages its own network security group (NSG) per
+    project, which grants everything instances need: SSH ingress (TCP port 22) from `0.0.0.0/0`,
+    unrestricted egress, and — so multi-node clusters work out of the box — all traffic between
+    instances in the group.
+
+    OCI network security groups are region-scoped, so a custom NSG is configured per region via
+    `network_security_group_ids`:
+
+    ```yaml
+    projects:
+    - name: main
+      backends:
+      - type: oci
+        creds:
+          type: default
+        network_security_group_ids:
+          eu-frankfurt-1: ocid1.networksecuritygroup.oc1..aaaaaaaa
+          us-ashburn-1: ocid1.networksecuritygroup.oc1..bbbbbbbb
+    ```
+
+    Regions not covered by `network_security_group_ids` fall back to dstack's auto-created network
+    security group.
+
+    OCI requires a network security group to belong to the same VCN as the instances it's attached
+    to. Since `dstack` provisions instances into its own default VCN per project
+    (`dstack-<project>-default-vcn`), your custom NSG must be created in that VCN.
+
+    When a custom NSG is used, `dstack` never adds, removes, or modifies its rules. Because the shared
+    subnet has no security list, the NSG is the sole security boundary for these instances — there is
+    no implicit SSH-from-anywhere or implicit outbound-all falling back from a security list. As a
+    result, your custom NSG is fully responsible for:
+
+    - **Ingress**, including SSH (TCP port 22) from wherever you connect.
+    - **Egress**, including outbound internet access. Without an egress rule (e.g. allow all to
+      `0.0.0.0/0`), instances will have no outbound connectivity and will fail to pull Docker images
+      and start runs.
+    - For multi-node clusters, **traffic between instances** in the group.
+
+    You can also override this per fleet or run using the `security_group` profile property.
+
 SSH fleets support the same features as [VM-based](#vm-based) backends.
 
 !!! info "What's next"

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

@@ -31,6 +31,7 @@
     ComputeWithPrivateGatewaySupport,
     ComputeWithPrivilegedSupport,
     ComputeWithReservationSupport,
+    ComputeWithSecurityGroupSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
@@ -124,6 +125,7 @@ class AWSCompute(
     ComputeWithGatewaySupport,
     ComputeWithPrivateGatewaySupport,
     ComputeWithVolumeSupport,
+    ComputeWithSecurityGroupSupport,
     Compute,
 ):
     def __init__(
@@ -329,12 +331,27 @@ def create_instance(
                 instance_type=instance_offer.instance.name,
                 image_config=self.config.os_images,
             )
-            security_group_id = self._create_security_group(
-                ec2_client=ec2_client,
-                region=instance_offer.region,
-                project_id=project_name,
-                vpc_id=vpc_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,
+                    region=instance_offer.region,
+                    project_id=project_name,
+                    vpc_id=vpc_id,
+                )
             try:
                 response = ec2_resource.create_instances(  # pyright: ignore[reportAttributeAccessIssue]
                     **aws_resources.create_instances_struct(
@@ -363,9 +380,18 @@ def create_instance(
                 )
             except botocore.exceptions.ClientError as e:
                 logger.warning("Got botocore.exceptions.ClientError: %s", e)
-                if e.response["Error"]["Code"] == "InvalidParameterValue":
+                error_code = e.response["Error"]["Code"]
+                if error_code == "InvalidParameterValue":
                     msg = e.response["Error"].get("Message", "")
                     raise ComputeError(f"Invalid AWS request: {msg}")
+                if error_code == "InvalidGroup.NotFound":
+                    # A misconfigured security group (e.g. wrong VPC/region) is not a
+                    # capacity issue, so surface it clearly instead of retrying other AZs.
+                    msg = e.response["Error"].get("Message", "")
+                    raise ComputeError(
+                        f"Security group not found for instance in region"
+                        f" {instance_offer.region}: {msg}"
+                    )
                 continue
             instance = response[0]
             # wait_until_running() is only needed so that instance is immediately ready for volume attach.

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,19 @@ 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_ids is None:
+            return
+        regions = config.regions if config.regions is not None else DEFAULT_REGIONS
+        unknown_regions = [r for r in config.security_group_ids if r not in regions]
+        if unknown_regions:
+            raise ServerClientError(
+                msg=(
+                    f"`security_group_ids` specifies regions not in `regions`: {unknown_regions}."
+                    " This is likely a typo — remove the extra keys or add them to `regions`"
+                )
+            )
+
     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,6 +96,37 @@ class AWSBackendConfig(CoreModel):
             )
         ),
     ] = None
+    security_group_name: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "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

@@ -45,6 +45,7 @@
     ComputeWithInstanceVolumesSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPrivilegedSupport,
+    ComputeWithSecurityGroupSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     get_gateway_user_data,
@@ -89,6 +90,7 @@ class AzureCompute(
     ComputeWithInstanceVolumesSupport,
     ComputeWithMultinodeSupport,
     ComputeWithGatewaySupport,
+    ComputeWithSecurityGroupSupport,
     Compute,
 ):
     def __init__(self, config: AzureConfig, credential: TokenCredential):
@@ -146,10 +148,14 @@ def create_instance(
             location=location,
             allocate_public_ip=allocate_public_ip,
         )
-        network_security_group = azure_utils.get_default_network_security_group_name(
-            resource_group=self.config.resource_group,
-            location=location,
-        )
+        network_security_group = instance_config.security_group
+        if network_security_group is None and self.config.network_security_group_names is not None:
+            network_security_group = self.config.network_security_group_names.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
@@ -256,7 +262,12 @@ def create_gateway(
             location=configuration.region,
             allocate_public_ip=True,
         )
-        network_security_group = azure_utils.get_default_network_security_group_name(
+        # Gateways always use the dedicated gateway NSG (created unconditionally in
+        # `_create_network_resources`, regardless of `network_security_group_names`),
+        # never the per-location default/custom instance NSG. This keeps gateway
+        # provisioning working even for locations where the default instance NSG is
+        # skipped because a custom one is configured for instances.
+        network_security_group = azure_utils.get_gateway_network_security_group_name(
             resource_group=self.config.resource_group,
             location=configuration.region,
         )