diff --git a/.gitignore b/.gitignore index 09c4dc3..2b96cb7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ *.sh +!update_kitrt.sh ./kitrt_code/ # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index a301a38..ef4a7fe 100644 --- a/README.md +++ b/README.md @@ -1,42 +1,63 @@ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) [![GitHub Stars](https://img.shields.io/github/stars/KiT-RT/CharmKiT)](https://github.com/KiT-RT/CharmKiT/stargazers) -[![Tests](https://github.com/KiT-RT/charm_kit/actions/workflows/tests.yml/badge.svg)](https://github.com/KiT-RT/charm_kit/actions/workflows/tests.yml) +[![Tests](https://github.com/KiT-RT/CharmKiT/actions/workflows/tests.yml/badge.svg)](https://github.com/KiT-RT/CharmKiT/actions/workflows/tests.yml) -# charm_kit: A wrapper for KiT-RT to conduct simulation sweeps fast +# CharmKiT: A wrapper for KiT-RT to conduct simulation sweeps fast -charm_kit is a benchmarking suite for the CharmNet project, providing automated parameter studies and test case management for the [KiT-RT PDE simulator](https://kit-rt.readthedocs.io/en/develop/index.html). It enables reproducible runs of radiative transfer test cases such as the lattice and hohlraum setups, using Python scripts to manage parameter sweeps, configuration, and result collection. charm_kit supports both high-performance computing (HPC) and local (no-HPC) execution modes, leveraging Singularity containers for reproducibility. +CharmKiT is a benchmarking suite for the CharmNet project, providing automated parameter studies and test case management for the [KiT-RT PDE simulator](https://kit-rt.readthedocs.io/en/develop/index.html). It enables reproducible runs of radiative transfer test cases such as the lattice and hohlraum setups, using Python scripts to manage parameter sweeps, configuration, and result collection. CharmKiT supports both high-performance computing (HPC) and local (no-HPC) execution modes, leveraging Apptainer/Singularity containers for reproducibility. +## Start here + +CharmKiT is the Python workflow layer for KiT-RT. Use `kitrt_code` directly when you want to build or modify the C++ solver; use CharmKiT when you want reproducible lattice or hohlraum sweeps, CSV-driven runs, SLURM submission, or QOI collection. + +| User goal | Use this | First command | +|---|---|---| +| Run one deterministic solver config | `kitrt_code` | `./build_omp/KiT-RT ...` | +| Run a local lattice sweep | CharmKiT | `poetry run charm-kit run lattice --singularity` | +| Run a local hohlraum sweep | CharmKiT | `poetry run charm-kit run hohlraum --singularity` | +| Submit a CPU SLURM sweep | CharmKiT | `poetry run charm-kit submit lattice --singularity` | +| Submit a CUDA SLURM sweep | CharmKiT | `poetry run charm-kit submit lattice --cuda` | + ## Installation Preliminaries: -1. Install [Singularity](https://docs.sylabs.io/guides/latest/user-guide/quick_start.html) on your system. +1. Install [Apptainer](https://apptainer.org/) or [Singularity](https://docs.sylabs.io/guides/latest/user-guide/quick_start.html). Set `KITRT_CONTAINER_RUNTIME=apptainer` or `KITRT_CONTAINER_RUNTIME=singularity` if the runtime is not on `PATH` or you need to force one. -2. Clone the `charm_kit` Github repository: - ``` - git clone git@github.com:ScSteffen/charm_kit.git +2. Clone the CharmKiT repository: + + ```bash + git clone git@github.com:KiT-RT/CharmKiT.git + cd CharmKiT ``` 3. Install Poetry and create the project environment: - ``` + + ```bash python3 -m pip install --user poetry poetry install ``` -4. Install [KiT-RT](https://github.com/KiT-RT/kitrt_code) as a submodule using the provided installer. (Requires root for container build.) - ``` - bash install_KiT-RT.sh +4. Install [KiT-RT](https://github.com/KiT-RT/kitrt_code) into `./kitrt_code/` using the provided installer: + + ```bash + bash install_kitrt.sh ``` - The repository is installed into `./kitrt_code/`. - The installer always builds the CPU Singularity image + binary. If a CUDA GPU is detected (`nvidia-smi`), it also builds the CUDA Singularity image and `build_singularity_cuda` binary automatically. - If updating KiT-RT: + + The installer builds or reuses the CPU container and binary. If a CUDA GPU is detected, it also prepares the CUDA container and `build_singularity_cuda` binary. To avoid root builds on an HPC system, provide existing SIF files in `kitrt_code/tools/singularity/`, set `KITRT_CONTAINER_BUILD=skip`, or set `KITRT_CPU_IMAGE_URI` / `KITRT_CUDA_IMAGE_URI` to an Apptainer/Singularity-compatible pull URI. + + ```bash + KITRT_CONTAINER_RUNTIME=apptainer KITRT_CONTAINER_BUILD=skip bash install_kitrt.sh ``` - bash update_KiT-RT.sh + +5. Update an existing KiT-RT checkout and rebuild inside existing containers: + + ```bash + bash update_kitrt.sh ``` - If on a cluster without root, build the container locally and upload it to `charm_kit/kitrt_code/tools/singularity/`. ## Testing @@ -48,9 +69,9 @@ poetry run pytest -q ``` -## How charm_kit Works +## How CharmKiT Works -charm_kit automates the setup, execution, and result collection for radiative transfer test cases using the KiT-RT solver. The workflow is managed by Python scripts (e.g., `run_hohlraum.py`, `run_lattice.py`) that: +CharmKiT automates the setup, execution, and result collection for radiative transfer test cases using the KiT-RT solver. The workflow is managed by Python scripts (e.g., `run_hohlraum.py`, `run_lattice.py`) that: - Define parameter sweeps for each test case (e.g., mesh size, quadrature order, absorption/scattering coefficients). - Generate the necessary configuration files for KiT-RT. @@ -60,26 +81,23 @@ charm_kit automates the setup, execution, and result collection for radiative tr Scripts support both HPC (SLURM) and local (no-HPC) execution. -## Running charm_kit Scripts - -charm_kit provides test-case drivers: - -- `run_lattice.py` -- `run_hohlraum.py` +## Running CharmKiT -Use Poetry to run commands in the project environment. - -Each script has its own CLI parser. They share execution flags, but design-parameter flags are test-case specific. Print all flags with: +CharmKiT provides a `charm-kit` CLI plus backward-compatible root scripts. Use Poetry to run commands in the project environment. ```bash +poetry run charm-kit run lattice --help +poetry run charm-kit run hohlraum --help poetry run python run_lattice.py --help poetry run python run_hohlraum.py --help ``` +`charm-kit run ` runs locally by default. `charm-kit submit ` adds `--slurm` and generates/submits SLURM scripts. The lattice and hohlraum runners share execution flags, but design-parameter flags are test-case specific. + Execution and I/O flags: - `--slurm`: Submit jobs through SLURM. -- `--singularity`: Run KiT-RT through the CPU Singularity image. +- `--singularity`: Run KiT-RT through the CPU container image. - `--cuda`: Run KiT-RT through the CUDA Singularity image (`--nv` is added automatically). - `--csv CSV`: Read design parameters from CSV and write QOIs back to that CSV. - `--config CONFIG`: Path to a TOML hyperparameter file. @@ -136,7 +154,7 @@ Precedence for hyperparameters is: ``` Uses image/executable: - `kitrt_code/tools/singularity/kit_rt.sif` and `./kitrt_code/build_singularity/KiT-RT`. + `kitrt_code/tools/singularity/kit_rt.sif` and `./kitrt_code/build_singularity/KiT-RT`. Set `KITRT_CONTAINER_RUNTIME=apptainer` to run with Apptainer instead of Singularity. 3. **Local mode + Singularity + GPU** @@ -149,7 +167,7 @@ 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 ...`. + `$KITRT_CONTAINER_RUNTIME 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=`. @@ -163,7 +181,7 @@ Precedence for hyperparameters is: Generated SLURM scripts call: `srun ./kitrt_code/build/KiT-RT ...`. -5. **SLURM mode + Singularity (CPU)** +5. **SLURM mode + Singularity/Apptainer (CPU)** ```bash poetry run python run_lattice.py --slurm --singularity @@ -172,12 +190,17 @@ Precedence for hyperparameters is: ``` Generated SLURM scripts call: - `singularity exec kitrt_code/tools/singularity/kit_rt.sif ./kitrt_code/build_singularity/KiT-RT ...`. + `srun "$KITRT_CONTAINER_RUNTIME" exec kitrt_code/tools/singularity/kit_rt.sif ./kitrt_code/build_singularity/KiT-RT ...`. + +6. **SLURM mode + Singularity/Apptainer + GPU** -### Not supported + ```bash + poetry run charm-kit submit lattice --cuda + # or + poetry run charm-kit submit hohlraum --cuda + ``` -- `--slurm --cuda` is intentionally blocked. - GPU mode is currently supported only for local Singularity runs (no SLURM). + Generated SLURM scripts use `--nv`, the CUDA SIF at `kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif`, and `mpirun -np "$KITRT_CUDA_MPI_RANKS"`. If `KITRT_CUDA_MPI_RANKS` is unset, the script falls back to `SLURM_GPUS_ON_NODE` and then `1`. Update `slurm_template.sh` with the account, partition, and GPU directives required by your site. # Test Case Descriptions diff --git a/install_kitrt.sh b/install_kitrt.sh index 832cc99..32985d4 100755 --- a/install_kitrt.sh +++ b/install_kitrt.sh @@ -2,6 +2,22 @@ set -euo pipefail KITRT_REPO_URL="${KITRT_REPO_URL:-https://github.com/KiT-RT/kitrt_code.git}" +KITRT_CONTAINER_BUILD="${KITRT_CONTAINER_BUILD:-auto}" + +detect_container_runtime() { + if [ -n "${KITRT_CONTAINER_RUNTIME:-}" ]; then + echo "${KITRT_CONTAINER_RUNTIME}" + elif command -v apptainer >/dev/null 2>&1; then + echo "apptainer" + elif command -v singularity >/dev/null 2>&1; then + echo "singularity" + else + echo "ERROR: install Apptainer or Singularity, or set KITRT_CONTAINER_RUNTIME." >&2 + return 1 + fi +} + +CONTAINER_RUNTIME="$(detect_container_runtime)" has_cuda_gpu() { if ! command -v nvidia-smi >/dev/null 2>&1; then @@ -10,35 +26,62 @@ has_cuda_gpu() { nvidia-smi -L >/dev/null 2>&1 } -# clone KiT-RT -git clone "${KITRT_REPO_URL}" kitrt_code +if [ -d kitrt_code/.git ]; then + git -C kitrt_code remote set-url origin "${KITRT_REPO_URL}" +else + git clone "${KITRT_REPO_URL}" kitrt_code +fi -# go to kitrt_code directory cd kitrt_code - -# keep origin synchronized with configured upstream and load submodules git remote set-url origin "${KITRT_REPO_URL}" 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 -# build CPU singularity container. This requires root privileges. -echo "Building CPU singularity container (sudo required)." -sudo ./build_container.sh cpu +ensure_image() { + local mode="$1" + local image="$2" + local uri="${3:-}" + + if [ -f "${image}" ]; then + echo "Using existing ${image}." + return 0 + fi + + if [ -n "${uri}" ]; then + "${CONTAINER_RUNTIME}" pull "${image}" "${uri}" + return 0 + fi + + case "${KITRT_CONTAINER_BUILD}" in + auto|sudo) + echo "Building ${image} with sudo. Set KITRT_CONTAINER_BUILD=skip or fakeroot on rootless clusters." + sudo "${CONTAINER_RUNTIME}" build "${image}" "kit_rt${mode}.def" + ;; + fakeroot) + "${CONTAINER_RUNTIME}" build --fakeroot "${image}" "kit_rt${mode}.def" + ;; + skip) + echo "ERROR: ${image} is missing. Provide it, set a pull URI, or allow a container build." >&2 + return 1 + ;; + *) + echo "ERROR: KITRT_CONTAINER_BUILD must be auto, sudo, fakeroot, or skip." >&2 + return 1 + ;; + esac +} -# compile CPU KiT-RT within the singularity container -singularity exec kit_rt.sif ./install_kitrt_singularity.sh +ensure_image "" "kit_rt.sif" "${KITRT_CPU_IMAGE_URI:-}" +"${CONTAINER_RUNTIME}" exec kit_rt.sif ./install_kitrt_singularity.sh -# optionally build and compile CUDA KiT-RT if a CUDA GPU is present if has_cuda_gpu; then - echo "CUDA GPU detected. Building CUDA singularity container and CUDA KiT-RT binary." - sudo ./build_container.sh cuda - singularity exec --nv kit_rt_MPI_cuda.sif ./install_kitrt_singularity_cuda.sh + echo "CUDA GPU detected. Preparing CUDA container and KiT-RT binary." + ensure_image "_MPI_cuda" "kit_rt_MPI_cuda.sif" "${KITRT_CUDA_IMAGE_URI:-}" + "${CONTAINER_RUNTIME}" exec --nv kit_rt_MPI_cuda.sif ./install_kitrt_singularity_cuda.sh else - echo "No CUDA GPU detected. Skipping CUDA singularity container and CUDA build." + echo "No CUDA GPU detected. Skipping CUDA container and CUDA build." fi -# go back to CharmKiT repo root cd ../../../ diff --git a/pyproject.toml b/pyproject.toml index f737746..9cb00f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.1.0" description = "Benchmarking workflows for KiT-RT lattice and hohlraum test cases." authors = ["charm_kit Developers"] readme = "README.md" -package-mode = false +packages = [{ include = "src" }] [tool.poetry.dependencies] python = ">=3.10,<3.13" @@ -15,6 +15,9 @@ pandas = "*" [tool.poetry.group.dev.dependencies] pytest = "*" +[tool.poetry.scripts] +charm-kit = "src.cli:main" + [tool.pytest.ini_options] pythonpath = ["."] diff --git a/run_hohlraum.py b/run_hohlraum.py index edd2112..27f8d48 100644 --- a/run_hohlraum.py +++ b/run_hohlraum.py @@ -22,14 +22,9 @@ def main(): hpc_operation = args.use_slurm # Flag when using HPC cluster use_cuda = args.cuda - if use_cuda and hpc_operation: - raise SystemExit( - "ERROR: --cuda cannot be combined with --slurm. " - "GPU mode is supported only with Singularity in non-SLURM runs." - ) singularity_hpc = args.use_singularity or use_cuda if use_cuda and not args.use_singularity: - print("CUDA mode requested; enabling Singularity execution.") + print("CUDA mode requested; enabling containerized execution.") config_path = args.config or "benchmarks/hohlraum/hyperparams.toml" hyper = load_toml_hyperparameters(config_path) print(f"Hohlraum hyperparameter config = {config_path}") diff --git a/run_lattice.py b/run_lattice.py index ff4240f..e1facb5 100644 --- a/run_lattice.py +++ b/run_lattice.py @@ -24,14 +24,9 @@ def main(): hpc_operation = args.use_slurm # Flag when using HPC cluster use_cuda = args.cuda - if use_cuda and hpc_operation: - raise SystemExit( - "ERROR: --cuda cannot be combined with --slurm. " - "GPU mode is supported only with Singularity in non-SLURM runs." - ) singularity_hpc = args.use_singularity or use_cuda if use_cuda and not args.use_singularity: - print("CUDA mode requested; enabling Singularity execution.") + print("CUDA mode requested; enabling containerized execution.") # print(f"Use rectangular_mesh = {args.rectangular_mesh}") rectangular_mesh = False diff --git a/slurm_template.sh b/slurm_template.sh index 9f1dadb..a59c71c 100644 --- a/slurm_template.sh +++ b/slurm_template.sh @@ -13,4 +13,11 @@ #SBATCH -o ./slurm_logs/singularity_hohlraum%j.txt #SBATCH -e ./slurm_logs/singularity_hohlraum_err_%j.txt #SBATCH --mail-type=FAIL -#SBATCH --mail-user=mai@mail.com \ No newline at end of file +#SBATCH --mail-user=mail@example.com + +# Uncomment and tune these lines for CUDA jobs on GPU partitions. +##SBATCH --gres=gpu:1 +##SBATCH --gpus-per-node=1 + +### command below +echo "CharmKiT command placeholder" \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..6acaefe --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +"""CharmKiT benchmark workflow helpers.""" diff --git a/src/cli.py b/src/cli.py new file mode 100644 index 0000000..de76781 --- /dev/null +++ b/src/cli.py @@ -0,0 +1,54 @@ +import sys + + +USAGE = """usage: + charm-kit run lattice [runner options] + charm-kit run hohlraum [runner options] + charm-kit submit lattice [runner options] + charm-kit submit hohlraum [runner options] + +Use "charm-kit run lattice --help" or "charm-kit run hohlraum --help" for +case-specific options. +""" + + +def _dispatch(case_name, runner_args): + if case_name == "lattice": + import run_lattice + + sys.argv = ["run_lattice.py", *runner_args] + return run_lattice.main() + if case_name == "hohlraum": + import run_hohlraum + + sys.argv = ["run_hohlraum.py", *runner_args] + return run_hohlraum.main() + raise SystemExit(f"ERROR: unknown case {case_name!r}\n\n{USAGE}") + + +def main(argv=None): + argv = list(sys.argv[1:] if argv is None else argv) + if not argv or argv[0] in {"-h", "--help"}: + print(USAGE) + return 0 + + command = argv.pop(0) + if command not in {"run", "submit"}: + raise SystemExit(f"ERROR: unknown command {command!r}\n\n{USAGE}") + if not argv: + raise SystemExit(f"ERROR: missing benchmark case\n\n{USAGE}") + + case_name = argv.pop(0) + runner_args = argv + if ( + command == "submit" + and "--slurm" not in runner_args + and "--use-slurm" not in runner_args + ): + runner_args = ["--slurm", *runner_args] + + return _dispatch(case_name, runner_args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/config_utils.py b/src/config_utils.py index 0df603d..f54b7dc 100644 --- a/src/config_utils.py +++ b/src/config_utils.py @@ -384,13 +384,59 @@ def update_half_lattice_mesh_file(n_cell, filepath, rectangular_mesh=False): return f"half_lattice_p{n_cell}.su2" +def _slurm_command_lines(unique_name, subfolder, singularity=True, use_cuda=False): + config_path = subfolder + unique_name + ".cfg" + if use_cuda: + return [ + 'CONTAINER_RUNTIME="${KITRT_CONTAINER_RUNTIME:-singularity}"\n', + 'CUDA_MPI_RANKS="${KITRT_CUDA_MPI_RANKS:-${SLURM_GPUS_ON_NODE:-1}}"\n', + 'srun "${CONTAINER_RUNTIME}" exec --nv ' + "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif " + 'mpirun -np "${CUDA_MPI_RANKS}" ' + "./kitrt_code/build_singularity_cuda/KiT-RT " + + config_path + + "\n", + ] + if singularity: + return [ + 'CONTAINER_RUNTIME="${KITRT_CONTAINER_RUNTIME:-singularity}"\n', + 'srun "${CONTAINER_RUNTIME}" exec ' + "kitrt_code/tools/singularity/kit_rt.sif " + "./kitrt_code/build_singularity/KiT-RT " + + config_path + + "\n", + ] + return ["srun ./kitrt_code/build/KiT-RT " + config_path + "\n"] + + +def _insert_slurm_command(lines, command_lines): + marker = "### command below" + updated = [] + inserted = False + skip_existing_command = False + + for line in lines: + if skip_existing_command: + skip_existing_command = False + continue + updated.append(line) + if line.strip() == marker: + updated.extend(command_lines) + inserted = True + skip_existing_command = True + + if not inserted: + if updated and updated[-1].strip(): + updated.append("\n") + updated.append(marker + "\n") + updated.extend(command_lines) + + return updated + + def write_slurm_file( output_slurm_dir, unique_name, subfolder, singularity=True, use_cuda=False ): - if use_cuda: - raise ValueError( - "CUDA mode with SLURM is not supported in this workflow." - ) basic_slurm_file = "./slurm_template.sh" # Ensure the output directory exists @@ -401,21 +447,16 @@ def write_slurm_file( with open(basic_slurm_file, "r") as file: lines = file.readlines() - # Replace the last line - if lines: - if singularity: - lines[-1] = ( - "singularity exec kitrt_code/tools/singularity/kit_rt.sif ./kitrt_code/build_singularity/KiT-RT " - + subfolder - + unique_name - + ".cfg\n" - ) - - else: - lines[-1] = "srun ./kitrt_code/build/KiT-RT " + subfolder + unique_name + ".cfg\n" + command_lines = _slurm_command_lines( + unique_name, + subfolder, + singularity=singularity, + use_cuda=use_cuda, + ) + lines = _insert_slurm_command(lines, command_lines) # Write the modified lines to the output file - with open(output_slurm_dir + unique_name + ".sh", "w") as file: + with open(os.path.join(output_slurm_dir, unique_name + ".sh"), "w") as file: file.writelines(lines) return 0 diff --git a/src/models/hohlraum.py b/src/models/hohlraum.py index d858123..cfb2791 100644 --- a/src/models/hohlraum.py +++ b/src/models/hohlraum.py @@ -35,10 +35,6 @@ def model(parameters): quiet = bool(parameters[0][13]) if len(parameters[0]) > 13 else False if use_cuda: singularity_hpc = 1 - if use_cuda and hpc_operation == 1: - raise ValueError( - "CUDA mode with SLURM is not supported in this workflow." - ) subfolder = "benchmarks/hohlraum/" base_config_file = subfolder + "hohlraum.cfg" diff --git a/src/models/lattice.py b/src/models/lattice.py index 14ab797..0c7869a 100644 --- a/src/models/lattice.py +++ b/src/models/lattice.py @@ -34,10 +34,6 @@ def model(parameters): quiet = bool(parameters[0][8]) if len(parameters[0]) > 8 else False if use_cuda: singularity_hpc = 1 - if use_cuda and hpc_operation == 1: - raise ValueError( - "CUDA mode with SLURM is not supported in this workflow." - ) subfolder = "benchmarks/lattice/" base_config_file = subfolder + "lattice.cfg" diff --git a/src/simulation_utils.py b/src/simulation_utils.py index d6298fc..1970883 100644 --- a/src/simulation_utils.py +++ b/src/simulation_utils.py @@ -65,6 +65,10 @@ def _resolve_cuda_mpi_ranks(quiet=False): return "1" +def _container_runtime(): + return os.environ.get("KITRT_CONTAINER_RUNTIME", "singularity") + + def _run_and_raise(command, mode_label, quiet=False): try: if quiet: @@ -127,7 +131,7 @@ def run_cpp_simulation_containerized(config_file, use_cuda=False, quiet=False): if use_cuda: mpi_ranks = _resolve_cuda_mpi_ranks(quiet=quiet) singularity_command = [ - "singularity", + _container_runtime(), "exec", "--nv", "kitrt_code/tools/singularity/kit_rt_MPI_cuda.sif", @@ -139,7 +143,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", diff --git a/tests/test_charm_kit_cli.py b/tests/test_charm_kit_cli.py new file mode 100644 index 0000000..a491caf --- /dev/null +++ b/tests/test_charm_kit_cli.py @@ -0,0 +1,19 @@ +from src import cli + + +def test_submit_injects_slurm_flag(monkeypatch): + calls = [] + monkeypatch.setattr(cli, "_dispatch", lambda case, args: calls.append((case, args)) or 0) + + assert cli.main(["submit", "lattice", "--cuda"]) == 0 + + assert calls == [("lattice", ["--slurm", "--cuda"])] + + +def test_run_preserves_runner_args(monkeypatch): + calls = [] + monkeypatch.setattr(cli, "_dispatch", lambda case, args: calls.append((case, args)) or 0) + + assert cli.main(["run", "hohlraum", "--config", "cfg.toml"]) == 0 + + assert calls == [("hohlraum", ["--config", "cfg.toml"])] diff --git a/tests/test_cli_and_runner_flags.py b/tests/test_cli_and_runner_flags.py index 52cfe4b..49bcc50 100644 --- a/tests/test_cli_and_runner_flags.py +++ b/tests/test_cli_and_runner_flags.py @@ -148,22 +148,60 @@ def test_parse_hohlraum_args_supports_current_flags(monkeypatch): assert not hasattr(args, "load_from_npz") -def test_lattice_main_rejects_cuda_with_slurm(monkeypatch): +def test_lattice_main_slurm_cuda_runs_stage_1_and_2(monkeypatch): monkeypatch.setattr( run_lattice, "parse_lattice_args", lambda: lattice_args(use_slurm=True, cuda=True) ) - with pytest.raises(SystemExit, match="--cuda cannot be combined with --slurm"): - run_lattice.main() + monkeypatch.setattr(run_lattice, "load_toml_hyperparameters", lambda _: {}) + monkeypatch.setattr( + run_lattice, + "create_lattice_samples_from_param_range", + lambda *_: (np.array([[10.0, 1.0, 0.01, 4]]), np.array(["a", "b", "c", "d"])), + ) + calls = [] + + def fake_call_models(*args, **kwargs): + calls.append((kwargs["hpc_operation_count"], kwargs["use_cuda"])) + return np.array([[1.0] * 8]) + + monkeypatch.setattr(run_lattice, "call_models", fake_call_models) + monkeypatch.setattr(run_lattice, "delete_slurm_scripts", lambda *_: None) + monkeypatch.setattr(run_lattice, "read_username_from_config", lambda *_: "alice") + monkeypatch.setattr(run_lattice, "execute_slurm_scripts", lambda *_: None) + monkeypatch.setattr(run_lattice, "wait_for_slurm_jobs", lambda **_: None) + monkeypatch.setattr(run_lattice.np, "savez", lambda *_, **__: None) + + assert run_lattice.main() == 0 + assert calls == [(1, True), (2, True)] -def test_hohlraum_main_rejects_cuda_with_slurm(monkeypatch): +def test_hohlraum_main_slurm_cuda_runs_stage_1_and_2(monkeypatch): monkeypatch.setattr( run_hohlraum, "parse_hohlraum_args", lambda: hohlraum_args(use_slurm=True, cuda=True), ) - with pytest.raises(SystemExit, match="--cuda cannot be combined with --slurm"): - run_hohlraum.main() + monkeypatch.setattr(run_hohlraum, "load_toml_hyperparameters", lambda _: {}) + monkeypatch.setattr( + run_hohlraum, + "create_hohlraum_samples_from_param_range", + lambda *_: (np.zeros((10, 1)), np.array(["p"] * 10)), + ) + calls = [] + + def fake_call_models(*args, **kwargs): + calls.append((kwargs["hpc_operation_count"], kwargs["use_cuda"])) + return np.array([[1.0] * len(run_hohlraum.get_qois_col_names())]) + + monkeypatch.setattr(run_hohlraum, "call_models", fake_call_models) + monkeypatch.setattr(run_hohlraum, "delete_slurm_scripts", lambda *_: None) + monkeypatch.setattr(run_hohlraum, "read_username_from_config", lambda *_: "alice") + monkeypatch.setattr(run_hohlraum, "execute_slurm_scripts", lambda *_: None) + monkeypatch.setattr(run_hohlraum, "wait_for_slurm_jobs", lambda **_: None) + monkeypatch.setattr(run_hohlraum.np, "savez", lambda *_, **__: None) + + assert run_hohlraum.main() == 0 + assert calls == [(1, True), (2, True)] def test_lattice_main_local_cuda_forces_containerized_mode(monkeypatch): diff --git a/tests/test_config_utils.py b/tests/test_config_utils.py new file mode 100644 index 0000000..d64a3ca --- /dev/null +++ b/tests/test_config_utils.py @@ -0,0 +1,37 @@ +from src.config_utils import write_slurm_file + + +def test_write_slurm_file_cuda_uses_gpu_container(tmp_path, monkeypatch): + template = tmp_path / "slurm_template.sh" + template.write_text( + "#SBATCH --mail-user=person@example.com\n" + "### command below\n" + "echo old command\n" + ) + monkeypatch.chdir(tmp_path) + + write_slurm_file( + "jobs", + "case_a", + "benchmarks/lattice/", + singularity=True, + use_cuda=True, + ) + + output = (tmp_path / "jobs" / "case_a.sh").read_text() + assert "#SBATCH --mail-user=person@example.com" in output + assert "echo old command" not in output + assert "kit_rt_MPI_cuda.sif" in output + assert "--nv" in output + assert "CUDA_MPI_RANKS" in output + assert "benchmarks/lattice/case_a.cfg" in output + + +def test_write_slurm_file_raw_uses_srun(tmp_path, monkeypatch): + (tmp_path / "slurm_template.sh").write_text("#SBATCH -J test\n") + monkeypatch.chdir(tmp_path) + + write_slurm_file("jobs", "case_b", "benchmarks/hohlraum/", singularity=False) + + output = (tmp_path / "jobs" / "case_b.sh").read_text() + assert "srun ./kitrt_code/build/KiT-RT benchmarks/hohlraum/case_b.cfg" in output diff --git a/tests/test_model_operations.py b/tests/test_model_operations.py index fe92976..9f45318 100644 --- a/tests/test_model_operations.py +++ b/tests/test_model_operations.py @@ -160,10 +160,12 @@ def test_lattice_model_postprocess_operation_reads_qois(lattice_mocks): assert lattice_mocks["container"] == [] -def test_lattice_model_rejects_cuda_with_slurm(): +def test_lattice_model_slurm_cuda_writes_script(lattice_mocks): params = [[10.0, 1.0, 0.01, 4, 1, 1, False, True, False]] - with pytest.raises(ValueError, match="CUDA mode with SLURM is not supported"): - lattice_model.model(params) + qois = lattice_model.model(params)[0] + assert qois == [0] * 8 + assert len(lattice_mocks["slurm"]) == 1 + assert lattice_mocks["slurm"][0][1]["use_cuda"] is True def test_hohlraum_model_local_raw_operation(hohlraum_mocks): @@ -202,7 +204,9 @@ def test_hohlraum_model_postprocess_operation_reads_qois(hohlraum_mocks): assert hohlraum_mocks["container"] == [] -def test_hohlraum_model_rejects_cuda_with_slurm(): +def test_hohlraum_model_slurm_cuda_writes_script(hohlraum_mocks): params = [[0.4, -0.4, 0.4, -0.4, -0.6, 0.6, 0.0, 0.0, 0.01, 6, 1, 1, True, False]] - with pytest.raises(ValueError, match="CUDA mode with SLURM is not supported"): - hohlraum_model.model(params) + qois = hohlraum_model.model(params)[0] + assert qois == [0] * 125 + assert len(hohlraum_mocks["slurm"]) == 1 + assert hohlraum_mocks["slurm"][0][1]["use_cuda"] is True