diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 92d0f72e94..0837718980 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -242,6 +242,49 @@ 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 that allows: SSH + ingress (TCP port 22) from `0.0.0.0/0`, unrestricted egress, and all traffic within the group so + that multi-node clusters work out of the box. + To use a security group you manage yourself instead, set `security_group_name` if you create a + security group with the same name in every configured region's VPC: + + ```yaml + projects: + - name: main + backends: + - type: aws + creds: + type: default + + security_group_name: my-security-group + ``` + + If your security groups have different names (or IDs are more convenient) per region, use + `security_group_ids` instead: + + ```yaml + projects: + - name: main + backends: + - type: aws + creds: + type: default + + security_group_ids: + us-east-1: sg-0a1b2c3d4e5f6g7h8 + us-west-2: sg-1b2c3d4e5f6g7h8i9 + ``` + + Regions not covered by `security_group_ids` fall back to `security_group_name` if set, or to + dstack's auto-created security group otherwise. Either way, `dstack` attaches the security group + to instances as-is and never adds, removes, or modifies its rules. You're responsible for its + rules, including SSH ingress (from wherever the `dstack` server and users connect), egress + (e.g. outbound access to pull Docker images), and, for multi-node clusters, traffic between + instances in the group. + + You can also override this per fleet or run using the `security_group` profile property. + ??? 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 +487,37 @@ 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 per project that allows: + SSH ingress (TCP port 22) from the internet, unrestricted egress, and all traffic within the + virtual network so that multi-node clusters work out of the box. + Azure NSG names must be unique within a resource group regardless of region, so a custom NSG is + configured per location via `network_security_group_names`. The values are plain NSG names + within the configured `resource_group` (not full Azure resource IDs), so NSGs in a different + resource group cannot be referenced: + + ```yaml + projects: + - name: main + backends: + - type: azure + creds: + type: default + regions: [westeurope, eastus] + network_security_group_names: + westeurope: my-network-security-group-we + eastus: my-network-security-group-eus + ``` + + Locations not covered by `network_security_group_names` fall back to dstack's auto-created network + security group. Either way, `dstack` attaches the network security group to instances as-is and + never adds, removes, or modifies its rules. You're responsible for its ingress rules, in + particular allowing SSH from wherever you connect — Azure's own default NSG rules already permit + egress and traffic within the virtual network (including between instances in a multi-node + cluster), unless your NSG's rules override those defaults. + + You can also override this per fleet or run using the `security_group` profile property. + ### GCP There are two ways to configure GCP: using a service account or using the default credentials. @@ -662,6 +736,33 @@ gcloud projects list --format="json(projectId)" Using private subnets assumes that both the `dstack` server and users can access the configured VPC's private subnets. Additionally, [Cloud NAT](https://cloud.google.com/nat/docs/overview) must be configured to provide access to external resources for provisioned instances. +??? info "Custom firewall rules" + By default, `dstack` creates a VPC firewall rule per network that allows SSH ingress (TCP port 22) + from `0.0.0.0/0`, scoped to the `dstack-runner-instance` target tag. + Unlike AWS/Azure/OCI, GCP firewall rules apply to the whole VPC rather than to a single attachable resource, + so there's no per-fleet override — if you manage your own firewall rules and don't want `dstack` creating + a rule that opens port 22 to `0.0.0.0/0`, disable this at the project level with `create_firewall_rules: false`: + + ```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. `create_firewall_rules: false` only affects the instance SSH rule — the separate + firewall rule `dstack` creates for gateways (allowing TCP ports 22, 80, and 443 from the internet, + scoped to the `dstack-gateway-instance` target tag) is always auto-managed when it's created at + all, since gateways are meant to be internet-reachable. Note that when a shared VPC + (`vpc_project_id`) is used, `dstack` never creates either firewall rule, regardless of + `create_firewall_rules` — you must have your own firewall rules in place either way. + ### Lambda 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 +1169,45 @@ 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 (NSG) per project that + allows: SSH ingress (TCP port 22) from `0.0.0.0/0`, unrestricted egress, and all traffic within + the group so that multi-node clusters work out of the box. These rules live on the NSG itself + (rather than an OCI security list) because `dstack` places every instance in a shared subnet that + has **no** security list attached, making the NSG the sole security boundary. + + OCI network security groups are region-scoped, so a custom NSG is configured per region via + `network_security_group_ids`: + + ```yaml + projects: + - name: main + backends: + - type: oci + creds: + type: default + network_security_group_ids: + eu-frankfurt-1: ocid1.networksecuritygroup.oc1..aaaaaaaa + us-ashburn-1: ocid1.networksecuritygroup.oc1..bbbbbbbb + ``` + + Regions not covered by `network_security_group_ids` fall back to dstack's auto-created network + security group. + + OCI requires a network security group to belong to the same VCN as the instances it's attached + to. Since `dstack` provisions instances into its own default VCN per project + (`dstack--default-vcn`), your custom NSG must be created in that VCN. + + When a custom NSG is used, `dstack` attaches it to instances as-is and never adds, removes, or + modifies its rules. Because the shared subnet has no security list, the NSG is the sole security + boundary for these instances — there is no implicit SSH-from-anywhere or implicit outbound-all + falling back from a security list. You're responsible for its rules, including SSH ingress (from + wherever you connect), egress (e.g. outbound access to pull Docker images — without an egress + rule, instances will have no outbound connectivity and will fail to start runs), and, for + multi-node clusters, traffic between instances in the group. + + You can also override this per fleet or run using the `security_group` profile property. + 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..6a5b22150e 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,27 @@ def create_instance( instance_type=instance_offer.instance.name, image_config=self.config.os_images, ) - security_group_id = self._create_security_group( - ec2_client=ec2_client, - region=instance_offer.region, - project_id=project_name, - vpc_id=vpc_id, - ) + security_group_id = instance_config.security_group + if security_group_id is None and self.config.security_group_ids is not None: + security_group_id = self.config.security_group_ids.get(instance_offer.region) + if security_group_id is None and self.config.security_group_name is not None: + security_group_id = aws_resources.get_security_group_id_by_name( + ec2_client=ec2_client, + name=self.config.security_group_name, + vpc_id=vpc_id, + ) + if security_group_id is None: + raise ComputeError( + f"Security group '{self.config.security_group_name}' not found in" + f" VPC {vpc_id} (region {instance_offer.region})" + ) + if security_group_id is None: + security_group_id = self._create_security_group( + ec2_client=ec2_client, + region=instance_offer.region, + project_id=project_name, + vpc_id=vpc_id, + ) try: response = ec2_resource.create_instances( # pyright: ignore[reportAttributeAccessIssue] **aws_resources.create_instances_struct( @@ -363,9 +380,18 @@ def create_instance( ) except botocore.exceptions.ClientError as e: logger.warning("Got botocore.exceptions.ClientError: %s", e) - if e.response["Error"]["Code"] == "InvalidParameterValue": + error_code = e.response["Error"]["Code"] + if error_code == "InvalidParameterValue": msg = e.response["Error"].get("Message", "") raise ComputeError(f"Invalid AWS request: {msg}") + if error_code == "InvalidGroup.NotFound": + # A misconfigured security group (e.g. wrong VPC/region) is not a + # capacity issue, so surface it clearly instead of retrying other AZs. + msg = e.response["Error"].get("Message", "") + raise ComputeError( + f"Security group not found for instance in region" + f" {instance_offer.region}: {msg}" + ) continue instance = response[0] # wait_until_running() is only needed so that instance is immediately ready for volume attach. diff --git a/src/dstack/_internal/core/backends/aws/configurator.py b/src/dstack/_internal/core/backends/aws/configurator.py index 8d6c8afe14..6ca43997db 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,19 @@ def _check_config_iam_instance_profile( f"Failed to check IAM instance profile {config.iam_instance_profile}" ) + def _check_config_security_group(self, config: AWSBackendConfigWithCreds): + if config.security_group_ids is None: + return + regions = config.regions if config.regions is not None else DEFAULT_REGIONS + unknown_regions = [r for r in config.security_group_ids if r not in regions] + if unknown_regions: + raise ServerClientError( + msg=( + f"`security_group_ids` specifies regions not in `regions`: {unknown_regions}." + " This is likely a typo — remove the extra keys or add them to `regions`" + ) + ) + def _check_config_vpc(self, session: Session, config: AWSBackendConfigWithCreds): allocate_public_ip = config.public_ips if config.public_ips is not None else True use_default_vpcs = config.default_vpcs if config.default_vpcs is not None else True diff --git a/src/dstack/_internal/core/backends/aws/models.py b/src/dstack/_internal/core/backends/aws/models.py index d76a4c9ab3..946a588978 100644 --- a/src/dstack/_internal/core/backends/aws/models.py +++ b/src/dstack/_internal/core/backends/aws/models.py @@ -96,6 +96,37 @@ class AWSBackendConfig(CoreModel): ) ), ] = None + security_group_name: Annotated[ + Optional[str], + Field( + description=( + "The name of an existing security group to use for instances instead of the one" + " `dstack` creates and manages automatically (`dstack_security_group_`)." + " 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 d2843c7f22..fb24addd21 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,10 +148,14 @@ def create_instance( location=location, allocate_public_ip=allocate_public_ip, ) - network_security_group = azure_utils.get_default_network_security_group_name( - resource_group=self.config.resource_group, - location=location, - ) + network_security_group = instance_config.security_group + if network_security_group is None and self.config.network_security_group_names is not None: + network_security_group = self.config.network_security_group_names.get(location) + if network_security_group is None: + network_security_group = azure_utils.get_default_network_security_group_name( + resource_group=self.config.resource_group, + location=location, + ) managed_identity_resource_group, managed_identity_name = parse_vm_managed_identity( self.config.vm_managed_identity @@ -256,7 +262,12 @@ def create_gateway( location=configuration.region, allocate_public_ip=True, ) - network_security_group = azure_utils.get_default_network_security_group_name( + # Gateways always use the dedicated gateway NSG (created unconditionally in + # `_create_network_resources`, regardless of `network_security_group_names`), + # never the per-location default/custom instance NSG. This keeps gateway + # provisioning working even for locations where the default instance NSG is + # skipped because a custom one is configured for instances. + network_security_group = azure_utils.get_gateway_network_security_group_name( resource_group=self.config.resource_group, location=configuration.region, ) diff --git a/src/dstack/_internal/core/backends/azure/configurator.py b/src/dstack/_internal/core/backends/azure/configurator.py index 31297e4ce3..b93e088db2 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 @@ -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,6 +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_names=config.network_security_group_names, ) return BackendRecord( config=AzureStoredConfig( @@ -282,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 ): @@ -340,6 +363,7 @@ def _create_network_resources( resource_group: str, locations: List[str], create_default_network: bool, + network_security_group_names: Optional[Dict[str, str]] = None, ): def func(location: str): network_manager = NetworkManager( @@ -352,11 +376,17 @@ 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 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_names` - 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..10d0536076 100644 --- a/src/dstack/_internal/core/backends/azure/models.py +++ b/src/dstack/_internal/core/backends/azure/models.py @@ -81,6 +81,23 @@ class AzureBackendConfig(CoreModel): ) ), ] = None + network_security_group_names: Annotated[ + Optional[Dict[str, str]], + Field( + description=( + "The mapping from Azure locations to the names of existing network security groups" + " 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" + " 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/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 7cb7c84b35..fd0d20c6ad 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -320,7 +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, + reservation=requirements.reservation, + security_group=requirements.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..9652ed86dc 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, @@ -562,6 +562,9 @@ def create_gateway( configuration: GatewayComputeConfiguration, ) -> GatewayProvisioningData: 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 4d06144ee8..87e3ad9eff 100644 --- a/src/dstack/_internal/core/backends/gcp/models.py +++ b/src/dstack/_internal/core/backends/gcp/models.py @@ -80,6 +80,19 @@ class GCPBackendConfig(CoreModel): ) ), ] = None + create_firewall_rules: Annotated[ + Optional[bool], + Field( + description=( + "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 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 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..d372f2c23c 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): @@ -133,18 +135,30 @@ 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 - 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 + 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", + 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 +172,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/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 12ce4e9f91..ece575467f 100644 --- a/src/dstack/_internal/core/backends/oci/models.py +++ b/src/dstack/_internal/core/backends/oci/models.py @@ -69,6 +69,28 @@ class OCIBackendConfig(CoreModel): ) ), ] = None + network_security_group_ids: Annotated[ + Optional[Dict[str, str]], + Field( + description=( + "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." + " 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. 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 class OCIBackendConfigWithCreds(OCIBackendConfig): diff --git a/src/dstack/_internal/core/backends/oci/resources.py b/src/dstack/_internal/core/backends/oci/resources.py index 91e1d890f0..b9ab80dee5 100644 --- a/src/dstack/_internal/core/backends/oci/resources.py +++ b/src/dstack/_internal/core/backends/oci/resources.py @@ -556,10 +556,47 @@ 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: + """ + 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( @@ -568,6 +605,7 @@ def get_or_create_subnet( compartment_id=compartment_id, display_name=name, vcn_id=vcn_id, + security_list_ids=[], ) ).data @@ -628,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", @@ -638,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/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/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 547f91d52c..eed8154750 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, ) @@ -957,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, ) @@ -1263,6 +1265,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/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 new file mode 100644 index 0000000000..dce562bd7c --- /dev/null +++ b/src/tests/_internal/core/backends/aws/test_compute.py @@ -0,0 +1,213 @@ +from unittest.mock import MagicMock, patch + +import botocore.exceptions +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.errors import ComputeError +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_name=None, security_group_ids=None) -> AWSConfig: + return AWSConfig( + creds=AWSAccessKeyCreds(access_key="test", secret_key="test"), + regions=["us-east-1"], + security_group_name=security_group_name, + security_group_ids=security_group_ids, + ) + + +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(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=region, + 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, + 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: + 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_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()) + 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() + + 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" + + 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() + + 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/aws/test_resources.py b/src/tests/_internal/core/backends/aws/test_resources.py index a719294cce..2ff12855e2 100644 --- a/src/tests/_internal/core/backends/aws/test_resources.py +++ b/src/tests/_internal/core/backends/aws/test_resources.py @@ -9,7 +9,9 @@ _is_valid_tag_key, _is_valid_tag_value, 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 @@ -276,6 +278,132 @@ 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 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 a0ce0afaef..68d3bb2a22 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/src/tests/_internal/core/backends/azure/test_compute.py @@ -1,8 +1,21 @@ +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, +) +from dstack._internal.server.testing.common import get_gateway_compute_configuration class TestVMImageVariant: @@ -71,3 +84,151 @@ 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_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_names=network_security_group_names, + ) + + +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=region, + 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", "nsg_names", "region", "expected"], + [ + # 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, nsg_names, region, 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_names=nsg_names), credential=Mock() + ) + compute.create_instance( + instance_offer=_offer(region=region), + instance_config=_instance_config(security_group=instance_sg), + placement_group=None, + ) + _, 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/azure/test_configurator.py b/src/tests/_internal/core/backends/azure/test_configurator.py index 1a329e0cd5..e121f95e1d 100644 --- a/src/tests/_internal/core/backends/azure/test_configurator.py +++ b/src/tests/_internal/core/backends/azure/test_configurator.py @@ -129,3 +129,101 @@ 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, 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=regions if regions is not None else ["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 _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_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_names={"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_names={ + "eastus": "my-nsg", + "westeurope": "my-other-nsg", + }, + ) + ) + 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/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/gcp/test_compute.py b/src/tests/_internal/core/backends/gcp/test_compute.py new file mode 100644 index 0000000000..2888678a5d --- /dev/null +++ b/src/tests/_internal/core/backends/gcp/test_compute.py @@ -0,0 +1,93 @@ +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_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_called_once() 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..1c0d68faa0 --- /dev/null +++ b/src/tests/_internal/core/backends/oci/test_compute.py @@ -0,0 +1,281 @@ +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 +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceConfiguration, + InstanceOfferWithAvailability, + InstanceType, + Resources, + SSHKey, +) + + +def _make_config(network_security_group_ids=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_ids=network_security_group_ids, + ) + + +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: + """ + 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.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" + ) + 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"}) + ) + 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" + # 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" + + 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" + ) + 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()) + 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" + 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"}) + ) + 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" + assert res.launch_instance.call_args.kwargs["subnet_id"] == "ocid1.subnet.oc1..subnet" + + +class TestGetOrCreateSubnet: + 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_subnet( + "dstack-test-project-default-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-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_subnet( + "dstack-test-project-default-subnet", + "ocid1.vcn.oc1..vcn", + "ocid1.compartment.oc1..compartment", + client, + ) + + 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" 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")