Moving the checkpointing util funtions to a spearate file to avoid circular dependency

[wang2yn84]

Description

The circular dependency is multihost_dataloading -> maxtext_utils -> checkpointing -> multihost_dataloading. That's why in the previous development, we removed the deps of maxtext_utils from multihost_dataloading to avoid the issue, but that's just temporarily. Moving all the maxtext_util functions that depends on checkpointing to a separate util function solves this issue.

Tests

Presubmit tests.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• [x] I have performed a self-review of my code.
• [x] I have necessary comments in my code, particularly in hard-to-understand areas.
• [x] I have run end-to-end tests tests and provided workload links above if applicable.
• [x] I have made or will make corresponding changes to the doc if needed.

[wang2yn84]

Thanks @wang2yn84 for the nice refactor. Was the circular dependency causing some type of error? Or is the change for cleaner code?

Hi Branden, in the current implementation, we have to carefully avoid depend on maxtext_utils from multihost_dataloading. There are couple of util functions in multihost_dataloading we should place in maxtext_utils but we can't.

[github-actions[bot]]

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

[github-actions[bot]]

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.