import argparse import sys from unittest import mock import pytest from freezegun import freeze_time from ray import tune from ray.air.constants import TRAINING_ITERATION from ray.tune.experiment.trial import Trial from ray.tune.experimental.output import ( AirVerbosity, TrainReporter, TuneTerminalReporter, _best_trial_str, _current_best_trial, _get_dict_as_table_data, _get_time_str, _get_trial_info, _get_trial_table_data, _get_trials_by_state, _infer_params, _infer_user_metrics, _max_len, ) from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME LAST_RESULT = { "custom_metrics": {}, "episode_media": {}, "info": { "learner": { "default_policy": { "allreduce_latency": 0.0, "grad_gnorm": 40.0, "cur_lr": 0.001, "total_loss": 93.35336303710938, "policy_loss": -18.39633560180664, "entropy": 0.5613694190979004, "entropy_coeff": 0.01, "var_gnorm": 23.452943801879883, "vf_loss": 223.5106201171875, "vf_explained_var": -0.0017577409744262695, "mean_IS": 0.9987365007400513, "var_IS": 0.0007558994111604989, }, } }, "sampler_results": { "episode_reward_max": 500.0, "episode_reward_min": 54.0, "episode_reward_mean": 214.45, }, "episode_reward_max": 500.0, "episode_reward_min": 54.0, "episode_reward_mean": 214.45, "episode_len_mean": 214.45, "episodes_this_iter": 66, "timesteps_total": 33000, } @freeze_time("Mar 27th, 2023", auto_tick_seconds=15) def test_get_time_str(): base = 1679875200 # 2023-03-27 00:00:00 assert _get_time_str(base, base) == ("2023-03-27 00:00:00", "0s") assert _get_time_str(base, base + 15) == ("2023-03-27 00:00:15", "15s") assert _get_time_str(base, base + 60) == ("2023-03-27 00:01:00", "1min 0s") assert _get_time_str(base, base + 65) == ("2023-03-27 00:01:05", "1min 5s") assert _get_time_str(base, base + 3600) == ( "2023-03-27 01:00:00", "1hr 0min 0s", ) assert _get_time_str(base, base + 3605) == ( "2023-03-27 01:00:05", "1hr 0min 5s", ) assert _get_time_str(base, base + 3660) == ( "2023-03-27 01:01:00", "1hr 1min 0s", ) assert _get_time_str(base, base + 86400) == ( "2023-03-28 00:00:00", "1d 0hr 0min 0s", ) def test_get_trials_by_state(): t1 = Trial(MOCK_TRAINABLE_NAME, stub=True) t1.set_status(Trial.RUNNING) t2 = Trial(MOCK_TRAINABLE_NAME, stub=True) t2.set_status(Trial.PENDING) trials = [t1, t2] assert _get_trials_by_state(trials) == {"RUNNING": [t1], "PENDING": [t2]} def test_infer_user_metrics(): t = Trial(MOCK_TRAINABLE_NAME, stub=True) t.run_metadata.last_result = LAST_RESULT result = [ "episode_reward_max", "episode_reward_min", "episode_len_mean", "episodes_this_iter", ] assert _infer_user_metrics([t]) == result def test_max_len(): assert _max_len("long_metrics_name", max_len=5) == "...me" assert _max_len("long_metrics_name", max_len=10) == "...cs_name" assert _max_len("long_metrics_name", max_len=9, wrap=True) == "long_metr\nics_name" assert _max_len("long_metrics_name", max_len=8, wrap=True) == "..._metr\nics_name" def test_current_best_trial(): t1 = Trial(MOCK_TRAINABLE_NAME, stub=True) t2 = Trial(MOCK_TRAINABLE_NAME, stub=True) t1.run_metadata.last_result = {"metric": 2} t2.run_metadata.last_result = {"metric": 1} assert _current_best_trial([t1, t2], metric="metric", mode="min") == (t2, "metric") def test_best_trial_str(): t = Trial(MOCK_TRAINABLE_NAME, stub=True) t.trial_id = "18ae7_00005" t.run_metadata.last_result = { "loss": 0.5918508041056858, "config": {"train_loop_config": {"lr": 0.059253447253394785}}, } assert ( _best_trial_str(t, "loss") == "Current best trial: 18ae7_00005 with loss=0.5918508041056858" " and params={'train_loop_config': {'lr': 0.059253447253394785}}" ) def test_get_trial_info(): t = Trial(MOCK_TRAINABLE_NAME, stub=True) t.trial_id = "af42b609" t.set_status(Trial.RUNNING) t.run_metadata.last_result = LAST_RESULT assert _get_trial_info( t, param_keys=[], metric_keys=[ "episode_reward_mean", "episode_reward_max", "episode_reward_min", "episode_len_mean", "episodes_this_iter", ], ) == ["mock_trainable_af42b609", "RUNNING", 214.45, 500.0, 54.0, 214.45, 66] def test_get_trial_table_data_less_than_20(): trials = [] for i in range(20): t = Trial(MOCK_TRAINABLE_NAME, stub=True) t.trial_id = str(i) t.set_status(Trial.RUNNING) t.run_metadata.last_result = {"episode_reward_mean": 100 + i} t.config = {"param": i} trials.append(t) table_data = _get_trial_table_data(trials, ["param"], ["episode_reward_mean"]) header = table_data.header assert header == ["Trial name", "status", "param", "reward"] table_data = table_data.data assert len(table_data) == 1 # only the running category assert len(table_data[0].trial_infos) == 20 assert not table_data[0].more_info def test_get_trial_table_data_more_than_20(): trials = [] # total of 30 trials. for status in [Trial.RUNNING, Trial.TERMINATED, Trial.PENDING]: for i in range(10): t = Trial(MOCK_TRAINABLE_NAME, stub=True) t.trial_id = str(i) t.set_status(status) t.run_metadata.last_result = {"episode_reward_mean": 100 + i} t.config = {"param": i} trials.append(t) table_data = _get_trial_table_data(trials, ["param"], ["episode_reward_mean"]) header = table_data.header assert header == ["Trial name", "status", "param", "reward"] table_data = table_data.data assert len(table_data) == 3 # only the running category for i in range(3): assert len(table_data[i].trial_infos) == 5 assert table_data[0].more_info == "5 more RUNNING" assert table_data[1].more_info == "5 more TERMINATED" assert table_data[2].more_info == "5 more PENDING" def test_infer_params(): assert _infer_params({}) == [] assert _infer_params({"some": "val"}) == [] assert _infer_params({"some": "val", "param": tune.uniform(0, 1)}) == ["param"] assert _infer_params({"some": "val", "param": tune.grid_search([0, 1])}) == [ "param" ] assert sorted( _infer_params( { "some": "val", "param": tune.grid_search([0, 1]), "other": tune.choice([0, 1]), } ) ) == ["other", "param"] def test_result_table_no_divison(): data = _get_dict_as_table_data( { "b": 6, "a": 8, "x": 19.123123123, "c": 5, "ignore": 9, "nested_ignore": {"value": 5}, "y": 20, "z": {"m": 4, "n": {"o": "p"}}, }, exclude={"ignore", "nested_ignore"}, ) assert data == [ ["a", 8], ["b", 6], ["c", 5], ["x", "19.12312"], ["y", 20], ["z/m", 4], ["z/n/o", "p"], ] def test_result_table_divison(): data = _get_dict_as_table_data( { "b": 6, "a": 8, "x": 19.123123123, "c": 5, "ignore": 9, "nested_ignore": {"value": 5}, "y": 20, "z": {"m": 4, "n": {"o": "p"}}, }, exclude={"ignore", "nested_ignore"}, upper_keys={"x", "y", "z", "z/m", "z/n/o"}, ) assert data == [ ["x", "19.12312"], ["y", 20], ["z/m", 4], ["z/n/o", "p"], ["a", 8], ["b", 6], ["c", 5], ] def test_result_include(): data = _get_dict_as_table_data( { "b": 6, "a": 8, "x": 19.123123123, "c": 5, "ignore": 9, "nested_ignore": {"value": 5}, "y": 20, "z": {"m": 4, "n": {"o": "p"}}, }, include={"y", "z"}, exclude={"z/n/o"}, ) assert data == [ ["y", 20], ["z/m", 4], ] def test_config_argparse(): parser = argparse.ArgumentParser() parser.add_argument("--bool-val", action="store_true", default=True) parser.add_argument("--foo", default="bar") args = parser.parse_args([]) data = _get_dict_as_table_data({"parsed_args": args}) assert data == [ ["parsed_args/bool_val", True], ["parsed_args/foo", "bar"], ] @pytest.mark.parametrize("progress_reporter_cls", [TrainReporter, TuneTerminalReporter]) def test_heartbeat_reset(progress_reporter_cls): """Test heartbeat functionality in train and tune. Tune prints a table every `heartbeat_freq` seconds. Train prints a heartbeat every `heartbeat_freq` seconds, but a result also resets the counter. """ # Train heartbeats are only reporter in VERBOSE reporter = progress_reporter_cls(verbosity=AirVerbosity.VERBOSE) reporter._print_heartbeat = mock.MagicMock() with freeze_time() as frozen: reporter.print_heartbeat([]) assert reporter._print_heartbeat.call_count == 1 # Tick until heartbeat freq. Next call to print_heartbeat should trigger frozen.tick(reporter._heartbeat_freq) reporter.print_heartbeat([]) assert reporter._print_heartbeat.call_count == 2 # Not quite there, yet. This should not trigger a heartbeat. frozen.tick(reporter._heartbeat_freq // 2) reporter.print_heartbeat([]) assert reporter._print_heartbeat.call_count == 2 # Let's report a result. This will reset the heartbeat timer reporter.on_trial_result( 0, [], Trial(MOCK_TRAINABLE_NAME, stub=True), {TRAINING_ITERATION: 1} ) # Progress another half heartbeat. In Tune this triggers a heartbeat, # but in train the heartbeat is reset on trial result. frozen.tick(reporter._heartbeat_freq // 2 + 1) reporter.print_heartbeat([]) if progress_reporter_cls == TrainReporter: # Thus, train shouldn't have reported assert reporter._print_heartbeat.call_count == 2 elif progress_reporter_cls == TuneTerminalReporter: # But Tune should have. assert reporter._print_heartbeat.call_count == 3 else: raise RuntimeError("Test faulty.") if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))