From e8aee37a26790935b0ebbaf384c2d18da97448a5 Mon Sep 17 00:00:00 2001 From: Steffen Date: Fri, 13 Feb 2026 16:29:31 -0500 Subject: [PATCH] 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 + )