added cuda mpi support

README.md

@@ -148,6 +148,10 @@ Precedence for hyperparameters is:
 
    Uses image/executable:
    `kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif` and `./kitrt_code/build_singularity_cuda/KiT-RT`.
+   CUDA runs are dispatched as:
+   `singularity exec --nv ... mpirun -np <gpu_count> ./kitrt_code/build_singularity_cuda/KiT-RT ...`.
+   `<gpu_count>` is auto-detected from `CUDA_VISIBLE_DEVICES` or `nvidia-smi`.
+   Override rank count with `KITRT_CUDA_MPI_RANKS=<N>`.
 
 4. **SLURM mode, raw (no Singularity)**
 

src/simulation_utils.py

@@ -4,6 +4,67 @@
 from src.general_utils import get_user_job_count
 
 
+def _get_cuda_visible_device_count():
+    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if visible_devices is None:
+        return None
+
+    parsed = [value.strip() for value in visible_devices.split(",") if value.strip()]
+    parsed = [value for value in parsed if value != "-1"]
+    return len(parsed)
+
+
+def _query_nvidia_smi_gpu_count():
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+    except FileNotFoundError:
+        return 0
+
+    if result.returncode != 0:
+        return 0
+
+    return len([line for line in (result.stdout or "").splitlines() if line.strip()])
+
+
+def _resolve_cuda_mpi_ranks(quiet=False):
+    override = os.environ.get("KITRT_CUDA_MPI_RANKS")
+    if override is not None:
+        try:
+            rank_count = int(override)
+        except ValueError as e:
+            raise RuntimeError(
+                "Invalid KITRT_CUDA_MPI_RANKS value. Expected positive integer, "
+                f"got: {override!r}"
+            ) from e
+        if rank_count < 1:
+            raise RuntimeError(
+                "Invalid KITRT_CUDA_MPI_RANKS value. Expected >= 1, "
+                f"got: {rank_count}"
+            )
+        return str(rank_count)
+
+    visible_count = _get_cuda_visible_device_count()
+    if visible_count is not None:
+        if visible_count >= 1:
+            return str(visible_count)
+        if not quiet:
+            print("CUDA_VISIBLE_DEVICES is empty; falling back to 1 MPI rank.")
+        return "1"
+
+    detected_gpu_count = _query_nvidia_smi_gpu_count()
+    if detected_gpu_count >= 1:
+        return str(detected_gpu_count)
+
+    if not quiet:
+        print("Could not detect available GPUs; falling back to 1 MPI rank.")
+    return "1"
+
+
 def _run_and_raise(command, mode_label, quiet=False):
     try:
         if quiet:
@@ -64,11 +125,15 @@ def run_cpp_simulation(config_file, quiet=False):
 def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False):
     # Path to the C++ executable
     if use_cuda:
+        mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet)
         singularity_command = [
             "singularity",
             "exec",
             "--nv",
             "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif",
+            "mpirun",
+            "-np",
+            mpi_ranks,
             "./kitrt_code/build_singularity_cuda/KiT-RT",
             config_file,
         ]

tests/test_simulation_utils.py

@@ -0,0 +1,101 @@
+import pytest
+from src import simulation_utils
+
+
+class _DummyResult:
+    def __init__(self, returncode=0, stderr="", stdout=""):
+        self.returncode = returncode
+        self.stderr = stderr
+        self.stdout = stdout
+
+
+def test_run_cpp_simulation_containerized_cuda_uses_mpi(monkeypatch):
+    calls = []
+
+    def fake_run(command, **kwargs):
+        calls.append((command, kwargs))
+        return _DummyResult(returncode=0)
+
+    monkeypatch.delenv("KITRT_CUDA_MPI_RANKS", raising=False)
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
+    monkeypatch.setattr(simulation_utils.subprocess, "run", fake_run)
+
+    simulation_utils.run_cpp_simulation_containerized(
+        "tests/input/validation_tests/SN_solver_hpc/lattice_hpc_200_cuda_order2.cfg",
+        use_cuda=True,
+        quiet=True,
+    )
+
+    assert len(calls) == 1
+    command, kwargs = calls[0]
+    assert command == [
+        "singularity",
+        "exec",
+        "--nv",
+        "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif",
+        "mpirun",
+        "-np",
+        "2",
+        "./kitrt_code/build_singularity_cuda/KiT-RT",
+        "tests/input/validation_tests/SN_solver_hpc/lattice_hpc_200_cuda_order2.cfg",
+    ]
+    assert kwargs.get("stdout") == simulation_utils.subprocess.PIPE
+    assert kwargs.get("stderr") == simulation_utils.subprocess.PIPE
+    assert kwargs.get("text") is True
+
+
+def test_run_cpp_simulation_containerized_cuda_respects_rank_override(monkeypatch):
+    calls = []
+
+    def fake_run(command, **kwargs):
+        calls.append(command)
+        return _DummyResult(returncode=0)
+
+    monkeypatch.setenv("KITRT_CUDA_MPI_RANKS", "4")
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
+    monkeypatch.setattr(simulation_utils.subprocess, "run", fake_run)
+
+    simulation_utils.run_cpp_simulation_containerized(
+        "benchmarks/lattice/example.cfg", use_cuda=True, quiet=True
+    )
+
+    assert len(calls) == 1
+    assert calls[0][6] == "4"
+
+
+def test_run_cpp_simulation_containerized_cuda_queries_nvidia_smi(monkeypatch):
+    calls = []
+
+    def fake_run(command, **kwargs):
+        calls.append((command, kwargs))
+        if command[:3] == [
+            "nvidia-smi",
+            "--query-gpu=index",
+            "--format=csv,noheader",
+        ]:
+            return _DummyResult(returncode=0, stdout="0\n1\n2\n")
+        return _DummyResult(returncode=0)
+
+    monkeypatch.delenv("KITRT_CUDA_MPI_RANKS", raising=False)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    monkeypatch.setattr(simulation_utils.subprocess, "run", fake_run)
+
+    simulation_utils.run_cpp_simulation_containerized(
+        "benchmarks/lattice/example.cfg", use_cuda=True, quiet=True
+    )
+
+    assert len(calls) == 2
+    assert calls[0][0] == [
+        "nvidia-smi",
+        "--query-gpu=index",
+        "--format=csv,noheader",
+    ]
+    assert calls[1][0][6] == "3"
+
+
+def test_run_cpp_simulation_containerized_cuda_rejects_bad_rank_override(monkeypatch):
+    monkeypatch.setenv("KITRT_CUDA_MPI_RANKS", "abc")
+    with pytest.raises(RuntimeError, match="KITRT_CUDA_MPI_RANKS"):
+        simulation_utils.run_cpp_simulation_containerized(
+            "benchmarks/lattice/example.cfg", use_cuda=True, quiet=True
+        )