[k8s_cloud_beta1] Adding support for ssh using kubectl port-forward to access k8s instance

sky/adaptors/kubernetes.py

@@ -11,6 +11,12 @@
 kubernetes = None
 urllib3 = None
 
+LOCAL_PORT_FOR_PORT_FORWARD = 23100
+PORT_FORWARD_PROXY_CMD_TEMPLATE = \
+    'kubernetes-port-forward-proxy-command.yaml.j2'
+PORT_FORWARD_PROXY_CMD_PATH = '~/.sky/port-forward-proxy-cmd.sh'
+KUBE_CONFIG_DEFAULT_PATH = '~/.kube/config'
+
 _configured = False
 _core_api = None
 _auth_api = None

sky/authentication.py

@@ -37,8 +37,11 @@
 
 from sky import clouds
 from sky import sky_logging
+from sky import skypilot_config
 from sky.adaptors import gcp
 from sky.adaptors import ibm
+from sky.adaptors import kubernetes
+from sky.backends import backend_utils
 from sky.skylet.providers.lambda_cloud import lambda_utils
 from sky.utils import common_utils
 from sky.utils import kubernetes_utils
@@ -377,9 +380,57 @@ def setup_scp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     return _replace_ssh_info_in_config(config, public_key)
 
 
+def _get_kubernetes_proxy_command(ingress: int, ipaddress: str,
+                                  ssh_setup_mode: str):
+    """ returns Proxycommand to use when establishing ssh connection
+    to the k8s instance through the jump pod.
+
+    Args:
+        ingress: int; the port number host machine is listening to
+        ipaddress: str; ip address of the host machine
+        ssh_setup_mode: str; networking mode for ssh session. It is either
+            'nodeport' or 'port-forward'
+    """
+    if ssh_setup_mode == 'nodeport':
+        proxy_command = (f'ssh -tt -i {PRIVATE_SSH_KEY_PATH} '
+                         '-o StrictHostKeyChecking=no '
+                         '-o UserKnownHostsFile=/dev/null '
+                         f'-o IdentitiesOnly=yes -p {ingress} '
+                         f'-W %h:%p sky@{ipaddress}')
+    # Setting kubectl port-forward/socat to establish ssh session using
+    # ClusterIP service to disallow any ports opened
+    else:
+        ssh_jump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
+        kube_config_path = os.path.expanduser(
+            kubernetes.KUBE_CONFIG_DEFAULT_PATH)
+        vars_to_fill = {
+            'ssh_jump_name': ssh_jump_name,
+            'ipaddress': ipaddress,
+            'local_port': ingress,
+            'kube_config_path': kube_config_path
+        }
+        port_forward_proxy_cmd_path = os.path.expanduser(
+            kubernetes.PORT_FORWARD_PROXY_CMD_PATH)
+        backend_utils.fill_template(kubernetes.PORT_FORWARD_PROXY_CMD_TEMPLATE,
+                                    vars_to_fill,
+                                    output_path=port_forward_proxy_cmd_path)
+        os.chmod(port_forward_proxy_cmd_path,
+                 os.stat(port_forward_proxy_cmd_path).st_mode | 0o111)
+        proxy_command = (f'ssh -tt -i {PRIVATE_SSH_KEY_PATH} '
+                         f'-o ProxyCommand=\'{port_forward_proxy_cmd_path}\' '
+                         '-o StrictHostKeyChecking=no '
+                         '-o UserKnownHostsFile=/dev/null '
+                         f'-o IdentitiesOnly=yes -p {ingress} '
+                         f'-W %h:%p sky@{ipaddress}')
+    return proxy_command
+
+
 def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    # Default ssh session is established with kubectl port-forwarding with
+    # ClusterIP service
+    ssh_setup_mode = skypilot_config.get_nested(('kubernetes', 'networking'),
+                                                'port-forward').lower()
     get_or_generate_keys()
-
     # Run kubectl command to add the public key to the cluster.
     public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
     key_label = clouds.Kubernetes.SKY_SSH_KEY_SECRET_NAME
@@ -407,19 +458,22 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     sshjump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
     sshjump_image = clouds.Kubernetes.IMAGE_CPU
     namespace = kubernetes_utils.get_current_kube_config_context_namespace()
+    ssh_jump_ip = clouds.Kubernetes.get_external_ip()
+    if ssh_setup_mode == 'nodeport':
+        service_type = 'NodePort'
+    # If ssh connection is establisehd with kubectl port-forward, the jump pod
+    # will run on ClusterIP service. This enables to establish ssh session
+    # without opening any ports on the Kubernetes cluster.
+    else:
+        service_type = 'ClusterIP'
 
     kubernetes_utils.setup_sshjump(sshjump_name, sshjump_image, key_label,
-                                   namespace)
-
-    ssh_jump_port = clouds.Kubernetes.get_port(sshjump_name)
+                                   namespace, service_type)
     ssh_jump_ip = clouds.Kubernetes.get_external_ip()
-
-    ssh_jump_proxy_command = (f'ssh -tt -i {PRIVATE_SSH_KEY_PATH} '
-                              '-o StrictHostKeyChecking=no '
-                              '-o UserKnownHostsFile=/dev/null '
-                              '-o IdentitiesOnly=yes '
-                              f'-p {ssh_jump_port} -W %h:%p sky@{ssh_jump_ip}')
-
-    config['auth']['ssh_proxy_command'] = ssh_jump_proxy_command
-
+    if ssh_setup_mode == 'nodeport':
+        ssh_jump_port = clouds.Kubernetes.get_port(sshjump_name)
+    else:
+        ssh_jump_port = kubernetes.LOCAL_PORT_FOR_PORT_FORWARD
+    config['auth']['ssh_proxy_command'] = _get_kubernetes_proxy_command(
+        ssh_jump_port, ssh_jump_ip, ssh_setup_mode)
     return config

sky/backends/backend_utils.py

@@ -1341,7 +1341,7 @@ def wait_until_ray_cluster_ready(
 
 def ssh_credential_from_yaml(cluster_yaml: str,
                              docker_user: Optional[str] = None
-                            ) -> Dict[str, str]:
+                            ) -> Dict[str, Any]:
     """Returns ssh_user, ssh_private_key and ssh_control name."""
     config = common_utils.read_yaml(cluster_yaml)
     auth_section = config['auth']
@@ -1357,6 +1357,10 @@ def ssh_credential_from_yaml(cluster_yaml: str,
     }
     if docker_user is not None:
         credentials['docker_user'] = docker_user
+    ssh_provider_module = config['provider']['module']
+    # If we are running ssh command on kubernetes node.
+    if 'kubernetes' in ssh_provider_module:
+        credentials['proxy_to_k8s'] = True
     return credentials
 
 

sky/backends/cloud_vm_ray_backend.py

@@ -2824,7 +2824,6 @@ def _sync_workdir(self, handle: CloudVmRayResourceHandle,
                 f'down rsync.{style.RESET_ALL}')
 
         log_path = os.path.join(self.log_dir, 'workdir_sync.log')
-
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
             handle.cluster_yaml, handle.docker_user)
 
@@ -4068,7 +4067,6 @@ def _set_tpu_name(self, handle: CloudVmRayResourceHandle,
         assert ip_list is not None, 'external_ips is not cached in handle'
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
             handle.cluster_yaml, handle.docker_user)
-
         runners = command_runner.SSHCommandRunner.make_runner_list(
             ip_list, port_list=None, **ssh_credentials)
 

sky/templates/kubernetes-port-forward-proxy-command.yaml.j2

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+set -uo pipefail
+
+# Specifies which config file kubectl should use.
+export KUBECONFIG="{{ kube_config_path }}"
+
+# If the port-forward command is already running on the jump pod, kill it
+pgrep -f "kubectl port-forward svc/{{ ssh_jump_name }} {{ local_port }}:22" > /dev/null
+if [ $? -eq 0 ]; then
+    pkill -f "kubectl port-forward svc/{{ ssh_jump_name }} {{ local_port }}:22"
+fi
+
+# Establishes connection between local port and the ssh jump pod
+kubectl port-forward svc/{{ ssh_jump_name }} {{ local_port }}:22 &
+K8S_PID=$!
+trap "kill -9 $K8S_PID" EXIT
+
+# checks if a conection to local_port of ipaddress is established 
+while ! nc -z {{ ipaddress }} {{ local_port }}; do
+    sleep 0.1
+done
+
+# Estalbishes two directional byte streams to handle stdin/stdout between
+# terminal and the jump pod
+socat - tcp:{{ ipaddress }}:{{ local_port }}

sky/templates/kubernetes-sshjump.yml.j2

@@ -50,14 +50,13 @@ service_spec:
         name: {{ name }}
         app: skypilot
     spec:
-        type: NodePort
+        type: {{ service_type }}
         selector:
             component: {{ name }}
         ports:
           - protocol: TCP
             port: 22
             targetPort: 22
-
 # The following ServiceAccount/Role/RoleBinding sets up an RBAC for life cycle
 # management of the jump pod/service
 service_account:

sky/utils/command_runner.py

@@ -42,13 +42,16 @@ def _ssh_control_path(ssh_control_filename: Optional[str]) -> Optional[str]:
     return path
 
 
-def ssh_options_list(ssh_private_key: Optional[str],
-                     ssh_control_name: Optional[str],
-                     *,
-                     ssh_proxy_command: Optional[str] = None,
-                     docker_ssh_proxy_command: Optional[str] = None,
-                     timeout: int = 30,
-                     port: int = 22) -> List[str]:
+def ssh_options_list(
+    ssh_private_key: Optional[str],
+    ssh_control_name: Optional[str],
+    proxy_to_k8s: Optional[bool] = False,
+    *,
+    ssh_proxy_command: Optional[str] = None,
+    docker_ssh_proxy_command: Optional[str] = None,
+    timeout: int = 30,
+    port: int = 22,
+) -> List[str]:
     """Returns a list of sane options for 'ssh'."""
     # Forked from Ray SSHOptions:
     # https://github.com/ray-project/ray/blob/master/python/ray/autoscaler/_private/command_runner.py
@@ -79,7 +82,12 @@ def ssh_options_list(ssh_private_key: Optional[str],
     }
     # SSH Control will have a severe delay when using docker_ssh_proxy_command.
     # TODO(tian): Investigate why.
-    if ssh_control_name is not None and docker_ssh_proxy_command is None:
+    # k8s instances are accessed with an ssh session using Proxycommand. The
+    # process running Proxycommand is kept running as long as the ssh session
+    # is running and the ControlMaster keeps the session, which results in
+    # 'ControlPersist' number of seconds delay per ssh commands ran.
+    if ssh_control_name is not None and docker_ssh_proxy_command is None \
+        and not proxy_to_k8s:
         arg_dict.update({
             # Control path: important optimization as we do multiple ssh in one
             # sky.launch().
@@ -136,6 +144,7 @@ def __init__(
         ssh_proxy_command: Optional[str] = None,
         port: int = 22,
         docker_user: Optional[str] = None,
+        proxy_to_k8s: Optional[bool] = False,
     ):
         """Initialize SSHCommandRunner.
 
@@ -159,12 +168,15 @@ def __init__(
             docker_user: The docker user to use for ssh. If specified, the
                 command will be run inside a docker container which have a ssh
                 server running at port sky.skylet.constants.DEFAULT_DOCKER_PORT.
+            proxy_to_k8s: bool; specifies either or not the ssh command will be
+                ran on k8s instance through a jump pod with proxy command.
         """
         self.ssh_private_key = ssh_private_key
         self.ssh_control_name = (
             None if ssh_control_name is None else hashlib.md5(
                 ssh_control_name.encode()).hexdigest()[:_HASH_MAX_LENGTH])
         self._ssh_proxy_command = ssh_proxy_command
+        self.proxy_to_k8s = proxy_to_k8s
         if docker_user is not None:
             assert port is None or port == 22, (
                 f'port must be None or 22 for docker_user, got {port}.')
@@ -190,6 +202,7 @@ def make_runner_list(
         ssh_private_key: str,
         ssh_control_name: Optional[str] = None,
         ssh_proxy_command: Optional[str] = None,
+        proxy_to_k8s: Optional[bool] = False,
         port_list: Optional[List[int]] = None,
         docker_user: Optional[str] = None,
     ) -> List['SSHCommandRunner']:
@@ -198,7 +211,7 @@ def make_runner_list(
             port_list = [22] * len(ip_list)
         return [
             SSHCommandRunner(ip, ssh_user, ssh_private_key, ssh_control_name,
-                             ssh_proxy_command, port, docker_user)
+                             ssh_proxy_command, port, docker_user, proxy_to_k8s)
             for ip, port in zip(ip_list, port_list)
         ]
 
@@ -228,7 +241,7 @@ def _ssh_base_command(self, *, ssh_mode: SshMode,
             ssh_proxy_command=self._ssh_proxy_command,
             docker_ssh_proxy_command=docker_ssh_proxy_command,
             port=self.port,
-        ) + [f'{self.ssh_user}@{self.ip}']
+            proxy_to_k8s=self.proxy_to_k8s) + [f'{self.ssh_user}@{self.ip}']
 
     def run(
             self,
@@ -382,13 +395,12 @@ def rsync(
         else:
             docker_ssh_proxy_command = None
         ssh_options = ' '.join(
-            ssh_options_list(
-                self.ssh_private_key,
-                self.ssh_control_name,
-                ssh_proxy_command=self._ssh_proxy_command,
-                docker_ssh_proxy_command=docker_ssh_proxy_command,
-                port=self.port,
-            ))
+            ssh_options_list(self.ssh_private_key,
+                             self.ssh_control_name,
+                             ssh_proxy_command=self._ssh_proxy_command,
+                             docker_ssh_proxy_command=docker_ssh_proxy_command,
+                             port=self.port,
+                             proxy_to_k8s=self.proxy_to_k8s))
         rsync_command.append(f'-e "ssh {ssh_options}"')
         # To support spaces in the path, we need to quote source and target.
         # rsync doesn't support '~' in a quoted local path, but it is ok to

sky/utils/command_runner.pyi

@@ -20,10 +20,13 @@ RSYNC_FILTER_OPTION: str
 RSYNC_EXCLUDE_OPTION: str
 
 
-def ssh_options_list(ssh_private_key: Optional[str],
-                     ssh_control_name: Optional[str],
-                     *,
-                     timeout: int = ...) -> List[str]:
+def ssh_options_list(
+    ssh_private_key: Optional[str],
+    ssh_control_name: Optional[str],
+    proxy_to_k8s: Optional[bool] = False,
+    *,
+    timeout: int = ...,
+) -> List[str]:
     ...
 
 
@@ -40,14 +43,18 @@ class SSHCommandRunner:
     ssh_control_name: Optional[str]
     docker_user: str
     port: int
+    proxy_to_k8s: Optional[bool]
 
-    def __init__(self,
-                 ip: str,
-                 ssh_user: str,
-                 ssh_private_key: str,
-                 ssh_control_name: Optional[str] = ...,
-                 port: str = ...,
-                 docker_user: Optional[str] = ...) -> None:
+    def __init__(
+        self,
+        ip: str,
+        ssh_user: str,
+        ssh_private_key: str,
+        ssh_control_name: Optional[str] = ...,
+        port: str = ...,
+        docker_user: Optional[str] = ...,
+        proxy_to_k8s: Optional[bool] = ...,
+    ) -> None:
         ...
 
     @staticmethod
@@ -59,6 +66,7 @@ class SSHCommandRunner:
         ssh_proxy_command: Optional[str] = ...,
         port_list: Optional[List[int]] = ...,
         docker_user: Optional[str] = ...,
+        proxy_to_k8s: Optional[bool] = ...,
     ) -> List['SSHCommandRunner']:
         ...