skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/provision/fluidstack/instance.py

@@ -1,13 +1,13 @@
 """FluidStack instance provisioning."""
 import os
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
-from sky import authentication as auth
 from sky import exceptions
 from sky import sky_logging
 from sky.provision import common
 from sky.provision.fluidstack import fluidstack_utils as utils
+from sky.utils import auth_utils
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import status_lib
@@ -26,7 +26,8 @@ logger = sky_logging.init_logger(__name__)
 
 def get_internal_ip(node_info: Dict[str, Any]) -> None:
     node_info['internal_ip'] = node_info['ip_address']
-    private_key_path, _ = auth.get_or_generate_keys()
+
+    private_key_path, _ = auth_utils.get_or_generate_keys()
     runner = command_runner.SSHCommandRunner(
         (node_info['ip_address'], 22),
         ssh_user='ubuntu',
@@ -77,10 +78,10 @@ def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
     return head_instance_id
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
-
+    del cluster_name  # unused
     pending_status = ['pending', 'provisioning']
     while True:
         instances = _filter_instances(cluster_name_on_cloud, pending_status)
@@ -286,11 +287,14 @@ def get_cluster_info(
 
 
 def query_instances(
+    cluster_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)
     instances = _filter_instances(cluster_name_on_cloud, None)
@@ -301,7 +305,8 @@ def query_instances(
         'failed': status_lib.ClusterStatus.INIT,
         'terminated': None,
     }
-    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
     for inst_id, inst in instances.items():
         if inst['status'] not in status_map:
             with ux_utils.print_exception_no_traceback():
@@ -310,7 +315,7 @@ def query_instances(
         status = status_map.get(inst['status'], None)
         if non_terminated_only and status is None:
             continue
-        statuses[inst_id] = status
+        statuses[inst_id] = (status, None)
     return statuses
 
 

sky/provision/gcp/__init__.py

@@ -1,6 +1,7 @@
 """GCP provisioner for SkyPilot."""
 
 from sky.provision.gcp.config import bootstrap_instances
+from sky.provision.gcp.instance import cleanup_custom_multi_network
 from sky.provision.gcp.instance import cleanup_ports
 from sky.provision.gcp.instance import get_cluster_info
 from sky.provision.gcp.instance import open_ports

sky/provision/gcp/config.py

@@ -5,11 +5,14 @@ import time
 import typing
 from typing import Any, Dict, List, Set, Tuple
 
+from typing_extensions import TypedDict
+
 from sky.adaptors import gcp
 from sky.clouds.utils import gcp_utils
 from sky.provision import common
 from sky.provision.gcp import constants
 from sky.provision.gcp import instance_utils
+from sky.utils import resources_utils
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +78,30 @@ def wait_for_compute_global_operation(project_name, operation, compute):
     return result
 
 
+def wait_for_compute_region_operation(project_name, region, operation, compute):
+    """Poll for region compute operation until finished."""
+    logger.info('wait_for_compute_region_operation: '
+                'Waiting for operation {} to finish...'.format(
+                    operation['name']))
+
+    for _ in range(constants.MAX_POLLS):
+        result = (compute.regionOperations().get(
+            project=project_name,
+            region=region,
+            operation=operation['name'],
+        ).execute())
+        if 'error' in result:
+            raise Exception(result['error'])
+
+        if result['status'] == 'DONE':
+            logger.info('wait_for_compute_region_operation: Operation done.')
+            break
+
+        time.sleep(constants.POLL_INTERVAL)
+
+    return result
+
+
 def _create_crm(gcp_credentials=None):
     return gcp.build('cloudresourcemanager',
                      'v1',
@@ -168,6 +195,7 @@ def bootstrap_instances(
     iam_role = _configure_iam_role(config, crm, iam)
     config.node_config.update(iam_role)
     config = _configure_subnet(region, cluster_name, config, compute)
+    config = _configure_placement_policy(region, cluster_name, config, compute)
 
     return config
 
@@ -248,7 +276,7 @@ def _is_permission_satisfied(service_account, crm, iam, required_permissions,
     # For example, `roles/iam.serviceAccountUser` can be granted at the
     # skypilot-v1 service account level, which can be checked with
     # service_account_policy = iam.projects().serviceAccounts().getIamPolicy(
-    #    resource=f'projects/{project_id}/serviceAcccounts/{email}').execute()
+    #    resource=f'projects/{project_id}/serviceAccounts/{email}').execute()
     # We now skip the check for `iam.serviceAccounts.actAs` permission for
     # simplicity as it can be granted at the service account level.
     def check_permissions(policy, required_permissions):
@@ -389,6 +417,9 @@ def _configure_iam_role(config: common.ProvisionConfig, crm, iam) -> dict:
     return iam_role
 
 
+AllowedList = TypedDict('AllowedList', {'IPProtocol': str, 'ports': List[str]})
+
+
 def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
                           compute):
     """Check if the firewall rules in the VPC are sufficient."""
@@ -440,7 +471,7 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
             }
         """
         source2rules: Dict[Tuple[str, str], Dict[str, Set[int]]] = {}
-        source2allowed_list: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
+        source2allowed_list: Dict[Tuple[str, str], List[AllowedList]] = {}
         for rule in rules:
             # Rules applied to specific VM (targetTags) may not work for the
             # current VM, so should be skipped.
@@ -506,7 +537,23 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
     return True
 
 
-def _create_rules(project_id: str, compute, rules, vpc_name):
+def _delete_rules(project_id: str, compute, rules, vpc_name: str):
+    for rule_ori in rules:
+        # Query firewall rule by its name (unique in a project).
+        rule_name = rule_ori['name'].format(VPC_NAME=vpc_name)
+        rule_list = _list_firewall_rules(project_id,
+                                         compute,
+                                         filter=f'(name={rule_name})')
+        for rule in rule_list:
+            logger.info(f'Deleting firewall rule {rule["name"]}')
+            _delete_firewall_rule(project_id, compute, rule['name'])
+
+
+def _create_rules(project_id: str,
+                  compute,
+                  rules,
+                  vpc_name,
+                  recreate: bool = True):
     opertaions = []
     for rule in rules:
         # Query firewall rule by its name (unique in a project).
@@ -516,7 +563,11 @@ def _create_rules(project_id: str, compute, rules, vpc_name):
                                          compute,
                                          filter=f'(name={rule_name})')
         if rule_list:
-            _delete_firewall_rule(project_id, compute, rule_name)
+            if recreate:
+                _delete_firewall_rule(project_id, compute, rule_name)
+            else:
+                logger.info(f'Rule {rule_name} already exists')
+                continue
 
         body = rule.copy()
         body['name'] = body['name'].format(VPC_NAME=vpc_name)
@@ -660,6 +711,149 @@ def get_usable_vpc_and_subnet(
     return usable_vpc_name, usable_subnet
 
 
+def get_gpu_direct_usable_vpcs_and_subnets(
+    cluster_name: str,
+    region: str,
+    config: common.ProvisionConfig,
+    compute,
+) -> List[Tuple[str, 'google.cloud.compute_v1.types.compute.Subnetwork']]:
+    """Return a list of usable VPCs and subnets for GPU Direct."""
+    project_id = config.provider_config['project_id']
+    vpc_subnet_pairs = []
+
+    # TODO(hailong): Determine the num_vpcs per different GPU Direct types
+    num_vpcs = constants.SKYPILOT_GPU_DIRECT_VPC_NUM
+
+    cidr_prefix = constants.SKYPILOT_GPU_DIRECT_VPC_CIDR_PREFIX
+    for i in range(num_vpcs):
+        vpc_name = get_gpu_direct_vpc_name(cluster_name, i)
+        subnet_name = f'{vpc_name}-sub'
+        subnet_cidr_range = f'{cidr_prefix}.{i}.0/24'
+        # Check if VPC exists
+        vpc_list = _list_vpcnets(project_id, compute, filter=f'name={vpc_name}')
+        if not vpc_list:
+            body = constants.VPC_TEMPLATE.copy()
+            body['mtu'] = 8244
+            body['autoCreateSubnetworks'] = False
+            body['name'] = vpc_name
+            body['selfLink'] = body['selfLink'].format(PROJ_ID=project_id,
+                                                       VPC_NAME=vpc_name)
+            _create_vpcnet(project_id, compute, body)
+        # Check if subnet exists
+        subnets = _list_subnets(project_id, region, compute, network=vpc_name)
+        if not subnets:
+            _create_subnet(project_id, region, compute, vpc_name, subnet_name,
+                           subnet_cidr_range)
+            subnets = _list_subnets(project_id,
+                                    region,
+                                    compute,
+                                    network=vpc_name)
+        # Apply firewall rules
+        # No need to recreate the rules if exist,
+        # as they are totally managed by SkyPilot,
+        # in this case, we can skip the rules creation during failover
+        _create_rules(project_id,
+                      compute,
+                      constants.FIREWALL_RULES_TEMPLATE,
+                      vpc_name,
+                      recreate=False)
+        vpc_subnet_pairs.append((vpc_name, subnets[0]))
+    return vpc_subnet_pairs
+
+
+def get_gpu_direct_vpc_name(cluster_name: str, i: int) -> str:
+    """Get the name of the GPU Direct VPC."""
+    if i == 0:
+        return f'{cluster_name}-mgmt-net'
+    else:
+        return f'{cluster_name}-data-net-{i}'
+
+
+def delete_gpu_direct_vpcs_and_subnets(
+    cluster_name: str,
+    project_id: str,
+    region: str,
+    keep_global_resources: bool = False,
+):
+    """Delete GPU Direct subnets, firewalls, and VPCs.
+
+    Args:
+        cluster_name: The name of the cluster.
+        project_id: The ID of the project.
+        region: The region of the cluster.
+        keep_global_resources: Whether to keep the global resources. If True,
+            only delete the subnets. Otherwise, delete all the firewalls,
+            subnets, and VPCs.
+    """
+    compute = _create_compute()
+
+    # TODO(hailong): Determine the num_vpcs per different GPU Direct types
+    num_vpcs = constants.SKYPILOT_GPU_DIRECT_VPC_NUM
+
+    for i in range(num_vpcs):
+        vpc_name = get_gpu_direct_vpc_name(cluster_name, i)
+        # Check if VPC exists
+        vpc_list = _list_vpcnets(project_id, compute, filter=f'name={vpc_name}')
+        if not vpc_list:
+            continue
+        for vpc in vpc_list:
+            subnets = _list_subnets(project_id,
+                                    region,
+                                    compute,
+                                    network=vpc['name'])
+            for subnet in subnets:
+                logger.info(f'Deleting subnet {subnet["name"]}')
+                _delete_subnet(project_id, region, compute, subnet['name'])
+
+            if not keep_global_resources:
+                # For failover, keep_global_resources would be true,
+                # we don't delete the rules and VPCs,
+                # which are global resources and can be reused.
+                _delete_rules(project_id, compute,
+                              constants.FIREWALL_RULES_TEMPLATE, vpc['name'])
+                logger.info(f'Deleting VPC {vpc["name"]}')
+                _delete_vpcnet(project_id, compute, vpc['name'])
+
+
+def _configure_placement_policy(region: str, cluster_name: str,
+                                config: common.ProvisionConfig, compute):
+    """Configure placement group for GPU Direct."""
+    node_config = config.node_config
+    project_id = config.provider_config['project_id']
+    group_placement_policy = config.provider_config.get('placement_policy',
+                                                        None)
+    # If the placement policy is not compact,
+    # or the managed instance group is specified,
+    # skip the placement policy creation.
+    # If placement policy is specified together with managed instance group,
+    # it will cause the following error:
+    # Reason: [{'code': 'UNSUPPORTED_OPERATION',
+    # 'message': 'Creating queued resource with
+    # resource policies is not supported.'}]
+    mig_configuration = config.provider_config.get('use_managed_instance_group',
+                                                   False)
+    if (group_placement_policy is None or group_placement_policy.lower() !=
+            constants.COMPACT_GROUP_PLACEMENT_POLICY or mig_configuration):
+        return config
+
+    policy_name = f'{cluster_name}-placement-policy'
+    resource_policy = {
+        'name': policy_name,
+        'groupPlacementPolicy': {
+            'collocation': constants.COLLOCATED_COLLOCATION,
+        }
+    }
+    # Try to get the placement policy first, if not found, create it
+    placement_policy = _get_placement_policy(project_id, region, compute,
+                                             policy_name)
+    if not placement_policy:
+        logger.info(f'Creating placement policy {policy_name}'
+                    f' for cluster {cluster_name}')
+        _create_placement_policy(project_id, region, compute, resource_policy)
+    node_config['resourcePolicies'] = [policy_name]
+    return config
+
+
 def _configure_subnet(region: str, cluster_name: str,
                       config: common.ProvisionConfig, compute):
     """Pick a reasonable subnet if not specified by the config."""
@@ -671,25 +865,56 @@ def _configure_subnet(region: str, cluster_name: str,
     if 'networkInterfaces' in node_config or 'networkConfig' in node_config:
         return config
 
-    # SkyPilot: make sure there's a usable VPC
-    _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region, config,
-                                                  compute)
-
-    default_interfaces = [{
-        'subnetwork': default_subnet['selfLink'],
-        'accessConfigs': [{
-            'name': 'External NAT',
-            'type': 'ONE_TO_ONE_NAT',
-        }]
-    }]
-    # Add gVNIC if specified in config
+    default_interfaces = []
+    enable_gpu_direct = config.provider_config.get('enable_gpu_direct', False)
     enable_gvnic = config.provider_config.get('enable_gvnic', False)
-    if enable_gvnic:
-        default_interfaces[0]['nicType'] = 'gVNIC'
+    network_tier = config.provider_config.get('network_tier', 'standard')
+    if (enable_gpu_direct or
+            network_tier == resources_utils.NetworkTier.BEST.value):
+        if not enable_gvnic:
+            logger.warning(
+                'Enable GPU Direct requires gvnic to be enabled, enabling gvnic'
+            )
+            config.provider_config['enable_gvnic'] = True
+            enable_gvnic = True
+        if 'machineType' not in node_config or node_config[
+                'machineType'] not in constants.GPU_DIRECT_TCPX_INSTANCE_TYPES:
+            raise ValueError(
+                'Enable GPU Direct requires machineType to be one of '
+                f'{constants.GPU_DIRECT_TCPX_INSTANCE_TYPES}')
+        logger.info(f'Enable GPU Direct for cluster {cluster_name} '
+                    f'with machineType {node_config["machineType"]}')
+        vpc_subnet_pairs = get_gpu_direct_usable_vpcs_and_subnets(
+            cluster_name, region, config, compute)
+        for _, subnet in vpc_subnet_pairs:
+            default_interfaces.append({
+                'subnetwork': subnet['selfLink'],
+                'accessConfigs': [{
+                    'name': 'External NAT',
+                    'type': 'ONE_TO_ONE_NAT',
+                }],
+                'nicType': 'gVNIC'
+            })
+    else:
+        # SkyPilot: make sure there's a usable VPC
+        _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region,
+                                                      config, compute)
+
+        default_interfaces = [{
+            'subnetwork': default_subnet['selfLink'],
+            'accessConfigs': [{
+                'name': 'External NAT',
+                'type': 'ONE_TO_ONE_NAT',
+            }]
+        }]
+        # Add gVNIC if specified in config
+        if enable_gvnic:
+            default_interfaces[0]['nicType'] = 'gVNIC'
     enable_external_ips = _enable_external_ips(config)
     if not enable_external_ips:
         # Removing this key means the VM will not be assigned an external IP.
-        default_interfaces[0].pop('accessConfigs')
+        for interface in default_interfaces:
+            interface.pop('accessConfigs')
 
     # The not applicable key will be removed during node creation
 
@@ -747,6 +972,14 @@ def _list_vpcnets(project_id: str, compute, filter=None):  # pylint: disable=red
             if 'items' in response else [])
 
 
+def _delete_vpcnet(project_id: str, compute, vpcnet_name: str):
+    operation = compute.networks().delete(
+        project=project_id,
+        network=vpcnet_name,
+    ).execute()
+    return wait_for_compute_global_operation(project_id, operation, compute)
+
+
 def _list_subnets(
         project_id: str,
         region: str,
@@ -840,3 +1073,52 @@ def _add_iam_policy_binding(service_account, policy, crm, iam):
     ).execute())
 
     return result
+
+
+def _create_subnet(project_id: str, region: str, compute, vpc_name: str,
+                   subnet_name: str, ip_cidr_range: str):
+    body = {
+        'name': subnet_name,
+        'ipCidrRange': ip_cidr_range,
+        'network': f'projects/{project_id}/global/networks/{vpc_name}',
+        'region': region,
+    }
+    operation = compute.subnetworks().insert(project=project_id,
+                                             region=region,
+                                             body=body).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _delete_subnet(project_id: str, region: str, compute, subnet_name: str):
+    operation = compute.subnetworks().delete(
+        project=project_id,
+        region=region,
+        subnetwork=subnet_name,
+    ).execute()
+    return wait_for_compute_region_operation(project_id, region, operation,
+                                             compute)
+
+
+def _create_placement_policy(project_id: str, region: str, compute,
+                             placement_policy: dict):
+    operation = compute.resourcePolicies().insert(
+        project=project_id, region=region, body=placement_policy).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _get_placement_policy(project_id: str, region: str, compute, name: str):
+    try:
+        placement_policy = (compute.resourcePolicies().get(
+            project=project_id,
+            region=region,
+            resourcePolicy=name,
+        ).execute())
+    except gcp.http_error_exception() as e:
+        if e.resp.status == 404:
+            return None
+        raise
+    return placement_policy