Skip to content

Releases: google/orbax

v0.1.5

17 Mar 22:46
Compare
Choose a tag to compare

Added

  • Use a more precise timestamp when generating temporary directory names to
    permit more than one concurrent checkpointing attempt per second.
  • Automatic import of nest_asyncio.

v0.1.4

15 Mar 20:51
Compare
Choose a tag to compare

Added

  • Support for generic transformation function in PyTreeCheckpointHandler.
  • Support n-digit checkpoint step format.

Fixed

  • Eliminate Flax dependency to fix circular dependency problem.

v0.1.3

03 Mar 23:17
2e3f137
Compare
Choose a tag to compare

Added

  • sharding option on `ArrayRestoreArgs

v0.1.2

17 Feb 17:07
Compare
Choose a tag to compare

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

30 Jan 19:50
Compare
Choose a tag to compare

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

04 Jan 17:42
Compare
Choose a tag to compare

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

10 Dec 01:34
Compare
Choose a tag to compare

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

05 Dec 21:33
Compare
Choose a tag to compare

Added

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

v0.0.19

01 Dec 21:40
Compare
Choose a tag to compare

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

16 Nov 22:19
Compare
Choose a tag to compare

What's Changed

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

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