[tune] add output flag for Tune CLI (#4322)

python/ray/tune/commands.py

@@ -43,12 +43,12 @@ DEFAULT_PROJECT_INFO_KEYS = (
 )
 
 try:
-    TERM_HEIGHT, TERM_WIDTH = subprocess.check_output(['stty', 'size']).split()
+    TERM_HEIGHT, TERM_WIDTH = subprocess.check_output(["stty", "size"]).split()
     TERM_HEIGHT, TERM_WIDTH = int(TERM_HEIGHT), int(TERM_WIDTH)
 except subprocess.CalledProcessError:
     TERM_HEIGHT, TERM_WIDTH = 100, 100
 
-EDITOR = os.getenv('EDITOR', 'vim')
+EDITOR = os.getenv("EDITOR", "vim")
 
 
 def _check_tabulate():
@@ -116,6 +116,7 @@ def _get_experiment_state(experiment_path, exit_on_fail=False):
 
 def list_trials(experiment_path,
                 sort=None,
+                output=None,
                 info_keys=DEFAULT_EXPERIMENT_INFO_KEYS,
                 result_keys=DEFAULT_RESULT_KEYS):
     """Lists trials in the directory subtree starting at the given path.
@@ -124,6 +125,7 @@ def list_trials(experiment_path,
         experiment_path (str): Directory where trials are located.
             Corresponds to Experiment.local_dir/Experiment.name.
         sort (str): Key to sort by.
+        output (str): Name of file where output is saved.
         info_keys (list): Keys that are displayed.
         result_keys (list): Keys of last result that are displayed.
     """
@@ -142,7 +144,7 @@ def list_trials(experiment_path,
     checkpoints_df = checkpoints_df[col_keys]
 
     if "last_update_time" in checkpoints_df:
-        with pd.option_context('mode.use_inf_as_null', True):
+        with pd.option_context("mode.use_inf_as_null", True):
             datetime_series = checkpoints_df["last_update_time"].dropna()
 
         datetime_series = datetime_series.apply(
@@ -162,9 +164,22 @@ def list_trials(experiment_path,
 
     print_format_output(checkpoints_df)
 
+    if output:
+        experiment_path = os.path.expanduser(experiment_path)
+        output_path = os.path.join(experiment_path, output)
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            checkpoints_df.to_pickle(output_path)
+        elif file_extension == ".csv":
+            checkpoints_df.to_csv(output_path, index=False)
+        else:
+            raise ValueError("Unsupported filetype: {}".format(output))
+        print("Output saved at:", output_path)
+
 
 def list_experiments(project_path,
                      sort=None,
+                     output=None,
                      info_keys=DEFAULT_PROJECT_INFO_KEYS):
     """Lists experiments in the directory subtree.
 
@@ -172,6 +187,7 @@ def list_experiments(project_path,
         project_path (str): Directory where experiments are located.
             Corresponds to Experiment.local_dir.
         sort (str): Key to sort by.
+        output (str): Name of file where output is saved.
         info_keys (list): Keys that are displayed.
     """
     _check_tabulate()
@@ -233,6 +249,17 @@ def list_experiments(project_path,
 
     print_format_output(info_df)
 
+    if output:
+        output_path = os.path.join(base, output)
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            info_df.to_pickle(output_path)
+        elif file_extension == ".csv":
+            info_df.to_csv(output_path, index=False)
+        else:
+            raise ValueError("Unsupported filetype: {}".format(output))
+        print("Output saved at:", output_path)
+
 
 def add_note(path, filename="note.txt"):
     """Opens a txt file at the given path where user can add and save notes.

python/ray/tune/scripts.py

@@ -14,19 +14,31 @@ def cli():
 @cli.command()
 @click.argument("experiment_path", required=True, type=str)
 @click.option(
-    '--sort', default=None, type=str, help='Select which column to sort on.')
-def list_trials(experiment_path, sort):
+    "--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Output information to a pickle file.")
+def list_trials(experiment_path, sort, output):
     """Lists trials in the directory subtree starting at the given path."""
-    commands.list_trials(experiment_path, sort)
+    commands.list_trials(experiment_path, sort, output)
 
 
 @cli.command()
 @click.argument("project_path", required=True, type=str)
 @click.option(
-    '--sort', default=None, type=str, help='Select which column to sort on.')
-def list_experiments(project_path, sort):
+    "--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Select filename to output information to.")
+def list_experiments(project_path, sort, output):
     """Lists experiments in the directory subtree."""
-    commands.list_experiments(project_path, sort)
+    commands.list_experiments(project_path, sort, output)
 
 
 @cli.command()
@@ -35,7 +47,7 @@ def list_experiments(project_path, sort):
     "--filename",
     default="note.txt",
     type=str,
-    help='Specify filename for note.')
+    help="Specify filename for note.")
 def add_note(path, filename):
     """Adds user notes as a text file at the given path."""
     commands.add_note(path, filename)