Expose ANTs thread count to CLI (#283)

Author: nx10

src/rbc/cli/all.py

@@ -103,6 +103,7 @@ def main(args: AllArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/cli/anatomical.py

@@ -57,6 +57,7 @@ def main(args: AnatomicalArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/cli/base.py

@@ -29,6 +29,7 @@ class BaseArgs:
     session_label: list[str]
     verbose: int
     tmp_dir: Path | None
+    ants_threads: int
 
     @classmethod
     def validate_namespace(cls, ns: argparse.Namespace) -> BaseArgs:
@@ -56,6 +57,12 @@ def validate_namespace(cls, ns: argparse.Namespace) -> BaseArgs:
                 f"Temporary path exists, but is not a directory: {tmp_dir}"
             )
 
+        ants_threads: int = ns.ants_threads
+        if ants_threads < 1:
+            raise ValueError(
+                f"ANTs thread count must be at least 1, got: {ants_threads}"
+            )
+
         return cls(
             input_dirs=tuple(ns.input_dirs),
             output_dir=ns.output_dir,
@@ -64,6 +71,7 @@ def validate_namespace(cls, ns: argparse.Namespace) -> BaseArgs:
             session_label=ns.session_label,
             verbose=ns.verbose,
             tmp_dir=tmp_dir,
+            ants_threads=ants_threads,
         )
 
 

src/rbc/cli/functional.py

@@ -81,6 +81,7 @@ def main(args: FunctionalArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/cli/longitudinal.py

@@ -57,6 +57,7 @@ def main(args: LongitudinalArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/cli/main.py

@@ -79,6 +79,15 @@ def _global_opts() -> argparse.ArgumentParser:
         default=None,
         help="Directory for intermediate files (default: system temp)",
     )
+    global_opts.add_argument(
+        "--ants-threads",
+        type=int,
+        default=1,
+        metavar="N",
+        help="Number of threads for ANTs (ITK) operations (default: 1). "
+        "Values above 1 speed up registration but increase memory usage "
+        "and may produce non-deterministic results",
+    )
     return global_opts
 
 

src/rbc/cli/metrics.py

@@ -89,6 +89,7 @@ def main(args: MetricsArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/cli/qc.py

@@ -55,6 +55,7 @@ def main(args: QCArgs) -> int:
             runner=args.runner,
             verbose=bool(args.verbose),
             tmp_dir=args.tmp_dir,
+            ants_threads=args.ants_threads,
         ),
     )
     return 0

src/rbc/orchestration/__init__.py

@@ -72,11 +72,13 @@ class RunnerConfig:
         runner: Execution backend (local, docker, podman, singularity).
         verbose: Enable verbose output (progress bars, info logging).
         tmp_dir: Temporary directory for intermediate files.
+        ants_threads: Number of threads for ANTs (ITK) operations.
     """
 
     runner: Literal["auto", "local", "docker", "podman", "singularity"] = "local"
     verbose: bool = False
     tmp_dir: Path | None = None
+    ants_threads: int = 1
 
 
 def init_runner(config: RunnerConfig) -> None:
@@ -88,4 +90,7 @@ def init_runner(config: RunnerConfig) -> None:
     ctx = setup_runner(
         runner=config.runner, verbose=config.verbose, tmp_dir=config.tmp_dir
     )
-    ctx.runner.environ = _DEFAULT_ENV_VARS
+    ctx.runner.environ = {
+        **_DEFAULT_ENV_VARS,
+        "ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS": str(config.ants_threads),
+    }

tests/unit/cli/test_all.py

@@ -41,6 +41,7 @@ def base_args(tmp_path: Path) -> argparse.Namespace:
         func_template=None,
         func_template_mask=None,
         func_template_ref=None,
+        ants_threads=1,
     )