Expose DataLoader persistent_workers and prefetch_factor in train_settings - #3550

@deruyter92 You may want to double-check the pydantic/config side, in case I forgot anything. Thanks!

Thanks for the review! Addressed all three points in the latest push: prefetch_factor now uses StrictPositiveInt from core config validation; the kwarg branching moved into a shared build_dataloader_kwargs helper with direct unit tests (kept one entry-level train test); and the config test gained a backward-compat case pinning that old dicts without the new keys still validate. deruyter92, welcome your eyes on the pydantic side whenever convenient.

Thanks for the detailed review — I agree, and I think dataloader_kwargs is the better design. Modelling DataLoader's API field by field would mean a new config version every time torch adds an argument, and re-implementing validation torch already does. I hadn't weighed that against the flexibility gain properly.

I didn't want to push my own version over yours, so I've sent the one thing that still seemed missing as a PR into your branch: #3555, based on jaap/config-extra-kwargs.

It checks persistent_workers / prefetch_factor against dataloader_workers while the config loads. With dataloader_workers: 0, DataLoader raises persistent_workers option needs num_workers > 0, but only once training starts, and the message names neither dataloader_kwargs nor the field to change — so a config that validates cleanly still fails later. reject_reserved_keys catches the keys; this catches the combination. Values that ask for nothing (persistent_workers: false, prefetch_factor: null) still pass, matching DataLoader.

Take it or drop it as you see fit — it merges into #3554, and I'll close #3550 since your version supersedes it. If you'd rather have it in a separate PR against main, say so and I'll retarget it.

On C-Achard's comments in #3550: the StrictPositiveInt note doesn't apply any more since persistent_workers/prefetch_factor are no longer typed fields, and the "tests pydantic" and heavyweight-patching points were both about tests that came with my version — your test file already reads better than mine did. The wandb_kwargs simplification you mentioned is still open; happy to take it if useful.

Superseded by #3554 and #3555: dataloader_kwargs replaces the two typed fields this PR added. Keeping the branch in case anything here is still wanted. Thanks for the review.

Superseded by #3554 and #3555: dataloader_kwargs replaces the two typed fields this PR added. Keeping the branch in case anything here is still wanted. Thanks for the review.

Instead of merging the add-on PRs, I'd prefer to merge this original PR since it contains most of the discussion and context.

I've merged my suggestions from #3554 into this branch. Glad you agree with the approach.

Applied both suggestions — the hint now names data.train.collate, the key apis/training.py actually reads (run_config["data"]["train"].get("collate")), in the mapping and in the test.

On the validator: that was #3555, and I closed it on @deruyter92's call rather than because it didn't work. He agreed with @Ayskin in #3554 — keep dataloader_kwargs a pass-through and let PyTorch own the validation, on the grounds that PyTorch raises in DataLoader.__init__, "right after the config validation, so the practical benefit would be a few seconds at best". Since you're asking for it now, that is a call between the two of you, and I'll follow it either way.

One thing that may shift it: in the current code the raise is not right after config validation. train() builds the runner, creates both datasets through loader.create_dataset(...), and only then constructs the training DataLoader, so on a large project the error arrives after dataset indexing rather than seconds later.

The check itself is written and tested (model_validator → reject_multiprocessing_only_kwargs; it rejects only options that actually ask for something, so persistent_workers: false and prefetch_factor: null still pass, as they do in torch). Say the word and I'll push it as one commit; otherwise the pass-through stays as it is and the docs keep stating the requirement.

One thing that may shift it: in the current code the raise is not right after config validation. train() builds the runner, creates both datasets through loader.create_dataset(...), and only then constructs the training DataLoader, so on a large project the error arrives after dataset indexing rather than seconds later.

Don't mean to keep this open if it was already agreed on, any thoughts @Ayskin @deruyter92 ?

@lu123bing after internal discussion we agreed that the current state is great as is. No need for further adjustments! Thanks for your contribution!

Heads-up on the Dev docs build job on this PR — unrelated to the diff. It fails with

TypeError: Highlight.__init__() missing 1 required positional argument: 'md'

on reference/deeplabcut/index.md: mkdocstrings 1.0.6 — the last release supporting Python 3.10, which is what the job pins — calls pymdownx.highlight.Highlight() without the md argument that pymdown-extensions 12.2 made required. The job installs the group unpinned (uv pip install --system --group dev-docs), so it now resolves pymdown-extensions==12.3.

Reproduced locally on Python 3.10 with the same install command; the build passes again with pymdown-extensions<12.2. Note uv.lock still pins pymdown-extensions 10.21.3 — the docs job doesn't use the lock.

Two options on your side: cap pymdown-extensions<12.2 in the dev-docs group (which needs a lockfile refresh), or bump the docs job to Python 3.11+, where mkdocstrings>=1.0.7 supports pymdown-extensions 12.2. I haven't changed anything — just flagging it so the red check isn't misread as fallout from this PR.