import json import os import pickle import random import sys import tempfile import time import unittest from functools import partial from typing import List, Optional from unittest.mock import MagicMock import numpy as np import pytest import ray from ray import cloudpickle, tune from ray._private.test_utils import object_memory_usage from ray.tune import ( Callback, Checkpoint, CheckpointConfig, FailureConfig, RunConfig, Trainable, ) from ray.tune.experiment import Trial from ray.tune.schedulers import PopulationBasedTraining from ray.tune.schedulers.pb2 import PB2 from ray.tune.schedulers.pb2_utils import UCB from ray.tune.schedulers.pbt import _filter_mutated_params_from_config from ray.tune.tests.execution.utils import create_execution_test_objects from ray.tune.tune_config import TuneConfig from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable from ray.tune.utils.util import flatten_dict # Import psutil after ray so the packaged version is used. import psutil MB = 1024**2 class MockParam(object): def __init__(self, params): self._params = params self._index = 0 def __call__(self, *args, **kwargs): val = self._params[self._index % len(self._params)] self._index += 1 return val class DummyTrial: def __init__(self, trial_id, finished=False, config: Optional[dict] = None): self.trial_id = trial_id self._finished = finished self.config = config if config is not None else {} def is_finished(self): return self._finished class PopulationBasedTrainingMemoryTest(unittest.TestCase): def setUp(self): ray.init(num_cpus=1, object_store_memory=100 * MB) def tearDown(self): ray.shutdown() def testMemoryCheckpointFree(self): class MyTrainable(Trainable): def setup(self, config): # Make sure this is large enough so ray uses object store # instead of in-process store. self.large_object = random.getrandbits(int(10e6)) self.iter = 0 self.a = config["a"] def step(self): self.iter += 1 return {"metric": self.iter + self.a} def save_checkpoint(self, checkpoint_dir): file_path = os.path.join(checkpoint_dir, "model.mock") with open(file_path, "wb") as fp: pickle.dump((self.large_object, self.iter, self.a), fp) def load_checkpoint(self, checkpoint_dir): file_path = os.path.join(checkpoint_dir, "model.mock") with open(file_path, "rb") as fp: self.large_object, self.iter, self.a = pickle.load(fp) class CheckObjectMemoryUsage(Callback): def on_trial_save( self, iteration: int, trials: List["Trial"], trial: "Trial", **info ): assert object_memory_usage() <= (12 * 80e6) param_a = MockParam([1, -1]) pbt = PopulationBasedTraining( time_attr="training_iteration", metric="metric", mode="max", perturbation_interval=1, hyperparam_mutations={"b": [-1]}, ) tune.run( MyTrainable, name="ray_demo", scheduler=pbt, stop={"training_iteration": 10}, num_samples=3, checkpoint_config=CheckpointConfig(checkpoint_frequency=3), fail_fast=True, config={"a": tune.sample_from(lambda _: param_a())}, callbacks=[CheckObjectMemoryUsage()], ) class PopulationBasedTrainingFileDescriptorTest(unittest.TestCase): def setUp(self): ray.init(num_cpus=2) os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "1" def tearDown(self): ray.shutdown() def testFileFree(self): class MyTrainable(Trainable): def setup(self, config): self.iter = 0 self.a = config["a"] def step(self): self.iter += 1 return {"metric": self.iter + self.a} def save_checkpoint(self, checkpoint_dir): file_path = os.path.join(checkpoint_dir, "model.mock") with open(file_path, "wb") as fp: pickle.dump((self.iter, self.a), fp) def load_checkpoint(self, checkpoint_dir): file_path = os.path.join(checkpoint_dir, "model.mock") with open(file_path, "rb") as fp: self.iter, self.a = pickle.load(fp) from ray.tune.callback import Callback class FileCheck(Callback): def __init__(self, verbose=False): self.iter_ = 0 self.process = psutil.Process() self.verbose = verbose def on_trial_result(self, *args, **kwargs): self.iter_ += 1 all_files = self.process.open_files() if self.verbose: print("Iteration", self.iter_) print("=" * 10) print("Object memory use: ", object_memory_usage()) print("Virtual Mem:", self.get_virt_mem() >> 30, "gb") print("File Descriptors:", len(all_files)) assert len(all_files) < 20 @classmethod def get_virt_mem(cls): return psutil.virtual_memory().used param_a = MockParam([1, -1]) pbt = PopulationBasedTraining( time_attr="training_iteration", metric="metric", mode="max", perturbation_interval=1, quantile_fraction=0.5, hyperparam_mutations={"b": [-1]}, ) checkpoint_config = CheckpointConfig( num_to_keep=3, checkpoint_frequency=2, ) tune.run( MyTrainable, name="ray_demo", scheduler=pbt, stop={"training_iteration": 10}, num_samples=4, checkpoint_config=checkpoint_config, verbose=False, fail_fast=True, config={"a": tune.sample_from(lambda _: param_a())}, callbacks=[FileCheck()], ) class PopulationBasedTrainingSynchTest(unittest.TestCase): def setUp(self): ray.init(num_cpus=2) def train_fn_sync(config): iter = 0 checkpoint = tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: checkpoint_path = os.path.join(checkpoint_dir, "checkpoint") with open(checkpoint_path, "rb") as fp: a, iter = pickle.load(fp) a = config["a"] # Use the new hyperparameter if perturbed. while True: iter += 1 with tempfile.TemporaryDirectory() as checkpoint_dir: checkpoint_path = os.path.join(checkpoint_dir, "checkpoint") with open(checkpoint_path, "wb") as fp: pickle.dump((a, iter), fp) # Different sleep times so that asynch test runs do not # randomly succeed. If well performing trials finish later, # then bad performing trials will already have continued # to train, which is exactly what we want to test when # comparing sync vs. async. time.sleep(a / 20) # Score gets better every iteration. tune.report( {"mean_accuracy": iter + a, "a": a}, checkpoint=Checkpoint.from_directory(checkpoint_dir), ) self.MockTrainingFuncSync = train_fn_sync def tearDown(self): ray.shutdown() def synchSetup(self, synch, param=None): if param is None: param = [10, 20, 40] scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="mean_accuracy", mode="max", perturbation_interval=1, log_config=True, hyperparam_mutations={"c": lambda: 1}, synch=synch, ) param_a = MockParam(param) random.seed(100) np.random.seed(100) analysis = tune.run( self.MockTrainingFuncSync, config={"a": tune.sample_from(lambda _: param_a()), "c": 1}, fail_fast=True, num_samples=3, scheduler=scheduler, name="testPBTSync", stop={"training_iteration": 3}, ) return analysis def testAsynchFail(self): analysis = self.synchSetup(False) self.assertTrue( any( analysis.dataframe(metric="mean_accuracy", mode="max")["mean_accuracy"] != 43 ) ) def testSynchPass(self): analysis = self.synchSetup(True) all_results = set( analysis.dataframe(metric="mean_accuracy", mode="max")["mean_accuracy"] ) self.assertEqual(all_results, {43}) def testSynchPassLast(self): analysis = self.synchSetup(True, param=[30, 20, 10]) all_results = set( analysis.dataframe(metric="mean_accuracy", mode="max")["mean_accuracy"] ) self.assertEqual(all_results, {33}) def testExploitWhileSavingTrial(self): """Tests a synch PBT failure mode where a trial misses its `SAVING_RESULT` event book-keeping due to being stopped by the PBT algorithm (to exploit another trial). Trials checkpoint ever N iterations, and the perturbation interval is every N iterations. (N = 2 in the test.) Raises a `TimeoutError` if hanging for a specified `timeout`. 1. Trial 0 comes in with training result 2. Trial 0 begins saving checkpoint (which may take a long time, 5s here) 3. Trial 1 comes in with result 4. Trial 1 forcefully stops Trial 0 via exploit, while trial_0.is_saving 5. Trial 0 should resume training properly with Trial 1's checkpoint """ class MockTrainable(tune.Trainable): def setup(self, config): self.reset_config(config) def step(self): time.sleep(self.training_time) return {"score": self.score} def save_checkpoint(self, checkpoint_dir): with open(os.path.join(checkpoint_dir, "checkpoint.json"), "w") as f: json.dump({"a": self.a}, f) time.sleep(self.saving_time) def load_checkpoint(self, checkpoint_dir): with open(os.path.join(checkpoint_dir, "checkpoint.json"), "r") as f: checkpoint_dict = json.load(f) self.a = checkpoint_dict["a"] def reset_config(self, new_config): self.a = new_config["a"] self.score = new_config["score"] self.training_time = new_config["training_time"] self.saving_time = new_config["saving_time"] return True perturbation_interval = 2 scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="score", mode="max", perturbation_interval=perturbation_interval, hyperparam_mutations={"a": tune.uniform(0, 1)}, synch=True, ) class TimeoutExceptionStopper(tune.stopper.TimeoutStopper): def stop_all(self): decision = super().stop_all() if decision: raise TimeoutError("Trials are hanging! Timeout reached...") return decision timeout = 30.0 training_times = [0.1, 0.15] saving_times = [5.0, 0.1] tuner = tune.Tuner( MockTrainable, param_space={ "a": tune.uniform(0, 1), "score": tune.grid_search([0, 1]), "training_time": tune.sample_from( lambda config: training_times[config["score"]] ), "saving_time": tune.sample_from( lambda config: saving_times[config["score"]] ), }, tune_config=TuneConfig( num_samples=1, scheduler=scheduler, ), run_config=RunConfig( stop=tune.stopper.CombinedStopper( tune.stopper.MaximumIterationStopper(5), TimeoutExceptionStopper(timeout), ), failure_config=FailureConfig(fail_fast=True), checkpoint_config=CheckpointConfig( # Match `checkpoint_interval` with `perturbation_interval` checkpoint_frequency=perturbation_interval, ), ), ) random.seed(100) np.random.seed(1000) results = tuner.fit() assert not results.errors class PopulationBasedTrainingConfigTest(unittest.TestCase): def setUp(self): ray.init(num_cpus=2) def tearDown(self): ray.shutdown() def testNoConfig(self): def MockTrainingFunc(config): a = config["a"] b = config["b"] c1 = config["c"]["c1"] c2 = config["c"]["c2"] while True: tune.report({"mean_accuracy": a * b * (c1 + c2)}) scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="mean_accuracy", mode="max", perturbation_interval=1, hyperparam_mutations={ "a": tune.uniform(0, 0.3), "b": [1, 2, 3], "c": { "c1": lambda: np.random.uniform(0.5), "c2": tune.choice([2, 3, 4]), }, }, ) tune.run( MockTrainingFunc, fail_fast=True, num_samples=4, scheduler=scheduler, name="testNoConfig", stop={"training_iteration": 3}, ) class PopulationBasedTrainingResumeTest(unittest.TestCase): def setUp(self): ray.init(num_cpus=2) def tearDown(self): ray.shutdown() def testPermutationContinuation(self): """ Tests continuation of runs after permutation. Sometimes, runs were continued from deleted checkpoints. This deterministic initialisation would fail when the fix was not applied. See issues #9036, #9036 """ class MockTrainable(tune.Trainable): def setup(self, config): self.iter = 0 self.a = config["a"] self.b = config["b"] self.c = config["c"] def step(self): self.iter += 1 return {"mean_accuracy": (self.a - self.iter) * self.b} def save_checkpoint(self, tmp_checkpoint_dir): checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.mock") with open(checkpoint_path, "wb") as fp: pickle.dump((self.a, self.b, self.iter), fp) def load_checkpoint(self, tmp_checkpoint_dir): checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.mock") with open(checkpoint_path, "rb") as fp: self.a, self.b, self.iter = pickle.load(fp) scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="mean_accuracy", mode="max", perturbation_interval=1, log_config=True, hyperparam_mutations={"c": lambda: 1}, ) param_a = MockParam([10, 20, 30, 40]) param_b = MockParam([1.2, 0.9, 1.1, 0.8]) random.seed(100) np.random.seed(1000) checkpoint_config = CheckpointConfig( num_to_keep=2, checkpoint_score_attribute="min-training_iteration", checkpoint_frequency=1, checkpoint_at_end=True, ) tune.run( MockTrainable, config={ "a": tune.sample_from(lambda _: param_a()), "b": tune.sample_from(lambda _: param_b()), "c": 1, }, fail_fast=True, num_samples=4, checkpoint_config=checkpoint_config, scheduler=scheduler, name="testPermutationContinuation", stop={"training_iteration": 3}, ) def testPermutationContinuationFunc(self): def MockTrainingFunc(config): iter = 0 a = config["a"] b = config["b"] if tune.get_checkpoint(): with tune.get_checkpoint().as_directory() as checkpoint_dir: checkpoint_path = os.path.join(checkpoint_dir, "model.mock") with open(checkpoint_path, "rb") as fp: a, b, iter = pickle.load(fp) while True: iter += 1 with tempfile.TemporaryDirectory() as checkpoint_dir: checkpoint_path = os.path.join(checkpoint_dir, "model.mock") with open(checkpoint_path, "wb") as fp: pickle.dump((a, b, iter), fp) tune.report( {"mean_accuracy": (a - iter) * b}, checkpoint=Checkpoint.from_directory(checkpoint_dir), ) scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="mean_accuracy", mode="max", perturbation_interval=1, log_config=True, hyperparam_mutations={"c": lambda: 1}, ) param_a = MockParam([10, 20, 30, 40]) param_b = MockParam([1.2, 0.9, 1.1, 0.8]) random.seed(100) np.random.seed(1000) checkpoint_config = CheckpointConfig( num_to_keep=2, checkpoint_score_attribute="min-training_iteration", ) tune.run( MockTrainingFunc, config={ "a": tune.sample_from(lambda _: param_a()), "b": tune.sample_from(lambda _: param_b()), "c": 1, }, fail_fast=True, num_samples=4, checkpoint_config=checkpoint_config, scheduler=scheduler, name="testPermutationContinuationFunc", stop={"training_iteration": 3}, ) def testBurnInPeriod(self): runner, *_ = create_execution_test_objects() storage_context = runner._storage scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="error", mode="min", perturbation_interval=5, hyperparam_mutations={"ignored": [1]}, burn_in_period=50, log_config=True, synch=True, ) class MockTrial(Trial): @property def checkpoint(self): return Checkpoint.from_directory("dummy") @property def status(self): return Trial.PAUSED @status.setter def status(self, status): pass register_mock_trainable() trials = [ MockTrial(MOCK_TRAINABLE_NAME, config=dict(num=i), storage=storage_context) for i in range(1, 5) ] trial1, trial2, trial3, trial4 = trials for trial in trials: trial.init_local_path() runner.add_trial(trial) for trial in trials: scheduler.on_trial_add(runner, trial) # Add initial results. scheduler.on_trial_result( runner, trial1, result=dict(training_iteration=1, error=50) ) scheduler.on_trial_result( runner, trial2, result=dict(training_iteration=1, error=50) ) scheduler.on_trial_result( runner, trial3, result=dict(training_iteration=1, error=10) ) scheduler.on_trial_result( runner, trial4, result=dict(training_iteration=1, error=100) ) # Add more results. Without burn-in, this would now exploit scheduler.on_trial_result( runner, trial1, result=dict(training_iteration=30, error=50) ) scheduler.on_trial_result( runner, trial2, result=dict(training_iteration=30, error=50) ) scheduler.on_trial_result( runner, trial3, result=dict(training_iteration=30, error=10) ) scheduler.on_trial_result( runner, trial4, result=dict(training_iteration=30, error=100) ) self.assertEqual(trial4.config["num"], 4) # Add more results. Since this is after burn-in, it should now exploit scheduler.on_trial_result( runner, trial1, result=dict(training_iteration=50, error=50) ) scheduler.on_trial_result( runner, trial2, result=dict(training_iteration=50, error=50) ) scheduler.on_trial_result( runner, trial3, result=dict(training_iteration=50, error=10) ) scheduler.on_trial_result( runner, trial4, result=dict(training_iteration=50, error=100) ) self.assertEqual(trial4.config["num"], 3) # Assert that trials do not hang after `burn_in_period` self.assertTrue(all(t.status == "PAUSED" for t in runner.get_trials())) self.assertTrue(scheduler.choose_trial_to_run(runner)) # Assert that trials do not hang when a terminated trial is added trial5 = Trial(MOCK_TRAINABLE_NAME, config=dict(num=5)) runner.add_trial(trial5) scheduler.on_trial_add(runner, trial5) trial5.set_status(Trial.TERMINATED) self.assertTrue(scheduler.choose_trial_to_run(runner)) class PopulationBasedTrainingLoggingTest(unittest.TestCase): def testFilterHyperparamConfig(self): filtered_params = _filter_mutated_params_from_config( { "training_loop_config": { "lr": 0.1, "momentum": 0.9, "batch_size": 32, "test_mode": True, "ignore_nested": { "b": 0.1, }, }, "other_config": { "a": 0.5, }, }, { "training_loop_config": { "lr": tune.uniform(0, 1), "momentum": tune.uniform(0, 1), } }, ) assert filtered_params == { "training_loop_config": {"lr": 0.1, "momentum": 0.9} }, filtered_params def testSummarizeHyperparamChanges(self): def test_config( hyperparam_mutations, old_config, resample_probability=0.25, print_summary=False, ): scheduler = PopulationBasedTraining( time_attr="training_iteration", hyperparam_mutations=hyperparam_mutations, resample_probability=resample_probability, ) new_config, operations = scheduler._get_new_config( None, DummyTrial("parent_id", config=old_config) ) old_params = _filter_mutated_params_from_config( old_config, hyperparam_mutations ) new_params = _filter_mutated_params_from_config( new_config, hyperparam_mutations ) summary = scheduler._summarize_hyperparam_changes( old_params, new_params, operations ) if print_summary: print(summary) return scheduler, new_config, operations # 1. Empty hyperparam_mutations (no hyperparams mutated) should raise an error with self.assertRaises(tune.TuneError): _, new_config, operations = test_config({}, {}) # 2. No nesting hyperparam_mutations = { "a": tune.uniform(0, 1), "b": list(range(5)), } scheduler, new_config, operations = test_config( hyperparam_mutations, {"a": 0.5, "b": 2} ) assert operations["a"] in [ f"* {factor}" for factor in scheduler._perturbation_factors ] + ["resample"] assert operations["b"] in ["shift left", "shift right", "resample"] # 3. With nesting hyperparam_mutations = { "a": tune.uniform(0, 1), "b": list(range(5)), "c": { "d": tune.uniform(2, 3), "e": {"f": [-1, 0, 1]}, }, } scheduler, new_config, operations = test_config( hyperparam_mutations, { "a": 0.5, "b": 2, "c": { "d": 2.5, "e": {"f": 0}, }, }, ) assert isinstance(operations["c"], dict) assert isinstance(operations["c"]["e"], dict) assert operations["c"]["d"] in [ f"* {factor}" for factor in scheduler._perturbation_factors ] + ["resample"] assert operations["c"]["e"]["f"] in ["shift left", "shift right", "resample"] # 4. Test shift that results in noop hyperparam_mutations = {"a": [1]} scheduler, new_config, operations = test_config( hyperparam_mutations, {"a": 1}, resample_probability=0 ) assert operations["a"] in ["shift left (noop)", "shift right (noop)"] # 5. Test that missing keys in inputs raises an error with self.assertRaises(AssertionError): scheduler._summarize_hyperparam_changes( {"a": 1, "b": {"c": 2}}, {"a": 1, "b": {}}, {"a": "noop", "b": {"c": "noop"}}, ) # It's ok to have missing operations (just fill in the ones that are present) scheduler._summarize_hyperparam_changes( {"a": 1, "b": {"c": 2}}, {"a": 1, "b": {"c": 2}}, {"a": "noop"} ) scheduler._summarize_hyperparam_changes( {"a": 1, "b": {"c": 2}}, {"a": 1, "b": {"c": 2}}, {} ) # Make sure that perturbation and logging work with extra keys that aren't # included in hyperparam_mutations (both should ignore the keys) hyperparam_mutations = { "train_loop_config": { "lr": tune.uniform(0, 1), "momentum": tune.uniform(0, 1), } } test_config( hyperparam_mutations, { "train_loop_config": { "lr": 0.1, "momentum": 0.9, "batch_size": 32, "test_mode": True, } }, resample_probability=0, ) class DummyState: def __init__(self, last_score): self.last_score = last_score class PopulationBasedTrainingNanScoreTest(unittest.TestCase): def test_pbt_with_nan_scores(self): # Create three trials: one with nan, two with valid scores t1 = DummyTrial("t1", config=MagicMock()) t2 = DummyTrial("t2", config=MagicMock()) t3 = DummyTrial("t3", config=MagicMock()) for scheduler_class in (PopulationBasedTraining, PB2): # Patch _trial_state with dummy states # Note: list.sort does not change the order if nan is present max_states = { t1: DummyState(last_score=20.0), t2: DummyState(last_score=float("nan")), t3: DummyState(last_score=10.0), } min_states = { t3: DummyState(last_score=10.0), t2: DummyState(last_score=float("nan")), t1: DummyState(last_score=20.0), } with self.subTest(scheduler_class=scheduler_class.__name__): if scheduler_class is PopulationBasedTraining: hp_kwargs = {"hyperparam_mutations": {"lr": [1e-4, 1e-3]}} else: hp_kwargs = {"hyperparam_bounds": {"lr": [1e-4, 1e-3]}} # test max mode max_scheduler = scheduler_class( metric="reward", mode="max", quantile_fraction=0.5, **hp_kwargs, ) max_scheduler._trial_state = max_states for t, state in max_states.items(): max_scheduler._save_trial_state( state, 100, {"reward": state.last_score, "time_total_s": 1}, t ) # Should not raise, but nan disrupts sorting max_bottom, max_top = max_scheduler._quantiles() max_other_trials = [ t for t in max_scheduler._trial_state if t not in max_bottom + max_top ] max_ordered_results = [ max_scheduler._trial_state[t].last_score for t in [*max_bottom, *max_other_trials, *max_top] ] self.assertIn(t1, max_top) self.assertIn(t2, max_other_trials) self.assertIn(t3, max_bottom) self.assertEqual(max_ordered_results[-1], 20) # Test min mode min_scheduler = scheduler_class( metric="reward", mode="min", quantile_fraction=0.5, **hp_kwargs, ) min_scheduler._trial_state = min_states for t, state in min_states.items(): min_scheduler._save_trial_state( state, 100, {"reward": state.last_score, "time_total_s": 1}, t ) min_bottom, min_top = min_scheduler._quantiles() min_other_trials = [ t for t in min_scheduler._trial_state if t not in min_bottom + min_top ] min_ordered_results = [ min_scheduler._trial_state[t].last_score for t in [*min_bottom, *min_other_trials, *min_top] ] self.assertIn(t1, min_bottom) self.assertIn(t2, min_other_trials) self.assertIn(t3, min_top) self.assertEqual(abs(min_ordered_results[-1]), 10) def _create_pb2_scheduler( metric="score", mode="max", perturbation_interval=1, hyperparam_bounds=None, custom_explore_fn=None, ) -> PB2: hyperparam_bounds = hyperparam_bounds or {"a": [0.0, 1.0]} return PB2( metric=metric, mode=mode, time_attr="training_iteration", perturbation_interval=perturbation_interval, quantile_fraction=0.25, hyperparam_bounds=hyperparam_bounds, custom_explore_fn=custom_explore_fn, ) def _save_trial_result(scheduler: PB2, trial: Trial, time: int, result: dict): scheduler._save_trial_state(scheduler._trial_state[trial], time, result, trial) def _result(time: int, val: float) -> dict: """Creates a dummy Tune result to report.""" return {"training_iteration": time, "score": val} def test_pb2_perturbation(monkeypatch): hyperparam_bounds = {"a": [1.0, 2.0]} pb2 = _create_pb2_scheduler( metric="score", mode="max", hyperparam_bounds=hyperparam_bounds ) mock_runner = MagicMock() # One trial at each end of the hyperparam bounds, one performing better than the # other. We expect a perturbed value to be closer to the better performing one. trials = [ Trial("pb2_test", stub=True, config={"a": 1.0}), Trial("pb2_test", stub=True, config={"a": 2.0}), ] for trial in trials: pb2.on_trial_add(mock_runner, trial) # Collect 10 timesteps of data # PB2 fits a model to estimate the increase in score between timesteps # Each timestep, trial 1's score increases by 10, trial 2's score increases by 20 for t in range(1, 11): for i, trial in enumerate(trials): _save_trial_result(pb2, trial, t, _result(time=t, val=t * (i + 1) * 10)) # Ignoring variance (kappa=0) and only optimizing for exploitation, # we expect the next point suggested to be close to higher-performing trial monkeypatch.setattr(ray.tune.schedulers.pb2_utils, "UCB", partial(UCB, kappa=0.0)) new_config, _ = pb2._get_new_config(trials[0], trials[1]) assert new_config["a"] > 1.5 assert pb2._quantiles() == ([trials[0]], [trials[1]]) def test_pb2_nested_hyperparams(): """Test that PB2 with nested hyperparams behaves the same as without nesting.""" hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [2.0, 4.0], "d": [4.0, 10.0]}} pb2_nested = _create_pb2_scheduler( metric="score", mode="max", hyperparam_bounds=hyperparam_bounds, ) pb2_flat = _create_pb2_scheduler( metric="score", mode="max", hyperparam_bounds=flatten_dict(hyperparam_bounds, delimiter=""), ) mock_runner = MagicMock() trials_nested = [Trial("pb2_test", stub=True) for _ in range(3)] trials_flat = [Trial("pb2_test", stub=True) for _ in range(3)] np.random.seed(2023) for trial_nested, trial_flat in zip(trials_nested, trials_flat): pb2_nested.on_trial_add(mock_runner, trial_nested) # Let PB2 generate the initial config randomly, then use the same # initial values for the flattened version flattened_init_config = flatten_dict(trial_nested.config, delimiter="") trial_flat.config = flattened_init_config pb2_flat.on_trial_add(mock_runner, trial_flat) # Make sure that config suggestions are the same for each timestep for t in range(1, 10): for i, (trial_nested, trial_flat) in enumerate(zip(trials_nested, trials_flat)): res = _result(time=t, val=t * (i + 1) * 10) _save_trial_result(pb2_nested, trial_nested, t, res) _save_trial_result(pb2_flat, trial_flat, t, res) # Reset seed before each _get_new_config call so both schedulers # get the same random sequence in optimize_acq np.random.seed(2023 + t) new_config, _ = pb2_nested._get_new_config(trials_nested[0], trials_nested[-1]) np.random.seed(2023 + t) new_config_flat, _ = pb2_flat._get_new_config(trials_flat[0], trials_flat[-1]) # Make sure the suggested config is still nested properly assert list(new_config.keys()) == ["a", "b"] assert list(new_config["b"].keys()) == ["c", "d"] assert np.allclose( list(flatten_dict(new_config, delimiter="").values()), list(new_config_flat.values()), ) def test_pb2_missing_hyperparam_init(): """Test that PB2 fills in all missing hyperparameters (those that are not specified in param_space).""" hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [2.0, 4.0], "d": [4.0, 10.0]}} pb2 = _create_pb2_scheduler(hyperparam_bounds=hyperparam_bounds) mock_runner = MagicMock() def validate_config(config, bounds): for param, bound in bounds.items(): if isinstance(bound, dict): validate_config(config[param], bound) else: low, high = bound assert config[param] >= low and config[param] < high trial = Trial("test_pb2", stub=True) pb2.on_trial_add(mock_runner, trial) validate_config(trial.config, hyperparam_bounds) trial = Trial("test_pb2", stub=True, config={"b": {"c": 3.0}}) pb2.on_trial_add(mock_runner, trial) validate_config(trial.config, hyperparam_bounds) assert trial.config["b"]["c"] == 3.0 def test_pb2_hyperparam_bounds_validation(): """Check that hyperparam bounds are validated (must be tuples of [low, high]).""" # Too many values hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [2.0, 4.0, 6.0]}} with pytest.raises(ValueError): _create_pb2_scheduler(hyperparam_bounds=hyperparam_bounds) # Ordering is wrong hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [4.0, 2.0]}} with pytest.raises(ValueError): _create_pb2_scheduler(hyperparam_bounds=hyperparam_bounds) def test_pb2_custom_explore_fn(): """Test custom post-processing on the config generated by PB2.""" hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [2.0, 4.0], "d": [4.0, 10.0]}} def explore(config): config["b"]["c"] = int(config["b"]["c"]) return config pb2 = _create_pb2_scheduler( hyperparam_bounds=hyperparam_bounds, custom_explore_fn=explore, ) mock_runner = MagicMock() trial = Trial("test_pb2", stub=True) pb2.on_trial_add(mock_runner, trial) _save_trial_result(pb2, trial, 1, _result(time=1, val=10)) new_config, _ = pb2._get_new_config(trial, trial) assert isinstance(new_config["b"]["c"], int) def test_pb2_custom_explore_fn_lambda(): """Test that a PB2 scheduler with a lambda explore fn can be serialized.""" hyperparam_bounds = {"a": [1.0, 2.0], "b": {"c": [2.0, 4.0], "d": [4.0, 10.0]}} pb2 = _create_pb2_scheduler( hyperparam_bounds=hyperparam_bounds, custom_explore_fn=lambda config: config, ) cloudpickle.dumps(pb2) if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))