From cf9cd5da9ddec3c9f8e6f58677c662fd14119a29 Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Sat, 29 Sep 2018 23:04:13 -0700
Subject: [PATCH] [ray] Add --new flag for ray attach (#2973)

* new flag

* yapf
---
 python/ray/autoscaler/commands.py | 16 ++++++++++++++--
 python/ray/scripts/scripts.py     |  6 ++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/python/ray/autoscaler/commands.py b/python/ray/autoscaler/commands.py
index a2be4a803..91d3e4b6b 100644
--- a/python/ray/autoscaler/commands.py
+++ b/python/ray/autoscaler/commands.py
@@ -210,7 +210,7 @@ def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
                                        provider.external_ip(head_node)))
 
 
-def attach_cluster(config_file, start, use_tmux, override_cluster_name):
+def attach_cluster(config_file, start, use_tmux, override_cluster_name, new):
     """Attaches to a screen for the specified cluster.
 
     Arguments:
@@ -218,8 +218,20 @@ def attach_cluster(config_file, start, use_tmux, override_cluster_name):
         start: whether to start the cluster if it isn't up
         use_tmux: whether to use tmux as multiplexer
         override_cluster_name: set the name of the cluster
+        new: whether to force a new screen
     """
-    cmd = "tmux attach || tmux new" if use_tmux else "screen -L -xRR"
+
+    if use_tmux:
+        if new:
+            cmd = "tmux new"
+        else:
+            cmd = "tmux attach || tmux new"
+    else:
+        if new:
+            cmd = "screen -L"
+        else:
+            cmd = "screen -L -xRR"
+
     exec_cluster(config_file, cmd, False, False, False, start,
                  override_cluster_name, None)
 
diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 09ef158dd..0826a1387 100644
--- a/python/ray/scripts/scripts.py
+++ b/python/ray/scripts/scripts.py
@@ -478,8 +478,10 @@ def teardown(cluster_config_file, yes, workers_only, cluster_name):
     required=False,
     type=str,
     help=("Override the configured cluster name."))
-def attach(cluster_config_file, start, tmux, cluster_name):
-    attach_cluster(cluster_config_file, start, tmux, cluster_name)
+@click.option(
+    "--new", "-N", is_flag=True, help=("Force creation of a new screen."))
+def attach(cluster_config_file, start, tmux, cluster_name, new):
+    attach_cluster(cluster_config_file, start, tmux, cluster_name, new)
 
 
 @cli.command()