skypilot-nightly 1.0.0.dev20250521__py3-none-any.whl → 1.0.0.dev20250522__py3-none-any.whl

sky/provision/gcp/volume_utils.py

@@ -0,0 +1,247 @@
+"""Utilities for GCP volumes."""
+from typing import Any, Dict, List, Optional
+
+from sky import clouds
+from sky import exceptions
+from sky import sky_logging
+from sky.adaptors import gcp
+from sky.provision.gcp import constants
+from sky.utils import resources_utils
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def get_data_disk_tier_mapping(
+    instance_type: Optional[str],) -> Dict[resources_utils.DiskTier, str]:
+    # Define the default mapping from disk tiers to disk types.
+    # Refer to https://cloud.google.com/compute/docs/disks/hyperdisks
+    # and https://cloud.google.com/compute/docs/disks/persistent-disks
+    tier2name = {
+        resources_utils.DiskTier.ULTRA: 'pd-extreme',
+        resources_utils.DiskTier.HIGH: 'pd-ssd',
+        resources_utils.DiskTier.MEDIUM: 'pd-balanced',
+        resources_utils.DiskTier.LOW: 'pd-standard',
+    }
+
+    if instance_type is None:
+        return tier2name
+
+    # Remap series-specific disk types.
+    series = instance_type.split('-')[0]
+
+    if series in ['a4', 'x4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+    elif series in ['m4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 112:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['c4', 'c4a', 'c4d']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['a3']:
+        if (instance_type.startswith('a3-ultragpu') or
+                instance_type.startswith('a3-megagpu') or
+                instance_type.startswith('a3-edgegpu')):
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+            tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+            tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+            tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        elif instance_type.startswith('a3-highgpu'):
+            tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+            if instance_type.startswith('a3-highgpu-8g'):
+                tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+                tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+                tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+            elif instance_type.startswith('a3-highgpu-4g'):
+                tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+            else:
+                tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+    elif series in ['c3d']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 60:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['c3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 88:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['n4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+    elif series in ['n2d', 'n1', 't2d', 't2a', 'e2', 'c2', 'c2d', 'a2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+    elif series in ['z3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['h3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['m3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['m2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+    elif series in ['m1']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 80:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['g2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['n2']:
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+        elif num_cpus >= 80:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+
+    return tier2name
+
+
+def validate_instance_volumes(
+    instance_type: Optional[str],
+    volumes: Optional[List[Dict[str, Any]]],
+) -> None:
+    if not volumes:
+        return
+    if instance_type is None:
+        logger.warning('Instance type is not specified,'
+                       ' skipping instance volume validation')
+        return
+    instance_volume_count = 0
+    for volume in volumes:
+        if volume['storage_type'] == resources_utils.StorageType.INSTANCE:
+            instance_volume_count += 1
+    if (instance_type in constants.SSD_AUTO_ATTACH_MACHINE_TYPES and
+            instance_volume_count >
+            constants.SSD_AUTO_ATTACH_MACHINE_TYPES[instance_type]):
+        raise exceptions.ResourcesUnavailableError(
+            f'The instance type {instance_type} supports'
+            f' {constants.SSD_AUTO_ATTACH_MACHINE_TYPES[instance_type]}'
+            f'  instance storage, but {instance_volume_count} are specified')
+    # TODO(hailong):
+    # check the instance storage count for the other instance types,
+    # refer to https://cloud.google.com/compute/docs/disks/local-ssd
+
+
+def translate_attach_mode(attach_mode: resources_utils.DiskAttachMode) -> str:
+    if attach_mode == resources_utils.DiskAttachMode.READ_ONLY:
+        return 'READ_ONLY'
+    return 'READ_WRITE'
+
+
+def check_volume_name_exist_in_region(
+        project_id: str, region: clouds.Region, use_mig: bool,
+        volume_name: str) -> Optional[Dict[str, Any]]:
+    """Check if the volume name exists and return the volume info."""
+    logger.debug(f'Checking volume {volume_name} in region {region}')
+    try:
+        compute = gcp.build('compute',
+                            'v1',
+                            credentials=None,
+                            cache_discovery=False)
+    except gcp.credential_error_exception():
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('Not able to build compute client') from None
+
+    # Get all the zones in the region
+    all_zones = compute.zones().list(project=project_id).execute()
+    region_zones = []
+    if 'items' in all_zones:
+        for zone in all_zones['items']:
+            if zone['region'].split('/')[-1] == region.name:
+                region_zones.append(zone['name'])
+    volume_info = None
+    for zone in region_zones:
+        try:
+            volume_info = compute.disks().get(project=project_id,
+                                              zone=zone,
+                                              disk=volume_name).execute()
+            if volume_info is not None:
+                if use_mig:
+                    # With MIG, instance template will be used, in this case,
+                    # the `selfLink` for zonal disk needs to be the volume name
+                    # Refer to https://cloud.google.com/compute/docs/
+                    # reference/rest/v1/instances/insert
+                    volume_info['selfLink'] = volume_name
+                volume_info['available_zones'] = [zone]
+                return volume_info
+        except gcp.http_error_exception() as e:
+            if e.resp.status == 403:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Not able to access the volume '
+                                     f'{volume_name!r}') from None
+            if e.resp.status == 404:
+                continue  # Try next zone
+            raise
+
+    # If not found in any zone, check region disk
+    try:
+        volume_info = compute.regionDisks().get(project=project_id,
+                                                region=region.name,
+                                                disk=volume_name).execute()
+        # 'replicaZones':
+        #  ['https://xxx/compute/v1/projects/sky-dev-465/zones/us-central1-a',
+        # 'https://xxx/compute/v1/projects/sky-dev-465/zones/us-central1-c']
+        if volume_info is not None and 'replicaZones' in volume_info:
+            replica_zones = [
+                zone.split('/')[-1] for zone in volume_info['replicaZones']
+            ]
+            volume_info['available_zones'] = replica_zones
+        return volume_info
+    except gcp.http_error_exception() as e:
+        if e.resp.status == 403:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Not able to access the volume '
+                                 f'{volume_name!r}') from None
+        if e.resp.status == 404:
+            logger.warning(
+                f'Volume {volume_name} is not found in region {region}.'
+                f' It will be created.')
+            return volume_info
+        raise
+
+
+def check_volume_zone_match(volume_name: str,
+                            zones: Optional[List[clouds.Zone]],
+                            available_zones: List[str]):
+    if zones is None:
+        return None
+    for zone in zones:
+        if zone.name in available_zones:
+            return None
+    with ux_utils.print_exception_no_traceback():
+        # Return a ResourcesUnavailableError to trigger failover
+        raise exceptions.ResourcesUnavailableError(
+            f'Volume {volume_name} not available in zones {zones}') from None

sky/resources.py

@@ -98,7 +98,7 @@ class Resources:
     """
     # If any fields changed, increment the version. For backward compatibility,
     # modify the __setstate__ method to handle the old version.
-    _VERSION = 23
+    _VERSION = 24
 
     def __init__(
         self,
@@ -120,6 +120,7 @@ class Resources:
         ports: Optional[Union[int, str, List[str], Tuple[str]]] = None,
         labels: Optional[Dict[str, str]] = None,
         autostop: Union[bool, int, Dict[str, Any], None] = None,
+        volumes: Optional[List[Dict[str, Any]]] = None,
         # Internal use only.
         # pylint: disable=invalid-name
         _docker_login_config: Optional[docker_utils.DockerLoginConfig] = None,
@@ -210,6 +211,7 @@ class Resources:
             not supported and will be ignored.
           autostop: the autostop configuration to use. For launched resources,
             may or may not correspond to the actual current autostop config.
+          volumes: the volumes to mount on the instance.
           _docker_login_config: the docker configuration to use. This includes
             the docker username, password, and registry server. If None, skip
             docker login.
@@ -337,6 +339,7 @@ class Resources:
         self._set_memory(memory)
         self._set_accelerators(accelerators, accelerator_args)
         self._set_autostop_config(autostop)
+        self._set_volumes(volumes)
 
     def validate(self):
         """Validate the resources and infer the missing fields if possible."""
@@ -347,6 +350,7 @@ class Resources:
         self._try_validate_managed_job_attributes()
         self._try_validate_image_id()
         self._try_validate_disk_tier()
+        self._try_validate_volumes()
         self._try_validate_ports()
         self._try_validate_labels()
 
@@ -566,6 +570,10 @@ class Resources:
     def labels(self) -> Optional[Dict[str, str]]:
         return self._labels
 
+    @property
+    def volumes(self) -> Optional[List[Dict[str, Any]]]:
+        return self._volumes
+
     @property
     def autostop_config(self) -> Optional[AutostopConfig]:
         """The requested autostop config.
@@ -759,6 +767,91 @@ class Resources:
     ) -> None:
         self._autostop_config = AutostopConfig.from_yaml_config(autostop)
 
+    def _set_volumes(
+        self,
+        volumes: Optional[List[Dict[str, Any]]],
+    ) -> None:
+        if not volumes:
+            self._volumes = None
+            return
+        valid_volumes = []
+        supported_tiers = [tier.value for tier in resources_utils.DiskTier]
+        supported_storage_types = [
+            storage_type.value for storage_type in resources_utils.StorageType
+        ]
+        supported_attach_modes = [
+            attach_mode.value for attach_mode in resources_utils.DiskAttachMode
+        ]
+        network_type = resources_utils.StorageType.NETWORK
+        read_write_mode = resources_utils.DiskAttachMode.READ_WRITE
+        for volume in volumes:
+            if 'path' not in volume:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(f'Invalid volume {volume!r}. '
+                                     f'Volume must have a "path" field.')
+            if 'storage_type' not in volume:
+                volume['storage_type'] = network_type
+            else:
+                if isinstance(volume['storage_type'], str):
+                    storage_type_str = str(volume['storage_type']).lower()
+                    if storage_type_str not in supported_storage_types:
+                        logger.warning(
+                            f'Invalid storage_type {storage_type_str!r}. '
+                            f'Set it to '
+                            f'{network_type.value}.')
+                        volume['storage_type'] = network_type
+                    else:
+                        volume['storage_type'] = resources_utils.StorageType(
+                            storage_type_str)
+            if 'auto_delete' not in volume:
+                volume['auto_delete'] = False
+            if 'attach_mode' in volume:
+                if isinstance(volume['attach_mode'], str):
+                    attach_mode_str = str(volume['attach_mode']).lower()
+                    if attach_mode_str not in supported_attach_modes:
+                        logger.warning(
+                            f'Invalid attach_mode {attach_mode_str!r}. '
+                            f'Set it to {read_write_mode.value}.')
+                        volume['attach_mode'] = read_write_mode
+                    else:
+                        volume['attach_mode'] = resources_utils.DiskAttachMode(
+                            attach_mode_str)
+            else:
+                volume['attach_mode'] = read_write_mode
+            if volume['storage_type'] == network_type:
+                if ('disk_size' in volume and
+                        round(volume['disk_size']) != volume['disk_size']):
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(f'Volume size must be an integer. '
+                                         f'Got: {volume["size"]}.')
+                if 'name' not in volume:
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(f'Network volume {volume["path"]} '
+                                         f'must have "name" field.')
+            elif 'name' in volume:
+                logger.info(f'Volume {volume["path"]} is a local disk. '
+                            f'The "name" field will be ignored.')
+                del volume['name']
+            if 'disk_tier' in volume:
+                if isinstance(volume['disk_tier'], str):
+                    disk_tier_str = str(volume['disk_tier']).lower()
+                    if disk_tier_str not in supported_tiers:
+                        logger.warning(
+                            f'Invalid disk_tier {disk_tier_str!r}. '
+                            f'Set it to {resources_utils.DiskTier.BEST.value}.')
+                        volume['disk_tier'] = resources_utils.DiskTier.BEST
+                    else:
+                        volume['disk_tier'] = resources_utils.DiskTier(
+                            disk_tier_str)
+            elif volume['storage_type'] == network_type:
+                logger.debug(
+                    f'No disk_tier specified for volume {volume["path"]}. '
+                    f'Set it to {resources_utils.DiskTier.BEST.value}.')
+                volume['disk_tier'] = resources_utils.DiskTier.BEST
+
+            valid_volumes.append(volume)
+        self._volumes = valid_volumes
+
     def is_launchable(self) -> bool:
         """Returns whether the resource is launchable."""
         return self.cloud is not None and self._instance_type is not None
@@ -1123,6 +1216,48 @@ class Resources:
                         f'Disk tier {self.disk_tier.value} is not supported '
                         f'for instance type {self.instance_type}.') from None
 
+    def _try_validate_volumes(self) -> None:
+        """Try to validate the volumes attribute.
+
+        Raises:
+            ValueError: if the attribute is invalid.
+        """
+        if self.volumes is None:
+            return
+        if self.cloud is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Cloud must be specified when '
+                                 'volumes are provided.')
+        if not self.cloud.is_same_cloud(clouds.GCP()):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Volumes are only supported for GCP'
+                                 f' not for {self.cloud}.')
+
+        need_region_or_zone = False
+        try:
+            for volume in self.volumes:
+                if ('name' in volume and volume['storage_type']
+                        == resources_utils.StorageType.NETWORK):
+                    need_region_or_zone = True
+                if 'disk_tier' not in volume:
+                    continue
+                # TODO(hailong): check instance local SSD
+                # support for instance_type.
+                # Refer to https://cloud.google.com/compute/docs/disks/local-ssd#machine-series-lssd # pylint: disable=line-too-long
+                self.cloud.check_disk_tier_enabled(self.instance_type,
+                                                   volume['disk_tier'])
+            if (need_region_or_zone and self._region is None and
+                    self._zone is None):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('When specifying the volume name, please'
+                                     ' also specify the region or zone.')
+        except exceptions.NotSupportedError:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Disk tier {volume["disk_tier"].value} is not '
+                    f'supported for instance type {self.instance_type}.'
+                ) from None
+
     def _try_validate_ports(self) -> None:
         """Try to validate the ports attribute.
 
@@ -1293,9 +1428,18 @@ class Resources:
             skypilot_config.get_nested(
                 (str(self.cloud).lower(), 'specific_reservations'), set()))
 
+        if isinstance(self.cloud, clouds.DummyCloud):
+            return self.cloud.get_reservations_available_resources(
+                instance_type='',
+                region='',
+                zone=None,
+                specific_reservations=specific_reservations)
+
         assert (self.cloud is not None and self.instance_type is not None and
-                self.region
-                is not None), ('Cloud, instance type, region must be specified')
+                self.region is not None), (
+                    f'Cloud, instance type, region must be specified. '
+                    f'Resources={self}, cloud={self.cloud}, '
+                    f'instance_type={self.instance_type}, region={self.region}')
         return self.cloud.get_reservations_available_resources(
             self.instance_type, self.region, self.zone, specific_reservations)
 
@@ -1483,6 +1627,7 @@ class Resources:
             ports=override.pop('ports', self.ports),
             labels=override.pop('labels', self.labels),
             autostop=override.pop('autostop', current_autostop_config),
+            volumes=override.pop('volumes', self.volumes),
             infra=override.pop('infra', None),
             _docker_login_config=override.pop('_docker_login_config',
                                               self._docker_login_config),
@@ -1523,6 +1668,12 @@ class Resources:
             features.add(clouds.CloudImplementationFeatures.IMAGE_ID)
         if self.ports is not None:
             features.add(clouds.CloudImplementationFeatures.OPEN_PORTS)
+        if self.volumes is not None:
+            for volume in self.volumes:
+                if 'disk_tier' in volume and volume[
+                        'disk_tier'] != resources_utils.DiskTier.BEST:
+                    features.add(
+                        clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER)
         return features
 
     @staticmethod
@@ -1692,6 +1843,7 @@ class Resources:
         resources_fields['ports'] = config.pop('ports', None)
         resources_fields['labels'] = config.pop('labels', None)
         resources_fields['autostop'] = config.pop('autostop', None)
+        resources_fields['volumes'] = config.pop('volumes', None)
         resources_fields['_docker_login_config'] = config.pop(
             '_docker_login_config', None)
         resources_fields['_docker_username_for_runpod'] = config.pop(
@@ -1742,6 +1894,21 @@ class Resources:
             config['disk_tier'] = self.disk_tier.value
         add_if_not_none('ports', self.ports)
         add_if_not_none('labels', self.labels)
+        if self.volumes is not None:
+            # Convert DiskTier/StorageType enum to string value for each volume
+            volumes = []
+            for volume in self.volumes:
+                volume_copy = volume.copy()
+                if 'disk_tier' in volume_copy:
+                    volume_copy['disk_tier'] = volume_copy['disk_tier'].value
+                if 'storage_type' in volume_copy:
+                    volume_copy['storage_type'] = volume_copy[
+                        'storage_type'].value
+                if 'attach_mode' in volume_copy:
+                    volume_copy['attach_mode'] = volume_copy[
+                        'attach_mode'].value
+                volumes.append(volume_copy)
+            config['volumes'] = volumes
         if self._autostop_config is not None:
             config['autostop'] = self._autostop_config.to_yaml_config()
         if self._docker_login_config is not None:
@@ -1902,6 +2069,9 @@ class Resources:
         if version < 23:
             self._autostop_config = None
 
+        if version < 24:
+            self._volumes = None
+
         self.__dict__.update(state)
 
 

sky/serve/server/core.py

@@ -141,8 +141,7 @@ def up(
     # Always apply the policy again here, even though it might have been applied
     # in the CLI. This is to ensure that we apply the policy to the final DAG
     # and get the mutated config.
-    dag, mutated_user_config = admin_policy_utils.apply(
-        task, use_mutated_config_in_current_request=False)
+    dag, mutated_user_config = admin_policy_utils.apply(task)
     task = dag.tasks[0]
 
     with rich_utils.safe_status(
@@ -352,8 +351,7 @@ def update(
     # and get the mutated config.
     # TODO(cblmemo,zhwu): If a user sets a new skypilot_config, the update
     # will not apply the config.
-    dag, _ = admin_policy_utils.apply(
-        task, use_mutated_config_in_current_request=False)
+    dag, _ = admin_policy_utils.apply(task)
     task = dag.tasks[0]
 
     assert task.service is not None

sky/server/common.py

@@ -12,7 +12,7 @@ import subprocess
 import sys
 import time
 import typing
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Literal, Optional
 from urllib import parse
 import uuid
 
@@ -116,6 +116,7 @@ class ApiServerStatus(enum.Enum):
     HEALTHY = 'healthy'
     UNHEALTHY = 'unhealthy'
     VERSION_MISMATCH = 'version_mismatch'
+    NEEDS_AUTH = 'needs_auth'
 
 
 @dataclasses.dataclass
@@ -127,15 +128,21 @@ class ApiServerInfo:
     commit: Optional[str] = None
 
 
+def get_api_cookie_jar_path() -> str:
+    return os.environ.get(server_constants.API_COOKIE_FILE_ENV_VAR,
+                          server_constants.API_COOKIE_FILE_DEFAULT_LOCATION)
+
+
 def get_api_cookie_jar() -> requests.cookies.RequestsCookieJar:
     """Returns the cookie jar used by the client to access the API server."""
-    cookie_file = os.environ.get(server_constants.API_COOKIE_FILE_ENV_VAR)
     cookie_jar = requests.cookies.RequestsCookieJar()
-    if cookie_file and os.path.exists(cookie_file):
+    cookie_file = get_api_cookie_jar_path()
+    if cookie_file:
         cookie_path = pathlib.Path(cookie_file).expanduser().resolve()
-        file_cookie_jar = MozillaCookieJar(cookie_path)
-        file_cookie_jar.load()
-        cookie_jar.update(file_cookie_jar)
+        if cookie_path.exists():
+            file_cookie_jar = MozillaCookieJar(cookie_path)
+            file_cookie_jar.load()
+            cookie_jar.update(file_cookie_jar)
     return cookie_jar
 
 
@@ -196,6 +203,7 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
             response = requests.get(f'{server_url}/api/health',
                                     timeout=2.5,
                                     cookies=get_api_cookie_jar())
+            logger.debug(f'Health check status: {response.status_code}')
             if response.status_code == 200:
                 try:
                     result = response.json()
@@ -217,9 +225,24 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
                         server_info.status = ApiServerStatus.VERSION_MISMATCH
                     return server_info
                 except (json.JSONDecodeError, AttributeError) as e:
+                    # Try to check if we got redirected to a login page.
+                    for prev_response in response.history:
+                        logger.debug(f'Previous response: {prev_response.url}')
+                        # Heuristic: check if the url looks like a login page or
+                        # oauth flow.
+                        if any(key in prev_response.url
+                               for key in ['login', 'oauth2']):
+                            logger.debug(
+                                f'URL {prev_response.url} looks like '
+                                'a login page or oauth flow, so try to '
+                                'get the cookie.')
+                            return ApiServerInfo(
+                                status=ApiServerStatus.NEEDS_AUTH)
                     logger.warning('Failed to parse API server response: '
                                    f'{str(e)}')
                     return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
+            elif response.status_code == 401:
+                return ApiServerInfo(status=ApiServerStatus.NEEDS_AUTH)
             else:
                 return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
         except requests.exceptions.Timeout:
@@ -369,7 +392,12 @@ def _start_api_server(deploy: bool = False,
                 f'SkyPilot API server started. {dashboard_msg}'))
 
 
-def check_server_healthy(endpoint: Optional[str] = None,) -> None:
+def check_server_healthy(
+    endpoint: Optional[str] = None
+) -> Literal[
+        # Use an incomplete list of Literals here to enforce raising for other
+        # enum values.
+        ApiServerStatus.HEALTHY, ApiServerStatus.NEEDS_AUTH]:
     """Check if the API server is healthy.
 
     Args:
@@ -379,6 +407,11 @@ def check_server_healthy(endpoint: Optional[str] = None,) -> None:
     Raises:
         RuntimeError: If the server is not healthy or the client version does
             not match the server version.
+
+    Returns:
+        ApiServerStatus: The status of the API server, unless the server is
+            unhealthy or the client version does not match the server version,
+            in which case an exception is raised.
     """
     endpoint = endpoint if endpoint is not None else get_server_url()
     api_server_info = get_api_server_status(endpoint)
@@ -441,6 +474,8 @@ def check_server_healthy(endpoint: Optional[str] = None,) -> None:
 
         hinted_for_server_install_version_mismatch = True
 
+    return api_server_status
+
 
 def _get_version_info_hint(server_info: ApiServerInfo) -> str:
     assert server_info.version is not None, 'Server version is None'
@@ -491,11 +526,13 @@ def get_skypilot_version_on_disk() -> str:
 def check_server_healthy_or_start_fn(deploy: bool = False,
                                      host: str = '127.0.0.1',
                                      foreground: bool = False):
+    api_server_status = None
     try:
-        check_server_healthy()
+        api_server_status = check_server_healthy()
     except exceptions.ApiServerConnectionError as exc:
         endpoint = get_server_url()
-        if not is_api_server_local():
+        if (not is_api_server_local() or
+                api_server_status == ApiServerStatus.NEEDS_AUTH):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ApiServerConnectionError(endpoint) from exc
         # Lock to prevent multiple processes from starting the server at the

sky/server/constants.py

@@ -26,6 +26,8 @@ CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS = 60
 
 # Environment variable for a file path to the API cookie file.
 API_COOKIE_FILE_ENV_VAR = f'{constants.SKYPILOT_ENV_VAR_PREFIX}API_COOKIE_FILE'
+# Default file if unset.
+API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
 
 # The path to the dashboard build output
 DASHBOARD_DIR = os.path.join(os.path.dirname(__file__), '..', 'dashboard',