From ba4b7e4cb91ca9d990ff96934ab6026ed8e5e05e Mon Sep 17 00:00:00 2001
From: Colin Campbell <colin.campbell@shopify.com>
Date: Fri, 30 Aug 2024 11:41:13 -0400
Subject: [PATCH] Fix ssh breaking if kube context namespace changes

---
 sky/provision/kubernetes/utils.py                  | 12 +++++++++---
 .../kubernetes-port-forward-proxy-command.sh       | 14 ++++++++++----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
index 8ac3ab1d4ca7..48982e4122f3 100644
--- a/sky/provision/kubernetes/utils.py
+++ b/sky/provision/kubernetes/utils.py
@@ -935,7 +935,8 @@ def construct_ssh_jump_command(
         ssh_jump_user: str = 'sky',
         proxy_cmd_path: Optional[str] = None,
         proxy_cmd_target_pod: Optional[str] = None,
-        current_kube_context: Optional[str] = None) -> str:
+        current_kube_context: Optional[str] = None,
+        current_kube_namespace: Optional[str] = None) -> str:
     ssh_jump_proxy_command = (f'ssh -tt -i {private_key_path} '
                               '-o StrictHostKeyChecking=no '
                               '-o UserKnownHostsFile=/dev/null '
@@ -949,9 +950,12 @@ def construct_ssh_jump_command(
         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 (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_context_flag}'
+                                   f'{kube_namespace_flag}\'')
     return ssh_jump_proxy_command
 
 
@@ -1017,13 +1021,15 @@ def get_ssh_proxy_command(
     else:
         ssh_jump_proxy_command_path = create_proxy_command_script()
         current_context = get_current_kube_config_context_name()
+        current_namespace = get_current_kube_config_context_namespace()
         ssh_jump_proxy_command = construct_ssh_jump_command(
             private_key_path,
             ssh_jump_ip,
             ssh_jump_user=constants.SKY_SSH_USER_PLACEHOLDER,
             proxy_cmd_path=ssh_jump_proxy_command_path,
             proxy_cmd_target_pod=k8s_ssh_target,
-            current_kube_context=current_context)
+            current_kube_context=current_context,
+            current_kube_namespace=current_namespace)
     return ssh_jump_proxy_command
 
 
diff --git a/sky/templates/kubernetes-port-forward-proxy-command.sh b/sky/templates/kubernetes-port-forward-proxy-command.sh
index 27580ffbe042..26a929421e56 100644
--- a/sky/templates/kubernetes-port-forward-proxy-command.sh
+++ b/sky/templates/kubernetes-port-forward-proxy-command.sh
@@ -3,12 +3,13 @@ set -uo pipefail
 
 # Check if pod name is passed as an argument
 if [ $# -lt 1 ]; then
-  echo "Usage: $0 <pod_name> [kube_context]" >&2
+  echo "Usage: $0 <pod_name> [kube_context] [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
 
 # Checks if socat is installed
 if ! command -v socat > /dev/null; then
@@ -27,12 +28,17 @@ fi
 # This is preferred because of socket re-use issues in kubectl port-forward,
 # see - https://github.com/kubernetes/kubernetes/issues/74551#issuecomment-769185879
 KUBECTL_OUTPUT=$(mktemp)
+KUBECTL_ARGS=()
+
 if [ -n "$KUBE_CONTEXT" ]; then
-  kubectl --context="$KUBE_CONTEXT" port-forward pod/"${POD_NAME}" :22 > "${KUBECTL_OUTPUT}" 2>&1 &
-else
-  kubectl port-forward pod/"${POD_NAME}" :22 > "${KUBECTL_OUTPUT}" 2>&1 &
+  KUBECTL_ARGS+=("--context=$KUBE_CONTEXT")
+fi
+if [ -n "$KUBE_NAMESPACE" ]; then
+  KUBECTL_ARGS+=("--namespace=$KUBE_NAMESPACE")
 fi
 
+kubectl "${KUBECTL_ARGS[@]}" port-forward pod/"${POD_NAME}" :22 > "${KUBECTL_OUTPUT}" 2>&1 &
+
 # Capture the PID for the backgrounded kubectl command
 K8S_PORT_FWD_PID=$!