[K8s] Zero config networking for Kubernetes

Dockerfile_k8s

@@ -45,6 +45,9 @@ 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
+
 # Set WORKDIR and initialize conda for sky user
 WORKDIR /home/sky
 RUN conda init

Dockerfile_k8s_gpu

@@ -53,6 +53,9 @@ 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
+
 # Set WORKDIR and initialize conda for sky user
 WORKDIR /home/sky
 RUN conda init

sky/authentication.py

@@ -37,10 +37,12 @@
 
 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.skylet.providers.lambda_cloud import lambda_utils
 from sky.utils import common_utils
+from sky.utils import kubernetes_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
@@ -377,6 +379,12 @@ def setup_scp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
 
 def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    # Default ssh session is established with kubectl port-forwarding with
+    # ClusterIP service
+    nodeport_mode = kubernetes_utils.KubernetesNetworkingMode.NODEPORT
+    port_forward_mode = kubernetes_utils.KubernetesNetworkingMode.PORTFORWARD
+    ssh_setup_mode = skypilot_config.get_nested(('kubernetes', 'networking'),
+                                                port_forward_mode.value)
     get_or_generate_keys()
 
     # Run kubectl command to add the public key to the cluster.
@@ -403,4 +411,36 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
             logger.error(suffix)
             raise
 
+    ssh_jump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
+    if ssh_setup_mode.lower() == nodeport_mode.value:
+        network_mode = nodeport_mode
+        service_type = kubernetes_utils.KubernetesServiceType.NODEPORT
+
+    elif ssh_setup_mode.lower() == port_forward_mode.value:
+        kubernetes_utils.check_port_forward_mode_dependencies()
+        network_mode = port_forward_mode
+        # Using `kubectl port-forward` creates a direct tunnel to jump pod and
+        # does not require opening any ports on Kubernetes nodes. As a result,
+        # the service can be a simple ClusterIP service which we access with
+        # `kubectl port-forward`.
+        service_type = kubernetes_utils.KubernetesServiceType.CLUSTERIP
+    else:
+        raise ValueError(f'Unsupported kubernetes networking mode: '
+                         f'{ssh_setup_mode}. The mode has to be either '
+                         f'\'{port_forward_mode.value}\' or '
+                         f'\'{nodeport_mode.value}\'. '
+                         'Please check: ~/.sky/config.yaml')
+    # 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_sshjump_svc(ssh_jump_name, namespace, service_type)
+
+    ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(
+        PRIVATE_SSH_KEY_PATH, ssh_jump_name, network_mode, namespace,
+        clouds.Kubernetes.PORT_FORWARD_PROXY_CMD_PATH,
+        clouds.Kubernetes.PORT_FORWARD_PROXY_CMD_TEMPLATE)
+
+    config['auth']['ssh_proxy_command'] = ssh_proxy_cmd
+
     return config

sky/backends/backend_utils.py

@@ -1353,7 +1353,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']
@@ -1369,6 +1369,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['disable_control_master'] = True
     return credentials
 
 

sky/backends/cloud_vm_ray_backend.py

@@ -2319,23 +2319,12 @@ def _update_cluster_region(self):
         self.launched_resources = self.launched_resources.copy(region=region)
 
     def update_ssh_ports(self, max_attempts: int = 1) -> None:
-        """Updates the cluster SSH ports cached in the handle."""
-        # TODO(romilb): Replace this with a call to the cloud class to get ports
-        # Use port 22 for everything except Kubernetes
-        if not isinstance(self.launched_resources.cloud, clouds.Kubernetes):
-            head_ssh_port = 22
-        else:
-            svc_name = f'{self.cluster_name_on_cloud}-ray-head-ssh'
-            retry_cnt = 0
-            while True:
-                try:
-                    head_ssh_port = clouds.Kubernetes.get_port(svc_name)
-                    break
-                except Exception:  # pylint: disable=broad-except
-                    retry_cnt += 1
-                    if retry_cnt >= max_attempts:
-                        raise
-            # TODO(romilb): Multinode doesn't work with Kubernetes yet.
+        """Fetches and sets the SSH ports for the cluster nodes.
+
+        Use this method to use any cloud-specific port fetching logic.
+        """
+        del max_attempts  # Unused.
+        head_ssh_port = 22
         self.stable_ssh_ports = ([head_ssh_port] + [22] *
                                  (self.num_node_ips - 1))
 
@@ -3011,37 +3000,12 @@ def _sync_file_mounts(
         self._execute_file_mounts(handle, all_file_mounts)
         self._execute_storage_mounts(handle, storage_mounts)
 
-    def _update_envs_for_k8s(self, handle: CloudVmRayResourceHandle,
-                             task: task_lib.Task) -> None:
-        """Update envs with env vars from Kubernetes if cloud is Kubernetes.
-
-        Kubernetes automatically populates containers with critical environment
-        variables, such as those for discovering services running in the
-        cluster and CUDA/nvidia environment variables. We need to update task
-        environment variables with these env vars. This is needed for GPU
-        support and service discovery.
-
-        See https://github.com/skypilot-org/skypilot/issues/2287 for
-        more details.
-        """
-        if isinstance(handle.launched_resources.cloud, clouds.Kubernetes):
-            temp_envs = copy.deepcopy(task.envs)
-            cloud_env_vars = handle.launched_resources.cloud.query_env_vars(
-                handle.cluster_name_on_cloud)
-            task.update_envs(cloud_env_vars)
-
-            # Re update the envs with the original envs to give priority to
-            # the original envs.
-            task.update_envs(temp_envs)
-
     def _setup(self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
                detach_setup: bool) -> None:
         start = time.time()
         style = colorama.Style
         fore = colorama.Fore
 
-        self._update_envs_for_k8s(handle, task)
-
         if task.setup is None:
             return
 
@@ -3350,7 +3314,6 @@ def _execute(
         # Check the task resources vs the cluster resources. Since `sky exec`
         # will not run the provision and _check_existing_cluster
         self.check_resources_fit_cluster(handle, task)
-        self._update_envs_for_k8s(handle, task)
 
         resources_str = backend_utils.get_task_resources_str(task)
 

sky/clouds/kubernetes.py

@@ -19,15 +19,18 @@
 
 logger = sky_logging.init_logger(__name__)
 
-_CREDENTIAL_PATH = '~/.kube/config'
+CREDENTIAL_PATH = '~/.kube/config'
 
 
 @clouds.CLOUD_REGISTRY.register
 class Kubernetes(clouds.Cloud):
     """Kubernetes."""
 
     SKY_SSH_KEY_SECRET_NAME = f'sky-ssh-{common_utils.get_user_hash()}'
-
+    SKY_SSH_JUMP_NAME = f'sky-sshjump-{common_utils.get_user_hash()}'
+    PORT_FORWARD_PROXY_CMD_TEMPLATE = \
+        'kubernetes-port-forward-proxy-command.sh.j2'
+    PORT_FORWARD_PROXY_CMD_PATH = '~/.sky/port-forward-proxy-cmd.sh'
     # Timeout for resource provisioning. This timeout determines how long to
     # wait for pod to be in pending status before giving up.
     # Larger timeout may be required for autoscaling clusters, since autoscaler
@@ -223,6 +226,9 @@ def make_deploy_resources_variables(
             'k8s_ssh_key_secret_name': self.SKY_SSH_KEY_SECRET_NAME,
             'k8s_acc_label_key': k8s_acc_label_key,
             'k8s_acc_label_value': k8s_acc_label_value,
+            'k8s_sshjump_name': self.SKY_SSH_JUMP_NAME,
+            # TODO(romilb): Create a lightweight image for SSH jump host
+            'k8s_sshjump_image': self.IMAGE_CPU,
             # TODO(romilb): Allow user to specify custom images
             'image_id': image,
         }
@@ -292,7 +298,7 @@ def _make(instance_list):
 
     @classmethod
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
-        if os.path.exists(os.path.expanduser(_CREDENTIAL_PATH)):
+        if os.path.exists(os.path.expanduser(CREDENTIAL_PATH)):
             # Test using python API
             try:
                 return kubernetes_utils.check_credentials()
@@ -301,10 +307,10 @@ def check_credentials(cls) -> Tuple[bool, Optional[str]]:
                         f'{common_utils.format_exception(e)}')
         else:
             return (False, 'Credentials not found - '
-                    f'check if {_CREDENTIAL_PATH} exists.')
+                    f'check if {CREDENTIAL_PATH} exists.')
 
     def get_credential_file_mounts(self) -> Dict[str, str]:
-        return {_CREDENTIAL_PATH: _CREDENTIAL_PATH}
+        return {CREDENTIAL_PATH: CREDENTIAL_PATH}
 
     def instance_type_exists(self, instance_type: str) -> bool:
         return kubernetes_utils.KubernetesInstanceType.is_valid_instance_type(
@@ -362,39 +368,6 @@ def query_status(cls, name: str, tag_filters: Dict[str, str],
         # If pods are not found, we don't add them to the return list
         return cluster_status
 
-    @classmethod
-    def query_env_vars(cls, name: str) -> Dict[str, str]:
-        namespace = kubernetes_utils.get_current_kube_config_context_namespace()
-        pod = kubernetes.core_api().list_namespaced_pod(
-            namespace,
-            label_selector=f'skypilot-cluster={name},ray-node-type=head'
-        ).items[0]
-        response = kubernetes.stream()(
-            kubernetes.core_api().connect_get_namespaced_pod_exec,
-            pod.metadata.name,
-            namespace,
-            command=['env'],
-            stderr=True,
-            stdin=False,
-            stdout=True,
-            tty=False,
-            _request_timeout=kubernetes.API_TIMEOUT)
-        # Split response by newline and filter lines containing '='
-        raw_lines = response.split('\n')
-        filtered_lines = [line for line in raw_lines if '=' in line]
-
-        # Split each line at the first '=' occurrence
-        lines = [line.split('=', 1) for line in filtered_lines]
-
-        # Construct the dictionary using only valid environment variable names
-        env_vars = {}
-        for line in lines:
-            key = line[0]
-            if common_utils.is_valid_env_var(key):
-                env_vars[key] = line[1]
-
-        return env_vars
-
     @classmethod
     def get_current_user_identity(cls) -> Optional[List[str]]:
         k8s = kubernetes.get_kubernetes()

sky/skylet/providers/kubernetes/config.py

@@ -55,6 +55,8 @@ def bootstrap_kubernetes(config: Dict[str, Any]) -> Dict[str, Any]:
 
     _configure_services(namespace, config['provider'])
 
+    config = _configure_ssh_jump(namespace, config)
+
     if not config['provider'].get('_operator'):
         # These steps are unecessary when using the Operator.
         _configure_autoscaler_service_account(namespace, config['provider'])
@@ -257,6 +259,40 @@ def _configure_autoscaler_role_binding(namespace: str,
     logger.info(log_prefix + created_msg(binding_field, name))
 
 
+def _configure_ssh_jump(namespace, config):
+    """Creates a SSH jump pod to connect to the cluster.
+
+    Also updates config['auth']['ssh_proxy_command'] to use the newly created
+    jump pod.
+    """
+    pod_cfg = config['available_node_types']['ray_head_default']['node_config']
+
+    sshjump_name = pod_cfg['metadata']['labels']['skypilot-sshjump']
+    sshjump_image = config['provider']['sshjump_image']
+
+    volumes = pod_cfg['spec']['volumes']
+    # find 'secret-volume' and get the secret name
+    secret_volume = next(filter(lambda x: x['name'] == 'secret-volume',
+                                volumes))
+    ssh_key_secret_name = secret_volume['secret']['secretName']
+
+    # TODO(romilb): We currently split SSH jump pod and svc creation. Service
+    #  is first created in authentication.py::setup_kubernetes_authentication
+    #  and then SSH jump pod creation happens here. This is because we need to
+    #  set the ssh_proxy_command in the ray YAML before we pass it to the
+    #  autoscaler. If in the future if we can write the ssh_proxy_command to the
+    #  cluster yaml through this method, then we should move the service
+    #  creation here.
+
+    # TODO(romilb): We should add a check here to make sure the service is up
+    #  and available before we create the SSH jump pod. If for any reason the
+    #  service is missing, we should raise an error.
+
+    kubernetes_utils.setup_sshjump_pod(sshjump_name, sshjump_image,
+                                       ssh_key_secret_name, namespace)
+    return config
+
+
 def _configure_services(namespace: str, provider_config: Dict[str,
                                                               Any]) -> None:
     service_field = 'services'