feat: add HF-backed checkpoint storage#63
Merged
Conversation
Contributor
Author
|
@mkeeler43 requesting your review on this one. This PR adds an HF-backed checkpoint backend for EngiOpt so new model weights can live on Hugging Face instead of saturating W&B artifact storage, while keeping backward-compatible restore support for existing W&B-hosted checkpoints. It also packages run config and metadata with the checkpoint so new HF-backed runs remain reproducible and tied back to the originating W&B run. |
Tightens the HF checkpoint backend so HuggingFace is the sole storage
backend for model weights. W&B remains in use for media, scalars, and
run lineage, but no longer receives model artifacts.
- CheckpointBackend -> Literal["hf", "none"]; drop "wandb"/"both" writes
- save_checkpoint_package: drop wandb_artifacts param and W&B-write branch
- resolve_named_checkpoint: add wandb_artifact_alias override so callers
with custom alias schemes (e.g. LVAE's seed_<s>_rec<r>_perf<p>) can
reach legacy artifacts during the transition
- _resolve_wandb_package / _build_wandb_artifact_path: parameterize alias
- 15 training scripts: drop the now-unused wandb_artifacts={...} kwarg
- README: document HF-only writes and the transitional W&B read fallback
The W&B read fallback (ModelSource="wandb") is preserved as a safety net
for runs trained before the HF cutover; it will be removed in a later
cleanup pass once in-use artifacts are migrated.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The HF-backed checkpoint flow now mirrors the W&B media flow exactly:
Training:
--track --wandb-entity X --wandb-project Y (media to W&B)
--save-model --hf-entity Z --hf-repo-prefix W (checkpoint to HF)
Evaluation:
--hf-entity Z --hf-repo-prefix W (pull from HF)
Removed flags (rolled into sensible defaults inside the call sites):
- --checkpoint-backend (always "hf"; --save-model is the on/off)
- --hf-private (always public; create private repos manually
if needed -- create_repo(exist_ok=True) won't
override an existing repo's visibility)
- --model-source (always "auto": HF first, W&B fallback)
- --local-model-dir (power-user; can be re-added later)
The W&B read-fallback is still preserved inside checkpoint_store as a
safety net for runs trained before the HF cutover; users no longer need
a CLI flag to reach it.
Also fixed a pre-existing reportUndefinedVariable in
run_pe_optimization.load_model_from_reference (model_source was
referenced without being a parameter; now passed as "auto").
Updated 15 training scripts + 14 evaluator scripts + README.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…nce calls load_model_from_reference does not accept model_source; my earlier simplification pass left two now-invalid kwargs at the call sites. The wrapper internally passes model_source="auto" so callers don't need to specify it. Fixes mypy CI failure on PR #71. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
refactor(checkpoints): HF-only writes; configurable W&B read alias
…ackend # Conflicts: # README.md # engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py
Contributor
Author
|
I re-checked this before merging. Local and remote branch state match exactly, CI is green, the branch is up to date with main with no merge conflicts, and I verified the HF checkpoint flow end to end with a fresh smoke train/upload/load test. I approve the merge and I am merging this now. |
AnnaDelbeke
pushed a commit
to AnnaDelbeke/EngiOpt
that referenced
this pull request
Jun 26, 2026
feat: add HF-backed checkpoint storage
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation
Follow-up