Repository: skypilot-org/skypilot
Pull request: [K8s] Custom Image Support for Kubernetes Instances
<!-- Describe the changes in this PR -->
This PR resolves #2599 

When users opt to provide their custom images for setting up the k8s instance, instead of relying on Skypilot's default image built with `Dockerfile_k8s`, certain assumptions and adjustments are made to ensure smooth operation:

**Base Distribution:** We expect users to provide images that are based on Debian-based distributions.
**User Privilege:** It's crucial for users to configure the default user's privilege correctly within their custom image by either setting the default user to have root privilege or installing `sudo`.
**Dependency Installation & SSH Setup:** Instead of baking these directly into the image, we've shifted the responsibility of installing dependencies and setting up SSH to the `ray` manifest. These processes are triggered after the pods are active and are managed within `node_provider.py` under the `create_node()` method.
**SSH user update:** We have default `ssh_user` of `sky` set in `kubernetes-ray.yml.j2`. However, as users' may have different user name, we need to update the `ssh_user` in the cluster yaml and the `SSHCommandRunner` used within NodeProvider. This is each done from `node_provider.py/create_node()` and `node_provider.py/get_command_runner()`.

By adopting these measures, we aim to offer greater flexibility for users bringing their own images while ensuring consistent system behavior.

Update: This PR supports public registries only, and private registries may be supported in near future. Until then, there will be a doc on how to setup for the private registry by the user.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`

Tested Base Images:
- [x] continuumio/miniconda3:22.11.1
- [x] nvidia/cuda:12.2.0-devel-ubuntu20.04
- [x] pytorch/pytorch:latest
- [x] pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
- [x] tensorflow/tensorflow:latest
- [x] tensorflow/tensorflow:nightly-gpu-jupyter
- [x] ubuntu:latest
- [x] gpuci/miniforge-cuda:11.4-devel-ubuntu18.04

Tested Registries:
- [x] AWS ECR public
- [x] GCP Artifact Registry public
- [x] Docker Registry public

Early failovers:
- [x] AWS ECR private handled with error message
- [x] Docker private registry handled with error message with early failure
- [x] GCP Registry failover when user's not configured for the GCP account
- [x] GCP Registry failover when user's configured GCP account does not have IAM authorization
- [x] Attempting to pull from non-existent repository
- [x] Attempting to pull from non-existent tag name with correct repository

Problem statement:
[k8s] Custom image support for Kubernetes
We need to support custom images for Kubernetes clusters, possibly set via `image_id: docker:<image_id>`. 

We'd need to take care of a few things/refactor our k8s provisioning:

1. Install and configure openssh as a part of setup to ensure ssh access
2. Make sure basic dependencies are installed, and install them if necessary/add notes to docs about which dependencies are required

Hints from issue discussion:
This issue is stale because it has been open 120 days with no activity. Remove stale label or comment or this will be closed in 10 days.

Commit message: multinode fixes

Patch to review:
diff --git a/Dockerfile_k8s b/Dockerfile_k8s
index cde67cac210..6cbc167265e 100644
--- a/Dockerfile_k8s
+++ b/Dockerfile_k8s
@@ -25,26 +25,21 @@ RUN useradd -m -s /bin/bash sky && \
 # Switch to sky user
 USER sky
 
-# Install SkyPilot pip dependencies
+# Install SkyPilot pip dependencies preemptively to speed up provisioning time
 RUN pip install wheel Click colorama cryptography jinja2 jsonschema && \
     pip install networkx oauth2client pandas pendulum PrettyTable && \
     pip install ray[default]==2.4.0 rich tabulate filelock && \
     pip install packaging 'protobuf<4.0.0' pulp && \
-    pip install awscli boto3 pycryptodome==3.12.0 && \
-    pip install docker kubernetes
+    pip install pycryptodome==3.12.0 && \
+    pip install docker kubernetes==28.1.0
 
 # Add /home/sky/.local/bin/ to PATH
 RUN echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc
 
-# Install SkyPilot. This is purposely separate from installing SkyPilot
-# dependencies to optimize rebuild time
+# Copy SkyPilot code base. This is required for the ssh jump pod to find the
+# lifecycle management scripts
 COPY --chown=sky . /skypilot/sky/
 
-# TODO(romilb): Installing SkyPilot may not be necessary since ray up will do it
-RUN cd /skypilot/ && \
-    sudo mv -v sky/setup_files/* . && \
-    pip install ".[aws]"
-
 # Set PYTHONUNBUFFERED=1 to have Python print to stdout/stderr immediately
 ENV PYTHONUNBUFFERED=1
 
diff --git a/Dockerfile_k8s_gpu b/Dockerfile_k8s_gpu
index 6f242b6f5ab..2d579cdbdd2 100644
--- a/Dockerfile_k8s_gpu
+++ b/Dockerfile_k8s_gpu
@@ -48,11 +48,6 @@ RUN echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc
 # dependencies to optimize rebuild time
 COPY --chown=sky . /skypilot/sky/
 
-# TODO(romilb): Installing SkyPilot may not be necessary since ray up will do it
-RUN cd /skypilot/ && \
-    sudo mv -v sky/setup_files/* . && \
-    pip install ".[aws]"
-
 # Set PYTHONUNBUFFERED=1 to have Python print to stdout/stderr immediately
 ENV PYTHONUNBUFFERED=1
 
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 4cea0232b4b..f4dd4e2333d 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -66,7 +66,6 @@
 # Exclude subnet mask from IP address regex.
 IP_ADDR_REGEX = r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}(?!/\d{1,2})\b'
 SKY_REMOTE_PATH = '~/.sky/wheels'
-SKY_USER_FILE_PATH = '~/.sky/generated'
 
 BOLD = '\033[1m'
 RESET_BOLD = '\033[0m'
@@ -160,7 +159,8 @@ def is_ip(s: str) -> bool:
 
 
 def _get_yaml_path_from_cluster_name(cluster_name: str,
-                                     prefix: str = SKY_USER_FILE_PATH) -> str:
+                                     prefix: str = constants.SKY_USER_FILE_PATH
+                                    ) -> str:
     output_path = pathlib.Path(
         prefix).expanduser().resolve() / f'{cluster_name}.yml'
     os.makedirs(output_path.parents[0], exist_ok=True)
@@ -410,7 +410,7 @@ class SSHConfigHelper(object):
 
     ssh_conf_path = '~/.ssh/config'
     ssh_conf_lock_path = os.path.expanduser('~/.sky/ssh_config.lock')
-    ssh_cluster_path = SKY_USER_FILE_PATH + '/ssh/{}'
+    ssh_cluster_path = constants.SKY_USER_FILE_PATH + '/ssh/{}'
 
     @classmethod
     def _get_generated_config(cls, autogen_comment: str, host_name: str,
@@ -743,6 +743,96 @@ def _restore_block(new_block: Dict[str, Any], old_block: Dict[str, Any]):
     return common_utils.dump_yaml_str(new_config)
 
 
+def _add_kubernetes_config_fields(yaml_path: str) -> None:
+    """Adds or updates fields in the YAML with fields from the ~/.sky/config's
+    kubernetes.pod_spec dict.
+
+    This can be used to add fields to the YAML that are not supported by
+    SkyPilot yet, or require simple configuration (e.g., adding an
+    imagePullSecrets field).
+
+    Note that new fields are added and existing ones are updated. Nested fields
+    are not completely replaced, instead their objects are merged. Similarly,
+    if a list is encountered in the config, it will be appended to the
+    destination list.
+    For example, if the YAML has the following:
+        ```
+        ...
+        node_config:
+            spec:
+                containers:
+                    - name: ray
+                    image: rayproject/ray:nightly
+        ```
+
+    and the config has the following:
+        ```
+        kubernetes:
+            pod_override:
+                spec:
+                    containers:
+                        imagePullSecrets:
+                            - name: my-secret
+        ```
+
+    then the resulting YAML will be:
+        ```
+        ...
+        node_config:
+            spec:
+                containers:
+                    - name: ray
+                    image: rayproject/ray:nightly
+                    imagePullSecrets:
+                        - name: my-secret
+        ```
+    """
+
+    def _merge_dicts(source, destination):
+        """Merge two dictionaries.
+
+         Updates nested dictionaries instead of replacing them.
+
+         If a list is encountered, it will be appended to the destination list.
+         An exception is when the key is 'containers', in which case the
+         first container in the list will be fetched and _merge_dict will be
+         called on it with the first container in the destination list.
+         """
+        for key, value in source.items():
+            if isinstance(value, dict) and key in destination:
+                _merge_dicts(value, destination[key])
+            elif isinstance(value, list) and key in destination:
+                assert isinstance(destination[key], list), \
+                    f'Expected {key} to be a list, found {destination[key]}'
+                if key == 'containers':
+                    # If the key is 'containers', we take the first and only
+                    # container in the list and merge it.
+                    assert len(value) == 1, \
+                        f'Expected only one container, found {value}'
+                    _merge_dicts(value[0], destination[key][0])
+                else:
+                    destination[key].extend(value)
+            else:
+                destination[key] = value
+
+    with open(yaml_path, 'r') as f:
+        yaml_content = f.read()
+    yaml_obj = yaml.safe_load(yaml_content)
+    kubernetes_config = skypilot_config.get_nested(
+        ('kubernetes', 'pod_override'), {})
+
+    # Merge the kubernetes config into the YAML for both head and worker nodes.
+    _merge_dicts(
+        kubernetes_config,
+        yaml_obj['available_node_types']['ray_head_default']['node_config'])
+    _merge_dicts(
+        kubernetes_config,
+        yaml_obj['available_node_types']['ray_worker_default']['node_config'])
+
+    # Write the updated YAML back to the file
+    common_utils.dump_yaml(yaml_path, yaml_obj)
+
+
 # TODO: too many things happening here - leaky abstraction. Refactor.
 @timeline.event
 def write_cluster_config(
@@ -949,6 +1039,10 @@ def write_cluster_config(
         return config_dict
     _add_auth_to_cluster_config(cloud, tmp_yaml_path)
 
+    # Add kubernetes config fields from ~/.sky/config
+    if isinstance(cloud, clouds.Kubernetes):
+        _add_kubernetes_config_fields(tmp_yaml_path)
+
     # Restore the old yaml content for backward compatibility.
     if os.path.exists(yaml_path) and keep_launch_fields_in_existing_config:
         with open(yaml_path, 'r') as f:
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index 340d63b8a9f..fe0d8c3faf9 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -61,10 +61,7 @@ class Kubernetes(clouds.Cloud):
         clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER: 'Custom disk '
                                                              'tiers are not '
                                                              'supported in '
-                                                             'Kubernetes.',
-        clouds.CloudImplementationFeatures.DOCKER_IMAGE: 'Docker image is not '
-                                                         'supported in '
-                                                         'Kubernetes.',
+                                                             'Kubernetes.'
     }
 
     IMAGE_CPU = 'skypilot:cpu-ubuntu-2004'
@@ -212,15 +209,16 @@ def make_deploy_resources_variables(
         acc_count = k.accelerator_count if k.accelerator_count else 0
         acc_type = k.accelerator_type if k.accelerator_type else None
 
-        # Select image based on whether we are using GPUs or not.
-        image_id = self.IMAGE_GPU if acc_count > 0 else self.IMAGE_CPU
-        # Get the container image ID from the service catalog.
-        # TODO(romilb): Note that currently we do not support custom images,
-        #  so the image_id should start with 'skypilot:'.
-        #  In the future we may want to get image_id from the resources object.
-        assert image_id.startswith('skypilot:')
-        image_id = service_catalog.get_image_id_from_tag(image_id,
-                                                         clouds='kubernetes')
+        if resources.image_id is not None:
+            # Use custom image specified in resources
+            image_id_with_region = resources.image_id['kubernetes']
+            image_id = image_id_with_region.lstrip('docker:')
+        else:
+            # Select image based on whether we are using GPUs or not.
+            image_id = self.IMAGE_GPU if acc_count > 0 else self.IMAGE_CPU
+            # Get the container image ID from the service catalog.
+            image_id = service_catalog.get_image_id_from_tag(
+                image_id, clouds='kubernetes')
         # TODO(romilb): Create a lightweight image for SSH jump host
         ssh_jump_image = service_catalog.get_image_id_from_tag(
             self.IMAGE_CPU, clouds='kubernetes')
diff --git a/sky/exceptions.py b/sky/exceptions.py
index c6be390f7ce..ae717b23bfb 100644
--- a/sky/exceptions.py
+++ b/sky/exceptions.py
@@ -13,6 +13,8 @@
 RSYNC_FILE_NOT_FOUND_CODE = 23
 # Arbitrarily chosen value. Used in SkyPilot's storage mounting scripts
 MOUNT_PATH_NON_EMPTY_CODE = 42
+# Arbitrarily chosen value. Used to provision Kubernetes instance in Skypilot
+INSUFFICIENT_PRIVILEGES_CODE = 52
 # Return code when git command is ran in a dir that is not git repo
 GIT_FATAL_EXIT_CODE = 128
 
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index 1ef51ca2c7b..5e7df330288 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -2,6 +2,7 @@
 
 SKY_LOGS_DIRECTORY = '~/sky_logs'
 SKY_REMOTE_WORKDIR = '~/sky_workdir'
+SKY_USER_FILE_PATH = '~/.sky/generated'
 
 # Default Ray port is 6379. Default Ray dashboard port is 8265.
 # Default Ray tempdir is /tmp/ray.
diff --git a/sky/skylet/providers/kubernetes/node_provider.py b/sky/skylet/providers/kubernetes/node_provider.py
index 2a2d5b2c2aa..0656fbb424d 100644
--- a/sky/skylet/providers/kubernetes/node_provider.py
+++ b/sky/skylet/providers/kubernetes/node_provider.py
@@ -1,5 +1,7 @@
 import copy
 import logging
+import os
+import re
 import time
 from typing import Dict
 from uuid import uuid4
@@ -10,8 +12,12 @@
 from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME
 from ray.autoscaler.tags import TAG_RAY_NODE_KIND
 
+from sky import exceptions
 from sky.adaptors import kubernetes
+from sky.backends import backend_utils
+from sky.skylet import constants
 from sky.skylet.providers.kubernetes import config
+from sky.utils import cluster_yaml_utils
 from sky.utils import common_utils
 from sky.utils import kubernetes_utils
 
@@ -73,6 +79,20 @@ def to_label_selector(tags):
     return label_selector
 
 
+def run_command_on_pods(node_name, node_namespace, command):
+    cmd_output = kubernetes.stream()(
+        kubernetes.core_api().connect_get_namespaced_pod_exec,
+        node_name,
+        node_namespace,
+        command=command,
+        stderr=True,
+        stdin=False,
+        stdout=True,
+        tty=False,
+        _request_timeout=kubernetes.API_TIMEOUT)
+    return cmd_output
+
+
 class KubernetesNodeProvider(NodeProvider):
 
     def __init__(self, provider_config, cluster_name):
@@ -186,6 +206,23 @@ def set_node_tags(self, node_ids, tags):
         # One more try
         self._set_node_tags(node_ids, tags)
 
+    def _recover_cluster_yaml_path(self, cluster_name_with_hash: str) -> str:
+        # 'cluster_name_with_hash' combines the cluster name and hash value,
+        # separated by a hyphen. By using 'slice_length', we remove the hash
+        # (and its preceding hyphen) to retrieve the original cluster name.
+        slice_length = -(common_utils.USER_HASH_LENGTH_IN_CLUSTER_NAME + 1)
+        cluster_name = cluster_name_with_hash[:slice_length]
+        cluster_yaml_path = (os.path.join(
+            os.path.expanduser(constants.SKY_USER_FILE_PATH),
+            f'{cluster_name}.yml'))
+        # Check if cluster_yaml_path exists. If not, we are running on
+        # the master node in a multi-node setup, in which case we must use the
+        # default ~/.sky/sky_ray.yml path.
+        if not os.path.exists(cluster_yaml_path):
+            cluster_yaml_path = os.path.expanduser(
+                cluster_yaml_utils.SKY_CLUSTER_YAML_REMOTE_PATH)
+        return cluster_yaml_path
+
     def _set_node_tags(self, node_id, tags):
         pod = kubernetes.core_api().read_namespaced_pod(node_id, self.namespace)
         pod.metadata.labels.update(tags)
@@ -264,7 +301,7 @@ def _raise_pod_scheduling_errors(self, new_nodes):
                                              f'Details: \'{event_message}\' ')
         raise config.KubernetesError(f'{timeout_err_msg}')
 
-    def _wait_for_pods_to_schedule(self, new_nodes):
+    def _wait_for_pods_to_schedule(self, new_nodes_with_jump_pod):
         """Wait for all pods to be scheduled.
 
         Wait for all pods including jump pod to be scheduled, and if it
@@ -275,7 +312,7 @@ def _wait_for_pods_to_schedule(self, new_nodes):
         start_time = time.time()
         while time.time() - start_time < self.timeout:
             all_pods_scheduled = True
-            for node in new_nodes:
+            for node in new_nodes_with_jump_pod:
                 # Iterate over each pod to check their status
                 pod = kubernetes.core_api().read_namespaced_pod(
                     node.metadata.name, self.namespace)
@@ -292,7 +329,7 @@ def _wait_for_pods_to_schedule(self, new_nodes):
 
         # Handle pod scheduling errors
         try:
-            self._raise_pod_scheduling_errors(new_nodes)
+            self._raise_pod_scheduling_errors(new_nodes_with_jump_pod)
         except config.KubernetesError:
             raise
         except Exception as e:
@@ -301,7 +338,7 @@ def _wait_for_pods_to_schedule(self, new_nodes):
                 'for pod scheduling failure. '
                 f'Error: {common_utils.format_exception(e)}') from None
 
-    def _wait_for_pods_to_run(self, new_nodes):
+    def _wait_for_pods_to_run(self, new_nodes_with_jump_pod):
         """Wait for pods and their containers to be ready.
 
         Pods may be pulling images or may be in the process of container
@@ -310,7 +347,7 @@ def _wait_for_pods_to_run(self, new_nodes):
         while True:
             all_pods_running = True
             # Iterate over each pod to check their status
-            for node in new_nodes:
+            for node in new_nodes_with_jump_pod:
                 pod = kubernetes.core_api().read_namespaced_pod(
                     node.metadata.name, self.namespace)
 
@@ -345,6 +382,61 @@ def _wait_for_pods_to_run(self, new_nodes):
                 break
             time.sleep(1)
 
+    def _check_user_privilege(self, new_nodes):
+        # Checks if the default user has sufficient privilege to set up
+        # the kubernetes instance pod.
+        check_k8s_user_sudo_cmd = [
+            '/bin/sh',
+            '-c',
+            (
+                'if [ $(id -u) -eq 0 ]; then'
+                # If user is root, create an alias for sudo used in skypilot setup
+                '  echo \'alias sudo=""\' >> ~/.bashrc; '
+                'else '
+                '  if command -v sudo >/dev/null 2>&1; then '
+                '    timeout 2 sudo -l >/dev/null 2>&1 || '
+                f'    ( echo {exceptions.INSUFFICIENT_PRIVILEGES_CODE!r}; ); '
+                '  else '
+                f'    ( echo {exceptions.INSUFFICIENT_PRIVILEGES_CODE!r}; ); '
+                '  fi; '
+                'fi')
+        ]
+
+        for new_node in new_nodes:
+            privilege_check = run_command_on_pods(new_node.metadata.name,
+                                                  self.namespace,
+                                                  check_k8s_user_sudo_cmd)
+            if privilege_check == str(exceptions.INSUFFICIENT_PRIVILEGES_CODE):
+                raise config.KubernetesError(
+                    'Insufficient system privileges detected. '
+                    'Ensure the default user has root access or '
+                    '"sudo" is installed and the user is added to the sudoers '
+                    'from the image.')
+
+    def _setup_ssh_in_pods(self, new_nodes):
+        # Setting up ssh for the pod instance. This is already setup for
+        # the jump pod so it does not need to be run for it.
+        set_k8s_ssh_cmd = [
+            '/bin/sh', '-c',
+            ('prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }; '
+             'export DEBIAN_FRONTEND=noninteractive;'
+             '$(prefix_cmd) apt-get update;'
+             '$(prefix_cmd) apt install openssh-server rsync -y; '
+             '$(prefix_cmd) mkdir -p /var/run/sshd; '
+             '$(prefix_cmd) sed -i "s/PermitRootLogin prohibit-password/PermitRootLogin yes/" /etc/ssh/sshd_config; '
+             '$(prefix_cmd) sed "s@session\\s*required\\s*pam_loginuid.so@session optional pam_loginuid.so@g" -i /etc/pam.d/sshd; '
+             'cd /etc/ssh/ && $(prefix_cmd) ssh-keygen -A; '
+             '$(prefix_cmd) mkdir -p ~/.ssh; '
+             '$(prefix_cmd) cp /etc/secret-volume/ssh-publickey ~/.ssh/authorized_keys; '
+             '$(prefix_cmd) service ssh restart')
+        ]
+
+
+        # TODO(romilb): We need logging and surface errors here.
+        for new_node in new_nodes:
+            run_command_on_pods(new_node.metadata.name, self.namespace,
+                                set_k8s_ssh_cmd)
+
     def _set_env_vars_in_pods(self, new_nodes):
         """Setting environment variables in pods.
 
@@ -363,22 +455,33 @@ def _set_env_vars_in_pods(self, new_nodes):
         """
         set_k8s_env_var_cmd = [
             '/bin/sh', '-c',
-            ('printenv | awk -F "=" \'{print "export " $1 "=\\047" $2 "\\047"}\' > ~/k8s_env_var.sh && '
+            ('prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; } && '
+             'printenv | awk -F "=" \'{print "export " $1 "=\\047" $2 "\\047"}\' > ~/k8s_env_var.sh && '
              'mv ~/k8s_env_var.sh /etc/profile.d/k8s_env_var.sh || '
-             'sudo mv ~/k8s_env_var.sh /etc/profile.d/k8s_env_var.sh')
+             '$(prefix_cmd) mv ~/k8s_env_var.sh /etc/profile.d/k8s_env_var.sh')
         ]
 
         for new_node in new_nodes:
-            kubernetes.stream()(
-                kubernetes.core_api().connect_get_namespaced_pod_exec,
-                new_node.metadata.name,
-                self.namespace,
-                command=set_k8s_env_var_cmd,
-                stderr=True,
-                stdin=False,
-                stdout=True,
-                tty=False,
-                _request_timeout=kubernetes.API_TIMEOUT)
+            run_command_on_pods(new_node.metadata.name, self.namespace,
+                                set_k8s_env_var_cmd)
+
+    def _update_ssh_user_config(self, new_nodes, cluster_name_with_hash):
+        get_k8s_ssh_user_cmd = ['/bin/sh', '-c', ('echo $(whoami)')]
+        for new_node in new_nodes:
+            ssh_user = run_command_on_pods(new_node.metadata.name,
+                                           self.namespace, get_k8s_ssh_user_cmd)
+
+        cluster_yaml_path = self._recover_cluster_yaml_path(
+            cluster_name_with_hash)
+        with open(cluster_yaml_path, 'r') as f:
+            content = f.read()
+
+        # Replacing the default ssh user name with the actual user name.
+        # This updates user name specified in user's custom image if it's used.
+        content = re.sub(r'ssh_user: \w+', f'ssh_user: {ssh_user}', content)
+
+        with open(cluster_yaml_path, 'w') as f:
+            f.write(content)
 
     def create_node(self, node_config, tags, count):
         conf = copy.deepcopy(node_config)
@@ -423,17 +526,35 @@ def create_node(self, node_config, tags, count):
         # Adding the jump pod to the new_nodes list as well so it can be
         # checked if it's scheduled and running along with other pod instances.
         ssh_jump_pod_name = conf['metadata']['labels']['skypilot-ssh-jump']
+        new_nodes_with_jump_pod = new_nodes[:]
         jump_pod = kubernetes.core_api().read_namespaced_pod(
             ssh_jump_pod_name, self.namespace)
-        new_nodes.append(jump_pod)
+        new_nodes_with_jump_pod.append(jump_pod)
+        node_names = [node.metadata.name for node in new_nodes_with_jump_pod]
 
         # Wait until the pods are scheduled and surface cause for error
         # if there is one
-        self._wait_for_pods_to_schedule(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Waiting for pods to schedule. Pods: {node_names}')
+        self._wait_for_pods_to_schedule(new_nodes_with_jump_pod)
         # Wait until the pods and their containers are up and running, and
         # fail early if there is an error
-        self._wait_for_pods_to_run(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Waiting for pods to run. Pods: {node_names}')
+        self._wait_for_pods_to_run(new_nodes_with_jump_pod)
+        logger.info(config.log_prefix +
+                    f'Checking if user in image has sufficient privileges.')
+        self._check_user_privilege(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Setting up SSH in pod.')
+        self._setup_ssh_in_pods(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Setting up environment variables in pod.')
         self._set_env_vars_in_pods(new_nodes)
+        cluster_name_with_hash = conf['metadata']['labels']['skypilot-cluster']
+        logger.info(config.log_prefix +
+                    f'Fetching and updating ssh username.')
+        self._update_ssh_user_config(new_nodes, cluster_name_with_hash)
 
     def terminate_node(self, node_id):
         logger.info(config.log_prefix + 'calling delete_namespaced_pod')
@@ -480,7 +601,7 @@ def get_command_runner(self,
                            log_prefix,
                            node_id,
                            auth_config,
-                           cluster_name,
+                           cluster_name_with_hash,
                            process_runner,
                            use_internal_ip,
                            docker_config=None):
@@ -492,7 +613,8 @@ def get_command_runner(self,
         node_id(str): the node ID.
         auth_config(dict): the authentication configs from the autoscaler
             yaml file.
-        cluster_name(str): the name of the cluster.
+        cluster_name_with_hash(str): the name of the cluster and hash value,
+            separated by a hyphen.
         process_runner(module): the module to use to run the commands
             in the CommandRunner. E.g., subprocess.
         use_internal_ip(bool): whether the node_id belongs to an internal ip
@@ -500,12 +622,22 @@ def get_command_runner(self,
         docker_config(dict): If set, the docker information of the docker
             container that commands should be run on.
         """
+        # For custom images, the username might differ across images.
+        # The 'ssh_user' is updated inplace in the YAML at the end of the
+        # 'create_node()' process in _update_ssh_user_config.
+        # Since the node provider is initialized with stale auth information,
+        # we need to reload the updated user from YAML.
+        cluster_yaml_path = self._recover_cluster_yaml_path(
+            cluster_name_with_hash)
+        ssh_credentials = backend_utils.ssh_credential_from_yaml(cluster_yaml_path)
+        auth_config['ssh_user'] = ssh_credentials['ssh_user']
+
         common_args = {
             'log_prefix': log_prefix,
             'node_id': node_id,
             'provider': self,
             'auth_config': auth_config,
-            'cluster_name': cluster_name,
+            'cluster_name': cluster_name_with_hash,
             'process_runner': process_runner,
             'use_internal_ip': use_internal_ip,
         }
diff --git a/sky/templates/kubernetes-ray.yml.j2 b/sky/templates/kubernetes-ray.yml.j2
index c8babb255ae..a08586701fb 100644
--- a/sky/templates/kubernetes-ray.yml.j2
+++ b/sky/templates/kubernetes-ray.yml.j2
@@ -190,10 +190,6 @@ available_node_types:
             name: dev-fuse
           securityContext:          # Required for FUSE mounting. TODO(romilb): See if we can grant a reduced set of privileges.
             privileged: true
-          lifecycle:
-            postStart:
-              exec:
-                command: ["/bin/bash", "-c", "mkdir -p ~/.ssh && cp /etc/secret-volume/ssh-publickey ~/.ssh/authorized_keys && sudo service ssh restart"]
           resources:
             requests:
               cpu: {{cpus}}
@@ -254,10 +250,6 @@ available_node_types:
           image: {{image_id}}
           command: ["/bin/bash", "-c", "--"]
           args: ["trap : TERM INT; sleep infinity & wait;"]
-          lifecycle:
-            postStart:
-              exec:
-                command: ["/bin/bash", "-c", "mkdir -p ~/.ssh && cp /etc/secret-volume/ssh-publickey ~/.ssh/authorized_keys && sudo service ssh restart"]
           ports:
           - containerPort: 22  # Used for SSH
           # This volume allocates shared memory for Ray to use for its plasma
@@ -288,17 +280,19 @@ setup_commands:
   # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
-  - mkdir -p ~/.ssh; touch ~/.ssh/config;
+  - sudo DEBIAN_FRONTEND=noninteractive apt install gcc patch pciutils rsync fuse curl python3 python3-pip -y;
+    mkdir -p ~/.ssh; sudo touch ~/.ssh/config;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
     (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
-    {{ conda_installation_commands }}
+    {{ conda_installation_commands }} conda init;
     source ~/.bashrc;
-    mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app && touch ~/.sudo_as_admin_successful;
+    mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app && sudo touch ~/.sudo_as_admin_successful;
     (pip3 list | grep skypilot && [ "$(cat {{sky_remote_path}}/current_sky_wheel_hash)" == "{{sky_wheel_hash}}" ]) || (pip3 uninstall skypilot -y; pip3 install "$(echo {{sky_remote_path}}/{{sky_wheel_hash}}/skypilot-{{sky_version}}*.whl)[remote]" && echo "{{sky_wheel_hash}}" > {{sky_remote_path}}/current_sky_wheel_hash || exit 1);
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
+    pip3 install ray==2.4.0 aiohttp opencensus prometheus_client aiohttp_cors kubernetes==28.1.0;
     python3 -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
     [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf');
 
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index c0d7a59a382..5a4546b7324 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -274,6 +274,10 @@ def dump_yaml_str(config):
     # https://github.com/yaml/pyyaml/issues/127
     class LineBreakDumper(yaml.SafeDumper):
 
+        # Disable anchors and aliases to keep the output clean.
+        # https://github.com/yaml/pyyaml/issues/535
+        ignore_aliases = lambda *args: True
+
         def write_line_break(self, data=None):
             super().write_line_break(data)
             if len(self.indents) == 1:
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index deb63f3151e..12725c6d5aa 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -591,6 +591,12 @@ def get_config_schema():
                             for type in kubernetes_enums.KubernetesPortMode
                         ]
                     },
+                    'pod_override': {
+                        'type': 'object',
+                        'required': [],
+                        # Allow arbitrary keys since validating pod spec is hard
+                        'additionalProperties': True,
+                    }
                 }
             },
             'oci': {


Code review comments:
diff_hunk: @@ -288,17 +280,19 @@ setup_commands:
   # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
-  - mkdir -p ~/.ssh; touch ~/.ssh/config;
+  - sudo DEBIAN_FRONTEND=noninteractive apt install gcc patch pciutils rsync fuse curl python3 python3-pip -y;
+    mkdir -p ~/.ssh; sudo touch ~/.ssh/config;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
     (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
-    {{ conda_installation_commands }}
+    {{ conda_installation_commands }} conda init;
     source ~/.bashrc;
-    mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app && touch ~/.sudo_as_admin_successful;
+    mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app && sudo touch ~/.sudo_as_admin_successful;
     (pip3 list | grep skypilot && [ "$(cat {{sky_remote_path}}/current_sky_wheel_hash)" == "{{sky_wheel_hash}}" ]) || (pip3 uninstall skypilot -y; pip3 install "$(echo {{sky_remote_path}}/{{sky_wheel_hash}}/skypilot-{{sky_version}}*.whl)[remote]" && echo "{{sky_wheel_hash}}" > {{sky_remote_path}}/current_sky_wheel_hash || exit 1);
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
+    pip3 install ray==2.4.0 aiohttp opencensus prometheus_client aiohttp_cors kubernetes==28.1.0;
path: sky/templates/kubernetes-ray.yml.j2
text: @user1:
Ray installation sould go before skypilot installation and needs some patch to make it work. Reference: https://github.com/skypilot-org/skypilot/blob/efdd07deb7943666d67f000f6e09701e5161ecae/sky/templates/aws-ray.yml.j2#L163

https://github.com/skypilot-org/skypilot/blob/efdd07deb7943666d67f000f6e09701e5161ecae/sky/templates/aws-ray.yml.j2#L170

diff_hunk: @@ -423,17 +526,35 @@ def create_node(self, node_config, tags, count):
         # Adding the jump pod to the new_nodes list as well so it can be
         # checked if it's scheduled and running along with other pod instances.
         ssh_jump_pod_name = conf['metadata']['labels']['skypilot-ssh-jump']
+        new_nodes_with_jump_pod = new_nodes[:]
         jump_pod = kubernetes.core_api().read_namespaced_pod(
             ssh_jump_pod_name, self.namespace)
-        new_nodes.append(jump_pod)
+        new_nodes_with_jump_pod.append(jump_pod)
+        node_names = [node.metadata.name for node in new_nodes_with_jump_pod]
 
         # Wait until the pods are scheduled and surface cause for error
         # if there is one
-        self._wait_for_pods_to_schedule(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Waiting for pods to schedule. Pods: {node_names}')
+        self._wait_for_pods_to_schedule(new_nodes_with_jump_pod)
         # Wait until the pods and their containers are up and running, and
         # fail early if there is an error
-        self._wait_for_pods_to_run(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Waiting for pods to run. Pods: {node_names}')
+        self._wait_for_pods_to_run(new_nodes_with_jump_pod)
+        logger.info(config.log_prefix +
+                    f'Checking if user in image has sufficient privileges.')
+        self._check_user_privilege(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Setting up SSH in pod.')
+        self._setup_ssh_in_pods(new_nodes)
+        logger.info(config.log_prefix +
+                    f'Setting up environment variables in pod.')
         self._set_env_vars_in_pods(new_nodes)
path: sky/skylet/providers/kubernetes/node_provider.py
text: @user1:
Do we need to setup envs for jump pod?

@author:
The envs are necessary for pods to interact with GPUs, so it's only necessary for the instance pods, not for jump pod.