skypilot-nightly 1.0.0.dev20250912__py3-none-any.whl → 1.0.0.dev20250914__py3-none-any.whl

sky/skylet/job_lib.py

@@ -31,8 +31,11 @@ from sky.utils.db import db_utils
 
 if typing.TYPE_CHECKING:
     import psutil
+
+    from sky.schemas.generated import jobsv1_pb2
 else:
     psutil = adaptors_common.LazyImport('psutil')
+    jobsv1_pb2 = adaptors_common.LazyImport('sky.schemas.generated.jobsv1_pb2')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -220,6 +223,45 @@ class JobStatus(enum.Enum):
         color = _JOB_STATUS_TO_COLOR[self]
         return f'{color}{self.value}{colorama.Style.RESET_ALL}'
 
+    @classmethod
+    def from_protobuf(
+            cls,
+            protobuf_value: 'jobsv1_pb2.JobStatus') -> Optional['JobStatus']:
+        """Convert protobuf JobStatus enum to Python enum value."""
+        protobuf_to_enum = {
+            jobsv1_pb2.JOB_STATUS_INIT: cls.INIT,
+            jobsv1_pb2.JOB_STATUS_PENDING: cls.PENDING,
+            jobsv1_pb2.JOB_STATUS_SETTING_UP: cls.SETTING_UP,
+            jobsv1_pb2.JOB_STATUS_RUNNING: cls.RUNNING,
+            jobsv1_pb2.JOB_STATUS_FAILED_DRIVER: cls.FAILED_DRIVER,
+            jobsv1_pb2.JOB_STATUS_SUCCEEDED: cls.SUCCEEDED,
+            jobsv1_pb2.JOB_STATUS_FAILED: cls.FAILED,
+            jobsv1_pb2.JOB_STATUS_FAILED_SETUP: cls.FAILED_SETUP,
+            jobsv1_pb2.JOB_STATUS_CANCELLED: cls.CANCELLED,
+            jobsv1_pb2.JOB_STATUS_UNSPECIFIED: None,
+        }
+        if protobuf_value not in protobuf_to_enum:
+            raise ValueError(
+                f'Unknown protobuf JobStatus value: {protobuf_value}')
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'jobsv1_pb2.JobStatus':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            JobStatus.INIT: jobsv1_pb2.JOB_STATUS_INIT,
+            JobStatus.PENDING: jobsv1_pb2.JOB_STATUS_PENDING,
+            JobStatus.SETTING_UP: jobsv1_pb2.JOB_STATUS_SETTING_UP,
+            JobStatus.RUNNING: jobsv1_pb2.JOB_STATUS_RUNNING,
+            JobStatus.FAILED_DRIVER: jobsv1_pb2.JOB_STATUS_FAILED_DRIVER,
+            JobStatus.SUCCEEDED: jobsv1_pb2.JOB_STATUS_SUCCEEDED,
+            JobStatus.FAILED: jobsv1_pb2.JOB_STATUS_FAILED,
+            JobStatus.FAILED_SETUP: jobsv1_pb2.JOB_STATUS_FAILED_SETUP,
+            JobStatus.CANCELLED: jobsv1_pb2.JOB_STATUS_CANCELLED,
+        }
+        if self not in enum_to_protobuf:
+            raise ValueError(f'Unknown JobStatus value: {self}')
+        return enum_to_protobuf[self]
+
 
 # We have two steps for job submissions:
 # 1. Client reserve a job id from the job table by adding a INIT state job.
@@ -475,6 +517,11 @@ def get_status(job_id: int) -> Optional[JobStatus]:
 
 @init_db
 def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
+    return message_utils.encode_payload(get_statuses(job_ids))
+
+
+@init_db
+def get_statuses(job_ids: List[int]) -> Dict[int, Optional[str]]:
     assert _DB is not None
     # Per-job lock is not required here, since the staled job status will not
     # affect the caller.
@@ -482,10 +529,52 @@ def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
     rows = _DB.cursor.execute(
         f'SELECT job_id, status FROM jobs WHERE job_id IN ({query_str})',
         job_ids)
-    statuses = {job_id: None for job_id in job_ids}
+    statuses: Dict[int, Optional[str]] = {job_id: None for job_id in job_ids}
     for (job_id, status) in rows:
         statuses[job_id] = status
-    return message_utils.encode_payload(statuses)
+    return statuses
+
+
+@init_db
+def get_jobs_info(user_hash: Optional[str] = None,
+                  all_jobs: bool = False) -> List['jobsv1_pb2.JobInfo']:
+    """Get detailed job information.
+
+    Similar to dump_job_queue but returns structured protobuf objects instead
+    of encoded strings.
+
+    Args:
+        user_hash: The user hash to show jobs for. Show all the users if None.
+        all_jobs: Whether to show all jobs, not just the pending/running ones.
+    """
+    assert _DB is not None
+
+    status_list: Optional[List[JobStatus]] = [
+        JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING
+    ]
+    if all_jobs:
+        status_list = None
+
+    jobs = _get_jobs(user_hash, status_list=status_list)
+    jobs_info = []
+    for job in jobs:
+        jobs_info.append(
+            jobsv1_pb2.JobInfo(
+                job_id=job['job_id'],
+                job_name=job['job_name'],
+                username=job['username'],
+                submitted_at=job['submitted_at'],
+                status=job['status'].to_protobuf(),
+                run_timestamp=job['run_timestamp'],
+                start_at=job['start_at']
+                if job['start_at'] is not None else -1.0,
+                end_at=job['end_at'] if job['end_at'] is not None else 0.0,
+                resources=job['resources'] or '',
+                pid=job['pid'],
+                log_path=os.path.join(constants.SKY_LOGS_DIRECTORY,
+                                      job['run_timestamp']),
+                metadata=json.dumps(job['metadata'])))
+    return jobs_info
 
 
 def load_statuses_payload(
@@ -527,13 +616,24 @@ def get_job_submitted_or_ended_timestamp_payload(job_id: int,
     `format_job_queue()`), because the job may stay in PENDING if the cluster is
     busy.
     """
+    return message_utils.encode_payload(
+        get_job_submitted_or_ended_timestamp(job_id, get_ended_time))
+
+
+@init_db
+def get_job_submitted_or_ended_timestamp(
+        job_id: int, get_ended_time: bool) -> Optional[float]:
+    """Get the job submitted timestamp.
+
+    Returns the raw timestamp or None if job doesn't exist.
+    """
     assert _DB is not None
     field = 'end_at' if get_ended_time else 'submitted_at'
     rows = _DB.cursor.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
                               (job_id,))
     for (timestamp,) in rows:
-        return message_utils.encode_payload(timestamp)
-    return message_utils.encode_payload(None)
+        return timestamp
+    return None
 
 
 def get_ray_port():
@@ -947,6 +1047,13 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
         Encoded job IDs that are actually cancelled. Caller should use
         message_utils.decode_payload() to parse.
     """
+    return message_utils.encode_payload(cancel_jobs(jobs, cancel_all,
+                                                    user_hash))
+
+
+def cancel_jobs(jobs: Optional[List[int]],
+                cancel_all: bool = False,
+                user_hash: Optional[str] = None) -> List[int]:
     job_records = []
     all_status = [JobStatus.PENDING, JobStatus.SETTING_UP, JobStatus.RUNNING]
     if jobs is None and not cancel_all:
@@ -1010,7 +1117,7 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
                 cancelled_ids.append(job['job_id'])
 
         scheduler.schedule_step()
-    return message_utils.encode_payload(cancelled_ids)
+    return cancelled_ids
 
 
 @init_db
@@ -1030,6 +1137,17 @@ def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
 
 @init_db
 def get_log_dir_for_jobs(job_ids: List[Optional[str]]) -> str:
+    """Returns the relative paths to the log files for jobs with globbing,
+    encoded."""
+    job_to_dir = get_job_log_dirs(job_ids)
+    job_to_dir_str: Dict[str, str] = {}
+    for job_id, log_dir in job_to_dir.items():
+        job_to_dir_str[str(job_id)] = log_dir
+    return message_utils.encode_payload(job_to_dir_str)
+
+
+@init_db
+def get_job_log_dirs(job_ids: List[int]) -> Dict[int, str]:
     """Returns the relative paths to the log files for jobs with globbing."""
     assert _DB is not None
     query_str = ' OR '.join(['job_id GLOB (?)'] * len(job_ids))
@@ -1038,16 +1156,16 @@ def get_log_dir_for_jobs(job_ids: List[Optional[str]]) -> str:
             SELECT * FROM jobs
             WHERE {query_str}""", job_ids)
     rows = _DB.cursor.fetchall()
-    job_to_dir = {}
+    job_to_dir: Dict[int, str] = {}
     for row in rows:
         job_id = row[JobInfoLoc.JOB_ID.value]
         if row[JobInfoLoc.LOG_PATH.value]:
-            job_to_dir[str(job_id)] = row[JobInfoLoc.LOG_PATH.value]
+            job_to_dir[job_id] = row[JobInfoLoc.LOG_PATH.value]
         else:
             run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
-            job_to_dir[str(job_id)] = os.path.join(constants.SKY_LOGS_DIRECTORY,
-                                                   run_timestamp)
-    return message_utils.encode_payload(job_to_dir)
+            job_to_dir[job_id] = os.path.join(constants.SKY_LOGS_DIRECTORY,
+                                              run_timestamp)
+    return job_to_dir
 
 
 class JobLibCodeGen:

sky/skylet/log_lib.py

@@ -406,9 +406,9 @@ def _follow_job_logs(file,
                     wait_last_logs = False
                     continue
                 status_str = status.value if status is not None else 'None'
-                print(ux_utils.finishing_message(
-                    f'Job finished (status: {status_str}).'),
-                      flush=True)
+                finish = ux_utils.finishing_message(
+                    f'Job finished (status: {status_str}).')
+                yield finish + '\n'
                 return
 
             time.sleep(SKY_LOG_TAILING_GAP_SECONDS)

sky/skylet/services.py

@@ -1,11 +1,19 @@
 """gRPC service implementations for skylet."""
 
+import os
+
 import grpc
 
 from sky import sky_logging
+from sky.jobs import state as managed_job_state
 from sky.schemas.generated import autostopv1_pb2
 from sky.schemas.generated import autostopv1_pb2_grpc
+from sky.schemas.generated import jobsv1_pb2
+from sky.schemas.generated import jobsv1_pb2_grpc
+from sky.serve import serve_state
 from sky.skylet import autostop_lib
+from sky.skylet import constants
+from sky.skylet import job_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -42,3 +50,198 @@ class AutostopServiceImpl(autostopv1_pb2_grpc.AutostopServiceServicer):
                 is_autostopping=is_autostopping)
         except Exception as e:  # pylint: disable=broad-except
             context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+
+class JobsServiceImpl(jobsv1_pb2_grpc.JobsServiceServicer):
+    """Implementation of the JobsService gRPC service."""
+
+    def AddJob(  # type: ignore[return]
+            self, request: jobsv1_pb2.AddJobRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.AddJobResponse:
+        try:
+            job_name = request.job_name if request.HasField('job_name') else '-'
+            job_id, log_dir = job_lib.add_job(job_name, request.username,
+                                              request.run_timestamp,
+                                              request.resources_str,
+                                              request.metadata)
+            return jobsv1_pb2.AddJobResponse(job_id=job_id, log_dir=log_dir)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def QueueJob(  # type: ignore[return]
+            self, request: jobsv1_pb2.QueueJobRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.QueueJobResponse:
+        try:
+            job_id = request.job_id
+            # Create log directory and file
+            remote_log_dir = os.path.expanduser(request.remote_log_dir)
+            os.makedirs(remote_log_dir, exist_ok=True)
+            remote_log_path = os.path.join(remote_log_dir, 'run.log')
+            open(remote_log_path, 'a').close()  # pylint: disable=unspecified-encoding
+
+            script_path = os.path.expanduser(request.script_path)
+            os.makedirs(os.path.dirname(script_path), exist_ok=True)
+
+            # If `codegen` is not provided, assume script is already
+            # uploaded to `script_path` via rsync.
+            if request.HasField('codegen'):
+                with open(script_path, 'w', encoding='utf-8') as f:
+                    f.write(request.codegen)
+                os.chmod(script_path, 0o755)
+
+            cd = f'cd {constants.SKY_REMOTE_WORKDIR}'
+            job_submit_cmd = (
+                # JOB_CMD_IDENTIFIER is used for identifying the process
+                # retrieved with pid is the same driver process.
+                f'{job_lib.JOB_CMD_IDENTIFIER.format(job_id)} && '
+                f'{cd} && {constants.SKY_PYTHON_CMD} -u {script_path}'
+                # Do not use &>, which is not POSIX and may not work.
+                # Note that the order of ">filename 2>&1" matters.
+                f' > {remote_log_path} 2>&1')
+            job_lib.scheduler.queue(job_id, job_submit_cmd)
+
+            if request.HasField('managed_job'):
+                managed_job = request.managed_job
+                pool = managed_job.pool if managed_job.HasField(
+                    'pool') else None
+                pool_hash = None
+                if pool is not None:
+                    pool_hash = serve_state.get_service_hash(pool)
+                # Add the managed job to job queue database.
+                managed_job_state.set_job_info(job_id, managed_job.name,
+                                               managed_job.workspace,
+                                               managed_job.entrypoint, pool,
+                                               pool_hash)
+                # Set the managed job to PENDING state to make sure that
+                # this managed job appears in the `sky jobs queue`, even
+                # if it needs to wait to be submitted.
+                # We cannot set the managed job to PENDING state in the
+                # job template (jobs-controller.yaml.j2), as it may need
+                # to wait for the run commands to be scheduled on the job
+                # controller in high-load cases.
+                for task in managed_job.tasks:
+                    managed_job_state.set_pending(job_id, task.task_id,
+                                                  task.name, task.resources_str,
+                                                  task.metadata_json)
+            return jobsv1_pb2.QueueJobResponse()
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def UpdateStatus(  # type: ignore[return]
+            self, request: jobsv1_pb2.UpdateStatusRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.UpdateStatusResponse:
+        try:
+            job_lib.update_status()
+            return jobsv1_pb2.UpdateStatusResponse()
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def GetJobQueue(  # type: ignore[return]
+            self, request: jobsv1_pb2.GetJobQueueRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.GetJobQueueResponse:
+        try:
+            user_hash = request.user_hash if request.HasField(
+                'user_hash') else None
+            all_jobs = request.all_jobs
+            jobs_info = job_lib.get_jobs_info(user_hash=user_hash,
+                                              all_jobs=all_jobs)
+            return jobsv1_pb2.GetJobQueueResponse(jobs=jobs_info)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def CancelJobs(  # type: ignore[return]
+            self, request: jobsv1_pb2.CancelJobsRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.CancelJobsResponse:
+        try:
+            job_ids = list(request.job_ids) if request.job_ids else []
+            user_hash = request.user_hash if request.HasField(
+                'user_hash') else None
+            cancelled_job_ids = job_lib.cancel_jobs(job_ids, request.cancel_all,
+                                                    user_hash)
+            return jobsv1_pb2.CancelJobsResponse(
+                cancelled_job_ids=cancelled_job_ids)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def FailAllInProgressJobs(  # type: ignore[return]
+        self, _: jobsv1_pb2.FailAllInProgressJobsRequest,
+        context: grpc.ServicerContext
+    ) -> jobsv1_pb2.FailAllInProgressJobsResponse:
+        try:
+            job_lib.fail_all_jobs_in_progress()
+            return jobsv1_pb2.FailAllInProgressJobsResponse()
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def TailLogs(
+            self,
+            request: jobsv1_pb2.TailLogsRequest,  # type: ignore[return]
+            context: grpc.ServicerContext):
+        # TODO(kevin): implement this
+        raise NotImplementedError('TailLogs is not implemented')
+
+    def GetJobStatus(  # type: ignore[return]
+            self, request: jobsv1_pb2.GetJobStatusRequest,
+            context: grpc.ServicerContext) -> jobsv1_pb2.GetJobStatusResponse:
+        try:
+            if request.job_ids:
+                job_ids = list(request.job_ids)
+            else:
+                latest_job_id = job_lib.get_latest_job_id()
+                job_ids = [latest_job_id] if latest_job_id is not None else []
+            job_statuses = job_lib.get_statuses(job_ids)
+            for job_id, status in job_statuses.items():
+                job_statuses[job_id] = job_lib.JobStatus(status).to_protobuf(
+                ) if status is not None else jobsv1_pb2.JOB_STATUS_UNSPECIFIED
+            return jobsv1_pb2.GetJobStatusResponse(job_statuses=job_statuses)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def GetJobSubmittedTimestamp(  # type: ignore[return]
+        self, request: jobsv1_pb2.GetJobSubmittedTimestampRequest,
+        context: grpc.ServicerContext
+    ) -> jobsv1_pb2.GetJobSubmittedTimestampResponse:
+        try:
+            job_id = request.job_id if request.HasField(
+                'job_id') else job_lib.get_latest_job_id()
+            timestamp = job_lib.get_job_submitted_or_ended_timestamp(
+                job_id, False)
+            if timestamp is None:
+                context.abort(grpc.StatusCode.NOT_FOUND,
+                              f'Job {job_id} not found')
+            return jobsv1_pb2.GetJobSubmittedTimestampResponse(
+                timestamp=timestamp)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def GetJobEndedTimestamp(  # type: ignore[return]
+        self, request: jobsv1_pb2.GetJobEndedTimestampRequest,
+        context: grpc.ServicerContext
+    ) -> jobsv1_pb2.GetJobEndedTimestampResponse:
+        try:
+            job_id = request.job_id if request.HasField(
+                'job_id') else job_lib.get_latest_job_id()
+            timestamp = job_lib.get_job_submitted_or_ended_timestamp(
+                job_id, True)
+            if timestamp is None:
+                context.abort(grpc.StatusCode.NOT_FOUND,
+                              f'Job {job_id} not found or not ended')
+            return jobsv1_pb2.GetJobEndedTimestampResponse(timestamp=timestamp)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def GetLogDirsForJobs(  # type: ignore[return]
+            self, request: jobsv1_pb2.GetLogDirsForJobsRequest,
+            context: grpc.ServicerContext
+    ) -> jobsv1_pb2.GetLogDirsForJobsResponse:
+        try:
+            if request.job_ids:
+                job_ids = list(request.job_ids)
+            else:
+                latest_job_id = job_lib.get_latest_job_id()
+                job_ids = [latest_job_id] if latest_job_id is not None else []
+            job_log_dirs = job_lib.get_job_log_dirs(job_ids)
+            return jobsv1_pb2.GetLogDirsForJobsResponse(
+                job_log_dirs=job_log_dirs)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))

sky/skylet/skylet.py

@@ -9,6 +9,7 @@ import grpc
 import sky
 from sky import sky_logging
 from sky.schemas.generated import autostopv1_pb2_grpc
+from sky.schemas.generated import jobsv1_pb2_grpc
 from sky.skylet import constants
 from sky.skylet import events
 from sky.skylet import services
@@ -50,6 +51,9 @@ def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
     autostopv1_pb2_grpc.add_AutostopServiceServicer_to_server(
         services.AutostopServiceImpl(), server)
 
+    jobsv1_pb2_grpc.add_JobsServiceServicer_to_server(
+        services.JobsServiceImpl(), server)
+
     listen_addr = f'127.0.0.1:{port}'
     server.add_insecure_port(listen_addr)
 

sky/templates/seeweb-ray.yml.j2

@@ -0,0 +1,108 @@
+cluster_name: {{ cluster_name_on_cloud }}
+
+max_workers: {{ num_nodes - 1 }}
+upscaling_speed: {{ num_nodes - 1 }}
+idle_timeout_minutes: 5
+
+provider:
+  type: external
+  module: sky.provision.seeweb
+  region: "{{ region }}"
+
+auth:
+  ssh_user: ecuser
+  ssh_private_key: {{ ssh_private_key }}
+
+available_node_types:
+  ray_head_default:
+    resources: {}
+    node_config:
+      plan: {{ instance_type }}
+      image: {{ image_id }}
+      location: {{ region }}
+      {% if seeweb_gpu_config is not none %}
+      gpu: {{ seeweb_gpu_config.gpu }}
+      gpu_label: "{{ seeweb_gpu_config.gpu_label }}"
+      {% endif %}
+      disk: {{ disk_size }}
+
+head_node_type: ray_head_default
+
+file_mounts: {
+  "~/.seeweb_cloud/seeweb_keys": "~/.seeweb_cloud/seeweb_keys",
+  "{{sky_ray_yaml_remote_path}}": "{{sky_ray_yaml_local_path}}",
+  "{{sky_remote_path}}/{{sky_wheel_hash}}": "{{sky_local_path}}",
+{%- for remote_path, local_path in credentials.items() %}
+  "{{remote_path}}": "{{local_path}}",
+{%- endfor %}
+  "~/.ssh/sky-cluster-key": "{{ssh_private_key}}",
+}
+
+rsync_exclude: []
+
+setup_commands:
+  - |
+    touch ~/.bashrc;
+    echo "127.0.0.1 $(hostname)" | sudo tee -a /etc/hosts || true;
+    echo "127.0.0.1 localhost" | sudo tee -a /etc/hosts || true;
+    sudo systemctl stop unattended-upgrades || true;
+    sudo systemctl disable unattended-upgrades || true;
+    sudo apt update && sudo apt install -y patch || sudo yum install -y patch || true;
+    {{ conda_installation_commands }}
+    {{ ray_skypilot_installation_commands }}
+
+head_start_ray_commands:
+  - |
+    retry_ray() {
+      local n=0; local max=30
+      until [ $n -ge $max ]; do
+        export SKYPILOT_NUM_GPUS=0
+        command -v nvidia-smi >/dev/null 2>&1 && \
+          SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
+
+        ray stop || true
+        RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 \
+        ray start --disable-usage-stats --head \
+          --port={{ ray_port }} --dashboard-port={{ ray_dashboard_port }} \
+          --object-manager-port=8076 \
+          --autoscaling-config=~/ray_bootstrap_config.yaml \
+          --num-gpus=$SKYPILOT_NUM_GPUS --temp-dir {{ ray_temp_dir }} && break
+
+        echo "[head] Ray failed to start ($((++n))/$max), retrying in 5s..."
+        sleep 5
+      done
+      [ $n -eq $max ] && { echo "Ray head failed"; exit 1; }
+    }
+    retry_ray
+
+worker_start_ray_commands:
+  - |
+    retry_ray() {
+      local n=0; local max=30
+      until [ $n -ge $max ]; do
+        SKYPILOT_NUM_GPUS=0
+        command -v nvidia-smi >/dev/null 2>&1 && \
+          SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
+
+        ray stop || true
+        RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 \
+        ray start --disable-usage-stats \
+          --address=$RAY_HEAD_IP:{{ ray_port }} \
+          --object-manager-port=8076 \
+          --num-gpus=$SKYPILOT_NUM_GPUS --temp-dir {{ ray_temp_dir }} && break
+
+        echo "[worker] Ray failed to start ($((++n))/$max), retrying in 5s..."
+        sleep 5
+      done
+      [ $n -eq $max ] && { echo "Ray worker failed"; exit 1; }
+    }
+    retry_ray
+
+head_node: {}
+worker_nodes: {}
+
+head_setup_commands: []
+worker_setup_commands: []
+
+cluster_synced_files: []
+file_mounts_sync_continuously: False

sky/utils/cluster_utils.py

@@ -144,6 +144,9 @@ class SSHConfigHelper(object):
             username = docker_user
 
         key_path = cls.generate_local_key_file(cluster_name, auth_config)
+        # Keep the unexpanded path for SSH config (with ~)
+        key_path_for_config = key_path
+        # Expand the path for internal operations that need absolute path
         key_path = os.path.expanduser(key_path)
         sky_autogen_comment = ('# Added by sky (use `sky stop/down '
                                f'{cluster_name}` to remove)')
@@ -208,8 +211,9 @@ class SSHConfigHelper(object):
             node_name = cluster_name if i == 0 else cluster_name + f'-worker{i}'
             # TODO(romilb): Update port number when k8s supports multinode
             codegen += cls._get_generated_config(
-                sky_autogen_comment, node_name, ip, username, key_path,
-                proxy_command, port, docker_proxy_command) + '\n'
+                sky_autogen_comment, node_name, ip, username,
+                key_path_for_config, proxy_command, port,
+                docker_proxy_command) + '\n'
 
         cluster_config_path = os.path.expanduser(
             cls.ssh_cluster_path.format(cluster_name))

sky/utils/controller_utils.py

@@ -228,15 +228,21 @@ def get_controller_for_pool(pool: bool) -> Controllers:
 def high_availability_specified(cluster_name: Optional[str]) -> bool:
     """Check if the controller high availability is specified in user config.
     """
-    # pylint: disable=import-outside-toplevel
-    from sky.jobs import utils as managed_job_utils
-    if managed_job_utils.is_consolidation_mode():
-        return True
-
     controller = Controllers.from_name(cluster_name)
     if controller is None:
         return False
 
+    if controller.value.controller_type == 'jobs':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.jobs import utils as managed_job_utils
+        if managed_job_utils.is_consolidation_mode():
+            return True
+    elif controller.value.controller_type == 'serve':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.serve import serve_utils
+        if serve_utils.is_consolidation_mode():
+            return True
+
     if skypilot_config.loaded():
         return skypilot_config.get_nested((controller.value.controller_type,
                                            'controller', 'high_availability'),