I changed prep_utils.py (line 266) so every ANTs subprocess now sets:

ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS
OMP_NUM_THREADS
OMP_THREAD_LIMIT

README.md

@@ -38,6 +38,11 @@ These defaults can be overridden with:
 - `HF_RUNTIME_REPO_ID`
 - `HF_RUNTIME_REVISION`
 
+Registration threading can be controlled with:
+
+- `PREP_ANTS_THREADS` to set the exact ANTs/ITK thread count
+- `PREP_ANTS_THREAD_CAP` to cap the automatic count (default: `8`)
+
 ## Experimental Mask Editor
 
 An isolated post-segmentation mask editor scaffold is available behind the

src/data_prep/prep_utils.py

@@ -263,12 +263,52 @@ def _command_timeout_seconds() -> int | None:
     return None if seconds <= 0 else seconds
 
 
+def _available_cpu_count() -> int:
+    try:
+        return max(1, len(os.sched_getaffinity(0)))
+    except Exception:
+        return max(1, os.cpu_count() or 1)
+
+
+def _ants_thread_count() -> int:
+    raw = os.environ.get("PREP_ANTS_THREADS", "").strip()
+    if raw:
+        try:
+            return max(1, int(float(raw)))
+        except ValueError:
+            print(f"[setup] invalid PREP_ANTS_THREADS={raw!r}; using automatic thread count")
+
+    cap_raw = os.environ.get("PREP_ANTS_THREAD_CAP", "8").strip()
+    try:
+        cap = max(1, int(float(cap_raw)))
+    except ValueError:
+        print(f"[setup] invalid PREP_ANTS_THREAD_CAP={cap_raw!r}; using 8")
+        cap = 8
+    return max(1, min(_available_cpu_count(), cap))
+
+
+def _subprocess_env() -> dict[str, str]:
+    env = os.environ.copy()
+    threads = str(_ants_thread_count())
+    env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"] = threads
+    env["OMP_NUM_THREADS"] = threads
+    env["OMP_THREAD_LIMIT"] = threads
+    return env
+
+
 def _run(cmd: list[str]):
     cmd_str = " ".join(map(str, cmd))
     timeout_sec = _command_timeout_seconds()
     timeout_label = f"{timeout_sec}s" if timeout_sec is not None else "disabled"
+    env = _subprocess_env()
+    threads = env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"]
     print(">>", cmd_str)
     print(f"[cmd] timeout={timeout_label}")
+    print(
+        f"[cmd] ants_threads={threads} "
+        f"available_cpus={_available_cpu_count()} "
+        f"thread_cap={os.environ.get('PREP_ANTS_THREAD_CAP', '8')}"
+    )
     t0 = time.monotonic()
     try:
         res = subprocess.run(
@@ -277,6 +317,7 @@ def _run(cmd: list[str]):
             stderr=subprocess.STDOUT,
             text=True,
             timeout=timeout_sec,
+            env=env,
         )
     except subprocess.TimeoutExpired as exc:
         elapsed = time.monotonic() - t0