Add Checkpointing of Trailblaze Algorithm #1159
Conversation
SimonBoothroyd
requested review from
andrrizzi,
hannahbrucemacdonald and
jchodera
| @@ -1962,7 +1963,7 @@ def _get_number_box_waters(self, *args): | |||
| # ============================================================================== | |||
|
|
|||
| def trailblaze_alchemical_protocol(thermodynamic_state, sampler_state, mcmc_move, state_parameters, | |||
| std_energy_threshold=0.5, threshold_tolerance=0.05, | |||
| checkpoint_path, std_energy_threshold=0.5, threshold_tolerance=0.05, | |||
There was a problem hiding this comment.
Do we want to make checkpointing optional, or required via the Python API? This approach seems to make it required, but it seems like it would be simple to make it optional.
There was a problem hiding this comment.
I agree it would be useful to have it optional in the function (although exposing it in the YAML will be slightly more complicated). It's possible tests are not passing also because you're defining a positional argument after a keyword one.
|
I'm going to merge these changes so that I can pick up from here and store also the configurations generated by trailblaze in a separate PR. |
|
Ah! I didn't notice the test failures. I'm going to make the changes directly here then. |
|
The failing tests have been already fixed in #1168 so I think we can merge safely. I'll correct eventual issues in the next PR if there are any. |
Description
This PR aims to implement a basic checkpointing system for the trailblazing algorithm - currently if the algorithm is interrupted while running (e.g. if an insufficient wall clock limit is set by a job submission script), the partially computed optimized lambda states are lost.
Status