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 [kube_context]" >&2 + echo "Usage: $0 [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=$!