Skip to content

refactor(checkpoints): HF-only writes; configurable W&B read alias#71

Merged
mkeeler43 merged 3 commits into
codex/hf-checkpoint-backendfrom
codex/hf-checkpoint-backend-tighten
Jun 10, 2026
Merged

refactor(checkpoints): HF-only writes; configurable W&B read alias#71
mkeeler43 merged 3 commits into
codex/hf-checkpoint-backendfrom
codex/hf-checkpoint-backend-tighten

Conversation

@mkeeler43

@mkeeler43 mkeeler43 commented May 28, 2026

Copy link
Copy Markdown
Contributor

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:

wandb login
huggingface-cli login   # or: export HF_TOKEN=...

Training:

python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id beams2d --seed 1 --track --save-model
  • --track logs scalars + media to W&B (default project: engiopt)
  • --save-model uploads the checkpoint to HF (default repo: IDEALLab/engiopt-cgan-cnn-2d)
  • Override --hf-entity / --hf-repo-prefix for personal experiments — same shape as --wandb-entity / --wandb-project

Evaluation / inference:

python engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py --problem-id beams2d --seed 1
  • Pulls the checkpoint from HF automatically
  • Falls back silently to the legacy W&B artifact for runs trained before this cutover

Flags removed (cleanup)

These flags from #63 are no longer surfaced to users; behavior is folded into sensible defaults:

Removed flag New behavior
--checkpoint-backend Always HF. --save-model is the on/off.
--hf-private Always public. Create a private repo manually if needed.
--model-source Always auto: HF first, W&B fallback.
--local-model-dir Power-user knob; can be re-added later.

The W&B read fallback inside checkpoint_store.py stays 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

  • HF train → eval round-trip on beams2d / cgan_cnn_2d (1 epoch, seed 999):
    • Checkpoint uploaded to IDEALLab/engiopt-cgan-cnn-2d/beams2d/seed_999/
    • Immutable run-scoped path …/seed_999/run_<wandb_run_id>/ also created
    • W&B run summary recorded the HF repo, package path, and revision
    • Eval auto-pulled the checkpoint, ran inference, wrote the metrics CSV
    • Test artifacts deleted afterward
  • CI: ruff, pre-commit, mypy all green

Forward-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_parts and wandb_artifact_alias — are already wired into save_checkpoint_package and resolve_named_checkpoint, so the LVAE migration is per-script editing only. No further changes to checkpoint_store.py are needed.

Commits

  1. HF-only writes, configurable W&B read alias — drop the "wandb"/"both" write modes; add wandb_artifact_alias so legacy reads work for non-standard alias schemes.
  2. Collapse CLI surface to W&B parallel — drop the 4 flags listed above; bake their old defaults into the call sites.
  3. mypy fix — drop two now-invalid model_source= kwargs on load_model_from_reference call sites that the earlier flag cleanup left behind.

mkeeler43 and others added 3 commits May 28, 2026 16:54
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>
@mkeeler43 mkeeler43 merged commit 1ac9cb5 into codex/hf-checkpoint-backend Jun 10, 2026
3 checks passed
@mkeeler43 mkeeler43 deleted the codex/hf-checkpoint-backend-tighten branch June 10, 2026 15:36
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant