[Tune] Synchronous Mode for PBT (#10283)

2026-07-09 13:25:22 +08:00 · 2020-08-31 00:00:47 -07:00
parent 05fe6dc278
commit afde3db4f0
4 changed files with 558 additions and 87 deletions
@@ -241,6 +241,7 @@ class _MockTrialRunner():
        return True

    def _pause_trial(self, trial):
+        self.trial_executor.save(trial, Checkpoint.MEMORY, None)
        trial.status = Trial.PAUSED

    def _launch_trial(self, trial):
@@ -702,6 +703,13 @@ class _MockTrial(Trial):
        self.custom_trial_name = None
        self.custom_dirname = None

+    def on_checkpoint(self, checkpoint):
+        self.restored_checkpoint = checkpoint.value
+
+    @property
+    def checkpoint(self):
+        return Checkpoint(Checkpoint.MEMORY, self.trainable_name, None)
+

 class PopulationBasedTestingSuite(unittest.TestCase):
    def setUp(self):
@@ -720,7 +728,8 @@ class PopulationBasedTestingSuite(unittest.TestCase):
                   require_attrs=True,
                   hyperparams=None,
                   hyperparam_mutations=None,
-                   step_once=True):
+                   step_once=True,
+                   synch=False):
        hyperparam_mutations = hyperparam_mutations or {
            "float_factor": lambda: 100.0,
            "int_factor": lambda: 10,
@@ -734,7 +743,9 @@ class PopulationBasedTestingSuite(unittest.TestCase):
            hyperparam_mutations=hyperparam_mutations,
            custom_explore_fn=explore,
            log_config=log_config,
-            require_attrs=require_attrs)
+            synch=synch,
+            require_attrs=require_attrs,
+        )
        runner = _MockTrialRunner(pbt)
        for i in range(num_trials):
            trial_hyperparams = hyperparams or {
@@ -746,10 +757,17 @@ class PopulationBasedTestingSuite(unittest.TestCase):
            trial = _MockTrial(i, trial_hyperparams)
            runner.add_trial(trial)
            trial.status = Trial.RUNNING
+        for i in range(num_trials):
+            trial = runner.trials[i]
            if step_once:
-                self.assertEqual(
-                    pbt.on_trial_result(runner, trial, result(10, 50 * i)),
-                    TrialScheduler.CONTINUE)
+                if synch:
+                    self.assertEqual(
+                        pbt.on_trial_result(runner, trial, result(10, 50 * i)),
+                        TrialScheduler.PAUSE)
+                else:
+                    self.assertEqual(
+                        pbt.on_trial_result(runner, trial, result(10, 50 * i)),
+                        TrialScheduler.CONTINUE)
        pbt.reset_stats()
        return pbt, runner

@@ -822,6 +840,32 @@ class PopulationBasedTestingSuite(unittest.TestCase):
        self.assertEqual(pbt._num_checkpoints, 2)
        self.assertEqual(pbt._num_perturbations, 0)

+    def testCheckpointMostPromisingTrialsSynch(self):
+        pbt, runner = self.basicSetup(synch=True)
+        trials = runner.get_trials()
+
+        # no checkpoint: haven't hit next perturbation interval yet
+        self.assertEqual(pbt.last_scores(trials), [0, 50, 100, 150, 200])
+        self.assertEqual(
+            pbt.on_trial_result(runner, trials[0], result(15, 200)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(pbt.last_scores(trials), [0, 50, 100, 150, 200])
+        self.assertEqual(pbt._num_checkpoints, 0)
+
+        # trials should be paused until all trials are synced.
+        for i in range(len(trials) - 1):
+            self.assertEqual(
+                pbt.on_trial_result(runner, trials[i], result(20, 200 + i)),
+                TrialScheduler.PAUSE)
+
+        self.assertEqual(pbt.last_scores(trials), [200, 201, 202, 203, 200])
+        self.assertEqual(pbt._num_checkpoints, 0)
+
+        self.assertEqual(
+            pbt.on_trial_result(runner, trials[-1], result(20, 204)),
+            TrialScheduler.PAUSE)
+        self.assertEqual(pbt._num_checkpoints, 2)
+
    def testPerturbsLowPerformingTrials(self):
        pbt, runner = self.basicSetup()
        trials = runner.get_trials()
@@ -852,6 +896,35 @@ class PopulationBasedTestingSuite(unittest.TestCase):
        self.assertIn(trials[0].restored_checkpoint, ["trial_3", "trial_4"])
        self.assertTrue("@perturbed" in trials[2].experiment_tag)

+    def testPerturbsLowPerformingTrialsSynch(self):
+        pbt, runner = self.basicSetup(synch=True)
+        trials = runner.get_trials()
+
+        # no perturbation: haven't hit next perturbation interval
+        self.assertEqual(
+            pbt.on_trial_result(runner, trials[-1], result(15, -100)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(pbt.last_scores(trials), [0, 50, 100, 150, 200])
+        self.assertTrue("@perturbed" not in trials[-1].experiment_tag)
+        self.assertEqual(pbt._num_perturbations, 0)
+
+        # Don't perturb until all trials are synched.
+        self.assertEqual(
+            pbt.on_trial_result(runner, trials[-1], result(20, -100)),
+            TrialScheduler.PAUSE)
+        self.assertEqual(pbt.last_scores(trials), [0, 50, 100, 150, -100])
+        self.assertTrue("@perturbed" not in trials[-1].experiment_tag)
+
+        # Synch all trials.
+        for i in range(len(trials) - 1):
+            self.assertEqual(
+                pbt.on_trial_result(runner, trials[i], result(20, -10 * i)),
+                TrialScheduler.PAUSE)
+        self.assertEqual(pbt.last_scores(trials), [0, -10, -20, -30, -100])
+        self.assertIn(trials[-1].restored_checkpoint, ["trial_0", "trial_1"])
+        self.assertIn(trials[-2].restored_checkpoint, ["trial_0", "trial_1"])
+        self.assertEqual(pbt._num_perturbations, 2)
+
    def testPerturbWithoutResample(self):
        pbt, runner = self.basicSetup(resample_prob=0.0)
        trials = runner.get_trials()
@@ -1088,6 +1161,27 @@ class PopulationBasedTestingSuite(unittest.TestCase):
            trials[i].status = Trial.PENDING
        self.assertEqual(pbt.choose_trial_to_run(runner), trials[3])

+    def testSchedulesMostBehindTrialToRunSynch(self):
+        pbt, runner = self.basicSetup(synch=True)
+        trials = runner.get_trials()
+        runner.process_action(
+            trials[0], pbt.on_trial_result(runner, trials[0], result(
+                800, 1000)))
+        runner.process_action(
+            trials[1], pbt.on_trial_result(runner, trials[1], result(
+                700, 1001)))
+        runner.process_action(
+            trials[2], pbt.on_trial_result(runner, trials[2], result(
+                600, 1002)))
+        runner.process_action(
+            trials[3], pbt.on_trial_result(runner, trials[3], result(
+                500, 1003)))
+        runner.process_action(
+            trials[4], pbt.on_trial_result(runner, trials[4], result(
+                700, 1004)))
+        self.assertIn(
+            pbt.choose_trial_to_run(runner), [trials[0], trials[1], trials[3]])
+
    def testPerturbationResetsLastPerturbTime(self):
        pbt, runner = self.basicSetup()
        trials = runner.get_trials()
@@ -1141,6 +1235,44 @@ class PopulationBasedTestingSuite(unittest.TestCase):
                check_policy(json.loads(line))
        shutil.rmtree(tmpdir)

+    def testLogConfigSynch(self):
+        def check_policy(policy):
+            self.assertIsInstance(policy[2], int)
+            self.assertIsInstance(policy[3], int)
+            self.assertIn(policy[0], ["0tag", "1tag"])
+            self.assertIn(policy[1], ["3tag", "4tag"])
+            self.assertIn(policy[2], [0, 1])
+            self.assertIn(policy[3], [3, 4])
+            for i in [4, 5]:
+                self.assertIsInstance(policy[i], dict)
+                for key in [
+                        "const_factor", "int_factor", "float_factor",
+                        "id_factor"
+                ]:
+                    self.assertIn(key, policy[i])
+                self.assertIsInstance(policy[i]["float_factor"], float)
+                self.assertIsInstance(policy[i]["int_factor"], int)
+                self.assertIn(policy[i]["const_factor"], [3])
+                self.assertIn(policy[i]["int_factor"], [8, 10, 12])
+                self.assertIn(policy[i]["float_factor"], [2.4, 2, 1.6])
+                self.assertIn(policy[i]["id_factor"], [3, 4, 100])
+
+        pbt, runner = self.basicSetup(
+            log_config=True, synch=True, step_once=False)
+        trials = runner.get_trials()
+        tmpdir = tempfile.mkdtemp()
+        for i, trial in enumerate(trials):
+            trial.local_dir = tmpdir
+            trial.last_result = {TRAINING_ITERATION: i}
+            pbt.on_trial_result(runner, trials[i], result(10, i))
+        log_files = ["pbt_global.txt", "pbt_policy_0.txt", "pbt_policy_1.txt"]
+        for log_file in log_files:
+            self.assertTrue(os.path.exists(os.path.join(tmpdir, log_file)))
+            raw_policy = open(os.path.join(tmpdir, log_file), "r").readlines()
+            for line in raw_policy:
+                check_policy(json.loads(line))
+        shutil.rmtree(tmpdir)
+
    def testReplay(self):
        # Returns unique increasing parameter mutations
        class _Counter:
@@ -1156,6 +1288,7 @@ class PopulationBasedTestingSuite(unittest.TestCase):
            perturbation_interval=5,
            log_config=True,
            step_once=False,
+            synch=False,
            hyperparam_mutations={
                "float_factor": lambda: 100.0,
                "int_factor": _Counter(1000)
@@ -1292,6 +1425,176 @@ class PopulationBasedTestingSuite(unittest.TestCase):

        shutil.rmtree(tmpdir)

+    def testReplaySynch(self):
+        # Returns unique increasing parameter mutations
+        class _Counter:
+            def __init__(self, start=0):
+                self.count = start - 1
+
+            def __call__(self, *args, **kwargs):
+                self.count += 1
+                return self.count
+
+        pbt, runner = self.basicSetup(
+            num_trials=4,
+            perturbation_interval=5,
+            log_config=True,
+            step_once=False,
+            synch=True,
+            hyperparam_mutations={
+                "float_factor": lambda: 100.0,
+                "int_factor": _Counter(1000)
+            })
+        trials = runner.get_trials()
+        tmpdir = tempfile.mkdtemp()
+
+        # Internal trial state to collect the real PBT history
+        class _TrialState:
+            def __init__(self, config):
+                self.step = 0
+                self.config = config
+                self.history = []
+
+            def forward(self, t):
+                while self.step < t:
+                    self.history.append(self.config)
+                    self.step += 1
+
+        trial_state = []
+        for i, trial in enumerate(trials):
+            trial.local_dir = tmpdir
+            trial.last_result = {TRAINING_ITERATION: 0}
+            trial_state.append(_TrialState(trial.config))
+
+        # Helper function to simulate stepping trial k a number of steps,
+        # and reporting a score at the end
+        def trial_step(k, steps, score, synced=False):
+            res = result(trial_state[k].step + steps, score)
+
+            trials[k].last_result = res
+            trial_state[k].forward(res[TRAINING_ITERATION])
+
+            trials[k].status = Trial.RUNNING
+            if not synced:
+                action = pbt.on_trial_result(runner, trials[k], res)
+                runner.process_action(trials[k], action)
+                return
+            else:
+                # Reached synchronization point
+                old_configs = [trial.config for trial in trials]
+                action = pbt.on_trial_result(runner, trials[k], res)
+                runner.process_action(trials[k], action)
+                new_configs = [trial.config for trial in trials]
+
+                for i in range(len(trials)):
+                    old_config = old_configs[i]
+                    new_config = new_configs[i]
+                    if old_config != new_config:
+                        # Copy history from source trial
+                        source = -1
+                        for m, cand in enumerate(trials):
+                            if cand.trainable_name == trials[
+                                    i].restored_checkpoint:
+                                source = m
+                                break
+                        assert source >= 0
+                        trial_state[i].history = trial_state[
+                            source].history.copy()
+                        trial_state[i].step = trial_state[source].step
+                        trial_state[i].config = new_config.copy()
+
+        # Initial steps
+        trial_step(0, 10, 0)
+        trial_step(1, 11, 10)
+        trial_step(2, 12, 0)
+        trial_step(3, 13, -1, synced=True)
+
+        # 3 <-- 1, new_t 11
+        # next_perturb_sync = 13
+
+        # Next block
+        trial_step(0, 17, -10)  # 20
+        trial_step(2, 15, -20)  # 20
+        trial_step(3, 16, 0)  # 20
+        trial_step(1, 7, 1, synced=True)  # 18
+
+        # 2 <-- 1, new_t=11+7=18
+        # next_perturb_sync = 20
+
+        # Next block
+        trial_step(2, 13, 0)  # 31
+        trial_step(3, 14, 10)  # 34
+        trial_step(0, 11, -1)  # 31
+        trial_step(1, 12, 0, synced=True)  # 30
+
+        # 0 <-- 3, new_t=11+9+14=34
+        # next_perturb_sync = 34
+
+        # Next block
+        trial_step(0, 6, 20)  # 40
+        trial_step(3, 9, -40)  # 43
+        trial_step(2, 8, -50)  # 39
+        trial_step(1, 7, 30, synced=True)  # 37
+
+        # 2 <-- 1, new_t=18+13+8=37
+        # next_perturb_sync = 43
+
+        # Playback trainable to collect configs at each step
+        class Playback(Trainable):
+            def setup(self, config):
+                self.config = config
+                self.replayed = []
+                self.iter = 0
+
+            def step(self):
+                self.iter += 1
+                self.replayed.append(self.config)
+                return {
+                    "reward": 0,
+                    "done": False,
+                    "replayed": self.replayed,
+                    TRAINING_ITERATION: self.iter
+                }
+
+            def reset_config(self, new_config):
+                self.config = new_config
+                return True
+
+            def save_checkpoint(self, tmp_checkpoint_dir):
+                return tmp_checkpoint_dir
+
+            def load_checkpoint(self, checkpoint):
+                pass
+
+        # Loop through all trials and check if PBT history is the
+        # same as the playback history
+        for i, trial in enumerate(trials):
+            if trial.trial_id in ["1"]:  # Did not exploit anything
+                continue
+
+            replay = PopulationBasedTrainingReplay(
+                os.path.join(tmpdir,
+                             "pbt_policy_{}.txt".format(trial.trial_id)))
+            analysis = tune.run(
+                Playback,
+                scheduler=replay,
+                stop={TRAINING_ITERATION: trial_state[i].step})
+
+            replayed = analysis.trials[0].last_result["replayed"]
+            self.assertSequenceEqual(trial_state[i].history, replayed)
+
+        # Trial 1 did not exploit anything and should raise an error
+        with self.assertRaises(ValueError):
+            replay = PopulationBasedTrainingReplay(
+                os.path.join(tmpdir,
+                             "pbt_policy_{}.txt".format(trials[1].trial_id)))
+            tune.run(
+                Playback,
+                scheduler=replay,
+                stop={TRAINING_ITERATION: trial_state[1].step})
+
+        shutil.rmtree(tmpdir)
+
    def testPostprocessingHook(self):
        def explore(new_config):
            new_config["id_factor"] = 42
@@ -4,64 +4,13 @@ import pickle
 import random
 import unittest
 import sys
+import time

 import ray
 from ray import tune
 from ray.tune.schedulers import PopulationBasedTraining


-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)
-        return tmp_checkpoint_dir
-
-    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)
-
-
-def MockTrainingFunc(config, checkpoint_dir=None):
-    iter = 0
-    a = config["a"]
-    b = config["b"]
-
-    if 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 tune.checkpoint_dir(step=iter) 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)
-
-
-def MockTrainingFunc2(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))
-
-
 class MockParam(object):
    def __init__(self, params):
        self._params = params
@@ -73,6 +22,77 @@ class MockParam(object):
        return val


+class PopulationBasedTrainingSynchTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=2)
+
+        def MockTrainingFuncSync(config, checkpoint_dir=None):
+            iter = 0
+
+            if 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 tune.checkpoint_dir(step=iter) as checkpoint_dir:
+                    checkpoint_path = os.path.join(checkpoint_dir,
+                                                   "checkpoint")
+                    with open(checkpoint_path, "wb") as fp:
+                        pickle.dump((a, iter), fp)
+                # Score gets better every iteration.
+                time.sleep(1)
+                tune.report(mean_accuracy=iter + a, a=a)
+
+        self.MockTrainingFuncSync = MockTrainingFuncSync
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def synchSetup(self, synch, param=[10, 20, 30]):
+        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()["mean_accuracy"] != 33))
+
+    def testSynchPass(self):
+        analysis = self.synchSetup(True)
+        self.assertTrue(all(analysis.dataframe()["mean_accuracy"] == 33))
+
+    def testSynchPassLast(self):
+        analysis = self.synchSetup(True, param=[30, 20, 10])
+        self.assertTrue(all(analysis.dataframe()["mean_accuracy"] == 33))
+
+
 class PopulationBasedTrainingConfigTest(unittest.TestCase):
    def setUp(self):
        ray.init()
@@ -81,6 +101,15 @@ class PopulationBasedTrainingConfigTest(unittest.TestCase):
        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",
@@ -97,7 +126,7 @@ class PopulationBasedTrainingConfigTest(unittest.TestCase):
        )

        tune.run(
-            MockTrainingFunc2,
+            MockTrainingFunc,
            fail_fast=True,
            num_samples=4,
            scheduler=scheduler,
@@ -120,6 +149,31 @@ class PopulationBasedTrainingResumeTest(unittest.TestCase):
        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)
+                return tmp_checkpoint_dir
+
+            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",
@@ -151,6 +205,25 @@ class PopulationBasedTrainingResumeTest(unittest.TestCase):
            stop={"training_iteration": 3})

    def testPermutationContinuationFunc(self):
+        def MockTrainingFunc(config, checkpoint_dir=None):
+            iter = 0
+            a = config["a"]
+            b = config["b"]
+
+            if 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 tune.checkpoint_dir(step=iter) 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)
+
        scheduler = PopulationBasedTraining(
            time_attr="training_iteration",
            metric="mean_accuracy",