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

sky/resources.py

@@ -1,6 +1,5 @@
 """Resources: compute requirements of Tasks."""
 import dataclasses
-import functools
 import textwrap
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
@@ -9,15 +8,17 @@ import colorama
 from sky import check as sky_check
 from sky import clouds
 from sky import exceptions
-from sky import jobs as managed_jobs
 from sky import sky_logging
 from sky import skypilot_config
 from sky.clouds import service_catalog
 from sky.provision import docker_utils
+from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.skylet import constants
 from sky.utils import accelerator_registry
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import log_utils
+from sky.utils import registry
 from sky.utils import resources_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
@@ -32,7 +33,7 @@ class Resources:
 
     This class is immutable once created (to ensure some validations are done
     whenever properties change). To update the property of an instance of
-    Resources, use `resources.copy(**new_properties)`.
+    Resources, use ``resources.copy(**new_properties)``.
 
     Used:
 
@@ -44,7 +45,7 @@ class Resources:
     """
     # If any fields changed, increment the version. For backward compatibility,
     # modify the __setstate__ method to handle the old version.
-    _VERSION = 18
+    _VERSION = 22
 
     def __init__(
         self,
@@ -55,7 +56,7 @@ class Resources:
         accelerators: Union[None, str, Dict[str, int]] = None,
         accelerator_args: Optional[Dict[str, str]] = None,
         use_spot: Optional[bool] = None,
-        job_recovery: Optional[str] = None,
+        job_recovery: Optional[Union[Dict[str, Union[str, int]], str]] = None,
         region: Optional[str] = None,
         zone: Optional[str] = None,
         image_id: Union[Dict[str, str], str, None] = None,
@@ -66,8 +67,10 @@ class Resources:
         # Internal use only.
         # pylint: disable=invalid-name
         _docker_login_config: Optional[docker_utils.DockerLoginConfig] = None,
+        _docker_username_for_runpod: Optional[str] = None,
         _is_image_managed: Optional[bool] = None,
         _requires_fuse: Optional[bool] = None,
+        _cluster_config_overrides: Optional[Dict[str, Any]] = None,
     ):
         """Initialize a Resources object.
 
@@ -110,6 +113,12 @@ class Resources:
             job to recover the cluster from preemption. Refer to
             `recovery_strategy module <https://github.com/skypilot-org/skypilot/blob/master/sky/jobs/recovery_strategy.py>`__ # pylint: disable=line-too-long
             for more details.
+            When a dict is provided, it can have the following fields:
+
+            - strategy: the recovery strategy to use.
+            - max_restarts_on_errors: the max number of restarts on user code
+              errors.
+
           region: the region to use.
           zone: the zone to use.
           image_id: the image ID to use. If a str, must be a string
@@ -140,6 +149,9 @@ class Resources:
           _docker_login_config: the docker configuration to use. This includes
             the docker username, password, and registry server. If None, skip
             docker login.
+          _docker_username_for_runpod: the login username for the docker
+            containers. This is used by RunPod to set the ssh user for the
+            docker containers.
           _requires_fuse: whether the task requires FUSE mounting support. This
             is used internally by certain cloud implementations to do additional
             setup for FUSE mounting. This flag also safeguards against using
@@ -152,18 +164,27 @@ class Resources:
         """
         self._version = self._VERSION
         self._cloud = cloud
-        self._region: Optional[str] = None
-        self._zone: Optional[str] = None
-        self._validate_and_set_region_zone(region, zone)
+        self._region: Optional[str] = region
+        self._zone: Optional[str] = zone
 
         self._instance_type = instance_type
 
         self._use_spot_specified = use_spot is not None
         self._use_spot = use_spot if use_spot is not None else False
-        self._job_recovery = None
+        self._job_recovery: Optional[Dict[str, Union[str, int]]] = None
         if job_recovery is not None:
-            if job_recovery.strip().lower() != 'none':
-                self._job_recovery = job_recovery.upper()
+            if isinstance(job_recovery, str):
+                job_recovery = {'strategy': job_recovery}
+            if 'strategy' not in job_recovery:
+                job_recovery['strategy'] = None
+
+            strategy_name = job_recovery['strategy']
+            if strategy_name == 'none':
+                self._job_recovery = None
+            else:
+                if strategy_name is not None:
+                    job_recovery['strategy'] = strategy_name.upper()
+                self._job_recovery = job_recovery
 
         if disk_size is not None:
             if round(disk_size) != disk_size:
@@ -174,8 +195,6 @@ class Resources:
         else:
             self._disk_size = _DEFAULT_DISK_SIZE_GB
 
-        # self._image_id is a dict of {region: image_id}.
-        # The key is None if the same image_id applies for all regions.
         self._image_id = image_id
         if isinstance(image_id, str):
             self._image_id = {self._region: image_id.strip()}
@@ -216,12 +235,25 @@ class Resources:
 
         self._docker_login_config = _docker_login_config
 
+        # TODO(andyl): This ctor param seems to be unused.
+        # We always use `Task.set_resources` and `Resources.copy` to set the
+        # `docker_username_for_runpod`. But to keep the consistency with
+        # `_docker_login_config`, we keep it here.
+        self._docker_username_for_runpod = _docker_username_for_runpod
+
         self._requires_fuse = _requires_fuse
 
+        self._cluster_config_overrides = _cluster_config_overrides
+        self._cached_repr = None
+
         self._set_cpus(cpus)
         self._set_memory(memory)
         self._set_accelerators(accelerators, accelerator_args)
 
+    def validate(self):
+        """Validate the resources and infer the missing fields if possible."""
+        self._try_canonicalize_accelerators()
+        self._try_validate_and_set_region_zone()
         self._try_validate_instance_type()
         self._try_validate_cpus_mem()
         self._try_validate_managed_job_attributes()
@@ -260,6 +292,8 @@ class Resources:
             >>> sky.Resources(disk_size=100)
             <Cloud>(disk_size=100)
         """
+        if self._cached_repr is not None:
+            return self._cached_repr
         accelerators = ''
         accelerator_args = ''
         if self.accelerators is not None:
@@ -319,7 +353,8 @@ class Resources:
         if self.cloud is not None:
             cloud_str = f'{self.cloud}'
 
-        return f'{cloud_str}({hardware_str})'
+        self._cached_repr = f'{cloud_str}({hardware_str})'
+        return self._cached_repr
 
     @property
     def repr_with_region_zone(self) -> str:
@@ -353,7 +388,7 @@ class Resources:
         return self._instance_type
 
     @property
-    @functools.lru_cache(maxsize=1)
+    @annotations.lru_cache(scope='global', maxsize=1)
     def cpus(self) -> Optional[str]:
         """Returns the number of vCPUs that each instance must have.
 
@@ -387,8 +422,8 @@ class Resources:
         return self._memory
 
     @property
-    @functools.lru_cache(maxsize=1)
-    def accelerators(self) -> Optional[Dict[str, int]]:
+    @annotations.lru_cache(scope='global', maxsize=1)
+    def accelerators(self) -> Optional[Dict[str, Union[int, float]]]:
         """Returns the accelerators field directly or by inferring.
 
         For example, Resources(AWS, 'p3.2xlarge') has its accelerators field
@@ -415,7 +450,7 @@ class Resources:
         return self._use_spot_specified
 
     @property
-    def job_recovery(self) -> Optional[str]:
+    def job_recovery(self) -> Optional[Dict[str, Union[str, int]]]:
         return self._job_recovery
 
     @property
@@ -448,10 +483,20 @@ class Resources:
             return False
         return self._requires_fuse
 
+    @property
+    def cluster_config_overrides(self) -> Dict[str, Any]:
+        if self._cluster_config_overrides is None:
+            return {}
+        return self._cluster_config_overrides
+
     @requires_fuse.setter
     def requires_fuse(self, value: Optional[bool]) -> None:
         self._requires_fuse = value
 
+    @property
+    def docker_username_for_runpod(self) -> Optional[str]:
+        return self._docker_username_for_runpod
+
     def _set_cpus(
         self,
         cpus: Union[None, int, float, str],
@@ -513,7 +558,7 @@ class Resources:
         if memory_gb <= 0:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'The "cpus" field should be positive. Found: {memory!r}')
+                    f'The "memory" field should be positive. Found: {memory!r}')
 
     def _set_accelerators(
         self,
@@ -546,36 +591,49 @@ class Resources:
                         with ux_utils.print_exception_no_traceback():
                             raise ValueError(parse_error) from None
 
-            # Canonicalize the accelerator names.
-            accelerators = {
-                accelerator_registry.canonicalize_accelerator_name(
-                    acc, self._cloud): acc_count
-                for acc, acc_count in accelerators.items()
-            }
-
             acc, _ = list(accelerators.items())[0]
             if 'tpu' in acc.lower():
                 if self.cloud is None:
-                    self._cloud = clouds.GCP()
-                assert self.cloud.is_same_cloud(
-                    clouds.GCP()), 'Cloud must be GCP.'
+                    if kubernetes_utils.is_tpu_on_gke(acc):
+                        self._cloud = clouds.Kubernetes()
+                    else:
+                        self._cloud = clouds.GCP()
+                assert (self.cloud.is_same_cloud(clouds.GCP()) or
+                        self.cloud.is_same_cloud(clouds.Kubernetes())), (
+                            'Cloud must be GCP or Kubernetes for TPU '
+                            'accelerators.')
+
                 if accelerator_args is None:
                     accelerator_args = {}
+
                 use_tpu_vm = accelerator_args.get('tpu_vm', True)
-                if self.instance_type is not None and use_tpu_vm:
-                    if self.instance_type != 'TPU-VM':
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                'Cannot specify instance type'
-                                f' (got "{self.instance_type}") for TPU VM.')
-                if 'runtime_version' not in accelerator_args:
-                    if use_tpu_vm:
-                        accelerator_args['runtime_version'] = 'tpu-vm-base'
-                    else:
-                        accelerator_args['runtime_version'] = '2.12.0'
-                    logger.info(
-                        'Missing runtime_version in accelerator_args, using'
-                        f' default ({accelerator_args["runtime_version"]})')
+                if (self.cloud.is_same_cloud(clouds.GCP()) and
+                        not kubernetes_utils.is_tpu_on_gke(acc)):
+                    if 'runtime_version' not in accelerator_args:
+
+                        def _get_default_runtime_version() -> str:
+                            if not use_tpu_vm:
+                                return '2.12.0'
+                            # TPU V5 requires a newer runtime version.
+                            if acc.startswith('tpu-v5'):
+                                return 'v2-alpha-tpuv5'
+                            # TPU V6e requires a newer runtime version.
+                            elif acc.startswith('tpu-v6e'):
+                                return 'v2-alpha-tpuv6e'
+                            return 'tpu-vm-base'
+
+                        accelerator_args['runtime_version'] = (
+                            _get_default_runtime_version())
+                        logger.info(
+                            'Missing runtime_version in accelerator_args, using'
+                            f' default ({accelerator_args["runtime_version"]})')
+
+                    if self.instance_type is not None and use_tpu_vm:
+                        if self.instance_type != 'TPU-VM':
+                            with ux_utils.print_exception_no_traceback():
+                                raise ValueError(
+                                    'Cannot specify instance type (got '
+                                    f'{self.instance_type!r}) for TPU VM.')
 
         self._accelerators = accelerators
         self._accelerator_args = accelerator_args
@@ -588,15 +646,30 @@ class Resources:
         assert self.is_launchable(), self
         return self.cloud.need_cleanup_after_preemption_or_failure(self)
 
-    def _validate_and_set_region_zone(self, region: Optional[str],
-                                      zone: Optional[str]) -> None:
+    def _try_canonicalize_accelerators(self) -> None:
+        """Try to canonicalize the accelerators attribute.
+
+        We don't canonicalize accelerators during creation of Resources object
+        because it may check Kubernetes accelerators online. It requires
+        Kubernetes credentias which may not be available locally when a remote
+        API server is used.
+        """
+        if self._accelerators is None:
+            return
+        self._accelerators = {
+            accelerator_registry.canonicalize_accelerator_name(
+                acc, self._cloud): acc_count
+            for acc, acc_count in self._accelerators.items()
+        }
+
+    def _try_validate_and_set_region_zone(self) -> None:
         """Try to validate and set the region and zone attribute.
 
         Raises:
             ValueError: if the attributes are invalid.
             exceptions.NoCloudAccessError: if no public cloud is enabled.
         """
-        if region is None and zone is None:
+        if self._region is None and self._zone is None:
             return
 
         if self._cloud is None:
@@ -608,13 +681,13 @@ class Resources:
             cloud_to_errors = {}
             for cloud in enabled_clouds:
                 try:
-                    cloud.validate_region_zone(region, zone)
+                    cloud.validate_region_zone(self._region, self._zone)
                 except ValueError as e:
                     cloud_to_errors[repr(cloud)] = e
                     continue
                 valid_clouds.append(cloud)
 
-            if len(valid_clouds) == 0:
+            if not valid_clouds:
                 if len(enabled_clouds) == 1:
                     cloud_str = f'for cloud {enabled_clouds[0]}'
                 else:
@@ -632,23 +705,24 @@ class Resources:
                             table.add_row([str(cloud), reason_str])
                         hint = table.get_string()
                     raise ValueError(
-                        f'Invalid (region {region!r}, zone {zone!r}) '
-                        f'{cloud_str}. Details:\n{hint}')
+                        f'Invalid (region {self._region!r}, zone '
+                        f'{self._zone!r}) {cloud_str}. Details:\n{hint}')
             elif len(valid_clouds) > 1:
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
-                        f'Cannot infer cloud from (region {region!r}, zone '
-                        f'{zone!r}). Multiple enabled clouds have region/zone '
-                        f'of the same names: {valid_clouds}. '
+                        f'Cannot infer cloud from (region {self._region!r}, '
+                        f'zone {self._zone!r}). Multiple enabled clouds '
+                        f'have region/zone of the same names: {valid_clouds}. '
                         f'To fix: explicitly specify `cloud`.')
             logger.debug(f'Cloud is not specified, using {valid_clouds[0]} '
-                         f'inferred from region {region!r} and zone {zone!r}')
+                         f'inferred from region {self._region!r} and zone '
+                         f'{self._zone!r}')
             self._cloud = valid_clouds[0]
 
         # Validate if region and zone exist in the catalog, and set the region
         # if zone is specified.
         self._region, self._zone = self._cloud.validate_region_zone(
-            region, zone)
+            self._region, self._zone)
 
     def get_valid_regions_for_launchable(self) -> List[clouds.Region]:
         """Returns a set of `Region`s that can provision this Resources.
@@ -726,7 +800,7 @@ class Resources:
             for cloud in enabled_clouds:
                 if cloud.instance_type_exists(self._instance_type):
                     valid_clouds.append(cloud)
-            if len(valid_clouds) == 0:
+            if not valid_clouds:
                 if len(enabled_clouds) == 1:
                     cloud_str = f'for cloud {enabled_clouds[0]}'
                 else:
@@ -797,14 +871,11 @@ class Resources:
         Raises:
             ValueError: if the attributes are invalid.
         """
-        if self._job_recovery is None:
+        if self._job_recovery is None or self._job_recovery['strategy'] is None:
             return
-        if self._job_recovery not in managed_jobs.RECOVERY_STRATEGIES:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    f'Spot recovery strategy {self._job_recovery} '
-                    'is not supported. The strategy should be among '
-                    f'{list(managed_jobs.RECOVERY_STRATEGIES.keys())}')
+        # Validate the job recovery strategy
+        registry.JOBS_RECOVERY_STRATEGY_REGISTRY.from_str(
+            self._job_recovery['strategy'])
 
     def extract_docker_image(self) -> Optional[str]:
         if self.image_id is None:
@@ -826,12 +897,6 @@ class Resources:
 
         if self.extract_docker_image() is not None:
             # TODO(tian): validate the docker image exists / of reasonable size
-            if self.accelerators is not None:
-                for acc in self.accelerators.keys():
-                    if acc.lower().startswith('tpu'):
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                'Docker image is not supported for TPU VM.')
             if self.cloud is not None:
                 self.cloud.check_features_are_supported(
                     self, {clouds.CloudImplementationFeatures.DOCKER_IMAGE})
@@ -920,12 +985,6 @@ class Resources:
         """
         if self.ports is None:
             return
-        if skypilot_config.get_nested(('aws', 'security_group_name'),
-                                      None) is not None:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    'Cannot specify ports when AWS security group name is '
-                    'specified.')
         if self.cloud is not None:
             self.cloud.check_features_are_supported(
                 self, {clouds.CloudImplementationFeatures.OPEN_PORTS})
@@ -956,21 +1015,23 @@ class Resources:
         """
         if not self._labels:
             return
-
-        if self.cloud is None:
-            # Because each cloud has its own label format, we cannot validate
-            # the labels without knowing the cloud.
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    'Cloud must be specified when labels are provided.')
-
-        # Check if the label key value pairs are valid.
+        if self.cloud is not None:
+            validated_clouds = [self.cloud]
+        else:
+            # If no specific cloud is set, validate label against ALL clouds.
+            # The label will be dropped if invalid for any one of the cloud
+            validated_clouds = sky_check.get_cached_enabled_clouds_or_refresh()
         invalid_table = log_utils.create_table(['Label', 'Reason'])
         for key, value in self._labels.items():
-            valid, err_msg = self.cloud.is_label_valid(key, value)
-            if not valid:
-                invalid_table.add_row([f'{key}: {value}', err_msg])
-        if len(invalid_table.rows) > 0:
+            for cloud in validated_clouds:
+                valid, err_msg = cloud.is_label_valid(key, value)
+                if not valid:
+                    invalid_table.add_row([
+                        f'{key}: {value}',
+                        f'Label rejected due to {cloud}: {err_msg}'
+                    ])
+                    break
+        if invalid_table.rows:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
                     'The following labels are invalid:'
@@ -1000,9 +1061,10 @@ class Resources:
     def get_spot_str(self) -> str:
         return '[Spot]' if self.use_spot else ''
 
-    def make_deploy_variables(self, cluster_name_on_cloud: str,
+    def make_deploy_variables(self, cluster_name: resources_utils.ClusterName,
                               region: clouds.Region,
                               zones: Optional[List[clouds.Zone]],
+                              num_nodes: int,
                               dryrun: bool) -> Dict[str, Optional[str]]:
         """Converts planned sky.Resources to resource variables.
 
@@ -1011,13 +1073,48 @@ class Resources:
         cloud.make_deploy_resources_variables() method, and the cloud-agnostic
         variables are generated by this method.
         """
-        cloud_specific_variables = self.cloud.make_deploy_resources_variables(
-            self, cluster_name_on_cloud, region, zones, dryrun)
+        # Initial setup commands
+        initial_setup_commands = []
+        if (skypilot_config.get_nested(
+            ('nvidia_gpus', 'disable_ecc'),
+                False,
+                override_configs=self.cluster_config_overrides) and
+                self.accelerators is not None):
+            initial_setup_commands = [constants.DISABLE_GPU_ECC_COMMAND]
+
         docker_image = self.extract_docker_image()
+
+        # Cloud specific variables
+        cloud_specific_variables = self.cloud.make_deploy_resources_variables(
+            self, cluster_name, region, zones, num_nodes, dryrun)
+
+        # TODO(andyl): Should we print some warnings if users' envs share
+        # same names with the cloud specific variables, but not enabled
+        # since it's not on the particular cloud?
+
+        # Docker run options
+        docker_run_options = skypilot_config.get_nested(
+            ('docker', 'run_options'),
+            default_value=[],
+            override_configs=self.cluster_config_overrides)
+        if isinstance(docker_run_options, str):
+            docker_run_options = [docker_run_options]
+        # Special accelerator runtime might require additional docker run
+        # options. e.g., for TPU, we need --privileged.
+        if 'docker_run_options' in cloud_specific_variables:
+            docker_run_options.extend(
+                cloud_specific_variables['docker_run_options'])
+        if docker_run_options and isinstance(self.cloud, clouds.Kubernetes):
+            logger.warning(
+                f'{colorama.Style.DIM}Docker run options are specified, '
+                'but ignored for Kubernetes: '
+                f'{" ".join(docker_run_options)}'
+                f'{colorama.Style.RESET_ALL}')
         return dict(
             cloud_specific_variables,
             **{
                 # Docker config
+                'docker_run_options': docker_run_options,
                 # Docker image. The image name used to pull the image, e.g.
                 # ubuntu:latest.
                 'docker_image': docker_image,
@@ -1027,7 +1124,9 @@ class Resources:
                     constants.DEFAULT_DOCKER_CONTAINER_NAME,
                 # Docker login config (if any). This helps pull the image from
                 # private registries.
-                'docker_login_config': self._docker_login_config
+                'docker_login_config': self._docker_login_config,
+                # Initial setup commands.
+                'initial_setup_commands': initial_setup_commands,
             })
 
     def get_reservations_available_resources(self) -> Dict[str, int]:
@@ -1169,17 +1268,17 @@ class Resources:
     def is_empty(self) -> bool:
         """Is this Resources an empty request (all fields None)?"""
         return all([
-            self.cloud is None,
+            self._cloud is None,
             self._instance_type is None,
             self._cpus is None,
-            self.memory is None,
-            self.accelerators is None,
-            self.accelerator_args is None,
+            self._memory is None,
+            self._accelerators is None,
+            self._accelerator_args is None,
             not self._use_spot_specified,
-            self.disk_size == _DEFAULT_DISK_SIZE_GB,
-            self.disk_tier is None,
+            self._disk_size == _DEFAULT_DISK_SIZE_GB,
+            self._disk_tier is None,
             self._image_id is None,
-            self.ports is None,
+            self._ports is None,
             self._docker_login_config is None,
         ])
 
@@ -1205,11 +1304,16 @@ class Resources:
             labels=override.pop('labels', self.labels),
             _docker_login_config=override.pop('_docker_login_config',
                                               self._docker_login_config),
+            _docker_username_for_runpod=override.pop(
+                '_docker_username_for_runpod',
+                self._docker_username_for_runpod),
             _is_image_managed=override.pop('_is_image_managed',
                                            self._is_image_managed),
             _requires_fuse=override.pop('_requires_fuse', self._requires_fuse),
+            _cluster_config_overrides=override.pop(
+                '_cluster_config_overrides', self._cluster_config_overrides),
         )
-        assert len(override) == 0
+        assert not override
         return resources
 
     def valid_on_region_zones(self, region: str, zones: List[str]) -> bool:
@@ -1337,7 +1441,7 @@ class Resources:
     def _from_yaml_config_single(cls, config: Dict[str, str]) -> 'Resources':
 
         resources_fields = {}
-        resources_fields['cloud'] = clouds.CLOUD_REGISTRY.from_str(
+        resources_fields['cloud'] = registry.CLOUD_REGISTRY.from_str(
             config.pop('cloud', None))
         resources_fields['instance_type'] = config.pop('instance_type', None)
         resources_fields['cpus'] = config.pop('cpus', None)
@@ -1364,9 +1468,13 @@ class Resources:
         resources_fields['labels'] = config.pop('labels', None)
         resources_fields['_docker_login_config'] = config.pop(
             '_docker_login_config', None)
+        resources_fields['_docker_username_for_runpod'] = config.pop(
+            '_docker_username_for_runpod', None)
         resources_fields['_is_image_managed'] = config.pop(
             '_is_image_managed', None)
         resources_fields['_requires_fuse'] = config.pop('_requires_fuse', None)
+        resources_fields['_cluster_config_overrides'] = config.pop(
+            '_cluster_config_overrides', None)
 
         if resources_fields['cpus'] is not None:
             resources_fields['cpus'] = str(resources_fields['cpus'])
@@ -1393,7 +1501,7 @@ class Resources:
         add_if_not_none('instance_type', self.instance_type)
         add_if_not_none('cpus', self._cpus)
         add_if_not_none('memory', self.memory)
-        add_if_not_none('accelerators', self.accelerators)
+        add_if_not_none('accelerators', self._accelerators)
         add_if_not_none('accelerator_args', self.accelerator_args)
 
         if self._use_spot_specified:
@@ -1410,6 +1518,11 @@ class Resources:
         if self._docker_login_config is not None:
             config['_docker_login_config'] = dataclasses.asdict(
                 self._docker_login_config)
+        if self._docker_username_for_runpod is not None:
+            config['_docker_username_for_runpod'] = (
+                self._docker_username_for_runpod)
+        add_if_not_none('_cluster_config_overrides',
+                        self._cluster_config_overrides)
         if self._is_image_managed is not None:
             config['_is_image_managed'] = self._is_image_managed
         if self._requires_fuse is not None:
@@ -1525,4 +1638,36 @@ class Resources:
         if version < 18:
             self._job_recovery = state.pop('_spot_recovery', None)
 
+        if version < 19:
+            self._cluster_config_overrides = state.pop(
+                '_cluster_config_overrides', None)
+
+        if version < 20:
+            # Pre-0.7.0, we used 'kubernetes' as the default region for
+            # Kubernetes clusters. With the introduction of support for
+            # multiple contexts, we now set the region to the context name.
+            # Since we do not have information on which context the cluster
+            # was run in, we default it to the current active context.
+            legacy_region = clouds.Kubernetes().LEGACY_SINGLETON_REGION
+            original_cloud = state.get('_cloud', None)
+            original_region = state.get('_region', None)
+            if (isinstance(original_cloud, clouds.Kubernetes) and
+                    original_region == legacy_region):
+                current_context = (
+                    kubernetes_utils.get_current_kube_config_context_name())
+                state['_region'] = current_context
+                # Also update the image_id dict if it contains the old region
+                if isinstance(state['_image_id'], dict):
+                    if legacy_region in state['_image_id']:
+                        state['_image_id'][current_context] = (
+                            state['_image_id'][legacy_region])
+                        del state['_image_id'][legacy_region]
+
+        if version < 21:
+            self._cached_repr = None
+
+        if version < 22:
+            self._docker_username_for_runpod = state.pop(
+                '_docker_username_for_runpod', None)
+
         self.__dict__.update(state)