diff --git a/python/ray/tune/commands.py b/python/ray/tune/commands.py index 5abf49858dd1..add8cf25ada5 100644 --- a/python/ray/tune/commands.py +++ b/python/ray/tune/commands.py @@ -43,12 +43,12 @@ ) 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. diff --git a/python/ray/tune/scripts.py b/python/ray/tune/scripts.py index 3b810ee3ae5c..97527e5ead98 100644 --- a/python/ray/tune/scripts.py +++ b/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)