train: add PACE H200 latent cotrain launcher

[ElmoPA]

No description provided.

[ElmoPA]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• train: add PACE H200 latent cotrain launcher #595 👈 (View in Graphite)
• config(flow-transfer): prepare new-data H16 cotrain #592 : 3 other dependent PRs (#593 , #594 , #596 )
• feat(pipeline): add U-socket arc-length policy #589 : 1 other dependent PR (#590 )
• feat(pipeline): port arc-length action tokens #585 : 1 other dependent PR (#586 )
• feat(pipeline): consolidate Fold, PushT, overlays, and MoE #579 : 2 other dependent PRs (#580 , #583 )
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Claude Code Review

Review

Summary

Adds a PACE H200 sbatch launcher for the Flow Transfer H16 latent cotrain arm, plus a small helper (read_successful_wandb_exit_code) and eval-resolver registration in verify_training_smoke.py, with an extensive companion test suite that exercises the launcher's embedded Python heredocs and Hydra composition.

Key concerns

1. _load_training_config uses raw eval as a resolver. The added fallback OmegaConf.register_new_resolver("eval", eval) re-introduces arbitrary code execution on config load and, more importantly, likely does not match the real egomimic.utils.hydra_resolvers registration (which the file already imports on line 22). If hydra_resolvers registers eval with different semantics (e.g., a sandboxed evaluator, or with replace=True), this fallback silently masks a broken import. Recommend: drop the fallback entirely and rely on the import; if the resolver isn't registered after import that's a real bug worth surfacing.

2. open_for_scan API assumption. read_successful_wandb_exit_code uses DataStore().open_for_scan(...) / scan_data() / close(). This is a private wandb internal API — please confirm it exists in the pinned wandb version and add a version pin note. The test mocks it so a real-world API mismatch would not be caught.

3. Launcher is environment-specific and unreviewable in isolation. The 1101-line sbatch hardcodes paths under /storage/project/r-dxu345-0/paphiwetsa3/..., dataset SHAs, episode counts, transfer job IDs, and a specific commit SHA (EXPECTED_HEAD). This is fine as a one-shot experiment launcher but:

   • It lives under scripts/train/ where it will accumulate. Consider a subdirectory like scripts/train/pace_h200/flow_transfer_newdata_h16/ and a short README explaining these are pinned per-experiment artifacts, not reusable templates.
   • The EXPECTED_LAUNCHER_SHA self-check means every future edit (including trivial ones) requires updating the expected SHA — that's the intent, but worth calling out.

4. Repository-scoped git status --porcelain check will fail against a shared checkout. test -z "$(git -C "$REPO" status --porcelain=v1 --untracked-files=all)" runs after training, but training writes training.log, wandb dirs, etc. under $RUN_DIR, which is outside $REPO, so this should be OK — but confirm no artifacts land in-tree (e.g., .hydra under CWD when hydra.run.dir is absolute).

5. --config-name=train_zarr_cartesian — verify this config still exists and hasn't been renamed; the test uses it via compose(config_name="train_zarr_cartesian", ...) so CI should catch drift, but only if this test actually runs in CI (it will try to initialize_config_dir and compose, which needs the experiment config and all its dependencies importable).

6. Test suite will not run without the full training deps. _compose triggers Hydra instantiation of the full pipeline config, which imports Torch, wandb, egomimic algo modules, etc. If tests/ is meant to run in a lightweight CI, this will fail. Mark with a pytest marker (@pytest.mark.slow or @pytest.mark.requires_gpu_configs) or gate on env presence.

7. torch.load(..., weights_only=False) on a checkpoint path (in the launcher heredoc) — fine for trusted local checkpoints, but flag that this is intentional; PyTorch will warn/error in future versions.

Suggestions

• Remove the if not OmegaConf.has_resolver("eval"): register(eval) fallback in verify_training_smoke.py; the hydra_resolvers import should be authoritative. If it's not, fix hydra_resolvers instead.
• Move per-experiment launchers under a dated/named subdirectory to prevent scripts/train/ bloat.
• In read_successful_wandb_exit_code, assert a non-None exit record more clearly, and consider catching store.close() inside a try/except or checking for scan errors before asserting exit_codes[-1] == 0 (the current message (stream_path, exit_codes) is a tuple assertion message which is fine but unusual).
• Add a top-of-file docstring to the sbatch documenting: intended one-shot use, pinned commit, pinned dataset, how to regenerate EXPECTED_LAUNCHER_SHA.
• Confirm the two new helpers (_load_training_config, read_successful_wandb_exit_code) don't break any existing callers of verify_training_smoke.verify_training_smoke (they shouldn't — additive — but double-check import-time side effects from import egomimic.utils.hydra_resolvers).

Verdict: Comment

The launcher itself is a self-contained, aggressively fail-closed experiment artifact and looks internally consistent. The concerns are mostly (a) the eval resolver fallback which I'd like removed, and (b) organizational (where these launchers live, whether the heavy tests run in CI). The changes to verify_training_smoke.py are the only pieces that affect shared code paths and deserve a second look before merge.

---

Reviewed by Claude · Review workflow