From 29eee7f9707f9ebda55ebafa834ecc709a7d9add Mon Sep 17 00:00:00 2001
From: Philipp Moritz <pcmoritz@gmail.com>
Date: Fri, 18 Oct 2019 16:50:46 -0700
Subject: [PATCH] Forward multiple ports for autoscaler (#5893)

---
 python/ray/autoscaler/commands.py |  2 +-
 python/ray/autoscaler/updater.py  | 13 +++++++++----
 python/ray/scripts/scripts.py     | 16 ++++++++++++----
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/python/ray/autoscaler/commands.py b/python/ray/autoscaler/commands.py
index 43ec645d3..9367baa29 100644
--- a/python/ray/autoscaler/commands.py
+++ b/python/ray/autoscaler/commands.py
@@ -347,7 +347,7 @@ def exec_cluster(config_file, cmd, docker, screen, tmux, stop, start,
         stop: whether to stop the cluster after command run
         start: whether to start the cluster if it isn't up
         override_cluster_name: set the name of the cluster
-        port_forward: port to forward
+        port_forward (int or list[int]): port(s) to forward
     """
     assert not (screen and tmux), "Can specify only one of `screen` or `tmux`."
 
diff --git a/python/ray/autoscaler/updater.py b/python/ray/autoscaler/updater.py
index b94200907..ee93f9734 100644
--- a/python/ray/autoscaler/updater.py
+++ b/python/ray/autoscaler/updater.py
@@ -57,10 +57,12 @@ class KubernetesCommandRunner(object):
         logger.info(self.log_prefix + "Running {}...".format(cmd))
 
         if port_forward:
+            if not isinstance(port_forward, list):
+                port_forward = [port_forward]
             port_forward_cmd = self.kubectl + [
-                "port-forward", self.node_id,
-                str(port_forward)
-            ]
+                "port-forward",
+                self.node_id,
+            ] + [str(fwd) for fwd in port_forward]
             port_forward_process = subprocess.Popen(port_forward_cmd)
             # Give port-forward a grace period to run and print output before
             # running the actual command. This is a little ugly, but it should
@@ -242,7 +244,10 @@ class SSHCommandRunner(object):
             ssh.append("-tt")
 
         if port_forward:
-            ssh += ["-L", "{}:localhost:{}".format(port_forward, port_forward)]
+            if not isinstance(port_forward, list):
+                port_forward = [port_forward]
+            for fwd in port_forward:
+                ssh += ["-L", "{}:localhost:{}".format(fwd, fwd)]
 
         final_cmd = ssh + self.get_default_ssh_options(timeout) + [
             "{}@{}".format(self.ssh_user, self.ssh_ip)
diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 0d8754fcd..0cee431c7 100644
--- a/python/ray/scripts/scripts.py
+++ b/python/ray/scripts/scripts.py
@@ -637,7 +637,11 @@ def rsync_up(cluster_config_file, source, target, cluster_name):
     type=str,
     help="Override the configured cluster name.")
 @click.option(
-    "--port-forward", required=False, type=int, help="Port to forward.")
+    "--port-forward",
+    required=False,
+    multiple=True,
+    type=int,
+    help="Port to forward. Use this multiple times to forward multiple ports.")
 @click.argument("script", required=True, type=str)
 @click.option("--args", required=False, type=str, help="Script args.")
 def submit(cluster_config_file, docker, screen, tmux, stop, start,
@@ -665,7 +669,7 @@ def submit(cluster_config_file, docker, screen, tmux, stop, start,
         command_parts += [args]
     cmd = " ".join(command_parts)
     exec_cluster(cluster_config_file, cmd, docker, screen, tmux, stop, False,
-                 cluster_name, port_forward)
+                 cluster_name, list(port_forward))
 
 
 @cli.command()
@@ -700,11 +704,15 @@ def submit(cluster_config_file, docker, screen, tmux, stop, start,
     type=str,
     help="Override the configured cluster name.")
 @click.option(
-    "--port-forward", required=False, type=int, help="Port to forward.")
+    "--port-forward",
+    required=False,
+    multiple=True,
+    type=int,
+    help="Port to forward. Use this multiple times to forward multiple ports.")
 def exec_cmd(cluster_config_file, cmd, docker, screen, tmux, stop, start,
              cluster_name, port_forward):
     exec_cluster(cluster_config_file, cmd, docker, screen, tmux, stop, start,
-                 cluster_name, port_forward)
+                 cluster_name, list(port_forward))
 
 
 @cli.command()