From 5caf27483ec8688a8d43ea8776effaaeaff1d405 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 10:52:40 -0400 Subject: [PATCH 1/6] Support custom security groups / NSGs / firewall rules in AWS, Azure, GCP, OCI backends Adds a way to point dstack at a pre-existing, user-managed network security resource instead of the one dstack creates and manages automatically (which always opens SSH to 0.0.0.0/0). Useful for private-network setups (e.g. a VPC only reachable via a Tailscale subnet router). - AWS: security_group_id (backend config) - Azure: network_security_group (backend config) - OCI: network_security_group_id (backend config) - GCP: create_firewall_rules: false (backend config; GCP firewall rules are VPC-wide, not an attachable per-instance resource, so there is no fleet-level override for GCP) - AWS/Azure/OCI also support a fleet/run-level override via the security_group profile property (mirrors the existing reservation field), which takes precedence over the backend config default. When a custom security group is configured, dstack attaches it as-is and never adds, removes, or modifies rules on it. Users are responsible for SSH reachability and, for multi-node clusters, for allowing traffic between instances in the group. Gateway security groups/NSGs (AWS, Azure, OCI) are unaffected by this change and keep dstack's existing auto-managed behavior. --- mkdocs/docs/concepts/backends.md | 86 ++++++++++++ .../_internal/core/backends/aws/compute.py | 16 ++- .../_internal/core/backends/aws/models.py | 12 ++ .../_internal/core/backends/azure/compute.py | 12 +- .../core/backends/azure/configurator.py | 17 ++- .../_internal/core/backends/azure/models.py | 12 ++ .../_internal/core/backends/base/compute.py | 19 +++ .../_internal/core/backends/features.py | 5 + .../_internal/core/backends/gcp/compute.py | 4 +- .../_internal/core/backends/gcp/models.py | 11 ++ .../_internal/core/backends/oci/compute.py | 25 ++-- .../_internal/core/backends/oci/models.py | 12 ++ src/dstack/_internal/core/models/fleets.py | 15 ++ src/dstack/_internal/core/models/instances.py | 1 + src/dstack/_internal/core/models/profiles.py | 15 ++ .../_internal/server/services/fleets.py | 2 + .../_internal/server/services/instances.py | 2 + .../core/backends/aws/test_compute.py | 123 ++++++++++++++++ .../core/backends/aws/test_resources.py | 82 +++++++++++ .../core/backends/azure/test_compute.py | 96 ++++++++++++- .../core/backends/azure/test_configurator.py | 33 +++++ .../core/backends/gcp/test_compute.py | 90 ++++++++++++ .../core/backends/oci/test_compute.py | 132 ++++++++++++++++++ 23 files changed, 795 insertions(+), 27 deletions(-) create mode 100644 src/tests/_internal/core/backends/aws/test_compute.py create mode 100644 src/tests/_internal/core/backends/gcp/test_compute.py create mode 100644 src/tests/_internal/core/backends/oci/test_compute.py diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 92d0f72e94..02b0d9c909 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -242,6 +242,28 @@ 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_id`: + + ```yaml + projects: + - name: main + backends: + - type: aws + creds: + type: default + + security_group_id: sg-0a1b2c3d4e5f6g7h8 + ``` + + 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. + + 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 +466,27 @@ 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). + To use a network security group you manage yourself instead, set `network_security_group`: + + ```yaml + projects: + - name: main + backends: + - type: azure + creds: + type: default + network_security_group: my-network-security-group + ``` + + 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. + + 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 +705,28 @@ 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 VPC firewall rules allowing inbound SSH (and, for gateways, HTTP/HTTPS) from + the internet, scoped to the `dstack-runner-instance` and `dstack-gateway-instance` target tags. + 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 + rules that open ports 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. + ### 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 +1133,27 @@ 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` 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`: + + ```yaml + projects: + - name: main + backends: + - type: oci + creds: + type: default + network_security_group_id: ocid1.networksecuritygroup.oc1..aaaaaaaa + ``` + + 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. + + 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" diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index f2420ecf16..267b886a2c 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/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,14 @@ 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 or self.config.security_group_id + 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( diff --git a/src/dstack/_internal/core/backends/aws/models.py b/src/dstack/_internal/core/backends/aws/models.py index d76a4c9ab3..242efce045 100644 --- a/src/dstack/_internal/core/backends/aws/models.py +++ b/src/dstack/_internal/core/backends/aws/models.py @@ -96,6 +96,18 @@ class AWSBackendConfig(CoreModel): ) ), ] = None + security_group_id: Annotated[ + Optional[str], + Field( + description=( + "The ID of an existing security group to use for instances instead of the one" + " `dstack` creates and manages automatically (`dstack_security_group_`)." + " 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 tags: Annotated[ Optional[Dict[str, str]], Field(description="The tags that will be assigned to resources created by `dstack`"), diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index d2843c7f22..be2e065946 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/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,9 +148,13 @@ 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 + or self.config.network_security_group + or 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( diff --git a/src/dstack/_internal/core/backends/azure/configurator.py b/src/dstack/_internal/core/backends/azure/configurator.py index 31297e4ce3..1b8f76522d 100644 --- a/src/dstack/_internal/core/backends/azure/configurator.py +++ b/src/dstack/_internal/core/backends/azure/configurator.py @@ -126,6 +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, ) return BackendRecord( config=AzureStoredConfig( @@ -340,6 +341,7 @@ def _create_network_resources( resource_group: str, locations: List[str], create_default_network: bool, + create_instance_network_security_group: bool = True, ): def func(location: str): network_manager = NetworkManager( @@ -352,11 +354,16 @@ 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), ) - network_manager.create_network_security_group( - resource_group=resource_group, - location=location, - name=azure_utils.get_default_network_security_group_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. + network_manager.create_network_security_group( + resource_group=resource_group, + location=location, + name=azure_utils.get_default_network_security_group_name( + resource_group, location + ), + ) network_manager.create_gateway_network_security_group( resource_group=resource_group, location=location, diff --git a/src/dstack/_internal/core/backends/azure/models.py b/src/dstack/_internal/core/backends/azure/models.py index 9881edcaea..c883616714 100644 --- a/src/dstack/_internal/core/backends/azure/models.py +++ b/src/dstack/_internal/core/backends/azure/models.py @@ -81,6 +81,18 @@ class AzureBackendConfig(CoreModel): ) ), ] = None + network_security_group: Annotated[ + Optional[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" + " 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`"), diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 7cb7c84b35..806d92605b 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -321,6 +321,7 @@ def run_job( ssh_keys=[SSHKey(public=project_ssh_public_key.strip())], volumes=volumes, reservation=job.job_spec.requirements.reservation, + security_group=run.run_spec.merged_profile.security_group, tags=run.run_spec.merged_profile.tags, ) instance_offer = instance_offer.copy() @@ -413,6 +414,24 @@ class ComputeWithReservationSupport: pass +class ComputeWithSecurityGroupSupport: + """ + Must be subclassed to support provisioning instances into a custom, user-managed + security group (or the backend's equivalent: Azure network security group, + OCI network security group). + + The following is expected from a backend that supports this: + + - `create_instance` respects `InstanceConfiguration.security_group` if set, and + attaches the given security group to the instance instead of creating and + managing dstack's own. + - The backend must not add, remove, or modify any rules on a user-supplied + security group. It is fully user-owned. + """ + + pass + + class ComputeWithPlacementGroupSupport(ABC): """ Must be subclassed and implemented to support placement groups. diff --git a/src/dstack/_internal/core/backends/features.py b/src/dstack/_internal/core/backends/features.py index d5c28728dc..3ecd266101 100644 --- a/src/dstack/_internal/core/backends/features.py +++ b/src/dstack/_internal/core/backends/features.py @@ -8,6 +8,7 @@ ComputeWithPrivateGatewaySupport, ComputeWithPrivilegedSupport, ComputeWithReservationSupport, + ComputeWithSecurityGroupSupport, ComputeWithVolumeSupport, ) from dstack._internal.core.backends.base.configurator import Configurator @@ -61,6 +62,10 @@ def _get_backends_with_compute_feature( configurator_classes=_configurator_classes, compute_feature_class=ComputeWithReservationSupport, ) +BACKENDS_WITH_SECURITY_GROUP_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithSecurityGroupSupport, +) BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature( configurator_classes=_configurator_classes, compute_feature_class=ComputeWithGatewaySupport, diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 3be1255d8e..b416b7791d 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -262,7 +262,7 @@ def create_instance( if len(zones) == 0: raise NoCapacityError("No eligible availability zones") # If a shared VPC is not used, we can create firewall rules for user - if self.config.vpc_project_id is None: + if self.config.vpc_project_id is None and self.config.create_firewall_rules is not False: gcp_resources.create_runner_firewall_rules( firewalls_client=self.firewalls_client, project_id=self.config.project_id, @@ -561,7 +561,7 @@ def create_gateway( self, configuration: GatewayComputeConfiguration, ) -> GatewayProvisioningData: - if self.config.vpc_project_id is None: + if self.config.vpc_project_id is None and self.config.create_firewall_rules is not False: gcp_resources.create_gateway_firewall_rules( firewalls_client=self.firewalls_client, project_id=self.config.project_id, diff --git a/src/dstack/_internal/core/backends/gcp/models.py b/src/dstack/_internal/core/backends/gcp/models.py index 4d06144ee8..d5477aca6e 100644 --- a/src/dstack/_internal/core/backends/gcp/models.py +++ b/src/dstack/_internal/core/backends/gcp/models.py @@ -80,6 +80,17 @@ class GCPBackendConfig(CoreModel): ) ), ] = None + create_firewall_rules: Annotated[ + Optional[bool], + Field( + description=( + "A flag to enable/disable `dstack` creating VPC firewall rules that allow SSH" + " (and, for gateways, HTTP/HTTPS) traffic from the internet." + " Set to `false` if you manage your own firewall rules and don't want `dstack`" + " creating rules that open ports to `0.0.0.0/0`. Defaults to `true`" + ) + ), + ] = None vm_service_account: Annotated[ Optional[str], Field(description="The service account to associate with provisioned VMs") ] = None diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index ceaffcf030..8241612b40 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -12,6 +12,7 @@ ComputeWithInstanceVolumesSupport, ComputeWithMultinodeSupport, ComputeWithPrivilegedSupport, + ComputeWithSecurityGroupSupport, generate_unique_instance_name, get_user_data, ) @@ -58,6 +59,7 @@ class OCICompute( ComputeWithPrivilegedSupport, ComputeWithInstanceVolumesSupport, ComputeWithMultinodeSupport, + ComputeWithSecurityGroupSupport, Compute, ): def __init__(self, config: OCIConfig): @@ -136,15 +138,20 @@ 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 = resources.get_or_create_security_group( - f"dstack-{instance_config.project_name}-default-security-group", - subnet.vcn_id, - self.config.compartment_id, - region.virtual_network_client, - ) - resources.update_security_group_rules_for_runner_instances( - security_group.id, region.virtual_network_client + security_group_id = ( + instance_config.security_group or self.config.network_security_group_id ) + if security_group_id is None: + security_group = resources.get_or_create_security_group( + f"dstack-{instance_config.project_name}-default-security-group", + subnet.vcn_id, + self.config.compartment_id, + region.virtual_network_client, + ) + resources.update_security_group_rules_for_runner_instances( + security_group.id, region.virtual_network_client + ) + security_group_id = security_group.id cloud_init_user_data = get_user_data( authorized_keys=instance_config.get_public_keys(), @@ -158,7 +165,7 @@ def create_instance( availability_domain=availability_domain, compartment_id=self.config.compartment_id, subnet_id=subnet.id, - security_group_id=security_group.id, + security_group_id=security_group_id, display_name=display_name, cloud_init_user_data=cloud_init_user_data, shape=instance_offer.instance.name, diff --git a/src/dstack/_internal/core/backends/oci/models.py b/src/dstack/_internal/core/backends/oci/models.py index 12ce4e9f91..d6dd01deec 100644 --- a/src/dstack/_internal/core/backends/oci/models.py +++ b/src/dstack/_internal/core/backends/oci/models.py @@ -69,6 +69,18 @@ class OCIBackendConfig(CoreModel): ) ), ] = None + network_security_group_id: Annotated[ + Optional[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" + " clusters, for allowing traffic between instances in the group" + ) + ), + ] = None class OCIBackendConfigWithCreds(OCIBackendConfig): diff --git a/src/dstack/_internal/core/models/fleets.py b/src/dstack/_internal/core/models/fleets.py index ce636ba8de..0b03823afa 100644 --- a/src/dstack/_internal/core/models/fleets.py +++ b/src/dstack/_internal/core/models/fleets.py @@ -240,6 +240,21 @@ class BackendFleetConfiguraionProps(CoreModel): ) ), ] = None + security_group: Annotated[ + Optional[str], + Field( + description=( + "The existing security group to use for instance provisioning instead of the one" + " `dstack` creates and manages automatically." + " Supported on AWS (security group ID), Azure (network security group name)," + " and OCI (network security group ID). Not supported on GCP;" + " use the GCP backend's `create_firewall_rules: false` setting 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 resources: Annotated[ Optional[ResourcesSpec], Field(description="The resources requirements"), diff --git a/src/dstack/_internal/core/models/instances.py b/src/dstack/_internal/core/models/instances.py index dfce209c32..f60a86aec7 100644 --- a/src/dstack/_internal/core/models/instances.py +++ b/src/dstack/_internal/core/models/instances.py @@ -156,6 +156,7 @@ class InstanceConfiguration(CoreModel): ssh_keys: List[SSHKey] instance_id: Optional[str] = None reservation: Optional[str] = None + security_group: Optional[str] = None volumes: Optional[List[Volume]] = None tags: Optional[Dict[str, str]] = None diff --git a/src/dstack/_internal/core/models/profiles.py b/src/dstack/_internal/core/models/profiles.py index 7a448486df..9f9fca98d9 100644 --- a/src/dstack/_internal/core/models/profiles.py +++ b/src/dstack/_internal/core/models/profiles.py @@ -339,6 +339,21 @@ class ProfileParams(CoreModel): ) ), ] = None + security_group: Annotated[ + Optional[str], + Field( + description=( + "The existing security group to use for instance provisioning instead of the one" + " `dstack` creates and manages automatically." + " Supported on AWS (security group ID), Azure (network security group name)," + " and OCI (network security group ID). Not supported on GCP;" + " use the GCP backend's `create_firewall_rules: false` setting 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 spot_policy: Annotated[ Optional[SpotPolicy], Field( diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 547f91d52c..a1f15e01dc 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -683,6 +683,7 @@ def create_fleet_instance_model( instance_num=instance_num, instance_id=instance_id, reservation=spec.merged_profile.reservation, + security_group=spec.merged_profile.security_group, blocks=spec.configuration.blocks, tags=spec.configuration.tags, ) @@ -1263,6 +1264,7 @@ def _check_can_update_fleet_configuration(current: FleetConfiguration, new: Flee ( "nodes", "reservation", + "security_group", "tags", "resources", "backends", diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index 913d3c9f44..b361785c7f 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -879,6 +879,7 @@ def create_instance_model( instance_name: str, instance_num: int, reservation: Optional[str], + security_group: Optional[str], blocks: Union[Literal["auto"], int], tags: Optional[Dict[str, str]], instance_id: Optional[uuid.UUID] = None, @@ -899,6 +900,7 @@ def create_instance_model( ssh_keys=[project_ssh_key], instance_id=str(instance_id), reservation=reservation, + security_group=security_group, tags=tags, ) now = common_utils.get_current_datetime() diff --git a/src/tests/_internal/core/backends/aws/test_compute.py b/src/tests/_internal/core/backends/aws/test_compute.py new file mode 100644 index 0000000000..332e9249a4 --- /dev/null +++ b/src/tests/_internal/core/backends/aws/test_compute.py @@ -0,0 +1,123 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from dstack._internal.core.backends.aws.compute import AWSCompute +from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds, AWSConfig +from dstack._internal.core.backends.base.compute import ComputeWithSecurityGroupSupport +from dstack._internal.core.backends.features import BACKENDS_WITH_SECURITY_GROUP_SUPPORT +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceConfiguration, + InstanceOfferWithAvailability, + InstanceType, + Resources, + SSHKey, +) + + +def _config(security_group_id=None) -> AWSConfig: + return AWSConfig( + creds=AWSAccessKeyCreds(access_key="test", secret_key="test"), + regions=["us-east-1"], + security_group_id=security_group_id, + ) + + +def _compute(config: AWSConfig) -> AWSCompute: + compute = AWSCompute(config) + compute.session = MagicMock() + # Bypass everything that would hit AWS before/after the security group resolution. + compute._get_maximum_efa_interfaces = MagicMock(return_value=0) + compute._get_vpc_id_subnets_ids_or_error = MagicMock(return_value=("vpc-1", ["subnet-1"])) + compute._get_subnets_availability_zones = MagicMock(return_value={"subnet-1": "az-1"}) + compute._get_image_id_and_username = MagicMock(return_value=("ami-1", "ubuntu")) + compute._create_security_group = MagicMock(return_value="sg-auto") + return compute + + +def _offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.AWS, + instance=InstanceType( + name="m5.large", + resources=Resources(cpus=2, memory_mib=8192, gpus=[], spot=False), + ), + region="us-east-1", + price=0.1, + availability=InstanceAvailability.AVAILABLE, + ) + + +def _instance_config(security_group=None) -> InstanceConfiguration: + return InstanceConfiguration( + project_name="main", + instance_name="test-instance", + user="test-user", + ssh_keys=[SSHKey(public="ssh-rsa test")], + security_group=security_group, + ) + + +def _run_create_instance(compute: AWSCompute, instance_config: InstanceConfiguration) -> str: + """Runs create_instance with the instances struct mocked and returns the + security_group_id that was passed into create_instances_struct.""" + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources.create_instances_struct" + ) as struct_mock: + struct_mock.return_value = {} + ec2_resource = compute.session.resource.return_value + instance_mock = MagicMock() + instance_mock.instance_id = "i-123" + instance_mock.capacity_reservation_id = None + ec2_resource.create_instances.return_value = [instance_mock] + compute.create_instance( + instance_offer=_offer(), + instance_config=instance_config, + placement_group=None, + ) + assert struct_mock.call_count == 1 + return struct_mock.call_args.kwargs["security_group_id"] + + +class TestAWSComputeSecurityGroupSupport: + def test_registered_as_security_group_backend(self): + assert issubclass(AWSCompute, ComputeWithSecurityGroupSupport) + assert BackendType.AWS in BACKENDS_WITH_SECURITY_GROUP_SUPPORT + + def test_auto_creates_group_when_no_custom_sg_configured(self): + compute = _compute(_config()) + security_group_id = _run_create_instance(compute, _instance_config()) + # dstack's auto-create-and-manage path is used. + compute._create_security_group.assert_called_once() + assert security_group_id == "sg-auto" + + def test_uses_project_level_security_group_id_without_managing_it(self): + compute = _compute(_config(security_group_id="sg-project")) + security_group_id = _run_create_instance(compute, _instance_config()) + # No auto-create / rule-management happens for a custom SG. + compute._create_security_group.assert_not_called() + assert security_group_id == "sg-project" + + def test_run_level_security_group_overrides_project_level(self): + compute = _compute(_config(security_group_id="sg-project")) + security_group_id = _run_create_instance( + compute, _instance_config(security_group="sg-run") + ) + compute._create_security_group.assert_not_called() + assert security_group_id == "sg-run" + + @pytest.mark.parametrize( + ["config_sg", "instance_sg", "expected"], + [ + [None, None, "sg-auto"], + ["sg-project", None, "sg-project"], + [None, "sg-run", "sg-run"], + ["sg-project", "sg-run", "sg-run"], + ], + ) + def test_precedence(self, config_sg, instance_sg, expected): + compute = _compute(_config(security_group_id=config_sg)) + security_group_id = _run_create_instance(compute, _instance_config(instance_sg)) + assert security_group_id == expected diff --git a/src/tests/_internal/core/backends/aws/test_resources.py b/src/tests/_internal/core/backends/aws/test_resources.py index a719294cce..0a5536587c 100644 --- a/src/tests/_internal/core/backends/aws/test_resources.py +++ b/src/tests/_internal/core/backends/aws/test_resources.py @@ -9,6 +9,7 @@ _is_valid_tag_key, _is_valid_tag_value, create_instances_struct, + create_security_group, get_image_id_and_username, validate_tags, ) @@ -276,6 +277,87 @@ def test_tenancy_merged_with_placement_group(self): assert struct["Placement"]["Tenancy"] == "dedicated" +class TestCreateSecurityGroup: + @pytest.fixture + def ec2_client_mock(self) -> Mock: + mock = Mock( + spec_set=[ + "describe_security_groups", + "create_security_group", + "authorize_security_group_ingress", + "authorize_security_group_egress", + ] + ) + # No existing group -> a new one is created and all rules are added. + mock.describe_security_groups.return_value = {"SecurityGroups": []} + mock.create_security_group.return_value = {"GroupId": "sg-new"} + return mock + + def test_creates_and_manages_group_when_missing(self, ec2_client_mock: Mock): + security_group_id = create_security_group( + ec2_client=ec2_client_mock, + project_id="my-project", + vpc_id="vpc-1", + ) + assert security_group_id == "sg-new" + ec2_client_mock.create_security_group.assert_called_once() + create_kwargs = ec2_client_mock.create_security_group.call_args.kwargs + assert create_kwargs["GroupName"] == "dstack_security_group_my_project" + assert create_kwargs["VpcId"] == "vpc-1" + # dstack manages the group: SSH (22) ingress + intra-group ingress are authorized, + # plus egress rules. + assert ec2_client_mock.authorize_security_group_ingress.call_count == 2 + assert ec2_client_mock.authorize_security_group_egress.call_count == 2 + ingress_rules = [ + call.kwargs["IpPermissions"][0] + for call in ec2_client_mock.authorize_security_group_ingress.call_args_list + ] + assert { + "FromPort": 22, + "ToPort": 22, + "IpProtocol": "tcp", + "IpRanges": [{"CidrIp": "0.0.0.0/0"}], + } in ingress_rules + + def test_reuses_existing_group_without_recreating(self, ec2_client_mock: Mock): + # An existing group with all managed rules already present -> no create/authorize calls. + ec2_client_mock.describe_security_groups.return_value = { + "SecurityGroups": [ + { + "GroupId": "sg-existing", + "IpPermissions": [ + { + "FromPort": 22, + "ToPort": 22, + "IpProtocol": "tcp", + "IpRanges": [{"CidrIp": "0.0.0.0/0"}], + }, + { + "IpProtocol": "-1", + "UserIdGroupPairs": [{"GroupId": "sg-existing"}], + }, + ], + "IpPermissionsEgress": [ + {"IpProtocol": "-1"}, + { + "IpProtocol": "-1", + "UserIdGroupPairs": [{"GroupId": "sg-existing"}], + }, + ], + } + ] + } + security_group_id = create_security_group( + ec2_client=ec2_client_mock, + project_id="my-project", + vpc_id="vpc-1", + ) + assert security_group_id == "sg-existing" + ec2_client_mock.create_security_group.assert_not_called() + ec2_client_mock.authorize_security_group_ingress.assert_not_called() + ec2_client_mock.authorize_security_group_egress.assert_not_called() + + class TestCreateNetworkInterfacesStruct: def test_non_efa_instance_single_interface(self): interfaces = _create_network_interfaces_struct( diff --git a/src/tests/_internal/core/backends/azure/test_compute.py b/src/tests/_internal/core/backends/azure/test_compute.py index a0ce0afaef..50ea8759f7 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/src/tests/_internal/core/backends/azure/test_compute.py @@ -1,8 +1,20 @@ +from unittest.mock import Mock, patch + import pytest from dstack._internal import settings -from dstack._internal.core.backends.azure.compute import VMImageVariant -from dstack._internal.core.models.instances import Gpu, InstanceType, Resources +from dstack._internal.core.backends.azure import utils as azure_utils +from dstack._internal.core.backends.azure.compute import AzureCompute, VMImageVariant +from dstack._internal.core.backends.azure.models import AzureClientCreds, AzureConfig +from dstack._internal.core.models.instances import ( + Gpu, + InstanceAvailability, + InstanceConfiguration, + InstanceOfferWithAvailability, + InstanceType, + Resources, + SSHKey, +) class TestVMImageVariant: @@ -71,3 +83,83 @@ def test_from_instance_type( ) def test_get_image_name(self, variant: VMImageVariant, expected_name: str): assert variant.get_image_name() == expected_name + + +def _config(network_security_group=None) -> AzureConfig: + return AzureConfig( + creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), + tenant_id="ten1", + subscription_id="sub1", + resource_group="my-rg", + regions=["eastus"], + network_security_group=network_security_group, + ) + + +def _offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend="azure", + instance=InstanceType( + name="Standard_DS1_v2", + resources=Resources(cpus=1, memory_mib=3500, gpus=[], spot=False), + ), + region="eastus", + price=0.1, + availability=InstanceAvailability.AVAILABLE, + ) + + +def _instance_config(security_group=None) -> InstanceConfiguration: + return InstanceConfiguration( + project_name="main", + instance_name="test-instance", + user="test-user", + ssh_keys=[SSHKey(public="ssh-rsa test")], + security_group=security_group, + ) + + +class TestAzureComputeNetworkSecurityGroup: + @pytest.mark.parametrize( + ["instance_sg", "config_nsg", "expected"], + [ + [None, None, azure_utils.get_default_network_security_group_name("my-rg", "eastus")], + [None, "config-nsg", "config-nsg"], + ["instance-nsg", None, "instance-nsg"], + # instance_config.security_group takes precedence over config.network_security_group + ["instance-nsg", "config-nsg", "instance-nsg"], + ], + ) + def test_create_instance_resolves_network_security_group( + self, instance_sg, config_nsg, expected + ): + with ( + patch("dstack._internal.core.backends.azure.compute.compute_mgmt"), + patch("dstack._internal.core.backends.azure.compute.network_mgmt"), + patch( + "dstack._internal.core.backends.azure.compute" + ".get_resource_group_network_subnet_or_error", + return_value=("net-rg", "net", "subnet"), + ), + patch("dstack._internal.core.backends.azure.compute._get_image_ref"), + patch( + "dstack._internal.core.backends.azure.compute._create_instance_and_wait" + ) as create_and_wait_mock, + patch( + "dstack._internal.core.backends.azure.compute._get_vm_public_private_ips", + return_value=("1.2.3.4", "10.0.0.1"), + ), + ): + vm_mock = Mock() + vm_mock.name = "test-vm-id" + create_and_wait_mock.return_value = vm_mock + compute = AzureCompute( + config=_config(network_security_group=config_nsg), credential=Mock() + ) + compute.create_instance( + instance_offer=_offer(), + instance_config=_instance_config(security_group=instance_sg), + placement_group=None, + ) + _, kwargs = create_and_wait_mock.call_args + assert kwargs["network_security_group"] == expected diff --git a/src/tests/_internal/core/backends/azure/test_configurator.py b/src/tests/_internal/core/backends/azure/test_configurator.py index 1a329e0cd5..11104fd117 100644 --- a/src/tests/_internal/core/backends/azure/test_configurator.py +++ b/src/tests/_internal/core/backends/azure/test_configurator.py @@ -129,3 +129,36 @@ def test_mixed_vpc_and_subnet_ids_covers_all_regions(self): subnet_ids={"eastus": "rg/net/subnet"}, ) self._check(config) + + +class TestCreateBackendNetworkSecurityGroup: + def _make_config(self, **kwargs): + return AzureBackendConfigWithCreds( + creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), + tenant_id="ten1", + subscription_id="sub1", + resource_group="my-rg", + regions=["eastus"], + **kwargs, + ) + + def _create_backend(self, config): + with ( + patch("dstack._internal.core.backends.azure.auth.authenticate") as authenticate_mock, + patch( + "dstack._internal.core.backends.azure.configurator.NetworkManager" + ) as NetworkManagerMock, + ): + authenticate_mock.return_value = Mock(), Mock() + AzureConfigurator().create_backend("proj", config) + return NetworkManagerMock.return_value + + def test_creates_instance_nsg_by_default(self): + network_manager = self._create_backend(self._make_config()) + network_manager.create_network_security_group.assert_called_once() + + def test_skips_instance_nsg_when_configured(self): + network_manager = self._create_backend(self._make_config(network_security_group="my-nsg")) + network_manager.create_network_security_group.assert_not_called() + # The gateway NSG is unaffected and still created. + network_manager.create_gateway_network_security_group.assert_called_once() diff --git a/src/tests/_internal/core/backends/gcp/test_compute.py b/src/tests/_internal/core/backends/gcp/test_compute.py new file mode 100644 index 0000000000..f7804be029 --- /dev/null +++ b/src/tests/_internal/core/backends/gcp/test_compute.py @@ -0,0 +1,90 @@ +from unittest.mock import Mock, patch + +import pytest + +from dstack._internal.core.backends.gcp.compute import GCPCompute +from dstack._internal.core.backends.gcp.models import GCPConfig, GCPServiceAccountCreds +from dstack._internal.core.errors import ComputeResourceNotFoundError + + +class _StopExecution(Exception): + """Raised to stop `create_*` right after the firewall guard is evaluated.""" + + +def _make_compute(create_firewall_rules=None, vpc_project_id=None) -> GCPCompute: + config = GCPConfig( + project_id="test-project", + vpc_project_id=vpc_project_id, + create_firewall_rules=create_firewall_rules, + creds=GCPServiceAccountCreds(data="creds"), + ) + # Bypass __init__ to avoid authenticating and creating real GCP clients. + compute = GCPCompute.__new__(GCPCompute) + compute.config = config + compute.firewalls_client = Mock() + return compute + + +def _run_create_instance(compute: GCPCompute): + instance_offer = Mock() + instance_offer.availability_zones = ["us-west1-a"] + instance_offer.region = "us-west1" + instance_offer.instance.resources.disk.size_mib = 102400 + instance_config = Mock() + instance_config.get_public_keys.return_value = [] + # Stop execution right after the firewall guard block. + compute._get_vpc_subnet = Mock(side_effect=_StopExecution) + with ( + patch( + "dstack._internal.core.backends.gcp.compute.generate_unique_instance_name", + return_value="instance-name", + ), + patch( + "dstack._internal.core.backends.gcp.resources.create_runner_firewall_rules" + ) as create_rules_mock, + pytest.raises(_StopExecution), + ): + compute.create_instance(instance_offer, instance_config, None) + return create_rules_mock + + +def _run_create_gateway(compute: GCPCompute): + configuration = Mock() + configuration.region = "us-west1" + compute.regions_client = Mock() + # Empty region list makes create_gateway raise right after the firewall guard. + compute.regions_client.list.return_value = [] + with ( + patch( + "dstack._internal.core.backends.gcp.resources.create_gateway_firewall_rules" + ) as create_rules_mock, + pytest.raises(ComputeResourceNotFoundError), + ): + compute.create_gateway(configuration) + return create_rules_mock + + +class TestCreateFirewallRules: + def test_runner_firewall_rules_created_by_default(self): + create_rules_mock = _run_create_instance(_make_compute()) + create_rules_mock.assert_called_once() + + def test_runner_firewall_rules_created_when_explicitly_enabled(self): + create_rules_mock = _run_create_instance(_make_compute(create_firewall_rules=True)) + create_rules_mock.assert_called_once() + + def test_runner_firewall_rules_skipped_when_disabled(self): + create_rules_mock = _run_create_instance(_make_compute(create_firewall_rules=False)) + create_rules_mock.assert_not_called() + + def test_gateway_firewall_rules_created_by_default(self): + create_rules_mock = _run_create_gateway(_make_compute()) + create_rules_mock.assert_called_once() + + def test_gateway_firewall_rules_created_when_explicitly_enabled(self): + create_rules_mock = _run_create_gateway(_make_compute(create_firewall_rules=True)) + create_rules_mock.assert_called_once() + + def test_gateway_firewall_rules_skipped_when_disabled(self): + create_rules_mock = _run_create_gateway(_make_compute(create_firewall_rules=False)) + create_rules_mock.assert_not_called() diff --git a/src/tests/_internal/core/backends/oci/test_compute.py b/src/tests/_internal/core/backends/oci/test_compute.py new file mode 100644 index 0000000000..00535d8ded --- /dev/null +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -0,0 +1,132 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from dstack._internal.core.backends.oci.compute import OCICompute +from dstack._internal.core.backends.oci.models import OCIClientCreds, OCIConfig +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceConfiguration, + InstanceOfferWithAvailability, + InstanceType, + Resources, + SSHKey, +) + + +def _make_config(network_security_group_id=None) -> OCIConfig: + return OCIConfig( + creds=OCIClientCreds( + user="user", + tenancy="tenancy", + key_content="key", + key_file=None, + pass_phrase=None, + fingerprint="fingerprint", + region="us-ashburn-1", + ), + regions=["us-ashburn-1"], + compartment_id="ocid1.compartment.oc1..compartment", + subnet_ids_per_region={"us-ashburn-1": "ocid1.subnet.oc1..subnet"}, + network_security_group_id=network_security_group_id, + ) + + +def _make_offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.OCI, + instance=InstanceType( + name="VM.Standard2.1", + resources=Resources(cpus=1, memory_mib=15000, gpus=[], spot=False), + ), + region="us-ashburn-1", + price=0.1, + availability=InstanceAvailability.AVAILABLE, + availability_zones=["AD-1"], + ) + + +def _make_instance_config(security_group=None) -> InstanceConfiguration: + return InstanceConfiguration( + project_name="test-project", + instance_name="test-instance", + user="test-user", + ssh_keys=[SSHKey(public="ssh-rsa AAAA")], + security_group=security_group, + ) + + +def _make_compute(config: OCIConfig) -> OCICompute: + with patch("dstack._internal.core.backends.oci.compute.make_region_clients_map") as m: + region = MagicMock() + subnet = MagicMock() + subnet.id = "ocid1.subnet.oc1..subnet" + subnet.vcn_id = "ocid1.vcn.oc1..vcn" + region.virtual_network_client.get_subnet.return_value.data = subnet + m.return_value = {"us-ashburn-1": region} + compute = OCICompute(config) + return compute + + +class TestOCIComputeSecurityGroup: + def _run_create_instance(self, compute, instance_config): + with patch("dstack._internal.core.backends.oci.compute.resources") as res: + res.VCN_CIDR = "10.0.0.0/16" + res.get_marketplace_listing_and_package.return_value = (MagicMock(), MagicMock()) + res.get_or_create_security_group.return_value.id = "ocid1.nsg.oc1..managed" + res.launch_instance.return_value.id = "ocid1.instance.oc1..instance" + compute.create_instance(_make_offer(), instance_config, placement_group=None) + return res + + def test_default_creates_and_syncs_managed_security_group(self): + compute = _make_compute(_make_config()) + res = self._run_create_instance(compute, _make_instance_config()) + + res.get_or_create_security_group.assert_called_once() + res.update_security_group_rules_for_runner_instances.assert_called_once() + assert ( + res.launch_instance.call_args.kwargs["security_group_id"] + == "ocid1.nsg.oc1..managed" + ) + + def test_project_level_custom_nsg_is_left_untouched(self): + compute = _make_compute(_make_config(network_security_group_id="ocid1.nsg.oc1..custom")) + res = self._run_create_instance(compute, _make_instance_config()) + + res.get_or_create_security_group.assert_not_called() + res.update_security_group_rules_for_runner_instances.assert_not_called() + res.update_security_group_rules.assert_not_called() + assert ( + res.launch_instance.call_args.kwargs["security_group_id"] + == "ocid1.nsg.oc1..custom" + ) + + def test_instance_level_custom_nsg_is_left_untouched(self): + compute = _make_compute(_make_config()) + res = self._run_create_instance( + compute, _make_instance_config(security_group="ocid1.nsg.oc1..run") + ) + + res.get_or_create_security_group.assert_not_called() + res.update_security_group_rules_for_runner_instances.assert_not_called() + res.update_security_group_rules.assert_not_called() + assert ( + res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" + ) + + def test_instance_level_overrides_project_level(self): + compute = _make_compute(_make_config(network_security_group_id="ocid1.nsg.oc1..project")) + res = self._run_create_instance( + compute, _make_instance_config(security_group="ocid1.nsg.oc1..run") + ) + + res.get_or_create_security_group.assert_not_called() + res.update_security_group_rules_for_runner_instances.assert_not_called() + assert ( + res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 3be52da25dddb54e1f41d935f1e36cf258cedfde Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 11:03:50 -0400 Subject: [PATCH 2/6] Support multi-region custom security groups/NSGs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- mkdocs/docs/concepts/backends.md | 58 ++++++--- .../_internal/core/backends/aws/compute.py | 15 ++- .../core/backends/aws/configurator.py | 7 ++ .../_internal/core/backends/aws/models.py | 23 +++- .../_internal/core/backends/aws/resources.py | 15 +++ .../_internal/core/backends/azure/compute.py | 10 +- .../core/backends/azure/configurator.py | 13 +- .../_internal/core/backends/azure/models.py | 15 ++- .../_internal/core/backends/oci/compute.py | 8 +- .../_internal/core/backends/oci/models.py | 14 ++- .../core/backends/aws/test_compute.py | 113 ++++++++++++++---- .../core/backends/aws/test_resources.py | 46 +++++++ .../core/backends/azure/test_compute.py | 43 ++++--- .../core/backends/azure/test_configurator.py | 37 +++++- .../core/backends/oci/test_compute.py | 35 +++++- 15 files changed, 357 insertions(+), 95 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 02b0d9c909..690ae64bff 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/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. diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 267b886a2c..def8e662ed 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/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, diff --git a/src/dstack/_internal/core/backends/aws/configurator.py b/src/dstack/_internal/core/backends/aws/configurator.py index 8d6c8afe14..63f7469ab6 100644 --- a/src/dstack/_internal/core/backends/aws/configurator.py +++ b/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 diff --git a/src/dstack/_internal/core/backends/aws/models.py b/src/dstack/_internal/core/backends/aws/models.py index 242efce045..946a588978 100644 --- a/src/dstack/_internal/core/backends/aws/models.py +++ b/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_`)." + " 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`"), diff --git a/src/dstack/_internal/core/backends/aws/resources.py b/src/dstack/_internal/core/backends/aws/resources.py index 5ee3f63191..4a744931c5 100644 --- a/src/dstack/_internal/core/backends/aws/resources.py +++ b/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, diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index be2e065946..aabf6e6e16 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/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 diff --git a/src/dstack/_internal/core/backends/azure/configurator.py b/src/dstack/_internal/core/backends/azure/configurator.py index 1b8f76522d..19efe7b1e7 100644 --- a/src/dstack/_internal/core/backends/azure/configurator.py +++ b/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, diff --git a/src/dstack/_internal/core/backends/azure/models.py b/src/dstack/_internal/core/backends/azure/models.py index c883616714..30a15fd6e1 100644 --- a/src/dstack/_internal/core/backends/azure/models.py +++ b/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" ) ), diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index 8241612b40..b4640d8ceb 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/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", diff --git a/src/dstack/_internal/core/backends/oci/models.py b/src/dstack/_internal/core/backends/oci/models.py index d6dd01deec..597763342d 100644 --- a/src/dstack/_internal/core/backends/oci/models.py +++ b/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" ) ), diff --git a/src/tests/_internal/core/backends/aws/test_compute.py b/src/tests/_internal/core/backends/aws/test_compute.py index 332e9249a4..49392b9df5 100644 --- a/src/tests/_internal/core/backends/aws/test_compute.py +++ b/src/tests/_internal/core/backends/aws/test_compute.py @@ -6,6 +6,7 @@ from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds, AWSConfig from dstack._internal.core.backends.base.compute import ComputeWithSecurityGroupSupport from dstack._internal.core.backends.features import BACKENDS_WITH_SECURITY_GROUP_SUPPORT +from dstack._internal.core.errors import ComputeError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( InstanceAvailability, @@ -17,11 +18,12 @@ ) -def _config(security_group_id=None) -> AWSConfig: +def _config(security_group_name=None, security_group_ids=None) -> AWSConfig: return AWSConfig( creds=AWSAccessKeyCreds(access_key="test", secret_key="test"), regions=["us-east-1"], - security_group_id=security_group_id, + security_group_name=security_group_name, + security_group_ids=security_group_ids, ) @@ -37,14 +39,14 @@ def _compute(config: AWSConfig) -> AWSCompute: return compute -def _offer() -> InstanceOfferWithAvailability: +def _offer(region="us-east-1") -> InstanceOfferWithAvailability: return InstanceOfferWithAvailability( backend=BackendType.AWS, instance=InstanceType( name="m5.large", resources=Resources(cpus=2, memory_mib=8192, gpus=[], spot=False), ), - region="us-east-1", + region=region, price=0.1, availability=InstanceAvailability.AVAILABLE, ) @@ -60,9 +62,15 @@ def _instance_config(security_group=None) -> InstanceConfiguration: ) -def _run_create_instance(compute: AWSCompute, instance_config: InstanceConfiguration) -> str: +def _run_create_instance( + compute: AWSCompute, + instance_config: InstanceConfiguration, + offer: InstanceOfferWithAvailability = None, +) -> str: """Runs create_instance with the instances struct mocked and returns the security_group_id that was passed into create_instances_struct.""" + if offer is None: + offer = _offer() with patch( "dstack._internal.core.backends.aws.compute.aws_resources.create_instances_struct" ) as struct_mock: @@ -73,7 +81,7 @@ def _run_create_instance(compute: AWSCompute, instance_config: InstanceConfigura instance_mock.capacity_reservation_id = None ec2_resource.create_instances.return_value = [instance_mock] compute.create_instance( - instance_offer=_offer(), + instance_offer=offer, instance_config=instance_config, placement_group=None, ) @@ -93,31 +101,86 @@ def test_auto_creates_group_when_no_custom_sg_configured(self): compute._create_security_group.assert_called_once() assert security_group_id == "sg-auto" - def test_uses_project_level_security_group_id_without_managing_it(self): - compute = _compute(_config(security_group_id="sg-project")) + def test_uses_security_group_ids_for_region_without_managing_it(self): + compute = _compute(_config(security_group_ids={"us-east-1": "sg-region"})) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources" + ".get_security_group_id_by_name" + ) as lookup_mock: + security_group_id = _run_create_instance(compute, _instance_config()) + assert security_group_id == "sg-region" + # The per-region ID is used directly: no auto-create, no name lookup. + compute._create_security_group.assert_not_called() + lookup_mock.assert_not_called() + + def test_security_group_ids_missing_region_falls_back_to_auto_create(self): + compute = _compute(_config(security_group_ids={"eu-west-1": "sg-eu"})) + # Offer region us-east-1 is not in the mapping -> fall back to auto-create. security_group_id = _run_create_instance(compute, _instance_config()) - # No auto-create / rule-management happens for a custom SG. + compute._create_security_group.assert_called_once() + assert security_group_id == "sg-auto" + + def test_security_group_ids_missing_region_falls_back_to_name(self): + compute = _compute( + _config( + security_group_name="my-sg", + security_group_ids={"eu-west-1": "sg-eu"}, + ) + ) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources" + ".get_security_group_id_by_name", + return_value="sg-by-name", + ) as lookup_mock: + security_group_id = _run_create_instance(compute, _instance_config()) + # Region not in the ID map -> fall back to name lookup, not auto-create. + assert security_group_id == "sg-by-name" + compute._create_security_group.assert_not_called() + lookup_mock.assert_called_once() + assert lookup_mock.call_args.kwargs["name"] == "my-sg" + assert lookup_mock.call_args.kwargs["vpc_id"] == "vpc-1" + + def test_security_group_name_triggers_lookup_and_uses_result(self): + compute = _compute(_config(security_group_name="my-sg")) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources" + ".get_security_group_id_by_name", + return_value="sg-by-name", + ) as lookup_mock: + security_group_id = _run_create_instance(compute, _instance_config()) + assert security_group_id == "sg-by-name" + compute._create_security_group.assert_not_called() + lookup_mock.assert_called_once() + + def test_security_group_name_not_found_raises(self): + compute = _compute(_config(security_group_name="missing-sg")) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources" + ".get_security_group_id_by_name", + return_value=None, + ): + with pytest.raises(ComputeError, match="missing-sg"): + _run_create_instance(compute, _instance_config()) + # No silent fall-through to auto-create. compute._create_security_group.assert_not_called() - assert security_group_id == "sg-project" - def test_run_level_security_group_overrides_project_level(self): - compute = _compute(_config(security_group_id="sg-project")) + def test_run_level_security_group_overrides_ids(self): + compute = _compute(_config(security_group_ids={"us-east-1": "sg-region"})) security_group_id = _run_create_instance( compute, _instance_config(security_group="sg-run") ) compute._create_security_group.assert_not_called() assert security_group_id == "sg-run" - @pytest.mark.parametrize( - ["config_sg", "instance_sg", "expected"], - [ - [None, None, "sg-auto"], - ["sg-project", None, "sg-project"], - [None, "sg-run", "sg-run"], - ["sg-project", "sg-run", "sg-run"], - ], - ) - def test_precedence(self, config_sg, instance_sg, expected): - compute = _compute(_config(security_group_id=config_sg)) - security_group_id = _run_create_instance(compute, _instance_config(instance_sg)) - assert security_group_id == expected + def test_run_level_security_group_overrides_name(self): + compute = _compute(_config(security_group_name="my-sg")) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources" + ".get_security_group_id_by_name" + ) as lookup_mock: + security_group_id = _run_create_instance( + compute, _instance_config(security_group="sg-run") + ) + assert security_group_id == "sg-run" + compute._create_security_group.assert_not_called() + lookup_mock.assert_not_called() diff --git a/src/tests/_internal/core/backends/aws/test_resources.py b/src/tests/_internal/core/backends/aws/test_resources.py index 0a5536587c..2ff12855e2 100644 --- a/src/tests/_internal/core/backends/aws/test_resources.py +++ b/src/tests/_internal/core/backends/aws/test_resources.py @@ -11,6 +11,7 @@ create_instances_struct, create_security_group, get_image_id_and_username, + get_security_group_id_by_name, validate_tags, ) from dstack._internal.core.errors import BackendError, ComputeResourceNotFoundError @@ -358,6 +359,51 @@ def test_reuses_existing_group_without_recreating(self, ec2_client_mock: Mock): ec2_client_mock.authorize_security_group_egress.assert_not_called() +class TestGetSecurityGroupIdByName: + def test_returns_id_when_found(self): + ec2_client_mock = Mock(spec_set=["describe_security_groups"]) + ec2_client_mock.describe_security_groups.return_value = { + "SecurityGroups": [{"GroupId": "sg-found"}] + } + result = get_security_group_id_by_name( + ec2_client=ec2_client_mock, + name="my-sg", + vpc_id="vpc-1", + ) + assert result == "sg-found" + ec2_client_mock.describe_security_groups.assert_called_once_with( + Filters=[ + {"Name": "group-name", "Values": ["my-sg"]}, + {"Name": "vpc-id", "Values": ["vpc-1"]}, + ] + ) + + def test_omits_vpc_filter_when_vpc_id_none(self): + ec2_client_mock = Mock(spec_set=["describe_security_groups"]) + ec2_client_mock.describe_security_groups.return_value = { + "SecurityGroups": [{"GroupId": "sg-found"}] + } + result = get_security_group_id_by_name( + ec2_client=ec2_client_mock, + name="my-sg", + vpc_id=None, + ) + assert result == "sg-found" + ec2_client_mock.describe_security_groups.assert_called_once_with( + Filters=[{"Name": "group-name", "Values": ["my-sg"]}] + ) + + def test_returns_none_when_not_found(self): + ec2_client_mock = Mock(spec_set=["describe_security_groups"]) + ec2_client_mock.describe_security_groups.return_value = {"SecurityGroups": []} + result = get_security_group_id_by_name( + ec2_client=ec2_client_mock, + name="my-sg", + vpc_id="vpc-1", + ) + assert result is None + + class TestCreateNetworkInterfacesStruct: def test_non_efa_instance_single_interface(self): interfaces = _create_network_interfaces_struct( diff --git a/src/tests/_internal/core/backends/azure/test_compute.py b/src/tests/_internal/core/backends/azure/test_compute.py index 50ea8759f7..b6a8975e97 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/src/tests/_internal/core/backends/azure/test_compute.py @@ -85,25 +85,25 @@ def test_get_image_name(self, variant: VMImageVariant, expected_name: str): assert variant.get_image_name() == expected_name -def _config(network_security_group=None) -> AzureConfig: +def _config(network_security_group_ids=None) -> AzureConfig: return AzureConfig( creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), tenant_id="ten1", subscription_id="sub1", resource_group="my-rg", - regions=["eastus"], - network_security_group=network_security_group, + regions=["eastus", "westeurope"], + network_security_group_ids=network_security_group_ids, ) -def _offer() -> InstanceOfferWithAvailability: +def _offer(region="eastus") -> InstanceOfferWithAvailability: return InstanceOfferWithAvailability( backend="azure", instance=InstanceType( name="Standard_DS1_v2", resources=Resources(cpus=1, memory_mib=3500, gpus=[], spot=False), ), - region="eastus", + region=region, price=0.1, availability=InstanceAvailability.AVAILABLE, ) @@ -121,17 +121,32 @@ def _instance_config(security_group=None) -> InstanceConfiguration: class TestAzureComputeNetworkSecurityGroup: @pytest.mark.parametrize( - ["instance_sg", "config_nsg", "expected"], + ["instance_sg", "nsg_ids", "region", "expected"], [ - [None, None, azure_utils.get_default_network_security_group_name("my-rg", "eastus")], - [None, "config-nsg", "config-nsg"], - ["instance-nsg", None, "instance-nsg"], - # instance_config.security_group takes precedence over config.network_security_group - ["instance-nsg", "config-nsg", "instance-nsg"], + # No mapping, no instance security_group -> auto-derived default per location. + [ + None, + None, + "eastus", + azure_utils.get_default_network_security_group_name("my-rg", "eastus"), + ], + # Location present in the mapping resolves to the mapped NSG name. + [None, {"eastus": "config-nsg"}, "eastus", "config-nsg"], + # Location absent from the mapping falls back to the auto-derived default name. + [ + None, + {"eastus": "config-nsg"}, + "westeurope", + azure_utils.get_default_network_security_group_name("my-rg", "westeurope"), + ], + # instance_config.security_group takes precedence when no mapping is set. + ["instance-nsg", None, "eastus", "instance-nsg"], + # instance_config.security_group takes precedence over the mapping. + ["instance-nsg", {"eastus": "config-nsg"}, "eastus", "instance-nsg"], ], ) def test_create_instance_resolves_network_security_group( - self, instance_sg, config_nsg, expected + self, instance_sg, nsg_ids, region, expected ): with ( patch("dstack._internal.core.backends.azure.compute.compute_mgmt"), @@ -154,10 +169,10 @@ def test_create_instance_resolves_network_security_group( vm_mock.name = "test-vm-id" create_and_wait_mock.return_value = vm_mock compute = AzureCompute( - config=_config(network_security_group=config_nsg), credential=Mock() + config=_config(network_security_group_ids=nsg_ids), credential=Mock() ) compute.create_instance( - instance_offer=_offer(), + instance_offer=_offer(region=region), instance_config=_instance_config(security_group=instance_sg), placement_group=None, ) diff --git a/src/tests/_internal/core/backends/azure/test_configurator.py b/src/tests/_internal/core/backends/azure/test_configurator.py index 11104fd117..a327fb8985 100644 --- a/src/tests/_internal/core/backends/azure/test_configurator.py +++ b/src/tests/_internal/core/backends/azure/test_configurator.py @@ -132,13 +132,13 @@ def test_mixed_vpc_and_subnet_ids_covers_all_regions(self): class TestCreateBackendNetworkSecurityGroup: - def _make_config(self, **kwargs): + def _make_config(self, regions=None, **kwargs): return AzureBackendConfigWithCreds( creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), tenant_id="ten1", subscription_id="sub1", resource_group="my-rg", - regions=["eastus"], + regions=regions if regions is not None else ["eastus"], **kwargs, ) @@ -153,12 +153,37 @@ def _create_backend(self, config): AzureConfigurator().create_backend("proj", config) return NetworkManagerMock.return_value + def _default_nsg_locations(self, network_manager): + return sorted( + call.kwargs["location"] + for call in network_manager.create_network_security_group.call_args_list + ) + def test_creates_instance_nsg_by_default(self): network_manager = self._create_backend(self._make_config()) network_manager.create_network_security_group.assert_called_once() - def test_skips_instance_nsg_when_configured(self): - network_manager = self._create_backend(self._make_config(network_security_group="my-nsg")) + def test_skips_instance_nsg_only_for_configured_locations(self): + # eastus has a custom NSG, westeurope does not - only westeurope gets a default NSG. + network_manager = self._create_backend( + self._make_config( + regions=["eastus", "westeurope"], + network_security_group_ids={"eastus": "my-nsg"}, + ) + ) + assert self._default_nsg_locations(network_manager) == ["westeurope"] + # The gateway NSG is unaffected and still created for every location. + assert network_manager.create_gateway_network_security_group.call_count == 2 + + def test_skips_instance_nsg_for_all_configured_locations(self): + network_manager = self._create_backend( + self._make_config( + regions=["eastus", "westeurope"], + network_security_group_ids={ + "eastus": "my-nsg", + "westeurope": "my-other-nsg", + }, + ) + ) network_manager.create_network_security_group.assert_not_called() - # The gateway NSG is unaffected and still created. - network_manager.create_gateway_network_security_group.assert_called_once() + assert network_manager.create_gateway_network_security_group.call_count == 2 diff --git a/src/tests/_internal/core/backends/oci/test_compute.py b/src/tests/_internal/core/backends/oci/test_compute.py index 00535d8ded..8ced03f83a 100644 --- a/src/tests/_internal/core/backends/oci/test_compute.py +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -15,7 +15,7 @@ ) -def _make_config(network_security_group_id=None) -> OCIConfig: +def _make_config(network_security_group_ids=None) -> OCIConfig: return OCIConfig( creds=OCIClientCreds( user="user", @@ -29,7 +29,7 @@ def _make_config(network_security_group_id=None) -> OCIConfig: regions=["us-ashburn-1"], compartment_id="ocid1.compartment.oc1..compartment", subnet_ids_per_region={"us-ashburn-1": "ocid1.subnet.oc1..subnet"}, - network_security_group_id=network_security_group_id, + network_security_group_ids=network_security_group_ids, ) @@ -90,8 +90,12 @@ def test_default_creates_and_syncs_managed_security_group(self): == "ocid1.nsg.oc1..managed" ) - def test_project_level_custom_nsg_is_left_untouched(self): - compute = _make_compute(_make_config(network_security_group_id="ocid1.nsg.oc1..custom")) + def test_per_region_custom_nsg_is_left_untouched(self): + compute = _make_compute( + _make_config( + network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..custom"} + ) + ) res = self._run_create_instance(compute, _make_instance_config()) res.get_or_create_security_group.assert_not_called() @@ -102,6 +106,21 @@ def test_project_level_custom_nsg_is_left_untouched(self): == "ocid1.nsg.oc1..custom" ) + def test_region_not_in_mapping_falls_back_to_managed(self): + compute = _make_compute( + _make_config( + network_security_group_ids={"us-phoenix-1": "ocid1.nsg.oc1..other"} + ) + ) + res = self._run_create_instance(compute, _make_instance_config()) + + res.get_or_create_security_group.assert_called_once() + res.update_security_group_rules_for_runner_instances.assert_called_once() + assert ( + res.launch_instance.call_args.kwargs["security_group_id"] + == "ocid1.nsg.oc1..managed" + ) + def test_instance_level_custom_nsg_is_left_untouched(self): compute = _make_compute(_make_config()) res = self._run_create_instance( @@ -115,8 +134,12 @@ def test_instance_level_custom_nsg_is_left_untouched(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" ) - def test_instance_level_overrides_project_level(self): - compute = _make_compute(_make_config(network_security_group_id="ocid1.nsg.oc1..project")) + def test_instance_level_overrides_per_region_mapping(self): + compute = _make_compute( + _make_config( + network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..project"} + ) + ) res = self._run_create_instance( compute, _make_instance_config(security_group="ocid1.nsg.oc1..run") ) From 5d690f0bef3b36bf3a98fea95de754f02b4412b9 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 12:27:10 -0400 Subject: [PATCH 3/6] Fix issues found in independent review An independent review of the custom-security-group feature (previous two commits) found several real gaps. Fixed all of them: - security_group was dropped when a run provisioned a new instance into an existing fleet (only the fleet-apply path honored it). Fixed by threading security_group through the Requirements pipeline exactly like reservation already is (Requirements.security_group, combine_fleet_and_run_profiles/ combine_fleet_and_run_requirements, and sourcing run_job's InstanceConfiguration from job.job_spec.requirements.security_group instead of the run's raw profile). - security_group was silently ignored when an offer resolved to a backend that doesn't support it (e.g. GCP). offers.py now narrows backend_types to BACKENDS_WITH_SECURITY_GROUP_SUPPORT when security_group is set, mirroring the existing reservation filtering. - AWS: the configurator forbade combining security_group_name with security_group_ids, but compute.py implements (and docs/tests described) a fallback from ids to name - the combination is now allowed. Also added validation catching region-key typos in security_group_ids, and a clearer ComputeError instead of a confusing NoCapacityError retry loop when a configured security group doesn't exist in the target VPC. - Azure: renamed network_security_group_ids to network_security_group_names since the values are NSG names (not IDs) scoped to the backend's resource_group, and added region-key typo validation. - GCP: create_firewall_rules no longer disables the gateway firewall rule, matching AWS/Azure/OCI where gateway security resources are always auto-managed regardless of the custom-security-group settings. - OCI: the shared subnet has no security_list_ids, so it inherits the VCN's permissive default security list (SSH open to 0.0.0.0/0, allow-all egress). Since OCI evaluates security lists and NSGs as a union of allows, a custom NSG could not actually restrict anything. Fixed by routing custom-NSG instances into a separate, dedicated VCN/subnet with no security list, so the NSG becomes the sole security boundary. The default VCN/subnet used by auto-managed-NSG instances is completely untouched. Also added region-key typo validation and corrected the docs to accurately attribute default SSH exposure to the security list, not the NSG. --- mkdocs/docs/concepts/backends.md | 44 +++++++--- .../_internal/core/backends/aws/compute.py | 11 ++- .../core/backends/aws/configurator.py | 11 ++- .../_internal/core/backends/azure/compute.py | 4 +- .../core/backends/azure/configurator.py | 30 ++++++- .../_internal/core/backends/azure/models.py | 8 +- .../_internal/core/backends/base/compute.py | 2 +- .../_internal/core/backends/gcp/compute.py | 5 +- .../_internal/core/backends/gcp/models.py | 8 +- .../_internal/core/backends/oci/compute.py | 20 ++++- .../core/backends/oci/configurator.py | 17 ++++ .../_internal/core/backends/oci/models.py | 7 +- .../_internal/core/backends/oci/resources.py | 86 +++++++++++++++++++ src/dstack/_internal/core/models/runs.py | 1 + .../_internal/server/services/fleets.py | 1 + .../services/jobs/configurators/base.py | 4 + .../_internal/server/services/offers.py | 6 ++ .../server/services/requirements/combine.py | 6 ++ .../core/backends/aws/test_compute.py | 27 ++++++ .../core/backends/aws/test_configurator.py | 67 +++++++++++++++ .../core/backends/azure/test_compute.py | 10 +-- .../core/backends/azure/test_configurator.py | 44 +++++++++- .../core/backends/gcp/test_compute.py | 7 +- .../core/backends/oci/test_compute.py | 79 ++++++++++++++++- .../core/backends/oci/test_configurator.py | 68 ++++++++++++++- .../_internal/server/routers/test_backends.py | 2 + .../_internal/server/routers/test_fleets.py | 6 ++ .../_internal/server/routers/test_runs.py | 6 ++ .../services/requirements/test_combine.py | 84 ++++++++++++++++++ .../_internal/server/services/test_offers.py | 22 +++++ 30 files changed, 646 insertions(+), 47 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 690ae64bff..cefa4a0cce 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -489,7 +489,9 @@ There are two ways to configure Azure: using a client secret or using the defaul 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_ids`: + 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: @@ -499,12 +501,12 @@ There are two ways to configure Azure: using a client secret or using the defaul creds: type: default regions: [westeurope, eastus] - network_security_group_ids: + network_security_group_names: westeurope: my-network-security-group-we eastus: my-network-security-group-eus ``` - Locations not covered by `network_security_group_ids` fall back to dstack's auto-created network + 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. @@ -730,11 +732,11 @@ gcloud projects list --format="json(projectId)" 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 VPC firewall rules allowing inbound SSH (and, for gateways, HTTP/HTTPS) from - the internet, scoped to the `dstack-runner-instance` and `dstack-gateway-instance` target tags. + 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 - rules that open ports to `0.0.0.0/0`, disable this at the project level with `create_firewall_rules: false`: + 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: @@ -749,7 +751,9 @@ gcloud projects list --format="json(projectId)" ``` You're then responsible for ensuring your VPC's own firewall rules allow whatever SSH and cluster traffic - `dstack` needs. + `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 @@ -1158,8 +1162,13 @@ 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). + By default, `dstack` places instances in a shared subnet whose OCI security list opens SSH + (TCP port 22) to `0.0.0.0/0` and permits all outbound traffic. On top of that, `dstack` creates + and manages its own network security group (NSG) per project, which only adds a rule allowing all + traffic within the group so multi-node clusters work out of the box. In other words, the + permissive SSH ingress and the outbound access come from the subnet's security list, not from the + auto-managed NSG. + OCI network security groups are region-scoped, so a custom NSG is configured per region via `network_security_group_ids`: @@ -1176,9 +1185,20 @@ There are two ways to configure OCI: using client credentials or using the defau ``` 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. + security group and shared subnet. + + When a custom NSG is used, `dstack` never adds, removes, or modifies its rules, and it places the + affected instances in a separate VCN and subnet that has **no** OCI security list (a second subnet + in the same VCN isn't possible here, since the default subnet already occupies the whole VCN's + address space). This makes the NSG the sole security boundary — there is no longer any implicit + SSH-from-anywhere or implicit outbound-all coming 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. diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index def8e662ed..6a5b22150e 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -380,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. diff --git a/src/dstack/_internal/core/backends/aws/configurator.py b/src/dstack/_internal/core/backends/aws/configurator.py index 63f7469ab6..6ca43997db 100644 --- a/src/dstack/_internal/core/backends/aws/configurator.py +++ b/src/dstack/_internal/core/backends/aws/configurator.py @@ -148,9 +148,16 @@ def _check_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: + 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="Only one of `security_group_name` and `security_group_ids` can be specified" + 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): diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index aabf6e6e16..b219560702 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -149,8 +149,8 @@ def create_instance( allocate_public_ip=allocate_public_ip, ) 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 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, diff --git a/src/dstack/_internal/core/backends/azure/configurator.py b/src/dstack/_internal/core/backends/azure/configurator.py index 19efe7b1e7..b93e088db2 100644 --- a/src/dstack/_internal/core/backends/azure/configurator.py +++ b/src/dstack/_internal/core/backends/azure/configurator.py @@ -104,6 +104,7 @@ def validate_config(self, config: AzureBackendConfigWithCreds, default_creds_ena self._check_config_resource_group(config=config, credential=credential) self._check_config_vm_managed_identity(config=config, credential=credential) self._check_config_vpc(config=config, credential=credential) + self._check_config_network_security_groups(config) def create_backend( self, project_name: str, config: AzureBackendConfigWithCreds @@ -126,7 +127,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, - network_security_group_ids=config.network_security_group_ids, + network_security_group_names=config.network_security_group_names, ) return BackendRecord( config=AzureStoredConfig( @@ -283,6 +284,27 @@ def _check_config_vpc( except BackendError as e: raise ServerClientError(e.args[0]) + def _check_config_network_security_groups(self, config: AzureBackendConfigWithCreds): + if not config.network_security_group_names: + return + # When `regions` is None, all regions are used and there is no feasible way to validate + # location keys against the full list of Azure regions here, so the check is skipped. + if config.regions is None: + return + configured_regions = set(config.regions) + unknown_locations = sorted( + location + for location in config.network_security_group_names + if location not in configured_regions + ) + if unknown_locations: + raise ServerClientError( + f"Locations {unknown_locations} in `network_security_group_names` are not in" + " `regions`. Otherwise, these network security groups would be silently ignored" + " and instances in those locations would fall back to dstack's auto-created" + " network security group." + ) + def _check_config_vm_managed_identity( self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential ): @@ -341,7 +363,7 @@ def _create_network_resources( resource_group: str, locations: List[str], create_default_network: bool, - network_security_group_ids: Optional[Dict[str, str]] = None, + network_security_group_names: Optional[Dict[str, str]] = None, ): def func(location: str): network_manager = NetworkManager( @@ -354,9 +376,9 @@ 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 location not in (network_security_group_ids or {}): + if location not in (network_security_group_names 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 + # location via `network_security_group_names` - dstack does not create or manage # it in that case. network_manager.create_network_security_group( resource_group=resource_group, diff --git a/src/dstack/_internal/core/backends/azure/models.py b/src/dstack/_internal/core/backends/azure/models.py index 30a15fd6e1..10d0536076 100644 --- a/src/dstack/_internal/core/backends/azure/models.py +++ b/src/dstack/_internal/core/backends/azure/models.py @@ -81,13 +81,15 @@ class AzureBackendConfig(CoreModel): ) ), ] = None - network_security_group_ids: Annotated[ + network_security_group_names: Annotated[ Optional[Dict[str, str]], Field( description=( "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." + " to use for instances instead of the one `dstack` creates and manages automatically." + " The values are plain NSG names (not full Azure resource IDs) and must refer to" + " network security groups within the configured `resource_group` — NSGs in a different" + " resource group cannot be referenced." " 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" diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 806d92605b..35589381b6 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -321,7 +321,7 @@ def run_job( ssh_keys=[SSHKey(public=project_ssh_public_key.strip())], volumes=volumes, reservation=job.job_spec.requirements.reservation, - security_group=run.run_spec.merged_profile.security_group, + security_group=job.job_spec.requirements.security_group, tags=run.run_spec.merged_profile.tags, ) instance_offer = instance_offer.copy() diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index b416b7791d..9652ed86dc 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -561,7 +561,10 @@ def create_gateway( self, configuration: GatewayComputeConfiguration, ) -> GatewayProvisioningData: - if self.config.vpc_project_id is None and self.config.create_firewall_rules is not False: + if self.config.vpc_project_id is None: + # Gateway firewall rules are intentionally not gated by `create_firewall_rules`: + # gateways are meant to be internet-reachable, and this keeps behavior consistent + # with AWS/Azure/OCI, where gateway security groups/NSGs are always auto-managed. gcp_resources.create_gateway_firewall_rules( firewalls_client=self.firewalls_client, project_id=self.config.project_id, diff --git a/src/dstack/_internal/core/backends/gcp/models.py b/src/dstack/_internal/core/backends/gcp/models.py index d5477aca6e..87e3ad9eff 100644 --- a/src/dstack/_internal/core/backends/gcp/models.py +++ b/src/dstack/_internal/core/backends/gcp/models.py @@ -84,10 +84,12 @@ class GCPBackendConfig(CoreModel): Optional[bool], Field( description=( - "A flag to enable/disable `dstack` creating VPC firewall rules that allow SSH" - " (and, for gateways, HTTP/HTTPS) traffic from the internet." + "A flag to enable/disable `dstack` creating a VPC firewall rule that allows SSH" + " traffic from the internet to instances." " Set to `false` if you manage your own firewall rules and don't want `dstack`" - " creating rules that open ports to `0.0.0.0/0`. Defaults to `true`" + " creating a rule that opens port 22 to `0.0.0.0/0`. Defaults to `true`." + " This does not affect the separate firewall rule `dstack` creates for gateways," + " which are meant to be internet-reachable and are always auto-managed" ) ), ] = None diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index b4640d8ceb..b2f2f9ae88 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -143,7 +143,8 @@ def create_instance( security_group_id = self.config.network_security_group_ids.get( instance_offer.region ) - if security_group_id is None: + using_custom_security_group = security_group_id is not None + if not using_custom_security_group: security_group = resources.get_or_create_security_group( f"dstack-{instance_config.project_name}-default-security-group", subnet.vcn_id, @@ -154,10 +155,25 @@ def create_instance( security_group.id, region.virtual_network_client ) security_group_id = security_group.id + firewall_allow_from_subnet = resources.VCN_CIDR + else: + # A user-managed (custom) NSG is in use. Place the instance in a + # dedicated VCN/subnet that has no OCI security list, so the NSG is + # the sole security boundary. This is a *separate* VCN, not a second + # subnet in the default one: the default subnet already occupies the + # default VCN's entire CIDR block, so a second subnet there would + # conflict. The default VCN/subnet is left completely untouched, + # so instances using dstack's auto-managed NSG are unaffected. + subnet = resources.set_up_restricted_network_resources_in_region( + compartment_id=self.config.compartment_id, + project_name=instance_config.project_name, + client=region.virtual_network_client, + ) + firewall_allow_from_subnet = resources.RESTRICTED_VCN_CIDR cloud_init_user_data = get_user_data( authorized_keys=instance_config.get_public_keys(), - firewall_allow_from_subnets=[resources.VCN_CIDR], + firewall_allow_from_subnets=[firewall_allow_from_subnet], ) display_name = generate_unique_instance_name(instance_config) diff --git a/src/dstack/_internal/core/backends/oci/configurator.py b/src/dstack/_internal/core/backends/oci/configurator.py index 4558e8bf96..2ac42c27f5 100644 --- a/src/dstack/_internal/core/backends/oci/configurator.py +++ b/src/dstack/_internal/core/backends/oci/configurator.py @@ -60,6 +60,7 @@ def validate_config(self, config: OCIBackendConfigWithCreds, default_creds_enabl get_subscribed_regions(config.creds).names except any_oci_exception as e: raise_invalid_credentials_error(fields=[["creds"]], details=e) + _check_config_network_security_groups(config) def create_backend( self, project_name: str, config: OCIBackendConfigWithCreds @@ -106,6 +107,22 @@ def _get_config(self, record: BackendRecord) -> OCIConfig: ) +def _check_config_network_security_groups(config: OCIBackendConfigWithCreds) -> None: + if config.network_security_group_ids is None or config.regions is None: + return + regions = set(config.regions) + unknown_regions = [r for r in config.network_security_group_ids if r not in regions] + if unknown_regions: + raise ServerClientError( + msg=( + f"`network_security_group_ids` specifies regions not in `regions`:" + f" {unknown_regions}. This is likely a typo — remove the extra keys or add" + " them to `regions`" + ), + fields=[["network_security_group_ids"]], + ) + + def _filter_supported_regions(subscribed_region_names: Set[str]) -> List[str]: available_regions = subscribed_region_names & SUPPORTED_REGIONS if not available_regions: diff --git a/src/dstack/_internal/core/backends/oci/models.py b/src/dstack/_internal/core/backends/oci/models.py index 597763342d..b30530125b 100644 --- a/src/dstack/_internal/core/backends/oci/models.py +++ b/src/dstack/_internal/core/backends/oci/models.py @@ -78,8 +78,11 @@ class OCIBackendConfig(CoreModel): " 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" + " security groups, and it places the affected instances in a separate subnet that" + " has no OCI security list, so the network security group becomes the sole security" + " boundary. You are fully responsible for the network security group's rules," + " including ingress (e.g. SSH), egress (e.g. outbound access to pull Docker images)," + " and, for multi-node clusters, traffic between instances in the group" ) ), ] = None diff --git a/src/dstack/_internal/core/backends/oci/resources.py b/src/dstack/_internal/core/backends/oci/resources.py index 91e1d890f0..30e15ed07d 100644 --- a/src/dstack/_internal/core/backends/oci/resources.py +++ b/src/dstack/_internal/core/backends/oci/resources.py @@ -38,6 +38,9 @@ ADD_SECURITY_RULES_MAX_CHUNK_SIZE = 25 LIST_OBJECTS_MAX_LIMIT = 1000 VCN_CIDR = "10.0.0.0/16" +# CIDR for the dedicated VCN used for custom-NSG instances (see `get_or_create_restricted_vcn`). +# Must not overlap `VCN_CIDR` in case the two VCNs are ever peered. +RESTRICTED_VCN_CIDR = "10.1.0.0/16" WAIT_FOR_COMPARTMENT_ATTEMPS = 36 WAIT_FOR_COMPARTMENT_DELAY = 5 @@ -572,6 +575,89 @@ def get_or_create_subnet( ).data +def get_or_create_restricted_vcn( + name: str, compartment_id: str, client: oci.core.VirtualNetworkClient +) -> oci.core.models.Vcn: + """ + Like `get_or_create_vcn`, but a separate VCN (own CIDR block) dedicated to + custom-NSG instances. A *separate* VCN is used - rather than a second subnet + in the existing default VCN - because the default subnet already occupies + the default VCN's entire CIDR block, and adding a second CIDR block to an + existing VCN requires an async OCI operation (`add_vcn_cidr`) that takes + the VCN out of service for subnet/route-table updates for its duration. + A brand new VCN avoids that entirely and keeps the default VCN/subnet used + by dstack's auto-managed instances completely untouched. + """ + query_results = chain_paginated_responses( + client.list_vcns, compartment_id=compartment_id, display_name=name + ) + if vcn := next(query_results, None): + return vcn + + return client.create_vcn( + oci.core.models.CreateVcnDetails( + cidr_blocks=[RESTRICTED_VCN_CIDR], + compartment_id=compartment_id, + display_name=name, + ) + ).data + + +def get_or_create_restricted_subnet( + name: str, vcn_id: str, compartment_id: str, client: oci.core.VirtualNetworkClient +) -> oci.core.models.Subnet: + """ + Like `get_or_create_subnet`, but creates the subnet with an empty list of + security lists (`security_list_ids=[]`) instead of letting OCI attach the + VCN's permissive default security list. Must be created in a VCN returned + by `get_or_create_restricted_vcn`, not the default VCN. + + This is used for instances that run with a user-managed (custom) network + security group. With no security list contributing rules, the NSG becomes + the sole source of truth for what traffic is allowed to and from these + instances. OCI evaluates security lists and NSGs as a union of allows, so an + empty security list list simply means "the security-list layer grants + nothing"; it does not deny anything on its own. + """ + query_results = chain_paginated_responses( + client.list_subnets, compartment_id=compartment_id, display_name=name + ) + if subnet := next(query_results, None): + return subnet + + return client.create_subnet( + oci.core.models.CreateSubnetDetails( + cidr_block=RESTRICTED_VCN_CIDR, + compartment_id=compartment_id, + display_name=name, + vcn_id=vcn_id, + security_list_ids=[], + ) + ).data + + +def set_up_restricted_network_resources_in_region( + compartment_id: str, project_name: str, client: oci.core.VirtualNetworkClient +) -> oci.core.models.Subnet: + """ + Like `set_up_network_resources_in_region`, but for the dedicated VCN/subnet + used by custom-NSG instances (see `get_or_create_restricted_vcn` and + `get_or_create_restricted_subnet`). Idempotent - safe to call on every + instance launch, mirroring how `get_or_create_security_group` is already + called on every launch for the default (non-custom-NSG) path. + """ + vcn = get_or_create_restricted_vcn( + f"dstack-{project_name}-restricted-vcn", compartment_id, client + ) + internet_gateway = get_or_create_internet_gateway( + f"dstack-{project_name}-restricted-internet-gateway", vcn.id, compartment_id, client + ) + update_route_table(vcn.default_route_table_id, internet_gateway.id, client) + return get_or_create_restricted_subnet( + f"dstack-{project_name}-restricted-subnet", vcn.id, compartment_id, client + ) + + def get_or_create_internet_gateway( name: str, vcn_id: str, compartment_id: str, client: oci.core.VirtualNetworkClient ) -> oci.core.models.InternetGateway: diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index 04f4c326d8..5f2110b4d9 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -222,6 +222,7 @@ class Requirements(CoreModel): max_price: Optional[float] = None spot: Optional[bool] = None reservation: Optional[str] = None + security_group: Optional[str] = None multinode: Optional[bool] = None """Backends can use `multinode` to filter out offers when some offers support multinode and some do not. """ diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index a1f15e01dc..eed8154750 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -958,6 +958,7 @@ def get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements: max_price=profile.max_price, spot=get_policy_map(profile.spot_policy, default=SpotPolicy.ONDEMAND), reservation=fleet_spec.configuration.reservation, + security_group=fleet_spec.configuration.security_group, multinode=fleet_spec.configuration.placement == InstanceGroupPlacement.CLUSTER, backend_options=profile.backend_options, ) diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 761745247d..ea9eb70f0c 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -130,6 +130,9 @@ def _spot_policy(self) -> SpotPolicy: def _reservation(self) -> Optional[str]: return self.run_spec.merged_profile.reservation + def _security_group(self) -> Optional[str]: + return self.run_spec.merged_profile.security_group + @abstractmethod def _ports(self) -> List[PortMapping]: pass @@ -338,6 +341,7 @@ def _requirements(self, jobs_per_replica: int) -> Requirements: max_price=self.run_spec.merged_profile.max_price, spot=None if spot_policy == SpotPolicy.AUTO else (spot_policy == SpotPolicy.SPOT), reservation=self._reservation(), + security_group=self._security_group(), multinode=jobs_per_replica > 1, backend_options=self.run_spec.merged_profile.backend_options, ) diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 6fd739f13e..b7e0584698 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -12,6 +12,7 @@ BACKENDS_WITH_MULTINODE_SUPPORT, BACKENDS_WITH_PRIVILEGED_SUPPORT, BACKENDS_WITH_RESERVATION_SUPPORT, + BACKENDS_WITH_SECURITY_GROUP_SUPPORT, ) from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( @@ -73,6 +74,11 @@ async def get_offers_by_requirements( backend_types = BACKENDS_WITH_RESERVATION_SUPPORT backend_types = [b for b in backend_types if b in BACKENDS_WITH_RESERVATION_SUPPORT] + if requirements.security_group is not None: + if backend_types is None: + backend_types = BACKENDS_WITH_SECURITY_GROUP_SUPPORT + backend_types = [b for b in backend_types if b in BACKENDS_WITH_SECURITY_GROUP_SUPPORT] + # For multi-node, restrict backend and region. # The default behavior is to provision all nodes in the same backend and region. if master_job_provisioning_data is not None: diff --git a/src/dstack/_internal/server/services/requirements/combine.py b/src/dstack/_internal/server/services/requirements/combine.py index 53cec3457e..4fef3f3237 100644 --- a/src/dstack/_internal/server/services/requirements/combine.py +++ b/src/dstack/_internal/server/services/requirements/combine.py @@ -41,6 +41,9 @@ def combine_fleet_and_run_profiles( reservation=get_single_value_optional( fleet_profile.reservation, run_profile.reservation ), + security_group=get_single_value_optional( + fleet_profile.security_group, run_profile.security_group + ), spot_policy=_combine_spot_policy_optional( fleet_profile.spot_policy, run_profile.spot_policy ), @@ -68,6 +71,9 @@ def combine_fleet_and_run_requirements( reservation=get_single_value_optional( fleet_requirements.reservation, run_requirements.reservation ), + security_group=get_single_value_optional( + fleet_requirements.security_group, run_requirements.security_group + ), multinode=fleet_requirements.multinode or run_requirements.multinode, backend_options=_combine_backend_options_optional( fleet_requirements.backend_options, run_requirements.backend_options diff --git a/src/tests/_internal/core/backends/aws/test_compute.py b/src/tests/_internal/core/backends/aws/test_compute.py index 49392b9df5..dce562bd7c 100644 --- a/src/tests/_internal/core/backends/aws/test_compute.py +++ b/src/tests/_internal/core/backends/aws/test_compute.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock, patch +import botocore.exceptions import pytest from dstack._internal.core.backends.aws.compute import AWSCompute @@ -184,3 +185,29 @@ def test_run_level_security_group_overrides_name(self): assert security_group_id == "sg-run" compute._create_security_group.assert_not_called() lookup_mock.assert_not_called() + + def test_invalid_security_group_raises_compute_error(self): + # A security group in the wrong VPC/region is a misconfiguration, not a + # capacity issue: surface a clear ComputeError instead of retrying AZs. + compute = _compute(_config(security_group_ids={"us-east-1": "sg-wrong-vpc"})) + error = botocore.exceptions.ClientError( + error_response={ + "Error": { + "Code": "InvalidGroup.NotFound", + "Message": "The security group 'sg-wrong-vpc' does not exist", + } + }, + operation_name="RunInstances", + ) + with patch( + "dstack._internal.core.backends.aws.compute.aws_resources.create_instances_struct", + return_value={}, + ): + ec2_resource = compute.session.resource.return_value + ec2_resource.create_instances.side_effect = error + with pytest.raises(ComputeError, match="Security group not found"): + compute.create_instance( + instance_offer=_offer(), + instance_config=_instance_config(), + placement_group=None, + ) diff --git a/src/tests/_internal/core/backends/aws/test_configurator.py b/src/tests/_internal/core/backends/aws/test_configurator.py index f18582c1fe..bf8b37c3f3 100644 --- a/src/tests/_internal/core/backends/aws/test_configurator.py +++ b/src/tests/_internal/core/backends/aws/test_configurator.py @@ -7,6 +7,7 @@ from dstack._internal.core.errors import ( BackendAuthError, BackendInvalidCredentialsError, + ServerClientError, ) @@ -33,3 +34,69 @@ def test_validate_config_invalid_creds(self): authenticate_mock.side_effect = BackendAuthError() AWSConfigurator().validate_config(config, default_creds_enabled=True) assert exc_info.value.fields == [["creds", "access_key"], ["creds", "secret_key"]] + + def test_validate_config_security_group_name_and_ids_together(self): + # Combining `security_group_name` (cross-region catch-all) with a per-region + # `security_group_ids` map is supported and must not raise at validation time. + config = AWSBackendConfigWithCreds( + creds=AWSAccessKeyCreds(access_key="valid", secret_key="valid"), + regions=["us-east-1", "us-west-1"], + security_group_name="my-sg", + security_group_ids={"us-east-1": "sg-123"}, + ) + with ( + patch("dstack._internal.core.backends.aws.auth.authenticate"), + patch("dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error"), + ): + AWSConfigurator().validate_config(config, default_creds_enabled=True) + + def test_validate_config_security_group_ids_unknown_region_raises(self): + config = AWSBackendConfigWithCreds( + creds=AWSAccessKeyCreds(access_key="valid", secret_key="valid"), + regions=["us-east-1"], + security_group_ids={"us-east1": "sg-123"}, + ) + with ( + patch("dstack._internal.core.backends.aws.auth.authenticate"), + patch("dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error"), + pytest.raises(ServerClientError) as exc_info, + ): + AWSConfigurator().validate_config(config, default_creds_enabled=True) + assert "us-east1" in str(exc_info.value) + + def test_validate_config_security_group_ids_known_region_passes(self): + config = AWSBackendConfigWithCreds( + creds=AWSAccessKeyCreds(access_key="valid", secret_key="valid"), + regions=["us-east-1"], + security_group_ids={"us-east-1": "sg-123"}, + ) + with ( + patch("dstack._internal.core.backends.aws.auth.authenticate"), + patch("dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error"), + ): + AWSConfigurator().validate_config(config, default_creds_enabled=True) + + def test_validate_config_security_group_ids_default_regions_passes(self): + # `regions` unset -> validated against DEFAULT_REGIONS. + config = AWSBackendConfigWithCreds( + creds=AWSAccessKeyCreds(access_key="valid", secret_key="valid"), + security_group_ids={"us-east-1": "sg-123"}, + ) + with ( + patch("dstack._internal.core.backends.aws.auth.authenticate"), + patch("dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error"), + ): + AWSConfigurator().validate_config(config, default_creds_enabled=True) + + def test_validate_config_security_group_ids_partial_coverage_passes(self): + # Partial coverage is intentional (falls back to name/auto-create); must not raise. + config = AWSBackendConfigWithCreds( + creds=AWSAccessKeyCreds(access_key="valid", secret_key="valid"), + regions=["us-east-1", "us-west-1"], + security_group_ids={"us-east-1": "sg-123"}, + ) + with ( + patch("dstack._internal.core.backends.aws.auth.authenticate"), + patch("dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error"), + ): + AWSConfigurator().validate_config(config, default_creds_enabled=True) diff --git a/src/tests/_internal/core/backends/azure/test_compute.py b/src/tests/_internal/core/backends/azure/test_compute.py index b6a8975e97..022ff9d69f 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/src/tests/_internal/core/backends/azure/test_compute.py @@ -85,14 +85,14 @@ def test_get_image_name(self, variant: VMImageVariant, expected_name: str): assert variant.get_image_name() == expected_name -def _config(network_security_group_ids=None) -> AzureConfig: +def _config(network_security_group_names=None) -> AzureConfig: return AzureConfig( creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), tenant_id="ten1", subscription_id="sub1", resource_group="my-rg", regions=["eastus", "westeurope"], - network_security_group_ids=network_security_group_ids, + network_security_group_names=network_security_group_names, ) @@ -121,7 +121,7 @@ def _instance_config(security_group=None) -> InstanceConfiguration: class TestAzureComputeNetworkSecurityGroup: @pytest.mark.parametrize( - ["instance_sg", "nsg_ids", "region", "expected"], + ["instance_sg", "nsg_names", "region", "expected"], [ # No mapping, no instance security_group -> auto-derived default per location. [ @@ -146,7 +146,7 @@ class TestAzureComputeNetworkSecurityGroup: ], ) def test_create_instance_resolves_network_security_group( - self, instance_sg, nsg_ids, region, expected + self, instance_sg, nsg_names, region, expected ): with ( patch("dstack._internal.core.backends.azure.compute.compute_mgmt"), @@ -169,7 +169,7 @@ def test_create_instance_resolves_network_security_group( vm_mock.name = "test-vm-id" create_and_wait_mock.return_value = vm_mock compute = AzureCompute( - config=_config(network_security_group_ids=nsg_ids), credential=Mock() + config=_config(network_security_group_names=nsg_names), credential=Mock() ) compute.create_instance( instance_offer=_offer(region=region), diff --git a/src/tests/_internal/core/backends/azure/test_configurator.py b/src/tests/_internal/core/backends/azure/test_configurator.py index a327fb8985..e121f95e1d 100644 --- a/src/tests/_internal/core/backends/azure/test_configurator.py +++ b/src/tests/_internal/core/backends/azure/test_configurator.py @@ -168,7 +168,7 @@ def test_skips_instance_nsg_only_for_configured_locations(self): network_manager = self._create_backend( self._make_config( regions=["eastus", "westeurope"], - network_security_group_ids={"eastus": "my-nsg"}, + network_security_group_names={"eastus": "my-nsg"}, ) ) assert self._default_nsg_locations(network_manager) == ["westeurope"] @@ -179,7 +179,7 @@ def test_skips_instance_nsg_for_all_configured_locations(self): network_manager = self._create_backend( self._make_config( regions=["eastus", "westeurope"], - network_security_group_ids={ + network_security_group_names={ "eastus": "my-nsg", "westeurope": "my-other-nsg", }, @@ -187,3 +187,43 @@ def test_skips_instance_nsg_for_all_configured_locations(self): ) network_manager.create_network_security_group.assert_not_called() assert network_manager.create_gateway_network_security_group.call_count == 2 + + +class TestCheckConfigNetworkSecurityGroups: + def _make_config(self, **kwargs): + return AzureBackendConfigWithCreds( + creds=AzureClientCreds(tenant_id="t", client_id="c", client_secret="s"), + tenant_id="ten1", + subscription_id="sub1", + resource_group="my-rg", + **kwargs, + ) + + def test_unknown_location_raises(self): + config = self._make_config( + regions=["eastus", "westeurope"], + network_security_group_names={"eastu": "my-nsg"}, + ) + with pytest.raises(ServerClientError, match="eastu"): + AzureConfigurator()._check_config_network_security_groups(config) + + def test_configured_location_passes(self): + config = self._make_config( + regions=["eastus", "westeurope"], + network_security_group_names={"eastus": "my-nsg", "westeurope": "my-other-nsg"}, + ) + AzureConfigurator()._check_config_network_security_groups(config) + + def test_partial_coverage_passes(self): + config = self._make_config( + regions=["eastus", "westeurope"], + network_security_group_names={"eastus": "my-nsg"}, + ) + AzureConfigurator()._check_config_network_security_groups(config) + + def test_regions_none_skips_check(self): + config = self._make_config( + regions=None, + network_security_group_names={"eastus": "my-nsg"}, + ) + AzureConfigurator()._check_config_network_security_groups(config) diff --git a/src/tests/_internal/core/backends/gcp/test_compute.py b/src/tests/_internal/core/backends/gcp/test_compute.py index f7804be029..2888678a5d 100644 --- a/src/tests/_internal/core/backends/gcp/test_compute.py +++ b/src/tests/_internal/core/backends/gcp/test_compute.py @@ -85,6 +85,9 @@ def test_gateway_firewall_rules_created_when_explicitly_enabled(self): create_rules_mock = _run_create_gateway(_make_compute(create_firewall_rules=True)) create_rules_mock.assert_called_once() - def test_gateway_firewall_rules_skipped_when_disabled(self): + def test_gateway_firewall_rules_not_affected_by_create_firewall_rules(self): + # `create_firewall_rules` only gates the runner (instance) firewall rule, not the + # gateway one - gateways are meant to be internet-reachable and are always auto-managed, + # consistent with how gateway security groups/NSGs are handled on AWS/Azure/OCI. create_rules_mock = _run_create_gateway(_make_compute(create_firewall_rules=False)) - create_rules_mock.assert_not_called() + create_rules_mock.assert_called_once() diff --git a/src/tests/_internal/core/backends/oci/test_compute.py b/src/tests/_internal/core/backends/oci/test_compute.py index 8ced03f83a..6a61eb9d89 100644 --- a/src/tests/_internal/core/backends/oci/test_compute.py +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -1,7 +1,5 @@ from unittest.mock import MagicMock, patch -import pytest - from dstack._internal.core.backends.oci.compute import OCICompute from dstack._internal.core.backends.oci.models import OCIClientCreds, OCIConfig from dstack._internal.core.models.backends.base import BackendType @@ -73,8 +71,12 @@ class TestOCIComputeSecurityGroup: def _run_create_instance(self, compute, instance_config): with patch("dstack._internal.core.backends.oci.compute.resources") as res: res.VCN_CIDR = "10.0.0.0/16" + res.RESTRICTED_VCN_CIDR = "10.1.0.0/16" res.get_marketplace_listing_and_package.return_value = (MagicMock(), MagicMock()) res.get_or_create_security_group.return_value.id = "ocid1.nsg.oc1..managed" + res.set_up_restricted_network_resources_in_region.return_value.id = ( + "ocid1.subnet.oc1..restricted" + ) res.launch_instance.return_value.id = "ocid1.instance.oc1..instance" compute.create_instance(_make_offer(), instance_config, placement_group=None) return res @@ -89,6 +91,12 @@ def test_default_creates_and_syncs_managed_security_group(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) + # The default/auto-managed-NSG path uses the original default subnet + # and never touches the restricted VCN/subnet. + res.set_up_restricted_network_resources_in_region.assert_not_called() + assert ( + res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" + ) def test_per_region_custom_nsg_is_left_untouched(self): compute = _make_compute( @@ -105,6 +113,17 @@ def test_per_region_custom_nsg_is_left_untouched(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..custom" ) + # A custom NSG routes the instance into a dedicated restricted VCN/subnet + # (no security list), never the default one. + res.set_up_restricted_network_resources_in_region.assert_called_once() + assert ( + res.set_up_restricted_network_resources_in_region.call_args.kwargs["project_name"] + == "test-project" + ) + assert ( + res.launch_instance.call_args.kwargs["subnet_id"] + == "ocid1.subnet.oc1..restricted" + ) def test_region_not_in_mapping_falls_back_to_managed(self): compute = _make_compute( @@ -120,6 +139,10 @@ def test_region_not_in_mapping_falls_back_to_managed(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) + res.set_up_restricted_network_resources_in_region.assert_not_called() + assert ( + res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" + ) def test_instance_level_custom_nsg_is_left_untouched(self): compute = _make_compute(_make_config()) @@ -133,6 +156,11 @@ def test_instance_level_custom_nsg_is_left_untouched(self): assert ( res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" ) + res.set_up_restricted_network_resources_in_region.assert_called_once() + assert ( + res.launch_instance.call_args.kwargs["subnet_id"] + == "ocid1.subnet.oc1..restricted" + ) def test_instance_level_overrides_per_region_mapping(self): compute = _make_compute( @@ -149,7 +177,50 @@ def test_instance_level_overrides_per_region_mapping(self): assert ( res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" ) + res.set_up_restricted_network_resources_in_region.assert_called_once() + assert ( + res.launch_instance.call_args.kwargs["subnet_id"] + == "ocid1.subnet.oc1..restricted" + ) + + +class TestGetOrCreateRestrictedSubnet: + def test_creates_subnet_with_empty_security_list_ids(self): + from dstack._internal.core.backends.oci import resources + + client = MagicMock() + client.list_subnets.return_value.data = [] + client.list_subnets.return_value.next_page = None + client.list_subnets.return_value.has_next_page = False + + resources.get_or_create_restricted_subnet( + "dstack-test-project-restricted-subnet", + "ocid1.vcn.oc1..vcn", + "ocid1.compartment.oc1..compartment", + client, + ) + client.create_subnet.assert_called_once() + details = client.create_subnet.call_args.args[0] + assert details.security_list_ids == [] + assert details.vcn_id == "ocid1.vcn.oc1..vcn" + assert details.display_name == "dstack-test-project-restricted-subnet" + + def test_returns_existing_subnet_without_creating(self): + from dstack._internal.core.backends.oci import resources + + existing = MagicMock() + client = MagicMock() + client.list_subnets.return_value.data = [existing] + client.list_subnets.return_value.next_page = None + client.list_subnets.return_value.has_next_page = False + + result = resources.get_or_create_restricted_subnet( + "dstack-test-project-restricted-subnet", + "ocid1.vcn.oc1..vcn", + "ocid1.compartment.oc1..compartment", + client, + ) -if __name__ == "__main__": - pytest.main([__file__, "-v"]) + assert result is existing + client.create_subnet.assert_not_called() diff --git a/src/tests/_internal/core/backends/oci/test_configurator.py b/src/tests/_internal/core/backends/oci/test_configurator.py index 5472b3048f..380be828f9 100644 --- a/src/tests/_internal/core/backends/oci/test_configurator.py +++ b/src/tests/_internal/core/backends/oci/test_configurator.py @@ -8,7 +8,10 @@ OCIBackendConfigWithCreds, OCIClientCreds, ) -from dstack._internal.core.errors import BackendInvalidCredentialsError +from dstack._internal.core.errors import ( + BackendInvalidCredentialsError, + ServerClientError, +) class TestOCIConfigurator: @@ -53,3 +56,66 @@ def test_validate_config_invalid_creds(self): regions_mock.side_effect = ClientError("Invalid credentials") OCIConfigurator().validate_config(config, default_creds_enabled=True) assert exc_info.value.fields == [["creds"]] + + def test_validate_config_nsg_region_typo(self): + config = OCIBackendConfigWithCreds( + creds=OCIClientCreds( + user="valid_user", + tenancy="valid_tenancy", + key_content="valid_key", + key_file=None, + pass_phrase=None, + fingerprint="valid_fingerprint", + region="us-ashburn-1", + ), + regions=["us-ashburn-1"], + network_security_group_ids={"us-ashburn-typo": "ocid1.nsg.oc1..custom"}, + ) + with ( + patch( + "dstack._internal.core.backends.oci.configurator.get_subscribed_regions" + ) as regions_mock, + pytest.raises(ServerClientError) as exc_info, + ): + regions_mock.return_value = Mock(names=["us-ashburn-1"]) + OCIConfigurator().validate_config(config, default_creds_enabled=True) + assert exc_info.value.fields == [["network_security_group_ids"]] + + def test_validate_config_nsg_region_valid(self): + config = OCIBackendConfigWithCreds( + creds=OCIClientCreds( + user="valid_user", + tenancy="valid_tenancy", + key_content="valid_key", + key_file=None, + pass_phrase=None, + fingerprint="valid_fingerprint", + region="us-ashburn-1", + ), + regions=["us-ashburn-1"], + network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..custom"}, + ) + with patch( + "dstack._internal.core.backends.oci.configurator.get_subscribed_regions" + ) as regions_mock: + regions_mock.return_value = Mock(names=["us-ashburn-1"]) + OCIConfigurator().validate_config(config, default_creds_enabled=True) + + def test_validate_config_nsg_no_regions_skips_check(self): + config = OCIBackendConfigWithCreds( + creds=OCIClientCreds( + user="valid_user", + tenancy="valid_tenancy", + key_content="valid_key", + key_file=None, + pass_phrase=None, + fingerprint="valid_fingerprint", + region="us-ashburn-1", + ), + network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..custom"}, + ) + with patch( + "dstack._internal.core.backends.oci.configurator.get_subscribed_regions" + ) as regions_mock: + regions_mock.return_value = Mock(names=["us-ashburn-1"]) + OCIConfigurator().validate_config(config, default_creds_enabled=True) diff --git a/src/tests/_internal/server/routers/test_backends.py b/src/tests/_internal/server/routers/test_backends.py index 9b548ab297..10db9d670d 100644 --- a/src/tests/_internal/server/routers/test_backends.py +++ b/src/tests/_internal/server/routers/test_backends.py @@ -828,6 +828,8 @@ async def test_returns_config_info(self, test_db, session: AsyncSession, client: "iam_instance_profile": None, "tags": None, "os_images": None, + "security_group_name": None, + "security_group_ids": None, "creds": json.loads(backend.auth.get_plaintext_or_error()), } diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index 04d0145dfe..54456b5b96 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/src/tests/_internal/server/routers/test_fleets.py @@ -949,6 +949,7 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async "type": "fleet", "name": "test-fleet", "reservation": None, + "security_group": None, "blocks": 1, "tags": None, "backend_options": None, @@ -972,6 +973,7 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async "name": "", "default": False, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -1070,6 +1072,7 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A "type": "fleet", "name": spec.configuration.name, "reservation": None, + "security_group": None, "blocks": 1, "tags": None, "backend_options": None, @@ -1093,6 +1096,7 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A "name": "", "default": False, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -1290,6 +1294,7 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "type": "fleet", "name": spec.configuration.name, "reservation": None, + "security_group": None, "blocks": 1, "tags": None, "backend_options": None, @@ -1313,6 +1318,7 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "name": "", "default": False, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index fdc5a8f355..299b70df90 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -215,6 +215,7 @@ def get_dev_env_run_plan_dict( "stop_criteria": None, "schedule": None, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -242,6 +243,7 @@ def get_dev_env_run_plan_dict( "stop_criteria": None, "schedule": None, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -298,6 +300,7 @@ def get_dev_env_run_plan_dict( "max_price": None, "spot": None, "reservation": None, + "security_group": None, "multinode": False, "backend_options": None, }, @@ -466,6 +469,7 @@ def get_dev_env_run_dict( "stop_criteria": None, "schedule": None, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -493,6 +497,7 @@ def get_dev_env_run_dict( "stop_criteria": None, "schedule": None, "reservation": None, + "security_group": None, "fleets": None, "tags": None, "backend_options": None, @@ -544,6 +549,7 @@ def get_dev_env_run_dict( "max_price": None, "spot": None, "reservation": None, + "security_group": None, "multinode": False, "backend_options": None, }, diff --git a/src/tests/_internal/server/services/requirements/test_combine.py b/src/tests/_internal/server/services/requirements/test_combine.py index 3680161e83..4c2339f427 100644 --- a/src/tests/_internal/server/services/requirements/test_combine.py +++ b/src/tests/_internal/server/services/requirements/test_combine.py @@ -132,6 +132,37 @@ def test_combines_profiles( ): assert combine_fleet_and_run_profiles(fleet_profile, run_profile) == expected_profile + @pytest.mark.parametrize( + argnames=["fleet_value", "run_value", "expected_value"], + argvalues=[ + pytest.param(None, None, None, id="both_unset"), + pytest.param("sg-1", None, "sg-1", id="only_fleet_set"), + pytest.param(None, "sg-2", "sg-2", id="only_run_set"), + pytest.param("sg-1", "sg-1", "sg-1", id="both_same"), + ], + ) + def test_combines_security_group( + self, + fleet_value: Optional[str], + run_value: Optional[str], + expected_value: Optional[str], + ): + combined = combine_fleet_and_run_profiles( + Profile(security_group=fleet_value), + Profile(security_group=run_value), + ) + assert combined is not None + assert combined.security_group == expected_value + + def test_incompatible_security_group_returns_none(self): + assert ( + combine_fleet_and_run_profiles( + Profile(security_group="sg-1"), + Profile(security_group="sg-2"), + ) + is None + ) + class TestCombineFleetAndRunRequirements: def test_returns_the_same_requirements_if_requirements_identical(self): @@ -212,6 +243,59 @@ def test_combines_requirements( == expected_requirements ) + @pytest.mark.parametrize( + argnames=["fleet_value", "run_value", "expected_value"], + argvalues=[ + pytest.param(None, None, None, id="both_unset"), + pytest.param("sg-1", None, "sg-1", id="only_fleet_set"), + pytest.param(None, "sg-2", "sg-2", id="only_run_set"), + pytest.param("sg-1", "sg-1", "sg-1", id="both_same"), + ], + ) + def test_combines_security_group( + self, + fleet_value: Optional[str], + run_value: Optional[str], + expected_value: Optional[str], + ): + combined = combine_fleet_and_run_requirements( + Requirements(resources=ResourcesSpec(), security_group=fleet_value), + Requirements(resources=ResourcesSpec(), security_group=run_value), + ) + assert combined is not None + assert combined.security_group == expected_value + + def test_incompatible_security_group_returns_none(self): + assert ( + combine_fleet_and_run_requirements( + Requirements(resources=ResourcesSpec(), security_group="sg-1"), + Requirements(resources=ResourcesSpec(), security_group="sg-2"), + ) + is None + ) + + def test_fleet_security_group_reaches_run_without_own_security_group(self): + # Regression test for Bug 1: a run provisioned into a fleet must inherit the + # fleet's `security_group` even when the run does not set its own. + from dstack._internal.core.models.fleets import FleetConfiguration, FleetNodesSpec + from dstack._internal.server.services.fleets import get_fleet_requirements + from dstack._internal.server.testing.common import get_fleet_spec + + fleet_spec = get_fleet_spec( + conf=FleetConfiguration( + name="test-fleet", + nodes=FleetNodesSpec(min=1, target=1, max=1), + security_group="sg-fleet", + ) + ) + fleet_requirements = get_fleet_requirements(fleet_spec) + assert fleet_requirements.security_group == "sg-fleet" + + run_requirements = Requirements(resources=ResourcesSpec()) + combined = combine_fleet_and_run_requirements(fleet_requirements, run_requirements) + assert combined is not None + assert combined.security_group == "sg-fleet" + def test_unconstrained_fleet_resources_pass_through_run_requirements(self): unconstrained_fleet = Requirements( resources=ResourcesSpec.unconstrained(), diff --git a/src/tests/_internal/server/services/test_offers.py b/src/tests/_internal/server/services/test_offers.py index 25ce8021ae..ac04cc77ce 100644 --- a/src/tests/_internal/server/services/test_offers.py +++ b/src/tests/_internal/server/services/test_offers.py @@ -61,6 +61,28 @@ async def test_returns_multinode_offers(self): m.assert_awaited_once() assert res == [(aws_backend_mock, aws_offer)] + @pytest.mark.asyncio + async def test_returns_only_security_group_supporting_offers(self): + profile = Profile(name="test") + requirements = Requirements(resources=ResourcesSpec(), security_group="sg-1") + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + aws_backend_mock = Mock() + aws_backend_mock.TYPE = BackendType.AWS + aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS) + aws_backend_mock.compute.return_value.get_offers.return_value = [aws_offer] + runpod_backend_mock = Mock() + runpod_backend_mock.TYPE = BackendType.RUNPOD + runpod_offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD) + runpod_backend_mock.compute.return_value.get_offers.return_value = [runpod_offer] + m.return_value = [aws_backend_mock, runpod_backend_mock] + res = await get_offers_by_requirements( + project=Mock(), + profile=profile, + requirements=requirements, + ) + m.assert_awaited_once() + assert res == [(aws_backend_mock, aws_offer)] + @pytest.mark.asyncio async def test_returns_volume_offers(self): profile = Profile(name="test") From 8434813d60d04a07aa55acdc047626cf608ef3e9 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 13:55:29 -0400 Subject: [PATCH 4/6] Fix issues found in second independent review - run_job(): source reservation/security_group from the `requirements` parameter (already fleet+run-combined) instead of job.job_spec.requirements (run-only). Without this, a run provisioning new capacity into a fleet with a fleet-level security_group/reservation would silently ignore it - the same latent bug reservation already had, inherited by security_group. - Azure create_gateway(): use get_gateway_network_security_group_name (the dedicated, always-created gateway NSG) instead of get_default_network_security_group_name. The default/per-location instance NSG can now be skipped when network_security_group_names covers that location, which would have broken gateway provisioning since it was referencing an NSG that might not exist. - OCI: redesign custom-NSG networking. The previous "separate restricted VCN" approach is fundamentally broken - OCI network security groups are VCN-scoped, so a user's NSG can never be attached to an instance in a different VCN than the one the NSG lives in. Fixed by using a single shared subnet for all instances (default-NSG and custom-NSG alike), with no OCI security list attached. dstack's auto-managed NSG now carries explicit SSH ingress and all-egress rules to compensate for the removed security list; custom NSGs remain fully hands-off, per the feature's contract. Existing subnets are migrated in place (security list detached) since the subnet is dstack-owned infrastructure, not a user-supplied resource. Updated docs to clarify a custom NSG must live in dstack's own default VCN. Regression: 388 core backend tests + base/azure/oci targeted suites + server routers/services (requirements, offers, fleets, runs, backends) all passing. Pre-existing unrelated failures (verda/vastai/nebius modules not installed in this environment) left untouched. --- mkdocs/docs/concepts/backends.md | 29 ++-- .../_internal/core/backends/azure/compute.py | 7 +- .../_internal/core/backends/base/compute.py | 4 +- .../_internal/core/backends/oci/compute.py | 27 ++-- .../_internal/core/backends/oci/models.py | 15 +- .../_internal/core/backends/oci/resources.py | 144 +++++++----------- .../core/backends/azure/test_compute.py | 54 +++++++ .../core/backends/base/test_compute.py | 90 +++++++++++ .../core/backends/oci/test_compute.py | 141 +++++++++++++---- 9 files changed, 356 insertions(+), 155 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index cefa4a0cce..a7f26b36d3 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1162,12 +1162,11 @@ There are two ways to configure OCI: using client credentials or using the defau ``` ??? info "Custom network security group" - By default, `dstack` places instances in a shared subnet whose OCI security list opens SSH - (TCP port 22) to `0.0.0.0/0` and permits all outbound traffic. On top of that, `dstack` creates - and manages its own network security group (NSG) per project, which only adds a rule allowing all - traffic within the group so multi-node clusters work out of the box. In other words, the - permissive SSH ingress and the outbound access come from the subnet's security list, not from the - auto-managed NSG. + 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`: @@ -1185,14 +1184,16 @@ There are two ways to configure OCI: using client credentials or using the defau ``` Regions not covered by `network_security_group_ids` fall back to dstack's auto-created network - security group and shared subnet. - - When a custom NSG is used, `dstack` never adds, removes, or modifies its rules, and it places the - affected instances in a separate VCN and subnet that has **no** OCI security list (a second subnet - in the same VCN isn't possible here, since the default subnet already occupies the whole VCN's - address space). This makes the NSG the sole security boundary — there is no longer any implicit - SSH-from-anywhere or implicit outbound-all coming from a security list. As a result, your custom - NSG is fully responsible for: + 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--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 diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index b219560702..fb24addd21 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -262,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, ) diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 35589381b6..fd0d20c6ad 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -320,8 +320,8 @@ def run_job( user=run.user, ssh_keys=[SSHKey(public=project_ssh_public_key.strip())], volumes=volumes, - reservation=job.job_spec.requirements.reservation, - security_group=job.job_spec.requirements.security_group, + reservation=requirements.reservation, + security_group=requirements.security_group, tags=run.run_spec.merged_profile.tags, ) instance_offer = instance_offer.copy() diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index b2f2f9ae88..a46e90f01c 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -135,6 +135,13 @@ def create_instance( listing, self.config.compartment_id, region.marketplace_client ) + # All instances - whether using dstack's auto-managed NSG or a custom + # one - share the same subnet. This is required because OCI NSGs are + # VCN-scoped: an NSG can only be attached to a VNIC whose subnet + # belongs to the same VCN the NSG lives in, and dstack only manages + # one VCN per project. The subnet has no security list attached (see + # `get_or_create_subnet`), so the NSG attached to each instance is its + # sole security boundary. subnet: oci.core.models.Subnet = region.virtual_network_client.get_subnet( self.config.subnet_ids_per_region[instance_offer.region] ).data @@ -143,8 +150,7 @@ def create_instance( security_group_id = self.config.network_security_group_ids.get( instance_offer.region ) - using_custom_security_group = security_group_id is not None - if not using_custom_security_group: + if security_group_id is None: security_group = resources.get_or_create_security_group( f"dstack-{instance_config.project_name}-default-security-group", subnet.vcn_id, @@ -155,25 +161,10 @@ def create_instance( security_group.id, region.virtual_network_client ) security_group_id = security_group.id - firewall_allow_from_subnet = resources.VCN_CIDR - else: - # A user-managed (custom) NSG is in use. Place the instance in a - # dedicated VCN/subnet that has no OCI security list, so the NSG is - # the sole security boundary. This is a *separate* VCN, not a second - # subnet in the default one: the default subnet already occupies the - # default VCN's entire CIDR block, so a second subnet there would - # conflict. The default VCN/subnet is left completely untouched, - # so instances using dstack's auto-managed NSG are unaffected. - subnet = resources.set_up_restricted_network_resources_in_region( - compartment_id=self.config.compartment_id, - project_name=instance_config.project_name, - client=region.virtual_network_client, - ) - firewall_allow_from_subnet = resources.RESTRICTED_VCN_CIDR cloud_init_user_data = get_user_data( authorized_keys=instance_config.get_public_keys(), - firewall_allow_from_subnets=[firewall_allow_from_subnet], + firewall_allow_from_subnets=[resources.VCN_CIDR], ) display_name = generate_unique_instance_name(instance_config) diff --git a/src/dstack/_internal/core/backends/oci/models.py b/src/dstack/_internal/core/backends/oci/models.py index b30530125b..ece575467f 100644 --- a/src/dstack/_internal/core/backends/oci/models.py +++ b/src/dstack/_internal/core/backends/oci/models.py @@ -77,12 +77,17 @@ class OCIBackendConfig(CoreModel): " 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." + " Because OCI requires a network security group to belong to the same VCN as the" + " instances it's attached to, and `dstack` provisions instances into its own" + " default VCN for each project, the network security group must be created in that" + " VCN (named `dstack--default-vcn`)." + " `dstack`'s default subnet in that VCN has no OCI security list attached, so" + " whichever network security group is attached to an instance - yours or dstack's" + " own - is the sole security boundary for it." " When set, `dstack` does not add, remove, or modify any rules on these network" - " security groups, and it places the affected instances in a separate subnet that" - " has no OCI security list, so the network security group becomes the sole security" - " boundary. You are fully responsible for the network security group's rules," - " including ingress (e.g. SSH), egress (e.g. outbound access to pull Docker images)," - " and, for multi-node clusters, traffic between instances in the group" + " security groups. You are fully responsible for their rules, including ingress" + " (e.g. SSH), egress (e.g. outbound access to pull Docker images), and, for" + " multi-node clusters, traffic between instances in the group" ) ), ] = None diff --git a/src/dstack/_internal/core/backends/oci/resources.py b/src/dstack/_internal/core/backends/oci/resources.py index 30e15ed07d..b9ab80dee5 100644 --- a/src/dstack/_internal/core/backends/oci/resources.py +++ b/src/dstack/_internal/core/backends/oci/resources.py @@ -38,9 +38,6 @@ ADD_SECURITY_RULES_MAX_CHUNK_SIZE = 25 LIST_OBJECTS_MAX_LIMIT = 1000 VCN_CIDR = "10.0.0.0/16" -# CIDR for the dedicated VCN used for custom-NSG instances (see `get_or_create_restricted_vcn`). -# Must not overlap `VCN_CIDR` in case the two VCNs are ever peered. -RESTRICTED_VCN_CIDR = "10.1.0.0/16" WAIT_FOR_COMPARTMENT_ATTEMPS = 36 WAIT_FOR_COMPARTMENT_DELAY = 5 @@ -559,75 +556,52 @@ def get_or_create_vcn( def get_or_create_subnet( name: str, vcn_id: str, compartment_id: str, client: oci.core.VirtualNetworkClient ) -> oci.core.models.Subnet: - query_results = chain_paginated_responses( - client.list_subnets, compartment_id=compartment_id, display_name=name - ) - if subnet := next(query_results, None): - return subnet - - return client.create_subnet( - oci.core.models.CreateSubnetDetails( - cidr_block=VCN_CIDR, - compartment_id=compartment_id, - display_name=name, - vcn_id=vcn_id, - ) - ).data - - -def get_or_create_restricted_vcn( - name: str, compartment_id: str, client: oci.core.VirtualNetworkClient -) -> oci.core.models.Vcn: """ - Like `get_or_create_vcn`, but a separate VCN (own CIDR block) dedicated to - custom-NSG instances. A *separate* VCN is used - rather than a second subnet - in the existing default VCN - because the default subnet already occupies - the default VCN's entire CIDR block, and adding a second CIDR block to an - existing VCN requires an async OCI operation (`add_vcn_cidr`) that takes - the VCN out of service for subnet/route-table updates for its duration. - A brand new VCN avoids that entirely and keeps the default VCN/subnet used - by dstack's auto-managed instances completely untouched. - """ - query_results = chain_paginated_responses( - client.list_vcns, compartment_id=compartment_id, display_name=name - ) - if vcn := next(query_results, None): - return vcn - - return client.create_vcn( - oci.core.models.CreateVcnDetails( - cidr_blocks=[RESTRICTED_VCN_CIDR], - compartment_id=compartment_id, - display_name=name, - ) - ).data - - -def get_or_create_restricted_subnet( - name: str, vcn_id: str, compartment_id: str, client: oci.core.VirtualNetworkClient -) -> oci.core.models.Subnet: - """ - Like `get_or_create_subnet`, but creates the subnet with an empty list of - security lists (`security_list_ids=[]`) instead of letting OCI attach the - VCN's permissive default security list. Must be created in a VCN returned - by `get_or_create_restricted_vcn`, not the default VCN. - - This is used for instances that run with a user-managed (custom) network - security group. With no security list contributing rules, the NSG becomes - the sole source of truth for what traffic is allowed to and from these - instances. OCI evaluates security lists and NSGs as a union of allows, so an - empty security list list simply means "the security-list layer grants - nothing"; it does not deny anything on its own. + The subnet is created with an empty list of security lists + (`security_list_ids=[]`) instead of letting OCI attach the VCN's permissive + default security list. All instances - whether they use dstack's + auto-managed network security group (NSG) or a user-supplied custom one - + live in this single shared subnet, so the NSG attached to each instance's + VNIC is the sole security boundary; there is no separate security-list + layer to reason about. + + A single shared subnet (rather than a second, NSG-only subnet) is required + because OCI NSGs are scoped to a single VCN: an NSG can only be attached to + a VNIC whose subnet belongs to the *same* VCN the NSG was created in. Since + dstack only manages one VCN, a user's custom NSG must live in - and a + custom-NSG instance must therefore be placed in a subnet within - that same + VCN. + + For instances using dstack's auto-managed NSG, the rules normally + contributed by the security list (SSH ingress, unrestricted egress) are + added directly to that NSG instead; see + `update_security_group_rules_for_runner_instances`. Instances using a + user-supplied NSG get no such compensating rules - per dstack's "fully + hands-off" contract, it never adds, removes, or otherwise modifies rules on + a user-supplied security group, so the user is fully responsible for + allowing the traffic their instances need (including SSH). """ query_results = chain_paginated_responses( client.list_subnets, compartment_id=compartment_id, display_name=name ) if subnet := next(query_results, None): + if subnet.security_list_ids: + # A subnet created before this fix still has the VCN's default + # security list attached. Since dstack owns this subnet + # (it is not user-supplied), it's safe to update it in place - + # detaching the security list so the NSG becomes the sole + # security boundary for every instance in it, matching newly + # created subnets. This is unrelated to dstack's "fully + # hands-off" contract for user-supplied *security groups*, which + # this does not touch. + subnet = client.update_subnet( + subnet.id, oci.core.models.UpdateSubnetDetails(security_list_ids=[]) + ).data return subnet return client.create_subnet( oci.core.models.CreateSubnetDetails( - cidr_block=RESTRICTED_VCN_CIDR, + cidr_block=VCN_CIDR, compartment_id=compartment_id, display_name=name, vcn_id=vcn_id, @@ -636,28 +610,6 @@ def get_or_create_restricted_subnet( ).data -def set_up_restricted_network_resources_in_region( - compartment_id: str, project_name: str, client: oci.core.VirtualNetworkClient -) -> oci.core.models.Subnet: - """ - Like `set_up_network_resources_in_region`, but for the dedicated VCN/subnet - used by custom-NSG instances (see `get_or_create_restricted_vcn` and - `get_or_create_restricted_subnet`). Idempotent - safe to call on every - instance launch, mirroring how `get_or_create_security_group` is already - called on every launch for the default (non-custom-NSG) path. - """ - vcn = get_or_create_restricted_vcn( - f"dstack-{project_name}-restricted-vcn", compartment_id, client - ) - internet_gateway = get_or_create_internet_gateway( - f"dstack-{project_name}-restricted-internet-gateway", vcn.id, compartment_id, client - ) - update_route_table(vcn.default_route_table_id, internet_gateway.id, client) - return get_or_create_restricted_subnet( - f"dstack-{project_name}-restricted-subnet", vcn.id, compartment_id, client - ) - - def get_or_create_internet_gateway( name: str, vcn_id: str, compartment_id: str, client: oci.core.VirtualNetworkClient ) -> oci.core.models.InternetGateway: @@ -714,8 +666,11 @@ def get_or_create_security_group( def update_security_group_rules_for_runner_instances( security_group_id: str, client: oci.core.VirtualNetworkClient ) -> None: - # These rules are combined with subnet's default Security List that allows - # ingress TCP on port 22 from anywhere + # The subnet these instances live in has no security list attached (see + # `get_or_create_subnet`), so this NSG must grant everything a runner + # instance needs on its own: SSH ingress from anywhere and unrestricted + # egress, in addition to allowing all traffic between instances that share + # this NSG. rules = [ SecurityRule( description="Allow all traffic within this security group", @@ -724,6 +679,23 @@ def update_security_group_rules_for_runner_instances( source=security_group_id, protocol="all", ), + SecurityRule( + description="Allow SSH ingress from anywhere", + direction=oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS, + source_type=oci.core.models.AddSecurityRuleDetails.SOURCE_TYPE_CIDR_BLOCK, + source="0.0.0.0/0", + protocol="6", # TCP + tcp_options=oci.core.models.TcpOptions( + destination_port_range=oci.core.models.PortRange(min=22, max=22) + ), + ), + SecurityRule( + description="Allow all egress traffic", + direction=oci.core.models.AddSecurityRuleDetails.DIRECTION_EGRESS, + destination_type=oci.core.models.AddSecurityRuleDetails.DESTINATION_TYPE_CIDR_BLOCK, + destination="0.0.0.0/0", + protocol="all", + ), ] update_security_group_rules(security_group_id, rules, client) diff --git a/src/tests/_internal/core/backends/azure/test_compute.py b/src/tests/_internal/core/backends/azure/test_compute.py index 022ff9d69f..68d3bb2a22 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/src/tests/_internal/core/backends/azure/test_compute.py @@ -15,6 +15,7 @@ Resources, SSHKey, ) +from dstack._internal.server.testing.common import get_gateway_compute_configuration class TestVMImageVariant: @@ -178,3 +179,56 @@ def test_create_instance_resolves_network_security_group( ) _, kwargs = create_and_wait_mock.call_args assert kwargs["network_security_group"] == expected + + +class TestAzureComputeGatewayNetworkSecurityGroup: + @pytest.mark.parametrize( + ["nsg_names", "region"], + [ + # No custom NSG configured anywhere. + [None, "eastus"], + # The instance-level default NSG is skipped for this exact location + # because a custom one is configured - the gateway must still use its + # own, always-created NSG, not the (now-nonexistent) instance default. + [{"eastus": "config-nsg"}, "eastus"], + [{"eastus": "config-nsg"}, "westeurope"], + ], + ) + def test_create_gateway_always_uses_gateway_network_security_group(self, nsg_names, region): + with ( + patch("dstack._internal.core.backends.azure.compute.compute_mgmt"), + patch("dstack._internal.core.backends.azure.compute.network_mgmt"), + patch( + "dstack._internal.core.backends.azure.compute" + ".get_resource_group_network_subnet_or_error", + return_value=("net-rg", "net", "subnet"), + ), + patch("dstack._internal.core.backends.azure.compute._get_gateway_image_ref"), + patch( + "dstack._internal.core.backends.azure.compute._create_instance_and_wait" + ) as create_and_wait_mock, + patch( + "dstack._internal.core.backends.azure.compute.get_gateway_user_data", + return_value="", + ), + patch( + "dstack._internal.core.backends.azure.compute._get_vm_public_private_ips", + return_value=("1.2.3.4", "10.0.0.1"), + ), + ): + vm_mock = Mock() + vm_mock.name = "test-gateway-vm-id" + create_and_wait_mock.return_value = vm_mock + compute = AzureCompute( + config=_config(network_security_group_names=nsg_names), credential=Mock() + ) + compute.create_gateway( + get_gateway_compute_configuration(region=region, backend="azure") + ) + _, kwargs = create_and_wait_mock.call_args + # Never the per-location default/custom instance NSG - always the + # dedicated gateway NSG, which is created unconditionally regardless + # of `network_security_group_names`. + assert kwargs["network_security_group"] == ( + azure_utils.get_gateway_network_security_group_name("my-rg", region) + ) diff --git a/src/tests/_internal/core/backends/base/test_compute.py b/src/tests/_internal/core/backends/base/test_compute.py index 7892a3f0f5..ab3308e8ec 100644 --- a/src/tests/_internal/core/backends/base/test_compute.py +++ b/src/tests/_internal/core/backends/base/test_compute.py @@ -1,10 +1,12 @@ import re from typing import Optional +from unittest.mock import MagicMock import gpuhunt import pytest from dstack._internal.core.backends.base.compute import ( + ComputeWithCreateInstanceSupport, GoArchType, generate_unique_backend_name, generate_unique_gateway_instance_name, @@ -12,6 +14,9 @@ generate_unique_volume_name, normalize_arch, ) +from dstack._internal.core.models.instances import InstanceConfiguration +from dstack._internal.core.models.resources import CPUSpec, ResourcesSpec +from dstack._internal.core.models.runs import Requirements from dstack._internal.server.testing.common import ( get_gateway_compute_configuration, get_instance_configuration, @@ -19,6 +24,91 @@ ) +class _FakeCreateInstanceCompute(ComputeWithCreateInstanceSupport): + """Minimal Compute stub that just records the `InstanceConfiguration` it was given.""" + + last_instance_config: Optional[InstanceConfiguration] = None + + def create_instance(self, instance_offer, instance_config, placement_group): + self.last_instance_config = instance_config + # `run_job()`'s return value isn't exercised by this test - it's returned + # as-is by `create_instance`, with no validation in `run_job()` itself. + return MagicMock() + + +class TestRunJobSourcesFromEffectiveRequirements: + """ + `run_job()` is called with an already fleet+run-combined `Requirements` object + (see `_get_effective_profile_and_requirements` / + `get_run_profile_and_requirements_in_fleet` in + `server/background/pipeline_tasks/jobs_submitted.py`). It must build the + `InstanceConfiguration` from that `requirements` parameter, not from + `job.job_spec.requirements`, which only ever reflects the run's own + requirements as computed at submission time and is never updated to include + a fleet's `reservation`/`security_group` when a run provisions new capacity + into an existing fleet. + """ + + def _resources(self) -> ResourcesSpec: + return ResourcesSpec(cpu=CPUSpec.parse("1")) + + def _run_job(self, effective_requirements: Requirements): + compute = _FakeCreateInstanceCompute() + run = MagicMock() + run.project_name = "test-project" + run.user = "test-user" + run.run_spec.merged_profile.tags = None + job = MagicMock() + job.job_spec.job_name = "test-run-0-0" + # The job's own (run-only) requirements deliberately differ from the + # effective (fleet+run-combined) ones passed as the `requirements` arg, + # to prove which one `run_job` actually uses. + job.job_spec.requirements = Requirements( + resources=self._resources(), + reservation="job-spec-reservation", + security_group="job-spec-security-group", + ) + instance_offer = MagicMock() + instance_offer.region = "us-east-1" + instance_offer.price = 1.0 + instance_offer.copy.return_value = instance_offer + compute.run_job( + run=run, + job=job, + instance_offer=instance_offer, + project_ssh_public_key="ssh-rsa AAAA", + project_ssh_private_key="private-key", + volumes=[], + placement_group=None, + requirements=effective_requirements, + ) + assert compute.last_instance_config is not None + return compute.last_instance_config + + def test_uses_effective_reservation_not_job_spec_reservation(self): + effective_requirements = Requirements( + resources=self._resources(), reservation="fleet-reservation" + ) + instance_config = self._run_job(effective_requirements) + assert instance_config.reservation == "fleet-reservation" + + def test_uses_effective_security_group_not_job_spec_security_group(self): + effective_requirements = Requirements( + resources=self._resources(), security_group="fleet-security-group" + ) + instance_config = self._run_job(effective_requirements) + assert instance_config.security_group == "fleet-security-group" + + def test_none_in_effective_requirements_is_respected(self): + # Even though job.job_spec.requirements sets both fields, an effective + # Requirements with neither set must result in neither being used - + # confirming `run_job` isn't merging the two, just using `requirements`. + effective_requirements = Requirements(resources=self._resources()) + instance_config = self._run_job(effective_requirements) + assert instance_config.reservation is None + assert instance_config.security_group is None + + class TestGenerateUniqueInstanceName: def test_generates_name(self): configuration = get_instance_configuration( diff --git a/src/tests/_internal/core/backends/oci/test_compute.py b/src/tests/_internal/core/backends/oci/test_compute.py index 6a61eb9d89..33361e0c2d 100644 --- a/src/tests/_internal/core/backends/oci/test_compute.py +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -1,5 +1,7 @@ from unittest.mock import MagicMock, patch +import oci + from dstack._internal.core.backends.oci.compute import OCICompute from dstack._internal.core.backends.oci.models import OCIClientCreds, OCIConfig from dstack._internal.core.models.backends.base import BackendType @@ -68,15 +70,18 @@ def _make_compute(config: OCIConfig) -> OCICompute: class TestOCIComputeSecurityGroup: + """ + All instances - default (auto-managed NSG) or custom-NSG - share the same + subnet/VCN, since OCI NSGs are VCN-scoped and only attach to VNICs in the + same VCN they live in. `create_instance` must never route a custom-NSG + instance to some other subnet. + """ + def _run_create_instance(self, compute, instance_config): with patch("dstack._internal.core.backends.oci.compute.resources") as res: res.VCN_CIDR = "10.0.0.0/16" - res.RESTRICTED_VCN_CIDR = "10.1.0.0/16" res.get_marketplace_listing_and_package.return_value = (MagicMock(), MagicMock()) res.get_or_create_security_group.return_value.id = "ocid1.nsg.oc1..managed" - res.set_up_restricted_network_resources_in_region.return_value.id = ( - "ocid1.subnet.oc1..restricted" - ) res.launch_instance.return_value.id = "ocid1.instance.oc1..instance" compute.create_instance(_make_offer(), instance_config, placement_group=None) return res @@ -91,9 +96,6 @@ def test_default_creates_and_syncs_managed_security_group(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) - # The default/auto-managed-NSG path uses the original default subnet - # and never touches the restricted VCN/subnet. - res.set_up_restricted_network_resources_in_region.assert_not_called() assert ( res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" ) @@ -113,16 +115,10 @@ def test_per_region_custom_nsg_is_left_untouched(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..custom" ) - # A custom NSG routes the instance into a dedicated restricted VCN/subnet - # (no security list), never the default one. - res.set_up_restricted_network_resources_in_region.assert_called_once() - assert ( - res.set_up_restricted_network_resources_in_region.call_args.kwargs["project_name"] - == "test-project" - ) + # A custom NSG uses the same shared default subnet as auto-managed + # instances - there is no separate subnet/VCN for custom-NSG instances. assert ( - res.launch_instance.call_args.kwargs["subnet_id"] - == "ocid1.subnet.oc1..restricted" + res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" ) def test_region_not_in_mapping_falls_back_to_managed(self): @@ -139,7 +135,6 @@ def test_region_not_in_mapping_falls_back_to_managed(self): res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) - res.set_up_restricted_network_resources_in_region.assert_not_called() assert ( res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" ) @@ -156,10 +151,8 @@ def test_instance_level_custom_nsg_is_left_untouched(self): assert ( res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" ) - res.set_up_restricted_network_resources_in_region.assert_called_once() assert ( - res.launch_instance.call_args.kwargs["subnet_id"] - == "ocid1.subnet.oc1..restricted" + res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" ) def test_instance_level_overrides_per_region_mapping(self): @@ -177,14 +170,12 @@ def test_instance_level_overrides_per_region_mapping(self): assert ( res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" ) - res.set_up_restricted_network_resources_in_region.assert_called_once() assert ( - res.launch_instance.call_args.kwargs["subnet_id"] - == "ocid1.subnet.oc1..restricted" + res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" ) -class TestGetOrCreateRestrictedSubnet: +class TestGetOrCreateSubnet: def test_creates_subnet_with_empty_security_list_ids(self): from dstack._internal.core.backends.oci import resources @@ -193,8 +184,8 @@ def test_creates_subnet_with_empty_security_list_ids(self): client.list_subnets.return_value.next_page = None client.list_subnets.return_value.has_next_page = False - resources.get_or_create_restricted_subnet( - "dstack-test-project-restricted-subnet", + resources.get_or_create_subnet( + "dstack-test-project-default-subnet", "ocid1.vcn.oc1..vcn", "ocid1.compartment.oc1..compartment", client, @@ -204,19 +195,20 @@ def test_creates_subnet_with_empty_security_list_ids(self): details = client.create_subnet.call_args.args[0] assert details.security_list_ids == [] assert details.vcn_id == "ocid1.vcn.oc1..vcn" - assert details.display_name == "dstack-test-project-restricted-subnet" + assert details.display_name == "dstack-test-project-default-subnet" def test_returns_existing_subnet_without_creating(self): from dstack._internal.core.backends.oci import resources existing = MagicMock() + existing.security_list_ids = [] client = MagicMock() client.list_subnets.return_value.data = [existing] client.list_subnets.return_value.next_page = None client.list_subnets.return_value.has_next_page = False - result = resources.get_or_create_restricted_subnet( - "dstack-test-project-restricted-subnet", + result = resources.get_or_create_subnet( + "dstack-test-project-default-subnet", "ocid1.vcn.oc1..vcn", "ocid1.compartment.oc1..compartment", client, @@ -224,3 +216,94 @@ def test_returns_existing_subnet_without_creating(self): assert result is existing client.create_subnet.assert_not_called() + client.update_subnet.assert_not_called() + + def test_existing_subnet_with_security_list_is_updated_to_remove_it(self): + """ + A subnet created before this fix still has the VCN's default security + list attached. Since this subnet is dstack-owned infrastructure (not a + user-supplied resource), it must be updated in place so the NSG becomes + the sole security boundary, matching subnets created after the fix. + """ + from dstack._internal.core.backends.oci import resources + + existing = MagicMock() + existing.id = "ocid1.subnet.oc1..existing" + existing.security_list_ids = ["ocid1.securitylist.oc1..default"] + client = MagicMock() + client.list_subnets.return_value.data = [existing] + client.list_subnets.return_value.next_page = None + client.list_subnets.return_value.has_next_page = False + updated = MagicMock() + client.update_subnet.return_value.data = updated + + result = resources.get_or_create_subnet( + "dstack-test-project-default-subnet", + "ocid1.vcn.oc1..vcn", + "ocid1.compartment.oc1..compartment", + client, + ) + + client.create_subnet.assert_not_called() + client.update_subnet.assert_called_once() + args, _ = client.update_subnet.call_args + assert args[0] == "ocid1.subnet.oc1..existing" + assert args[1].security_list_ids == [] + assert result is updated + + +class TestUpdateSecurityGroupRulesForRunnerInstances: + def test_adds_ssh_ingress_and_egress_rules(self): + from dstack._internal.core.backends.oci import resources + + client = MagicMock() + client.list_network_security_group_security_rules.return_value.data = [] + client.list_network_security_group_security_rules.return_value.next_page = None + client.list_network_security_group_security_rules.return_value.has_next_page = False + + resources.update_security_group_rules_for_runner_instances( + "ocid1.nsg.oc1..managed", client + ) + + client.add_network_security_group_security_rules.assert_called_once() + (_, details), _ = client.add_network_security_group_security_rules.call_args + rules = details.security_rules + directions = {rule.direction for rule in rules} + assert ( + oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS in directions + ) + assert oci.core.models.AddSecurityRuleDetails.DIRECTION_EGRESS in directions + + ssh_rules = [ + rule + for rule in rules + if rule.direction == oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS + and rule.source_type + == oci.core.models.AddSecurityRuleDetails.SOURCE_TYPE_CIDR_BLOCK + ] + assert len(ssh_rules) == 1 + assert ssh_rules[0].source == "0.0.0.0/0" + assert ssh_rules[0].tcp_options.destination_port_range.min == 22 + assert ssh_rules[0].tcp_options.destination_port_range.max == 22 + + egress_rules = [ + rule + for rule in rules + if rule.direction == oci.core.models.AddSecurityRuleDetails.DIRECTION_EGRESS + ] + assert len(egress_rules) == 1 + assert egress_rules[0].destination == "0.0.0.0/0" + assert ( + egress_rules[0].destination_type + == oci.core.models.AddSecurityRuleDetails.DESTINATION_TYPE_CIDR_BLOCK + ) + + intra_group_rules = [ + rule + for rule in rules + if rule.direction == oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS + and rule.source_type + == oci.core.models.AddSecurityRuleDetails.SOURCE_TYPE_NETWORK_SECURITY_GROUP + ] + assert len(intra_group_rules) == 1 + assert intra_group_rules[0].source == "ocid1.nsg.oc1..managed" From 95b5e4489cb645e681d64921ab66425377c95958 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 16:34:07 -0400 Subject: [PATCH 5/6] Apply ruff-format to last commit --- .../_internal/core/backends/oci/compute.py | 4 +- .../core/backends/oci/test_compute.py | 58 +++++-------------- 2 files changed, 16 insertions(+), 46 deletions(-) diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index a46e90f01c..d372f2c23c 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -147,9 +147,7 @@ def create_instance( ).data 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 - ) + 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", diff --git a/src/tests/_internal/core/backends/oci/test_compute.py b/src/tests/_internal/core/backends/oci/test_compute.py index 33361e0c2d..1c0d68faa0 100644 --- a/src/tests/_internal/core/backends/oci/test_compute.py +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -93,51 +93,36 @@ def test_default_creates_and_syncs_managed_security_group(self): res.get_or_create_security_group.assert_called_once() res.update_security_group_rules_for_runner_instances.assert_called_once() assert ( - res.launch_instance.call_args.kwargs["security_group_id"] - == "ocid1.nsg.oc1..managed" - ) - assert ( - res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" + res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" def test_per_region_custom_nsg_is_left_untouched(self): compute = _make_compute( - _make_config( - network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..custom"} - ) + _make_config(network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..custom"}) ) res = self._run_create_instance(compute, _make_instance_config()) res.get_or_create_security_group.assert_not_called() res.update_security_group_rules_for_runner_instances.assert_not_called() res.update_security_group_rules.assert_not_called() - assert ( - res.launch_instance.call_args.kwargs["security_group_id"] - == "ocid1.nsg.oc1..custom" - ) + assert res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..custom" # A custom NSG uses the same shared default subnet as auto-managed # instances - there is no separate subnet/VCN for custom-NSG instances. - assert ( - res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" - ) + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" def test_region_not_in_mapping_falls_back_to_managed(self): compute = _make_compute( - _make_config( - network_security_group_ids={"us-phoenix-1": "ocid1.nsg.oc1..other"} - ) + _make_config(network_security_group_ids={"us-phoenix-1": "ocid1.nsg.oc1..other"}) ) res = self._run_create_instance(compute, _make_instance_config()) res.get_or_create_security_group.assert_called_once() res.update_security_group_rules_for_runner_instances.assert_called_once() assert ( - res.launch_instance.call_args.kwargs["security_group_id"] - == "ocid1.nsg.oc1..managed" - ) - assert ( - res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" + res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..managed" ) + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" def test_instance_level_custom_nsg_is_left_untouched(self): compute = _make_compute(_make_config()) @@ -148,18 +133,12 @@ def test_instance_level_custom_nsg_is_left_untouched(self): res.get_or_create_security_group.assert_not_called() res.update_security_group_rules_for_runner_instances.assert_not_called() res.update_security_group_rules.assert_not_called() - assert ( - res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" - ) - assert ( - res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" - ) + assert res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" def test_instance_level_overrides_per_region_mapping(self): compute = _make_compute( - _make_config( - network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..project"} - ) + _make_config(network_security_group_ids={"us-ashburn-1": "ocid1.nsg.oc1..project"}) ) res = self._run_create_instance( compute, _make_instance_config(security_group="ocid1.nsg.oc1..run") @@ -167,12 +146,8 @@ def test_instance_level_overrides_per_region_mapping(self): res.get_or_create_security_group.assert_not_called() res.update_security_group_rules_for_runner_instances.assert_not_called() - assert ( - res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" - ) - assert ( - res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" - ) + assert res.launch_instance.call_args.kwargs["security_group_id"] == "ocid1.nsg.oc1..run" + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" class TestGetOrCreateSubnet: @@ -269,17 +244,14 @@ def test_adds_ssh_ingress_and_egress_rules(self): (_, details), _ = client.add_network_security_group_security_rules.call_args rules = details.security_rules directions = {rule.direction for rule in rules} - assert ( - oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS in directions - ) + assert oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS in directions assert oci.core.models.AddSecurityRuleDetails.DIRECTION_EGRESS in directions ssh_rules = [ rule for rule in rules if rule.direction == oci.core.models.AddSecurityRuleDetails.DIRECTION_INGRESS - and rule.source_type - == oci.core.models.AddSecurityRuleDetails.SOURCE_TYPE_CIDR_BLOCK + and rule.source_type == oci.core.models.AddSecurityRuleDetails.SOURCE_TYPE_CIDR_BLOCK ] assert len(ssh_rules) == 1 assert ssh_rules[0].source == "0.0.0.0/0" From 6d49ee0645630f323ef6486a2909fc8bf6279148 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Tue, 14 Jul 2026 10:25:55 -0400 Subject: [PATCH 6/6] docs: unify custom security group / NSG / firewall rule wording across backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses PR review feedback: use the maintainer's suggested template consistently across AWS, Azure, and OCI ('creates and manages its own per project that allows: SSH ingress (TCP port 22) from , unrestricted egress, and all traffic within the so that multi-node clusters work out of the box'), swapping only backend-specific terminology. GCP keeps a distinct structure since it's a VPC-wide firewall rule rather than a per-instance-attachable resource. Also fixes accuracy issues in the 'responsibility for a custom one' paragraphs found during review: - GCP: the firewall rule is per VPC network, not per dstack project; documents that dstack creates neither the instance nor gateway firewall rule when a shared VPC is configured, and corrects the gateway rule's ports (22/80/443, not just HTTP/HTTPS). - OCI: converted the bulleted responsibility list to prose matching the AWS/Azure phrasing. - Azure: notes that a custom NSG only needs an explicit SSH-ingress rule, since Azure's own default NSG rules already permit egress and intra-VNet traffic on any NSG unless overridden — kept in the 'what you're responsible for' paragraph, while the opening 'default behavior' sentence matches the uniform template per the maintainer's request. - AWS: removed a doubled 'from ... from'. --- mkdocs/docs/concepts/backends.md | 63 +++++++++++++++++--------------- 1 file changed, 34 insertions(+), 29 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index a7f26b36d3..0837718980 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -243,8 +243,9 @@ There are two ways to configure AWS: using an access key or using the default cr 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). + By default, `dstack` creates and manages its own security group per project that allows: SSH + ingress (TCP port 22) from `0.0.0.0/0`, unrestricted egress, and all traffic within the group so + that 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: @@ -277,9 +278,10 @@ There are two ways to configure AWS: using an access key or using the default cr 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. + to instances as-is and never adds, removes, or modifies its rules. You're responsible for its + rules, including SSH ingress (from wherever the `dstack` server and users connect), egress + (e.g. outbound access to pull Docker images), and, 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. @@ -486,8 +488,9 @@ There are two ways to configure Azure: using a client secret or using the defaul 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). + By default, `dstack` creates and manages its own network security group per project that allows: + SSH ingress (TCP port 22) from the internet, unrestricted egress, and all traffic within the + virtual network so that 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 @@ -508,8 +511,10 @@ There are two ways to configure Azure: using a client secret or using the defaul 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. + never adds, removes, or modifies its rules. You're responsible for its ingress rules, in + particular allowing SSH from wherever you connect — Azure's own default NSG rules already permit + egress and traffic within the virtual network (including between instances in a multi-node + cluster), unless your NSG's rules override those defaults. You can also override this per fleet or run using the `security_group` profile property. @@ -732,8 +737,8 @@ gcloud projects list --format="json(projectId)" 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. + By default, `dstack` creates a VPC firewall rule per network that allows SSH ingress (TCP port 22) + from `0.0.0.0/0`, 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`: @@ -751,9 +756,12 @@ gcloud projects list --format="json(projectId)" ``` 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. + `dstack` needs. `create_firewall_rules: false` only affects the instance SSH rule — the separate + firewall rule `dstack` creates for gateways (allowing TCP ports 22, 80, and 443 from the internet, + scoped to the `dstack-gateway-instance` target tag) is always auto-managed when it's created at + all, since gateways are meant to be internet-reachable. Note that when a shared VPC + (`vpc_project_id`) is used, `dstack` never creates either firewall rule, regardless of + `create_firewall_rules` — you must have your own firewall rules in place either way. ### Lambda @@ -1162,11 +1170,11 @@ There are two ways to configure OCI: using client credentials or using the defau ``` ??? 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. + By default, `dstack` creates and manages its own network security group (NSG) per project that + allows: SSH ingress (TCP port 22) from `0.0.0.0/0`, unrestricted egress, and all traffic within + the group so that multi-node clusters work out of the box. These rules live on the NSG itself + (rather than an OCI security list) because `dstack` places every instance in a shared subnet that + has **no** security list attached, making the NSG the sole security boundary. OCI network security groups are region-scoped, so a custom NSG is configured per region via `network_security_group_ids`: @@ -1190,16 +1198,13 @@ There are two ways to configure OCI: using client credentials or using the defau to. Since `dstack` provisions instances into its own default VCN per project (`dstack--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. + When a custom NSG is used, `dstack` attaches it to instances as-is and 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. You're responsible for its rules, including SSH ingress (from + wherever you connect), egress (e.g. outbound access to pull Docker images — without an egress + rule, instances will have no outbound connectivity and will fail to start runs), and, 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.