v0.1.3

Added

• sharding option on `ArrayRestoreArgs

v0.1.2

Added

• Add "standard user recipe" to documentation.
• Add unit tests using mock to simulate preemption.
• Logging to increase transparency around why checkpoints are kept vs. deleted.
• Expand on uses of restore_args in colab.
• Expose utils_test.
• Add msgpack_utils to move toward eliminating Flax dependency.
• CheckpointManager starts a background thread to finalize checkpoints so that
  checkpoints are finalized as soon as possible in async case.

Changed

• Remove CheckpointManager update API.
• Remove support for deprecated GDA.
• Add tmp suffix on step directory creation in CheckpointManager.save.

Fixed

• Preemption when using keep_time_interval caused the most recent steps before
  preemption to be kept, despite not falling on the keep time interval.

v0.1.1

Added

• A util function that constructs restore_args from a target PyTree.
• CheckpointManager delete API, which allows deleting an existing step.
• Made dev dependencies optional to minimize import overhead.

Changed

• Refactored higher-level utils in checkpoint_utils, which provides user-convenience functions.
• Guard option to create top-level directory behind create option.
• Remove support for Python 3.7.

v0.1.0

Added

• Check for metric file in addition to item directory in CheckpointManager.
• Additional logs to indicate save/restore completion.
• Support for None leaves in PyTree save/restore.
• ArrayCheckpointHandler for individual arrays/scalars.
• read: bool option on all_steps to force read from storage location instead
  of using cached steps.
• Simplified "Getting Started" section in the docs.
• CheckpointManager creates the top level directory if it does not yet exist.
• Write msgpack bytes asynchronously.

Changed

• Removed some unused test_utils methods for filtering empty nodes.
• Update docs on PyTreeCheckpointHandler.
• Removed unneeded AbstractCheckpointManager.

Fixed

• Usage of bytes_limiter to prevent too many bytes from being read during a
  single restore call.
• Temp checkpoint cleanup when using a step prefix (i.e. 'checkpoint_0').

v0.0.23

Added

• Option to customize metadata file name for Tensorstore.

Fixed

• Restore failure on GCS due to misidentification of checkpoint as
  "not finalized".

v0.0.22

Added

• Added CHANGELOG.md for version updates (additions and changes), ingested by
  auto-publish functionality.

v0.0.19

What's Changed

• [Orbax] Log unmatched keys from the new tree by @copybara-service in #147
• Fix auto-publish tool, replace setup.py with pyproject.toml standard, disable duplicate builds. by @copybara-service in #154
• Add support for tracking checkpoint metrics with Orbax in T5X. by @copybara-service in #153

Full Changelog: v0.0.18...v0.0.19

v0.0.18

What's Changed

• Add auto-publish tool for Orbax. by @copybara-service in #142

Full Changelog: v0.0.16...v0.0.18

Release

What's Changed

• Add update API for CheckpointManager. At the moment, this will only support updating metrics. We will consider whether to support updating any part of the checkpoint in the future. by @copybara-service in #137

Full Changelog: v0.0.15...v0.0.16

v0.0.15

What's Changed

• Add missing import to JAX gda_serialization. by @copybara-service in #131
• Increase orbax.checkpoint version. by @copybara-service in #135

Full Changelog: v0.0.14...v0.0.15