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: 35 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 (6)
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" |
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