skypilot-nightly 1.0.0.dev20240910__py3-none-any.whl → 1.0.0.dev20240911__py3-none-any.whl

sky/provision/kubernetes/utils.py

@@ -58,7 +58,11 @@ KIND_CONTEXT_NAME = 'kind-skypilot'  # Context name used by sky local up
 
 # Port-forward proxy command constants
 PORT_FORWARD_PROXY_CMD_TEMPLATE = 'kubernetes-port-forward-proxy-command.sh'
-PORT_FORWARD_PROXY_CMD_PATH = '~/.sky/kubernetes-port-forward-proxy-command.sh'
+# We add a version suffix to the port-forward proxy command to ensure backward
+# compatibility and avoid overwriting the older version.
+PORT_FORWARD_PROXY_CMD_VERSION = 2
+PORT_FORWARD_PROXY_CMD_PATH = ('~/.sky/kubernetes-port-forward-proxy-command-'
+                               f'v{PORT_FORWARD_PROXY_CMD_VERSION}.sh')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -355,6 +359,10 @@ def get_kubernetes_nodes() -> List[Any]:
 
 
 def get_kubernetes_pods() -> List[Any]:
+    """Gets the kubernetes pods in the current namespace and current context.
+
+    Used for computing cluster resource usage.
+    """
     try:
         ns = get_current_kube_config_context_namespace()
         pods = kubernetes.core_api().list_namespaced_pod(
@@ -572,37 +580,39 @@ def get_gpu_label_key_value(acc_type: str, check_mode=False) -> Tuple[str, str]:
                 f'to set up GPUs.{suffix}')
 
 
-def get_head_ssh_port(cluster_name: str, namespace: str) -> int:
+def get_head_ssh_port(cluster_name: str, namespace: str,
+                      context: Optional[str]) -> int:
     svc_name = f'{cluster_name}-head-ssh'
-    return get_port(svc_name, namespace)
+    return get_port(svc_name, namespace, context)
 
 
-def get_port(svc_name: str, namespace: str) -> int:
+def get_port(svc_name: str, namespace: str, context: Optional[str]) -> int:
     """Gets the nodeport of the specified service.
 
     Args:
         svc_name (str): Name of the kubernetes service. Note that this may be
             different from the cluster name.
         namespace (str): Kubernetes namespace to look for the service in.
+        context (str): Kubernetes context to use.
     """
-    head_service = kubernetes.core_api().read_namespaced_service(
+    head_service = kubernetes.core_api(context).read_namespaced_service(
         svc_name, namespace)
     return head_service.spec.ports[0].node_port
 
 
-def get_external_ip(
-        network_mode: Optional[kubernetes_enums.KubernetesNetworkingMode]):
+def get_external_ip(network_mode: Optional[
+    kubernetes_enums.KubernetesNetworkingMode], context: Optional[str]) -> str:
     if network_mode == kubernetes_enums.KubernetesNetworkingMode.PORTFORWARD:
         return '127.0.0.1'
     # Return the IP address of the first node with an external IP
-    nodes = kubernetes.core_api().list_node().items
+    nodes = kubernetes.core_api(context).list_node().items
     for node in nodes:
         if node.status.addresses:
             for address in node.status.addresses:
                 if address.type == 'ExternalIP':
                     return address.address
     # If no external IP is found, use the API server IP
-    api_host = kubernetes.core_api().api_client.configuration.host
+    api_host = kubernetes.core_api(context).api_client.configuration.host
     parsed_url = urlparse(api_host)
     return parsed_url.hostname
 
@@ -620,7 +630,9 @@ def check_credentials(timeout: int = kubernetes.API_TIMEOUT) -> \
     """
     try:
         ns = get_current_kube_config_context_namespace()
-        kubernetes.core_api().list_namespaced_pod(ns, _request_timeout=timeout)
+        context = get_current_kube_config_context_name()
+        kubernetes.core_api(context).list_namespaced_pod(
+            ns, _request_timeout=timeout)
     except ImportError:
         # TODO(romilb): Update these error strs to also include link to docs
         #  when docs are ready.
@@ -948,14 +960,14 @@ def construct_ssh_jump_command(
         proxy_cmd_path = os.path.expanduser(proxy_cmd_path)
         # adding execution permission to the proxy command script
         os.chmod(proxy_cmd_path, os.stat(proxy_cmd_path).st_mode | 0o111)
-        kube_context_flag = f' {current_kube_context}' if (current_kube_context
-                                                           is not None) else ''
-        kube_namespace_flag = f' {current_kube_namespace}' if (
+        kube_context_flag = f'-c {current_kube_context} ' if (
+            current_kube_context is not None) else ''
+        kube_namespace_flag = f'-n {current_kube_namespace} ' if (
             current_kube_namespace is not None) else ''
         ssh_jump_proxy_command += (f' -o ProxyCommand=\'{proxy_cmd_path} '
-                                   f'{proxy_cmd_target_pod}'
                                    f'{kube_context_flag}'
-                                   f'{kube_namespace_flag}\'')
+                                   f'{kube_namespace_flag}'
+                                   f'{proxy_cmd_target_pod}\'')
     return ssh_jump_proxy_command
 
 
@@ -963,7 +975,8 @@ def get_ssh_proxy_command(
         k8s_ssh_target: str,
         network_mode: kubernetes_enums.KubernetesNetworkingMode,
         private_key_path: Optional[str] = None,
-        namespace: Optional[str] = None) -> str:
+        namespace: Optional[str] = None,
+        context: Optional[str] = None) -> str:
     """Generates the SSH proxy command to connect to the pod.
 
     Uses a jump pod if the network mode is NODEPORT, and direct port-forwarding
@@ -1011,11 +1024,11 @@ def get_ssh_proxy_command(
             Required for NODEPORT networking mode.
     """
     # Fetch IP to connect to for the jump svc
-    ssh_jump_ip = get_external_ip(network_mode)
+    ssh_jump_ip = get_external_ip(network_mode, context)
     assert private_key_path is not None, 'Private key path must be provided'
     if network_mode == kubernetes_enums.KubernetesNetworkingMode.NODEPORT:
         assert namespace is not None, 'Namespace must be provided for NodePort'
-        ssh_jump_port = get_port(k8s_ssh_target, namespace)
+        ssh_jump_port = get_port(k8s_ssh_target, namespace, context)
         ssh_jump_proxy_command = construct_ssh_jump_command(
             private_key_path, ssh_jump_ip, ssh_jump_port=ssh_jump_port)
     else:
@@ -1061,7 +1074,7 @@ def create_proxy_command_script() -> str:
     return port_fwd_proxy_cmd_path
 
 
-def setup_ssh_jump_svc(ssh_jump_name: str, namespace: str,
+def setup_ssh_jump_svc(ssh_jump_name: str, namespace: str, context: str,
                        service_type: kubernetes_enums.KubernetesServiceType):
     """Sets up Kubernetes service resource to access for SSH jump pod.
 
@@ -1083,13 +1096,14 @@ def setup_ssh_jump_svc(ssh_jump_name: str, namespace: str,
 
     # Create service
     try:
-        kubernetes.core_api().create_namespaced_service(namespace,
-                                                        content['service_spec'])
+        kubernetes.core_api(context).create_namespaced_service(
+            namespace, content['service_spec'])
     except kubernetes.api_exception() as e:
         # SSH Jump Pod service already exists.
         if e.status == 409:
-            ssh_jump_service = kubernetes.core_api().read_namespaced_service(
-                name=ssh_jump_name, namespace=namespace)
+            ssh_jump_service = kubernetes.core_api(
+                context).read_namespaced_service(name=ssh_jump_name,
+                                                 namespace=namespace)
             curr_svc_type = ssh_jump_service.spec.type
             if service_type.value == curr_svc_type:
                 # If the currently existing SSH Jump service's type is identical
@@ -1101,9 +1115,9 @@ def setup_ssh_jump_svc(ssh_jump_name: str, namespace: str,
                 # If a different type of service type for SSH Jump pod compared
                 # to user's configuration for networking mode exists, we remove
                 # existing servie to create a new one following user's config
-                kubernetes.core_api().delete_namespaced_service(
+                kubernetes.core_api(context).delete_namespaced_service(
                     name=ssh_jump_name, namespace=namespace)
-                kubernetes.core_api().create_namespaced_service(
+                kubernetes.core_api(context).create_namespaced_service(
                     namespace, content['service_spec'])
                 port_forward_mode = (
                     kubernetes_enums.KubernetesNetworkingMode.PORTFORWARD.value)
@@ -1132,7 +1146,7 @@ def setup_ssh_jump_svc(ssh_jump_name: str, namespace: str,
 
 
 def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
-                       ssh_key_secret: str, namespace: str):
+                       ssh_key_secret: str, namespace: str, context: str):
     """Sets up Kubernetes RBAC and pod for SSH jump host.
 
     Our Kubernetes implementation uses a SSH jump pod to reach SkyPilot clusters
@@ -1161,7 +1175,7 @@ def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
 
     # ServiceAccount
     try:
-        kubernetes.core_api().create_namespaced_service_account(
+        kubernetes.core_api(context).create_namespaced_service_account(
             namespace, content['service_account'])
     except kubernetes.api_exception() as e:
         if e.status == 409:
@@ -1174,7 +1188,8 @@ def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
         logger.info('Created SSH Jump ServiceAccount.')
     # Role
     try:
-        kubernetes.auth_api().create_namespaced_role(namespace, content['role'])
+        kubernetes.auth_api(context).create_namespaced_role(
+            namespace, content['role'])
     except kubernetes.api_exception() as e:
         if e.status == 409:
             logger.info(
@@ -1185,7 +1200,7 @@ def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
         logger.info('Created SSH Jump Role.')
     # RoleBinding
     try:
-        kubernetes.auth_api().create_namespaced_role_binding(
+        kubernetes.auth_api(context).create_namespaced_role_binding(
             namespace, content['role_binding'])
     except kubernetes.api_exception() as e:
         if e.status == 409:
@@ -1198,8 +1213,8 @@ def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
         logger.info('Created SSH Jump RoleBinding.')
     # Pod
     try:
-        kubernetes.core_api().create_namespaced_pod(namespace,
-                                                    content['pod_spec'])
+        kubernetes.core_api(context).create_namespaced_pod(
+            namespace, content['pod_spec'])
     except kubernetes.api_exception() as e:
         if e.status == 409:
             logger.info(
@@ -1211,7 +1226,7 @@ def setup_ssh_jump_pod(ssh_jump_name: str, ssh_jump_image: str,
         logger.info(f'Created SSH Jump Host {ssh_jump_name}.')
 
 
-def clean_zombie_ssh_jump_pod(namespace: str, node_id: str):
+def clean_zombie_ssh_jump_pod(namespace: str, context: str, node_id: str):
     """Analyzes SSH jump pod and removes if it is in a bad state
 
     Prevents the existence of a dangling SSH jump pod. This could happen
@@ -1231,7 +1246,8 @@ def clean_zombie_ssh_jump_pod(namespace: str, node_id: str):
 
     # Get the SSH jump pod name from the head pod
     try:
-        pod = kubernetes.core_api().read_namespaced_pod(node_id, namespace)
+        pod = kubernetes.core_api(context).read_namespaced_pod(
+            node_id, namespace)
     except kubernetes.api_exception() as e:
         if e.status == 404:
             logger.warning(f'Failed to get pod {node_id},'
@@ -1240,7 +1256,7 @@ def clean_zombie_ssh_jump_pod(namespace: str, node_id: str):
     else:
         ssh_jump_name = pod.metadata.labels.get('skypilot-ssh-jump')
     try:
-        ssh_jump_pod = kubernetes.core_api().read_namespaced_pod(
+        ssh_jump_pod = kubernetes.core_api(context).read_namespaced_pod(
             ssh_jump_name, namespace)
         cont_ready_cond = find(ssh_jump_pod.status.conditions,
                                lambda c: c.type == 'ContainersReady')
@@ -1251,9 +1267,9 @@ def clean_zombie_ssh_jump_pod(namespace: str, node_id: str):
             # ssh jump pod, lets remove it and the service. Otherwise, main
             # container is ready and its lifecycle management script takes
             # care of the cleaning.
-            kubernetes.core_api().delete_namespaced_pod(ssh_jump_name,
-                                                        namespace)
-            kubernetes.core_api().delete_namespaced_service(
+            kubernetes.core_api(context).delete_namespaced_pod(
+                ssh_jump_name, namespace)
+            kubernetes.core_api(context).delete_namespaced_service(
                 ssh_jump_name, namespace)
     except kubernetes.api_exception() as e:
         # We keep the warning in debug to avoid polluting the `sky launch`
@@ -1265,7 +1281,7 @@ def clean_zombie_ssh_jump_pod(namespace: str, node_id: str):
         # We encountered an issue while checking ssh jump pod. To be on
         # the safe side, lets remove its service so the port is freed
         try:
-            kubernetes.core_api().delete_namespaced_service(
+            kubernetes.core_api(context).delete_namespaced_service(
                 ssh_jump_name, namespace)
         except kubernetes.api_exception():
             pass
@@ -1521,10 +1537,10 @@ def merge_custom_metadata(original_metadata: Dict[str, Any]) -> None:
     merge_dicts(custom_metadata, original_metadata)
 
 
-def check_nvidia_runtime_class() -> bool:
+def check_nvidia_runtime_class(context: Optional[str] = None) -> bool:
     """Checks if the 'nvidia' RuntimeClass exists in the cluster"""
     # Fetch the list of available RuntimeClasses
-    runtime_classes = kubernetes.node_api().list_runtime_class()
+    runtime_classes = kubernetes.node_api(context).list_runtime_class()
 
     # Check if 'nvidia' RuntimeClass exists
     nvidia_exists = any(
@@ -1532,7 +1548,7 @@ def check_nvidia_runtime_class() -> bool:
     return nvidia_exists
 
 
-def check_secret_exists(secret_name: str, namespace: str) -> bool:
+def check_secret_exists(secret_name: str, namespace: str, context: str) -> bool:
     """Checks if a secret exists in a namespace
 
     Args:
@@ -1541,7 +1557,7 @@ def check_secret_exists(secret_name: str, namespace: str) -> bool:
     """
 
     try:
-        kubernetes.core_api().read_namespaced_secret(
+        kubernetes.core_api(context).read_namespaced_secret(
             secret_name, namespace, _request_timeout=kubernetes.API_TIMEOUT)
     except kubernetes.api_exception() as e:
         if e.status == 404:
@@ -1551,17 +1567,18 @@ def check_secret_exists(secret_name: str, namespace: str) -> bool:
         return True
 
 
-def create_namespace(namespace: str) -> None:
+def create_namespace(namespace: str, context: Optional[str]) -> None:
     """Creates a namespace in the cluster.
 
     If the namespace already exists, logs a message and does nothing.
 
     Args:
         namespace: Name of the namespace to create
+        context: Name of the context to use. Can be none to use default context.
     """
     kubernetes_client = kubernetes.kubernetes.client
     try:
-        kubernetes.core_api().read_namespace(namespace)
+        kubernetes.core_api(context).read_namespace(namespace)
     except kubernetes.api_exception() as e:
         if e.status != 404:
             raise
@@ -1572,7 +1589,7 @@ def create_namespace(namespace: str) -> None:
     merge_custom_metadata(ns_metadata)
     namespace_obj = kubernetes_client.V1Namespace(metadata=ns_metadata)
     try:
-        kubernetes.core_api().create_namespace(namespace_obj)
+        kubernetes.core_api(context).create_namespace(namespace_obj)
     except kubernetes.api_exception() as e:
         if e.status == 409:
             logger.info(f'Namespace {namespace} already exists in the cluster.')
@@ -1729,3 +1746,13 @@ def get_kubernetes_node_info() -> Dict[str, KubernetesNodeInfo]:
             free={'nvidia.com/gpu': int(accelerators_available)})
 
     return node_info_dict
+
+
+def get_namespace_from_config(provider_config: Dict[str, Any]) -> str:
+    return provider_config.get('namespace',
+                               get_current_kube_config_context_namespace())
+
+
+def get_context_from_config(provider_config: Dict[str, Any]) -> str:
+    return provider_config.get('context',
+                               get_current_kube_config_context_name())

sky/skylet/log_lib.py

@@ -208,9 +208,12 @@ def run_with_log(
                 str(proc.pid),
             ]
 
+            # We do not need to set `start_new_session=True` here, as the
+            # daemon script will detach itself from the parent process with
+            # fork to avoid being killed by ray job. See the reason we
+            # daemonize the process in `sky/skylet/subprocess_daemon.py`.
             subprocess.Popen(
                 daemon_cmd,
-                start_new_session=True,
                 # Suppress output
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,

sky/skylet/subprocess_daemon.py

@@ -1,17 +1,44 @@
 """Sky subprocess daemon.
-
 Wait for parent_pid to exit, then SIGTERM (or SIGKILL if needed) the child
 processes of proc_pid.
 """
-
 import argparse
+import os
 import sys
 import time
 
 import psutil
 
-if __name__ == '__main__':
 
+def daemonize():
+    """Detaches the process from its parent process with double-forking.
+
+    This detachment is crucial in the context of SkyPilot and Ray job. When
+    'sky cancel' is executed, it uses Ray's stop job API to terminate the job.
+    Without daemonization, this subprocess_daemon process would be terminated
+    along with its parent process, ray::task, which is launched with Ray job.
+    Daemonization ensures this process survives the 'sky cancel' command,
+    allowing it to prevent orphaned processes of Ray job.
+    """
+    # First fork: Creates a child process identical to the parent
+    if os.fork() > 0:
+        # Parent process exits, allowing the child to run independently
+        sys.exit()
+
+    # Continues to run from first forked child process.
+    # Detach from parent environment.
+    os.setsid()
+
+    # Second fork: Creates a grandchild process
+    if os.fork() > 0:
+        # First child exits, orphaning the grandchild
+        sys.exit()
+    # Continues execution in the grandchild process
+    # This process is now fully detached from the original parent and terminal
+
+
+if __name__ == '__main__':
+    daemonize()
     parser = argparse.ArgumentParser()
     parser.add_argument('--parent-pid', type=int, required=True)
     parser.add_argument('--proc-pid', type=int, required=True)
@@ -28,29 +55,34 @@ if __name__ == '__main__':
     if process is None:
         sys.exit()
 
+    children = []
     if parent_process is not None:
         # Wait for either parent or target process to exit.
         while process.is_running() and parent_process.is_running():
+            try:
+                # process.children() must be called while the target process
+                # is alive, as it will return an empty list if the target
+                # process has already terminated.
+                tmp_children = process.children(recursive=True)
+                if tmp_children:
+                    children = tmp_children
+            except psutil.NoSuchProcess:
+                pass
             time.sleep(1)
+    children.append(process)
 
-    try:
-        children = process.children(recursive=True)
-        children.append(process)
-    except psutil.NoSuchProcess:
-        sys.exit()
-
-    for pid in children:
+    for child in children:
         try:
-            pid.terminate()
+            child.terminate()
         except psutil.NoSuchProcess:
-            pass
+            continue
 
     # Wait 30s for the processes to exit gracefully.
     time.sleep(30)
 
     # SIGKILL if they're still running.
-    for pid in children:
+    for child in children:
         try:
-            pid.kill()
+            child.kill()
         except psutil.NoSuchProcess:
-            pass
+            continue

sky/templates/kubernetes-port-forward-proxy-command.sh

@@ -1,15 +1,40 @@
 #!/usr/bin/env bash
 set -uo pipefail
 
+KUBE_CONTEXT=""
+KUBE_NAMESPACE=""
+
+# Parse flags
+while getopts ":c:n:" opt; do
+  case ${opt} in
+    c)
+      KUBE_CONTEXT="$OPTARG"
+      ;;
+    n)
+      KUBE_NAMESPACE="$OPTARG"
+      ;;
+    \?)
+      echo "Invalid option: -$OPTARG" >&2
+      echo "Usage: $0 <pod_name> [-c kube_context] [-n kube_namespace]" >&2
+      exit 1
+      ;;
+    :)
+      echo "Option -$OPTARG requires an argument." >&2
+      exit 1
+      ;;
+  esac
+done
+
+# Shift the processed options away so that $1 becomes the pod name
+shift $((OPTIND -1))
+
 # Check if pod name is passed as an argument
 if [ $# -lt 1 ]; then
-  echo "Usage: $0 <pod_name> [kube_context] [kube_namespace]" >&2
+  echo "Usage: $0 <pod_name> [-c kube_context] [-n kube_namespace]" >&2
   exit 1
 fi
 
-POD_NAME="$1"  # The first argument is the name of the pod
-KUBE_CONTEXT="${2:-}"  # The second argument is the kube context, default is empty
-KUBE_NAMESPACE="${3:-}"  # The third argument is the kube namespace, default is empty
+POD_NAME="$1"  # The first positional argument is the name of the pod
 
 # Checks if socat is installed
 if ! command -v socat > /dev/null; then

sky/templates/kubernetes-ray.yml.j2

@@ -21,6 +21,11 @@ provider:
     # The namespace to create the Ray cluster in.
     namespace: {{k8s_namespace}}
 
+    # The kubecontext used to connect to the Kubernetes cluster.
+    {% if k8s_context is not none %}
+    context: {{k8s_context}}
+    {% endif %}
+
     # This should be one of KubernetesPortMode
     port_mode: {{k8s_port_mode}}
 

sky/templates/lambda-ray.yml.j2

@@ -89,13 +89,13 @@ setup_commands:
 # Increment the following for catching performance bugs easier:
 #   current num items (num SSH connections): 2
 head_start_ray_commands:
-  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
+  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --head --port={{ray_port}} --min-worker-port 11002 --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
     {{dump_port_command}}; {{ray_head_wait_initialized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --address=$RAY_HEAD_IP:{{ray_port}} --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
+  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --address=$RAY_HEAD_IP:{{ray_port}} --min-worker-port 11002 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []

sky/utils/command_runner.py

@@ -649,13 +649,13 @@ class KubernetesCommandRunner(CommandRunner):
 
     def __init__(
         self,
-        node: Tuple[str, str],
+        node: Tuple[Tuple[str, str], str],
         **kwargs,
     ):
         """Initialize KubernetesCommandRunner.
 
         Example Usage:
-            runner = KubernetesCommandRunner((namespace, pod_name))
+            runner = KubernetesCommandRunner((namespace, context), pod_name))
             runner.run('ls -l')
             runner.rsync(source, target, up=True)
 
@@ -664,7 +664,11 @@ class KubernetesCommandRunner(CommandRunner):
         """
         del kwargs
         super().__init__(node)
-        self.namespace, self.pod_name = node
+        (self.namespace, self.context), self.pod_name = node
+
+    @property
+    def node_id(self) -> str:
+        return f'{self.context}-{self.namespace}-{self.pod_name}'
 
     @timeline.event
     def run(
@@ -719,9 +723,11 @@ class KubernetesCommandRunner(CommandRunner):
         if connect_timeout is None:
             connect_timeout = _DEFAULT_CONNECT_TIMEOUT
         kubectl_args = [
-            '--pod-running-timeout', f'{connect_timeout}s', '-n',
-            self.namespace, self.pod_name
+            '--pod-running-timeout', f'{connect_timeout}s', '-n', self.namespace
         ]
+        if self.context:
+            kubectl_args += ['--context', self.context]
+        kubectl_args += [self.pod_name]
         if ssh_mode == SshMode.LOGIN:
             assert isinstance(cmd, list), 'cmd must be a list for login mode.'
             base_cmd = ['kubectl', 'exec', '-it', *kubectl_args, '--']
@@ -821,7 +827,7 @@ class KubernetesCommandRunner(CommandRunner):
         self._rsync(
             source,
             target,
-            node_destination=f'{self.pod_name}@{self.namespace}',
+            node_destination=f'{self.pod_name}@{self.namespace}+{self.context}',
             up=up,
             rsh_option=helper_path,
             log_path=log_path,

sky/utils/command_runner.pyi

@@ -204,7 +204,7 @@ class KubernetesCommandRunner(CommandRunner):
 
     def __init__(
         self,
-        node: Tuple[str, str],
+        node: Tuple[Tuple[str, str], str],
     ) -> None:
         ...
 

sky/utils/kubernetes/rsync_helper.sh

@@ -1,7 +1,16 @@
-# When using pod@namespace, rsync passes args as: {us} -l pod namespace 
+# When using pod@namespace+context, rsync passes args as: {us} -l pod namespace+context
+# We need to split the pod@namespace+context into pod, namespace and context
+# For backward compatibility, we use + as the separator between namespace and context and add handling when context is not provided
 shift
 pod=$1
 shift
-namespace=$1
+namespace_context=$1
+namespace=$(echo $namespace_context | cut -d+ -f1)
+context=$(echo $namespace_context | grep '+' >/dev/null && echo $namespace_context | cut -d+ -f2- || echo "")
+context_lower=$(echo "$context" | tr '[:upper:]' '[:lower:]')
 shift
-kubectl exec -i $pod -n $namespace -- "$@"
+if [ -z "$context" ] || [ "$context_lower" = "none" ]; then
+    kubectl exec -i $pod -n $namespace -- "$@"
+else
+    kubectl exec -i $pod -n $namespace --context=$context -- "$@"
+fi

{skypilot_nightly-1.0.0.dev20240910.dist-info → skypilot_nightly-1.0.0.dev20240911.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20240910
+Version: 1.0.0.dev20240911
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0