From 09849f8493cbef714101d2fec08a24c828ebb4d0 Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@berkeley.edu>
Date: Sun, 15 Oct 2023 21:23:09 -0700
Subject: [PATCH] [k8s] Add netcat to dependency checks (#2712)

* add netcat dep

* nc -h
---
 docs/source/reference/kubernetes/index.rst        |  6 +++---
 .../kubernetes-port-forward-proxy-command.sh.j2   |  6 ++++++
 sky/utils/kubernetes_utils.py                     | 15 ++++++++++-----
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/docs/source/reference/kubernetes/index.rst b/docs/source/reference/kubernetes/index.rst
index 3d5c0c7d086..7517c70d211 100644
--- a/docs/source/reference/kubernetes/index.rst
+++ b/docs/source/reference/kubernetes/index.rst
@@ -50,15 +50,15 @@ Submitting SkyPilot tasks to Kubernetes Clusters
 
 Once your cluster administrator has :ref:`setup a Kubernetes cluster <kubernetes-setup>` and provided you with a kubeconfig file:
 
-0. Make sure `kubectl <https://kubernetes.io/docs/tasks/tools/>`_ and ``socat`` are installed on your local machine.
+0. Make sure `kubectl <https://kubernetes.io/docs/tasks/tools/>`_, ``socat`` and ``nc`` (netcat) are installed on your local machine.
 
    .. code-block:: console
 
      $ # MacOS
-     $ brew install kubectl socat
+     $ brew install kubectl socat netcat
 
      $ # Linux (may have socat already installed)
-     $ sudo apt-get install kubectl socat
+     $ sudo apt-get install kubectl socat netcat
 
 
 1. Place your kubeconfig file at ``~/.kube/config``.
diff --git a/sky/templates/kubernetes-port-forward-proxy-command.sh.j2 b/sky/templates/kubernetes-port-forward-proxy-command.sh.j2
index 13cef3e8b44..74c6a236313 100644
--- a/sky/templates/kubernetes-port-forward-proxy-command.sh.j2
+++ b/sky/templates/kubernetes-port-forward-proxy-command.sh.j2
@@ -7,6 +7,12 @@ if ! command -v socat > /dev/null; then
   exit
 fi
 
+# Checks if netcat is installed (may not be present in many docker images)
+if ! command -v nc > /dev/null; then
+  echo "Using 'port-forward' mode to run ssh session on Kubernetes instances requires 'nc' to be installed. Please install 'nc' (netcat)." >&2
+  exit
+fi
+
 # Establishes connection between local port and the ssh jump pod using kube port-forward
 # Instead of specifying a port, we let kubectl select a random port.
 # This is preferred because of socket re-use issues in kubectl port-forward,
diff --git a/sky/utils/kubernetes_utils.py b/sky/utils/kubernetes_utils.py
index afd5b0affb2..5943df86c1a 100644
--- a/sky/utils/kubernetes_utils.py
+++ b/sky/utils/kubernetes_utils.py
@@ -991,19 +991,24 @@ def fill_ssh_jump_template(ssh_key_secret: str, ssh_jump_image: str,
 
 def check_port_forward_mode_dependencies() -> None:
     """Checks if 'socat' is installed"""
-    for name, option in [('socat', '-V')]:
+    # We store the dependency list as a list of lists. Each inner list
+    # contains the name of the dependency, the command to check if it is
+    # installed, and the package name to install it.
+    dependency_list = [['socat', ['socat', '-V'], 'socat'],
+                       ['nc', ['nc', '-h'], 'netcat']]
+    for name, check_cmd, install_cmd in dependency_list:
         try:
-            subprocess.run([name, option],
+            subprocess.run(check_cmd,
                            stdout=subprocess.DEVNULL,
                            stderr=subprocess.DEVNULL,
                            check=True)
-        except FileNotFoundError:
+        except (FileNotFoundError, subprocess.CalledProcessError):
             with ux_utils.print_exception_no_traceback():
                 raise RuntimeError(
                     f'`{name}` is required to setup Kubernetes cloud with '
                     f'`{KubernetesNetworkingMode.PORTFORWARD.value}` default '
                     'networking mode and it is not installed. '
                     'On Debian/Ubuntu, install it with:\n'
-                    f'  $ sudo apt install {name}\n'
+                    f'  $ sudo apt install {install_cmd}\n'
                     f'On MacOS, install it with: \n'
-                    f'  $ brew install {name}') from None
+                    f'  $ brew install {install_cmd}') from None