HuggingFace Hub integration: opt-in weight upload after training, default weight download for inference, tag dedup guards

[zachtheyek]

Closes #128

Stacking note

This PR's base is integration/pr0910-base — coordinator scaffolding that merges the finished train chain (stage machine + run manifest, #122/#126) and inference chain (streaming inference + inference_cadences manifest + viz suite, #120/#125, and PR-08's #129). Only the commits on feature/huggingface-integration are this PR's diff; review it against the integration base.

What this adds (PLAN PR-10)

Dependency — huggingface_hub>=1.21,<1.22 in requirements-container.txt + environment.yml (pip section). Version per SECURITY.md's Version Selection Policy: latest stable at pin time is 1.23.0 (2026-07-09); two minors back is 1.21 (1.21.0, 2026-06-25), which is newer than the newest ≥6-month-old release (pre-1.6, before 2026-01-12), so 1.21 is the target; no known advisory affects it. Clusters need a container image rebuild (singularity build aetherscan-ngc25.02.sif aetherscan.def) before production use. Pre-rebuild workaround (validated on blpc3): ./utils/run_container.sh python -m pip install --user "huggingface_hub==1.21.0", then launch with SINGULARITYENV_PYTHONNOUSERSITE= — the NGC image sets PYTHONNOUSERSITE=1, which disables the user site the --user install lands in.

Auth — HF_TOKEN via the gitignored .env, forwarded by utils/run_container.sh (new entry in the env allowlist), documented in the README .env example. Never logged. Downloads hit a public repo and need no token; only --hf-upload needs one.

Config/CLI — new HFConfig (repo_id = zachtheyek/aetherscan, upload_after_training = off, revision = None) + checkpoint.force_tag; flags: train --hf-upload/--no-hf-upload, both modes --hf-repo-id, inference --hf-revision, both modes --force-tag (tri-state BooleanOptionalAction like the other bool flags; effective default off). to_dict() updated; README CLI Reference regenerated verbatim; docs/CONFIG_AND_CLI.md tables/audit counts updated.

Training upload — new manifest-tracked hf_upload stage after final_save: stages the four artifacts under stable names (vae_encoder.keras, vae_decoder.keras, random_forest.joblib, config.json) plus a generated model card (pipeline description, tag, config summary, val ROC AUC/AP from rf_eval_artifacts when available, library versions, GitHub link/citation), commit message = save_tag, then create_tag(save_tag). Creates the public repo on first upload (create_repo(exist_ok=True)). Non-critical: failure is logged + recorded in the run manifest, never fails the run; re-running the identical command retries just this stage, and re-running with the upload disabled clears a stale recorded failure. Hub progress bars are force-disabled and StreamToLogger answers file-protocol probes (isatty/writable/readable/fileno), so transfers never depend on tty probing of the redirected streams.

Inference download — --encoder-path/--rf-path/--config-path become optional as an all-or-none trio. Resolution order: explicit local paths → --hf-revision → latest semver v* tag → latest final_vX tag → error with guidance. Downloads are revision-pinned via hf_hub_download; the cached trio is written onto args before validation, so the existing apply_saved_config + strategy-scoped model-load path runs unchanged, and the resolved revision lands in config.hf.revision for provenance.

Tag dedup guards (fail-early) — enforced right before command dispatch (post-validation, post-DB-init):

• Train, explicitly-provided tags only: collisions with vae_encoder_{tag}.keras, config_{tag}.json, or non-superseded training_stats rows hard-exit with the options listed — unless a resumable run-state manifest exists (same-tag retries keep working).
• HF-side: with --hf-upload, the save_tag is checked against existing repo tags at startup, not after ~30 h of training; a failed check (network) only warns.
• Inference: a completed prior run (its saved config_{tag}.json) or stale legacy-path DB rows collide; in-progress streaming runs (manifest rows, no config JSON) resume as before.
• Default datetime tags are immune by construction; --force-tag consciously overrides (and moves the HF tag on upload).

Verification

Unit suite — PYTHONPATH=src pytest -m "not gpu and not cluster" -q → 519 passed, 3 skipped (local arm64 venv, TF 2.17.1; baseline on the integration base was 445 passed). New tests cover upload arg shapes/staged filenames, revision resolution order, semver sorting, the dedup guard matrix, model-card content, hf_upload stage-machine semantics, StreamToLogger's file-protocol probes, and the progress-bar disabling — huggingface_hub fully mocked, no network in the suite. pre-commit run --all-files green on every commit.

Live round-trip on blpc3 (completed) — huggingface_hub==1.21.0 installed in the container via the pre-rebuild workaround above; HF_TOKEN present in .env (never printed/logged).

1. Train smoke with --hf-upload (5-GPU smoke config, fresh tag test_v26): training completed; the first upload attempt exposed a real bug — 'StreamToLogger' object has no attribute 'isatty' from upload_folder's progress machinery probing the redirected stdout. Failure containment worked exactly as designed: the run finished, hf_upload was recorded failed in the manifest, and the process exited nonzero with re-run guidance. Fixed in 7cd19e1 (both StreamToLogger probes and disabled Hub progress bars).
2. Manifest-driven stage retry (real-world validation of the PR-04 stage-resume design): relaunching the identical command after the fix skipped all five completed stages and re-ran only the upload:

   INFO | Resuming training run from manifest .../run_state_test_v26.json (attempt 2, completed rounds: [1, 2], stages done: ['vae_rounds', 'vae_plots', 'rf_train', 'rf_plots', 'final_save'], stages failed: ['hf_upload'])
   INFO | Stage 'vae_rounds' already complete — skipping
   INFO | Stage 'vae_plots' already complete — skipping
   INFO | Stage 'rf_train' already complete — loaded persisted RF model
   INFO | Stage 'rf_plots' already complete — skipping
   INFO | Stage 'final_save' already complete — skipping
   INFO | Uploaded run 'test_v26' to https://huggingface.co/zachtheyek/aetherscan and tagged the commit
   INFO | Stage 'hf_upload' complete (recorded in run manifest)
   INFO | Training completed successfully!
   

3. Hub state verified (public read, no token): repo zachtheyek/aetherscan PUBLIC with vae_encoder.keras, vae_decoder.keras, random_forest.joblib, config.json, README.md (generated model card with correct frontmatter + this run's config summary) on main; tag test_v26 pointing at the main head.
4. Inference smoke pulling from the Hub with NO local model paths: inference --hf-revision test_v26 --inference-files subset_test.csv resolved and downloaded the revision-pinned trio into the container's ~/.cache/huggingface (HOME bind-mount caching confirmed working — no HF_HOME override needed), fed the downloaded config.json through apply_saved_config, loaded the encoder inside strategy.scope() from the cache path, and completed: 2 cadences, 94,667 snippets processed, 0 candidates, exit 0 (~22 min wall for the 2-cadence subset).
5. Cleanup: tag test_v26 deleted from the Hub via delete_tag (repo + main-branch artifacts left in place; tags now []); cluster scratch//tmp//dev/shm cleaned, repo back on master.

AI disclosure

Implemented with Claude Code (Anthropic), end-to-end: design from the v1 plan, implementation, tests, live cluster validation, and this PR. Human-specified requirements and review by @zachtheyek.

🤖 Generated with Claude Code

[claude]

Claude finished @zachtheyek's task in 2m 29s —— View job

---

PR Review: HuggingFace Hub Integration

• Fetch base branch and review diff
• Review core module: hf_hub.py
• Review config/CLI changes
• Review tag guards: tag_guards.py
• Review integration points (main.py, train.py, inference.py)
• Review tests
• Review dependency and docs changes
• Post review summary with inline comments

---

Summary

This is a well-structured, thoroughly tested feature addition. The architecture — lazy imports for the HF dependency, thin seams for testability, non-critical upload stage that never fails a training run, clean resolution-order for inference artifacts — is sound and integrates cleanly with the existing stage machine and manifest-based resume system.

Strengths

• Testability: The _hf_api() / _hf_hub_download() lazy-import seams make the module fully unit-testable without network access. The _FakeHfApi stub in test_hf_hub.py is thorough and records call shapes for assertion. 64 new tests covering the full matrix (revision resolution, semver sorting, upload staging, tag guards, stage machine semantics).
• Fail-safe design: Upload failures are recorded in the manifest but never fail the run. The _clear_stage_failure mechanism for when users opt out after a failed upload is a nice touch — prevents a stale failure from blocking indefinitely.
• Inference artifact resolution: The resolve_inference_artifacts() design — writing downloaded cache paths onto args before validation — elegantly reuses the entire existing validation + apply_saved_config + model-load pipeline unchanged.
• Tag dedup guards: Good placement (post-validation, post-DB-init, pre-work) and clean separation of concerns between local and HF-side checks. Default datetime tags being immune by construction avoids false positives.
• Documentation: CLI reference regenerated, .env example updated, CONFIG_AND_CLI.md audit counts updated.

Issues and Suggestions

I've left detailed inline comments on the relevant lines. The key items:

1. TOCTOU on HF tag creation (hf_hub.py:390-395): Between the startup tag-existence check and the actual create_tag ~30 h later, the tag could have been created externally. With force=False, this would produce an opaque HF API error. Consider a targeted try/except on create_tag with a message pointing at --force-tag.

2. Download error UX (main.py:757-762): Raw huggingface_hub exceptions (HTTP 404, auth errors) can be cryptic for users. A wrapping exception with guidance ("is the revision correct?") would improve the operator experience.

3. Non-atomic tag move (hf_hub.py:390-395): The delete_tag → create_tag sequence in the force=True path is not atomic — a concurrent reader could see no tag. Single-user/single-run, so not a real risk today, but worth a comment.

4. sys.exit(1) in guards (tag_guards.py:94-105): Style preference — raising a custom exception instead of calling sys.exit directly would keep the exit decision at the call site and simplify testing. Not a blocker.

None of these are blockers. The PR is in good shape for merge once the live round-trip smoke test (blocked on HF_TOKEN on blpc3, as noted in the PR description) is completed as a follow-up.

Security

• HF_TOKEN never logged, loaded from gitignored .env, forwarded via container env allowlist — correct.
• Public repo downloads need no token — verified in the code.
• No secrets in model card or committed files.
• Dependency pin (>=1.21,<1.22) follows SECURITY.md's version selection policy.

Test Coverage

Good coverage across all new code paths: revision selection/resolution, artifact download routing, upload staging/arg shapes, model card content, tag guard matrix (local + HF, resume exemptions, force override), stage machine integration (opt-in/skip/failure/clear semantics), CLI validation (all-or-none trio, HF repo ID format, force-tag routing). huggingface_hub is fully mocked throughout.

[claude]

Nit: hf_tag_exists downloads the entire tag list and throws it away after a single in check. The dedup guard calls this once, so it's fine in practice, but if there's ever a second call site it'd be worth caching the tag list or switching to a targeted API check.

Also, the list_hf_tags call could theoretically be paginated for repos with very many tags — list_repo_refs doesn't paginate (it returns all refs), so this is safe with the current huggingface_hub API. Just flagging for future awareness.

[zachtheyek]

Agreed on both counts. Keeping the simple form for now: the guard is the only call site and runs once per launch, and list_repo_refs returns all refs unpaginated in huggingface_hub 1.21. If a second call site appears we'll cache the tag list (the inference resolution path already reuses one list_hf_tags result).

[claude]

There's a subtle race window between upload_folder and create_tag: if the process dies between these two calls, the commit exists on the Hub but is untagged. The next retry would then upload a duplicate commit and tag it, which is fine but leaves an orphaned commit.

More importantly, in the force=True path, delete_tag + create_tag is not atomic — a concurrent reader could see no tag between the two calls. Since this is single-user / single-run, it's not a real concern, but worth a comment noting the non-atomicity for the future.

Also: with force=False, if the tag already exists (e.g. a previous partial run committed but then successfully tagged), create_tag will raise. The startup _guard_hf_tag check should have caught this, but there's a TOCTOU gap (~30 h of training between check and upload). Might be worth a targeted try/except on create_tag with a clear error message pointing at --force-tag.

[zachtheyek]

Fixed in 8cf9832: create_tag now has a targeted 409-conflict handler that names the TOCTOU window and points at --force-tag (the artifacts stay uploaded, only the tag is missing), and the force path carries a comment documenting the non-atomic delete_tag -> create_tag sequence and why a crash between the calls is recoverable (the stage stays un-done in the run manifest, so a re-run recreates the tag). The orphaned-commit-on-retry case is accepted as-is — a duplicate commit on main is harmless and the tag lands on the newest one.

[claude]

Minor: download_inference_artifacts builds a tuple via a generator expression, which means the three downloads are strictly sequential. That's probably the right call (avoids hammering the Hub API), but worth noting that parallelizing them (e.g. via ThreadPoolExecutor) could shave a few seconds off cold-cache inference startup if this ever becomes a bottleneck.

[zachtheyek]

Keeping the downloads sequential: the trio is dominated by the encoder file, so parallelizing saves little, and cold-cache startup is a one-time cost per revision (later runs hit the local cache). Will revisit with a ThreadPoolExecutor if it ever shows up in startup timings.

[claude]

_exit_on_collisions calls sys.exit(1) directly. This is consistent with the module's documented contract ("hard-stop"), but it means the function is untestable without pytest.raises(SystemExit) — which the tests do correctly use. Just noting that returning a sentinel or raising a custom exception instead of sys.exit would make the guard testable without that ceremony and keep the exit decision at the call site. A minor style preference, not a blocker.

[zachtheyek]

Keeping sys.exit(1) here, with rationale: the guard's callers live in main.py's startup path, which already uses sys.exit(1) pervasively for fail-fast conditions (init failures, GPU strategy failures, resolution failures), and the guard logs its multi-line collision list + options at the failure site where the context lives. A custom exception whose only consumer immediately exits would add a type + call-site handling without changing behavior, and the existing tests' pytest.raises(SystemExit) is the standard idiom for exactly this contract. Happy to revisit if the guard ever grows a non-CLI caller.

[claude]

Good placement — resolving artifacts before apply_saved_config / validate_args is the right sequencing. One observation: if the HF download fails (network error, repo doesn't exist), the user gets a sys.exit(1) with the raw exception message. The error from resolve_hf_revision already includes actionable guidance (the three alternatives), which is great. But hf_hub_download failures from huggingface_hub itself can be opaque (HTTP 404, auth errors). Consider wrapping the download path in a try/except that re-raises a more helpful RuntimeError with "is the revision correct? does the repo carry the expected files?" guidance.

[zachtheyek]

Fixed in 8cf9832: both failure surfaces in the resolution path are now wrapped with operator guidance — download_inference_artifacts re-raises raw huggingface_hub errors as a RuntimeError covering revision/repo-id correctness, repo visibility / HF_TOKEN, and network reachability, and resolve_hf_revision wraps tag-listing failures the same way. Unit tests cover both wrappers.

[claude]

The _HF_REPO_ID_PATTERN check is a useful guardrail. Just noting that HuggingFace also allows repo IDs with dots and special characters in the namespace component — the current regex [A-Za-z0-9][A-Za-z0-9._-]* should cover all valid cases. Looks correct.

[zachtheyek]

Thanks for double-checking the character classes — agreed the pattern covers the valid namespace/name forms (dots, dashes, underscores), so no change needed.

[zachtheyek]

Review addressed in 8cf9832 (all four items):

1. create_tag TOCTOU — targeted HTTP-409 handler: a tag created between the startup dedup check and the post-upload create_tag now raises a clear error pointing at --force-tag, noting the artifacts are uploaded but untagged. Unit-tested with a mocked 409.
2. Download error UX — raw huggingface_hub errors from the inference resolution path (tag listing and artifact downloads) are wrapped in RuntimeErrors with operator guidance (revision/repo-id correctness, repo visibility / HF_TOKEN, network). Unit-tested.
3. Non-atomic force path — comment added documenting delete_tag -> create_tag non-atomicity and the manifest-driven recovery.
4. sys.exit in tag_guards — kept as-is; rationale in the inline reply (consistent with main.py's fail-fast startup style, rich logging at the failure site, and the change would only add indirection for a consumer that immediately exits).

Affected tests + pre-commit run --all-files green (tests/unit/test_hf_hub.py + tests/unit/test_tag_guards.py: 54 passed). Live round-trip is in progress on blpc3 — PR body will be updated with the evidence when it completes.

[zachtheyek]

Live round-trip completed on blpc3 — full evidence now in the PR body's Verification section. Highlights:

• The first --hf-upload train smoke (tag test_v26) exposed a real bug: 'StreamToLogger' object has no attribute 'isatty' (the Hub client's progress machinery probing the pipeline's redirected stdout). Failure containment worked as designed — training completed, the failure was recorded in the run manifest, nonzero exit with re-run guidance.
• Fixed in 7cd19e1 with two independent layers: StreamToLogger now answers file-protocol probes (isatty/writable/readable/fileno), and hf_hub.py force-disables Hub progress bars before any transfer. Unit tests for both.
• Relaunching the identical command was a real-world validation of the stage-resume design: all five completed stages skipped via the manifest, only hf_upload re-ran, and it succeeded — repo created public with the four stable-name artifacts + generated model card, tag test_v26 on the main head.
• The inference smoke then pulled --hf-revision test_v26 with no local model paths: revision-pinned download into the container's HOME cache, downloaded config through apply_saved_config, encoder loaded in strategy.scope(), 2 cadences / 94,667 snippets processed, exit 0.
• test_v26 tag deleted from the Hub afterwards (delete_tag); repo + main-branch artifacts left in place. Cluster cleaned (scratch, /tmp, /dev/shm; repo back on master).

Full local suite after all fixes: 519 passed, 3 skipped.

[zachtheyek]

SECURITY.md now documents HF_TOKEN (commit b5c8e7e). Added it alongside SLACK_BOT_TOKEN: a new row in the Tokens and API Keys table, a note covering its opt-in upload role (train --hf-upload), public-repo downloads needing no token, .env sourcing plus utils/run_container.sh --env forwarding, never-log guidance, and a Token Rotation subsection mirroring the Slack one.

[zachtheyek]

@claude please perform a second-pass code review of this PR. Re-review the current diff for any remaining correctness bugs, security issues, or significant quality problems, and post your findings. If the diff is clean, please say so explicitly. (Final verification round before merge.)

[claude]

Second-pass code review — PR #130 (HuggingFace Hub integration)

Verdict: the diff is clean. I reviewed all new/modified files for correctness bugs, security issues, and significant quality problems. No blocking issues found.

Security ✅

• HF_TOKEN is never logged, never committed, read implicitly by huggingface_hub from the environment — correct.
• The run_container.sh allowlist forwards it as ${HF_TOKEN:-} (empty string when unset) — safe.
• Downloads hit a public repo with no token; only uploads require auth — documented accurately.
• No secrets in the model card or config.json (the serialized config contains no env values).

Correctness ✅

• Ordering in main.py: resolve_inference_artifacts → apply_saved_config → validate_args → enforce_tag_guards → command dispatch. The HF download populates the config path before apply_saved_config reads it — correct dependency chain.
• download_inference_artifacts: The tuple unpacking (encoder, rf, config) matches the download order in the generator (HF_ENCODER_FILENAME, HF_RF_FILENAME, HF_CONFIG_FILENAME) — consistent with resolve_inference_artifacts's assignment args.encoder_path, args.rf_path, args.config_path = ....
• Tag dedup guards: The exemption logic (manifest file → resumable run, so skip the guard) is correct. --force-tag overrides are applied after the collision check, not before — correct priority.
• upload_run_to_hf: delete_tag → create_tag non-atomicity under --force-tag is documented and acceptable (the manifest records the stage as undone, so a crash is retryable).
• _is_tag_conflict: Properly checks for HTTP 409 via the response object chain, with fallback to False if the import fails.
• Stage machine: hf_upload is correctly non-critical (catch-all, record failure, continue). Stale failure clearing when the user re-runs without --hf-upload prevents a permanent nonzero exit — nice edge case.
• Validation: The all-or-none inference artifact trio check (0 < n_provided < len(artifact_specs)) correctly rejects partial sets while allowing both the zero-provided (Hub download) and all-provided (local) cases.
• StreamToLogger: fileno() raises io.UnsupportedOperation (the io.StringIO convention, subclasses both OSError and ValueError) — correct for the probing patterns in tqdm/rich/click.

Minor observations (non-blocking, no action needed)

1. list_hf_tags imports RepositoryNotFoundError directly (outside a try/except), unlike the lazy-import seams _hf_api()/_hf_hub_download(). This is fine because list_hf_tags is only reachable when HF features are active (the call sites in _guard_hf_tag and resolve_hf_revision both gate on --hf-upload / no-local-paths), and _guard_hf_tag wraps the call in its own try/except anyway.
2. The _HF_REPO_ID_PATTERN regex is a reasonable sanity check; the Hub's own server-side validation is the authority.

Test coverage

74 new test cases covering: revision selection/sorting, inference artifact resolution (all branches), upload staging/arg shapes/filenames, model card content, tag dedup guard matrix (train + inference × collision/clean/resume/force), StreamToLogger probes, stage machine semantics for the new hf_upload stage, and progress-bar disabling. All mocked — no network in the suite.

LGTM — ready to merge. 🫡