Make training runs self-describing in MLflow

.pre-commit-config.yaml

@@ -1,6 +1,7 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.13
+    # Keep in lockstep with the dev-group pin in pyproject.toml (>=0.15.1,<0.16)
+    rev: v0.15.22
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

pyproject.toml

@@ -82,7 +82,9 @@ dev = [
     "pytest-timer>=1.0.0",
     "pytest-xdist>=3.6.1",
     "pytest>=8.3.1",
-    "ruff>=0.14.4",
+    # <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",
     "types-PyYAML",
     "mlflow>=3.14.0",
     "jaxonnxruntime>=0.3",

src/lanfactory/cli/jax_train.py

@@ -28,6 +28,7 @@
 import typer
 from lanfactory.cli.utils import (
     _get_train_network_config,
+    log_training_run_identity,
 )
 from torch.utils.data import DataLoader
 
@@ -471,6 +472,33 @@ def main(
         ),
     )
 
+    if mlflow_tracking_enabled:
+        # Identity params/tags: make the run self-describing (model, bounds,
+        # run_uuid join key). Best-effort inside the helper.
+        log_training_run_identity(
+            model=extra_config["model"],
+            network_type=network_config["network_type"],
+            backend="jax",
+            run_uuid=RUN_ID,
+            config_path=config_path,
+            training_data_folder=training_data_folder,
+            n_training_files=n_training_files,
+            dataset=train_dataset,
+        )
+        try:
+            import mlflow
+
+            # The network config is required to reconstruct the network but was
+            # previously only written to disk, never logged — the MLflow
+            # artifact set alone could not rebuild a network.
+            mlflow.log_artifact(
+                str(networks_path / file_name_suffix), artifact_path="training_output"
+            )
+            if config_path is not None and Path(config_path).is_file():
+                mlflow.log_artifact(str(config_path), artifact_path="training_output")
+        except Exception as e:
+            logger.error("Failed to log config artifacts to MLflow: %s", e)
+
     # Load network
     net = lanfactory.trainers.JaxMLPFactory(
         network_config=deepcopy(network_config),
@@ -494,6 +522,9 @@ def main(
         mlflow_on=mlflow_tracking_enabled,
         save_outputs=True,
         verbose=1,
+        # Pass explicitly: inference from train_output_type mislabels OPN as
+        # cpn (both train on logits).
+        network_type=network_config["network_type"],
     )
     # -------------------------------------------------------------
 

src/lanfactory/cli/torch_train.py

@@ -29,6 +29,7 @@
 import lanfactory
 from lanfactory.cli.utils import (
     _get_train_network_config,
+    log_training_run_identity,
 )
 
 app = typer.Typer()
@@ -475,6 +476,33 @@ def main(
         ),
     )
 
+    if mlflow_tracking_enabled:
+        # Identity params/tags: make the run self-describing (model, bounds,
+        # run_uuid join key). Best-effort inside the helper.
+        log_training_run_identity(
+            model=extra_config["model"],
+            network_type=network_config["network_type"],
+            backend="torch",
+            run_uuid=RUN_ID,
+            config_path=config_path,
+            training_data_folder=training_data_folder,
+            n_training_files=n_training_files,
+            dataset=train_dataset,
+        )
+        try:
+            import mlflow
+
+            # The network config is required to reconstruct the network but was
+            # previously only written to disk, never logged — the MLflow
+            # artifact set alone could not rebuild a network.
+            mlflow.log_artifact(
+                str(networks_path / file_name_suffix), artifact_path="training_output"
+            )
+            if config_path is not None and Path(config_path).is_file():
+                mlflow.log_artifact(str(config_path), artifact_path="training_output")
+        except Exception as e:
+            logger.error("Failed to log config artifacts to MLflow: %s", e)
+
     # Load network
     net = lanfactory.trainers.TorchMLP(
         network_config=deepcopy(network_config),

src/lanfactory/cli/utils.py

@@ -185,3 +185,111 @@ def _get_train_network_config(yaml_config_path: str | Path | None = None, net_in
     config["extra_fields"] = {"model": basic_config["MODEL"]}
 
     return config
+
+
+def log_training_run_identity(
+    *,
+    model: str,
+    network_type: str,
+    backend: str,
+    run_uuid: str,
+    config_path: Path | str | None,
+    training_data_folder: Path | str | None,
+    n_training_files: int,
+    dataset=None,
+) -> None:
+    """Log identity params/tags that make a training run self-describing.
+
+    Every field a catalog needs to answer "which network is this?" is logged on
+    the run itself rather than being recoverable only from experiment-name
+    conventions or by unpickling artifacts. Schema documented in HSSMSpine
+    ``_docs/mlflow-schema.md``.
+
+    Params vs tags split is deliberate, for resume safety: MLflow rejects
+    re-logging a param key with a *different* value, and a run resumed via
+    ``--mlflow-run-id`` re-runs this function with per-invocation values
+    (fresh ``run_uuid``, possibly different data folder / file count).
+
+    - **Params** hold facts immutable for a given network (model,
+      network_type, backend, input_dim, param_space, bounds): a resume
+      re-logs identical values, which MLflow permits.
+    - **Tags** hold per-invocation values (``run_uuid``, training data folder,
+      files used, config sha, lanfactory version): tags are mutable, so a
+      resume overwrites them and the run always reflects its *latest*
+      invocation — whose ``run_uuid`` is the one in the newest artifact
+      filenames, keeping the MLflow<->disk join correct.
+
+    Note the effective file count is tagged ``n_training_files_used``: the
+    trainer separately bulk-logs ``train_config`` whose ``n_training_files``
+    is the configured *cap*, and reusing that param key with the effective
+    value would make MLflow reject the trainer's entire param batch.
+
+    Best-effort: failures are logged, never raised — training must not die on
+    a tracking hiccup. No-op when no MLflow run is active.
+    """
+    import hashlib
+    import json
+    import os
+    from importlib.metadata import PackageNotFoundError, version
+
+    try:
+        import mlflow
+
+        if mlflow.active_run() is None:
+            return
+
+        params: dict = {
+            "model": model,
+            "network_type": network_type,
+            "backend": backend,
+        }
+
+        # The dataset has read a training file: it knows the exact input
+        # dimensionality, and (when the data pickles embed model_config) the
+        # parameter space and training bounds — the fields HSSM consumers
+        # need to use a network correctly.
+        if dataset is not None:
+            input_dim = getattr(dataset, "input_dim", None)
+            if input_dim is not None:
+                params["input_dim"] = int(input_dim)
+            model_config = getattr(dataset, "data_model_config", None)
+            if isinstance(model_config, dict):
+                if "params" in model_config:
+                    params["param_space"] = json.dumps(list(model_config["params"]))
+                if "param_bounds" in model_config:
+                    bounds_json = json.dumps(
+                        np.asarray(model_config["param_bounds"]).tolist()
+                    )
+                    params["param_bounds_json"] = bounds_json
+                    params["param_bounds_sha256"] = hashlib.sha256(
+                        bounds_json.encode()
+                    ).hexdigest()
+
+        mlflow.log_params(params)
+
+        tags = {
+            "schema_version": "1",
+            "phase": "train",
+            "run_uuid": run_uuid,
+            "n_training_files_used": str(n_training_files),
+        }
+        if training_data_folder is not None:
+            tags["training_data_folder"] = str(training_data_folder)
+        if config_path is not None and Path(config_path).is_file():
+            tags["config_sha256"] = hashlib.sha256(
+                Path(config_path).read_bytes()
+            ).hexdigest()
+        try:
+            tags["lanfactory_version"] = version("lanfactory")
+        except PackageNotFoundError:  # pragma: no cover - not hit under uv run
+            pass
+        for env_key, tag in (
+            ("SLURM_JOB_ID", "slurm_job_id"),
+            ("SLURM_ARRAY_JOB_ID", "slurm_array_job_id"),
+            ("SLURM_ARRAY_TASK_ID", "slurm_array_task_id"),
+        ):
+            if os.getenv(env_key):
+                tags[tag] = os.environ[env_key]
+        mlflow.set_tags(tags)
+    except Exception as e:  # noqa: BLE001 - tracking must never kill training
+        logger.error("Failed to log run identity to MLflow: %s", e)

src/lanfactory/trainers/jax_mlp.py

@@ -500,6 +500,7 @@ def train_and_evaluate(
         mlflow_on: bool = False,
         save_outputs: bool = True,
         verbose: int = 1,
+        network_type: str | None = None,
     ) -> train_state.TrainState:
         """Train and evaluate JAXMLP model.
         Arguments
@@ -517,6 +518,11 @@ def train_and_evaluate(
                 Whether to save all files or not.
             verbose (int):
                 The verbosity level.
+            network_type (str | None):
+                The network type ('lan', 'cpn', 'opn', ...), used in output
+                filenames. When None it is inferred from train_output_type —
+                which cannot distinguish cpn from opn (both use logits), so
+                callers that know the type should pass it.
         Returns
         -------
             flax.core.frozen_dict.FrozenDict:
@@ -528,18 +534,22 @@ def train_and_evaluate(
         if mlflow_on:
             self.__try_mlflow(run_id=run_id)
 
-        # Identify network type:
-        if self.model.train_output_type == "logprob":
-            network_type = "lan"
-        elif self.model.train_output_type == "logits":
-            network_type = "cpn"
-        else:
-            network_type = "unknown"
-            print(
-                'Model type identified as "unknown" because '
-                "the training_output_type attribute"
-                ' of the supplied jax model is neither "logprob", nor "logits"'
-            )
+        # Identify network type. Inference from train_output_type is a
+        # fallback only: logits cannot distinguish cpn from opn, which used to
+        # mislabel OPN artifacts as cpn.
+        if network_type is None:
+            if self.model.train_output_type == "logprob":
+                network_type = "lan"
+            elif self.model.train_output_type == "logits":
+                network_type = "cpn"
+            else:  # pragma: no cover - unreachable: JaxMLP.setup() raises on
+                # train_output_type values outside network_type_dict
+                network_type = "unknown"
+                print(
+                    'Model type identified as "unknown" because '
+                    "the training_output_type attribute"
+                    ' of the supplied jax model is neither "logprob", nor "logits"'
+                )
 
         # Initialize Training history
         training_history = pd.DataFrame(
@@ -580,6 +590,21 @@ def train_and_evaluate(
             # Collect loss in training history
             training_history.values[epoch, :] = [int(epoch), float(test_loss)]
 
+            if self.mlflow_on:
+                try:
+                    # Per-epoch metrics under the cross-backend schema names
+                    # (HSSMSpine _docs/mlflow-schema.md), matching torchtrain.
+                    # The per-100-step `loss` metric above is kept as-is.
+                    mlflow.log_metrics(
+                        {
+                            "train_loss": float(train_loss),
+                            "val_loss": float(test_loss),
+                        },
+                        step=int(epoch),
+                    )
+                except Exception:
+                    pass
+
             print(
                 "Epoch: {} / {}, test_loss: {}".format(
                     epoch, self.train_config["n_epochs"], test_loss
@@ -618,8 +643,10 @@ def train_and_evaluate(
             pickle.dump(
                 {
                     "train_data_generator_config": self.train_dl.dataset.data_generator_config,
+                    "train_data_model_config": self.train_dl.dataset.data_model_config,
                     "train_data_file_ids": self.train_dl.dataset.file_ids,
                     "valid_data_generator_config": self.valid_dl.dataset.data_generator_config,
+                    "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"),