Make training runs self-describing in MLflow - #108
Conversation
Model identity was never logged on training runs: MODEL is deliberately kept out of train_config (cli/utils.py extra_fields), and train_config is the only bulk param log — so "which network is this?" was answerable only via experiment-name conventions or by unpickling artifacts. - New shared log_training_run_identity() (cli/utils.py), called by both jaxtrain and torchtrain: logs model, network_type, backend, run_uuid (the uuid1 in every artifact filename — the MLflow<->disk join key), input_dim, param_space, param_bounds_json + sha256, training_data_folder, n_training_files, config_sha256, lanfactory_version; tags schema_version=1, phase=train, slurm ids from env. Best-effort: tracking never kills training. - Log network_config.pickle and the source YAML as run artifacts. Previously the pickle was written to disk but never logged, so the MLflow artifact set alone could not reconstruct a network. - Carry model_config from training data into DatasetTorch and on into data_details.pickle (train/valid_data_model_config). The training pickles embed param_bounds; the trained-network folder previously dropped them — the one field HSSM-facing consumers most need. - Unify per-epoch metrics across backends: jax now logs train_loss/val_loss per epoch (torch adds train_loss alongside its existing loss/val_loss). The jax per-100-step `loss` metric is unchanged. - Fix OPN mislabeling: jax train_and_evaluate() gains a network_type parameter passed from the CLI. The old inference from train_output_type maps logits -> "cpn" unconditionally, so OPN artifacts were labeled cpn; inference is kept only as a fallback for direct trainer use. Schema documented in HSSMSpine _docs/mlflow-schema.md (forthcoming). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three findings from adversarial review, each reproduced against a live store: - n_training_files identity param (effective count) collided with the trainer's later log_params(train_config) (configured cap): MLflow rejects changed param values, silently dropping the ENTIRE train_config batch. The effective count is now the tag n_training_files_used. - Resumed runs (--mlflow-run-id) regenerate RUN_ID, so re-logging run_uuid as a param raised and dropped the whole identity batch, and the recorded join key went stale. Split: immutable network facts (model, network_type, backend, input_dim, param_space, bounds) stay params — identical re-log is permitted; per-invocation values (run_uuid, training_data_folder, n_training_files_used, config_sha256, lanfactory_version) are now tags, which a resume overwrites so the run always reflects its latest artifacts. Regression test logs identity twice in one run. - torch train_loss logged the LAST minibatch loss at the cumulative batch step while jax logs the EPOCH-MEAN at step=epoch — same metric name, different semantics, contradicting the comment claiming they match. torch now logs train_loss as the epoch mean at step=epoch; legacy loss/val_loss series unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ssm-simulators training pickles are written with cloudpickle and can carry lambda boundary functions inside model_config (e.g. race_no_bias_angle_2). Retaining that dict on DatasetTorch made _save_data_details' stdlib pickle.dump raise. Unpicklable values are now repr'd at capture; the catalog-relevant fields (params, param_bounds, choices) pass through. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Warning Review limit reached
Next review available in: 32 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThe training CLIs now log MLflow run identity and configuration artifacts. JAX and Torch trainers persist model configuration metadata and epoch metrics. JAX accepts an explicit network type. Ruff now has a bounded development dependency range. ChangesTraining metadata
Development tooling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingCLI
participant TrainingTrainer
participant MLflow
TrainingCLI->>MLflow: log run identity and configuration artifacts
TrainingCLI->>TrainingTrainer: start training with network type
TrainingTrainer->>MLflow: log epoch train_loss and val_loss
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Adds a shared, CLI-level MLflow identity logging schema so training runs from both torch and jax backends become self-describing for downstream catalog/registry tooling, while also ensuring training-data model_config (incl. param bounds) is preserved through to data_details artifacts.
Changes:
- Introduces
log_training_run_identityhelper to log immutable identity as MLflow params and per-invocation identity as tags (resume-safe), and logs key config artifacts to MLflow. - Propagates
model_configfrom training data into datasets anddata_detailsartifacts, including sanitization for stdlib-pickle compatibility. - Fixes JAX output filename labeling by threading
network_type, and aligns cross-backendtrain_lossmetric semantics to per-epoch values.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/test_run_identity.py | Adds regression tests for MLflow identity logging, model_config retention/sanitization, and JAX network_type passthrough. |
| src/lanfactory/trainers/torch_mlp.py | Captures/sanitizes model_config from training pickles; logs epoch-mean train_loss; includes model_config in data_details. |
| src/lanfactory/trainers/jax_mlp.py | Adds explicit network_type passthrough for output naming; logs per-epoch schema metrics; includes model_config in data_details. |
| src/lanfactory/cli/utils.py | Adds the shared log_training_run_identity MLflow helper implementing the params/tags schema. |
| src/lanfactory/cli/torch_train.py | Calls identity logger and logs network config + YAML config as MLflow artifacts. |
| src/lanfactory/cli/jax_train.py | Calls identity logger, logs config artifacts, and passes network_type into the JAX trainer. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # the provenance HSSM-facing consumers need. Populated from the first | ||
| # data file when present; training pickles embed it alongside | ||
| # generator_config (ssm-simulators lan_mlp.py). | ||
| self.data_model_config: str = "None" |
Unpinned ruff resolved to 0.16 in CI (no tracked lockfile) and flags 111 pre-existing errors on main alone — unrelated to any PR content. Pin matches ssm-simulators' identical fix (>=0.15.1,<0.16); the 0.16 rule migration should land as its own deliberate PR. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pyproject.toml`:
- Around line 85-87: Update the Ruff revision in the pre-commit configuration to
a release within the supported range >=0.15.1 and <0.16, aligning it with the
dependency constraint in pyproject.toml; do not retain v0.14.13 unless an
explicit rationale is documented.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 38e9ac58-660d-4baa-97fa-2d8317920ccf
📒 Files selected for processing (7)
pyproject.tomlsrc/lanfactory/cli/jax_train.pysrc/lanfactory/cli/torch_train.pysrc/lanfactory/cli/utils.pysrc/lanfactory/trainers/jax_mlp.pysrc/lanfactory/trainers/torch_mlp.pytests/test_run_identity.py
| # <0.16: ruff 0.16 flags ~111 pre-existing errors repo-wide (same drift that | ||
| # broke ssm-simulators CI, pinned identically there); migrate deliberately. | ||
| "ruff>=0.15.1,<0.16", |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Project Ruff constraint:"
rg -n 'ruff' pyproject.toml
echo "Pre-commit Ruff revision:"
rg -n -A3 -B2 'ruff-pre-commit|rev:' .pre-commit-config.yamlRepository: lnccbrown/LANfactory
Length of output: 853
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "tool.ruff lint/format config in pyproject.toml:"
sed -n '110,145p' pyproject.toml
echo
echo "uv-related Ruff usage:"
rg -n 'uv .*ruff|ruff|pre-commit|tox|.ci|workflow|github/workflows' -S pyproject.toml .pre-commit-config.yaml .github || trueRepository: lnccbrown/LANfactory
Length of output: 2970
🌐 Web query:
ruff-pre-commit v0.14.13 Ruff version bundled
💡 Result:
The ruff-pre-commit repository follows the versioning of the Ruff tool itself [1]. Therefore, ruff-pre-commit v0.14.13 bundles Ruff version 0.14.13 [2][3][4].
Citations:
- 1: https://github.com/astral-sh/ruff-pre-commit
- 2: astral-sh/ruff-pre-commit@v0.14.12...v0.14.13
- 3: Update prek dependencies astral-sh/ruff#22868
- 4: chore(deps): update pre-commit hook astral-sh/ruff-pre-commit to v0.14.13 mw-root/uv-lock-report#101
Align the pre-commit Ruff revision with the supported range.
pyproject.toml requires Ruff >=0.15.1,<0.16, while .pre-commit-config.yaml uses ruff-pre-commit v0.14.13, which bundles Ruff 0.14.13. CI runs uv run ruff from the project dependency, so local pre-commit hooks can lint/format differently. Update the pre-commit revision to a Ruff release in >=0.15.1,<0.16, or document why the lower tooling version remains intentional.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pyproject.toml` around lines 85 - 87, Update the Ruff revision in the
pre-commit configuration to a release within the supported range >=0.15.1 and
<0.16, aligning it with the dependency constraint in pyproject.toml; do not
retain v0.14.13 unless an explicit rationale is documented.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/lanfactory/trainers/torch_mlp.py:93
data_model_configis annotated asstrbut is assigned a dict whenmodel_configis present in the training data. This makes the attribute’s type inaccurate and can break static analysis and downstream assumptions about its shape.
self.data_generator_config: str = "None"
# model_config from the training data (param_bounds, params, choices):
# the provenance HSSM-facing consumers need. Populated from the first
# data file when present; training pickles embed it alongside
# generator_config (ssm-simulators lan_mlp.py).
self.data_model_config: str = "None"
src/lanfactory/trainers/torch_mlp.py:723
epoch_loss_sum += float(loss)runs on every minibatch even whenmlflow_onis false. On GPU, converting a tensor tofloatcan force a device sync each step, slowing training. Sinceepoch_loss_sumis only used for the MLflowtrain_lossmetric, guard the accumulation behindmlflow_on.
# Log training progress
self._log_training_progress(epoch, cnt, loss, verbose)
epoch_loss_sum += float(loss)
cnt += 1
step_cnt += 1
pyproject.toml:89
tests/test_run_identity.pyimports and usescloudpickle, butcloudpickleis not listed as a direct dependency in the dev/test dependency group. Relying on transitive installation can make CI fail if upstream dependencies change; declare it explicitly.
"ruff>=0.15.1,<0.16",
"types-PyYAML",
"mlflow>=3.14.0",
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
Also mark the two genuinely unreachable defensive branches no-cover: the jax network_type 'unknown' fallback (JaxMLP.setup() raises on any train_output_type outside network_type_dict before a trainer can exist) and the PackageNotFoundError guard (lanfactory is always installed under uv run). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (2)
src/lanfactory/trainers/torch_mlp.py:92
data_model_configis annotated asstrbut is later assigned a dict (via_picklable_copy(init_file["model_config"])). This makes the attribute type incorrect and can break type checking / IDE expectations for code that consumesDatasetTorch.data_model_config.
# model_config from the training data (param_bounds, params, choices):
# the provenance HSSM-facing consumers need. Populated from the first
# data file when present; training pickles embed it alongside
# generator_config (ssm-simulators lan_mlp.py).
self.data_model_config: str = "None"
src/lanfactory/trainers/jax_mlp.py:653
pickle.dump(..., open(...))leaves the file handle unclosed if an exception occurs (and can leak descriptors). Use a context manager when writingdata_details.pickle, matching the torch trainer’s_save_data_detailsimplementation.
"valid_data_model_config": self.valid_dl.dataset.data_model_config,
"valid_data_file_ids": self.valid_dl.dataset.file_ids,
},
open(data_details_path, "wb"),
)
What
PR-0.2 of the ecosystem's "MLflow as database of record" workstream (companion to lnccbrown/ssm-simulators#319). Training runs from both backends now carry the identity a network catalog needs, logged at the CLI level via a shared
log_training_run_identityhelper:model,network_type,backend,input_dim,param_space,param_bounds_json+param_bounds_sha256schema_version=1,phase=train,run_uuid(the uuid in every artifact filename — the MLflow ↔ disk join key),n_training_files_used,training_data_folder,config_sha256,lanfactory_version, SLURM ids--mlflow-run-idresumes (regression-tested)network_config.pickle+ the source YAML are now logged as run artifacts (previously written to disk but never logged)param_boundsno longer lost:DatasetTorchretainsmodel_configfrom the training data, so bounds reachdata_details.pickle— with values sanitized for stdlib pickle (ssm-simulators pickles are cloudpickle-written and can carry lambda boundary functions; regression-tested)network_typeis threaded fromnetwork_configinstead of being inferred from the output type (inference kept as fallback only)train_loss: torch now logs epoch-meantrain_lossatstep=epoch(matching jax) instead of last-minibatch-at-step-count; legacyloss/val_lossstreams unchangedEcosystem impact
Together with #319 this makes every producer run in the ecosystem self-describing under a shared
schema_version=1schema (to be documented in HSSMSpine_docs/mlflow-schema.md). Downstream, LAN_pipeline_minimal's publish/registry tooling resolves artifacts via therun_uuidtag. Additive only; no API breaks.Review provenance
Adversarially reviewed pre-open (multi-agent, findings empirically reproduced). Fixed here: param-collision between identity logging and the trainer's
train_configbulk-log (dropped the entire param batch), resume re-log raising and stranding the join key, train_loss semantics divergence, lambda-carrying model_config breaking_save_data_details.Commands run
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes