chore: import upstream snapshot with attribution

python/ray/tune/tests/test_commands.py

@@ -0,0 +1,196 @@
+import os
+import random
+import subprocess
+import sys
+import time
+from unittest import mock
+
+import click
+import pytest
+
+import ray
+from ray import tune
+from ray.train.tests.util import create_dict_checkpoint
+from ray.tune.cli import commands
+from ray.tune.result import CONFIG_PREFIX
+from ray.tune.utils.mock_trainable import MyTrainableClass
+
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+
+class Capturing:
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        self.captured = []
+        return self
+
+    def __exit__(self, *args):
+        self.captured.extend(self._stringio.getvalue().splitlines())
+        del self._stringio  # free up some memory
+        sys.stdout = self._stdout
+
+
+@pytest.fixture
+def start_ray():
+    ray.init(log_to_driver=False)
+    yield
+    ray.shutdown()
+
+
+def test_time(start_ray, tmpdir, monkeypatch):
+    experiment_name = "test_time"
+    num_samples = 2
+
+    def train_fn(config):
+        for i in range(3):
+            with create_dict_checkpoint({"dummy": "data"}) as checkpoint:
+                ray.tune.report(
+                    {
+                        "epoch": i,
+                        "a": random.random(),
+                        "b/c": random.random(),
+                        "d": random.random(),
+                    },
+                    checkpoint=checkpoint,
+                )
+
+    tuner = tune.Tuner(
+        train_fn,
+        param_space={f"hp{i}": tune.uniform(0, 1) for i in range(100)},
+        tune_config=tune.TuneConfig(num_samples=num_samples),
+        run_config=ray.tune.RunConfig(name=experiment_name),
+    )
+    results = tuner.fit()
+    times = []
+    for _ in range(5):
+        start = time.time()
+        subprocess.check_call(["tune", "ls", results.experiment_path])
+        times += [time.time() - start]
+
+    print("Average CLI time: ", sum(times) / len(times))
+    assert sum(times) / len(times) < 5, "CLI is taking too long!"
+
+
+@mock.patch(
+    "ray.tune.cli.commands.print_format_output",
+    wraps=ray.tune.cli.commands.print_format_output,
+)
+def test_ls(mock_print_format_output, start_ray, tmpdir):
+    """This test captures output of list_trials."""
+    experiment_name = "test_ls"
+    experiment_path = os.path.join(str(tmpdir), experiment_name)
+    num_samples = 3
+    tune.run(
+        MyTrainableClass,
+        name=experiment_name,
+        stop={"training_iteration": 1},
+        num_samples=num_samples,
+        storage_path=str(tmpdir),
+    )
+
+    columns = ["episode_reward_mean", "training_iteration", "trial_id"]
+    limit = 2
+    commands.list_trials(experiment_path, info_keys=columns, limit=limit)
+
+    # The dataframe that is printed as a table is the first arg of the last
+    # call made to `ray.tune.cli.commands.print_format_output`.
+    mock_print_format_output.assert_called()
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(columns), df
+    assert len(df.index) == limit, df
+
+    commands.list_trials(
+        experiment_path,
+        sort=["trial_id"],
+        info_keys=("trial_id", "training_iteration"),
+        filter_op="training_iteration == 1",
+    )
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(["trial_id", "training_iteration"])
+    assert len(df.index) == num_samples
+
+    with pytest.raises(click.ClickException):
+        commands.list_trials(
+            experiment_path, sort=["trial_id"], info_keys=("training_iteration",)
+        )
+
+    with pytest.raises(click.ClickException):
+        commands.list_trials(experiment_path, info_keys=("asdf",))
+
+
+@mock.patch(
+    "ray.tune.cli.commands.print_format_output",
+    wraps=ray.tune.cli.commands.print_format_output,
+)
+def test_ls_with_cfg(mock_print_format_output, start_ray, tmpdir):
+    experiment_name = "test_ls_with_cfg"
+    experiment_path = os.path.join(str(tmpdir), experiment_name)
+    tune.run(
+        MyTrainableClass,
+        name=experiment_name,
+        stop={"training_iteration": 1},
+        config={"test_variable": tune.grid_search(list(range(5)))},
+        storage_path=str(tmpdir),
+    )
+
+    columns = [CONFIG_PREFIX + "/test_variable", "trial_id"]
+    limit = 4
+
+    commands.list_trials(experiment_path, info_keys=columns, limit=limit)
+
+    # The dataframe that is printed as a table is the first arg of the last
+    # call made to `ray.tune.cli.commands.print_format_output`.
+    mock_print_format_output.assert_called()
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(columns), df
+    assert len(df.index) == limit, df
+
+
+def test_lsx(start_ray, tmpdir):
+    """This test captures output of list_experiments."""
+    project_path = str(tmpdir)
+    num_experiments = 3
+    for i in range(num_experiments):
+        experiment_name = "test_lsx{}".format(i)
+        tune.run(
+            MyTrainableClass,
+            name=experiment_name,
+            stop={"training_iteration": 1},
+            num_samples=1,
+            storage_path=project_path,
+        )
+
+    limit = 2
+    with Capturing() as output:
+        commands.list_experiments(
+            project_path, info_keys=("total_trials",), limit=limit
+        )
+    lines = output.captured
+    assert "total_trials" in lines[1]
+    assert lines[1].count("|") == 2
+    assert len(lines) == 3 + limit + 1
+
+    with Capturing() as output:
+        commands.list_experiments(
+            project_path,
+            sort=["total_trials"],
+            info_keys=("total_trials",),
+            filter_op="total_trials == 1",
+        )
+    lines = output.captured
+    assert sum("1" in line for line in lines) >= num_experiments
+    assert len(lines) == 3 + num_experiments + 1
+
+
+if __name__ == "__main__":
+    # Make click happy in bazel.
+    os.environ["LC_ALL"] = "en_US.UTF-8"
+    os.environ["LANG"] = "en_US.UTF-8"
+    sys.exit(pytest.main([__file__]))