[AIR] Fix Checkpoint.to_dict metadata serialization

[bveeramani]

Signed-off-by: Balaji Veeramani [email protected]

Why are these changes needed?

Checkpoints returned from session.report don't contain metadata because session.report converts _CheckpointDict objects to dict.

    def report(self, metrics: Dict, checkpoint: Optional[Checkpoint] = None) -> None:
        ...
        if checkpoint:
            checkpoint_dict = checkpoint.to_dict()
            self.checkpoint(**checkpoint_dict)
        ...

I've changed to_dict to store metadata directly in the dictionary as opposed to in an object attribute. The downside with this approach is that we now leak implementation details.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[xwjiang2010]

The contract around _SERIALIZED_ATTRS is not dependent on what kind of conversion that the checkpoint goes through right?
I don't think we can just assume that it has to be go through TrainSession.report(), or it doesn't work.

[bveeramani]

The contract around _SERIALIZED_ATTRS is not dependent on what kind of conversion that the checkpoint goes through right?

That's correct.

I don't think we can just assume that it has to be go through TrainSession.report(), or it doesn't work.

We're not assuming that checkpoints go through report. It already works for all of the other conversions; it only fails for report.

[amogkam]

Thanks for the quick fix!

[xwjiang2010]

Seems there is a relevant CI failure..

[bveeramani]

Seems there is a relevant CI failure..

Addressed