refactor(checkpoints): HF-only writes; configurable W&B read alias#71
Merged
mkeeler43 merged 3 commits intoJun 10, 2026
Merged
Conversation
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>
AnnaDelbeke
pushed a commit
to AnnaDelbeke/EngiOpt
that referenced
this pull request
Jun 26, 2026
…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 IDEALLab#71. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
AnnaDelbeke
pushed a commit
to AnnaDelbeke/EngiOpt
that referenced
this pull request
Jun 26, 2026
…end-tighten refactor(checkpoints): HF-only writes; configurable W&B read alias
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
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
Stacks on #63 to deliver the policy we agreed on: HuggingFace is the only place model checkpoints are stored. W&B continues to handle scalars, media, and run history.
What the new UX looks like
One-time setup:
Training:
--tracklogs scalars + media to W&B (default project:engiopt)--save-modeluploads the checkpoint to HF (default repo:IDEALLab/engiopt-cgan-cnn-2d)--hf-entity/--hf-repo-prefixfor personal experiments — same shape as--wandb-entity/--wandb-projectEvaluation / inference:
Flags removed (cleanup)
These flags from #63 are no longer surfaced to users; behavior is folded into sensible defaults:
--checkpoint-backend--save-modelis the on/off.--hf-private--model-source--local-model-dirThe W&B read fallback inside
checkpoint_store.pystays as a safety net while existing artifacts are still on W&B. A later cleanup PR will remove it once everything is migrated.What's verified
beams2d/cgan_cnn_2d(1 epoch, seed 999):IDEALLab/engiopt-cgan-cnn-2d/beams2d/seed_999/…/seed_999/run_<wandb_run_id>/also createdForward-compat for the upcoming LVAE work
The future LVAE-minimal branch saves with custom subpaths (
rec{r}/perf{p}/{filter}/seed_{s}) and reads legacy W&B artifacts under the LVAE alias scheme (seed_{s}_rec{r}_perf{p}…). Both hooks —extra_path_partsandwandb_artifact_alias— are already wired intosave_checkpoint_packageandresolve_named_checkpoint, so the LVAE migration is per-script editing only. No further changes tocheckpoint_store.pyare needed.Commits
"wandb"/"both"write modes; addwandb_artifact_aliasso legacy reads work for non-standard alias schemes.model_source=kwargs onload_model_from_referencecall sites that the earlier flag cleanup left behind.