Implement complete Keras-Orbax checkpoint integration by amitsrivastava78 · Pull Request #22002 · keras-team/keras (original) (raw)

@amitsrivastava78

This enables users to save and load complete model state (weights, optimizer, metrics) using keras.saving.load_model() with Orbax checkpoints.

[gemini-code-assist[bot]](/apps/gemini-code-assist)

@amitsrivastava78

@amitsrivastava78

@amitsrivastava78

@amitsrivastava78

@amitsrivastava78

Tests consolidated:

@amitsrivastava78

Performance improvements:

Code simplification:

Results:

@amitsrivastava78

Problem:

Solution:

Results:

@amitsrivastava78

Problem:

Solution:

Root Cause Analysis:

Fixes:

  1. Automatic cleanup: del ensures cleanup during garbage collection
  2. Explicit cleanup: try/finally blocks in async/sync tests
  3. Defense in depth: Both normal and abnormal termination scenarios covered

Results:

@amitsrivastava78

Test file descriptor leak fixes in OrbaxCheckpoint:

hertschuh

@amitsrivastava78

@amitsrivastava78

@amitsrivastava78

@amitsrivastava78

The overwrite parameter was never used (always defaulted to False) and is unnecessary with our preservation policy and save_decision_policy handling checkpoint management automatically.

@amitsrivastava78

The force_sync parameter was never used (always defaulted to None) and added unnecessary complexity. Sync vs async behavior is already controlled by the save_on_background constructor parameter, making this override unnecessary.

@amitsrivastava78

Orbax's checkpointer.close() already waits for pending async operations to complete before closing (per its API contract). The explicit wait_until_finished() call was redundant and added in later commits unnecessarily.

Reverting to the simpler original pattern where close() handles the wait.

hertschuh

@amitsrivastava78

@amitsrivastava78

hertschuh

@amitsrivastava78

…test structure

hertschuh

jerryxyj added a commit to jerryxyj/keras that referenced this pull request

Feb 14, 2026

@jerryxyj

This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters

[ Show hidden characters]({{ revealButtonHref }})