From 5caf27483ec8688a8d43ea8776effaaeaff1d405 Mon Sep 17 00:00:00 2001 From: James Boydell Date: Mon, 13 Jul 2026 10:52:40 -0400 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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/5] 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"