skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/provision/gcp/config.py

@@ -397,7 +397,7 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
     operation = compute.networks().getEffectiveFirewalls(project=project_id,
                                                          network=vpc_name)
     response = operation.execute()
-    if len(response) == 0:
+    if not response:
         return False
     effective_rules = response['firewalls']
 
@@ -515,7 +515,7 @@ def _create_rules(project_id: str, compute, rules, vpc_name):
         rule_list = _list_firewall_rules(project_id,
                                          compute,
                                          filter=f'(name={rule_name})')
-        if len(rule_list) > 0:
+        if rule_list:
             _delete_firewall_rule(project_id, compute, rule_name)
 
         body = rule.copy()
@@ -624,7 +624,7 @@ def get_usable_vpc_and_subnet(
     vpc_list = _list_vpcnets(project_id,
                              compute,
                              filter=f'name={constants.SKYPILOT_VPC_NAME}')
-    if len(vpc_list) == 0:
+    if not vpc_list:
         body = constants.VPC_TEMPLATE.copy()
         body['name'] = body['name'].format(VPC_NAME=constants.SKYPILOT_VPC_NAME)
         body['selfLink'] = body['selfLink'].format(
@@ -670,9 +670,14 @@ def _configure_subnet(region: str, cluster_name: str,
         'accessConfigs': [{
             'name': 'External NAT',
             'type': 'ONE_TO_ONE_NAT',
-        }],
+        }]
     }]
-    if config.provider_config.get('use_internal_ips', False):
+    # Add gVNIC if specified in config
+    enable_gvnic = config.provider_config.get('enable_gvnic', False)
+    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')
 
@@ -686,14 +691,19 @@ def _configure_subnet(region: str, cluster_name: str,
         node_config['networkConfig'] = copy.deepcopy(default_interfaces)[0]
         # TPU doesn't have accessConfigs
         node_config['networkConfig'].pop('accessConfigs', None)
-        if config.provider_config.get('use_internal_ips', False):
-            node_config['networkConfig']['enableExternalIps'] = False
-        else:
-            node_config['networkConfig']['enableExternalIps'] = True
+        node_config['networkConfig']['enableExternalIps'] = enable_external_ips
 
     return config
 
 
+def _enable_external_ips(config: common.ProvisionConfig) -> bool:
+    force_enable_external_ips = config.provider_config.get(
+        'force_enable_external_ips', False)
+    use_internal_ips = config.provider_config.get('use_internal_ips', False)
+
+    return force_enable_external_ips or not use_internal_ips
+
+
 def _delete_firewall_rule(project_id: str, compute, name):
     operation = (compute.firewalls().delete(project=project_id,
                                             firewall=name).execute())

sky/provision/gcp/constants.py

@@ -142,7 +142,7 @@ FIREWALL_RULES_TEMPLATE = [
 ]
 
 # A list of permissions required to run SkyPilot on GCP.
-# Keep this in sync with https://skypilot.readthedocs.io/en/latest/cloud-setup/cloud-permissions/gcp.html # pylint: disable=line-too-long
+# Keep this in sync with https://docs.skypilot.co/en/latest/cloud-setup/cloud-permissions/gcp.html # pylint: disable=line-too-long
 VM_MINIMAL_PERMISSIONS = [
     'compute.disks.create',
     'compute.disks.list',
@@ -214,3 +214,9 @@ POLL_INTERVAL = 1
 MAX_POLLS = 60 // POLL_INTERVAL
 # Stopping instances can take several minutes, so we increase the timeout
 MAX_POLLS_STOP = MAX_POLLS * 8
+
+# MIG constants
+MANAGED_INSTANCE_GROUP_CONFIG = 'managed-instance-group'
+DEFAULT_MANAGED_INSTANCE_GROUP_PROVISION_TIMEOUT = 900  # 15 minutes
+MIG_NAME_PREFIX = 'sky-mig-'
+INSTANCE_TEMPLATE_NAME_PREFIX = 'sky-it-'

sky/provision/gcp/instance.py

@@ -7,20 +7,16 @@ import time
 from typing import Any, Callable, Dict, Iterable, List, Optional, Type
 
 from sky import sky_logging
-from sky import status_lib
 from sky.adaptors import gcp
 from sky.provision import common
+from sky.provision import constants as provision_constants
 from sky.provision.gcp import constants
 from sky.provision.gcp import instance_utils
 from sky.utils import common_utils
+from sky.utils import status_lib
 
 logger = sky_logging.init_logger(__name__)
 
-TAG_SKYPILOT_HEAD_NODE = 'skypilot-head-node'
-# Tag uniquely identifying all nodes of a cluster
-TAG_RAY_CLUSTER_NAME = 'ray-cluster-name'
-TAG_RAY_NODE_KIND = 'ray-node-type'
-
 _INSTANCE_RESOURCE_NOT_FOUND_PATTERN = re.compile(
     r'The resource \'projects/.*/zones/.*/instances/.*\' was not found')
 
@@ -56,6 +52,8 @@ def _filter_instances(
 # non_terminated_only=True?
 # Will there be callers who would want this to be False?
 # stop() and terminate() for example already implicitly assume non-terminated.
+# Currently, even with non_terminated_only=False, we may not have a dict entry
+# for terminated instances, if they have already been fully deleted.
 @common_utils.retry
 def query_instances(
     cluster_name_on_cloud: str,
@@ -66,7 +64,9 @@ def query_instances(
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     zone = provider_config['availability_zone']
     project_id = provider_config['project_id']
-    label_filters = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    label_filters = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
 
     handler: Type[
         instance_utils.GCPInstance] = instance_utils.GCPComputeInstance
@@ -124,15 +124,15 @@ def _wait_for_operations(
             logger.debug(
                 f'wait_for_compute_{op_type}_operation: '
                 f'Waiting for operation {operation["name"]} to finish...')
-            handler.wait_for_operation(operation, project_id, zone)
+            handler.wait_for_operation(operation, project_id, zone=zone)
 
 
 def _get_head_instance_id(instances: List) -> Optional[str]:
     head_instance_id = None
     for inst in instances:
         labels = inst.get('labels', {})
-        if (labels.get(TAG_RAY_NODE_KIND) == 'head' or
-                labels.get(TAG_SKYPILOT_HEAD_NODE) == '1'):
+        if (labels.get(provision_constants.TAG_RAY_NODE_KIND) == 'head' or
+                labels.get(provision_constants.TAG_SKYPILOT_HEAD_NODE) == '1'):
             head_instance_id = inst['name']
             break
     return head_instance_id
@@ -158,12 +158,16 @@ def _run_instances(region: str, cluster_name_on_cloud: str,
     resource: Type[instance_utils.GCPInstance]
     if node_type == instance_utils.GCPNodeType.COMPUTE:
         resource = instance_utils.GCPComputeInstance
+    elif node_type == instance_utils.GCPNodeType.MIG:
+        resource = instance_utils.GCPManagedInstanceGroup
     elif node_type == instance_utils.GCPNodeType.TPU:
         resource = instance_utils.GCPTPUVMInstance
     else:
         raise ValueError(f'Unknown node type {node_type}')
 
-    filter_labels = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    filter_labels = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
 
     # wait until all stopping instances are stopped/terminated
     while True:
@@ -264,12 +268,16 @@ def _run_instances(region: str, cluster_name_on_cloud: str,
     if config.resume_stopped_nodes and to_start_count > 0 and stopped_instances:
         resumed_instance_ids = [n['name'] for n in stopped_instances]
         if resumed_instance_ids:
-            for instance_id in resumed_instance_ids:
-                resource.start_instance(instance_id, project_id,
-                                        availability_zone)
-                resource.set_labels(project_id, availability_zone, instance_id,
-                                    labels)
-        to_start_count -= len(resumed_instance_ids)
+            resumed_instance_ids = resource.start_instances(
+                cluster_name_on_cloud, project_id, availability_zone,
+                resumed_instance_ids, labels)
+        # In MIG case, the resumed_instance_ids will include the previously
+        # PENDING and RUNNING instances. To avoid double counting, we need to
+        # remove them from the resumed_instance_ids.
+        ready_instances = set(resumed_instance_ids)
+        ready_instances |= set([n['name'] for n in running_instances])
+        ready_instances |= set([n['name'] for n in pending_instances])
+        to_start_count = config.count - len(ready_instances)
 
         if head_instance_id is None:
             head_instance_id = resource.create_node_tag(
@@ -281,9 +289,14 @@ def _run_instances(region: str, cluster_name_on_cloud: str,
 
     if to_start_count > 0:
         errors, created_instance_ids = resource.create_instances(
-            cluster_name_on_cloud, project_id, availability_zone,
-            config.node_config, labels, to_start_count,
-            head_instance_id is None)
+            cluster_name_on_cloud,
+            project_id,
+            availability_zone,
+            config.node_config,
+            labels,
+            to_start_count,
+            total_count=config.count,
+            include_head_node=head_instance_id is None)
         if errors:
             error = common.ProvisionerError('Failed to launch instances.')
             error.errors = errors
@@ -387,7 +400,9 @@ def get_cluster_info(
     assert provider_config is not None, cluster_name_on_cloud
     zone = provider_config['availability_zone']
     project_id = provider_config['project_id']
-    label_filters = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    label_filters = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
 
     handlers: List[Type[instance_utils.GCPInstance]] = [
         instance_utils.GCPComputeInstance
@@ -415,7 +430,7 @@ def get_cluster_info(
         project_id,
         zone,
         {
-            **label_filters, TAG_RAY_NODE_KIND: 'head'
+            **label_filters, provision_constants.TAG_RAY_NODE_KIND: 'head'
         },
         lambda h: [h.RUNNING_STATE],
     )
@@ -441,14 +456,16 @@ def stop_instances(
     assert provider_config is not None, cluster_name_on_cloud
     zone = provider_config['availability_zone']
     project_id = provider_config['project_id']
-    label_filters = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    label_filters = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
 
     tpu_node = provider_config.get('tpu_node')
     if tpu_node is not None:
         instance_utils.delete_tpu_node(project_id, zone, tpu_node)
 
     if worker_only:
-        label_filters[TAG_RAY_NODE_KIND] = 'worker'
+        label_filters[provision_constants.TAG_RAY_NODE_KIND] = 'worker'
 
     handlers: List[Type[instance_utils.GCPInstance]] = [
         instance_utils.GCPComputeInstance
@@ -510,9 +527,18 @@ def terminate_instances(
     if tpu_node is not None:
         instance_utils.delete_tpu_node(project_id, zone, tpu_node)
 
-    label_filters = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    use_mig = provider_config.get('use_managed_instance_group', False)
+    if use_mig:
+        # Deleting the MIG will also delete the instances.
+        instance_utils.GCPManagedInstanceGroup.delete_mig(
+            project_id, zone, cluster_name_on_cloud)
+        return
+
+    label_filters = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
     if worker_only:
-        label_filters[TAG_RAY_NODE_KIND] = 'worker'
+        label_filters[provision_constants.TAG_RAY_NODE_KIND] = 'worker'
 
     handlers: List[Type[instance_utils.GCPInstance]] = [
         instance_utils.GCPComputeInstance
@@ -555,7 +581,9 @@ def open_ports(
     project_id = provider_config['project_id']
     firewall_rule_name = provider_config['firewall_rule']
 
-    label_filters = {TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    label_filters = {
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud
+    }
     handlers: List[Type[instance_utils.GCPInstance]] = [
         instance_utils.GCPComputeInstance,
         instance_utils.GCPTPUVMInstance,
@@ -606,13 +634,6 @@ def cleanup_ports(
     del ports  # Unused.
     assert provider_config is not None, cluster_name_on_cloud
     project_id = provider_config['project_id']
-    if 'ports' in provider_config:
-        # Backward compatibility for old provider config.
-        # TODO(tian): remove this after 2 minor releases, 0.6.0.
-        for port in provider_config['ports']:
-            firewall_rule_name = f'user-ports-{cluster_name_on_cloud}-{port}'
-            instance_utils.GCPComputeInstance.delete_firewall_rule(
-                project_id, firewall_rule_name)
     if 'firewall_rule' in provider_config:
         firewall_rule_name = provider_config['firewall_rule']
         instance_utils.GCPComputeInstance.delete_firewall_rule(