From e8aee37a26790935b0ebbaf384c2d18da97448a5 Mon Sep 17 00:00:00 2001 From: Steffen Date: Fri, 13 Feb 2026 16:29:31 -0500 Subject: [PATCH 1/5] added cuda mpi support --- README.md | 4 + explore_exploit_hohlraum.py | 202 --------------------------------- src/simulation_utils.py | 65 +++++++++++ tests/test_simulation_utils.py | 101 +++++++++++++++++ 4 files changed, 170 insertions(+), 202 deletions(-) delete mode 100644 explore_exploit_hohlraum.py create mode 100644 tests/test_simulation_utils.py diff --git a/README.md b/README.md index 4fb827c..a301a38 100644 --- a/README.md +++ b/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 ./kitrt_code/build_singularity_cuda/KiT-RT ...`. + `` is auto-detected from `CUDA_VISIBLE_DEVICES` or `nvidia-smi`. + Override rank count with `KITRT_CUDA_MPI_RANKS=`. 4. **SLURM mode, raw (no Singularity)** diff --git a/explore_exploit_hohlraum.py b/explore_exploit_hohlraum.py deleted file mode 100644 index ae222c5..0000000 --- a/explore_exploit_hohlraum.py +++ /dev/null @@ -1,202 +0,0 @@ -import pyapprox.multifidelity.etc as etc -from pyapprox.variables.joint import IndependentMarginalsVariable -from pyapprox.multifidelity.groupacv import get_model_subsets -import numpy as np -from scipy import stats -import time -import pickle -import glob -import os - -# import model from client_hohlraum -np.random.seed(1) - -from src.models.hohlraum import get_qois_col_names, model -from src.general_utils import parse_args -from src.general_utils import ( - create_hohlraum_samples_from_param_range, - load_hohlraum_samples_from_npz, - delete_slurm_scripts, -) -from src.simulation_utils import execute_slurm_scripts, wait_for_slurm_jobs -from src.config_utils import read_username_from_config - - -np.random.seed(1) -# import model from client_hohlraum - - -def generate_model(cell_length, nquad, hpc_operation, singularity_hpc, qoi_idx): - """ - Return a function that can be passed to the AETC algorithm - cell_length (float): the model hyperparameter - nquad (int): the model hyperparameter - hpc_operation_count (int): parameter to indicate how the model is run - singularity_hpc (bool): parameter to indicate how the model is run - qoi_idx (int): Index of the qoi whose mean we want to estimate - """ - - # np.ndarray (nvars, nsamples) -> np.ndarary (nsamples, nqoi) - def eval_model(samples): - size = samples.shape - extras = np.zeros((4, size[1])) - extras[2, :] = cell_length - extras[3, :] = nquad - design_params = np.concatenate((samples, extras), axis=0) - - if hpc_operation: - print("==== Execute HPC version ====") - directory = "./benchmarks/hohlraum/slurm_scripts/" - user = read_username_from_config("./slurm_config.txt") - - delete_slurm_scripts(directory) # delete existing slurm files for hohlraum - call_models( - design_params, hpc_operation_count=1, singularity_hpc=singularity_hpc - ) - wait_for_slurm_jobs(user=user, sleep_interval=10) - if user: - print("Executing slurm scripts with user " + user) - execute_slurm_scripts(directory, user) - wait_for_slurm_jobs(user=user, sleep_interval=10) - [ - os.remove(file) - for file in glob.glob("./benchmarks/hohlraum/result/*.vtk") - ] - [ - os.remove(file) - for file in glob.glob("./benchmarks/hohlraum/mesh/*.su2") - ] - - else: - print("Username could not be read from slurm config file.") - exit(1) - time.sleep(10) - - qois = call_models(design_params, hpc_operation_count=2) - else: - qois = call_models(design_params, hpc_operation_count=0) - - qois = np.array(qois) - # wall_times = qois[0] - qoi_result = qois[:, qoi_idx].reshape(-1, 1) - return qoi_result - - def call_models(design_params, hpc_operation_count, singularity_hpc=True): - qois = [] - for column in design_params.T: - input = column.tolist() - input.append(hpc_operation_count) - input.append(singularity_hpc) - res = model([input]) - qois.append(res[0]) - - return np.array(qois) - - return eval_model - - -def main(): - - args = parse_args() - print(f"HPC mode = { not args.no_hpc}") - print(f"Load from npz = {args.load_from_npz}") - print(f"HPC with singularity = { not args.no_singularity_hpc}") - - hpc_operation = False # not args.no_hpc # Flag when using HPC cluster - singularity_hpc = False # not args.no_singularity_hpc - - # Define the model parameters that declare the fidelity - # Each element is one fidelity - # Highest fidelity comes first - cell_lengths = np.linspace(2.5e-3, 1e-2, 9) - # nquads = [4* i + 10 for i in range( len(cell_lengths))]# 20 * np.ones(len(cell_lengths)) - # [4* i + 10 for i in range( len(cell_lengths))] # - nquads = sorted([4 * i + 10 for i in range(len(cell_lengths))], reverse=True) - - print(nquads) - print(cell_lengths) - print("======GO=======") - # Generate the models into an array - # Include the correct HPC parameters - models = np.array( - [ - generate_model( - cell_length=cl, - nquad=nq, - hpc_operation=hpc_operation, - singularity_hpc=singularity_hpc, - qoi_idx=1, - ) - for (cl, nq) in zip(cell_lengths, nquads) - ] - ) - - # Update the ranges of the variables to match what charm_kit needs - # marginals = [stats.uniform(0.1, 0.2)]*4 + [stats.uniform(0.025, 0.05)]*2 - marginals = [ # first entry: lower bound, second entry: length of intervall - stats.uniform(0.3, 0.2), # left red top - stats.uniform(-0.5, 0.2), # left red bottom - stats.uniform(0.3, 0.2), # right red top - stats.uniform(-0.5, 0.2), # right red bottom - stats.uniform(-0.625, 0.05), # left red width - stats.uniform(0.575, 0.05), # right red width - ] - variable = IndependentMarginalsVariable(marginals) - sample_variable = variable.rvs - - # Run each model 1 time to estimate the model cost in seconds - # In the future, this will be done inside of the AETC algorithm to progressively get better estimates - n_time_samples = 1 - costs = [] - for model in models: - s = time.time() - model(sample_variable(n_time_samples)) - e = time.time() - costs.append((e - s) / n_time_samples) - - print("Costs: ", [f"{costs[i]:.2E}" for i, _ in enumerate(models)]) - - # This is the total budget in seconds that the total model run time will not exceed - # Note that since the model costs are estimates, the algorithm may take longer than the given budget - # This budget does not include the overhead that the AETC takes to run - budget = 60 * 60 * 24 * 2 # 345600 - - # Construct the estimator - optim_options = {"method": "cvxpy", "solver": "CLARABEL"} - estimator = etc.AETCBLUE( - models, sample_variable, costs=costs, opt_options=optim_options - ) - - # Explore phase - lf_subsets = get_model_subsets(len(models) - 1, 4) - samples, values, result = estimator.explore(budget, lf_subsets) - result_dict = estimator._explore_result_to_dict(result) - np.save("./hohlraum_explore_values.npy", values) - np.save("./hohlraum_explore_result.npy", result_dict) - - # Exploit phase - samples_per_model, best_subset = estimator.get_exploit_samples(result) - - # Run the `best_subset` models here using the input samples: `samples_per_model` - values_per_model = [ - models[s](samples) for s, samples in zip(best_subset, samples_per_model) - ] - # np.save("./hohlraum_exploit_values.npy", values_per_model) - - # Create a dictionary to map each array with a unique name - array_dict = {f"array_{i}": array for i, array in enumerate(values_per_model)} - - # Save the dictionary of arrays to an npz file - np.savez("./hohlraum_exploit_values.npz", **array_dict) - - # Run the estimator - mean = estimator.find_exploit_mean(values_per_model, result) - - # Print or save the results - print("AETC Algorithm ran successfully") - print("Mean of qoi:", mean) - print("Result Dictionary:", result) - - -if __name__ == "__main__": - main() diff --git a/src/simulation_utils.py b/src/simulation_utils.py index 6b76c34..d6298fc 100644 --- a/src/simulation_utils.py +++ b/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, ] diff --git a/tests/test_simulation_utils.py b/tests/test_simulation_utils.py new file mode 100644 index 0000000..6b5b30f --- /dev/null +++ b/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 + ) From 96bfeea39eef397de3714944a828bac4255d7986 Mon Sep 17 00:00:00 2001 From: Steffen Date: Wed, 18 Feb 2026 16:47:04 -0500 Subject: [PATCH 2/5] streamline installer --- install_kitrt.sh | 44 +++++++- src/simulation_utils.py | 229 +++++++++++++++++++++++++++++++++++----- 2 files changed, 241 insertions(+), 32 deletions(-) diff --git a/install_kitrt.sh b/install_kitrt.sh index 832cc99..d16c85a 100755 --- a/install_kitrt.sh +++ b/install_kitrt.sh @@ -10,8 +10,29 @@ has_cuda_gpu() { nvidia-smi -L >/dev/null 2>&1 } -# clone KiT-RT -git clone "${KITRT_REPO_URL}" kitrt_code +has_rocm_gpu() { + if command -v rocm-smi >/dev/null 2>&1; then + rocm-smi --showid 2>/dev/null | grep -q "^GPU\\[" + return + fi + + if command -v rocminfo >/dev/null 2>&1; then + rocminfo 2>/dev/null | grep -qi "gfx" + return + fi + + return 1 +} + +# clone KiT-RT once and reuse the existing checkout on subsequent runs +if [ -d "kitrt_code/.git" ]; then + echo "Existing kitrt_code checkout detected. Skipping clone + checkout." +elif [ -d "kitrt_code" ]; then + echo "Directory kitrt_code exists but is not a git repository. Please remove or rename it first." >&2 + exit 1 +else + git clone "${KITRT_REPO_URL}" kitrt_code +fi # go to kitrt_code directory cd kitrt_code @@ -22,7 +43,11 @@ git submodule update --init --recursive # navigate to directory where the singularity scripts are located cd tools/singularity -chmod +x build_container.sh install_kitrt_singularity.sh install_kitrt_singularity_cuda.sh +chmod +x \ + build_container.sh \ + install_kitrt_singularity.sh \ + install_kitrt_singularity_cuda.sh \ + install_kitrt_singularity_rocm.sh # build CPU singularity container. This requires root privileges. echo "Building CPU singularity container (sudo required)." @@ -40,5 +65,18 @@ else echo "No CUDA GPU detected. Skipping CUDA singularity container and CUDA build." fi +# optionally build and compile ROCm KiT-RT if a ROCm GPU is present +if has_rocm_gpu; then + if [ -f "kit_rt_MPI_rocm72.def" ]; then + echo "ROCm GPU detected. Building ROCm singularity container and ROCm KiT-RT binary." + sudo ./build_container.sh rocm + singularity exec --rocm kit_rt_MPI_rocm72.sif ./install_kitrt_singularity_rocm.sh + else + echo "ROCm GPU detected, but kit_rt_MPI_rocm72.def was not found. Skipping ROCm build." + fi +else + echo "No ROCm GPU detected. Skipping ROCm singularity container and ROCm build." +fi + # go back to CharmKiT repo root cd ../../../ diff --git a/src/simulation_utils.py b/src/simulation_utils.py index d6298fc..a606c92 100644 --- a/src/simulation_utils.py +++ b/src/simulation_utils.py @@ -1,11 +1,13 @@ -import subprocess +import glob import os +import shutil +import subprocess import time from src.general_utils import get_user_job_count -def _get_cuda_visible_device_count(): - visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") +def _get_visible_device_count(env_var_name): + visible_devices = os.environ.get(env_var_name) if visible_devices is None: return None @@ -14,6 +16,18 @@ def _get_cuda_visible_device_count(): return len(parsed) +def _get_cuda_visible_device_count(): + return _get_visible_device_count("CUDA_VISIBLE_DEVICES") + + +def _get_rocm_visible_device_count(): + for env_var_name in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + visible_count = _get_visible_device_count(env_var_name) + if visible_count is not None: + return visible_count + return None + + def _query_nvidia_smi_gpu_count(): try: result = subprocess.run( @@ -31,22 +45,61 @@ def _query_nvidia_smi_gpu_count(): return len([line for line in (result.stdout or "").splitlines() if line.strip()]) -def _resolve_cuda_mpi_ranks(quiet=False): +def _query_rocm_smi_gpu_count(): + try: + result = subprocess.run( + ["rocm-smi", "--showid"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + except FileNotFoundError: + return 0 + + if result.returncode != 0: + return 0 + + gpu_ids = set() + for line in (result.stdout or "").splitlines(): + stripped = line.strip() + if not stripped.startswith("GPU["): + continue + + closing_bracket = stripped.find("]") + if closing_bracket <= 4: + continue + + gpu_id = stripped[4:closing_bracket] + if gpu_id.isdigit(): + gpu_ids.add(gpu_id) + + return len(gpu_ids) + + +def _resolve_gpu_mpi_ranks_override(): override = os.environ.get("KITRT_CUDA_MPI_RANKS") + if override is None: + return 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) + + +def _resolve_cuda_mpi_ranks(quiet=False): + override = _resolve_gpu_mpi_ranks_override() 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) + return override visible_count = _get_cuda_visible_device_count() if visible_count is not None: @@ -65,6 +118,72 @@ def _resolve_cuda_mpi_ranks(quiet=False): return "1" +def _resolve_rocm_mpi_ranks(quiet=False): + override = _resolve_gpu_mpi_ranks_override() + if override is not None: + return override + + visible_count = _get_rocm_visible_device_count() + if visible_count is not None: + if visible_count >= 1: + return str(visible_count) + if not quiet: + print( + "HIP/ROCR visible-device mask is empty; falling back to 1 MPI rank." + ) + return "1" + + detected_gpu_count = _query_rocm_smi_gpu_count() + if detected_gpu_count >= 1: + return str(detected_gpu_count) + + if not quiet: + print("Could not detect available ROCm GPUs; falling back to 1 MPI rank.") + return "1" + + +def _is_rocm_installed(): + return ( + shutil.which("rocm-smi") is not None + or shutil.which("rocminfo") is not None + or os.path.isdir("/opt/rocm") + ) + + +def _find_rocm_container_image(): + preferred_images = [ + "kitrt_code/tools/singularity/kit_rt_MPI_rocm72.sif", + "kitrt_code/tools/singularity/kit_rt_MPI_rocm.sif", + ] + for image_path in preferred_images: + if os.path.exists(image_path): + return image_path + + rocm_images = sorted( + glob.glob("kitrt_code/tools/singularity/kit_rt_MPI_rocm*.sif") + ) + if rocm_images: + return rocm_images[0] + return None + + +def _find_rocm_executable(): + preferred_executables = [ + "./kitrt_code/build_singularity_rocm72/KiT-RT", + "./kitrt_code/build_singularity_rocm/KiT-RT", + ] + for executable_path in preferred_executables: + if os.path.exists(executable_path): + return executable_path + + rocm_executables = sorted( + glob.glob("./kitrt_code/build_singularity_rocm*/KiT-RT") + ) + if rocm_executables: + return rocm_executables[0] + return None + + def _run_and_raise(command, mode_label, quiet=False): try: if quiet: @@ -125,18 +244,70 @@ 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, - ] + # Keep the existing public flag, but select CUDA or ROCm at runtime. + if _query_nvidia_smi_gpu_count() >= 1: + 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, + ] + elif _is_rocm_installed(): + rocm_image = _find_rocm_container_image() + if rocm_image is None: + raise RuntimeError( + "ROCm runtime detected, but no KiT-RT ROCm Singularity image was found " + "under kitrt_code/tools/singularity/ (expected kit_rt_MPI_rocm*.sif)." + ) + + rocm_executable = _find_rocm_executable() + if rocm_executable is None: + raise RuntimeError( + "ROCm runtime detected, but no KiT-RT ROCm executable was found " + "under kitrt_code/build_singularity_rocm*/KiT-RT." + ) + + if not quiet: + print( + "CUDA GPUs not detected; using ROCm KiT-RT container and executable." + ) + + mpi_ranks = _resolve_rocm_mpi_ranks(quiet=quiet) + singularity_command = [ + "singularity", + "exec", + "--rocm", + rocm_image, + "mpirun", + "-np", + mpi_ranks, + rocm_executable, + config_file, + ] + else: + if not quiet: + print( + "CUDA GPUs were not detected and ROCm runtime is not available; " + "trying CUDA container path." + ) + 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, + ] else: singularity_command = [ "singularity", From e47d7d8959543b70f66bfdec370fc0d4e20ec129 Mon Sep 17 00:00:00 2001 From: Steffen Date: Wed, 18 Feb 2026 17:41:09 -0500 Subject: [PATCH 3/5] compatibility with apptainer --- src/simulation_utils.py | 38 ++++++++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/src/simulation_utils.py b/src/simulation_utils.py index a606c92..b592217 100644 --- a/src/simulation_utils.py +++ b/src/simulation_utils.py @@ -184,6 +184,34 @@ def _find_rocm_executable(): return None +def _resolve_container_runtime(): + runtime_override = os.environ.get("KITRT_CONTAINER_RUNTIME") + if runtime_override: + runtime = runtime_override.strip().lower() + if runtime not in ("apptainer", "singularity"): + raise RuntimeError( + "Invalid KITRT_CONTAINER_RUNTIME value. Expected " + "'apptainer' or 'singularity', " + f"got: {runtime_override!r}" + ) + if shutil.which(runtime) is None: + raise RuntimeError( + f"KITRT_CONTAINER_RUNTIME is set to {runtime!r}, " + "but that executable was not found in PATH." + ) + return runtime + + if shutil.which("apptainer") is not None: + return "apptainer" + if shutil.which("singularity") is not None: + return "singularity" + + raise RuntimeError( + "Containerized KiT-RT run failed: neither 'apptainer' nor " + "'singularity' was found in PATH." + ) + + def _run_and_raise(command, mode_label, quiet=False): try: if quiet: @@ -242,13 +270,15 @@ def run_cpp_simulation(config_file, quiet=False): def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): + container_runtime = _resolve_container_runtime() + # Path to the C++ executable if use_cuda: # Keep the existing public flag, but select CUDA or ROCm at runtime. if _query_nvidia_smi_gpu_count() >= 1: mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ - "singularity", + container_runtime, "exec", "--nv", "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif", @@ -280,7 +310,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): mpi_ranks = _resolve_rocm_mpi_ranks(quiet=quiet) singularity_command = [ - "singularity", + container_runtime, "exec", "--rocm", rocm_image, @@ -298,7 +328,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): ) mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ - "singularity", + container_runtime, "exec", "--nv", "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif", @@ -310,7 +340,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): ] else: singularity_command = [ - "singularity", + container_runtime, "exec", "kitrt_code/tools/singularity/kit_rt.sif", "./kitrt_code/build_singularity/KiT-RT", From e554175d74d5dfa36c7e8e7d9afabb670d96261e Mon Sep 17 00:00:00 2001 From: Steffen Date: Wed, 18 Feb 2026 19:58:23 -0500 Subject: [PATCH 4/5] gmsh now always uses container --- src/config_utils.py | 135 +++++++++++++++++++++++++++------------- src/simulation_utils.py | 31 +++++---- 2 files changed, 111 insertions(+), 55 deletions(-) diff --git a/src/config_utils.py b/src/config_utils.py index 0df603d..83b3083 100644 --- a/src/config_utils.py +++ b/src/config_utils.py @@ -3,6 +3,49 @@ import shutil import subprocess from src.general_utils import replace_next_line +from src.simulation_utils import _resolve_container_runtime + + +def _run_gmsh(input_geo, output_mesh, mesh_format): + container_runtime = _resolve_container_runtime() + cpu_image = "kitrt_code/tools/singularity/kit_rt.sif" + if not os.path.exists(cpu_image): + raise RuntimeError( + "CPU KiT-RT container image not found: " + "kitrt_code/tools/singularity/kit_rt.sif" + ) + + gmsh_command = [ + container_runtime, + "exec", + cpu_image, + "gmsh", + input_geo, + "-2", + "-format", + mesh_format, + "-save_all", + "-o", + output_mesh, + ] + + result = subprocess.run( + gmsh_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True + ) + if result.returncode != 0: + err = (result.stderr or result.stdout or "").strip() + raise RuntimeError( + f"gmsh mesh generation failed in container with return code {result.returncode}. " + f"Command: {' '.join(gmsh_command)}. " + f"Stderr: {err or 'no output'}" + ) + + +def _normalize_numeric_label(value): + try: + return format(float(value), ".15g") + except (TypeError, ValueError): + return str(value) def read_config_file(config_file): @@ -107,13 +150,12 @@ def update_quarter_hohlraum_mesh_file(n_cell, filepath): if os.path.exists(filename_con): os.remove(filename_con) - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - os.system( - f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}" - ) - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + _run_gmsh(filename_geo_backup, filename_vtk, "vtk") + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) return f"quarter_hohlraum_p{n_cell}.su2" @@ -143,13 +185,12 @@ def update_sym_hohlraum_mesh_file(n_cell, filepath): os.remove(filename_con) print("saving mesh with n_cell = ", n_cell) - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - os.system( - f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}" - ) - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + _run_gmsh(filename_geo_backup, filename_vtk, "vtk") + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) return f"sym_hohlraum_n{n_cell}.su2" @@ -231,11 +272,12 @@ def update_var_hohlraum_mesh_file( # # os.remove(slurm_script_path) # # else: - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - # os.system(f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}") - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + # os.system(f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}") + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) return unique_name + ".su2" @@ -308,20 +350,22 @@ def update_var_quarter_hohlraum_mesh_file( # os.remove(slurm_script_path) # # else: - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - # os.system(f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}") - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + # os.system(f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}") + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) return unique_name + ".su2" def update_lattice_mesh_file(n_cell, filepath, rectangular_mesh=False): + n_cell_label = _normalize_numeric_label(n_cell) filename_geo = filepath + "lattice_triangular.geo" filename_geo_backup = filepath + "lattice_triangular_backup.geo" - filename_su2 = filepath + f"lattice_p{n_cell}.su2" - filename_con = filepath + f"lattice_p{n_cell}.con" + filename_su2 = filepath + f"lattice_p{n_cell_label}.su2" + filename_con = filepath + f"lattice_p{n_cell_label}.con" if not os.path.exists(filename_su2): shutil.copy(filename_geo, filename_geo_backup) @@ -332,30 +376,32 @@ def update_lattice_mesh_file(n_cell, filepath, rectangular_mesh=False): with open(filename_geo_backup, "w") as file: for line in lines: if line.startswith("cl_fine"): - line = f"cl_fine = {n_cell};\n" + line = f"cl_fine = {n_cell_label};\n" file.write(line) # Remove the .con file if os.path.exists(filename_con): os.remove(filename_con) - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) - return f"lattice_p{n_cell}.su2" + return f"lattice_p{n_cell_label}.su2" def update_half_lattice_mesh_file(n_cell, filepath, rectangular_mesh=False): + n_cell_label = _normalize_numeric_label(n_cell) if rectangular_mesh: filename_geo = filepath + "half_lattice_rectangular.geo" else: filename_geo = filepath + "half_lattice_homogeneous.geo" filename_geo_backup = filepath + "half_lattice_backup.geo" - filename_su2 = filepath + f"half_lattice_p{n_cell}.su2" - filename_vtk = filepath + f"half_lattice_p{n_cell}.vtk" - filename_con = filepath + f"half_lattice_p{n_cell}.con" + filename_su2 = filepath + f"half_lattice_p{n_cell_label}.su2" + filename_vtk = filepath + f"half_lattice_p{n_cell_label}.vtk" + filename_con = filepath + f"half_lattice_p{n_cell_label}.con" if not os.path.exists(filename_su2): shutil.copy(filename_geo, filename_geo_backup) @@ -366,22 +412,23 @@ def update_half_lattice_mesh_file(n_cell, filepath, rectangular_mesh=False): with open(filename_geo_backup, "w") as file: for line in lines: if line.startswith("cl_fine"): - line = f"cl_fine = {n_cell};\n" + line = f"cl_fine = {n_cell_label};\n" file.write(line) # Remove the .con file if os.path.exists(filename_con): os.remove(filename_con) - os.system( - f"gmsh {filename_geo_backup} -2 -format su2 -save_all -o {filename_su2}" - ) - # os.system( - # f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}" - # ) - os.remove(filename_geo_backup) + try: + _run_gmsh(filename_geo_backup, filename_su2, "su2") + # os.system( + # f"gmsh {filename_geo_backup} -2 -format vtk -save_all -o {filename_vtk}" + # ) + finally: + if os.path.exists(filename_geo_backup): + os.remove(filename_geo_backup) - return f"half_lattice_p{n_cell}.su2" + return f"half_lattice_p{n_cell_label}.su2" def write_slurm_file( diff --git a/src/simulation_utils.py b/src/simulation_utils.py index b592217..42f51af 100644 --- a/src/simulation_utils.py +++ b/src/simulation_utils.py @@ -28,6 +28,20 @@ def _get_rocm_visible_device_count(): return None +def _has_cuda_gpu(): + visible_count = _get_cuda_visible_device_count() + if visible_count is not None: + return visible_count >= 1 + return _query_nvidia_smi_gpu_count() >= 1 + + +def _has_rocm_gpu(): + visible_count = _get_rocm_visible_device_count() + if visible_count is not None: + return visible_count >= 1 + return _query_rocm_smi_gpu_count() >= 1 + + def _query_nvidia_smi_gpu_count(): try: result = subprocess.run( @@ -275,7 +289,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): # Path to the C++ executable if use_cuda: # Keep the existing public flag, but select CUDA or ROCm at runtime. - if _query_nvidia_smi_gpu_count() >= 1: + if _has_cuda_gpu(): mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ container_runtime, @@ -288,7 +302,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): "./kitrt_code/build_singularity_cuda/KiT-RT", config_file, ] - elif _is_rocm_installed(): + elif _is_rocm_installed() and _has_rocm_gpu(): rocm_image = _find_rocm_container_image() if rocm_image is None: raise RuntimeError( @@ -323,19 +337,14 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): else: if not quiet: print( - "CUDA GPUs were not detected and ROCm runtime is not available; " - "trying CUDA container path." + "CUDA GPUs were not detected and no ROCm GPU fallback is available; " + "running CPU KiT-RT container path." ) - mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ container_runtime, "exec", - "--nv", - "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif", - "mpirun", - "-np", - mpi_ranks, - "./kitrt_code/build_singularity_cuda/KiT-RT", + "kitrt_code/tools/singularity/kit_rt.sif", + "./kitrt_code/build_singularity/KiT-RT", config_file, ] else: From e401efa817929c1b9c0185f7e3b2f6615569a139 Mon Sep 17 00:00:00 2001 From: Steffen Date: Wed, 18 Feb 2026 20:10:36 -0500 Subject: [PATCH 5/5] update sim utils --- src/simulation_utils.py | 26 +++++++++----------------- 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/src/simulation_utils.py b/src/simulation_utils.py index 42f51af..a5b1f1c 100644 --- a/src/simulation_utils.py +++ b/src/simulation_utils.py @@ -21,7 +21,7 @@ def _get_cuda_visible_device_count(): def _get_rocm_visible_device_count(): - for env_var_name in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + for env_var_name in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"): visible_count = _get_visible_device_count(env_var_name) if visible_count is not None: return visible_count @@ -288,7 +288,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): # Path to the C++ executable if use_cuda: - # Keep the existing public flag, but select CUDA or ROCm at runtime. + # CUDA requested: use CUDA when available, otherwise try ROCm. if _has_cuda_gpu(): mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ @@ -302,7 +302,13 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): "./kitrt_code/build_singularity_cuda/KiT-RT", config_file, ] - elif _is_rocm_installed() and _has_rocm_gpu(): + else: + if not (_is_rocm_installed() and _has_rocm_gpu()): + raise RuntimeError( + "CUDA mode was requested, but no CUDA device was detected and no ROCm " + "device is available. Cannot continue." + ) + rocm_image = _find_rocm_container_image() if rocm_image is None: raise RuntimeError( @@ -326,7 +332,6 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): singularity_command = [ container_runtime, "exec", - "--rocm", rocm_image, "mpirun", "-np", @@ -334,19 +339,6 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): rocm_executable, config_file, ] - else: - if not quiet: - print( - "CUDA GPUs were not detected and no ROCm GPU fallback is available; " - "running CPU KiT-RT container path." - ) - singularity_command = [ - container_runtime, - "exec", - "kitrt_code/tools/singularity/kit_rt.sif", - "./kitrt_code/build_singularity/KiT-RT", - config_file, - ] else: singularity_command = [ container_runtime,