fixup: StrictPositiveInt, shared dataloader-kwargs helper, compat test

deruyter92 · Oct 13, 2026 · Oct 13, 2026 · Oct 16, 2026 · Oct 16, 2026 · Oct 16, 2026
commit 658becf43b5eee02a0c871dbe245d4d989eccad3
diff --git a/deeplabcut/pose_estimation_pytorch/apis/training.py b/deeplabcut/pose_estimation_pytorch/apis/training.py
@@ -41,6 +41,30 @@
 from deeplabcut.pose_estimation_pytorch.task import Task
 
 
+def build_dataloader_kwargs(
+    num_workers: int,
+    persistent_workers: bool = False,
+    prefetch_factor: int | None = None,
+) -> dict:
+    """Builds extra DataLoader kwargs for the throughput settings.
+
+    ``persistent_workers``/``prefetch_factor`` only apply with
+    ``num_workers > 0``; otherwise they are dropped with a warning.
+    """
+    kwargs: dict = {}
+    if num_workers > 0:
+        if persistent_workers:
+            kwargs["persistent_workers"] = True
+        if prefetch_factor is not None:
+            kwargs["prefetch_factor"] = prefetch_factor
+    elif persistent_workers or prefetch_factor is not None:
+        logging.warning(
+            "dataloader_persistent_workers/prefetch_factor require "
+            "dataloader_workers > 0; ignoring them for single-process loading."
+        )
+    return kwargs
+
+
 def train(
     loader: Loader,
     run_config: dict,
@@ -155,17 +179,11 @@ def train(
     pin_memory = run_config["train_settings"]["dataloader_pin_memory"]
     persistent_workers = run_config["train_settings"].get("dataloader_persistent_workers", False)
     prefetch_factor = run_config["train_settings"].get("dataloader_prefetch_factor")
-    dataloader_kwargs: dict = {}
-    if num_workers > 0:
-        if persistent_workers:
-            dataloader_kwargs["persistent_workers"] = True
-        if prefetch_factor is not None:
-            dataloader_kwargs["prefetch_factor"] = prefetch_factor
-    elif persistent_workers or prefetch_factor is not None:
-        logging.warning(
-            "dataloader_persistent_workers/prefetch_factor require "
-            "dataloader_workers > 0; ignoring them for single-process loading."
-        )
+    dataloader_kwargs = build_dataloader_kwargs(
+        num_workers,
+        persistent_workers=persistent_workers,
+        prefetch_factor=prefetch_factor,
+    )
     train_dataloader = DataLoader(
         train_dataset,
         batch_size=batch_size,

diff --git a/deeplabcut/pose_estimation_pytorch/config/training.py b/deeplabcut/pose_estimation_pytorch/config/training.py
@@ -10,9 +10,8 @@
 #
 """Training configuration classes for DeepLabCut pose estimation models."""
 
-from pydantic import Field
-
 from deeplabcut.core.config import DLCBaseConfig
+from deeplabcut.core.config.validation import StrictPositiveInt
 from deeplabcut.core.weight_init import WeightInitialization
 
 
@@ -39,7 +38,7 @@ class TrainSettingsConfig(DLCBaseConfig):
     dataloader_workers: int = 0
     dataloader_pin_memory: bool = False
     dataloader_persistent_workers: bool = False
-    dataloader_prefetch_factor: int | None = Field(default=None, ge=1)
+    dataloader_prefetch_factor: StrictPositiveInt | None = None
     display_iters: int = 500
     epochs: int = 200
     seed: int = 42

diff --git a/tests/pose_estimation_pytorch/apis/test_apis_training.py b/tests/pose_estimation_pytorch/apis/test_apis_training.py
@@ -16,7 +16,7 @@
 
 import pytest
 
-from deeplabcut.pose_estimation_pytorch.apis.training import train
+from deeplabcut.pose_estimation_pytorch.apis.training import build_dataloader_kwargs, train
 from deeplabcut.pose_estimation_pytorch.config import make_pytorch_pose_config
 from deeplabcut.pose_estimation_pytorch.config.pose import PoseConfig
 from deeplabcut.pose_estimation_pytorch.task import Task
@@ -92,34 +92,20 @@ def test_train_passes_dataloader_throughput_settings(
     assert train_kwargs["prefetch_factor"] == 4
 
 
-@patch("deeplabcut.pose_estimation_pytorch.apis.training.DataLoader")
-@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_transforms", return_value=Mock())
-@patch("deeplabcut.pose_estimation_pytorch.apis.training.PoseModel.build", return_value=Mock())
-@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_training_runner", return_value=Mock())
-def test_train_ignores_throughput_settings_without_workers(
-    mock_build_runner: Mock,
-    mock_build_model: Mock,
-    mock_build_transforms: Mock,
-    mock_dataloader: Mock,
-    tmp_path: Path,
-    caplog: pytest.LogCaptureFixture,
-) -> None:
-    """Throughput settings are dropped (with a warning) for single-process loading."""
-    run_config = _minimal_run_config(tmp_path)
-    run_config["train_settings"]["dataloader_workers"] = 0
-    run_config["train_settings"]["dataloader_persistent_workers"] = True
+def test_build_dataloader_kwargs_forwards_throughput_settings() -> None:
+    """Helper passes the settings through when workers are used."""
+    assert build_dataloader_kwargs(4, persistent_workers=True, prefetch_factor=4) == {
+        "persistent_workers": True,
+        "prefetch_factor": 4,
+    }
+    assert build_dataloader_kwargs(4) == {}
 
-    loader = Mock()
-    loader.model_folder = tmp_path
-    loader.model_cfg = run_config
-    loader.create_dataset = Mock(
-        side_effect=[Mock(__len__=Mock(return_value=2)), Mock(__len__=Mock(return_value=1))]
-    )
 
+def test_build_dataloader_kwargs_ignores_throughput_settings_without_workers(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Helper drops the settings (with a warning) for single-process loading."""
     with caplog.at_level(logging.WARNING):
-        train(loader=loader, run_config=run_config, task=Task.BOTTOM_UP, device="cpu", snapshot_path=None)
+        assert build_dataloader_kwargs(0, persistent_workers=True, prefetch_factor=4) == {}
 
-    train_kwargs = mock_dataloader.call_args_list[0].kwargs
-    assert "persistent_workers" not in train_kwargs
-    assert "prefetch_factor" not in train_kwargs
     assert "dataloader_workers > 0" in caplog.text
diff --git a/tests/pose_estimation_pytorch/config/test_train_settings.py b/tests/pose_estimation_pytorch/config/test_train_settings.py
@@ -42,3 +42,16 @@ def test_dataloader_prefetch_factor_must_be_positive():
     """prefetch_factor=0 is rejected at config load, not deep in training."""
     with pytest.raises(ValidationError):
         TrainSettingsConfig(dataloader_prefetch_factor=0)
+
+
+def test_old_configs_without_throughput_settings_still_validate():
+    """Configs written before the new keys existed load with safe defaults.
+
+    This pins the backward-compatibility contract that train() relies on
+    when it reads the settings with .get() fallbacks.
+    """
+    cfg = TrainSettingsConfig.model_validate({"batch_size": 4})
+
+    assert cfg.batch_size == 4
+    assert cfg.dataloader_persistent_workers is False
+    assert cfg.dataloader_prefetch_factor is None