From cc3d22c944595431edf2cceed847fac80a904323 Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Wed, 5 Aug 2026 19:59:01 -0400
Subject: [PATCH 1/6] feat: make training runs self-describing in MLflow
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 src/lanfactory/cli/jax_train.py      |  31 +++
 src/lanfactory/cli/torch_train.py    |  28 +++
 src/lanfactory/cli/utils.py          |  89 +++++++++
 src/lanfactory/trainers/jax_mlp.py   |  50 +++--
 src/lanfactory/trainers/torch_mlp.py |  19 +-
 tests/test_run_identity.py           | 269 +++++++++++++++++++++++++++
 6 files changed, 473 insertions(+), 13 deletions(-)
 create mode 100644 tests/test_run_identity.py

diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py
index ffde030..213b5ca 100644
--- a/src/lanfactory/cli/jax_train.py
+++ b/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"],
     )
     # -------------------------------------------------------------
 
diff --git a/src/lanfactory/cli/torch_train.py b/src/lanfactory/cli/torch_train.py
index 84b1855..dea8f60 100644
--- a/src/lanfactory/cli/torch_train.py
+++ b/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),
diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py
index 7933973..ec5789f 100644
--- a/src/lanfactory/cli/utils.py
+++ b/src/lanfactory/cli/utils.py
@@ -185,3 +185,92 @@ 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. ``run_uuid`` is the join key
+    between the MLflow run and the artifact filenames on disk. Schema
+    documented in HSSMSpine ``_docs/mlflow-schema.md``.
+
+    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,
+            "run_uuid": run_uuid,
+            "n_training_files": n_training_files,
+        }
+        if training_data_folder is not None:
+            params["training_data_folder"] = str(training_data_folder)
+
+        if config_path is not None and Path(config_path).is_file():
+            params["config_sha256"] = hashlib.sha256(
+                Path(config_path).read_bytes()
+            ).hexdigest()
+
+        try:
+            params["lanfactory_version"] = version("lanfactory")
+        except PackageNotFoundError:
+            pass
+
+        # 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"}
+        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)
diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py
index 3c3968c..597eccb 100755
--- a/src/lanfactory/trainers/jax_mlp.py
+++ b/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,21 @@ 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:
+                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 +589,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 +642,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"),
diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py
index d439bb7..05b11cf 100755
--- a/src/lanfactory/trainers/torch_mlp.py
+++ b/src/lanfactory/trainers/torch_mlp.py
@@ -66,6 +66,11 @@ def __init__(
         self.label_key = label_key
         self.out_framework = out_framework
         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"
 
         self.tmp_data: dict = {}
 
@@ -126,6 +131,9 @@ def __init_file_shape(self) -> None:
         if "generator_config" in init_file:
             self.data_generator_config = init_file["generator_config"]
 
+        if "model_config" in init_file:
+            self.data_model_config = init_file["model_config"]
+
         if len(self.file_shape_dict["labels"]) > 1:
             self.label_dim = self.file_shape_dict["labels"][1]
         else:
@@ -715,8 +723,15 @@ def train_and_evaluate(
 
             if mlflow_on:
                 try:
+                    # train_loss duplicates loss under the cross-backend schema
+                    # name (HSSMSpine _docs/mlflow-schema.md); loss is kept for
+                    # continuity with existing dashboards.
                     mlflow.log_metrics(
-                        {"loss": float(loss), "val_loss": float(val_loss)},
+                        {
+                            "loss": float(loss),
+                            "train_loss": float(loss),
+                            "val_loss": float(val_loss),
+                        },
                         step=step_cnt,
                     )
                 except Exception as e:
@@ -817,8 +832,10 @@ def _save_data_details(
             pickle.dump(
                 {
                     "train_data_generator_config": train_dl.dataset.data_generator_config,
+                    "train_data_model_config": train_dl.dataset.data_model_config,
                     "train_datafile_ids": train_dl.dataset.file_ids,
                     "valid_data_generator_config": valid_dl.dataset.data_generator_config,
+                    "valid_data_model_config": valid_dl.dataset.data_model_config,
                     "valid_datafile_ids": valid_dl.dataset.file_ids,
                 },
                 f,
diff --git a/tests/test_run_identity.py b/tests/test_run_identity.py
new file mode 100644
index 0000000..99af45f
--- /dev/null
+++ b/tests/test_run_identity.py
@@ -0,0 +1,269 @@
+"""Tests for self-describing training runs (feat/mlflow-self-describing-training).
+
+Covers:
+- ``log_training_run_identity``: the identity params/tags every training run
+  must carry (model, network_type, backend, run_uuid, bounds, ...)
+- ``DatasetTorch`` retaining ``model_config`` from training data files
+- ``_save_data_details`` carrying model_config into data_details.pickle
+- the jax trainer's ``network_type`` passthrough (OPN was mislabeled as cpn)
+"""
+
+import contextlib
+import hashlib
+import json
+import pickle
+import shutil
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+mlflow = pytest.importorskip("mlflow")
+
+import lanfactory
+from lanfactory.cli.utils import log_training_run_identity
+
+
+@pytest.fixture(scope="function", autouse=True)
+def cleanup_mlflow():
+    """Reset MLflow state after each test (mirrors test_mlflow_integration)."""
+    original_uri = mlflow.get_tracking_uri()
+    yield
+    if mlflow.active_run() is not None:
+        mlflow.end_run()
+    mlruns_path = Path.cwd() / "mlruns"
+    if mlruns_path.exists():
+        shutil.rmtree(mlruns_path)
+    with contextlib.suppress(Exception):
+        mlflow.set_tracking_uri(original_uri)
+
+
+@pytest.fixture
+def tmp_tracking(tmp_path):
+    """Isolated sqlite tracking backend."""
+    db = tmp_path / "tracking.db"
+    uri = f"sqlite:///{db.absolute()}"
+    mlflow.set_tracking_uri(uri)
+    return uri
+
+
+MODEL_CONFIG = {
+    "name": "ddm",
+    "params": ["v", "a", "z", "t"],
+    "param_bounds": [[-3.0, 0.3, 0.1, 0.0], [3.0, 2.5, 0.9, 2.0]],
+}
+
+
+def make_training_pickle(path: Path, features_key="lan_data", label_key="lan_labels"):
+    """A minimal training-data pickle shaped like ssm-simulators output."""
+    n_samples, n_features = 64, 6
+    data = {
+        features_key: np.random.randn(n_samples, n_features).astype(np.float32),
+        label_key: np.random.randn(n_samples).astype(np.float32),
+        "generator_config": {"model": "ddm", "generator_approach": "lan"},
+        "model_config": MODEL_CONFIG,
+    }
+    with open(path, "wb") as f:
+        pickle.dump(data, f)
+    return path
+
+
+class TestLogTrainingRunIdentity:
+    def _log_and_fetch(self, tmp_tracking, tmp_path, **overrides):
+        config_yaml = tmp_path / "train.yaml"
+        config_yaml.write_text("NETWORK_TYPE: lan\nMODEL: ddm\n")
+
+        dataset = SimpleNamespace(input_dim=6, data_model_config=MODEL_CONFIG)
+
+        kwargs = {
+            "model": "ddm",
+            "network_type": "lan",
+            "backend": "jax",
+            "run_uuid": "abc123",
+            "config_path": config_yaml,
+            "training_data_folder": tmp_path / "data",
+            "n_training_files": 42,
+            "dataset": dataset,
+        }
+        kwargs.update(overrides)
+
+        mlflow.set_experiment("identity-train-test")
+        with mlflow.start_run() as run:
+            log_training_run_identity(**kwargs)
+            run_id = run.info.run_id
+        return mlflow.tracking.MlflowClient().get_run(run_id)
+
+    def test_identity_params(self, tmp_tracking, tmp_path):
+        run = self._log_and_fetch(tmp_tracking, tmp_path)
+        p = run.data.params
+
+        assert p["model"] == "ddm"
+        assert p["network_type"] == "lan"
+        assert p["backend"] == "jax"
+        assert p["run_uuid"] == "abc123"
+        assert p["n_training_files"] == "42"
+        assert p["input_dim"] == "6"
+        assert json.loads(p["param_space"]) == ["v", "a", "z", "t"]
+        assert len(p["config_sha256"]) == 64
+        assert "lanfactory_version" in p
+
+    def test_param_bounds_json_and_sha_consistent(self, tmp_tracking, tmp_path):
+        run = self._log_and_fetch(tmp_tracking, tmp_path)
+        p = run.data.params
+
+        bounds = json.loads(p["param_bounds_json"])
+        assert bounds == MODEL_CONFIG["param_bounds"]
+        expected_sha = hashlib.sha256(p["param_bounds_json"].encode()).hexdigest()
+        assert p["param_bounds_sha256"] == expected_sha
+
+    def test_schema_tags(self, tmp_tracking, tmp_path):
+        run = self._log_and_fetch(tmp_tracking, tmp_path)
+        assert run.data.tags["schema_version"] == "1"
+        assert run.data.tags["phase"] == "train"
+
+    def test_no_dataset_still_logs_core_identity(self, tmp_tracking, tmp_path):
+        run = self._log_and_fetch(tmp_tracking, tmp_path, dataset=None)
+        p = run.data.params
+        assert p["model"] == "ddm"
+        assert "input_dim" not in p
+        assert "param_bounds_json" not in p
+
+    def test_dataset_without_model_config_skips_bounds(self, tmp_tracking, tmp_path):
+        # DatasetTorch defaults data_model_config to the string "None" when the
+        # training files carry no model_config — must not crash or log junk.
+        dataset = SimpleNamespace(input_dim=6, data_model_config="None")
+        run = self._log_and_fetch(tmp_tracking, tmp_path, dataset=dataset)
+        p = run.data.params
+        assert p["input_dim"] == "6"
+        assert "param_bounds_json" not in p
+
+    def test_noop_without_active_run(self, tmp_tracking, tmp_path):
+        # Must be a silent no-op, not an error.
+        assert mlflow.active_run() is None
+        log_training_run_identity(
+            model="ddm",
+            network_type="lan",
+            backend="jax",
+            run_uuid="x",
+            config_path=None,
+            training_data_folder=None,
+            n_training_files=1,
+        )
+
+
+class TestDatasetModelConfigRetention:
+    def test_model_config_retained_from_training_file(self, tmp_path):
+        f = make_training_pickle(tmp_path / "training_data_x.pickle")
+        dataset = lanfactory.trainers.DatasetTorch(
+            file_ids=[f], batch_size=16, features_key="lan_data", label_key="lan_labels"
+        )
+        assert dataset.data_model_config == MODEL_CONFIG
+        assert dataset.data_generator_config == {
+            "model": "ddm",
+            "generator_approach": "lan",
+        }
+
+    def test_absent_model_config_leaves_default(self, tmp_path):
+        f = tmp_path / "training_data_y.pickle"
+        with open(f, "wb") as fh:
+            pickle.dump(
+                {
+                    "lan_data": np.random.randn(64, 6).astype(np.float32),
+                    "lan_labels": np.random.randn(64).astype(np.float32),
+                },
+                fh,
+            )
+        dataset = lanfactory.trainers.DatasetTorch(
+            file_ids=[f], batch_size=16, features_key="lan_data", label_key="lan_labels"
+        )
+        assert dataset.data_model_config == "None"
+
+
+class TestDataDetailsCarriesModelConfig:
+    def test_torch_save_data_details_includes_model_config(self, tmp_path):
+        from lanfactory.trainers.torch_mlp import ModelTrainerTorchMLP
+
+        stub = SimpleNamespace(
+            dataset=SimpleNamespace(
+                data_generator_config={"model": "ddm"},
+                data_model_config=MODEL_CONFIG,
+                file_ids=["a.pickle"],
+            )
+        )
+        out = tmp_path / "details_data_details.pickle"
+        ModelTrainerTorchMLP._save_data_details(stub, stub, str(out))
+
+        with open(out, "rb") as f:
+            details = pickle.load(f)
+        assert details["train_data_model_config"] == MODEL_CONFIG
+        assert details["valid_data_model_config"] == MODEL_CONFIG
+
+
+class TestJaxNetworkTypePassthrough:
+    def _train_tiny(self, tmp_path, network_type_arg):
+        """One-epoch micro-training run; returns the produced filenames."""
+        from torch.utils.data import DataLoader
+
+        f = make_training_pickle(
+            tmp_path / "training_data_z.pickle",
+            features_key="opn_data",
+            label_key="opn_labels",
+        )
+        dataset = lanfactory.trainers.DatasetTorch(
+            file_ids=[f], batch_size=16, features_key="opn_data", label_key="opn_labels"
+        )
+        dl = DataLoader(dataset, batch_size=None)
+
+        net = lanfactory.trainers.JaxMLPFactory(
+            network_config={
+                "layer_sizes": [8, 1],
+                "activations": ["tanh", "linear"],
+                "train_output_type": "logits",
+                "network_type": "opn",
+            },
+            train=True,
+        )
+        # n_epochs must be >= 2: the warmup_cosine_decay schedule uses
+        # warmup_steps = len(dataset) and decay_steps = len * n_epochs, and
+        # optax requires decay_steps > warmup_steps.
+        trainer = lanfactory.trainers.ModelTrainerJaxMLP(
+            train_config={
+                "n_epochs": 2,
+                "loss": "bcelogit",
+                "optimizer": "adam",
+                "learning_rate": 0.001,
+                "lr_scheduler": None,
+                "lr_scheduler_params": {},
+                "weight_decay": 0.0,
+                "train_output_type": "logits",
+            },
+            train_dl=dl,
+            valid_dl=dl,
+            model=net,
+            seed=42,
+        )
+        out_dir = tmp_path / "nets"
+        trainer.train_and_evaluate(
+            output_folder=out_dir,
+            output_file_id="ddm",
+            run_id="runx",
+            mlflow_on=False,
+            save_outputs=True,
+            verbose=0,
+            network_type=network_type_arg,
+        )
+        return [p.name for p in out_dir.iterdir()]
+
+    def test_explicit_opn_names_files_opn(self, tmp_path):
+        names = self._train_tiny(tmp_path, network_type_arg="opn")
+        assert names, "no output files produced"
+        assert all("_opn_" in n for n in names), names
+
+    def test_fallback_infers_cpn_for_logits(self, tmp_path):
+        # Documents the legacy inference: without an explicit network_type,
+        # logits-trained networks are labeled cpn — the bug the passthrough
+        # exists to avoid.
+        names = self._train_tiny(tmp_path, network_type_arg=None)
+        assert names, "no output files produced"
+        assert all("_cpn_" in n for n in names), names

From 87165e7846d55d177e38c4551571c6b7f6d01422 Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Wed, 5 Aug 2026 20:33:43 -0400
Subject: [PATCH 2/6] fix: resume-safe identity logging; honest cross-backend
 train_loss
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 src/lanfactory/cli/utils.py          | 55 +++++++++++++++++++---------
 src/lanfactory/trainers/torch_mlp.py | 20 ++++++----
 tests/test_run_identity.py           | 50 ++++++++++++++++++++++---
 3 files changed, 94 insertions(+), 31 deletions(-)

diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py
index ec5789f..7b21744 100644
--- a/src/lanfactory/cli/utils.py
+++ b/src/lanfactory/cli/utils.py
@@ -202,9 +202,27 @@ def log_training_run_identity(
 
     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. ``run_uuid`` is the join key
-    between the MLflow run and the artifact filenames on disk. Schema
-    documented in HSSMSpine ``_docs/mlflow-schema.md``.
+    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.
@@ -224,21 +242,7 @@ def log_training_run_identity(
             "model": model,
             "network_type": network_type,
             "backend": backend,
-            "run_uuid": run_uuid,
-            "n_training_files": n_training_files,
         }
-        if training_data_folder is not None:
-            params["training_data_folder"] = str(training_data_folder)
-
-        if config_path is not None and Path(config_path).is_file():
-            params["config_sha256"] = hashlib.sha256(
-                Path(config_path).read_bytes()
-            ).hexdigest()
-
-        try:
-            params["lanfactory_version"] = version("lanfactory")
-        except PackageNotFoundError:
-            pass
 
         # The dataset has read a training file: it knows the exact input
         # dimensionality, and (when the data pickles embed model_config) the
@@ -263,7 +267,22 @@ def log_training_run_identity(
 
         mlflow.log_params(params)
 
-        tags = {"schema_version": "1", "phase": "train"}
+        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:
+            pass
         for env_key, tag in (
             ("SLURM_JOB_ID", "slurm_job_id"),
             ("SLURM_ARRAY_JOB_ID", "slurm_array_job_id"),
diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py
index 05b11cf..e106019 100755
--- a/src/lanfactory/trainers/torch_mlp.py
+++ b/src/lanfactory/trainers/torch_mlp.py
@@ -671,6 +671,7 @@ def train_and_evaluate(
         for epoch in range(self.train_config["n_epochs"]):
             cnt = 0
             epoch_s_t = time()
+            epoch_loss_sum = 0.0
 
             # Training loop
             for xb, yb in self.train_dl:
@@ -692,6 +693,7 @@ def train_and_evaluate(
                 # Log training progress
                 self._log_training_progress(epoch, cnt, loss, verbose)
 
+                epoch_loss_sum += float(loss)
                 cnt += 1
                 step_cnt += 1
 
@@ -723,17 +725,19 @@ def train_and_evaluate(
 
             if mlflow_on:
                 try:
-                    # train_loss duplicates loss under the cross-backend schema
-                    # name (HSSMSpine _docs/mlflow-schema.md); loss is kept for
-                    # continuity with existing dashboards.
+                    # Legacy metrics unchanged (loss = last minibatch, at the
+                    # cumulative batch step). The cross-backend schema metric
+                    # train_loss (HSSMSpine _docs/mlflow-schema.md) is the
+                    # EPOCH-MEAN training loss at step=epoch, matching the jax
+                    # trainer's semantics so the two backends are comparable.
                     mlflow.log_metrics(
-                        {
-                            "loss": float(loss),
-                            "train_loss": float(loss),
-                            "val_loss": float(val_loss),
-                        },
+                        {"loss": float(loss), "val_loss": float(val_loss)},
                         step=step_cnt,
                     )
+                    mlflow.log_metrics(
+                        {"train_loss": epoch_loss_sum / max(cnt, 1)},
+                        step=int(epoch),
+                    )
                 except Exception as e:
                     logger.error(f"Unexpected mlflow error: {e}")
 
diff --git a/tests/test_run_identity.py b/tests/test_run_identity.py
index 99af45f..f68d9a2 100644
--- a/tests/test_run_identity.py
+++ b/tests/test_run_identity.py
@@ -94,19 +94,58 @@ def _log_and_fetch(self, tmp_tracking, tmp_path, **overrides):
             run_id = run.info.run_id
         return mlflow.tracking.MlflowClient().get_run(run_id)
 
-    def test_identity_params(self, tmp_tracking, tmp_path):
+    def test_immutable_facts_are_params(self, tmp_tracking, tmp_path):
         run = self._log_and_fetch(tmp_tracking, tmp_path)
         p = run.data.params
 
         assert p["model"] == "ddm"
         assert p["network_type"] == "lan"
         assert p["backend"] == "jax"
-        assert p["run_uuid"] == "abc123"
-        assert p["n_training_files"] == "42"
         assert p["input_dim"] == "6"
         assert json.loads(p["param_space"]) == ["v", "a", "z", "t"]
-        assert len(p["config_sha256"]) == 64
-        assert "lanfactory_version" in p
+        # per-invocation values must NOT be params (resume safety)
+        for key in ("run_uuid", "n_training_files", "training_data_folder"):
+            assert key not in p
+
+    def test_per_invocation_values_are_tags(self, tmp_tracking, tmp_path):
+        run = self._log_and_fetch(tmp_tracking, tmp_path)
+        t = run.data.tags
+
+        assert t["run_uuid"] == "abc123"
+        assert t["n_training_files_used"] == "42"
+        assert "training_data_folder" in t
+        assert len(t["config_sha256"]) == 64
+        assert "lanfactory_version" in t
+
+    def test_resume_relogs_without_error_and_updates_run_uuid(
+        self, tmp_tracking, tmp_path
+    ):
+        """A resumed run re-logs identity with a fresh run_uuid: params re-log
+        identical values (allowed), tags overwrite — no MlflowException, and
+        the run reflects the latest invocation's artifacts."""
+        config_yaml = tmp_path / "train.yaml"
+        config_yaml.write_text("NETWORK_TYPE: lan\nMODEL: ddm\n")
+        dataset = SimpleNamespace(input_dim=6, data_model_config=MODEL_CONFIG)
+
+        mlflow.set_experiment("identity-resume-test")
+        with mlflow.start_run() as run:
+            for uuid_value, n_files in (("first", 2), ("second", 5)):
+                log_training_run_identity(
+                    model="ddm",
+                    network_type="lan",
+                    backend="jax",
+                    run_uuid=uuid_value,
+                    config_path=config_yaml,
+                    training_data_folder=tmp_path / "data",
+                    n_training_files=n_files,
+                    dataset=dataset,
+                )
+            run_id = run.info.run_id
+
+        fetched = mlflow.tracking.MlflowClient().get_run(run_id)
+        assert fetched.data.tags["run_uuid"] == "second"
+        assert fetched.data.tags["n_training_files_used"] == "5"
+        assert fetched.data.params["model"] == "ddm"
 
     def test_param_bounds_json_and_sha_consistent(self, tmp_tracking, tmp_path):
         run = self._log_and_fetch(tmp_tracking, tmp_path)
@@ -128,6 +167,7 @@ def test_no_dataset_still_logs_core_identity(self, tmp_tracking, tmp_path):
         assert p["model"] == "ddm"
         assert "input_dim" not in p
         assert "param_bounds_json" not in p
+        assert run.data.tags["run_uuid"] == "abc123"
 
     def test_dataset_without_model_config_skips_bounds(self, tmp_tracking, tmp_path):
         # DatasetTorch defaults data_model_config to the string "None" when the

From 2a812ea82dad90ab0f871978dc97419965c133e9 Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Wed, 5 Aug 2026 21:05:22 -0400
Subject: [PATCH 3/6] fix: sanitize retained model_config for stdlib pickle

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>
---
 src/lanfactory/trainers/torch_mlp.py | 26 +++++++++++++++-
 tests/test_run_identity.py           | 44 ++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py
index e106019..38257ed 100755
--- a/src/lanfactory/trainers/torch_mlp.py
+++ b/src/lanfactory/trainers/torch_mlp.py
@@ -23,6 +23,25 @@
 logger = logging.getLogger(__name__)
 
 
+def _picklable_copy(config):
+    """Return a copy of a config dict safe for stdlib pickle.
+
+    Values that stdlib pickle rejects (e.g. local lambdas used as boundary
+    functions in ssm-simulators model configs) are replaced by their repr, so
+    provenance survives without carrying live callables.
+    """
+    if not isinstance(config, dict):
+        return config
+    out = {}
+    for key, value in config.items():
+        try:
+            pickle.dumps(value)
+            out[key] = value
+        except Exception:  # noqa: BLE001 - any unpicklable value gets repr'd
+            out[key] = repr(value)
+    return out
+
+
 class DatasetTorch(torch.utils.data.Dataset):
     """Dataset class for TorchMLP training.
 
@@ -132,7 +151,12 @@ def __init_file_shape(self) -> None:
             self.data_generator_config = init_file["generator_config"]
 
         if "model_config" in init_file:
-            self.data_model_config = init_file["model_config"]
+            # Sanitized at capture: ssm-simulators model_configs can carry
+            # callables (boundary/simulator lambdas) which the training data
+            # pickles tolerate (cloudpickle) but stdlib pickle in
+            # _save_data_details cannot. The catalog-relevant fields (params,
+            # param_bounds, choices, ...) are plain data and pass through.
+            self.data_model_config = _picklable_copy(init_file["model_config"])
 
         if len(self.file_shape_dict["labels"]) > 1:
             self.label_dim = self.file_shape_dict["labels"][1]
diff --git a/tests/test_run_identity.py b/tests/test_run_identity.py
index f68d9a2..761653c 100644
--- a/tests/test_run_identity.py
+++ b/tests/test_run_identity.py
@@ -307,3 +307,47 @@ def test_fallback_infers_cpn_for_logits(self, tmp_path):
         names = self._train_tiny(tmp_path, network_type_arg=None)
         assert names, "no output files produced"
         assert all("_cpn_" in n for n in names), names
+
+
+class TestUnpicklableModelConfig:
+    """ssm-simulators model_configs can carry lambdas (boundary functions).
+
+    Training pickles tolerate them (cloudpickle), but data_details is written
+    with stdlib pickle — the sanitized copy must keep plain fields and
+    stringify callables. Regression test for an order-dependent failure found
+    by the full suite (race_no_bias_angle_2's lambda boundary).
+    """
+
+    def test_lambda_fields_are_stringified_and_data_details_saves(self, tmp_path):
+        import cloudpickle
+        from lanfactory.trainers.torch_mlp import ModelTrainerTorchMLP
+
+        config_with_lambda = dict(MODEL_CONFIG)
+        config_with_lambda["boundary"] = lambda t: 1.0
+
+        f = tmp_path / "training_data_lambda.pickle"
+        with open(f, "wb") as fh:
+            cloudpickle.dump(
+                {
+                    "lan_data": np.random.randn(64, 6).astype(np.float32),
+                    "lan_labels": np.random.randn(64).astype(np.float32),
+                    "generator_config": {"model": "race_no_bias_angle_2"},
+                    "model_config": config_with_lambda,
+                },
+                fh,
+            )
+
+        dataset = lanfactory.trainers.DatasetTorch(
+            file_ids=[f], batch_size=16, features_key="lan_data", label_key="lan_labels"
+        )
+        # plain fields survive, the lambda is stringified
+        assert dataset.data_model_config["params"] == MODEL_CONFIG["params"]
+        assert isinstance(dataset.data_model_config["boundary"], str)
+
+        # and the data_details write (stdlib pickle) must succeed
+        stub = SimpleNamespace(dataset=dataset)
+        out = tmp_path / "d_data_details.pickle"
+        ModelTrainerTorchMLP._save_data_details(stub, stub, str(out))
+        with open(out, "rb") as fh:
+            details = pickle.load(fh)
+        assert details["train_data_model_config"]["params"] == MODEL_CONFIG["params"]

From dba61a129bd944cbaad9b2d0970f18d4f16d712a Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Sat, 8 Aug 2026 20:18:18 -0400
Subject: [PATCH 4/6] ci: pin ruff below 0.16 to stop lint drift
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 pyproject.toml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7277bf6..ebec8f3 100755
--- a/pyproject.toml
+++ b/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",

From 9efc052f5d07909ff9c2adf441bbac2dfda0931a Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Sat, 8 Aug 2026 20:45:31 -0400
Subject: [PATCH 5/6] test: cover SLURM tag capture and non-dict config
 passthrough

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>
---
 src/lanfactory/cli/utils.py        |  2 +-
 src/lanfactory/trainers/jax_mlp.py |  3 ++-
 tests/test_run_identity.py         | 24 ++++++++++++++++++++++--
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py
index 7b21744..311dc05 100644
--- a/src/lanfactory/cli/utils.py
+++ b/src/lanfactory/cli/utils.py
@@ -281,7 +281,7 @@ def log_training_run_identity(
             ).hexdigest()
         try:
             tags["lanfactory_version"] = version("lanfactory")
-        except PackageNotFoundError:
+        except PackageNotFoundError:  # pragma: no cover - not hit under uv run
             pass
         for env_key, tag in (
             ("SLURM_JOB_ID", "slurm_job_id"),
diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py
index 597eccb..70befb4 100755
--- a/src/lanfactory/trainers/jax_mlp.py
+++ b/src/lanfactory/trainers/jax_mlp.py
@@ -542,7 +542,8 @@ def train_and_evaluate(
                 network_type = "lan"
             elif self.model.train_output_type == "logits":
                 network_type = "cpn"
-            else:
+            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 '
diff --git a/tests/test_run_identity.py b/tests/test_run_identity.py
index 761653c..d254a57 100644
--- a/tests/test_run_identity.py
+++ b/tests/test_run_identity.py
@@ -19,11 +19,11 @@
 import numpy as np
 import pytest
 
-mlflow = pytest.importorskip("mlflow")
-
 import lanfactory
 from lanfactory.cli.utils import log_training_run_identity
 
+mlflow = pytest.importorskip("mlflow")
+
 
 @pytest.fixture(scope="function", autouse=True)
 def cleanup_mlflow():
@@ -117,6 +117,19 @@ def test_per_invocation_values_are_tags(self, tmp_tracking, tmp_path):
         assert len(t["config_sha256"]) == 64
         assert "lanfactory_version" in t
 
+    def test_slurm_env_ids_become_tags(self, tmp_tracking, tmp_path, monkeypatch):
+        # On Oscar these are the join keys between MLflow runs and sacct
+        # accounting; a run launched outside SLURM simply omits them.
+        monkeypatch.setenv("SLURM_JOB_ID", "12345")
+        monkeypatch.setenv("SLURM_ARRAY_JOB_ID", "12300")
+        monkeypatch.setenv("SLURM_ARRAY_TASK_ID", "7")
+        run = self._log_and_fetch(tmp_tracking, tmp_path)
+        t = run.data.tags
+
+        assert t["slurm_job_id"] == "12345"
+        assert t["slurm_array_job_id"] == "12300"
+        assert t["slurm_array_task_id"] == "7"
+
     def test_resume_relogs_without_error_and_updates_run_uuid(
         self, tmp_tracking, tmp_path
     ):
@@ -318,6 +331,13 @@ class TestUnpicklableModelConfig:
     by the full suite (race_no_bias_angle_2's lambda boundary).
     """
 
+    def test_non_dict_config_passes_through(self):
+        # DatasetTorch's data_model_config defaults to the string "None";
+        # the sanitizer must hand non-dict values back untouched.
+        from lanfactory.trainers.torch_mlp import _picklable_copy
+
+        assert _picklable_copy("None") == "None"
+
     def test_lambda_fields_are_stringified_and_data_details_saves(self, tmp_path):
         import cloudpickle
         from lanfactory.trainers.torch_mlp import ModelTrainerTorchMLP

From e83d25de75d3e033460714831a6fbe9c5e7a24c7 Mon Sep 17 00:00:00 2001
From: Alexander Fengler <alexanderfengler@gmx.de>
Date: Sun, 9 Aug 2026 00:59:20 -0400
Subject: [PATCH 6/6] review: widen dataset config annotations; align
 pre-commit ruff with pin
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 .pre-commit-config.yaml              | 3 ++-
 src/lanfactory/trainers/torch_mlp.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 94eb114..e6dfe31 100644
--- a/.pre-commit-config.yaml
+++ b/.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]
diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py
index 38257ed..352e8a5 100755
--- a/src/lanfactory/trainers/torch_mlp.py
+++ b/src/lanfactory/trainers/torch_mlp.py
@@ -84,12 +84,12 @@ def __init__(
         self.features_key = features_key
         self.label_key = label_key
         self.out_framework = out_framework
-        self.data_generator_config: str = "None"
+        self.data_generator_config: str | dict = "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"
+        self.data_model_config: str | dict = "None"
 
         self.tmp_data: dict = {}