From be6567e6fd8fdaab0bf1f72e2255890e94a4f2b3 Mon Sep 17 00:00:00 2001
From: Kristian Hartikainen <kristian.hartikainen@gmail.com>
Date: Mon, 3 Dec 2018 21:39:43 -0800
Subject: [PATCH] Tweak/exec attach info (#3447)

* Add custom cluster name to exec info

* Update submit info to match exec info
---
 python/ray/scripts/scripts.py | 37 +++++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 13fa63efb..60b05e422 100644
--- a/python/ray/scripts/scripts.py
+++ b/python/ray/scripts/scripts.py
@@ -591,9 +591,21 @@ def submit(cluster_config_file, screen, tmux, stop, start, cluster_name,
     cmd = " ".join(["python", target] + list(script_args))
     exec_cluster(cluster_config_file, cmd, screen, tmux, stop, False,
                  cluster_name, port_forward)
-    if tmux:
-        logger.info("Use `ray attach {} --tmux` "
-                    "to check on command status.".format(cluster_config_file))
+
+    if tmux or screen:
+        attach_command_parts = ["ray attach", cluster_config_file]
+        if cluster_name is not None:
+            attach_command_parts.append(
+                "--cluster-name={}".format(cluster_name))
+        if tmux:
+            attach_command_parts.append("--tmux")
+        elif screen:
+            attach_command_parts.append("--screen")
+
+        attach_command = " ".join(attach_command_parts)
+        attach_info = "Use `{}` to check on command status.".format(
+            attach_command)
+        logger.info(attach_info)
 
 
 @cli.command()
@@ -627,11 +639,24 @@ def submit(cluster_config_file, screen, tmux, stop, start, cluster_name,
 def exec_cmd(cluster_config_file, cmd, screen, tmux, stop, start, cluster_name,
              port_forward):
     assert not (screen and tmux), "Can specify only one of `screen` or `tmux`."
+
     exec_cluster(cluster_config_file, cmd, screen, tmux, stop, start,
                  cluster_name, port_forward)
-    if tmux:
-        logger.info("Use `ray attach {} --tmux` "
-                    "to check on command status.".format(cluster_config_file))
+
+    if tmux or screen:
+        attach_command_parts = ["ray attach", cluster_config_file]
+        if cluster_name is not None:
+            attach_command_parts.append(
+                "--cluster-name={}".format(cluster_name))
+        if tmux:
+            attach_command_parts.append("--tmux")
+        elif screen:
+            attach_command_parts.append("--screen")
+
+        attach_command = " ".join(attach_command_parts)
+        attach_info = "Use `{}` to check on command status.".format(
+            attach_command)
+        logger.info(attach_info)
 
 
 @cli.command()