[k8s] Share SSH jump pod

Dockerfile_k8s

@@ -5,7 +5,7 @@ FROM continuumio/miniconda3:23.3.1-0
 
 # Initialize conda for root user, install ssh and other local dependencies
 RUN apt update -y && \
-    apt install gcc rsync sudo patch openssh-server pciutils nano fuse -y && \
+    apt install gcc rsync sudo patch openssh-server pciutils nano fuse socat netcat -y && \
     rm -rf /var/lib/apt/lists/* && \
     apt remove -y python3 && \
     conda init

Dockerfile_k8s_gpu

@@ -7,7 +7,7 @@ FROM rayproject/ray:2.4.0-py310-gpu
 # We remove cuda lists to avoid conflicts with the cuda version installed by ray
 RUN sudo rm -rf /etc/apt/sources.list.d/cuda* && \
     sudo apt update -y && \
-    sudo apt install gcc rsync sudo patch openssh-server pciutils nano fuse unzip -y && \
+    sudo apt install gcc rsync sudo patch openssh-server pciutils nano fuse unzip socat netcat -y && \
     sudo rm -rf /var/lib/apt/lists/* && \
     sudo apt remove -y python3 && \
     conda init
@@ -39,7 +39,8 @@ RUN pip install wheel Click colorama cryptography jinja2 jsonschema && \
     pip install rich tabulate filelock && \
     pip install packaging 'protobuf<4.0.0' pulp && \
     pip install pycryptodome==3.12.0 && \
-    pip install docker kubernetes==28.1.0
+    pip install docker kubernetes==28.1.0 && \
+    pip install -U boto3==1.34.46 # Ray base image has a stale boto3 version - https://github.com/skypilot-org/skypilot/issues/3158
 
 # Add /home/sky/.local/bin/ to PATH
 RUN echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc

docs/source/reference/kubernetes/kubernetes-setup.rst

@@ -379,6 +379,7 @@ To use this mode:
     .. code-block:: bash
 
       # Patch the nginx ingress service with an external IP. Can be any node's IP if using NodePort service.
+      # Replace <IP> in the following command with the IP you select.
       $ kubectl patch svc ingress-nginx-controller -n ingress-nginx -p '{"spec": {"externalIPs": ["<IP>"]}}'
 
     If the ``EXTERNAL-IP`` field is left as ``<none>``, SkyPilot will use ``localhost`` as the external IP for the Ingress,

sky/authentication.py

@@ -41,6 +41,7 @@
 from sky import skypilot_config
 from sky.adaptors import gcp
 from sky.adaptors import ibm
+from sky.adaptors import kubernetes
 from sky.adaptors import runpod
 from sky.clouds.utils import lambda_utils
 from sky.provision.fluidstack import fluidstack_utils
@@ -398,29 +399,29 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
                 from None
     get_or_generate_keys()
 
-    # Run kubectl command to add the public key to the cluster.
+    # Add the user's public key to the SkyPilot cluster.
     public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
-    key_label = clouds.Kubernetes.SKY_SSH_KEY_SECRET_NAME
-    cmd = f'kubectl create secret generic {key_label} ' \
-          f'--from-file=ssh-publickey={public_key_path}'
-    try:
-        subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
-    except subprocess.CalledProcessError as e:
-        output = e.output.decode('utf-8')
-        suffix = f'\nError message: {output}'
-        if 'already exists' in output:
-            logger.debug(
-                f'Key {key_label} already exists in the cluster, using it...')
-        elif any(err in output for err in ['connection refused', 'timeout']):
-            with ux_utils.print_exception_no_traceback():
-                raise ConnectionError(
-                    'Failed to connect to the cluster. Check if your '
-                    'cluster is running, your kubeconfig is correct '
-                    'and you can connect to it using: '
-                    f'kubectl get namespaces.{suffix}') from e
-        else:
-            logger.error(suffix)
-            raise
+    secret_name = clouds.Kubernetes.SKY_SSH_KEY_SECRET_NAME
+    secret_field_name = clouds.Kubernetes.SKY_SSH_KEY_SECRET_FIELD_NAME
+    namespace = kubernetes_utils.get_current_kube_config_context_namespace()
+    k8s = kubernetes.get_kubernetes()
+    with open(public_key_path, 'r', encoding='utf-8') as f:
+        public_key = f.read()
+        if not public_key.endswith('\n'):
+            public_key += '\n'
+        secret_metadata = k8s.client.V1ObjectMeta(name=secret_name,
+                                                  labels={'parent': 'skypilot'})
+        secret = k8s.client.V1Secret(
+            metadata=secret_metadata,
+            string_data={secret_field_name: public_key})
+    if kubernetes_utils.check_secret_exists(secret_name, namespace):
+        logger.debug(f'Key {secret_name} exists in the cluster, patching it...')
+        kubernetes.core_api().patch_namespaced_secret(secret_name, namespace,
+                                                      secret)
+    else:
+        logger.debug(
+            f'Key {secret_name} does not exist in the cluster, creating it...')
+        kubernetes.core_api().create_namespaced_secret(namespace, secret)
 
     ssh_jump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
     if network_mode == nodeport_mode:
@@ -438,7 +439,6 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     # Setup service for SSH jump pod. We create the SSH jump service here
     # because we need to know the service IP address and port to set the
     # ssh_proxy_command in the autoscaler config.
-    namespace = kubernetes_utils.get_current_kube_config_context_namespace()
     kubernetes_utils.setup_ssh_jump_svc(ssh_jump_name, namespace, service_type)
 
     ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(

sky/clouds/kubernetes.py

@@ -28,8 +28,10 @@
 class Kubernetes(clouds.Cloud):
     """Kubernetes."""
 
-    SKY_SSH_KEY_SECRET_NAME = f'sky-ssh-{common_utils.get_user_hash()}'
-    SKY_SSH_JUMP_NAME = f'sky-ssh-jump-{common_utils.get_user_hash()}'
+    SKY_SSH_KEY_SECRET_NAME = 'sky-ssh-keys'
+    SKY_SSH_KEY_SECRET_FIELD_NAME = \
+        f'ssh-publickey-{common_utils.get_user_hash()}'
+    SKY_SSH_JUMP_NAME = 'sky-ssh-jump-pod'
     PORT_FORWARD_PROXY_CMD_TEMPLATE = \
         'kubernetes-port-forward-proxy-command.sh.j2'
     PORT_FORWARD_PROXY_CMD_PATH = '~/.sky/port-forward-proxy-cmd.sh'

sky/provision/kubernetes/instance.py

@@ -338,11 +338,11 @@ def _setup_ssh_in_pods(namespace: str, new_nodes: List) -> None:
             '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) chown -R $(whoami) ~/.ssh;'
             '$(prefix_cmd) chmod 700 ~/.ssh; '
             '$(prefix_cmd) chmod 644 ~/.ssh/authorized_keys; '
+            '$(prefix_cmd) cat /etc/secret-volume/ssh-publickey* > '
+            '~/.ssh/authorized_keys; '
             '$(prefix_cmd) service ssh restart; '
             # Eliminate the error
             # `mesg: ttyname failed: inappropriate ioctl for device`.

sky/provision/kubernetes/utils.py

@@ -1150,3 +1150,22 @@ def check_nvidia_runtime_class() -> bool:
     nvidia_exists = any(
         rc.metadata.name == 'nvidia' for rc in runtime_classes.items)
     return nvidia_exists
+
+
+def check_secret_exists(secret_name: str, namespace: str) -> bool:
+    """Checks if a secret exists in a namespace
+
+    Args:
+        secret_name: Name of secret to check
+        namespace: Namespace to check
+    """
+
+    try:
+        kubernetes.core_api().read_namespaced_secret(
+            secret_name, namespace, _request_timeout=kubernetes.API_TIMEOUT)
+    except kubernetes.api_exception() as e:
+        if e.status == 404:
+            return False
+        raise
+    else:
+        return True

sky/skylet/providers/kubernetes/node_provider.py

@@ -427,7 +427,7 @@ def _setup_ssh_in_pods(self, new_nodes):
              '$(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) cat /etc/secret-volume/ssh-publickey* > ~/.ssh/authorized_keys; '
              '$(prefix_cmd) service ssh restart')
         ]
 

sky/templates/kubernetes-ssh-jump.yml.j2

@@ -26,7 +26,7 @@ pod_spec:
             lifecycle:
                 postStart:
                     exec:
-                        command: ["/bin/bash", "-c", "mkdir -p ~/.ssh && cp /etc/secret-volume/ssh-publickey ~/.ssh/authorized_keys && sudo service ssh restart"]
+                        command: ["/bin/bash", "-c", "mkdir -p ~/.ssh && cat /etc/secret-volume/ssh-publickey* > ~/.ssh/authorized_keys && sudo service ssh restart"]
             env:
               - name: MY_POD_NAME
                 valueFrom:
@@ -74,7 +74,7 @@ role:
     rules:
       - apiGroups: [""]
         resources: ["pods", "pods/status", "pods/exec", "services"]
-        verbs: ["get", "list", "create", "delete"]    
+        verbs: ["get", "list", "create", "delete"]
 role_binding:
     apiVersion: rbac.authorization.k8s.io/v1
     kind: RoleBinding