Repository: ray-project/ray

Issue:
[tune] PBT perturbing after first iteration KeyError
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux
- **Ray installed from (source or binary)**: binary
- **Ray version**: 0.7.1
- **Python version**: 3.7.3
- **Exact command to reproduce**: run PBT with `perturbation_interval=1` and `time_attr="training_iteration"`.

### Describe the problem
Occasionally, with the above described config PBT attempts to perturb config before `trial.last_result` is populated. This results in `KeyError` and the worker dies.

### Source code / logs
```python
Traceback (most recent call last):
  File "/home/magnus/.cache/pypoetry/virtualenvs/document-classification-py3.7/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 458, in _process_trial
    self, trial, result)
  File "/home/magnus/.cache/pypoetry/virtualenvs/document-classification-py3.7/lib/python3.7/site-packages/ray/tune/schedulers/pbt.py", line 217, in on_trial_result
    self._exploit(trial_runner.trial_executor, trial, trial_to_clone)
  File "/home/magnus/.cache/pypoetry/virtualenvs/document-classification-py3.7/lib/python3.7/site-packages/ray/tune/schedulers/pbt.py", line 279, in _exploit
    trial_to_clone, new_config)
  File "/home/magnus/.cache/pypoetry/virtualenvs/document-classification-py3.7/lib/python3.7/site-packages/ray/tune/schedulers/pbt.py", line 244, in _log_config_on_step
    trial.last_result[TRAINING_ITERATION],
KeyError: 'training_iteration'
```



Discussion:
Can you post what your `run` command looks like?
(not actual code we're running, since it's scattered all over modules, but something like this)
```python
experiment = ray.tune.Experiment(
            name=experiment_name,
            run=SupervisedTrainable,
            resources_per_trial={"cpu": cpu_fraction, "gpu": gpu_fraction},
            stop=experiment_config.training.stopping_criteria,
            num_samples=4,
            config=resolved_config,
            checkpoint_at_end=True,
            checkpoint_freq=experiment_config.training.checkpoint_freq,
        )
scheduler = PopulationBasedTraining(
        time_attr="training_iteration",
        reward_attr="val/accuracy",
        perturbation_interval=1,
        hyperparam_mutations=hparam_mutations,
    )
ray.tune.run(
    experiment,
    scheduler=scheduler,
    reuse_actors=False
    )
```
OK, I see what the issue is; I'll open a PR to fix.

Resolving patch:
diff --git a/python/ray/tune/schedulers/pbt.py b/python/ray/tune/schedulers/pbt.py
--- a/python/ray/tune/schedulers/pbt.py
+++ b/python/ray/tune/schedulers/pbt.py
@@ -268,8 +268,8 @@ def _log_config_on_step(self, trial_state, new_state, trial,
             "pbt_policy_" + trial_to_clone_id + ".txt")
         policy = [
             trial_name, trial_to_clone_name,
-            trial.last_result[TRAINING_ITERATION],
-            trial_to_clone.last_result[TRAINING_ITERATION],
+            trial.last_result.get(TRAINING_ITERATION, 0),
+            trial_to_clone.last_result.get(TRAINING_ITERATION, 0),
             trial_to_clone.config, new_config
         ]
         # Log to global file.


Test patch:
diff --git a/python/ray/tune/tests/test_trial_scheduler.py b/python/ray/tune/tests/test_trial_scheduler.py
--- a/python/ray/tune/tests/test_trial_scheduler.py
+++ b/python/ray/tune/tests/test_trial_scheduler.py
@@ -622,10 +622,15 @@ def tearDown(self):
         ray.shutdown()
         _register_all()  # re-register the evicted objects
 
-    def basicSetup(self, resample_prob=0.0, explore=None, log_config=False):
+    def basicSetup(self,
+                   resample_prob=0.0,
+                   explore=None,
+                   perturbation_interval=10,
+                   log_config=False,
+                   step_once=True):
         pbt = PopulationBasedTraining(
             time_attr="training_iteration",
-            perturbation_interval=10,
+            perturbation_interval=perturbation_interval,
             resample_probability=resample_prob,
             quantile_fraction=0.25,
             hyperparam_mutations={
@@ -646,9 +651,10 @@ def basicSetup(self, resample_prob=0.0, explore=None, log_config=False):
                 })
             runner.add_trial(trial)
             trial.status = Trial.RUNNING
-            self.assertEqual(
-                pbt.on_trial_result(runner, trial, result(10, 50 * i)),
-                TrialScheduler.CONTINUE)
+            if step_once:
+                self.assertEqual(
+                    pbt.on_trial_result(runner, trial, result(10, 50 * i)),
+                    TrialScheduler.CONTINUE)
         pbt.reset_stats()
         return pbt, runner
 
@@ -959,6 +965,24 @@ def explore(new_config):
         self.assertEqual(trials[0].config["id_factor"], 42)
         self.assertEqual(trials[0].config["float_factor"], 43)
 
+    def testFastPerturb(self):
+        pbt, runner = self.basicSetup(
+            perturbation_interval=1, step_once=False, log_config=True)
+        trials = runner.get_trials()
+
+        tmpdir = tempfile.mkdtemp()
+        for i, trial in enumerate(trials):
+            trial.local_dir = tmpdir
+            trial.last_result = {}
+        pbt.on_trial_result(runner, trials[0], result(1, 10))
+        self.assertEqual(
+            pbt.on_trial_result(runner, trials[2], result(1, 200)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(pbt._num_checkpoints, 1)
+
+        pbt._exploit(runner.trial_executor, trials[1], trials[2])
+        shutil.rmtree(tmpdir)
+
 
 class AsyncHyperBandSuite(unittest.TestCase):
     def setUp(self):