diff --git a/python/ray/autoscaler/commands.py b/python/ray/autoscaler/commands.py index be028afc3..2364811d8 100644 --- a/python/ray/autoscaler/commands.py +++ b/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, diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py index cacf2bfe0..0aeec706d 100644 --- a/python/ray/scripts/scripts.py +++ b/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, diff --git a/python/ray/tests/test_ray_init.py b/python/ray/tests/test_ray_init.py index cba08a6cd..119e12aeb 100644 --- a/python/ray/tests/test_ray_init.py +++ b/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 diff --git a/python/ray/worker.py b/python/ray/worker.py index c9fffa23b..6c63119d1 100644 --- a/python/ray/worker.py +++ b/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,