BUG: `multi_voxel_fit` decorator forwards orchestration kwargs into the per-voxel fit function

[oesteban]

Summary

Warning

Provenance. This issue was fundamentally diagnosed by Claude Code. I had independently run into this problem before I started using Claude Code, and I have reviewed and revised the entire report.

Exact prompt that drove the related-work research

write a DIPY issue for the kwargs-leak bug so I can copy and paste into their repo.

- Make sure you add: permanent links to code sections involved,
   e.g., `https://github.com/dipy/dipy/blob/d9e9e30cedf89562b88cff120b0e69efb028d912/dipy/reconst/dti.py#L52-L54`
   if something happened on those four lines.
- Be brief and to the point.
- Link the nifreeze issues originating the concern.
- Clearly write how you confirmed the bug exists.

BUT MOST IMPORTANT of ALL: **CHECK ALL OPEN AND CLOSED PRs and ISSUES**
to map out all the related work in the repo.

I then executed a follow-up prompt to extract the evironment information required in the issue template.

Summary

When a @multi_voxel_fit-decorated fit/multi_fit is called with a parallel engine (engine="joblib", "dask", or "ray"), the decorator forwards the orchestration kwargs (engine, n_jobs, vox_per_chunk, verbose) into the per-voxel fit function, which doesn't accept them. Every reconstruction model is affected; DKI is the concrete example below. This makes DIPY's own multi-voxel parallelization unusable from the public API.

Reproduce (DIPY 1.10.0, pip)

import numpy as np
from dipy.core.gradients import gradient_table
from dipy.reconst.dki import DiffusionKurtosisModel

rng = np.random.default_rng(0)
bvals = np.r_[0, np.full(15, 1000), np.full(15, 2000), np.full(15, 3000)].astype(float)
bvecs = np.zeros((46, 3)); bvecs[1:] = rng.standard_normal((45, 3))
bvecs[1:] /= np.linalg.norm(bvecs[1:], axis=1, keepdims=True)
gtab = gradient_table(bvals, bvecs=bvecs)
data = rng.random((100, bvals.size)) * 100 + 100

m = DiffusionKurtosisModel(gtab)
m.multi_fit(data, engine="joblib", n_jobs=2)   # the @multi_voxel_fit-decorated entry point
m.fit(data, engine="joblib", n_jobs=2)         # convenience wrapper

multi_fit(engine="joblib"): TypeError: ls_fit_dki() got an unexpected keyword argument 'engine'
fit(engine="joblib"):       TypeError: DiffusionKurtosisModel.fit() got an unexpected keyword argument 'engine'

(The second is secondary — in 1.10.0 DKIModel.fit has signature (data, *, mask=None) and doesn't forward engine; multi_fit is the decorated path and the one that hits the leak.)

Root cause

In multi_voxel_fit's parallel branch, the per-chunk kwargs are built from a full kwargs.copy() with the orchestration keys never removed, and _parallel_fit_worker forwards them straight into single_voxel_fit(...):

1.10.0 (4eb161f)

• multi_voxel.py#L115-L125 chunks kwargs built unstripped:
  

  dipy/dipy/reconst/multi_voxel.py

  Lines 115 to 125 in 4eb161f

  	kwargs_chunks = []
  	for ii in range(0, data_to_fit.shape[0], vox_per_chunk):
  	kw = kwargs.copy()
  	if weights_is_array:
  	kw["weights"] = weights_to_fit[ii : ii + vox_per_chunk]
  	kwargs_chunks.append(kw)
  	parallel_kwargs = {}
  	for kk in ["n_jobs", "vox_per_chunk", "engine", "verbose"]:
  	if kk in kwargs:
  	parallel_kwargs[kk] = kwargs[kk]

• #L15-L35, worker forwards to single_voxel_fit:
  

  dipy/dipy/reconst/multi_voxel.py

  Lines 15 to 35 in 4eb161f

  	def _parallel_fit_worker(vox_data, single_voxel_fit, **kwargs):
  	"""
  	Works on a chunk of voxel data to create a list of
  	single voxel fits.
  	Parameters
  	----------
  	vox_data : ndarray, shape (n_voxels, ...)
  	The data to fit.
  	single_voxel_fit : callable
  	The fit function to use on each voxel.
  	"""
  	vox_weights = kwargs.pop("weights", None)
  	if type(vox_weights) is np.ndarray:
  	return [
  	single_voxel_fit(data, **(dict({"weights": weights}, **kwargs)))
  	for data, weights in zip(vox_data, vox_weights)
  	]
  	else:
  	return [single_voxel_fit(data, **kwargs) for data in vox_data]

master (d9e9e30, after the recent rework)

• kw = kwargs.copy() unstripped at #L324-L344: same pattern persists:
  

  dipy/dipy/reconst/multi_voxel.py

  Lines 324 to 344 in d9e9e30

  	kwargs_chunks = []
  	for ii in range(0, data_to_fit.shape[0], vox_per_chunk):
  	kw = kwargs.copy()
  	if batched:
  	kw["_batched"] = True
  	if use_raw:
  	kw["_raw"] = True
  	if weights_is_array:
  	kw["weights"] = weights_to_fit[ii : ii + vox_per_chunk]
  	kwargs_chunks.append(kw)
  	parallel_kwargs = {}
  	for kk in [
  	"n_jobs",
  	"vox_per_chunk",
  	"engine",
  	"verbose",
  	"inflight_cap",
  	]:
  	if kk in kwargs:
  	parallel_kwargs[kk] = kwargs[kk]

• #L101-L113, worker pops only _sobj_*/_batched/weights, not the orchestration keys, then forwards:
  

  dipy/dipy/reconst/multi_voxel.py

  Lines 101 to 113 in d9e9e30

  	batched = kwargs.pop("_batched", False)
  	vox_weights = kwargs.pop("weights", None)
  	if batched:
  	if type(vox_weights) is np.ndarray:
  	return fit_func(vox_data, weights=vox_weights, **kwargs)
  	return fit_func(vox_data, **kwargs)
  	if type(vox_weights) is np.ndarray:
  	return [
  	fit_func(data, **(dict({"weights": weights}, **kwargs)))
  	for data, weights in zip(vox_data, vox_weights)
  	]
  	return [fit_func(data, **kwargs) for data in vox_data]

The batched-serial path on master already strips these keys (if k not in ("engine", "n_jobs", "vox_per_chunk", "verbose")), so the parallel branch just needs the same treatment — strip them from kwargs_chunks (or pop them in _parallel_fit_worker) before calling single_voxel_fit.

How I confirmed

1. Runtime: the repro above on dipy==1.10.0 (pip) yields the two TypeErrors.
2. Source: read the decorator on both 4eb161f (1.10.0) and d9e9e30 (master HEAD) — the orchestration kwargs are copied into per-voxel kwargs and never removed before dispatch in the parallel branch.

Platform

Linux 6.17.0-35-generic x86_64 GNU/Linux

DIPY version

1.10.0 (pip)

Environment

dipy==1.10.0
numpy==2.3.5
scipy==1.15.1
joblib==1.4.2
dask==2025.5.1
nibabel==5.3.2
threadpoolctl==3.5.0
tqdm==4.67.1
# ray: not installed — engine="ray" unavailable; bug reproduces with engine="joblib"/"dask"

Python version

3.12.8

Additional information

Related work in this repo

• Embed parallelization into the multi_voxel_fit decorator. #2593 (merged) — embedded parallelization into the multi_voxel_fit decorator (origin of engine/paramap).
• NF: FORCE recon model #3844 (merged, 2026-03) — "NF: FORCE recon model" / multivoxel chunk integration + serialization rework = current master state; leak still present after it.
• ENH: Ray-aware CSD and MSMT-CSD workflows (#3503, #3617) #3814 (merged) + Create new interface for MTMS CSD and standard CSD that use Ray properly #3503 — Ray-aware CSD/MSMT-CSD interfaces.
• RF: Parallel quantize_evecs #3605 (merged) — parallel quantize_evecs via paramap.
• 22f767b — prior multi_voxel parallel bug (single-CPU ZeroDivisionError), already fixed.
• WIP : Multiprocessing - implemented a parallel_voxel_fit decorator  #1135 (closed, 2016) — early parallel_voxel_fit attempt.
• No open issue/PR currently addresses this leak.

Origin

Surfaced while parallelizing DKI for head-motion estimation in NiFreeze: nipreps/nifreeze#442 (earlier attempt nipreps/nifreeze#142; interim fix nipreps/nifreeze#443, which works around it by avoiding the engine path).