skypilot-nightly 1.0.0.dev20250528__py3-none-any.whl → 1.0.0.dev20250530__py3-none-any.whl

sky/jobs/server/core.py

@@ -91,6 +91,7 @@ def launch(
     dag_utils.maybe_infer_and_fill_dag_and_task_names(dag)
 
     task_names = set()
+    priority = None
     for task_ in dag.tasks:
         if task_.name in task_names:
             with ux_utils.print_exception_no_traceback():
@@ -100,6 +101,20 @@ def launch(
                     'name only and comment out the task names (so that they '
                     'will be auto-generated) .')
         task_names.add(task_.name)
+        if task_.job_priority is not None:
+            if (priority is not None and priority != task_.job_priority):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'Multiple tasks in the DAG have different priorities. '
+                        'Either specify a priority in only one task, or set '
+                        'the same priority for each task.')
+            priority = task_.job_priority
+
+    if priority is None:
+        priority = managed_job_constants.DEFAULT_PRIORITY
+
+    if priority < 0 or priority > 1000:
+        raise ValueError(f'Priority must be between 0 and 1000, got {priority}')
 
     dag_utils.fill_default_config_in_dag_for_job_launch(dag)
 
@@ -186,6 +201,7 @@ def launch(
                 service_catalog_common.get_modified_catalog_file_mounts(),
             'dashboard_setup_cmd': managed_job_constants.DASHBOARD_SETUP_CMD,
             'dashboard_user_id': common.SERVER_ID,
+            'priority': priority,
             **controller_utils.shared_controller_vars_to_fill(
                 controller,
                 remote_user_config_path=remote_user_config_path,

sky/jobs/state.py

@@ -120,7 +120,9 @@ def create_table(cursor, conn):
         dag_yaml_path TEXT,
         env_file_path TEXT,
         user_hash TEXT,
-        workspace TEXT DEFAULT NULL)""")
+        workspace TEXT DEFAULT NULL,
+        priority INTEGER DEFAULT 500,
+        entrypoint TEXT DEFAULT NULL)""")
 
     db_utils.add_column_to_table(cursor, conn, 'job_info', 'schedule_state',
                                  'TEXT')
@@ -142,6 +144,15 @@ def create_table(cursor, conn):
                                  'workspace',
                                  'TEXT DEFAULT NULL',
                                  value_to_replace_existing_entries='default')
+
+    db_utils.add_column_to_table(cursor,
+                                 conn,
+                                 'job_info',
+                                 'priority',
+                                 'INTEGER',
+                                 value_to_replace_existing_entries=500)
+
+    db_utils.add_column_to_table(cursor, conn, 'job_info', 'entrypoint', 'TEXT')
     conn.commit()
 
 
@@ -199,6 +210,8 @@ columns = [
     'env_file_path',
     'user_hash',
     'workspace',
+    'priority',
+    'entrypoint',
 ]
 
 
@@ -215,7 +228,7 @@ class ManagedJobStatus(enum.Enum):
     reset to INIT or SETTING_UP multiple times (depending on the preemptions).
 
     However, a managed job only has one ManagedJobStatus on the jobs controller.
-        ManagedJobStatus = [PENDING, SUBMITTED, STARTING, RUNNING, ...]
+        ManagedJobStatus = [PENDING, STARTING, RUNNING, ...]
     Mapping from JobStatus to ManagedJobStatus:
         INIT            ->  STARTING/RECOVERING
         SETTING_UP      ->  RUNNING
@@ -235,10 +248,14 @@ class ManagedJobStatus(enum.Enum):
     # PENDING: Waiting for the jobs controller to have a slot to run the
     # controller process.
     PENDING = 'PENDING'
+    # SUBMITTED: This state used to be briefly set before immediately changing
+    # to STARTING. Its use was removed in #5682. We keep it for backwards
+    # compatibility, so we can still parse old jobs databases that may have jobs
+    # in this state.
+    # TODO(cooperc): remove this in v0.12.0
+    DEPRECATED_SUBMITTED = 'SUBMITTED'
     # The submitted_at timestamp of the managed job in the 'spot' table will be
     # set to the time when the job controller begins running.
-    # SUBMITTED: The jobs controller starts the controller process.
-    SUBMITTED = 'SUBMITTED'
     # STARTING: The controller process is launching the cluster for the managed
     # job.
     STARTING = 'STARTING'
@@ -314,7 +331,6 @@ class ManagedJobStatus(enum.Enum):
 
 _SPOT_STATUS_TO_COLOR = {
     ManagedJobStatus.PENDING: colorama.Fore.BLUE,
-    ManagedJobStatus.SUBMITTED: colorama.Fore.BLUE,
     ManagedJobStatus.STARTING: colorama.Fore.BLUE,
     ManagedJobStatus.RUNNING: colorama.Fore.GREEN,
     ManagedJobStatus.RECOVERING: colorama.Fore.CYAN,
@@ -326,6 +342,8 @@ _SPOT_STATUS_TO_COLOR = {
     ManagedJobStatus.FAILED_CONTROLLER: colorama.Fore.RED,
     ManagedJobStatus.CANCELLING: colorama.Fore.YELLOW,
     ManagedJobStatus.CANCELLED: colorama.Fore.YELLOW,
+    # TODO(cooperc): backwards compatibility, remove this in v0.12.0
+    ManagedJobStatus.DEPRECATED_SUBMITTED: colorama.Fore.BLUE,
 }
 
 
@@ -342,8 +360,12 @@ class ManagedJobScheduleState(enum.Enum):
     - LAUNCHING -> ALIVE: The launch attempt was completed. It may have
       succeeded or failed. The job controller is not allowed to sky.launch again
       without transitioning to ALIVE_WAITING and then LAUNCHING.
+    - LAUNCHING -> ALIVE_BACKOFF: The launch failed to find resources, and is
+      in backoff waiting for resources.
     - ALIVE -> ALIVE_WAITING: The job controller wants to sky.launch again,
       either for recovery or to launch a subsequent task.
+    - ALIVE_BACKOFF -> ALIVE_WAITING: The backoff period has ended, and the job
+      controller wants to try to launch again.
     - ALIVE_WAITING -> LAUNCHING: The scheduler has determined that the job
       controller may launch again.
     - LAUNCHING, ALIVE, or ALIVE_WAITING -> DONE: The job controller is exiting
@@ -357,6 +379,7 @@ class ManagedJobScheduleState(enum.Enum):
     state or vice versa. (In fact, schedule state is defined on the job and
     status on the task.)
     - INACTIVE or WAITING should only be seen when a job is PENDING.
+    - ALIVE_BACKOFF should only be seen when a job is STARTING.
     - ALIVE_WAITING should only be seen when a job is RECOVERING, has multiple
       tasks, or needs to retry launching.
     - LAUNCHING and ALIVE can be seen in many different statuses.
@@ -382,6 +405,9 @@ class ManagedJobScheduleState(enum.Enum):
     # The job is running sky.launch, or soon will, using a limited number of
     # allowed launch slots.
     LAUNCHING = 'LAUNCHING'
+    # The job is alive, but is in backoff waiting for resources - a special case
+    # of ALIVE.
+    ALIVE_BACKOFF = 'ALIVE_BACKOFF'
     # The controller for the job is running, but it's not currently launching.
     ALIVE = 'ALIVE'
     # The job is in a terminal state. (Not necessarily SUCCEEDED.)
@@ -389,14 +415,15 @@ class ManagedJobScheduleState(enum.Enum):
 
 
 # === Status transition functions ===
-def set_job_info(job_id: int, name: str, workspace: str):
+def set_job_info(job_id: int, name: str, workspace: str, entrypoint: str):
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\
             INSERT INTO job_info
-            (spot_job_id, name, schedule_state, workspace)
-            VALUES (?, ?, ?, ?)""",
-            (job_id, name, ManagedJobScheduleState.INACTIVE.value, workspace))
+            (spot_job_id, name, schedule_state, workspace, entrypoint)
+            VALUES (?, ?, ?, ?, ?)""",
+            (job_id, name, ManagedJobScheduleState.INACTIVE.value, workspace,
+             entrypoint))
 
 
 def set_pending(job_id: int, task_id: int, task_name: str, resources_str: str):
@@ -411,11 +438,11 @@ def set_pending(job_id: int, task_id: int, task_name: str, resources_str: str):
              ManagedJobStatus.PENDING.value))
 
 
-def set_submitted(job_id: int, task_id: int, run_timestamp: str,
-                  submit_time: float, resources_str: str,
-                  specs: Dict[str, Union[str,
-                                         int]], callback_func: CallbackType):
-    """Set the task to submitted.
+def set_starting(job_id: int, task_id: int, run_timestamp: str,
+                 submit_time: float, resources_str: str,
+                 specs: Dict[str, Union[str,
+                                        int]], callback_func: CallbackType):
+    """Set the task to starting state.
 
     Args:
         job_id: The managed job ID.
@@ -432,6 +459,7 @@ def set_submitted(job_id: int, task_id: int, run_timestamp: str,
     # make it easier to find them based on one of the values.
     # Also, using the earlier timestamp should be closer to the term
     # `submit_at`, which represents the time the managed task is submitted.
+    logger.info('Launching the spot cluster...')
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\
@@ -445,19 +473,54 @@ def set_submitted(job_id: int, task_id: int, run_timestamp: str,
             task_id=(?) AND
             status=(?) AND
             end_at IS null""",
-            (resources_str, submit_time, ManagedJobStatus.SUBMITTED.value,
+            (resources_str, submit_time, ManagedJobStatus.STARTING.value,
              run_timestamp, json.dumps(specs), job_id, task_id,
              ManagedJobStatus.PENDING.value))
         if cursor.rowcount != 1:
             raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to submitted. '
+                'Failed to set the task to starting. '
                 f'({cursor.rowcount} rows updated)')
+    # SUBMITTED is no longer used, but we keep it for backward compatibility.
+    # TODO(cooperc): remove this in v0.12.0
     callback_func('SUBMITTED')
+    callback_func('STARTING')
 
 
-def set_starting(job_id: int, task_id: int, callback_func: CallbackType):
-    """Set the task to starting state."""
-    logger.info('Launching the spot cluster...')
+def set_backoff_pending(job_id: int, task_id: int):
+    """Set the task to PENDING state if it is in backoff.
+
+    This should only be used to transition from STARTING or RECOVERING back to
+    PENDING.
+    """
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute(
+            """\
+            UPDATE spot SET status=(?)
+            WHERE spot_job_id=(?) AND
+            task_id=(?) AND
+            status IN (?, ?) AND
+            end_at IS null""", (ManagedJobStatus.PENDING.value, job_id, task_id,
+                                ManagedJobStatus.STARTING.value,
+                                ManagedJobStatus.RECOVERING.value))
+        logger.debug('back to PENDING')
+        if cursor.rowcount != 1:
+            raise exceptions.ManagedJobStatusError(
+                'Failed to set the task back to pending. '
+                f'({cursor.rowcount} rows updated)')
+    # Do not call callback_func here, as we don't use the callback for PENDING.
+
+
+def set_restarting(job_id: int, task_id: int, recovering: bool):
+    """Set the task back to STARTING or RECOVERING from PENDING.
+
+    This should not be used for the initial transition from PENDING to STARTING.
+    In that case, use set_starting instead. This function should only be used
+    after using set_backoff_pending to transition back to PENDING during
+    launch retry backoff.
+    """
+    target_status = ManagedJobStatus.STARTING.value
+    if recovering:
+        target_status = ManagedJobStatus.RECOVERING.value
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\
@@ -465,13 +528,15 @@ def set_starting(job_id: int, task_id: int, callback_func: CallbackType):
             WHERE spot_job_id=(?) AND
             task_id=(?) AND
             status=(?) AND
-            end_at IS null""", (ManagedJobStatus.STARTING.value, job_id,
-                                task_id, ManagedJobStatus.SUBMITTED.value))
+            end_at IS null""",
+            (target_status, job_id, task_id, ManagedJobStatus.PENDING.value))
+        logger.debug(f'back to {target_status}')
         if cursor.rowcount != 1:
             raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to starting. '
+                f'Failed to set the task back to {target_status}. '
                 f'({cursor.rowcount} rows updated)')
-    callback_func('STARTING')
+    # Do not call callback_func here, as it should only be invoked for the
+    # initial (pre-`set_backoff_pending`) transition to STARTING or RECOVERING.
 
 
 def set_started(job_id: int, task_id: int, start_time: float,
@@ -947,6 +1012,21 @@ def get_managed_jobs(job_id: Optional[int] = None) -> List[Dict[str, Any]]:
                 job_dict['schedule_state'])
             if job_dict['job_name'] is None:
                 job_dict['job_name'] = job_dict['task_name']
+
+            # Add YAML content and command for managed jobs
+            dag_yaml_path = job_dict.get('dag_yaml_path')
+            if dag_yaml_path:
+                try:
+                    with open(dag_yaml_path, 'r', encoding='utf-8') as f:
+                        job_dict['dag_yaml'] = f.read()
+                except (FileNotFoundError, IOError, OSError):
+                    job_dict['dag_yaml'] = None
+
+                # Generate a command that could be used to launch this job
+                # Format: sky jobs launch <yaml_path>
+            else:
+                job_dict['dag_yaml'] = None
+
             jobs.append(job_dict)
         return jobs
 
@@ -1004,16 +1084,16 @@ def get_local_log_file(job_id: int, task_id: Optional[int]) -> Optional[str]:
 
 
 def scheduler_set_waiting(job_id: int, dag_yaml_path: str, env_file_path: str,
-                          user_hash: str) -> None:
+                          user_hash: str, priority: int) -> None:
     """Do not call without holding the scheduler lock."""
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         updated_count = cursor.execute(
             'UPDATE job_info SET '
             'schedule_state = (?), dag_yaml_path = (?), env_file_path = (?), '
-            '  user_hash = (?) '
+            '  user_hash = (?), priority = (?) '
             'WHERE spot_job_id = (?) AND schedule_state = (?)',
             (ManagedJobScheduleState.WAITING.value, dag_yaml_path,
-             env_file_path, user_hash, job_id,
+             env_file_path, user_hash, priority, job_id,
              ManagedJobScheduleState.INACTIVE.value)).rowcount
         assert updated_count == 1, (job_id, updated_count)
 
@@ -1043,15 +1123,28 @@ def scheduler_set_alive(job_id: int) -> None:
         assert updated_count == 1, (job_id, updated_count)
 
 
-def scheduler_set_alive_waiting(job_id: int) -> None:
+def scheduler_set_alive_backoff(job_id: int) -> None:
     """Do not call without holding the scheduler lock."""
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         updated_count = cursor.execute(
             'UPDATE job_info SET '
             'schedule_state = (?) '
             'WHERE spot_job_id = (?) AND schedule_state = (?)',
+            (ManagedJobScheduleState.ALIVE_BACKOFF.value, job_id,
+             ManagedJobScheduleState.LAUNCHING.value)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def scheduler_set_alive_waiting(job_id: int) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state IN (?, ?)',
             (ManagedJobScheduleState.ALIVE_WAITING.value, job_id,
-             ManagedJobScheduleState.ALIVE.value)).rowcount
+             ManagedJobScheduleState.ALIVE.value,
+             ManagedJobScheduleState.ALIVE_BACKOFF.value)).rowcount
         assert updated_count == 1, (job_id, updated_count)
 
 
@@ -1099,32 +1192,53 @@ def get_num_alive_jobs() -> int:
         return cursor.execute(
             'SELECT COUNT(*) '
             'FROM job_info '
-            'WHERE schedule_state IN (?, ?, ?)',
+            'WHERE schedule_state IN (?, ?, ?, ?)',
             (ManagedJobScheduleState.ALIVE_WAITING.value,
              ManagedJobScheduleState.LAUNCHING.value,
-             ManagedJobScheduleState.ALIVE.value)).fetchone()[0]
+             ManagedJobScheduleState.ALIVE.value,
+             ManagedJobScheduleState.ALIVE_BACKOFF.value)).fetchone()[0]
 
 
 def get_waiting_job() -> Optional[Dict[str, Any]]:
     """Get the next job that should transition to LAUNCHING.
 
+    Selects the highest-priority (lowest numerical value) WAITING or
+    ALIVE_WAITING job, provided its priority value is less than or equal to any
+    currently LAUNCHING or ALIVE_BACKOFF job.
+
     Backwards compatibility note: jobs submitted before #4485 will have no
     schedule_state and will be ignored by this SQL query.
     """
     with db_utils.safe_cursor(_DB_PATH) as cursor:
-        row = cursor.execute(
+        # Get the highest-priority (lowest numerical value) WAITING or
+        # ALIVE_WAITING job whose priority value is less than or equal to
+        # the highest priority (numerically smallest) LAUNCHING or
+        # ALIVE_BACKOFF job's priority.
+        waiting_job_row = cursor.execute(
             'SELECT spot_job_id, schedule_state, dag_yaml_path, env_file_path '
             'FROM job_info '
-            'WHERE schedule_state in (?, ?) '
-            'ORDER BY spot_job_id LIMIT 1',
+            'WHERE schedule_state IN (?, ?) '
+            'AND priority <= COALESCE('
+            '    (SELECT MIN(priority) '
+            '     FROM job_info '
+            '     WHERE schedule_state IN (?, ?)), '
+            '    1000'
+            ')'
+            'ORDER BY priority ASC, spot_job_id ASC LIMIT 1',
             (ManagedJobScheduleState.WAITING.value,
-             ManagedJobScheduleState.ALIVE_WAITING.value)).fetchone()
+             ManagedJobScheduleState.ALIVE_WAITING.value,
+             ManagedJobScheduleState.LAUNCHING.value,
+             ManagedJobScheduleState.ALIVE_BACKOFF.value)).fetchone()
+
+        if waiting_job_row is None:
+            return None
+
         return {
-            'job_id': row[0],
-            'schedule_state': ManagedJobScheduleState(row[1]),
-            'dag_yaml_path': row[2],
-            'env_file_path': row[3],
-        } if row is not None else None
+            'job_id': waiting_job_row[0],
+            'schedule_state': ManagedJobScheduleState(waiting_job_row[1]),
+            'dag_yaml_path': waiting_job_row[2],
+            'env_file_path': waiting_job_row[3],
+        }
 
 
 def get_workspace(job_id: int) -> str:

sky/jobs/utils.py

@@ -953,6 +953,22 @@ def dump_managed_job_queue() -> str:
             job['region'] = '-'
             job['zone'] = '-'
 
+        # Add details about schedule state / backoff.
+        state_details = None
+        if job['schedule_state'] == 'ALIVE_BACKOFF':
+            state_details = 'In backoff, waiting for resources'
+        elif job['schedule_state'] in ('WAITING', 'ALIVE_WAITING'):
+            state_details = 'Waiting for other jobs to launch'
+
+        if state_details and job['failure_reason']:
+            job['details'] = f'{state_details} - {job["failure_reason"]}'
+        elif state_details:
+            job['details'] = state_details
+        elif job['failure_reason']:
+            job['details'] = f'Failure: {job["failure_reason"]}'
+        else:
+            job['details'] = None
+
     return message_utils.encode_payload(jobs)
 
 
@@ -981,7 +997,7 @@ def _get_job_status_from_tasks(
         # Use the first non-succeeded status.
         if managed_task_status != managed_job_state.ManagedJobStatus.SUCCEEDED:
             # TODO(zhwu): we should not blindly use the first non-
-            # succeeded as the status could be changed to SUBMITTED
+            # succeeded as the status could be changed to PENDING
             # when going from one task to the next one, which can be
             # confusing.
             break
@@ -1063,6 +1079,7 @@ def format_job_table(
         'TASK',
         *(['WORKSPACE'] if show_workspace else []),
         'NAME',
+        'PRIORITY',
         *user_cols,
         'REQUESTED',
         'SUBMITTED',
@@ -1092,7 +1109,10 @@ def format_job_table(
         # by the task_id.
         jobs[get_hash(task)].append(task)
 
-    def generate_details(failure_reason: Optional[str]) -> str:
+    def generate_details(details: Optional[str],
+                         failure_reason: Optional[str]) -> str:
+        if details is not None:
+            return details
         if failure_reason is not None:
             return f'Failure: {failure_reason}'
         return '-'
@@ -1131,6 +1151,7 @@ def format_job_table(
             submitted_at = None
             end_at: Optional[int] = 0
             recovery_cnt = 0
+            priority = job_tasks[0].get('priority', '-')
             managed_job_status, current_task_id = _get_job_status_from_tasks(
                 job_tasks)
             for task in job_tasks:
@@ -1166,6 +1187,7 @@ def format_job_table(
                 '',
                 *([''] if show_workspace else []),
                 job_name,
+                str(priority),
                 *user_values,
                 '-',
                 submitted,
@@ -1175,13 +1197,14 @@ def format_job_table(
                 status_str,
             ]
             if show_all:
+                details = job_tasks[current_task_id].get('details')
                 failure_reason = job_tasks[current_task_id]['failure_reason']
                 job_values.extend([
                     '-',
                     '-',
                     '-',
                     job_tasks[0]['schedule_state'],
-                    generate_details(failure_reason),
+                    generate_details(details, failure_reason),
                 ])
             if tasks_have_k8s_user:
                 job_values.insert(0, job_tasks[0].get('user', '-'))
@@ -1195,11 +1218,13 @@ def format_job_table(
             submitted = log_utils.readable_time_duration(task['submitted_at'])
             user_values = get_user_column_values(task)
             task_workspace = '-' if len(job_tasks) > 1 else workspace
+            priority = task.get('priority', '-')
             values = [
                 task['job_id'] if len(job_tasks) == 1 else ' \u21B3',
                 task['task_id'] if len(job_tasks) > 1 else '-',
                 *([task_workspace] if show_workspace else []),
                 task['task_name'],
+                str(priority),
                 *user_values,
                 task['resources'],
                 # SUBMITTED
@@ -1244,7 +1269,8 @@ def format_job_table(
                     infra.formatted_str(),
                     task['cluster_resources'],
                     schedule_state,
-                    generate_details(task['failure_reason']),
+                    generate_details(task.get('details'),
+                                     task['failure_reason']),
                 ])
             if tasks_have_k8s_user:
                 values.insert(0, task.get('user', '-'))
@@ -1362,13 +1388,15 @@ class ManagedJobCodeGen:
 
     @classmethod
     def set_pending(cls, job_id: int, managed_job_dag: 'dag_lib.Dag',
-                    workspace) -> str:
+                    workspace: str, entrypoint: str) -> str:
         dag_name = managed_job_dag.name
         # Add the managed job to queue table.
         code = textwrap.dedent(f"""\
             set_job_info_kwargs = {{'workspace': {workspace!r}}}
             if managed_job_version < 4:
                 set_job_info_kwargs = {{}}
+            if managed_job_version >= 5:
+                set_job_info_kwargs['entrypoint'] = {entrypoint!r}
             managed_job_state.set_job_info(
                 {job_id}, {dag_name!r}, **set_job_info_kwargs)
             """)

sky/provision/kubernetes/utils.py

@@ -15,6 +15,7 @@ from urllib.parse import urlparse
 import sky
 from sky import clouds
 from sky import exceptions
+from sky import global_user_state
 from sky import models
 from sky import sky_logging
 from sky import skypilot_config
@@ -2810,7 +2811,7 @@ def set_autodown_annotations(handle: 'backends.CloudVmRayResourceHandle',
     tags = {
         provision_constants.TAG_RAY_CLUSTER_NAME: handle.cluster_name_on_cloud,
     }
-    ray_config = common_utils.read_yaml(handle.cluster_yaml)
+    ray_config = global_user_state.get_cluster_yaml_dict(handle.cluster_yaml)
     provider_config = ray_config['provider']
     namespace = get_namespace_from_config(provider_config)
     context = get_context_from_config(provider_config)

sky/provision/provisioner.py

@@ -15,6 +15,7 @@ import colorama
 import sky
 from sky import clouds
 from sky import exceptions
+from sky import global_user_state
 from sky import provision
 from sky import sky_logging
 from sky import skypilot_config
@@ -118,7 +119,7 @@ def bulk_provision(
         Cloud specific exceptions: If the provisioning process failed, cloud-
             specific exceptions will be raised by the cloud APIs.
     """
-    original_config = common_utils.read_yaml(cluster_yaml)
+    original_config = global_user_state.get_cluster_yaml_dict(cluster_yaml)
     head_node_type = original_config['head_node_type']
     bootstrap_config = provision_common.ProvisionConfig(
         provider_config=original_config['provider'],
@@ -413,9 +414,11 @@ def wait_for_ssh(cluster_info: provision_common.ClusterInfo,
 
 def _post_provision_setup(
         cloud_name: str, cluster_name: resources_utils.ClusterName,
-        cluster_yaml: str, provision_record: provision_common.ProvisionRecord,
+        handle_cluster_yaml: str,
+        provision_record: provision_common.ProvisionRecord,
         custom_resource: Optional[str]) -> provision_common.ClusterInfo:
-    config_from_yaml = common_utils.read_yaml(cluster_yaml)
+    config_from_yaml = global_user_state.get_cluster_yaml_dict(
+        handle_cluster_yaml)
     provider_config = config_from_yaml.get('provider')
     cluster_info = provision.get_cluster_info(cloud_name,
                                               provision_record.region,
@@ -446,7 +449,7 @@ def _post_provision_setup(
     # TODO(suquark): Move wheel build here in future PRs.
     # We don't set docker_user here, as we are configuring the VM itself.
     ssh_credentials = backend_utils.ssh_credential_from_yaml(
-        cluster_yaml, ssh_user=cluster_info.ssh_user)
+        handle_cluster_yaml, ssh_user=cluster_info.ssh_user)
     docker_config = config_from_yaml.get('docker', {})
 
     with rich_utils.safe_status(
@@ -657,7 +660,8 @@ def _post_provision_setup(
 @timeline.event
 def post_provision_runtime_setup(
         cloud_name: str, cluster_name: resources_utils.ClusterName,
-        cluster_yaml: str, provision_record: provision_common.ProvisionRecord,
+        handle_cluster_yaml: str,
+        provision_record: provision_common.ProvisionRecord,
         custom_resource: Optional[str],
         log_dir: str) -> provision_common.ClusterInfo:
     """Run internal setup commands after provisioning and before user setup.
@@ -675,11 +679,12 @@ def post_provision_runtime_setup(
     with provision_logging.setup_provision_logging(log_dir):
         try:
             logger.debug(_TITLE.format('System Setup After Provision'))
-            return _post_provision_setup(cloud_name,
-                                         cluster_name,
-                                         cluster_yaml=cluster_yaml,
-                                         provision_record=provision_record,
-                                         custom_resource=custom_resource)
+            return _post_provision_setup(
+                cloud_name,
+                cluster_name,
+                handle_cluster_yaml=handle_cluster_yaml,
+                provision_record=provision_record,
+                custom_resource=custom_resource)
         except Exception:  # pylint: disable=broad-except
             logger.error(
                 ux_utils.error_message(

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 = 24
+    _VERSION = 25
 
     def __init__(
         self,
@@ -294,6 +294,8 @@ class Resources:
                 }
         else:
             self._image_id = image_id
+        if isinstance(self._cloud, clouds.Kubernetes):
+            _maybe_add_docker_prefix_to_image_id(self._image_id)
         self._is_image_managed = _is_image_managed
 
         if isinstance(disk_tier, str):
@@ -2075,6 +2077,10 @@ class Resources:
         if version < 24:
             self._volumes = None
 
+        if version < 25:
+            if isinstance(state.get('_cloud', None), clouds.Kubernetes):
+                _maybe_add_docker_prefix_to_image_id(state['_image_id'])
+
         self.__dict__.update(state)
 
 
@@ -2111,3 +2117,12 @@ class LaunchableResources(Resources):
         """
         self.assert_launchable()
         return typing.cast(LaunchableResources, super().copy(**override))
+
+
+def _maybe_add_docker_prefix_to_image_id(
+        image_id_dict: Optional[Dict[Optional[str], str]]) -> None:
+    if image_id_dict is None:
+        return
+    for k, v in image_id_dict.items():
+        if not v.startswith('docker:'):
+            image_id_dict[k] = f'docker:{v}'