dstack 0.18.40rc1__py3-none-any.whl → 0.18.42__py3-none-any.whl

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

@@ -109,34 +109,21 @@ class AWSCompute(Compute):
             configurable_disk_size=CONFIGURABLE_DISK_SIZE,
             extra_filter=filter,
         )
-        regions = set(i.region for i in offers)
-
-        def get_quotas(client: botocore.client.BaseClient) -> Dict[str, int]:
-            region_quotas = {}
-            for page in client.get_paginator("list_service_quotas").paginate(ServiceCode="ec2"):
-                for q in page["Quotas"]:
-                    if "On-Demand" in q["QuotaName"]:
-                        region_quotas[q["UsageMetric"]["MetricDimensions"]["Class"]] = q["Value"]
-            return region_quotas
-
-        quotas = {}
-        with ThreadPoolExecutor(max_workers=8) as executor:
-            future_to_region = {}
-            for region in regions:
-                future = executor.submit(
-                    get_quotas, self.session.client("service-quotas", region_name=region)
-                )
-                future_to_region[future] = region
-            for future in as_completed(future_to_region):
-                quotas[future_to_region[future]] = future.result()
+        regions = list(set(i.region for i in offers))
+        regions_to_quotas = _get_regions_to_quotas(self.session, regions)
+        regions_to_zones = _get_regions_to_zones(self.session, regions)
 
         availability_offers = []
         for offer in offers:
             availability = InstanceAvailability.UNKNOWN
-            if not _has_quota(quotas[offer.region], offer.instance.name):
+            if not _has_quota(regions_to_quotas[offer.region], offer.instance.name):
                 availability = InstanceAvailability.NO_QUOTA
             availability_offers.append(
-                InstanceOfferWithAvailability(**offer.dict(), availability=availability)
+                InstanceOfferWithAvailability(
+                    **offer.dict(),
+                    availability=availability,
+                    availability_zones=regions_to_zones[offer.region],
+                )
             )
         return availability_offers
 
@@ -161,9 +148,9 @@ class AWSCompute(Compute):
         ec2_resource = self.session.resource("ec2", region_name=instance_offer.region)
         ec2_client = self.session.client("ec2", region_name=instance_offer.region)
         allocate_public_ip = self.config.allocate_public_ips
-        availability_zones = None
-        if instance_config.availability_zone is not None:
-            availability_zones = [instance_config.availability_zone]
+        zones = instance_offer.availability_zones
+        if zones is not None and len(zones) == 0:
+            raise NoCapacityError("No eligible availability zones")
 
         tags = {
             "Name": instance_config.instance_name,
@@ -174,7 +161,7 @@ class AWSCompute(Compute):
         tags = merge_tags(tags=tags, backend_tags=self.config.tags)
 
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
-        max_efa_interfaces = get_maximum_efa_interfaces(
+        max_efa_interfaces = _get_maximum_efa_interfaces(
             ec2_client=ec2_client, instance_type=instance_offer.instance.name
         )
         enable_efa = max_efa_interfaces > 0
@@ -185,7 +172,7 @@ class AWSCompute(Compute):
                 config=self.config,
                 region=instance_offer.region,
                 allocate_public_ip=allocate_public_ip,
-                availability_zones=availability_zones,
+                availability_zones=zones,
             )
             subnet_id_to_az_map = aws_resources.get_subnets_availability_zones(
                 ec2_client=ec2_client,
@@ -210,11 +197,11 @@ class AWSCompute(Compute):
         except botocore.exceptions.ClientError as e:
             logger.warning("Got botocore.exceptions.ClientError: %s", e)
             raise NoCapacityError()
-        tried_availability_zones = set()
+        tried_zones = set()
         for subnet_id, az in subnet_id_to_az_map.items():
-            if az in tried_availability_zones:
+            if az in tried_zones:
                 continue
-            tried_availability_zones.add(az)
+            tried_zones.add(az)
             try:
                 logger.debug("Trying provisioning %s in %s", instance_offer.instance.name, az)
                 image_id, username = aws_resources.get_image_id_and_username(
@@ -240,6 +227,7 @@ class AWSCompute(Compute):
                         allocate_public_ip=allocate_public_ip,
                         placement_group_name=instance_config.placement_group_name,
                         enable_efa=enable_efa,
+                        max_efa_interfaces=max_efa_interfaces,
                         reservation_id=instance_config.reservation,
                         is_capacity_block=is_capacity_block,
                     )
@@ -283,6 +271,7 @@ class AWSCompute(Compute):
         project_ssh_private_key: str,
         volumes: List[Volume],
     ) -> JobProvisioningData:
+        # TODO: run_job is the same for vm-based backends, refactor
         instance_config = InstanceConfiguration(
             project_name=run.project_name,
             instance_name=get_instance_name(run, job),  # TODO: generate name
@@ -290,15 +279,25 @@ class AWSCompute(Compute):
                 SSHKey(public=project_ssh_public_key.strip()),
             ],
             user=run.user,
+            volumes=volumes,
             reservation=run.run_spec.configuration.reservation,
         )
+        instance_offer = instance_offer.copy()
         if len(volumes) > 0:
             volume = volumes[0]
             if (
                 volume.provisioning_data is not None
                 and volume.provisioning_data.availability_zone is not None
             ):
-                instance_config.availability_zone = volume.provisioning_data.availability_zone
+                if instance_offer.availability_zones is None:
+                    instance_offer.availability_zones = [
+                        volume.provisioning_data.availability_zone
+                    ]
+                instance_offer.availability_zones = [
+                    z
+                    for z in instance_offer.availability_zones
+                    if z == volume.provisioning_data.availability_zone
+                ]
         return self.create_instance(instance_offer, instance_config)
 
     def create_placement_group(
@@ -544,14 +543,16 @@ class AWSCompute(Compute):
         }
         tags = merge_tags(tags=tags, backend_tags=self.config.tags)
 
-        zone = aws_resources.get_availability_zone(
+        zones = aws_resources.get_availability_zones(
             ec2_client=ec2_client, region=volume.configuration.region
         )
-        if zone is None:
+        if volume.configuration.availability_zone is not None:
+            zones = [z for z in zones if z == volume.configuration.availability_zone]
+        if len(zones) == 0:
             raise ComputeError(
                 f"Failed to find availability zone in region {volume.configuration.region}"
             )
-
+        zone = zones[0]
         volume_type = "gp3"
 
         logger.debug("Creating EBS volume %s", volume.configuration.name)
@@ -570,7 +571,6 @@ class AWSCompute(Compute):
 
         size = response["Size"]
         iops = response["Iops"]
-
         return VolumeProvisioningData(
             backend=BackendType.AWS,
             volume_id=response["VolumeId"],
@@ -635,11 +635,12 @@ class AWSCompute(Compute):
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
         logger.debug("Detaching EBS volume %s from instance %s", volume.volume_id, instance_id)
+        attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
         try:
             ec2_client.detach_volume(
                 VolumeId=volume.volume_id,
                 InstanceId=instance_id,
-                Device=get_or_error(volume.attachment_data).device_name,
+                Device=attachment_data.device_name,
                 Force=force,
             )
         except botocore.exceptions.ClientError as e:
@@ -672,23 +673,6 @@ class AWSCompute(Compute):
         return True
 
 
-def get_maximum_efa_interfaces(ec2_client: botocore.client.BaseClient, instance_type: str) -> int:
-    try:
-        response = ec2_client.describe_instance_types(
-            InstanceTypes=[instance_type],
-            Filters=[{"Name": "network-info.efa-supported", "Values": ["true"]}],
-        )
-    except botocore.exceptions.ClientError as e:
-        if e.response.get("Error", {}).get("Code") == "InvalidInstanceType":
-            # "The following supplied instance types do not exist: [<instance_type>]"
-            return 0
-        raise
-    instance_types = response["InstanceTypes"]
-    if not instance_types:
-        return 0
-    return instance_types[0]["NetworkInfo"]["EfaInfo"]["MaximumEfaInterfaces"]
-
-
 def get_vpc_id_subnet_id_or_error(
     ec2_client: botocore.client.BaseClient,
     config: AWSConfig,
@@ -770,6 +754,30 @@ def _get_vpc_id_subnet_id_by_vpc_name_or_error(
     )
 
 
+def _get_regions_to_quotas(
+    session: boto3.Session, regions: List[str]
+) -> Dict[str, Dict[str, int]]:
+    def get_region_quotas(client: botocore.client.BaseClient) -> Dict[str, int]:
+        region_quotas = {}
+        for page in client.get_paginator("list_service_quotas").paginate(ServiceCode="ec2"):
+            for q in page["Quotas"]:
+                if "On-Demand" in q["QuotaName"]:
+                    region_quotas[q["UsageMetric"]["MetricDimensions"]["Class"]] = q["Value"]
+        return region_quotas
+
+    regions_to_quotas = {}
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        future_to_region = {}
+        for region in regions:
+            future = executor.submit(
+                get_region_quotas, session.client("service-quotas", region_name=region)
+            )
+            future_to_region[future] = region
+        for future in as_completed(future_to_region):
+            regions_to_quotas[future_to_region[future]] = future.result()
+    return regions_to_quotas
+
+
 def _has_quota(quotas: Dict[str, int], instance_name: str) -> bool:
     if instance_name.startswith("p"):
         return quotas.get("P/OnDemand", 0) > 0
@@ -778,6 +786,22 @@ def _has_quota(quotas: Dict[str, int], instance_name: str) -> bool:
     return quotas.get("Standard/OnDemand", 0) > 0
 
 
+def _get_regions_to_zones(session: boto3.Session, regions: List[str]) -> Dict[str, List[str]]:
+    regions_to_zones = {}
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        future_to_region = {}
+        for region in regions:
+            future = executor.submit(
+                aws_resources.get_availability_zones,
+                session.client("ec2", region_name=region),
+                region,
+            )
+            future_to_region[future] = region
+        for future in as_completed(future_to_region):
+            regions_to_zones[future_to_region[future]] = future.result()
+    return regions_to_zones
+
+
 def _supported_instances(offer: InstanceOffer) -> bool:
     for family in [
         "t2.small",
@@ -798,6 +822,23 @@ def _supported_instances(offer: InstanceOffer) -> bool:
     return False
 
 
+def _get_maximum_efa_interfaces(ec2_client: botocore.client.BaseClient, instance_type: str) -> int:
+    try:
+        response = ec2_client.describe_instance_types(
+            InstanceTypes=[instance_type],
+            Filters=[{"Name": "network-info.efa-supported", "Values": ["true"]}],
+        )
+    except botocore.exceptions.ClientError as e:
+        if e.response.get("Error", {}).get("Code") == "InvalidInstanceType":
+            # "The following supplied instance types do not exist: [<instance_type>]"
+            return 0
+        raise
+    instance_types = response["InstanceTypes"]
+    if not instance_types:
+        return 0
+    return instance_types[0]["NetworkInfo"]["EfaInfo"]["MaximumEfaInterfaces"]
+
+
 def _get_instance_ip(instance: Any, public_ip: bool) -> str:
     if public_ip:
         return instance.public_ip_address

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

@@ -140,6 +140,7 @@ def create_instances_struct(
     allocate_public_ip: bool = True,
     placement_group_name: Optional[str] = None,
     enable_efa: bool = False,
+    max_efa_interfaces: int = 0,
     reservation_id: Optional[str] = None,
     is_capacity_block: bool = False,
 ) -> Dict[str, Any]:
@@ -183,7 +184,7 @@ def create_instances_struct(
     # AWS allows specifying either NetworkInterfaces for specific subnet_id
     # or instance-level SecurityGroupIds in case of no specific subnet_id, not both.
     if subnet_id is not None:
-        # Even if the instance type supports multiple cards, we always request only one interface
+        # If the instance type supports multiple cards, we request multiple interfaces only if not allocate_public_ip
         # due to the limitation: "AssociatePublicIpAddress [...] You cannot specify more than one
         # network interface in the request".
         # Error message: "(InvalidParameterCombination) when calling the RunInstances operation:
@@ -199,9 +200,28 @@ def create_instances_struct(
                 "DeviceIndex": 0,
                 "SubnetId": subnet_id,
                 "Groups": [security_group_id],
-                "InterfaceType": "efa" if enable_efa else "interface",
+                "InterfaceType": "efa" if max_efa_interfaces > 0 else "interface",
             },
         ]
+
+        if max_efa_interfaces > 1 and allocate_public_ip is False:
+            for i in range(1, max_efa_interfaces):
+                # Set to efa-only to use interfaces exclusively for GPU-to-GPU communication
+                interface_type = "efa-only"
+                if instance_type == "p5.48xlarge":
+                    # EFA configuration for P5 instances:
+                    # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-acc-inst-types.html#efa-for-p5
+                    interface_type = "efa" if i % 4 == 0 else "efa-only"
+                struct["NetworkInterfaces"].append(
+                    {
+                        "AssociatePublicIpAddress": allocate_public_ip,
+                        "NetworkCardIndex": i,
+                        "DeviceIndex": 1,
+                        "SubnetId": subnet_id,
+                        "Groups": [security_group_id],
+                        "InterfaceType": interface_type,
+                    }
+                )
     else:
         struct["SecurityGroupIds"] = [security_group_id]
 
@@ -370,16 +390,6 @@ def get_subnets_ids_for_vpc(
     return subnets_ids
 
 
-def get_availability_zone(ec2_client: botocore.client.BaseClient, region: str) -> Optional[str]:
-    zone_names = get_availability_zones(
-        ec2_client=ec2_client,
-        region=region,
-    )
-    if len(zone_names) == 0:
-        return None
-    return zone_names[0]
-
-
 def get_availability_zones(ec2_client: botocore.client.BaseClient, region: str) -> List[str]:
     response = ec2_client.describe_availability_zones(
         Filters=[

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

@@ -133,6 +133,8 @@ class AzureCompute(Compute):
         }
         tags = merge_tags(tags=tags, backend_tags=self.config.tags)
 
+        # TODO: Support custom availability_zones.
+        # Currently, VMs are regional, which means they don't have zone info.
         vm = _launch_instance(
             compute_client=self._compute_client,
             subscription_id=self.config.subscription_id,

dstack/_internal/core/backends/base/compute.py

@@ -371,7 +371,16 @@ def get_docker_commands(
         "rm -rf /run/sshd && mkdir -p /run/sshd && chown root:root /run/sshd",
         "rm -rf /var/empty && mkdir -p /var/empty && chown root:root /var/empty",
         # start sshd
-        f"/usr/sbin/sshd -p {DSTACK_RUNNER_SSH_PORT} -o PidFile=none -o PasswordAuthentication=no -o AllowTcpForwarding=yes -o PermitUserEnvironment=yes",
+        (
+            "/usr/sbin/sshd"
+            f" -p {DSTACK_RUNNER_SSH_PORT}"
+            " -o PidFile=none"
+            " -o PasswordAuthentication=no"
+            " -o AllowTcpForwarding=yes"
+            " -o PermitUserEnvironment=yes"
+            " -o ClientAliveInterval=30"
+            " -o ClientAliveCountMax=4"
+        ),
         # restore ld.so variables
         'if [ -n "$_LD_LIBRARY_PATH" ]; then export LD_LIBRARY_PATH="$_LD_LIBRARY_PATH"; fi',
         'if [ -n "$_LD_PRELOAD" ]; then export LD_PRELOAD="$_LD_PRELOAD"; fi',
@@ -381,7 +390,16 @@ def get_docker_commands(
     commands += [
         f"curl --connect-timeout 60 --max-time 240 --retry 1 --output {DSTACK_RUNNER_BINARY_PATH} {url}",
         f"chmod +x {DSTACK_RUNNER_BINARY_PATH}",
-        f"{DSTACK_RUNNER_BINARY_PATH} --log-level 6 start --http-port {DSTACK_RUNNER_HTTP_PORT} --temp-dir /tmp/runner --home-dir /root --working-dir /workflow",
+        (
+            f"{DSTACK_RUNNER_BINARY_PATH}"
+            " --log-level 6"
+            " start"
+            f" --http-port {DSTACK_RUNNER_HTTP_PORT}"
+            f" --ssh-port {DSTACK_RUNNER_SSH_PORT}"
+            " --temp-dir /tmp/runner"
+            " --home-dir /root"
+            " --working-dir /workflow"
+        ),
     ]
 
     return commands

dstack/_internal/core/backends/gcp/compute.py

@@ -94,21 +94,25 @@ class GCPCompute(Compute):
             for quota in region.quotas:
                 quotas[region.name][quota.metric] = quota.limit - quota.usage
 
-        seen_region_offers = set()
+        offer_keys_to_offers = {}
         offers_with_availability = []
         for offer in offers:
             region = offer.region[:-2]  # strip zone
             key = (_unique_instance_name(offer.instance), region)
-            if key in seen_region_offers:
+            if key in offer_keys_to_offers:
+                offer_keys_to_offers[key].availability_zones.append(offer.region)
                 continue
-            seen_region_offers.add(key)
             availability = InstanceAvailability.NO_QUOTA
             if _has_gpu_quota(quotas[region], offer.instance.resources):
                 availability = InstanceAvailability.UNKNOWN
             # todo quotas: cpu, memory, global gpu, tpu
-            offers_with_availability.append(
-                InstanceOfferWithAvailability(**offer.dict(), availability=availability)
+            offer_with_availability = InstanceOfferWithAvailability(
+                **offer.dict(),
+                availability=availability,
+                availability_zones=[offer.region],
             )
+            offer_keys_to_offers[key] = offer_with_availability
+            offers_with_availability.append(offer_with_availability)
             offers_with_availability[-1].region = region
 
         return offers_with_availability
@@ -156,10 +160,10 @@ class GCPCompute(Compute):
             )
         authorized_keys = instance_config.get_public_keys()
 
-        zones = _get_instance_zones(instance_offer)
-        if instance_config.availability_zone:
-            zones = [z for z in zones if z == instance_config.availability_zone]
-
+        # get_offers always fills instance_offer.availability_zones
+        zones = get_or_error(instance_offer.availability_zones)
+        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:
             gcp_resources.create_runner_firewall_rules(
@@ -371,6 +375,7 @@ class GCPCompute(Compute):
         project_ssh_private_key: str,
         volumes: List[Volume],
     ) -> JobProvisioningData:
+        # TODO: run_job is the same for vm-based backends, refactor
         instance_config = InstanceConfiguration(
             project_name=run.project_name,
             instance_name=get_instance_name(run, job),  # TODO: generate name
@@ -379,14 +384,24 @@ class GCPCompute(Compute):
             ],
             user=run.user,
             volumes=volumes,
+            reservation=run.run_spec.configuration.reservation,
         )
+        instance_offer = instance_offer.copy()
         if len(volumes) > 0:
             volume = volumes[0]
             if (
                 volume.provisioning_data is not None
                 and volume.provisioning_data.availability_zone is not None
             ):
-                instance_config.availability_zone = volume.provisioning_data.availability_zone
+                if instance_offer.availability_zones is None:
+                    instance_offer.availability_zones = [
+                        volume.provisioning_data.availability_zone
+                    ]
+                instance_offer.availability_zones = [
+                    z
+                    for z in instance_offer.availability_zones
+                    if z == volume.provisioning_data.availability_zone
+                ]
         return self.create_instance(instance_offer, instance_config)
 
     def create_gateway(
@@ -497,15 +512,18 @@ class GCPCompute(Compute):
         raise ComputeError(f"Persistent disk {volume.configuration.volume_id} not found")
 
     def create_volume(self, volume: Volume) -> VolumeProvisioningData:
-        zone = gcp_resources.get_availability_zone(
+        zones = gcp_resources.get_availability_zones(
             regions_client=self.regions_client,
             project_id=self.config.project_id,
             region=volume.configuration.region,
         )
-        if zone is None:
+        if volume.configuration.availability_zone is not None:
+            zones = [z for z in zones if z == volume.configuration.availability_zone]
+        if len(zones) == 0:
             raise ComputeError(
                 f"Failed to find availability zone in region {volume.configuration.region}"
             )
+        zone = zones[0]
 
         labels = {
             "owner": "dstack",
@@ -648,6 +666,7 @@ class GCPCompute(Compute):
             instance_id,
         )
         zone = get_or_error(volume.provisioning_data).availability_zone
+        attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
         # This method has no information if the instance is a TPU or a VM,
         # so we first try to see if there is a TPU with such name
         try:
@@ -676,7 +695,7 @@ class GCPCompute(Compute):
                 project=self.config.project_id,
                 zone=get_or_error(volume.provisioning_data).availability_zone,
                 instance=instance_id,
-                device_name=get_or_error(volume.attachment_data).device_name,
+                device_name=attachment_data.device_name,
             )
             gcp_resources.wait_for_extended_operation(operation, "persistent disk detachment")
         logger.debug(
@@ -759,17 +778,6 @@ def _unique_instance_name(instance: InstanceType) -> str:
     return f"{name}-{gpu.name}-{gpu.memory_mib}"
 
 
-def _get_instance_zones(instance_offer: InstanceOffer) -> List[str]:
-    zones = []
-    for offer in get_catalog_offers(backend=BackendType.GCP):
-        if _unique_instance_name(instance_offer.instance) != _unique_instance_name(offer.instance):
-            continue
-        if offer.region[:-2] != instance_offer.region:
-            continue
-        zones.append(offer.region)
-    return zones
-
-
 def _get_tpu_startup_script(authorized_keys: List[str]) -> str:
     commands = get_shim_commands(
         authorized_keys=authorized_keys, is_privileged=True, pjrt_device="TPU"

dstack/_internal/core/backends/gcp/resources.py

@@ -31,21 +31,6 @@ supported_accelerators = [
 ]
 
 
-def get_availability_zone(
-    regions_client: compute_v1.RegionsClient,
-    project_id: str,
-    region: str,
-) -> Optional[str]:
-    zones = get_availability_zones(
-        regions_client=regions_client,
-        project_id=project_id,
-        region=region,
-    )
-    if len(zones) == 0:
-        return None
-    return zones[0]
-
-
 def get_availability_zones(
     regions_client: compute_v1.RegionsClient,
     project_id: str,

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

@@ -76,7 +76,13 @@ class OCICompute(Compute):
             else:
                 availability = InstanceAvailability.NO_QUOTA
             offers_with_availability.append(
-                InstanceOfferWithAvailability(**offer.dict(), availability=availability)
+                InstanceOfferWithAvailability(
+                    **offer.dict(),
+                    availability=availability,
+                    availability_zones=shapes_availability[offer.region].get(
+                        offer.instance.name, []
+                    ),
+                )
             )
 
         return offers_with_availability
@@ -111,11 +117,9 @@ class OCICompute(Compute):
     ) -> JobProvisioningData:
         region = self.regions[instance_offer.region]
 
-        availability_domain = resources.choose_available_domain(
-            instance_offer.instance.name, self.shapes_quota, region, self.config.compartment_id
-        )
-        if availability_domain is None:
+        if not instance_offer.availability_zones:
             raise NoCapacityError("Shape unavailable in all availability domains")
+        availability_domain = instance_offer.availability_zones[0]
 
         listing, package = resources.get_marketplace_listing_and_package(
             cuda=len(instance_offer.instance.resources.gpus) > 0,
@@ -170,6 +174,7 @@ class OCICompute(Compute):
             hostname=None,
             internal_ip=None,
             region=instance_offer.region,
+            availability_zone=availability_domain,
             price=instance_offer.price,
             username="ubuntu",
             ssh_port=22,