Make training runs self-describing in MLflow - #108
Merged
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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 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 uses a bounded, synchronized development version 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
🚥 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 |
Contributor
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" |
Member
Author
There was a problem hiding this comment.
Fixed in e83d25d — annotated both data_model_config and the pre-existing sibling data_generator_config as str | dict to match runtime behavior.
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
Contributor
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>
Contributor
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"),
)
Copilot: data_model_config/data_generator_config hold dicts once a training file supplies them — annotate str | dict. CodeRabbit: pre-commit pinned ruff v0.14.13, below the dev-group floor; bump to v0.15.22 in lockstep with pyproject (>=0.15.1,<0.16). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (1)
src/lanfactory/trainers/jax_mlp.py:536
train_and_evaluate()gates MLflow logging viaself.mlflow_on, but this flag is only ever set toTruein__try_mlflow()and is never reset toFalsewhenmlflow_on=False. If the sameModelTrainerJaxMLPinstance is reused across multiple calls, a prior MLflow-enabled run can leaveself.mlflow_on=True, causing unexpected metric/artifact logging even when the caller disables MLflow for a later run. Reset the instance flag based on themlflow_onargument at the start of the method.
if mlflow_on:
self.__try_mlflow(run_id=run_id)
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.
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
Chores