best way to save trainer state?

[pchalasani]

Here's a simplified, abstract version of my question: Suppose I have two datasets A, B. I want to do this sequence:

Initialize states of trainer, modelA, modelB.
While A, B not exhausted, do:

• resume trainer from last state on data A
• train modelA for next 5 batches on data A
• save state of trainer and modelA
• resume trainer from last state on data B
• train modelB for next 5 batches on data B
• save state of trainer and modelB

// now do something with modelA, modelB,... immaterial for this question.

I know I can save and load "model state" using checkpoints.
But I'm not too clear on what is the best way to save + resume trainer state, which includes things like early stopping conditions, epoch, step numbers etc.
Any pointers appreciated.