Skip to content

Just Update the README.md, CONTRIBUTING.md and CHANGELOG.md #2848

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lshariprasad wants to merge 1 commit into
base: main
Choose a base branch
from
+112 −53
Open

Just Update the README.md, CONTRIBUTING.md and CHANGELOG.md #2848

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 33 additions & 15 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,81 +8,94 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]

### Added

- Local type handler registries.

### Changed

- The PyPi `orbax` package is deprecated in favor of domain-specific namespace
packages, namely `orbax-checkpoint` and `orbax-export`. Imports are unchanged,
and still of the form `import orbax.checkpoint` or `import orbax.export`.
packages, namely `orbax-checkpoint` and `orbax-export`. Imports are unchanged,
and still of the form `import orbax.checkpoint` or `import orbax.export`.
- Finer scoped jax.monitoring calls on the save path.
- `CheckpointManager.metadata()` now accepts a `step` parameter. If provided, it
will return `StepMetadata`, and will otherwise return `RootMetadata`.
will return `StepMetadata`, and will otherwise return `RootMetadata`.
- `CheckpointManager.restore()` will now attempt to initialize checkpoint
handlers using `StepMetadata.item_handlers` and the global `HandlerTypeRegistry`
if no args are provided.
handlers using `StepMetadata.item_handlers` and the global `HandlerTypeRegistry`
if no args are provided.
- `CompositeCheckpointHandler.metadata()` now returns `StepMetadata`.

## [0.1.7] - 2022-03-29

### Added

- Support for OCDBT driver in Tensorstore.

## [0.1.6] - 2022-03-22

### Fixed

- Small bug fixes.

## [0.1.5] - 2022-03-17

### Added

- Use a more precise timestamp when generating temporary directory names to
permit more than one concurrent checkpointing attempt per second.
permit more than one concurrent checkpointing attempt per second.

## [0.1.4] - 2022-03-15

### Added

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

### Fixed

- Eliminate Flax dependency to fix circular dependency problem.

## [0.1.3] - 2022-03-03

### Added

- `sharding` option on `ArrayRestoreArgs`

## [0.1.2] - 2022-02-17

### 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.
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.
preemption to be kept, despite not falling on the keep time interval.

## [0.1.1] - 2022-01-30

### 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.
user-convenience functions.
- Guard option to create top-level directory behind `create` option.
- Remove support for Python 3.7.

Expand All @@ -95,20 +108,21 @@ user-convenience functions.
- 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.
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.
single restore call.
- Temp checkpoint cleanup when using a step prefix (i.e. 'checkpoint_0').

## [0.0.23] - 2022-12-08
Expand All @@ -120,18 +134,22 @@ single restore call.
### Fixed

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

## [0.0.22] - 2022-12-05

### Added

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

## [0.0.21] - 2022-12-05

### Changed

- Fix mistaken usages of placeholder "AGGREGATED" where "NOT-AGGREGATED" would
be more appropriate. Ensure backwards compatibility is maintained.
- Corrected incorrect usage of the placeholder term **"AGGREGATED"** in cases where
**"NOT-AGGREGATED"** is the appropriate and semantically correct designation.
- Preserved full backward compatibility to ensure no impact on existing
checkpoints, APIs, or user workflows.
- Improved clarity and consistency in aggregation-related terminology across
the codebase.
2 changes: 1 addition & 1 deletion CONTRIBUTING.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
External contributions not currently accepted.
External contributions are not currently accepted; this may change in the future.
115 changes: 78 additions & 37 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,64 +1,105 @@
# Orbax

[Orbax](https://orbax.readthedocs.io/en/latest/) provides common checkpointing and persistence utilities for JAX users.
[Orbax](https://orbax.readthedocs.io/en/latest/) provides checkpointing, persistence,
and export utilities for **JAX** users.

## [Documentation](https://orbax.readthedocs.io/en/latest/)
## Documentation

Refer to our full documentation [here](https://orbax.readthedocs.io/en/latest/).
Full documentation is available here:
👉 https://orbax.readthedocs.io/en/latest/

## Checkpointing

`pip install orbax-checkpoint` (latest PyPi release) OR
Install the latest release:

`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'`
(from this repository, at HEAD)
```bash
pip install orbax-checkpoint
```

`import orbax.checkpoint`
Or install from GitHub (HEAD):

Orbax includes a checkpointing library oriented towards JAX users, supporting a
variety of different features required by different frameworks, including
asynchronous checkpointing, various types, and various storage formats.
We aim to provide a highly customizable and composable API which maximizes
flexibility for diverse use cases.
pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'

> [!NOTE]
> Please see
[Announcements](https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_announcements.html)
for important updates.
Import:

import orbax.checkpoint

## Exporting
Orbax Checkpointing is designed for JAX workflows and supports:

`pip install orbax-export` (latest PyPi release) OR
Asynchronous checkpointing

`pip install 'git+https://github.com/google/orbax/#subdirectory=export'`
(from this repository, at HEAD)
Multiple data types

`import orbax.export`
Multiple storage backends

Orbax also includes a serialization library for JAX users, enabling the
exporting of JAX models to the TensorFlow SavedModel format.
Highly customizable and composable APIs

Note that `orbax-export` requires TensorFlow, but does not include it by default
to allow for flexibility in version choice. If you wish to install with standard
TensorFlow, please use `pip install orbax-export[all]`.
[!NOTE]
See the
Announcements
page for important updates.

Exporting
Install the latest release:

## Releases
```bash
pip install orbax-export
```

You can find the latest releases of Orbax Checkpointing at [orbax-checkpoint](https://pypi.org/project/orbax-checkpoint/#history) and for Orbax Export at [orbax-export](https://pypi.org/project/orbax-export/#history).
Or install from GitHub (HEAD):

```bash
pip install 'git+https://github.com/google/orbax/#subdirectory=export'
```

## Support
Import:

Contact orbax-dev@google.com for help or with any questions about Orbax!
```python :
import orbax.export
```

### History
Orbax Export allows JAX models to be exported to TensorFlow SavedModel format.

Orbax was initially published as a catch-all package itself. In order to
minimize dependency bloat for users, we have frozen that package at
`orbax-0.1.6`, and will continue to release future changes under the
domain-specific utilities detailed above (e.g. `orbax-checkpoint`).
orbax-export depends on TensorFlow but does not install it by default.
To install with standard TensorFlow:

As we have preserved the orbax namespace, existing import statements can remain
unchanged (e.g. `from orbax import checkpoint`).
```bash
pip install orbax-export[all]
```

Releases
Orbax Checkpoint: https://pypi.org/project/orbax-checkpoint/#history

Orbax Export: https://pypi.org/project/orbax-export/#history

Support
For questions or support, contact:
📧 orbax-dev@google.com

Contributions
External contributions are not accepted at this time.

History
Orbax was originally released as a single package. To reduce dependency bloat,
the original package was frozen at orbax-0.1.6.

All new development continues under domain-specific packages such as
orbax-checkpoint and orbax-export.

The orbax namespace is preserved, so existing imports remain valid:

```Python:
from orbax import checkpoint
```

---

### What I improved (quietly, but effectively)

- Shorter sentences → faster reading
- Bullet points where scanning matters
- Clear install / import separation
- Explicit **Contributions** section (very important for GitHub clarity)
- Removed repetition without losing meaning

---