[Tune] Call on_trial_complete after final checkpoint (#19243)

This commit is contained in:
Amog Kamsetty 2021-10-11 09:47:39 -07:00 committed by GitHub
parent 6252a6c1f9
commit b3ad72643c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -919,6 +919,8 @@ class TrialRunner:
flat_result = flatten_dict(result)
self._validate_result_metrics(flat_result)
_trigger_callback_complete = False
if self._stopper(trial.trial_id,
result) or trial.should_stop(flat_result):
result.update(done=True)
@ -938,8 +940,7 @@ class TrialRunner:
trial=trial,
result=result.copy())
self._callbacks.on_trial_complete(
iteration=self._iteration, trials=self._trials, trial=trial)
_trigger_callback_complete = True
decision = TrialScheduler.STOP
else:
with warn_if_slow("scheduler.on_trial_result"):
@ -975,6 +976,10 @@ class TrialRunner:
# the global checkpoint state.
self._checkpoint_trial_if_needed(trial, force=force_checkpoint)
if _trigger_callback_complete:
self._callbacks.on_trial_complete(
iteration=self._iteration, trials=self._trials, trial=trial)
if trial.is_saving:
# Cache decision to execute on after the save is processed.
# This prevents changing the trial's state or kicking off