[autoscaler] port-forward for attach + redis_port (#7145)

* port-forward

* fixport

* force redis port in init mode

* test

* Update python/ray/tests/test_ray_init.py

python/ray/autoscaler/commands.py

@@ -310,7 +310,7 @@ def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
 
 
 def attach_cluster(config_file, start, use_screen, use_tmux,
-                   override_cluster_name, new):
+                   override_cluster_name, new, port_forward):
     """Attaches to a screen for the specified cluster.
 
     Arguments:
@@ -320,6 +320,7 @@ def attach_cluster(config_file, start, use_screen, use_tmux,
         use_tmux: whether to use tmux as multiplexer
         override_cluster_name: set the name of the cluster
         new: whether to force a new screen
+        port_forward (int or list[int]): port(s) to forward
     """
 
     if use_tmux:
@@ -339,7 +340,7 @@ def attach_cluster(config_file, start, use_screen, use_tmux,
         cmd = "$SHELL"
 
     exec_cluster(config_file, cmd, False, False, False, False, start,
-                 override_cluster_name, None)
+                 override_cluster_name, port_forward)
 
 
 def exec_cluster(config_file,

python/ray/scripts/scripts.py

@@ -668,8 +668,18 @@ def monitor(cluster_config_file, lines, cluster_name):
     help="Override the configured cluster name.")
 @click.option(
     "--new", "-N", is_flag=True, help="Force creation of a new screen.")
-def attach(cluster_config_file, start, screen, tmux, cluster_name, new):
-    attach_cluster(cluster_config_file, start, screen, tmux, cluster_name, new)
+@click.option(
+    "--port-forward",
+    "-p",
+    required=False,
+    multiple=True,
+    type=int,
+    help="Port to forward. Use this multiple times to forward multiple ports.")
+def attach(cluster_config_file, start, screen, tmux, cluster_name, new,
+           port_forward):
+    port_forward = [(port, port) for port in list(port_forward)]
+    attach_cluster(cluster_config_file, start, screen, tmux, cluster_name, new,
+                   port_forward)
 
 
 @cli.command()
@@ -744,6 +754,7 @@ def rsync_up(cluster_config_file, source, target, cluster_name, all_nodes):
     help="Override the configured cluster name.")
 @click.option(
     "--port-forward",
+    "-p",
     required=False,
     multiple=True,
     type=int,
@@ -821,6 +832,7 @@ def submit(cluster_config_file, docker, screen, tmux, stop, start,
     help="Override the configured cluster name.")
 @click.option(
     "--port-forward",
+    "-p",
     required=False,
     multiple=True,
     type=int,

python/ray/tests/test_ray_init.py

@@ -58,6 +58,20 @@ class TestRedisPassword:
         object_id = f.remote()
         ray.get(object_id)
 
+    def test_redis_port(self, shutdown_only):
+        @ray.remote
+        def f():
+            return 1
+
+        info = ray.init(redis_port=1234, redis_password="testpassword")
+        address = info["redis_address"]
+        redis_ip, redis_port = address.split(":")
+        assert redis_port == "1234"
+
+        redis_client = redis.StrictRedis(
+            host=redis_ip, port=redis_port, password="testpassword")
+        assert redis_client.ping()
+
 
 if __name__ == "__main__":
     import pytest

python/ray/worker.py

@@ -535,6 +535,7 @@ def print_failed_task(task_status):
 
 def init(address=None,
          redis_address=None,
+         redis_port=None,
          num_cpus=None,
          num_gpus=None,
          memory=None,
@@ -592,6 +593,8 @@ def init(address=None,
             raylet, a plasma store, a plasma manager, and some workers.
             It will also kill these processes when Python exits.
         redis_address (str): Deprecated; same as address.
+        redis_port (int): The port that the primary Redis shard should listen
+            to. If None, then a random port will be chosen.
         num_cpus (int): Number of cpus the user wishes all raylets to
             be configured with.
         num_gpus (int): Number of gpus the user wishes all raylets to
@@ -714,6 +717,7 @@ def init(address=None,
         # In this case, we need to start a new cluster.
         ray_params = ray.parameter.RayParams(
             redis_address=redis_address,
+            redis_port=redis_port,
             node_ip_address=node_ip_address,
             object_id_seed=object_id_seed,
             local_mode=local_mode,