From 8902fdfbfa401ecdbf8cea19684a307ce6309580 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 11:02:43 +0800 Subject: [PATCH 01/14] [skip ci] ci: add minimal SAI GPU validation --- .ci/sai/README.md | 155 +++ .ci/sai/build.sbatch.in | 45 + .ci/sai/case.sbatch.in | 41 + .ci/sai/gpu-matrix.ini | 290 ++++++ .ci/sai/known_hosts | 1 + .ci/sai/modules.sh | 14 + .ci/sai/mpirun_with_mapping.sh | 7 + .ci/sai/sai.py | 895 ++++++++++++++++++ .ci/sai/slurm.py | 87 ++ .ci/sai/tests/test_sai.py | 333 +++++++ .github/workflows/sai-gpu.yml | 255 +++++ .../19_NO_Si48_CUSOLVERMP_TDDFT_GPU/INPUT | 37 + .../19_NO_Si48_CUSOLVERMP_TDDFT_GPU/KPT | 4 + .../19_NO_Si48_CUSOLVERMP_TDDFT_GPU/README | 8 + .../19_NO_Si48_CUSOLVERMP_TDDFT_GPU/STRU | 68 ++ 15 files changed, 2240 insertions(+) create mode 100644 .ci/sai/README.md create mode 100644 .ci/sai/build.sbatch.in create mode 100644 .ci/sai/case.sbatch.in create mode 100644 .ci/sai/gpu-matrix.ini create mode 100644 .ci/sai/known_hosts create mode 100644 .ci/sai/modules.sh create mode 100755 .ci/sai/mpirun_with_mapping.sh create mode 100644 .ci/sai/sai.py create mode 100644 .ci/sai/slurm.py create mode 100644 .ci/sai/tests/test_sai.py create mode 100644 .github/workflows/sai-gpu.yml create mode 100644 tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/INPUT create mode 100644 tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/KPT create mode 100644 tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/README create mode 100644 tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/STRU diff --git a/.ci/sai/README.md b/.ci/sai/README.md new file mode 100644 index 00000000000..2ea54b4c54a --- /dev/null +++ b/.ci/sai/README.md @@ -0,0 +1,155 @@ +# SAI GPU CI manual + +## Purpose and scope + +This workflow rebuilds ABACUS at one exact commit on the SAI GPU cluster and +runs the matrix in `gpu-matrix.ini` through Slurm. The current matrix has 49 +cases. It reports build, resource, and case states and saves the raw logs. + +This is validation, not a benchmark or a general remote-shell service. The +client requires a committed checkout: it bundles the checked-out `HEAD` and +checks that it equals `--source-sha`. Uncommitted files are not sent. + +**Warning:** approved candidate code is compiled and executed as the configured +SAI SSH user. Use an account and project root with the intended permissions. + +## GitHub setup + +Create these GitHub Environments for the `SAI GPU validation` workflow: + +- `sai-ssh-scheduled`: no required reviewers. The daily schedule uses this + environment. +- `sai-ssh-manual`: require the reviewers who should approve a remote run. + Manual dispatches and the pull-request comment trigger use this environment. + +In each environment, configure the SSH secret and connection variables: + +- Secret: `SAI_SSH_PRIVATE_KEY` (the private key used by the runner). +- Variables: `SAI_PROJECT_ROOT`, `SAI_SSH_HOST`, `SAI_SSH_PORT`, and + `SAI_SSH_USER`. + +`SAI_PROJECT_ROOT` should be an absolute path below the SAI user's home, for +example `/home///abacus_sai_gpu_ci`. A manual dispatch may provide +the optional `project_root` input to override this variable for that run. + +## Triggers + +- **Manual:** open Actions, choose `SAI GPU validation`, and select + **Run workflow**. `source_sha` is required and must be a 40-character commit + SHA. The workflow must be dispatched from the repository default branch. + `project_root` is optional. The `sai-ssh-manual` approval is requested before + the build job starts. +- **Daily:** the schedule is `30 20 * * *` (20:30 UTC). It tests the default + branch of `deepmodeling/abacus-develop` and uses `sai-ssh-scheduled`. +- **Pull request:** comment exactly `/abacus-ci sai-gpu` on an open pull + request. The commenter needs Triage, Write, Maintain, or Admin permission. + The pull request head commit is tested after `sai-ssh-manual` approval. + +Other issue comments are ignored. + +## Run locally + +Run from the repository checkout whose `HEAD` is the commit to test. The +following uses the committed `HEAD` SHA and the same required options as CI: + +```bash +SOURCE_SHA="$(git rev-parse HEAD)" +RUN_ID="$(date +%s)" +python3 .ci/sai/sai.py local \ + --ssh-config "$HOME/.ssh/sai-config" \ + --target sai-ci \ + --project-root /home///abacus_sai_gpu_ci \ + --source-repository "$PWD" \ + --source-sha "$SOURCE_SHA" \ + --namespace manual \ + --run-id "$RUN_ID" \ + --run-attempt 1 \ + --artifacts ./sai-artifacts +``` + +The SSH config must define the `sai-ci` target (host, port, user, key, and +known-host checking). `--source-repository` must be the checkout used to +compute `SOURCE_SHA`. The command waits for Slurm completion and exits with a +non-zero status when a case or infrastructure result is not fully successful. + +For all available subcommands and options, run: + +```bash +python3 .ci/sai/sai.py --help +python3 .ci/sai/sai.py local --help +``` + +## Matrix configuration + +`gpu-matrix.ini` is validated before jobs are submitted. + +- `[cluster]`: Slurm `partition`, absolute `mapping_root` for the MPI mapping + script, `disable_nccl_ib` (`true` or `false`), and `poll_seconds` (1-300). +- `[build]`: build-job `qos`, `nodes`, `tasks_per_node`, `gpus_per_node`, and + `time_seconds`. +- `[resource.NAME]`: the same allocation fields plus `parallelism`, the + maximum number of array tasks running at once. Each resource must have a + case. There is one rank per GPU and no resource may exceed 16 GPUs. +- `[case.NNN]`: contiguous, zero-padded sections with `suite`, `name`, + `resource`, and `runner` (`autotest` or `cusolvermp`). + +Resource component labels are generated, not configured separately. A +single-node resource is shown as `N GPU` or `N GPUs`; a multi-node resource is +shown as `N nodes / M GPUs`. Thus `gpu1`, `gpu2`, and `gpu4` display `1 GPU`, +`2 GPUs`, and `4 GPUs`; `gpu8x2` displays `2 nodes / 16 GPUs`. +`case.049` is `15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU`; it uses +`gpu8x2` and the `cusolvermp` runner. + +## Results and retention + +On SAI, a run is created below: + +``` +/runs//-/ +``` + +Its `results/` directory contains `result.json`, `summary.md`, build and case +logs, Slurm output, module/tool records, and status files. The coordinator and +working data are alongside it while the run is active. After results are +collected, CI archives `results/` and `jobs/` as: + +``` +/archives//-.tar.gz +``` + +The client removes archived files older than 72 hours when preparing a later +run, and removes the active run after archiving. On the GitHub runner, +`ARTIFACT_ROOT` is `${runner.temp}/sai-artifacts`; it contains the collected +`results/`, `run.json`, and `client.log`. CI uploads that directory as +`sai-gpu--` and retains it for 30 days. A pull-request comment +links to the Actions run and the uploaded raw files. + +## Troubleshooting + +**SSH fails.** Check `SAI_SSH_HOST`, `SAI_SSH_PORT`, and `SAI_SSH_USER`, that +the key matches the account, and that the target is reachable. CI uses the +committed `.ci/sai/known_hosts` with strict host-key checking. Test the same +target with the SSH config before retrying; do not disable host-key checking. + +**A module cannot be loaded.** `modules.sh` sources Lmod, purges modules, and +loads `cmake/3.31.6` and the configured ABACUS dependency module. Ask the SAI +administrator to provide or update those modules. Modules provide the +compiler, CUDA, MPI, and library dependencies; do not add library paths to CI +(`LD_LIBRARY_PATH`, `CPATH`, or `CMAKE_PREFIX_PATH`) or hard-code site paths. + +**CMake or linking fails.** Inspect `results/configure.log`, `build.log`, +`install.log`, `CMakeCache.txt`, `tools.txt`, and `ldd.txt`. The build uses Unix +Makefiles, CUDA architecture 70, CUDA MPI, cuSOLVERMP, cuBLASMP, and +NCCL parallel-device options. A missing runtime library causes the `ldd` check +to fail; fix the module environment rather than adding a CI path. + +**A job stays pending or times out.** Check the selected partition and QoS, +GPU availability, node and task limits, and the `time_seconds` value for that +resource. Slurm output is in `results/`; an allocation or queue delay is an +infrastructure issue, not a case failure. + +**MPI/PMIx initialization fails.** Autotest retries once after a recognized +`PMIX_ERR_FILE_OPEN_FAILURE` or `PMIX_ERR_OUT_OF_RESOURCE` initialization +failure. If it persists, inspect the attempt logs and the loaded MPI module, +Slurm allocation, and mapping file. The `cusolvermp` runner does not use this +retry. diff --git a/.ci/sai/build.sbatch.in b/.ci/sai/build.sbatch.in new file mode 100644 index 00000000000..b2dc143f486 --- /dev/null +++ b/.ci/sai/build.sbatch.in @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +#SBATCH --job-name=@JOB_NAME@ +#SBATCH --partition=@PARTITION@ +#SBATCH --qos=@QOS@ +#SBATCH --nodes=@NODES@ +#SBATCH --ntasks=@TASKS@ +#SBATCH --ntasks-per-node=@TASKS_PER_NODE@ +#SBATCH --gpus-per-node=@GPUS_PER_NODE@ +#SBATCH --time=@TIME@ +#SBATCH --output=@OUTPUT@ +#SBATCH --export=NIL + +set -euo pipefail + +export HOME=@HOME@ +export USER=${SLURM_JOB_USER:?} +export LOGNAME=$USER +export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + +source @CONTROL@/modules.sh +module -t list 2>&1 | tee @RESULTS@/modules.txt +command -v cmake make mpicxx nvcc | tee @RESULTS@/tools.txt +nvidia-smi dmon -s pucvmte -o T > @RESULTS@/dmon-${SLURM_JOB_ID}.log & + +cmake -S @SOURCE@ -B @BUILD@ -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=@INSTALL@ \ + -DCMAKE_CUDA_ARCHITECTURES=70 \ + -DENABLE_LIBXC=ON \ + -DUSE_CUDA=ON \ + -DUSE_CUDA_MPI=ON \ + -DENABLE_CUSOLVERMP=ON \ + -DENABLE_CUBLASMP=ON \ + -DENABLE_NCCL_PARALLEL_DEVICE=ON \ + -DBUILD_TESTING=OFF \ + -DGIT_SUBMODULE=OFF 2>&1 | tee @RESULTS@/configure.log +cmake --build @BUILD@ --parallel 32 2>&1 | tee @RESULTS@/build.log +cmake --install @BUILD@ 2>&1 | tee @RESULTS@/install.log + +@INSTALL@/bin/abacus --info | tee @RESULTS@/abacus-info.txt +ldd @INSTALL@/bin/abacus | tee @RESULTS@/ldd.txt +! grep -q 'not found' @RESULTS@/ldd.txt +cp @BUILD@/CMakeCache.txt @RESULTS@/CMakeCache.txt + +exit 0 diff --git a/.ci/sai/case.sbatch.in b/.ci/sai/case.sbatch.in new file mode 100644 index 00000000000..c9b0048eeb7 --- /dev/null +++ b/.ci/sai/case.sbatch.in @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +#SBATCH --job-name=@JOB_NAME@ +#SBATCH --partition=@PARTITION@ +#SBATCH --qos=@QOS@ +#SBATCH --nodes=@NODES@ +#SBATCH --ntasks=@TASKS@ +#SBATCH --ntasks-per-node=@TASKS_PER_NODE@ +#SBATCH --gpus-per-node=@GPUS_PER_NODE@ +#SBATCH --time=@TIME@ +#SBATCH --array=@ARRAY@ +#SBATCH --dependency=afterok:@BUILD_JOB@ +#SBATCH --output=@OUTPUT@ +#SBATCH --export=NIL + +set -euo pipefail + +export HOME=@HOME@ +export USER=${SLURM_JOB_USER:?} +export LOGNAME=$USER +export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +export OMP_NUM_THREADS=1 + +# shellcheck source=/dev/null +source @MAPPING_ROOT@/${SLURM_JOB_PARTITION}.bash +source @CONTROL@/modules.sh +SAI_SYSTEM_MPIRUN=$(command -v mpirun) +export SAI_SYSTEM_MPIRUN +export MAP_OPT +nvidia-smi dmon -s pucvmte -o T > @RESULTS@/dmon-${SLURM_JOB_ID}.log & + +SAI_DISABLE_NCCL_IB=@DISABLE_NCCL_IB@ +if [[ $SAI_DISABLE_NCCL_IB == true ]]; then + export NCCL_IB_DISABLE=1 +else + unset NCCL_IB_DISABLE +fi + +python3 @CONTROL@/sai.py worker \ + @SOURCE@ @CONTROL@ @INSTALL@ @RESULTS@ @MANIFEST@ + +exit $? diff --git a/.ci/sai/gpu-matrix.ini b/.ci/sai/gpu-matrix.ini new file mode 100644 index 00000000000..80d654f11b2 --- /dev/null +++ b/.ci/sai/gpu-matrix.ini @@ -0,0 +1,290 @@ +[cluster] +partition = 16V100 +mapping_root = /opt/sai_config/mps_mapping.d +disable_nccl_ib = false +poll_seconds = 10 + +[build] +qos = huge-gpu +nodes = 1 +tasks_per_node = 4 +gpus_per_node = 4 +time_seconds = 3600 + +[resource.gpu1] +qos = flood-1o2gpu +nodes = 1 +tasks_per_node = 1 +gpus_per_node = 1 +time_seconds = 900 +parallelism = 2 + +[resource.gpu2] +qos = flood-1o2gpu +nodes = 1 +tasks_per_node = 2 +gpus_per_node = 2 +time_seconds = 900 +parallelism = 16 + +[resource.gpu4] +qos = flood-gpu +nodes = 1 +tasks_per_node = 4 +gpus_per_node = 4 +time_seconds = 900 +parallelism = 16 + +[resource.gpu8x2] +qos = flood-gpu +nodes = 2 +tasks_per_node = 8 +gpus_per_node = 8 +time_seconds = 2400 +parallelism = 1 + +[case.001] +suite = 11_PW_GPU +name = scf_bpcg +resource = gpu2 +runner = autotest +[case.002] +suite = 11_PW_GPU +name = scf_cg +resource = gpu4 +runner = autotest +[case.003] +suite = 11_PW_GPU +name = scf_cg_single +resource = gpu4 +runner = autotest +[case.004] +suite = 11_PW_GPU +name = scf_dav +resource = gpu4 +runner = autotest +[case.005] +suite = 11_PW_GPU +name = scf_dav_sub +resource = gpu4 +runner = autotest +[case.006] +suite = 11_PW_GPU +name = scf_out_wf +resource = gpu1 +runner = autotest +[case.007] +suite = 11_PW_GPU +name = scf_out_wf_norm +resource = gpu4 +runner = autotest +[case.008] +suite = 12_NAO_Gamma_GPU +name = 001_NO_BiSeCuO_GPU +resource = gpu4 +runner = autotest +[case.009] +suite = 12_NAO_Gamma_GPU +name = 002_NO_H2O_GPU +resource = gpu4 +runner = autotest +[case.010] +suite = 12_NAO_Gamma_GPU +name = 003_NO_H2_DZP_GPU +resource = gpu4 +runner = autotest +[case.011] +suite = 12_NAO_Gamma_GPU +name = 004_NO_H2_DZP_S2_GPU +resource = gpu4 +runner = autotest +[case.012] +suite = 12_NAO_Gamma_GPU +name = 005_NO_H2_SZ_GPU +resource = gpu4 +runner = autotest +[case.013] +suite = 12_NAO_Gamma_GPU +name = 006_NO_H2_SZ_S2_GPU +resource = gpu4 +runner = autotest +[case.014] +suite = 12_NAO_Gamma_GPU +name = 007_NO_H_DZP_GPU +resource = gpu4 +runner = autotest +[case.015] +suite = 12_NAO_Gamma_GPU +name = 008_NO_H_DZP_S2_GPU +resource = gpu4 +runner = autotest +[case.016] +suite = 12_NAO_Gamma_GPU +name = 009_NO_Si2_DZP_GPU +resource = gpu4 +runner = autotest +[case.017] +suite = 12_NAO_Gamma_GPU +name = 010_NO_Si2_DZP_NEQ_GPU +resource = gpu4 +runner = autotest +[case.018] +suite = 12_NAO_Gamma_GPU +name = 011_NO_Si2_DZP_NEQ_S2_GPU +resource = gpu4 +runner = autotest +[case.019] +suite = 12_NAO_Gamma_GPU +name = 012_NO_Si2_DZP_S2_GPU +resource = gpu4 +runner = autotest +[case.020] +suite = 12_NAO_Gamma_GPU +name = 013_NO_Si2_TZDP_GPU +resource = gpu4 +runner = autotest +[case.021] +suite = 12_NAO_Gamma_GPU +name = 014_NO_Si2_TZDP_NEQ_GPU +resource = gpu4 +runner = autotest +[case.022] +suite = 12_NAO_Gamma_GPU +name = 015_NO_Si2_TZDP_NEQ_S2_GPU +resource = gpu4 +runner = autotest +[case.023] +suite = 12_NAO_Gamma_GPU +name = 016_NO_Si2_TZDP_S2_GPU +resource = gpu4 +runner = autotest +[case.024] +suite = 13_NAO_multik_GPU +name = 001_NO_KP_BiSeCuO_GPU +resource = gpu4 +runner = autotest +[case.025] +suite = 13_NAO_multik_GPU +name = 002_NO_KP_Si2_DZP_NEQ_S2_GPU +resource = gpu4 +runner = autotest +[case.026] +suite = 13_NAO_multik_GPU +name = 003_NO_KP_Si2_TZDP_S2_GPU +resource = gpu4 +runner = autotest +[case.027] +suite = 15_rtTDDFT_GPU +name = 01_NO_KP_ocp_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.028] +suite = 15_rtTDDFT_GPU +name = 02_NO_CH_OW_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.029] +suite = 15_rtTDDFT_GPU +name = 03_NO_CO_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.030] +suite = 15_rtTDDFT_GPU +name = 04_NO_CO_ocp_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.031] +suite = 15_rtTDDFT_GPU +name = 05_NO_cur_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.032] +suite = 15_rtTDDFT_GPU +name = 06_NO_dir_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.033] +suite = 15_rtTDDFT_GPU +name = 07_NO_EDM_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.034] +suite = 15_rtTDDFT_GPU +name = 09_NO_HEAV_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.035] +suite = 15_rtTDDFT_GPU +name = 10_NO_HHG_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.036] +suite = 15_rtTDDFT_GPU +name = 11_NO_O3_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.037] +suite = 15_rtTDDFT_GPU +name = 12_NO_re_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.038] +suite = 15_rtTDDFT_GPU +name = 14_NO_TRAP_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.039] +suite = 15_rtTDDFT_GPU +name = 15_NO_TRI_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.040] +suite = 15_rtTDDFT_GPU +name = 16_NO_vel_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.041] +suite = 15_rtTDDFT_GPU +name = 17_NO_vel_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.042] +suite = 15_rtTDDFT_GPU +name = 18_NO_hyb_TDDFT_GPU +resource = gpu4 +runner = autotest +[case.043] +suite = 16_SDFT_GPU +name = 001_PW_KG_100_GPU +resource = gpu2 +runner = autotest +[case.044] +suite = 16_SDFT_GPU +name = 002_PW_SKG_MALL_GPU +resource = gpu2 +runner = autotest +[case.045] +suite = 16_SDFT_GPU +name = 003_PW_MD_SDFT_ALL_GPU +resource = gpu2 +runner = autotest +[case.046] +suite = 16_SDFT_GPU +name = 004_PW_SDFT_ALL_GPU +resource = gpu2 +runner = autotest +[case.047] +suite = 16_SDFT_GPU +name = 005_PW_SDFT_MALL_BPCG_GPU +resource = gpu2 +runner = autotest +[case.048] +suite = 16_SDFT_GPU +name = 006_PW_SDFT_MALL_GPU +resource = gpu2 +runner = autotest +[case.049] +suite = 15_rtTDDFT_GPU +name = 19_NO_Si48_CUSOLVERMP_TDDFT_GPU +resource = gpu8x2 +runner = cusolvermp diff --git a/.ci/sai/known_hosts b/.ci/sai/known_hosts new file mode 100644 index 00000000000..52a041a2068 --- /dev/null +++ b/.ci/sai/known_hosts @@ -0,0 +1 @@ +[c0.sai.ai-4s.com]:12022 ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBBurYQExmMdgat4VU1Twmu2pqxsBjLP6eff9oc6/8N3KUCtoe8HU9eONghETLECdA7A+O9DpefDRba0AcSXTfQc= diff --git a/.ci/sai/modules.sh b/.ci/sai/modules.sh new file mode 100644 index 00000000000..03f70a0d7d2 --- /dev/null +++ b/.ci/sai/modules.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash + +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-} +export CPATH=${CPATH:-} +export CMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH:-} + +source /etc/profile.d/lmod.sh +module purge +module load cmake/3.31.6 +module load abacus/develop-git-079fd0c-260724-sm70-auto + +# CMake consumes the search paths exported by the loaded modules. +export CMAKE_LIBRARY_PATH=${LIBRARY_PATH:-} +export CMAKE_INCLUDE_PATH=${CPATH:-} diff --git a/.ci/sai/mpirun_with_mapping.sh b/.ci/sai/mpirun_with_mapping.sh new file mode 100755 index 00000000000..cc4b6cd9b4c --- /dev/null +++ b/.ci/sai/mpirun_with_mapping.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash + +set -euo pipefail +: "${SAI_SYSTEM_MPIRUN:?}" +: "${MAP_OPT:?}" + +exec "$SAI_SYSTEM_MPIRUN" --map-by "$MAP_OPT" "$@" diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py new file mode 100644 index 00000000000..ffe57e9299f --- /dev/null +++ b/.ci/sai/sai.py @@ -0,0 +1,895 @@ +#!/usr/bin/env python3 +"""Run the ABACUS GPU matrix on SAI from GitHub or a workstation.""" + +import argparse +import configparser +import fcntl +import json +import os +import re +import shlex +import shutil +import subprocess +import sys +import tarfile +import tempfile +import time +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from pathlib import PurePosixPath +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +from slurm import Slurm, SlurmError + + +ROOT = Path(__file__).resolve().parent +NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]*\Z") +SHA = re.compile(r"[0-9a-f]{40}\Z") +PMIX = re.compile(br"PMIX_ERR_(?:FILE_OPEN_FAILURE|OUT_OF_RESOURCE)") +MAX_RESULT_MEMBERS = 10000 +MAX_RESULT_BYTES = 2 * 1024**3 +MAX_REPORT_BYTES = 1024**2 + + +@dataclass(frozen=True) +class Resource: + name: str + qos: str + nodes: int + tasks_per_node: int + gpus_per_node: int + time_seconds: int + parallelism: int = 1 + + @property + def tasks(self) -> int: + return self.nodes * self.tasks_per_node + + @property + def label(self) -> str: + total = self.nodes * self.gpus_per_node + if self.nodes == 1: + return "{} GPU{}".format(total, "" if total == 1 else "s") + return "{} nodes / {} GPUs".format(self.nodes, total) + + +@dataclass(frozen=True) +class Case: + suite: str + name: str + resource: str + runner: str + + @property + def case_id(self) -> str: + return self.suite + "/" + self.name + + +@dataclass(frozen=True) +class Config: + partition: str + mapping_root: Path + disable_nccl_ib: bool + poll_seconds: int + build: Resource + resources: Mapping[str, Resource] + cases: Tuple[Case, ...] + + +def _integer(section: Mapping[str, str], key: str, low: int, high: int) -> int: + try: + value = int(section[key]) + except (KeyError, ValueError) as error: + raise ValueError("invalid {}".format(key)) from error + if not low <= value <= high: + raise ValueError("{} is outside {}..{}".format(key, low, high)) + return value + + +def _resource(name: str, section: Mapping[str, str], array: bool) -> Resource: + expected = { + "qos", "nodes", "tasks_per_node", "gpus_per_node", "time_seconds", + } | ({"parallelism"} if array else set()) + if set(section) != expected or not NAME.fullmatch(name): + raise ValueError("invalid resource {}".format(name)) + qos = section["qos"] + if not NAME.fullmatch(qos): + raise ValueError("invalid QoS") + profile = Resource( + name, qos, _integer(section, "nodes", 1, 2), + _integer(section, "tasks_per_node", 1, 8), + _integer(section, "gpus_per_node", 1, 8), + _integer(section, "time_seconds", 1, 3600), + _integer(section, "parallelism", 1, 32) if array else 1, + ) + if profile.tasks_per_node != profile.gpus_per_node or profile.tasks > 16: + raise ValueError("{} must use one rank per GPU and at most 16 GPUs".format(name)) + return profile + + +def load_config(path: Path = ROOT / "gpu-matrix.ini") -> Config: + parser = configparser.ConfigParser(interpolation=None, strict=True) + with Path(path).open(encoding="utf-8") as stream: + parser.read_file(stream) + resources = [name for name in parser.sections() if name.startswith("resource.")] + cases = [name for name in parser.sections() if name.startswith("case.")] + expected_cases = ["case.{:03d}".format(index) for index in range(1, len(cases) + 1)] + if not resources or cases != expected_cases: + raise ValueError("resources and contiguous case sections are required") + known = {"cluster", "build", *resources, *cases} + if set(parser.sections()) != known or set(parser["cluster"]) != { + "partition", "mapping_root", "disable_nccl_ib", "poll_seconds", + }: + raise ValueError("unexpected configuration section or key") + partition = parser["cluster"]["partition"] + mapping = Path(parser["cluster"]["mapping_root"]) + disable = parser["cluster"]["disable_nccl_ib"].lower() + if not NAME.fullmatch(partition) or not mapping.is_absolute() or disable not in ("true", "false"): + raise ValueError("invalid cluster configuration") + profiles: "OrderedDict[str, Resource]" = OrderedDict() + for section_name in resources: + name = section_name[len("resource."):] + profiles[name] = _resource(name, parser[section_name], True) + matrix: List[Case] = [] + for section_name in cases: + section = parser[section_name] + if set(section) != {"suite", "name", "resource", "runner"}: + raise ValueError("invalid {}".format(section_name)) + case = Case(section["suite"], section["name"], section["resource"], section["runner"]) + if not all(NAME.fullmatch(value) for value in (case.suite, case.name, case.resource)): + raise ValueError("invalid case name") + if case.resource not in profiles or case.runner not in ("autotest", "cusolvermp"): + raise ValueError("invalid case resource or runner") + matrix.append(case) + if len({case.case_id for case in matrix}) != len(matrix): + raise ValueError("duplicate case") + if any(not any(case.resource == name for case in matrix) for name in profiles): + raise ValueError("every resource needs a case") + return Config( + partition, mapping, disable == "true", + _integer(parser["cluster"], "poll_seconds", 1, 300), + _resource("build", parser["build"], False), profiles, tuple(matrix), + ) + + +def _atomic(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(".{}.{}.tmp".format(path.name, os.getpid())) + text = value if isinstance(value, str) else json.dumps(value, indent=2, sort_keys=True) + "\n" + temporary.write_text(text, encoding="utf-8") + os.replace(str(temporary), str(path)) + + +def _time(seconds: int) -> str: + hours, remainder = divmod(seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return "{:02d}:{:02d}:{:02d}".format(hours, minutes, seconds) + + +def _render(template: Path, destination: Path, values: Mapping[str, Any]) -> None: + text = template.read_text(encoding="utf-8") + for name, value in values.items(): + text = text.replace("@{}@".format(name), str(value)) + if re.search(r"@[A-Z_]+@", text): + raise ValueError("unresolved Slurm template value") + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(text, encoding="utf-8") + destination.chmod(0o755) + + +def _job_values( + profile: Resource, config: Config, run: Path, label: str, output: Path, +) -> Dict[str, Any]: + return { + "JOB_NAME": "abacus-{}-{}".format(run.name, label), + "PARTITION": config.partition, "QOS": profile.qos, + "NODES": profile.nodes, "TASKS": profile.tasks, + "TASKS_PER_NODE": profile.tasks_per_node, + "GPUS_PER_NODE": profile.gpus_per_node, + "TIME": _time(profile.time_seconds), "OUTPUT": output, + "HOME": Path.home(), "CONTROL": run / "control", + "SOURCE": run / "source", "BUILD": run / "build", + "INSTALL": run / "install", "RESULTS": run / "results", + } + + +def _component(name: str, label: str, state: str, job: str = "", slurm: str = "", code: str = "") -> Dict[str, str]: + return { + "name": name, "label": label, "state": state, "job_id": job, + "slurm_state": slurm, "exit_code": code, + } + + +def _result_row(case: Case, state: str, **values: Any) -> Dict[str, Any]: + row = { + "case_id": case.case_id, "resource": case.resource, + "runner": case.runner, "state": state, "exit_code": None, + "slurm_state": "", "slurm_exit_code": "", "job_id": "", + "elapsed_seconds": 0, + } + row.update(values) + return row + + +def _save_result(run: Path, components: Sequence[Mapping[str, Any]], rows: Sequence[Mapping[str, Any]]) -> int: + passed = sum(row["state"] == "PASS" for row in rows) + failed = sum(row["state"] in ("FAIL", "TIMEOUT") for row in rows) + result = { + "protocol": 1, "total": len(rows), "passed": passed, "failed": failed, + "infrastructure": len(rows) - passed - failed, + "components": list(components), "cases": list(rows), + } + root = run / "results" + _atomic(root / "result.json", result) + _atomic(root / "summary.md", _result_markdown(result)) + return 0 if rows and passed == len(rows) else 1 + + +def _result_markdown(result: Mapping[str, Any]) -> str: + lines = [ + "# SAI GPU result", "", + "Passed: **{}**; failed: **{}**; infrastructure: **{}**".format( + result["passed"], result["failed"], result["infrastructure"] + ), "", "| Component | State | Slurm job |", "| --- | --- | --- |", + ] + lines.extend("| {} | {} | {} |".format(item["label"], item["state"], item["job_id"]) for item in result["components"]) + lines.extend(("", "| Case | Resource | State | Slurm job |", "| --- | --- | --- | --- |")) + lines.extend("| {} | {} | {} | {} |".format(row["case_id"], row["resource"], row["state"], row["job_id"]) for row in result["cases"]) + return "\n".join(lines) + "\n" + + +def remote_run(run: Path) -> int: + run = run.resolve() + config = load_config(run / "control" / "gpu-matrix.ini") + results = run / "results" + scripts = run / "jobs" + results.mkdir(parents=True, exist_ok=True) + slurm = Slurm(config.poll_seconds) + components: List[Mapping[str, Any]] = [] + rows: List[Mapping[str, Any]] = [] + returncode = 2 + try: + build_script = scripts / "build.sbatch" + _render( + run / "control" / "build.sbatch.in", build_script, + _job_values(config.build, config, run, "build", results / "build-%j.out"), + ) + build_job = slurm.submit(build_script) + build_state, build_exit = slurm.wait((build_job,))[build_job] + build_ok = (build_state, build_exit) == ("COMPLETED", "0:0") + components.append(_component( + "build", "Compile", "PASS" if build_ok else "FAIL", + build_job, build_state, build_exit, + )) + if not build_ok: + components.extend(_component(name, profile.label, "SKIPPED") for name, profile in config.resources.items()) + rows = [ + _result_row( + case, "INFRA", job_id=build_job, slurm_state=build_state, + slurm_exit_code=build_exit, + ) + for resource in config.resources + for case in config.cases if case.resource == resource + ] + returncode = _save_result(run, components, rows) + return returncode + + jobs: Dict[str, str] = {} + grouped: Dict[str, List[Case]] = {} + for name, profile in config.resources.items(): + grouped[name] = [case for case in config.cases if case.resource == name] + manifest = scripts / (name + ".tsv") + manifest.write_text("\n".join("\t".join((case.case_id, case.suite, case.name, case.resource, case.runner)) for case in grouped[name]) + "\n", encoding="utf-8") + values = _job_values(profile, config, run, name, results / (name + "-%A_%a.out")) + values.update({ + "ARRAY": "0-{}%{}".format(len(grouped[name]) - 1, min(profile.parallelism, len(grouped[name]))), + "BUILD_JOB": build_job, "MAPPING_ROOT": config.mapping_root, + "DISABLE_NCCL_IB": str(config.disable_nccl_ib).lower(), + "MANIFEST": manifest, + }) + script = scripts / (name + ".sbatch") + _render(run / "control" / "case.sbatch.in", script, values) + jobs[name] = slurm.submit(script, len(grouped[name])) + accounting = slurm.wait(tuple(jobs.values())) + + for name, profile in config.resources.items(): + group_states = [] + for index, case in enumerate(grouped[name]): + job = "{}_{}".format(jobs[name], index) + slurm_state, slurm_exit = accounting[job] + status = results / "status" / (case.case_id.replace("/", "__") + ".json") + try: + data = json.loads(status.read_text(encoding="utf-8")) + state = data["state"] + if state not in ("PASS", "FAIL", "TIMEOUT", "INFRA"): + raise ValueError + except (OSError, ValueError, KeyError, json.JSONDecodeError): + data, state = {}, "INFRA" + if state == "PASS" and (slurm_state, slurm_exit) != ("COMPLETED", "0:0"): + state = "INFRA" + rows.append(_result_row( + case, state, exit_code=data.get("exit_code"), + elapsed_seconds=data.get("elapsed_seconds", 0), + slurm_state=slurm_state, slurm_exit_code=slurm_exit, + job_id=job, + )) + group_states.append(state) + state = "PASS" if all(item == "PASS" for item in group_states) else ( + "FAIL" if any(item in ("FAIL", "TIMEOUT") for item in group_states) else "INFRA" + ) + components.append(_component(name, profile.label, state, jobs[name])) + returncode = _save_result(run, components, rows) + return returncode + except Exception as error: + slurm.cancel() + _atomic(results / "coordinator-error.txt", str(error) + "\n") + raise + finally: + _atomic(results / "done.json", {"returncode": returncode}) + + +def _stream(command: Sequence[str], cwd: Path, log: Path) -> int: + with log.open("wb") as output: + process = subprocess.Popen(command, cwd=str(cwd), stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + assert process.stdout is not None + for block in iter(lambda: process.stdout.read(65536), b""): + output.write(block) + output.flush() + sys.stdout.buffer.write(block) + sys.stdout.buffer.flush() + return process.wait() + + +def _pmix(log: Path) -> bool: + data = log.read_bytes() + return bool(PMIX.search(data)) and b"MPI_Init_thread" in data and b"PMIx_Init failed" in data + + +def worker(source: Path, control: Path, install: Path, results: Path, manifest: Path) -> int: + task_id = int(os.environ["SLURM_ARRAY_TASK_ID"]) + fields = manifest.read_text(encoding="utf-8").splitlines()[task_id].split("\t") + if len(fields) != 5: + raise ValueError("invalid case manifest") + case_id, suite, name, resource, runner = fields + artifacts = results / "cases" / case_id.replace("/", "__") + artifacts.mkdir(parents=True, exist_ok=True) + source = source.resolve() + work = results.parent / "work" / case_id.replace("/", "__") + tests = work / "tests" + case = tests / suite / name + shutil.rmtree(str(work), ignore_errors=True) + (tests / suite).mkdir(parents=True) + os.symlink(str(source / "tests" / "integrate"), str(tests / "integrate")) + os.symlink(str(source / "tests" / "PP_ORB"), str(tests / "PP_ORB")) + shutil.copytree(str(source / "tests" / suite / name), str(case)) + launcher = work / "launcher" + launcher.mkdir() + os.symlink(str(control / "mpirun_with_mapping.sh"), str(launcher / "mpirun")) + os.environ["PATH"] = str(launcher) + os.pathsep + os.environ["PATH"] + abacus = install / "bin" / "abacus" + listing = subprocess.run(("ldd", str(abacus)), text=True, capture_output=True) + (artifacts / "ldd.txt").write_text(listing.stdout + listing.stderr, encoding="utf-8") + if listing.returncode or "not found" in listing.stdout: + raise RuntimeError("ABACUS has unresolved runtime libraries") + start = time.time() + returncode = 2 + final_pmix = False + try: + if runner == "autotest": + cases_file = case.parent / "CASES.task.txt" + cases_file.write_text(name + "\n", encoding="utf-8") + command = ( + "timeout", "--signal=TERM", "--kill-after=30s", "10m", + "bash", str(tests / "integrate" / "Autotest.sh"), + "-a", str(abacus), "-n", os.environ["SLURM_NTASKS"], + "-o", "1", "-f", cases_file.name, "-r", "^{}$".format(name), + ) + for attempt in (1, 2): + shutil.rmtree(str(case / "OUT.autotest"), ignore_errors=True) + for filename in ("log.txt", "result.out", "warning.log"): + try: + (case / filename).unlink() + except FileNotFoundError: + pass + log = artifacts / "attempt-{}.log".format(attempt) + returncode = _stream(command, case.parent, log) + final_pmix = returncode != 0 and _pmix(log) + if returncode == 0 or attempt == 2 or not final_pmix: + break + time.sleep(10) + else: + returncode = _stream(( + "timeout", "--signal=TERM", "--kill-after=30s", "35m", + "mpirun", "-np", os.environ["SLURM_NTASKS"], str(abacus), + ), case, artifacts / "cusolvermp.log") + for filename in ("log.txt", "result.out", "warning.log"): + if (case / filename).is_file(): + shutil.copy2(str(case / filename), str(artifacts / filename)) + state = "PASS" if returncode == 0 else ( + "TIMEOUT" if returncode in (124, 137, 143) else + "INFRA" if final_pmix else "FAIL" + ) + _atomic(results / "status" / (case_id.replace("/", "__") + ".json"), { + "case_id": case_id, "resource": resource, "runner": runner, + "state": state, "exit_code": returncode, + "elapsed_seconds": int(time.time() - start), + }) + return returncode + except Exception: + _atomic(results / "status" / (case_id.replace("/", "__") + ".json"), { + "case_id": case_id, "resource": resource, "runner": runner, + "state": "INFRA", "exit_code": 2, + "elapsed_seconds": int(time.time() - start), + }) + raise + + +def _command(command: Sequence[str], cwd: Optional[Path] = None, stdout: Any = subprocess.PIPE) -> subprocess.CompletedProcess: + result = subprocess.run(command, cwd=str(cwd) if cwd else None, text=True, stdout=stdout, stderr=subprocess.PIPE) + if result.returncode: + raise RuntimeError(result.stderr.strip() or "command failed: {}".format(command[0])) + return result + + +def _retry(command: Sequence[str], cwd: Optional[Path] = None, stdout: Any = subprocess.PIPE) -> subprocess.CompletedProcess: + error: Optional[Exception] = None + for attempt in range(3): + try: + return _command(command, cwd, stdout) + except RuntimeError as caught: + error = caught + if attempt < 2: + time.sleep(5 * (attempt + 1)) + raise RuntimeError(str(error)) + + +def _cache(project: Path) -> Path: + return project / "cache" / "repository.git" + + +def remote_prepare(project: Path, run: Path) -> None: + project = project.resolve() + run = run.resolve() + if Path.home().resolve() not in project.parents or project not in run.parents: + raise ValueError("project and run must be below HOME") + run_paths = tuple(run / name for name in ("source", "results", "jobs", "input")) + if any(path.exists() for path in run_paths): + raise ValueError("run directory already exists") + for path in (*run_paths, project / "archives"): + path.mkdir(parents=True, exist_ok=True) + cache = _cache(project) + cache.parent.mkdir(parents=True, exist_ok=True) + with (cache.parent / "lock").open("w") as stream: + fcntl.flock(stream, fcntl.LOCK_EX) + if not cache.exists(): + _command(("git", "init", "--bare", str(cache))) + latest = cache.parent / "latest" + base = latest.read_text(encoding="utf-8").strip() if latest.is_file() else "" + if not SHA.fullmatch(base or ""): + base = "" + cleanup_archives(project) + print(json.dumps({"run_root": str(run), "cache_sha": base})) + + +def remote_receive(project: Path, run: Path, source_sha: str) -> None: + if not SHA.fullmatch(source_sha): + raise ValueError("invalid source SHA") + cache = _cache(project.resolve()) + bundle = run.resolve() / "input" / "source.bundle" + lock = cache.parent / "lock" + with lock.open("w") as stream: + fcntl.flock(stream, fcntl.LOCK_EX) + if bundle.is_file(): + heads = _command(("git", "bundle", "list-heads", str(bundle))).stdout.split() + if len(heads) != 2 or heads != [source_sha, "HEAD"]: + raise ValueError("bundle does not advertise the requested SHA") + _command(("git", "--git-dir", str(cache), "bundle", "verify", str(bundle))) + _command(( + "git", "--git-dir", str(cache), "fetch", str(bundle), + "HEAD:refs/sai/{}".format(source_sha), + )) + else: + _command(("git", "--git-dir", str(cache), "cat-file", "-e", source_sha + "^{commit}")) + source = run / "source" + shutil.rmtree(str(source)) + source.mkdir() + archive = subprocess.Popen(("git", "--git-dir", str(cache), "archive", source_sha), stdout=subprocess.PIPE) + assert archive.stdout is not None + extract = subprocess.run(("tar", "-xf", "-", "-C", str(source)), stdin=archive.stdout) + archive.stdout.close() + if archive.wait() or extract.returncode: + raise RuntimeError("unable to extract source tree") + _atomic(cache.parent / "latest", source_sha + "\n") + + +def collect(run: Path) -> None: + with tarfile.open(fileobj=sys.stdout.buffer, mode="w|gz") as archive: + archive.add(str(run / "results"), arcname="results") + + +def cleanup_archives(project: Path) -> int: + removed = 0 + root = project / "archives" + now = time.time() + if root.is_dir(): + for path in root.glob("*/*.tar.gz"): + if now - path.stat().st_mtime > 72 * 3600: + path.unlink() + removed += 1 + return removed + + +def archive_run(project: Path, run: Path) -> Path: + relative = run.resolve().relative_to((project / "runs").resolve()) + destination = project / "archives" / relative.parent / (relative.name + ".tar.gz") + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(".tmp") + with tarfile.open(str(temporary), "w:gz") as archive: + archive.add(str(run / "results"), arcname="results") + archive.add(str(run / "jobs"), arcname="jobs") + os.replace(str(temporary), str(destination)) + shutil.rmtree(str(run)) + return destination + + +def _ssh(config: Path, target: str, command: Sequence[str], check: bool = True) -> subprocess.CompletedProcess: + result = subprocess.run(("ssh", "-F", str(config), target, shlex.join(command)), text=True, capture_output=True) + if check and result.returncode: + raise RuntimeError(result.stderr.strip() or "SSH command failed") + return result + + +def _extract_results(stream: Any, destination: Path) -> None: + members = 0 + size = 0 + with tarfile.open(fileobj=stream, mode="r|gz") as archive: + for member in archive: + path = PurePosixPath(member.name) + if path.is_absolute() or not path.parts or path.parts[0] != "results" or \ + ".." in path.parts or not (member.isdir() or member.isfile()): + raise ValueError("unsafe result archive member: {}".format(member.name)) + members += 1 + size += member.size + if members > MAX_RESULT_MEMBERS or size > MAX_RESULT_BYTES: + raise ValueError("result archive is too large") + archive.extract(member, str(destination)) + + +def _bundle_revision(repository: Path, base: str, source_sha: str) -> Optional[str]: + if base == source_sha: + return None + if SHA.fullmatch(base) and subprocess.run( + ("git", "merge-base", "--is-ancestor", base, "HEAD"), cwd=str(repository), + ).returncode == 0: + return base + "..HEAD" + return "HEAD" + + +def _validate_result(result: Any) -> Mapping[str, Any]: + config = load_config() + if type(result) is not dict or set(result) != { + "protocol", "total", "passed", "failed", "infrastructure", + "components", "cases", + } or type(result["protocol"]) is not int or result["protocol"] != 1: + raise ValueError("invalid result protocol") + + count_names = ("passed", "failed", "infrastructure", "total") + counts = {name: result[name] for name in count_names} + if any(type(value) is not int or value < 0 for value in counts.values()): + raise ValueError("invalid result counts") + + components = result["components"] + expected_components = [("build", "Compile")] + [ + (name, profile.label) for name, profile in config.resources.items() + ] + if type(components) is not list or len(components) != len(expected_components): + raise ValueError("invalid result components") + for item, identity in zip(components, expected_components): + if type(item) is not dict or set(item) != { + "name", "label", "state", "job_id", "slurm_state", "exit_code", + } or any(type(item[name]) is not str for name in item) or \ + (item["name"], item["label"]) != identity: + raise ValueError("invalid result component") + + rows = result["cases"] + expected_cases = [ + case for resource in config.resources + for case in config.cases if case.resource == resource + ] + if type(rows) is not list or len(rows) != len(expected_cases): + raise ValueError("invalid result cases") + for row, case in zip(rows, expected_cases): + if type(row) is not dict or set(row) != { + "case_id", "resource", "runner", "state", "exit_code", + "slurm_state", "slurm_exit_code", "job_id", "elapsed_seconds", + } or any(type(row[name]) is not str for name in ( + "case_id", "resource", "runner", "state", "slurm_state", + "slurm_exit_code", "job_id", + )) or (type(row["exit_code"]) is not int and row["exit_code"] is not None) or \ + type(row["elapsed_seconds"]) is not int or row["elapsed_seconds"] < 0 or \ + (row["case_id"], row["resource"], row["runner"]) != ( + case.case_id, case.resource, case.runner, + ) or row["state"] not in ("PASS", "FAIL", "TIMEOUT", "INFRA"): + raise ValueError("invalid result case") + + passed = sum(row["state"] == "PASS" for row in rows) + failed = sum(row["state"] in ("FAIL", "TIMEOUT") for row in rows) + actual_counts = (passed, failed, len(rows) - passed - failed, len(rows)) + if tuple(counts[name] for name in count_names) != actual_counts: + raise ValueError("inconsistent result counts") + + build_state = components[0]["state"] + if build_state == "PASS": + for component, resource in zip(components[1:], config.resources): + states = [row["state"] for row in rows if row["resource"] == resource] + expected = "PASS" if all(state == "PASS" for state in states) else ( + "FAIL" if any(state in ("FAIL", "TIMEOUT") for state in states) else "INFRA" + ) + if component["state"] != expected: + raise ValueError("inconsistent component state") + elif build_state != "FAIL" or any( + component["state"] != "SKIPPED" for component in components[1:] + ) or any(row["state"] != "INFRA" for row in rows): + raise ValueError("invalid build result") + return result + + +def _read_result(path: Path) -> Mapping[str, Any]: + if path.stat().st_size > MAX_REPORT_BYTES: + raise ValueError("result report is too large") + return _validate_result(json.loads(path.read_text(encoding="utf-8"))) + + +def local_run(args: argparse.Namespace) -> int: + repository = args.source_repository.resolve() + project = Path(args.project_root) + if not project.is_absolute() or any(not NAME.fullmatch(part) for part in project.parts[1:]): + raise ValueError("project root must be a simple absolute path") + if not all(NAME.fullmatch(value) for value in (args.target, args.namespace, args.run_id, args.run_attempt)): + raise ValueError("invalid target or run name") + actual = _command(("git", "rev-parse", "HEAD"), repository).stdout.strip() + if actual != args.source_sha or not SHA.fullmatch(actual): + raise ValueError("source SHA must be the checked-out HEAD") + run = project / "runs" / args.namespace / "{}-{}".format(args.run_id, args.run_attempt) + remote_control = str(run / "control") + _ssh(args.ssh_config, args.target, ("mkdir", "-p", remote_control)) + _retry(( + "rsync", "-a", "--delete", "--info=stats2", "-e", + "ssh -F {}".format(args.ssh_config), str(ROOT) + "/", + "{}:{}/".format(args.target, remote_control), + )) + prepared = _ssh(args.ssh_config, args.target, ( + "python3", remote_control + "/sai.py", "prepare", args.project_root, str(run), + )) + data = json.loads(prepared.stdout) + base = data.get("cache_sha", "") + with tempfile.TemporaryDirectory() as directory: + bundle = Path(directory) / "source.bundle" + revision = _bundle_revision(repository, base, args.source_sha) + if revision: + _command(("git", "bundle", "create", str(bundle), revision), repository) + _retry(( + "rsync", "-a", "--partial", "--info=progress2", "-e", + "ssh -F {}".format(args.ssh_config), str(bundle), + "{}:{}/input/source.bundle".format(args.target, run), + ), stdout=None) + _ssh(args.ssh_config, args.target, ( + "python3", remote_control + "/sai.py", "receive", + args.project_root, str(run), args.source_sha, + )) + launch = "nohup python3 {} remote-run {} > {}/results/coordinator.log 2>&1 < /dev/null &".format( + shlex.quote(remote_control + "/sai.py"), shlex.quote(str(run)), shlex.quote(str(run)) + ) + _ssh(args.ssh_config, args.target, ("bash", "-lc", launch)) + failures = 0 + while True: + status = _ssh(args.ssh_config, args.target, ( + "bash", "-lc", "test -f {0}/results/done.json && cat {0}/results/done.json || true".format(shlex.quote(str(run))), + ), check=False) + if status.returncode: + failures += 1 + if failures == 10: + raise RuntimeError("lost contact with SAI") + elif status.stdout.strip(): + done = json.loads(status.stdout) + if type(done) is not dict or set(done) != {"returncode"} or \ + type(done["returncode"]) is not int or done["returncode"] not in (0, 1, 2): + raise ValueError("invalid remote completion status") + break + else: + failures = 0 + time.sleep(30) + args.artifacts.mkdir(parents=True, exist_ok=True) + process = subprocess.Popen(( + "ssh", "-F", str(args.ssh_config), args.target, + shlex.join(("python3", remote_control + "/sai.py", "collect", str(run))), + ), stdout=subprocess.PIPE) + assert process.stdout is not None + _extract_results(process.stdout, args.artifacts) + if process.wait(): + raise RuntimeError("unable to collect SAI results") + _atomic(args.artifacts / "run.json", { + "project_root": args.project_root, "run_root": str(run), + "source_sha": args.source_sha, + }) + result_path = args.artifacts / "results" / "result.json" + if result_path.is_file(): + result = _read_result(result_path) + expected_returncode = 0 if result["passed"] == result["total"] else 1 + else: + expected_returncode = 2 + if done["returncode"] != expected_returncode: + raise ValueError("completion status does not match result") + return done["returncode"] + + +def report(args: argparse.Namespace) -> int: + if not args.result.is_file(): + components = [{"name": "infrastructure", "label": "Infrastructure", "state": "INFRA"}] + values = {"available": "false", "passed": "", "failed": "", "infrastructure": "", "total": ""} + else: + result = _read_result(args.result) + components = [{key: item[key] for key in ("name", "label", "state")} for item in result["components"]] + counts = {name: result[name] for name in ("passed", "failed", "infrastructure", "total")} + values = {"available": "true", **{name: str(value) for name, value in counts.items()}} + if args.summary: + args.summary.write_text(_result_markdown(result), encoding="utf-8") + values["components"] = json.dumps(components, separators=(",", ":")) + with args.output.open("a", encoding="utf-8") as stream: + for name, value in values.items(): + stream.write("{}={}\n".format(name, value)) + return 0 + + +def _gh(path: str, method: str = "GET", fields: Optional[Mapping[str, str]] = None) -> Any: + command = ["gh", "api", "--method", method, path] + for name, value in (fields or {}).items(): + command.extend(("-f", "{}={}".format(name, value))) + return json.loads(_command(command).stdout) + + +def github_admit() -> int: + event = os.environ["GITHUB_EVENT_NAME"] + repository = os.environ["GITHUB_REPOSITORY"] + values = {"accepted": "true", "pr_number": "", "check_id": ""} + if event == "schedule": + upstream = os.environ.get("SAI_UPSTREAM_REPOSITORY", "deepmodeling/abacus-develop") + metadata = _gh("repos/{}".format(upstream)) + commit = _gh("repos/{}/commits/{}".format(upstream, metadata["default_branch"])) + values.update(source_repository=upstream, source_sha=commit["sha"], namespace="daily") + elif event == "workflow_dispatch": + values.update( + source_repository=repository, + source_sha=os.environ["MANUAL_SOURCE_SHA"], namespace="manual", + ) + else: + event_data = json.loads(Path(os.environ["GITHUB_EVENT_PATH"]).read_text(encoding="utf-8")) + user = event_data["comment"]["user"]["login"] + permission = _gh("repos/{}/collaborators/{}/permission".format(repository, user)) + if permission.get("permission") not in ("admin", "maintain", "write", "triage"): + raise ValueError("commenter needs Triage permission") + number = str(event_data["issue"]["number"]) + pull = _gh("repos/{}/pulls/{}".format(repository, number)) + if pull["state"] != "open": + raise ValueError("pull request is not open") + values.update( + source_repository=pull["head"]["repo"]["full_name"], + source_sha=pull["head"]["sha"], namespace="pr-" + number, + pr_number=number, + ) + check = _gh("repos/{}/check-runs".format(repository), "POST", { + "name": "SAI GPU validation", "head_sha": pull["head"]["sha"], + "status": "in_progress", + }) + values["check_id"] = str(check["id"]) + if not SHA.fullmatch(values["source_sha"]): + raise ValueError("invalid source SHA") + with Path(os.environ["GITHUB_OUTPUT"]).open("a", encoding="utf-8") as stream: + for name, value in values.items(): + stream.write("{}={}\n".format(name, value)) + return 0 + + +def github_finish() -> int: + repository = os.environ["GITHUB_REPOSITORY"] + result = os.environ["SAI_RESULT"] + conclusion = "success" if result == "success" else "failure" + check_id = os.environ["CHECK_ID"] + if check_id: + _gh("repos/{}/check-runs/{}".format(repository, check_id), "PATCH", { + "status": "completed", "conclusion": conclusion, + }) + pr = os.environ["PR_NUMBER"] + if pr: + body = ( + "## SAI GPU validation: {}\n\n" + "GPU cases: **{} passed, {} failed, {} infrastructure**.\n\n" + "[Open the Actions run]({}) | [Download raw test files]({})\n\n" + "Candidate: `{}`" + ).format( + result, os.environ.get("GPU_PASSED") or "?", os.environ.get("GPU_FAILED") or "?", + os.environ.get("GPU_INFRASTRUCTURE") or "?", os.environ["RUN_URL"], + os.environ.get("ARTIFACT_URL", os.environ["RUN_URL"]), os.environ["SOURCE_SHA"], + ) + _gh("repos/{}/issues/{}/comments".format(repository, pr), "POST", {"body": body}) + return 0 + + +def parser() -> argparse.ArgumentParser: + main = argparse.ArgumentParser(description=__doc__) + commands = main.add_subparsers(dest="command", required=True) + commands.add_parser("config", help="validate and summarize gpu-matrix.ini") + remote = commands.add_parser("remote-run") + remote.add_argument("run", type=Path) + worker_parser = commands.add_parser("worker") + for name in ("source", "control", "install", "results", "manifest"): + worker_parser.add_argument(name, type=Path) + prepare = commands.add_parser("prepare") + prepare.add_argument("project", type=Path) + prepare.add_argument("run", type=Path) + receive = commands.add_parser("receive") + receive.add_argument("project", type=Path) + receive.add_argument("run", type=Path) + receive.add_argument("source_sha") + collect_parser = commands.add_parser("collect") + collect_parser.add_argument("run", type=Path) + archive = commands.add_parser("archive") + archive.add_argument("project", type=Path) + archive.add_argument("run", type=Path) + local = commands.add_parser("local") + local.add_argument("--ssh-config", required=True, type=Path) + local.add_argument("--target", required=True) + local.add_argument("--project-root", required=True) + local.add_argument("--source-repository", required=True, type=Path) + local.add_argument("--source-sha", required=True) + local.add_argument("--namespace", required=True) + local.add_argument("--run-id", required=True) + local.add_argument("--run-attempt", required=True) + local.add_argument("--artifacts", required=True, type=Path) + report_parser = commands.add_parser("report") + report_parser.add_argument("--result", required=True, type=Path) + report_parser.add_argument("--output", required=True, type=Path) + report_parser.add_argument("--summary", type=Path) + commands.add_parser("github-admit") + commands.add_parser("github-finish") + return main + + +def main() -> int: + if sys.version_info < (3, 8): + raise RuntimeError("Python 3.8 or newer is required") + args = parser().parse_args() + if args.command == "config": + config = load_config() + print(json.dumps({"resources": list(config.resources), "cases": len(config.cases)})) + return 0 + if args.command == "remote-run": + return remote_run(args.run) + if args.command == "worker": + return worker(args.source, args.control, args.install, args.results, args.manifest) + if args.command == "prepare": + remote_prepare(args.project, args.run) + elif args.command == "receive": + remote_receive(args.project, args.run, args.source_sha) + elif args.command == "collect": + collect(args.run) + elif args.command == "archive": + print(archive_run(args.project, args.run)) + elif args.command == "local": + return local_run(args) + elif args.command == "report": + return report(args) + elif args.command == "github-admit": + return github_admit() + elif args.command == "github-finish": + return github_finish() + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except (OSError, RuntimeError, ValueError, SlurmError) as error: + print("sai-ci: {}".format(error), file=sys.stderr) + sys.exit(2) diff --git a/.ci/sai/slurm.py b/.ci/sai/slurm.py new file mode 100644 index 00000000000..b38d900e159 --- /dev/null +++ b/.ci/sai/slurm.py @@ -0,0 +1,87 @@ +"""Small Slurm adapter for the SAI GPU matrix.""" + +import re +import subprocess +import time +from pathlib import Path +from typing import Dict, Optional, Sequence, Tuple + + +Terminal = Tuple[str, str] +TERMINAL = { + "BOOT_FAIL", "CANCELLED", "COMPLETED", "DEADLINE", "FAILED", + "NODE_FAIL", "OUT_OF_MEMORY", "PREEMPTED", "REVOKED", "TIMEOUT", +} + + +class SlurmError(RuntimeError): + pass + + +class Slurm: + def __init__(self, poll_seconds: int = 10) -> None: + self.poll_seconds = poll_seconds + self.jobs: Dict[str, Optional[int]] = {} + + @staticmethod + def _run(command: Sequence[str]) -> str: + result = subprocess.run(command, text=True, capture_output=True) + if result.returncode: + raise SlurmError(result.stderr.strip() or result.stdout.strip()) + return result.stdout + + def submit(self, script: Path, array_count: Optional[int] = None) -> str: + output = self._run(("sbatch", "--parsable", str(script))).strip() + match = re.fullmatch(r"([0-9]+)(?:;[A-Za-z0-9_.-]+)?", output) + if not match: + raise SlurmError("invalid sbatch output: {!r}".format(output)) + job = match.group(1) + self.jobs[job] = array_count + return job + + def wait(self, jobs: Sequence[str]) -> Dict[str, Terminal]: + ids = ",".join(jobs) + failures = 0 + while True: + try: + if not self._run(("squeue", "--noheader", "--jobs=" + ids)).strip(): + break + failures = 0 + except SlurmError: + failures += 1 + if failures == 6: + raise + time.sleep(self.poll_seconds) + + for _ in range(30): + rows = self._accounting(ids) + required = [] + for job in jobs: + count = self.jobs[job] + required.extend( + [job] if count is None else + ["{}_{}".format(job, index) for index in range(count)] + ) + if all(job in rows for job in required): + return {job: rows[job] for job in required} + time.sleep(self.poll_seconds) + raise SlurmError("Slurm accounting did not become complete") + + def _accounting(self, jobs: str) -> Dict[str, Terminal]: + output = self._run(( + "sacct", "--noheader", "--allocations", "--parsable2", + "--jobs=" + jobs, "--format=JobID,State,ExitCode", + )) + rows: Dict[str, Terminal] = {} + for line in output.splitlines(): + fields = line.strip().split("|") + if len(fields) != 3 or "." in fields[0]: + continue + state = fields[1].split()[0].rstrip("+") if fields[1].strip() else "" + if re.fullmatch(r"[0-9]+(?:_[0-9]+)?", fields[0]) and state in TERMINAL: + rows[fields[0]] = state, fields[2] + return rows + + def cancel(self) -> None: + if self.jobs: + subprocess.run(("scancel", *self.jobs), check=False) diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py new file mode 100644 index 00000000000..9acc1ac719a --- /dev/null +++ b/.ci/sai/tests/test_sai.py @@ -0,0 +1,333 @@ +import argparse +import io +import importlib.util +import json +import subprocess +import sys +import tempfile +import tarfile +import unittest +from pathlib import Path +from unittest import mock + + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +SPEC = importlib.util.spec_from_file_location("sai", ROOT / "sai.py") +assert SPEC and SPEC.loader +sai = importlib.util.module_from_spec(SPEC) +sys.modules["sai"] = sai +SPEC.loader.exec_module(sai) +import slurm # noqa: E402 + + +def valid_result(): + config = sai.load_config() + components = [sai._component("build", "Compile", "PASS")] + components.extend(sai._component(name, profile.label, "PASS") for name, profile in config.resources.items()) + rows = [ + sai._result_row(case, "PASS", exit_code=0) + for resource in config.resources + for case in config.cases if case.resource == resource + ] + return { + "protocol": 1, "total": len(rows), "passed": len(rows), + "failed": 0, "infrastructure": 0, "components": components, + "cases": rows, + } + + +class ConfigTests(unittest.TestCase): + def test_current_matrix_is_loaded_from_ini(self): + config = sai.load_config() + self.assertEqual(len(config.cases), 49) + self.assertEqual(list(config.resources), ["gpu1", "gpu2", "gpu4", "gpu8x2"]) + self.assertEqual(config.resources["gpu4"].label, "4 GPUs") + self.assertEqual(config.resources["gpu8x2"].label, "2 nodes / 16 GPUs") + self.assertEqual(config.cases[-1].runner, "cusolvermp") + + def test_resource_names_are_not_hardcoded(self): + text = (ROOT / "gpu-matrix.ini").read_text(encoding="utf-8") + text = text.replace("resource.gpu1", "resource.single", 1) + text = text.replace("resource = gpu1", "resource = single", 1) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "matrix.ini" + path.write_text(text, encoding="utf-8") + self.assertIn("single", sai.load_config(path).resources) + + def test_unknown_resource_and_noncontiguous_case_fail(self): + original = (ROOT / "gpu-matrix.ini").read_text(encoding="utf-8") + for text in ( + original.replace("resource = gpu1", "resource = missing", 1), + original.replace("[case.002]", "[case.999]", 1), + ): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "matrix.ini" + path.write_text(text, encoding="utf-8") + with self.assertRaises(ValueError): + sai.load_config(path) + + +class TemplateTests(unittest.TestCase): + def test_resource_template_is_complete_and_has_no_cpu_request(self): + config = sai.load_config() + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) / "runs" / "manual" / "1-1" + destination = run / "jobs" / "gpu4.sbatch" + values = sai._job_values( + config.resources["gpu4"], config, run, "gpu4", + run / "results" / "gpu4-%A_%a.out", + ) + values.update({ + "ARRAY": "0-3%4", "BUILD_JOB": "123", + "MAPPING_ROOT": config.mapping_root, + "DISABLE_NCCL_IB": "false", "MANIFEST": run / "jobs" / "gpu4.tsv", + }) + sai._render(ROOT / "case.sbatch.in", destination, values) + text = destination.read_text(encoding="utf-8") + self.assertIn("#SBATCH --nodes=1", text) + self.assertIn("#SBATCH --array=0-3%4", text) + self.assertIn("#SBATCH --dependency=afterok:123", text) + self.assertNotRegex(text, r"cpus-per-task|--mem(?:ory)?|--nodelist") + self.assertNotRegex(text, r"@[A-Z_]+@") + + def test_modules_do_not_spell_dependency_paths(self): + text = (ROOT / "modules.sh").read_text(encoding="utf-8") + self.assertIn("module load abacus/", text) + self.assertIn("CMAKE_LIBRARY_PATH=${LIBRARY_PATH:-}", text) + self.assertIn("CMAKE_INCLUDE_PATH=${CPATH:-}", text) + self.assertNotRegex(text, r"CUSOLVERMP_PATH|CUBLASMP_PATH|NCCL_PATH|/lib/lib") + + +class SlurmTests(unittest.TestCase): + def test_submit_and_accounting_require_each_array_task(self): + responses = [ + mock.Mock(returncode=0, stdout="101\n", stderr=""), + mock.Mock(returncode=0, stdout="", stderr=""), + mock.Mock( + returncode=0, + stdout="101|COMPLETED|0:0\n101_0|COMPLETED|0:0\n101_1|FAILED|1:0\n", + stderr="", + ), + ] + with mock.patch("slurm.subprocess.run", side_effect=responses): + client = slurm.Slurm(poll_seconds=0) + job = client.submit(Path("case.sbatch"), array_count=2) + states = client.wait((job,)) + self.assertEqual(states["101_0"], ("COMPLETED", "0:0")) + self.assertEqual(states["101_1"], ("FAILED", "1:0")) + + def test_pass_requires_successful_slurm_accounting(self): + config = sai.Config( + "gpu", Path("/opt/sai_config/mps_mapping.d"), False, 1, + sai.Resource("build", "flood-gpu", 1, 1, 1, 60), + {"one": sai.Resource("one", "flood-gpu", 1, 1, 1, 60)}, + (sai.Case("suite", "case", "one", "autotest"),), + ) + client = mock.Mock() + client.submit.side_effect = ["100", "101"] + client.wait.side_effect = [ + {"100": ("COMPLETED", "0:0")}, + {"101_0": ("FAILED", "1:0")}, + ] + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) + (run / "jobs").mkdir() + status = run / "results" / "status" / "suite__case.json" + status.parent.mkdir(parents=True) + status.write_text(json.dumps({"state": "PASS", "exit_code": 0}), encoding="utf-8") + with mock.patch("sai.load_config", return_value=config), \ + mock.patch("sai.Slurm", return_value=client), \ + mock.patch("sai._render"): + self.assertEqual(sai.remote_run(run), 1) + row = json.loads((run / "results" / "result.json").read_text())["cases"][0] + self.assertEqual(row["state"], "INFRA") + self.assertEqual(row["slurm_exit_code"], "1:0") + + +class TransferTests(unittest.TestCase): + @staticmethod + def _git(root, *arguments): + return subprocess.run(("git", *arguments), cwd=str(root), check=True, text=True, capture_output=True) + + def test_full_then_incremental_bundle(self): + with tempfile.TemporaryDirectory() as directory: + home = Path(directory) + repository = home / "local" + repository.mkdir() + self._git(repository, "init") + self._git(repository, "config", "user.email", "ci@example.invalid") + self._git(repository, "config", "user.name", "CI") + (repository / "value.txt").write_text("one\n", encoding="utf-8") + self._git(repository, "add", ".") + self._git(repository, "commit", "-m", "one") + first = self._git(repository, "rev-parse", "HEAD").stdout.strip() + + project = home / "project" + run1 = project / "runs" / "manual" / "1-1" + (run1 / "control").mkdir(parents=True) + with mock.patch("sai.Path.home", return_value=home): + sai.remote_prepare(project, run1) + self._git(repository, "bundle", "create", str(run1 / "input" / "source.bundle"), "HEAD") + sai.remote_receive(project, run1, first) + self.assertEqual((run1 / "source" / "value.txt").read_text(), "one\n") + + (repository / "value.txt").write_text("two\n", encoding="utf-8") + self._git(repository, "commit", "-am", "two") + second = self._git(repository, "rev-parse", "HEAD").stdout.strip() + run2 = project / "runs" / "manual" / "2-1" + (run2 / "control").mkdir(parents=True) + with mock.patch("sai.Path.home", return_value=home): + sai.remote_prepare(project, run2) + self._git( + repository, "bundle", "create", str(run2 / "input" / "source.bundle"), + first + "..HEAD", + ) + sai.remote_receive(project, run2, second) + self.assertEqual((run2 / "source" / "value.txt").read_text(), "two\n") + + run3 = project / "runs" / "manual" / "3-1" + (run3 / "control").mkdir(parents=True) + with mock.patch("sai.Path.home", return_value=home): + sai.remote_prepare(project, run3) + self.assertIsNone(sai._bundle_revision(repository, second, second)) + sai.remote_receive(project, run3, second) + self.assertEqual((run3 / "source" / "value.txt").read_text(), "two\n") + + def test_prepare_rejects_reused_run(self): + with tempfile.TemporaryDirectory() as directory: + home = Path(directory) + project = home / "project" + run = project / "runs" / "manual" / "1-1" + (run / "control").mkdir(parents=True) + with mock.patch("sai.Path.home", return_value=home): + sai.remote_prepare(project, run) + with self.assertRaisesRegex(ValueError, "already exists"): + sai.remote_prepare(project, run) + + def test_transfer_retry_is_bounded(self): + completed = mock.Mock(returncode=0, stdout="done", stderr="") + failed = mock.Mock(returncode=1, stdout="", stderr="disconnected") + with mock.patch("sai.subprocess.run", side_effect=[failed, failed, completed]) as run, \ + mock.patch("sai.time.sleep"): + self.assertEqual(sai._retry(("rsync", "source", "target")).stdout, "done") + self.assertEqual(run.call_count, 3) + + def test_result_archive_rejects_traversal_and_links(self): + for name, link in (("../sai-ssh/id_ed25519", None), ("results/key", "../key")): + payload = io.BytesIO() + with tarfile.open(fileobj=payload, mode="w:gz") as archive: + member = tarfile.TarInfo(name) + if link is None: + member.size = 3 + archive.addfile(member, io.BytesIO(b"key")) + else: + member.type = tarfile.SYMTYPE + member.linkname = link + archive.addfile(member) + payload.seek(0) + with tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError): + sai._extract_results(payload, Path(directory)) + + def test_result_archive_extracts_regular_results(self): + payload = io.BytesIO() + with tarfile.open(fileobj=payload, mode="w:gz") as archive: + member = tarfile.TarInfo("results/result.json") + member.size = 3 + archive.addfile(member, io.BytesIO(b"{}\n")) + payload.seek(0) + with tempfile.TemporaryDirectory() as directory: + sai._extract_results(payload, Path(directory)) + self.assertEqual((Path(directory) / "results" / "result.json").read_text(), "{}\n") + + def test_result_archive_rejects_oversized_content(self): + payload = io.BytesIO() + with tarfile.open(fileobj=payload, mode="w:gz") as archive: + member = tarfile.TarInfo("results/large") + member.size = 3 + archive.addfile(member, io.BytesIO(b"abc")) + payload.seek(0) + with tempfile.TemporaryDirectory() as directory, \ + mock.patch("sai.MAX_RESULT_BYTES", 2), \ + self.assertRaisesRegex(ValueError, "too large"): + sai._extract_results(payload, Path(directory)) + + def test_result_archive_limits_directory_members(self): + payload = io.BytesIO() + with tarfile.open(fileobj=payload, mode="w:gz") as archive: + for name in ("results/one", "results/two"): + member = tarfile.TarInfo(name) + member.type = tarfile.DIRTYPE + archive.addfile(member) + payload.seek(0) + with tempfile.TemporaryDirectory() as directory, \ + mock.patch("sai.MAX_RESULT_MEMBERS", 1), \ + self.assertRaisesRegex(ValueError, "too large"): + sai._extract_results(payload, Path(directory)) + + +class ResultTests(unittest.TestCase): + def test_report_publishes_dynamic_components(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + result = valid_result() + path = root / "result.json" + path.write_text(json.dumps(result), encoding="utf-8") + args = argparse.Namespace(result=path, output=root / "output", summary=root / "summary") + sai.report(args) + output = args.output.read_text(encoding="utf-8") + self.assertIn("available=true", output) + self.assertIn('"name":"gpu8x2"', output) + self.assertTrue(args.summary.read_text().startswith("# SAI GPU result\n")) + + def test_report_rejects_untrusted_counts(self): + invalid = ( + {"passed": "x[$(printf ARITH_EXEC >&2)0]", "failed": 0, "infrastructure": 0, "total": 1}, + {"passed": True, "failed": 0, "infrastructure": 0, "total": 1}, + {"passed": -1, "failed": 1, "infrastructure": 0, "total": 0}, + {"passed": 1, "failed": 1, "infrastructure": 0, "total": 1}, + ) + for counts in invalid: + with self.subTest(counts=counts), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + path = root / "result.json" + result = valid_result() + result.update(counts) + path.write_text(json.dumps(result), encoding="utf-8") + args = argparse.Namespace(result=path, output=root / "output", summary=None) + with self.assertRaisesRegex(ValueError, "result counts"): + sai.report(args) + self.assertFalse(args.output.exists()) + + def test_report_rejects_wrong_matrix_identity(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + result = valid_result() + result["cases"][0]["case_id"] = "invented/case" + path = root / "result.json" + path.write_text(json.dumps(result), encoding="utf-8") + args = argparse.Namespace(result=path, output=root / "output", summary=None) + with self.assertRaisesRegex(ValueError, "result case"): + sai.report(args) + + def test_pmix_signature_requires_all_markers(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "log" + path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE MPI_Init_thread PMIx_Init failed") + self.assertTrue(sai._pmix(path)) + path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE") + self.assertFalse(sai._pmix(path)) + + +class PolicyTests(unittest.TestCase): + def test_workflow_uses_trusted_control_and_protected_environment(self): + text = (ROOT.parents[1] / ".github" / "workflows" / "sai-gpu.yml").read_text(encoding="utf-8") + self.assertIn("ref: ${{ env.CONTROL_SHA }}", text) + self.assertIn("sai-ssh-scheduled", text) + self.assertIn("sai-ssh-manual", text) + self.assertIn("/abacus-ci sai-gpu", text) + self.assertNotIn("cpus-per-task", text) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/sai-gpu.yml b/.github/workflows/sai-gpu.yml new file mode 100644 index 00000000000..10feb9a1561 --- /dev/null +++ b/.github/workflows/sai-gpu.yml @@ -0,0 +1,255 @@ +name: SAI GPU validation + +on: + issue_comment: + types: [created] + workflow_dispatch: + inputs: + source_sha: + description: "Approved commit to execute on SAI" + required: true + type: string + project_root: + description: "Optional SAI project root" + required: false + type: string + schedule: + - cron: "30 20 * * *" + +permissions: {} + +concurrency: + group: sai-gpu-${{ github.event_name == 'schedule' && 'daily' || github.run_id }} + cancel-in-progress: false + +jobs: + admit: + name: Authorize request + if: >- + github.event_name != 'issue_comment' || + (github.event.issue.pull_request && + github.event.comment.body == '/abacus-ci sai-gpu') + runs-on: ubuntu-24.04 + permissions: + checks: write + contents: read + pull-requests: read + outputs: + accepted: ${{ steps.request.outputs.accepted }} + check_id: ${{ steps.request.outputs.check_id }} + control_sha: ${{ steps.control.outputs.sha }} + namespace: ${{ steps.request.outputs.namespace }} + pr_number: ${{ steps.request.outputs.pr_number }} + source_repository: ${{ steps.request.outputs.source_repository }} + source_sha: ${{ steps.request.outputs.source_sha }} + steps: + - name: Require default branch for manual runs + if: github.event_name == 'workflow_dispatch' + env: + DEFAULT_BRANCH: ${{ github.event.repository.default_branch }} + run: test "$GITHUB_REF_NAME" = "$DEFAULT_BRANCH" + + - name: Checkout trusted control code + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + ref: ${{ github.event.repository.default_branch }} + path: control + persist-credentials: false + + - name: Pin control commit + id: control + run: echo "sha=$(git -C control rev-parse HEAD)" >> "$GITHUB_OUTPUT" + + - name: Resolve candidate + id: request + env: + GH_TOKEN: ${{ github.token }} + MANUAL_SOURCE_SHA: ${{ inputs.source_sha || '' }} + SAI_UPSTREAM_REPOSITORY: deepmodeling/abacus-develop + run: python3 control/.ci/sai/sai.py github-admit + + rebuild-and-test: + name: Build and run on SAI + needs: admit + if: needs.admit.outputs.accepted == 'true' + runs-on: ubuntu-24.04 + timeout-minutes: 240 + environment: + name: ${{ github.event_name == 'schedule' && 'sai-ssh-scheduled' || 'sai-ssh-manual' }} + permissions: + contents: read + outputs: + artifact_url: ${{ steps.upload.outputs.artifact-url }} + available: ${{ steps.summary.outputs.available }} + components: ${{ steps.summary.outputs.components }} + failed: ${{ steps.summary.outputs.failed }} + infrastructure: ${{ steps.summary.outputs.infrastructure }} + passed: ${{ steps.summary.outputs.passed }} + total: ${{ steps.summary.outputs.total }} + env: + CONTROL_SHA: ${{ needs.admit.outputs.control_sha }} + PROJECT_ROOT_INPUT: ${{ inputs.project_root || '' }} + RUN_NAMESPACE: ${{ needs.admit.outputs.namespace }} + SAI_PROJECT_ROOT: ${{ vars.SAI_PROJECT_ROOT }} + SAI_SSH_HOST: ${{ vars.SAI_SSH_HOST }} + SAI_SSH_PORT: ${{ vars.SAI_SSH_PORT }} + SAI_SSH_USER: ${{ vars.SAI_SSH_USER }} + SOURCE_REPOSITORY: ${{ needs.admit.outputs.source_repository }} + SOURCE_SHA: ${{ needs.admit.outputs.source_sha }} + steps: + - name: Checkout pinned control code + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + ref: ${{ env.CONTROL_SHA }} + path: control + persist-credentials: false + + - name: Checkout candidate source + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + fetch-depth: 0 + path: source + persist-credentials: false + + - name: Configure SSH + env: + SAI_SSH_PRIVATE_KEY: ${{ secrets.SAI_SSH_PRIVATE_KEY }} + run: | + set -euo pipefail + root="$RUNNER_TEMP/sai-ssh" + mkdir -m 700 "$root" + printf '%s\n' "$SAI_SSH_PRIVATE_KEY" > "$root/id_ed25519" + chmod 600 "$root/id_ed25519" + cp control/.ci/sai/known_hosts "$root/known_hosts" + chmod 600 "$root/known_hosts" + cat > "$root/config" </dev/null + echo "SAI_SSH_CONFIG=$root/config" >> "$GITHUB_ENV" + echo "ARTIFACT_ROOT=$RUNNER_TEMP/sai-artifacts" >> "$GITHUB_ENV" + + - name: Run shared SAI client + id: client + continue-on-error: true + run: | + set -euo pipefail + project_root=${PROJECT_ROOT_INPUT:-$SAI_PROJECT_ROOT} + : "${project_root:?Set SAI_PROJECT_ROOT in the GitHub Environment}" + mkdir -p "$ARTIFACT_ROOT" + set +e + python3 control/.ci/sai/sai.py local \ + --ssh-config "$SAI_SSH_CONFIG" \ + --target sai-ci \ + --project-root "$project_root" \ + --source-repository "$GITHUB_WORKSPACE/source" \ + --source-sha "$SOURCE_SHA" \ + --namespace "$RUN_NAMESPACE" \ + --run-id "$GITHUB_RUN_ID" \ + --run-attempt "$GITHUB_RUN_ATTEMPT" \ + --artifacts "$ARTIFACT_ROOT" \ + 2>&1 | tee "$ARTIFACT_ROOT/client.log" + rc=${PIPESTATUS[0]} + set -e + echo "exit_code=$rc" >> "$GITHUB_OUTPUT" + exit "$rc" + + - name: Publish result summary + id: summary + if: always() + run: | + python3 control/.ci/sai/sai.py report \ + --result "$ARTIFACT_ROOT/results/result.json" \ + --output "$GITHUB_OUTPUT" \ + --summary "$GITHUB_STEP_SUMMARY" + + - name: Upload raw results + id: upload + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: sai-gpu-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.ARTIFACT_ROOT }} + if-no-files-found: warn + retention-days: 30 + + - name: Archive uploaded run on SAI + if: always() && steps.upload.outcome == 'success' + run: | + set -euo pipefail + test -f "$ARTIFACT_ROOT/run.json" || exit 0 + project=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["project_root"])' "$ARTIFACT_ROOT/run.json") + run=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["run_root"])' "$ARTIFACT_ROOT/run.json") + ssh -F "$SAI_SSH_CONFIG" sai-ci \ + python3 "$run/control/sai.py" archive "$project" "$run" + + - name: Remove SSH credentials + if: always() + run: rm -rf "$RUNNER_TEMP/sai-ssh" + + - name: Validate result protocol + if: always() + env: + AVAILABLE: ${{ steps.summary.outputs.available }} + run: | + set -euo pipefail + test "$AVAILABLE" = true + + component-status: + name: SAI / ${{ matrix.component.label }} + needs: [admit, rebuild-and-test] + if: always() && needs.admit.outputs.accepted == 'true' + runs-on: ubuntu-24.04 + strategy: + fail-fast: false + matrix: + component: ${{ fromJSON(needs.rebuild-and-test.outputs.components || '[{"name":"infrastructure","label":"Infrastructure","state":"INFRA"}]') }} + steps: + - name: Report component state + env: + STATE: ${{ matrix.component.state }} + run: test "$STATE" = PASS + + report-pr: + name: Report result to pull request + needs: [admit, rebuild-and-test, component-status] + if: always() && needs.admit.outputs.accepted == 'true' && needs.admit.outputs.pr_number != '' + runs-on: ubuntu-24.04 + permissions: + checks: write + contents: read + issues: write + steps: + - name: Checkout pinned reporter + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + ref: ${{ needs.admit.outputs.control_sha }} + path: control + persist-credentials: false + + - name: Complete check and comment + env: + ARTIFACT_URL: ${{ needs.rebuild-and-test.outputs.artifact_url }} + CHECK_ID: ${{ needs.admit.outputs.check_id }} + GH_TOKEN: ${{ github.token }} + GPU_FAILED: ${{ needs.rebuild-and-test.outputs.failed }} + GPU_INFRASTRUCTURE: ${{ needs.rebuild-and-test.outputs.infrastructure }} + GPU_PASSED: ${{ needs.rebuild-and-test.outputs.passed }} + PR_NUMBER: ${{ needs.admit.outputs.pr_number }} + RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} + SAI_RESULT: ${{ needs.rebuild-and-test.result == 'success' && needs.component-status.result == 'success' && 'success' || 'failure' }} + SOURCE_SHA: ${{ needs.admit.outputs.source_sha }} + run: python3 control/.ci/sai/sai.py github-finish diff --git a/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/INPUT b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/INPUT new file mode 100644 index 00000000000..34984ea35c2 --- /dev/null +++ b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/INPUT @@ -0,0 +1,37 @@ +INPUT_PARAMETERS + +suffix si48_cusolvermp +pseudo_dir ../../PP_ORB +orbital_dir ../../PP_ORB + +calculation md +esolver_type tddft +md_type nve +md_nstep 1 +estep_per_md 2 +td_dt 0.005 +md_tfirst 0 + +td_vext 1 +td_vext_dire 3 +td_stype 2 +td_ttype 3 +td_tstart 1 +td_tend 2 +td_heavi_t0 1 +td_heavi_amp 0.05 +out_current 1 + +basis_type lcao +gamma_only 0 +ecutwfc 20 +scf_nmax 50 +scf_thr 1.0e-6 +device gpu +ks_solver cusolvermp + +mixing_type broyden +mixing_beta 0.3 +mixing_gg0 0.0 +cal_force 1 +cal_stress 0 diff --git a/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/KPT b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/KPT new file mode 100644 index 00000000000..c289c0158aa --- /dev/null +++ b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/README b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/README new file mode 100644 index 00000000000..c72b12fe649 --- /dev/null +++ b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/README @@ -0,0 +1,8 @@ +This Si48 solid smoke case is reconstructed from the repository's Si64 PEXSI +supercell and Si RT-TDDFT GPU input. It matches the 48-atom system size used to +validate the multi-GPU RT-TDDFT implementation in PR #7026, but it is not a +copy of that unpublished benchmark input. + +The case runs only two electronic propagation steps. It checks that the +multi-node cuSolverMp RT-TDDFT path completes; it is not a physical or +performance reference calculation. diff --git a/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/STRU b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/STRU new file mode 100644 index 00000000000..704123b22e5 --- /dev/null +++ b/tests/15_rtTDDFT_GPU/19_NO_Si48_CUSOLVERMP_TDDFT_GPU/STRU @@ -0,0 +1,68 @@ +ATOMIC_SPECIES +Si 28.085 Si_ONCV_PBE-1.0.upf + +NUMERICAL_ORBITAL +Si_gga_6au_100Ry_2s2p1d.orb + +LATTICE_CONSTANT +10.2 + +LATTICE_VECTORS +1.5 0.0 0.0 +0.0 2.0 0.0 +0.0 0.0 2.0 + +ATOMIC_POSITIONS +Cartesian + +Si +0.0 +48 +0.00 0.00 0.00 0 0 0 +0.25 0.25 0.25 0 0 0 +0.00 0.50 0.50 0 0 0 +0.25 0.75 0.75 0 0 0 +0.50 0.00 0.50 0 0 0 +0.75 0.25 0.75 0 0 0 +0.50 0.50 0.00 0 0 0 +0.75 0.75 0.25 0 0 0 +1.00 0.00 0.00 0 0 0 +1.25 0.25 0.25 0 0 0 +1.00 0.50 0.50 0 0 0 +1.25 0.75 0.75 0 0 0 +0.00 1.00 0.00 0 0 0 +0.25 1.25 0.25 0 0 0 +0.00 1.50 0.50 0 0 0 +0.25 1.75 0.75 0 0 0 +0.50 1.00 0.50 0 0 0 +0.75 1.25 0.75 0 0 0 +0.50 1.50 0.00 0 0 0 +0.75 1.75 0.25 0 0 0 +1.00 1.00 0.00 0 0 0 +1.25 1.25 0.25 0 0 0 +1.00 1.50 0.50 0 0 0 +1.25 1.75 0.75 0 0 0 +0.00 0.00 1.00 0 0 0 +0.25 0.25 1.25 0 0 0 +0.00 0.50 1.50 0 0 0 +0.25 0.75 1.75 0 0 0 +0.50 0.00 1.50 0 0 0 +0.75 0.25 1.75 0 0 0 +0.50 0.50 1.00 0 0 0 +0.75 0.75 1.25 0 0 0 +1.00 0.00 1.00 0 0 0 +1.25 0.25 1.25 0 0 0 +1.00 0.50 1.50 0 0 0 +1.25 0.75 1.75 0 0 0 +0.00 1.00 1.00 0 0 0 +0.25 1.25 1.25 0 0 0 +0.00 1.50 1.50 0 0 0 +0.25 1.75 1.75 0 0 0 +0.50 1.00 1.50 0 0 0 +0.75 1.25 1.75 0 0 0 +0.50 1.50 1.00 0 0 0 +0.75 1.75 1.25 0 0 0 +1.00 1.00 1.00 0 0 0 +1.25 1.25 1.25 0 0 0 +1.00 1.50 1.50 0 0 0 +1.25 1.75 1.75 0 0 0 From 12f8efc127436341ae80ab0e5389b347480791e5 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 11:42:40 +0800 Subject: [PATCH 02/14] [skip ci] ci: transfer source bundle in parallel --- .ci/sai/README.md | 5 +++ .ci/sai/sai.py | 81 ++++++++++++++++++++++++++++++++++----- .ci/sai/tests/test_sai.py | 31 +++++++++++---- 3 files changed, 99 insertions(+), 18 deletions(-) diff --git a/.ci/sai/README.md b/.ci/sai/README.md index 2ea54b4c54a..3d210e2769e 100644 --- a/.ci/sai/README.md +++ b/.ci/sai/README.md @@ -72,6 +72,11 @@ known-host checking). `--source-repository` must be the checkout used to compute `SOURCE_SHA`. The command waits for Slurm completion and exits with a non-zero status when a case or infrastructure result is not fully successful. +Source is sent as one compressed Git bundle split into eight parallel rsync +transfers. The remote side checks the merged SHA-256 and the Git bundle before +updating its cache. Later runs send only commits after the cached SHA; a run at +the same SHA sends no source data. + For all available subcommands and options, run: ```bash diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index ffe57e9299f..5df96301e5b 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -4,6 +4,7 @@ import argparse import configparser import fcntl +import hashlib import json import os import re @@ -15,6 +16,7 @@ import tempfile import time from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from pathlib import PurePosixPath @@ -30,6 +32,7 @@ MAX_RESULT_MEMBERS = 10000 MAX_RESULT_BYTES = 2 * 1024**3 MAX_REPORT_BYTES = 1024**2 +BUNDLE_PARTS = 8 @dataclass(frozen=True) @@ -448,6 +451,49 @@ def _cache(project: Path) -> Path: return project / "cache" / "repository.git" +def _split_bundle(bundle: Path, destination: Path) -> Tuple[List[Path], str]: + size = bundle.stat().st_size + if size < BUNDLE_PARTS: + raise ValueError("source bundle is too small") + digest = hashlib.sha256() + parts = [] + with bundle.open("rb") as source: + for index in range(BUNDLE_PARTS): + amount = size * (index + 1) // BUNDLE_PARTS - size * index // BUNDLE_PARTS + data = source.read(amount) + digest.update(data) + part = destination / "source.bundle.part.{:02d}".format(index) + part.write_bytes(data) + parts.append(part) + if source.read(1): + raise RuntimeError("unable to split source bundle") + return parts, digest.hexdigest() + + +def _assemble_bundle(root: Path, expected: str) -> Path: + if not re.fullmatch(r"[0-9a-f]{64}", expected): + raise ValueError("invalid bundle checksum") + parts = [root / "source.bundle.part.{:02d}".format(index) for index in range(BUNDLE_PARTS)] + if set(root.glob("source.bundle.part.*")) != set(parts) or not all(path.is_file() for path in parts): + raise ValueError("incomplete source bundle") + temporary = root / ".source.bundle.tmp" + digest = hashlib.sha256() + with temporary.open("wb") as output: + for part in parts: + with part.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + output.write(block) + if digest.hexdigest() != expected: + temporary.unlink() + raise ValueError("source bundle checksum mismatch") + bundle = root / "source.bundle" + os.replace(str(temporary), str(bundle)) + for part in parts: + part.unlink() + return bundle + + def remote_prepare(project: Path, run: Path) -> None: project = project.resolve() run = run.resolve() @@ -472,15 +518,18 @@ def remote_prepare(project: Path, run: Path) -> None: print(json.dumps({"run_root": str(run), "cache_sha": base})) -def remote_receive(project: Path, run: Path, source_sha: str) -> None: +def remote_receive(project: Path, run: Path, source_sha: str, bundle_checksum: str) -> None: if not SHA.fullmatch(source_sha): raise ValueError("invalid source SHA") cache = _cache(project.resolve()) - bundle = run.resolve() / "input" / "source.bundle" + input_root = run.resolve() / "input" + bundle = input_root / "source.bundle" + if bundle_checksum != "-": + bundle = _assemble_bundle(input_root, bundle_checksum) lock = cache.parent / "lock" with lock.open("w") as stream: fcntl.flock(stream, fcntl.LOCK_EX) - if bundle.is_file(): + if bundle_checksum != "-": heads = _command(("git", "bundle", "list-heads", str(bundle))).stdout.split() if len(heads) != 2 or heads != [source_sha, "HEAD"]: raise ValueError("bundle does not advertise the requested SHA") @@ -490,6 +539,8 @@ def remote_receive(project: Path, run: Path, source_sha: str) -> None: "HEAD:refs/sai/{}".format(source_sha), )) else: + if any(input_root.iterdir()): + raise ValueError("unexpected bundle data for cache hit") _command(("git", "--git-dir", str(cache), "cat-file", "-e", source_sha + "^{commit}")) source = run / "source" shutil.rmtree(str(source)) @@ -667,16 +718,25 @@ def local_run(args: argparse.Namespace) -> int: with tempfile.TemporaryDirectory() as directory: bundle = Path(directory) / "source.bundle" revision = _bundle_revision(repository, base, args.source_sha) + checksum = "-" if revision: _command(("git", "bundle", "create", str(bundle), revision), repository) - _retry(( - "rsync", "-a", "--partial", "--info=progress2", "-e", - "ssh -F {}".format(args.ssh_config), str(bundle), - "{}:{}/input/source.bundle".format(args.target, run), - ), stdout=None) + parts, checksum = _split_bundle(bundle, Path(directory)) + + def upload(item: Tuple[int, Path]) -> None: + index, part = item + print("[bundle {}/{}] {}".format(index + 1, BUNDLE_PARTS, part.name), flush=True) + _retry(( + "rsync", "-a", "--partial", "--info=name1,progress2", "-e", + "ssh -F {}".format(args.ssh_config), str(part), + "{}:{}/input/".format(args.target, run), + ), stdout=None) + + with ThreadPoolExecutor(max_workers=BUNDLE_PARTS) as pool: + list(pool.map(upload, enumerate(parts))) _ssh(args.ssh_config, args.target, ( "python3", remote_control + "/sai.py", "receive", - args.project_root, str(run), args.source_sha, + args.project_root, str(run), args.source_sha, checksum, )) launch = "nohup python3 {} remote-run {} > {}/results/coordinator.log 2>&1 < /dev/null &".format( shlex.quote(remote_control + "/sai.py"), shlex.quote(str(run)), shlex.quote(str(run)) @@ -832,6 +892,7 @@ def parser() -> argparse.ArgumentParser: receive.add_argument("project", type=Path) receive.add_argument("run", type=Path) receive.add_argument("source_sha") + receive.add_argument("bundle_checksum") collect_parser = commands.add_parser("collect") collect_parser.add_argument("run", type=Path) archive = commands.add_parser("archive") @@ -871,7 +932,7 @@ def main() -> int: if args.command == "prepare": remote_prepare(args.project, args.run) elif args.command == "receive": - remote_receive(args.project, args.run, args.source_sha) + remote_receive(args.project, args.run, args.source_sha, args.bundle_checksum) elif args.command == "collect": collect(args.run) elif args.command == "archive": diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index 9acc1ac719a..d7f98ab1dda 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -150,6 +150,14 @@ class TransferTests(unittest.TestCase): def _git(root, *arguments): return subprocess.run(("git", *arguments), cwd=str(root), check=True, text=True, capture_output=True) + def _bundle(self, repository, run, revision): + bundle = run / "source.bundle.local" + self._git(repository, "bundle", "create", str(bundle), revision) + parts, checksum = sai._split_bundle(bundle, run / "input") + bundle.unlink() + self.assertEqual(len(parts), 8) + return checksum + def test_full_then_incremental_bundle(self): with tempfile.TemporaryDirectory() as directory: home = Path(directory) @@ -168,8 +176,8 @@ def test_full_then_incremental_bundle(self): (run1 / "control").mkdir(parents=True) with mock.patch("sai.Path.home", return_value=home): sai.remote_prepare(project, run1) - self._git(repository, "bundle", "create", str(run1 / "input" / "source.bundle"), "HEAD") - sai.remote_receive(project, run1, first) + checksum = self._bundle(repository, run1, "HEAD") + sai.remote_receive(project, run1, first, checksum) self.assertEqual((run1 / "source" / "value.txt").read_text(), "one\n") (repository / "value.txt").write_text("two\n", encoding="utf-8") @@ -179,11 +187,8 @@ def test_full_then_incremental_bundle(self): (run2 / "control").mkdir(parents=True) with mock.patch("sai.Path.home", return_value=home): sai.remote_prepare(project, run2) - self._git( - repository, "bundle", "create", str(run2 / "input" / "source.bundle"), - first + "..HEAD", - ) - sai.remote_receive(project, run2, second) + checksum = self._bundle(repository, run2, first + "..HEAD") + sai.remote_receive(project, run2, second, checksum) self.assertEqual((run2 / "source" / "value.txt").read_text(), "two\n") run3 = project / "runs" / "manual" / "3-1" @@ -191,9 +196,19 @@ def test_full_then_incremental_bundle(self): with mock.patch("sai.Path.home", return_value=home): sai.remote_prepare(project, run3) self.assertIsNone(sai._bundle_revision(repository, second, second)) - sai.remote_receive(project, run3, second) + sai.remote_receive(project, run3, second, "-") self.assertEqual((run3 / "source" / "value.txt").read_text(), "two\n") + def test_bundle_merge_rejects_corruption(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + bundle = root / "bundle" + bundle.write_bytes(bytes(range(256)) * 4) + parts, checksum = sai._split_bundle(bundle, root) + parts[0].write_bytes(b"corrupt") + with self.assertRaisesRegex(ValueError, "checksum mismatch"): + sai._assemble_bundle(root, checksum) + def test_prepare_rejects_reused_run(self): with tempfile.TemporaryDirectory() as directory: home = Path(directory) From 2b48930beafce9ee0f9563e7eca10a44c74f9438 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 11:49:18 +0800 Subject: [PATCH 03/14] [skip ci] ci: retry initial SAI connection --- .ci/sai/sai.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index 5df96301e5b..4a169bad64a 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -704,7 +704,10 @@ def local_run(args: argparse.Namespace) -> int: raise ValueError("source SHA must be the checked-out HEAD") run = project / "runs" / args.namespace / "{}-{}".format(args.run_id, args.run_attempt) remote_control = str(run / "control") - _ssh(args.ssh_config, args.target, ("mkdir", "-p", remote_control)) + _retry(( + "ssh", "-F", str(args.ssh_config), args.target, + shlex.join(("mkdir", "-p", remote_control)), + )) _retry(( "rsync", "-a", "--delete", "--info=stats2", "-e", "ssh -F {}".format(args.ssh_config), str(ROOT) + "/", From 973dee7d09a299d5ab61bb541b8540e431847dd9 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 13:03:08 +0800 Subject: [PATCH 04/14] ci: initialize Lmod preload state --- .ci/sai/modules.sh | 1 + .ci/sai/tests/test_sai.py | 1 + 2 files changed, 2 insertions(+) diff --git a/.ci/sai/modules.sh b/.ci/sai/modules.sh index 03f70a0d7d2..87b47437fa5 100644 --- a/.ci/sai/modules.sh +++ b/.ci/sai/modules.sh @@ -1,6 +1,7 @@ #!/usr/bin/env bash export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-} +export LD_PRELOAD=${LD_PRELOAD:-} export CPATH=${CPATH:-} export CMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH:-} diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index d7f98ab1dda..76135b158cb 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -94,6 +94,7 @@ def test_resource_template_is_complete_and_has_no_cpu_request(self): def test_modules_do_not_spell_dependency_paths(self): text = (ROOT / "modules.sh").read_text(encoding="utf-8") self.assertIn("module load abacus/", text) + self.assertIn("LD_PRELOAD=${LD_PRELOAD:-}", text) self.assertIn("CMAKE_LIBRARY_PATH=${LIBRARY_PATH:-}", text) self.assertIn("CMAKE_INCLUDE_PATH=${CPATH:-}", text) self.assertNotRegex(text, r"CUSOLVERMP_PATH|CUBLASMP_PATH|NCCL_PATH|/lib/lib") From 2a7236a1b891e3b25fa61d1491f6fb445abee50b Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 17:03:07 +0800 Subject: [PATCH 05/14] ci: retry MPI daemon startup once --- .ci/sai/sai.py | 27 +++++++++++++++++++-------- .ci/sai/tests/test_sai.py | 10 +++++++--- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index 4a169bad64a..21d40226d34 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -344,9 +344,13 @@ def _stream(command: Sequence[str], cwd: Path, log: Path) -> int: return process.wait() -def _pmix(log: Path) -> bool: +def _mpi_startup_failure(log: Path) -> bool: data = log.read_bytes() - return bool(PMIX.search(data)) and b"MPI_Init_thread" in data and b"PMIx_Init failed" in data + return ( + bool(PMIX.search(data)) and b"MPI_Init_thread" in data and b"PMIx_Init failed" in data + ) or ( + b"srun returned non-zero exit status (512)" in data and b"per-node daemon" in data + ) def worker(source: Path, control: Path, install: Path, results: Path, manifest: Path) -> int: @@ -377,7 +381,7 @@ def worker(source: Path, control: Path, install: Path, results: Path, manifest: raise RuntimeError("ABACUS has unresolved runtime libraries") start = time.time() returncode = 2 - final_pmix = False + final_startup_failure = False try: if runner == "autotest": cases_file = case.parent / "CASES.task.txt" @@ -397,21 +401,28 @@ def worker(source: Path, control: Path, install: Path, results: Path, manifest: pass log = artifacts / "attempt-{}.log".format(attempt) returncode = _stream(command, case.parent, log) - final_pmix = returncode != 0 and _pmix(log) - if returncode == 0 or attempt == 2 or not final_pmix: + final_startup_failure = returncode != 0 and _mpi_startup_failure(log) + if returncode == 0 or attempt == 2 or not final_startup_failure: break time.sleep(10) else: - returncode = _stream(( + command = ( "timeout", "--signal=TERM", "--kill-after=30s", "35m", "mpirun", "-np", os.environ["SLURM_NTASKS"], str(abacus), - ), case, artifacts / "cusolvermp.log") + ) + for attempt in (1, 2): + log = artifacts / ("cusolvermp.log" if attempt == 1 else "cusolvermp-retry.log") + returncode = _stream(command, case, log) + final_startup_failure = returncode != 0 and _mpi_startup_failure(log) + if returncode == 0 or attempt == 2 or not final_startup_failure: + break + time.sleep(10) for filename in ("log.txt", "result.out", "warning.log"): if (case / filename).is_file(): shutil.copy2(str(case / filename), str(artifacts / filename)) state = "PASS" if returncode == 0 else ( "TIMEOUT" if returncode in (124, 137, 143) else - "INFRA" if final_pmix else "FAIL" + "INFRA" if final_startup_failure else "FAIL" ) _atomic(results / "status" / (case_id.replace("/", "__") + ".json"), { "case_id": case_id, "resource": resource, "runner": runner, diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index 76135b158cb..bb6b0d79bba 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -326,13 +326,17 @@ def test_report_rejects_wrong_matrix_identity(self): with self.assertRaisesRegex(ValueError, "result case"): sai.report(args) - def test_pmix_signature_requires_all_markers(self): + def test_mpi_startup_failure_requires_complete_signature(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "log" path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE MPI_Init_thread PMIx_Init failed") - self.assertTrue(sai._pmix(path)) + self.assertTrue(sai._mpi_startup_failure(path)) path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE") - self.assertFalse(sai._pmix(path)) + self.assertFalse(sai._mpi_startup_failure(path)) + path.write_bytes(b"srun returned non-zero exit status (512) from launching the per-node daemon") + self.assertTrue(sai._mpi_startup_failure(path)) + path.write_bytes(b"srun returned non-zero exit status (512)") + self.assertFalse(sai._mpi_startup_failure(path)) class PolicyTests(unittest.TestCase): From 4caa0a4a79a738fec0a1719252d31cd9eac25625 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 17:16:08 +0800 Subject: [PATCH 06/14] ci: retry SAI result download --- .ci/sai/sai.py | 27 +++++++++++++++++++++------ .ci/sai/tests/test_sai.py | 15 +++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index 21d40226d34..d6a0886af12 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -458,6 +458,20 @@ def _retry(command: Sequence[str], cwd: Optional[Path] = None, stdout: Any = sub raise RuntimeError(str(error)) +def _retry_download(command: Sequence[str], destination: Path) -> None: + error: Optional[Exception] = None + for attempt in range(3): + try: + with destination.open("wb") as output: + _command(command, stdout=output) + return + except RuntimeError as caught: + error = caught + if attempt < 2: + time.sleep(5 * (attempt + 1)) + raise RuntimeError(str(error)) + + def _cache(project: Path) -> Path: return project / "cache" / "repository.git" @@ -775,14 +789,15 @@ def upload(item: Tuple[int, Path]) -> None: failures = 0 time.sleep(30) args.artifacts.mkdir(parents=True, exist_ok=True) - process = subprocess.Popen(( + command = ( "ssh", "-F", str(args.ssh_config), args.target, shlex.join(("python3", remote_control + "/sai.py", "collect", str(run))), - ), stdout=subprocess.PIPE) - assert process.stdout is not None - _extract_results(process.stdout, args.artifacts) - if process.wait(): - raise RuntimeError("unable to collect SAI results") + ) + archive = args.artifacts / ".results.tar.gz" + _retry_download(command, archive) + with archive.open("rb") as stream: + _extract_results(stream, args.artifacts) + archive.unlink() _atomic(args.artifacts / "run.json", { "project_root": args.project_root, "run_root": str(run), "source_sha": args.source_sha, diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index bb6b0d79bba..c594b84caed 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -229,6 +229,21 @@ def test_transfer_retry_is_bounded(self): self.assertEqual(sai._retry(("rsync", "source", "target")).stdout, "done") self.assertEqual(run.call_count, 3) + def test_download_retry_replaces_partial_file(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "result.tar.gz" + + def download(_command, _cwd=None, stdout=None): + stdout.write(b"partial" if download.calls == 0 else b"complete") + download.calls += 1 + if download.calls == 1: + raise RuntimeError("connection closed") + + download.calls = 0 + with mock.patch("sai._command", side_effect=download), mock.patch("sai.time.sleep"): + sai._retry_download(("ssh", "collect"), path) + self.assertEqual(path.read_bytes(), b"complete") + def test_result_archive_rejects_traversal_and_links(self): for name, link in (("../sai-ssh/id_ed25519", None), ("results/key", "../key")): payload = io.BytesIO() From ec41a5f435c2251ba0e3c7b621cf9ae1fa7713fe Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 17:41:44 +0800 Subject: [PATCH 07/14] ci: stabilize SAI multinode launch and cleanup --- .ci/sai/case.sbatch.in | 2 ++ .ci/sai/tests/test_sai.py | 2 ++ .github/workflows/sai-gpu.yml | 8 ++++++-- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/.ci/sai/case.sbatch.in b/.ci/sai/case.sbatch.in index c9b0048eeb7..ad50b6205a0 100644 --- a/.ci/sai/case.sbatch.in +++ b/.ci/sai/case.sbatch.in @@ -26,6 +26,8 @@ source @CONTROL@/modules.sh SAI_SYSTEM_MPIRUN=$(command -v mpirun) export SAI_SYSTEM_MPIRUN export MAP_OPT +export OMPI_MCA_plm_slurm_args=--external-launcher +export PRTE_MCA_plm_slurm_args=--external-launcher nvidia-smi dmon -s pucvmte -o T > @RESULTS@/dmon-${SLURM_JOB_ID}.log & SAI_DISABLE_NCCL_IB=@DISABLE_NCCL_IB@ diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index c594b84caed..c963fc835b5 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -88,6 +88,8 @@ def test_resource_template_is_complete_and_has_no_cpu_request(self): self.assertIn("#SBATCH --nodes=1", text) self.assertIn("#SBATCH --array=0-3%4", text) self.assertIn("#SBATCH --dependency=afterok:123", text) + self.assertIn("OMPI_MCA_plm_slurm_args=--external-launcher", text) + self.assertIn("PRTE_MCA_plm_slurm_args=--external-launcher", text) self.assertNotRegex(text, r"cpus-per-task|--mem(?:ory)?|--nodelist") self.assertNotRegex(text, r"@[A-Z_]+@") diff --git a/.github/workflows/sai-gpu.yml b/.github/workflows/sai-gpu.yml index 10feb9a1561..7c0ba4f582c 100644 --- a/.github/workflows/sai-gpu.yml +++ b/.github/workflows/sai-gpu.yml @@ -193,8 +193,12 @@ jobs: test -f "$ARTIFACT_ROOT/run.json" || exit 0 project=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["project_root"])' "$ARTIFACT_ROOT/run.json") run=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["run_root"])' "$ARTIFACT_ROOT/run.json") - ssh -F "$SAI_SSH_CONFIG" sai-ci \ - python3 "$run/control/sai.py" archive "$project" "$run" + for attempt in 1 2 3; do + ssh -F "$SAI_SSH_CONFIG" sai-ci \ + "test ! -d '$run' || python3 '$run/control/sai.py' archive '$project' '$run'" && exit 0 + test "$attempt" = 3 || sleep 5 + done + exit 1 - name: Remove SSH credentials if: always() From ad092476cce2c1a180559c6ce30dd142be84cbc3 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 18:05:07 +0800 Subject: [PATCH 08/14] ci: export loaded modules to SAI job steps --- .ci/sai/case.sbatch.in | 1 + .ci/sai/tests/test_sai.py | 1 + 2 files changed, 2 insertions(+) diff --git a/.ci/sai/case.sbatch.in b/.ci/sai/case.sbatch.in index ad50b6205a0..b1a0735897c 100644 --- a/.ci/sai/case.sbatch.in +++ b/.ci/sai/case.sbatch.in @@ -26,6 +26,7 @@ source @CONTROL@/modules.sh SAI_SYSTEM_MPIRUN=$(command -v mpirun) export SAI_SYSTEM_MPIRUN export MAP_OPT +export SLURM_EXPORT_ENV=ALL export OMPI_MCA_plm_slurm_args=--external-launcher export PRTE_MCA_plm_slurm_args=--external-launcher nvidia-smi dmon -s pucvmte -o T > @RESULTS@/dmon-${SLURM_JOB_ID}.log & diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index c963fc835b5..60b0281511a 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -88,6 +88,7 @@ def test_resource_template_is_complete_and_has_no_cpu_request(self): self.assertIn("#SBATCH --nodes=1", text) self.assertIn("#SBATCH --array=0-3%4", text) self.assertIn("#SBATCH --dependency=afterok:123", text) + self.assertIn("SLURM_EXPORT_ENV=ALL", text) self.assertIn("OMPI_MCA_plm_slurm_args=--external-launcher", text) self.assertIn("PRTE_MCA_plm_slurm_args=--external-launcher", text) self.assertNotRegex(text, r"cpus-per-task|--mem(?:ory)?|--nodelist") From 749e3645c7cabc412a768e647a530142dda7518a Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 18:38:09 +0800 Subject: [PATCH 09/14] ci: retain SAI cleanup metadata on client failure --- .ci/sai/README.md | 8 +++----- .ci/sai/sai.py | 10 +++++----- .ci/sai/tests/test_sai.py | 27 +++++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 10 deletions(-) diff --git a/.ci/sai/README.md b/.ci/sai/README.md index 3d210e2769e..ba1e3118e96 100644 --- a/.ci/sai/README.md +++ b/.ci/sai/README.md @@ -153,8 +153,6 @@ GPU availability, node and task limits, and the `time_seconds` value for that resource. Slurm output is in `results/`; an allocation or queue delay is an infrastructure issue, not a case failure. -**MPI/PMIx initialization fails.** Autotest retries once after a recognized -`PMIX_ERR_FILE_OPEN_FAILURE` or `PMIX_ERR_OUT_OF_RESOURCE` initialization -failure. If it persists, inspect the attempt logs and the loaded MPI module, -Slurm allocation, and mapping file. The `cusolvermp` runner does not use this -retry. +**MPI/PMIx initialization fails.** Both runners retry once after a recognized +MPI startup failure. If it persists, inspect the attempt logs and the loaded +MPI module, Slurm allocation, and mapping file. diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index d6a0886af12..26ca2d7af56 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -728,6 +728,11 @@ def local_run(args: argparse.Namespace) -> int: if actual != args.source_sha or not SHA.fullmatch(actual): raise ValueError("source SHA must be the checked-out HEAD") run = project / "runs" / args.namespace / "{}-{}".format(args.run_id, args.run_attempt) + args.artifacts.mkdir(parents=True, exist_ok=True) + _atomic(args.artifacts / "run.json", { + "project_root": args.project_root, "run_root": str(run), + "source_sha": args.source_sha, + }) remote_control = str(run / "control") _retry(( "ssh", "-F", str(args.ssh_config), args.target, @@ -788,7 +793,6 @@ def upload(item: Tuple[int, Path]) -> None: else: failures = 0 time.sleep(30) - args.artifacts.mkdir(parents=True, exist_ok=True) command = ( "ssh", "-F", str(args.ssh_config), args.target, shlex.join(("python3", remote_control + "/sai.py", "collect", str(run))), @@ -798,10 +802,6 @@ def upload(item: Tuple[int, Path]) -> None: with archive.open("rb") as stream: _extract_results(stream, args.artifacts) archive.unlink() - _atomic(args.artifacts / "run.json", { - "project_root": args.project_root, "run_root": str(run), - "source_sha": args.source_sha, - }) result_path = args.artifacts / "results" / "result.json" if result_path.is_file(): result = _read_result(result_path) diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index 60b0281511a..10516d932b6 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -247,6 +247,33 @@ def download(_command, _cwd=None, stdout=None): sai._retry_download(("ssh", "collect"), path) self.assertEqual(path.read_bytes(), b"complete") + def test_local_run_records_cleanup_metadata_before_remote_work(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_sha = "a" * 40 + args = argparse.Namespace( + source_repository=root, + project_root="/project", + target="sai-ci", + namespace="manual", + run_id="1", + run_attempt="1", + source_sha=source_sha, + artifacts=root / "artifacts", + ssh_config=root / "ssh-config", + ) + revision = mock.Mock(stdout=source_sha + "\n") + with mock.patch("sai._command", return_value=revision), \ + mock.patch("sai._retry", side_effect=RuntimeError("disconnected")), \ + self.assertRaisesRegex(RuntimeError, "disconnected"): + sai.local_run(args) + metadata = json.loads((args.artifacts / "run.json").read_text()) + self.assertEqual(metadata, { + "project_root": "/project", + "run_root": "/project/runs/manual/1-1", + "source_sha": source_sha, + }) + def test_result_archive_rejects_traversal_and_links(self): for name, link in (("../sai-ssh/id_ed25519", None), ("results/key", "../key")): payload = io.BytesIO() From 5bf08b64e3480430bafa796f9fd17d482c546b2c Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Mon, 27 Jul 2026 19:23:28 +0800 Subject: [PATCH 10/14] ci: allow SAI results to comment on pull requests --- .ci/sai/tests/test_sai.py | 1 + .github/workflows/sai-gpu.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/tests/test_sai.py index 10516d932b6..c895a155aa4 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/tests/test_sai.py @@ -391,6 +391,7 @@ def test_workflow_uses_trusted_control_and_protected_environment(self): self.assertIn("sai-ssh-scheduled", text) self.assertIn("sai-ssh-manual", text) self.assertIn("/abacus-ci sai-gpu", text) + self.assertIn("pull-requests: write", text) self.assertNotIn("cpus-per-task", text) diff --git a/.github/workflows/sai-gpu.yml b/.github/workflows/sai-gpu.yml index 7c0ba4f582c..b27406ae1a3 100644 --- a/.github/workflows/sai-gpu.yml +++ b/.github/workflows/sai-gpu.yml @@ -236,6 +236,7 @@ jobs: checks: write contents: read issues: write + pull-requests: write steps: - name: Checkout pinned reporter uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 From 4b4a4311bb4e625066ade7885d3e1427f103131b Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Tue, 28 Jul 2026 00:48:59 +0800 Subject: [PATCH 11/14] ci: expose one documented SAI run command --- .ci/sai/README.md | 6 +-- .ci/sai/sai.py | 87 ++++++++++++++++++++++++--------- .ci/sai/{tests => }/test_sai.py | 53 ++++++++++++++++++-- .github/workflows/sai-gpu.yml | 2 +- 4 files changed, 117 insertions(+), 31 deletions(-) rename .ci/sai/{tests => }/test_sai.py (88%) diff --git a/.ci/sai/README.md b/.ci/sai/README.md index ba1e3118e96..87a05e7de88 100644 --- a/.ci/sai/README.md +++ b/.ci/sai/README.md @@ -55,7 +55,7 @@ following uses the committed `HEAD` SHA and the same required options as CI: ```bash SOURCE_SHA="$(git rev-parse HEAD)" RUN_ID="$(date +%s)" -python3 .ci/sai/sai.py local \ +python3 .ci/sai/sai.py run \ --ssh-config "$HOME/.ssh/sai-config" \ --target sai-ci \ --project-root /home///abacus_sai_gpu_ci \ @@ -77,11 +77,11 @@ transfers. The remote side checks the merged SHA-256 and the Git bundle before updating its cache. Later runs send only commits after the cached SHA; a run at the same SHA sends no source data. -For all available subcommands and options, run: +For the public commands and all `run` options, use: ```bash python3 .ci/sai/sai.py --help -python3 .ci/sai/sai.py local --help +python3 .ci/sai/sai.py run --help ``` ## Matrix configuration diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index 26ca2d7af56..cb02dfa3edb 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -717,7 +717,7 @@ def _read_result(path: Path) -> Mapping[str, Any]: return _validate_result(json.loads(path.read_text(encoding="utf-8"))) -def local_run(args: argparse.Namespace) -> int: +def run(args: argparse.Namespace) -> int: repository = args.source_repository.resolve() project = Path(args.project_root) if not project.is_absolute() or any(not NAME.fullmatch(part) for part in project.parts[1:]): @@ -907,42 +907,81 @@ def github_finish() -> int: def parser() -> argparse.ArgumentParser: main = argparse.ArgumentParser(description=__doc__) - commands = main.add_subparsers(dest="command", required=True) - commands.add_parser("config", help="validate and summarize gpu-matrix.ini") - remote = commands.add_parser("remote-run") + commands = main.add_subparsers( + title="commands", dest="command", required=True, metavar="{run,config}", + ) + client = commands.add_parser( + "run", help="build and run the GPU matrix on SAI", + description="Transfer one committed ABACUS revision, build it, and run the GPU matrix on SAI.", + ) + client.add_argument( + "--ssh-config", required=True, type=Path, + help="SSH config file containing the target host alias", + ) + client.add_argument( + "--target", required=True, + help="host alias in the SSH config", + ) + client.add_argument( + "--project-root", required=True, + help="absolute SAI directory for caches, runs, and archives", + ) + client.add_argument( + "--source-repository", required=True, type=Path, + help="local ABACUS Git checkout to transfer", + ) + client.add_argument( + "--source-sha", required=True, + help="40-character commit SHA; it must equal the checkout's HEAD", + ) + client.add_argument( + "--namespace", required=True, + help="group for the remote run, such as manual, daily, or pr-123", + ) + client.add_argument( + "--run-id", required=True, + help="unique identifier for this run", + ) + client.add_argument( + "--run-attempt", required=True, + help="attempt number within the run ID", + ) + client.add_argument( + "--artifacts", required=True, type=Path, + help="local directory for downloaded results and client logs", + ) + commands.add_parser( + "config", help="validate and summarize gpu-matrix.ini", + description="Validate gpu-matrix.ini and print its resource and case counts.", + ) + + def internal(name: str) -> argparse.ArgumentParser: + return commands.add_parser(name) + + remote = internal("remote-run") remote.add_argument("run", type=Path) - worker_parser = commands.add_parser("worker") + worker_parser = internal("worker") for name in ("source", "control", "install", "results", "manifest"): worker_parser.add_argument(name, type=Path) - prepare = commands.add_parser("prepare") + prepare = internal("prepare") prepare.add_argument("project", type=Path) prepare.add_argument("run", type=Path) - receive = commands.add_parser("receive") + receive = internal("receive") receive.add_argument("project", type=Path) receive.add_argument("run", type=Path) receive.add_argument("source_sha") receive.add_argument("bundle_checksum") - collect_parser = commands.add_parser("collect") + collect_parser = internal("collect") collect_parser.add_argument("run", type=Path) - archive = commands.add_parser("archive") + archive = internal("archive") archive.add_argument("project", type=Path) archive.add_argument("run", type=Path) - local = commands.add_parser("local") - local.add_argument("--ssh-config", required=True, type=Path) - local.add_argument("--target", required=True) - local.add_argument("--project-root", required=True) - local.add_argument("--source-repository", required=True, type=Path) - local.add_argument("--source-sha", required=True) - local.add_argument("--namespace", required=True) - local.add_argument("--run-id", required=True) - local.add_argument("--run-attempt", required=True) - local.add_argument("--artifacts", required=True, type=Path) - report_parser = commands.add_parser("report") + report_parser = internal("report") report_parser.add_argument("--result", required=True, type=Path) report_parser.add_argument("--output", required=True, type=Path) report_parser.add_argument("--summary", type=Path) - commands.add_parser("github-admit") - commands.add_parser("github-finish") + internal("github-admit") + internal("github-finish") return main @@ -966,8 +1005,8 @@ def main() -> int: collect(args.run) elif args.command == "archive": print(archive_run(args.project, args.run)) - elif args.command == "local": - return local_run(args) + elif args.command == "run": + return run(args) elif args.command == "report": return report(args) elif args.command == "github-admit": diff --git a/.ci/sai/tests/test_sai.py b/.ci/sai/test_sai.py similarity index 88% rename from .ci/sai/tests/test_sai.py rename to .ci/sai/test_sai.py index c895a155aa4..3036291ff46 100644 --- a/.ci/sai/tests/test_sai.py +++ b/.ci/sai/test_sai.py @@ -11,7 +11,7 @@ from unittest import mock -ROOT = Path(__file__).resolve().parents[1] +ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT)) SPEC = importlib.util.spec_from_file_location("sai", ROOT / "sai.py") assert SPEC and SPEC.loader @@ -68,6 +68,52 @@ def test_unknown_resource_and_noncontiguous_case_fail(self): sai.load_config(path) +class CliTests(unittest.TestCase): + def test_top_level_help_only_exposes_public_commands(self): + help_text = subprocess.run( + (sys.executable, str(ROOT / "sai.py"), "--help"), + check=True, text=True, capture_output=True, + ).stdout + self.assertIn("{run,config}", help_text) + self.assertIn("build and run the GPU matrix on SAI", help_text) + self.assertNotIn("remote-run", help_text) + self.assertNotIn("github-admit", help_text) + + def test_run_help_describes_every_option(self): + help_text = subprocess.run( + (sys.executable, str(ROOT / "sai.py"), "run", "--help"), + check=True, text=True, capture_output=True, + ).stdout + for option in ( + "--ssh-config", "--target", "--project-root", "--source-repository", + "--source-sha", "--namespace", "--run-id", "--run-attempt", "--artifacts", + ): + self.assertIn(option, help_text) + self.assertIn("40-character commit SHA", help_text) + self.assertIn("checkout's", help_text) + + def test_run_options_parse(self): + arguments = [ + "run", "--ssh-config", "/tmp/ssh/config", "--target", "sai-ci", + "--project-root", "/home/user/abacus_sai_gpu_ci", + "--source-repository", "/tmp/abacus", "--source-sha", "a" * 40, + "--namespace", "manual", "--run-id", "42", "--run-attempt", "1", + "--artifacts", "/tmp/results", + ] + args = sai.parser().parse_args(arguments) + self.assertEqual(args.command, "run") + self.assertEqual(args.ssh_config, Path("/tmp/ssh/config")) + self.assertEqual(args.source_repository, Path("/tmp/abacus")) + self.assertEqual(args.artifacts, Path("/tmp/results")) + self.assertEqual(args.source_sha, "a" * 40) + + for index in range(1, len(arguments), 2): + with self.subTest(option=arguments[index]), \ + mock.patch("sys.stderr", new_callable=io.StringIO), \ + self.assertRaises(SystemExit): + sai.parser().parse_args(arguments[:index] + arguments[index + 2:]) + + class TemplateTests(unittest.TestCase): def test_resource_template_is_complete_and_has_no_cpu_request(self): config = sai.load_config() @@ -247,7 +293,7 @@ def download(_command, _cwd=None, stdout=None): sai._retry_download(("ssh", "collect"), path) self.assertEqual(path.read_bytes(), b"complete") - def test_local_run_records_cleanup_metadata_before_remote_work(self): + def test_run_records_cleanup_metadata_before_remote_work(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source_sha = "a" * 40 @@ -266,7 +312,7 @@ def test_local_run_records_cleanup_metadata_before_remote_work(self): with mock.patch("sai._command", return_value=revision), \ mock.patch("sai._retry", side_effect=RuntimeError("disconnected")), \ self.assertRaisesRegex(RuntimeError, "disconnected"): - sai.local_run(args) + sai.run(args) metadata = json.loads((args.artifacts / "run.json").read_text()) self.assertEqual(metadata, { "project_root": "/project", @@ -391,6 +437,7 @@ def test_workflow_uses_trusted_control_and_protected_environment(self): self.assertIn("sai-ssh-scheduled", text) self.assertIn("sai-ssh-manual", text) self.assertIn("/abacus-ci sai-gpu", text) + self.assertIn("sai.py run", text) self.assertIn("pull-requests: write", text) self.assertNotIn("cpus-per-task", text) diff --git a/.github/workflows/sai-gpu.yml b/.github/workflows/sai-gpu.yml index b27406ae1a3..13f4df46ee2 100644 --- a/.github/workflows/sai-gpu.yml +++ b/.github/workflows/sai-gpu.yml @@ -151,7 +151,7 @@ jobs: : "${project_root:?Set SAI_PROJECT_ROOT in the GitHub Environment}" mkdir -p "$ARTIFACT_ROOT" set +e - python3 control/.ci/sai/sai.py local \ + python3 control/.ci/sai/sai.py run \ --ssh-config "$SAI_SSH_CONFIG" \ --target sai-ci \ --project-root "$project_root" \ From 1fb9db6c95f88863df9c702c44cffd37d973bead Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Tue, 28 Jul 2026 01:32:10 +0800 Subject: [PATCH 12/14] ci: provide defaults for local SAI runs --- .ci/sai/README.md | 29 +++++++++--------------- .ci/sai/sai.py | 41 +++++++++++++++++++++------------- .ci/sai/test_sai.py | 54 +++++++++++++++++++++++++++++++++++++++------ 3 files changed, 84 insertions(+), 40 deletions(-) diff --git a/.ci/sai/README.md b/.ci/sai/README.md index 87a05e7de88..8ae246542da 100644 --- a/.ci/sai/README.md +++ b/.ci/sai/README.md @@ -49,28 +49,21 @@ Other issue comments are ignored. ## Run locally -Run from the repository checkout whose `HEAD` is the commit to test. The -following uses the committed `HEAD` SHA and the same required options as CI: +Run from the repository checkout whose committed `HEAD` is the revision to +test. If the SSH config defines the default `sai-ci` target, no options are +required: ```bash -SOURCE_SHA="$(git rev-parse HEAD)" -RUN_ID="$(date +%s)" -python3 .ci/sai/sai.py run \ - --ssh-config "$HOME/.ssh/sai-config" \ - --target sai-ci \ - --project-root /home///abacus_sai_gpu_ci \ - --source-repository "$PWD" \ - --source-sha "$SOURCE_SHA" \ - --namespace manual \ - --run-id "$RUN_ID" \ - --run-attempt 1 \ - --artifacts ./sai-artifacts +python3 .ci/sai/sai.py run ``` -The SSH config must define the `sai-ci` target (host, port, user, key, and -known-host checking). `--source-repository` must be the checkout used to -compute `SOURCE_SHA`. The command waits for Slurm completion and exits with a -non-zero status when a case or infrastructure result is not fully successful. +By default, the command uses `~/.ssh/config`, the `sai-ci` host alias, the +current Git checkout and its `HEAD`, and `~/abacus_sai_gpu_ci` below the remote +user's home. It writes downloaded results to `./sai-artifacts`. Override any +value shown by `run --help`; for example, use `--target SAI-abacus` for a +different local host alias. The command waits for Slurm completion and exits +with a non-zero status when a case or infrastructure result is not fully +successful. Source is sent as one compressed Git bundle split into eight parallel rsync transfers. The remote side checks the merged SHA-256 and the Git bundle before diff --git a/.ci/sai/sai.py b/.ci/sai/sai.py index cb02dfa3edb..a8096036386 100644 --- a/.ci/sai/sai.py +++ b/.ci/sai/sai.py @@ -718,15 +718,25 @@ def _read_result(path: Path) -> Mapping[str, Any]: def run(args: argparse.Namespace) -> int: + args.ssh_config = args.ssh_config.expanduser() + args.source_repository = args.source_repository.expanduser() + args.artifacts = args.artifacts.expanduser() repository = args.source_repository.resolve() - project = Path(args.project_root) - if not project.is_absolute() or any(not NAME.fullmatch(part) for part in project.parts[1:]): - raise ValueError("project root must be a simple absolute path") if not all(NAME.fullmatch(value) for value in (args.target, args.namespace, args.run_id, args.run_attempt)): raise ValueError("invalid target or run name") actual = _command(("git", "rev-parse", "HEAD"), repository).stdout.strip() - if actual != args.source_sha or not SHA.fullmatch(actual): + if args.source_sha == "HEAD": + args.source_sha = actual + if args.source_sha != actual or not SHA.fullmatch(actual): raise ValueError("source SHA must be the checked-out HEAD") + if args.project_root == "~" or args.project_root.startswith("~/"): + command = shlex.join(("python3", "-c", "from pathlib import Path; print(Path.home())")) + home = _retry(("ssh", "-F", str(args.ssh_config), args.target, command)).stdout.strip() + relative = "" if args.project_root == "~" else args.project_root[2:] + args.project_root = str(Path(home) / relative) + project = Path(args.project_root) + if not project.is_absolute() or any(not NAME.fullmatch(part) for part in project.parts[1:]): + raise ValueError("project root must be a simple absolute path") run = project / "runs" / args.namespace / "{}-{}".format(args.run_id, args.run_attempt) args.artifacts.mkdir(parents=True, exist_ok=True) _atomic(args.artifacts / "run.json", { @@ -913,41 +923,42 @@ def parser() -> argparse.ArgumentParser: client = commands.add_parser( "run", help="build and run the GPU matrix on SAI", description="Transfer one committed ABACUS revision, build it, and run the GPU matrix on SAI.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) client.add_argument( - "--ssh-config", required=True, type=Path, + "--ssh-config", type=Path, default=Path("~/.ssh/config"), help="SSH config file containing the target host alias", ) client.add_argument( - "--target", required=True, + "--target", default="sai-ci", help="host alias in the SSH config", ) client.add_argument( - "--project-root", required=True, - help="absolute SAI directory for caches, runs, and archives", + "--project-root", default="~/abacus_sai_gpu_ci", + help="SAI directory for caches, runs, and archives", ) client.add_argument( - "--source-repository", required=True, type=Path, + "--source-repository", type=Path, default=Path("."), help="local ABACUS Git checkout to transfer", ) client.add_argument( - "--source-sha", required=True, - help="40-character commit SHA; it must equal the checkout's HEAD", + "--source-sha", default="HEAD", + help="commit to test; it must resolve to the checkout's HEAD", ) client.add_argument( - "--namespace", required=True, + "--namespace", default="manual", help="group for the remote run, such as manual, daily, or pr-123", ) client.add_argument( - "--run-id", required=True, + "--run-id", default=str(int(time.time())), help="unique identifier for this run", ) client.add_argument( - "--run-attempt", required=True, + "--run-attempt", default="1", help="attempt number within the run ID", ) client.add_argument( - "--artifacts", required=True, type=Path, + "--artifacts", type=Path, default=Path("sai-artifacts"), help="local directory for downloaded results and client logs", ) commands.add_parser( diff --git a/.ci/sai/test_sai.py b/.ci/sai/test_sai.py index 3036291ff46..c19aa9943e3 100644 --- a/.ci/sai/test_sai.py +++ b/.ci/sai/test_sai.py @@ -89,8 +89,12 @@ def test_run_help_describes_every_option(self): "--source-sha", "--namespace", "--run-id", "--run-attempt", "--artifacts", ): self.assertIn(option, help_text) - self.assertIn("40-character commit SHA", help_text) - self.assertIn("checkout's", help_text) + normalized = " ".join(help_text.split()) + for default in ( + "~/.ssh/config", "sai-ci", "~/abacus_sai_gpu_ci", "HEAD", "manual", + "sai-artifacts", + ): + self.assertIn("default: {}".format(default), normalized) def test_run_options_parse(self): arguments = [ @@ -107,11 +111,18 @@ def test_run_options_parse(self): self.assertEqual(args.artifacts, Path("/tmp/results")) self.assertEqual(args.source_sha, "a" * 40) - for index in range(1, len(arguments), 2): - with self.subTest(option=arguments[index]), \ - mock.patch("sys.stderr", new_callable=io.StringIO), \ - self.assertRaises(SystemExit): - sai.parser().parse_args(arguments[:index] + arguments[index + 2:]) + def test_run_defaults_parse(self): + with mock.patch("sai.time.time", return_value=42): + args = sai.parser().parse_args(["run"]) + self.assertEqual(args.ssh_config, Path("~/.ssh/config")) + self.assertEqual(args.target, "sai-ci") + self.assertEqual(args.project_root, "~/abacus_sai_gpu_ci") + self.assertEqual(args.source_repository, Path(".")) + self.assertEqual(args.source_sha, "HEAD") + self.assertEqual(args.namespace, "manual") + self.assertEqual(args.run_id, "42") + self.assertEqual(args.run_attempt, "1") + self.assertEqual(args.artifacts, Path("sai-artifacts")) class TemplateTests(unittest.TestCase): @@ -320,6 +331,35 @@ def test_run_records_cleanup_metadata_before_remote_work(self): "source_sha": source_sha, }) + def test_run_resolves_head_and_remote_home_defaults(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_sha = "a" * 40 + args = argparse.Namespace( + source_repository=root, + project_root="~/abacus_sai_gpu_ci", + target="sai-ci", + namespace="manual", + run_id="1", + run_attempt="1", + source_sha="HEAD", + artifacts=root / "artifacts", + ssh_config=Path("~/.ssh/config"), + ) + revision = mock.Mock(stdout=source_sha + "\n") + home = mock.Mock(stdout="/home/user\n") + with mock.patch("sai.Path.home", return_value=Path("/home/local")), \ + mock.patch("sai._command", return_value=revision), \ + mock.patch("sai._retry", side_effect=[home, RuntimeError("disconnected")]), \ + self.assertRaisesRegex(RuntimeError, "disconnected"): + sai.run(args) + metadata = json.loads((args.artifacts / "run.json").read_text()) + self.assertEqual(metadata, { + "project_root": "/home/user/abacus_sai_gpu_ci", + "run_root": "/home/user/abacus_sai_gpu_ci/runs/manual/1-1", + "source_sha": source_sha, + }) + def test_result_archive_rejects_traversal_and_links(self): for name, link in (("../sai-ssh/id_ed25519", None), ("results/key", "../key")): payload = io.BytesIO() From 0a6e1a0fb695076e507e4852c083a52420d65f36 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Tue, 28 Jul 2026 02:05:44 +0800 Subject: [PATCH 13/14] ci: make remote GPU validation site-neutral --- .ci/sai/mpirun_with_mapping.sh | 7 - .ci/{sai => slurm}/README.md | 70 +++---- .ci/{sai => slurm}/build.sbatch.in | 0 .ci/{sai => slurm}/case.sbatch.in | 10 +- .ci/{sai => slurm}/gpu-matrix.ini | 0 .ci/{sai => slurm}/known_hosts | 0 .ci/{sai => slurm}/modules.sh | 0 .ci/slurm/mpirun_with_mapping.sh | 7 + .ci/{sai/sai.py => slurm/runner.py} | 38 ++-- .ci/{sai => slurm}/slurm.py | 2 +- .ci/{sai/test_sai.py => slurm/test_runner.py} | 178 +++++++++--------- .../{sai-gpu.yml => gpu-validation.yml} | 76 ++++---- 12 files changed, 197 insertions(+), 191 deletions(-) delete mode 100755 .ci/sai/mpirun_with_mapping.sh rename .ci/{sai => slurm}/README.md (67%) rename .ci/{sai => slurm}/build.sbatch.in (100%) rename .ci/{sai => slurm}/case.sbatch.in (84%) rename .ci/{sai => slurm}/gpu-matrix.ini (100%) rename .ci/{sai => slurm}/known_hosts (100%) rename .ci/{sai => slurm}/modules.sh (100%) create mode 100755 .ci/slurm/mpirun_with_mapping.sh rename .ci/{sai/sai.py => slurm/runner.py} (97%) rename .ci/{sai => slurm}/slurm.py (98%) rename .ci/{sai/test_sai.py => slurm/test_runner.py} (76%) rename .github/workflows/{sai-gpu.yml => gpu-validation.yml} (78%) diff --git a/.ci/sai/mpirun_with_mapping.sh b/.ci/sai/mpirun_with_mapping.sh deleted file mode 100755 index cc4b6cd9b4c..00000000000 --- a/.ci/sai/mpirun_with_mapping.sh +++ /dev/null @@ -1,7 +0,0 @@ -#!/usr/bin/env bash - -set -euo pipefail -: "${SAI_SYSTEM_MPIRUN:?}" -: "${MAP_OPT:?}" - -exec "$SAI_SYSTEM_MPIRUN" --map-by "$MAP_OPT" "$@" diff --git a/.ci/sai/README.md b/.ci/slurm/README.md similarity index 67% rename from .ci/sai/README.md rename to .ci/slurm/README.md index 8ae246542da..e014cf1e303 100644 --- a/.ci/sai/README.md +++ b/.ci/slurm/README.md @@ -1,8 +1,8 @@ -# SAI GPU CI manual +# Remote GPU validation ## Purpose and scope -This workflow rebuilds ABACUS at one exact commit on the SAI GPU cluster and +This workflow rebuilds ABACUS at one exact commit on a remote GPU cluster and runs the matrix in `gpu-matrix.ini` through Slurm. The current matrix has 49 cases. It reports build, resource, and case states and saves the raw logs. @@ -11,56 +11,62 @@ client requires a committed checkout: it bundles the checked-out `HEAD` and checks that it equals `--source-sha`. Uncommitted files are not sent. **Warning:** approved candidate code is compiled and executed as the configured -SAI SSH user. Use an account and project root with the intended permissions. +remote SSH user. Use an account and project root with the intended permissions. + +## Reference deployment + +The maintained deployment currently runs on the SAI cluster. Its queue, +module, MPI mapping, and SSH host-key settings are kept in the configuration +files in this directory and can be adapted for another Slurm site. ## GitHub setup -Create these GitHub Environments for the `SAI GPU validation` workflow: +Create these GitHub Environments for the `GPU validation` workflow: -- `sai-ssh-scheduled`: no required reviewers. The daily schedule uses this +- `gpu-ci-scheduled`: no required reviewers. The daily schedule uses this environment. -- `sai-ssh-manual`: require the reviewers who should approve a remote run. +- `gpu-ci-manual`: require the reviewers who should approve a remote run. Manual dispatches and the pull-request comment trigger use this environment. In each environment, configure the SSH secret and connection variables: -- Secret: `SAI_SSH_PRIVATE_KEY` (the private key used by the runner). -- Variables: `SAI_PROJECT_ROOT`, `SAI_SSH_HOST`, `SAI_SSH_PORT`, and - `SAI_SSH_USER`. +- Secret: `REMOTE_SSH_PRIVATE_KEY` (the private key used by the runner). +- Variables: `REMOTE_PROJECT_ROOT`, `REMOTE_SSH_HOST`, `REMOTE_SSH_PORT`, and + `REMOTE_SSH_USER`. -`SAI_PROJECT_ROOT` should be an absolute path below the SAI user's home, for -example `/home///abacus_sai_gpu_ci`. A manual dispatch may provide +`REMOTE_PROJECT_ROOT` should be an absolute path below the remote user's home, for +example `/home///abacus_gpu_ci`. A manual dispatch may provide the optional `project_root` input to override this variable for that run. ## Triggers -- **Manual:** open Actions, choose `SAI GPU validation`, and select +- **Manual:** open Actions, choose `GPU validation`, and select **Run workflow**. `source_sha` is required and must be a 40-character commit SHA. The workflow must be dispatched from the repository default branch. - `project_root` is optional. The `sai-ssh-manual` approval is requested before + `project_root` is optional. The `gpu-ci-manual` approval is requested before the build job starts. - **Daily:** the schedule is `30 20 * * *` (20:30 UTC). It tests the default - branch of `deepmodeling/abacus-develop` and uses `sai-ssh-scheduled`. -- **Pull request:** comment exactly `/abacus-ci sai-gpu` on an open pull + branch of `deepmodeling/abacus-develop` and uses `gpu-ci-scheduled`. +- **Pull request:** comment exactly `/abacus-ci gpu` on an open pull request. The commenter needs Triage, Write, Maintain, or Admin permission. - The pull request head commit is tested after `sai-ssh-manual` approval. + The pull request head commit is tested after `gpu-ci-manual` approval. Other issue comments are ignored. ## Run locally Run from the repository checkout whose committed `HEAD` is the revision to -test. If the SSH config defines the default `sai-ci` target, no options are +test. If the SSH config defines the default `gpu-ci` target, no options are required: ```bash -python3 .ci/sai/sai.py run +python3 .ci/slurm/runner.py run ``` -By default, the command uses `~/.ssh/config`, the `sai-ci` host alias, the -current Git checkout and its `HEAD`, and `~/abacus_sai_gpu_ci` below the remote -user's home. It writes downloaded results to `./sai-artifacts`. Override any -value shown by `run --help`; for example, use `--target SAI-abacus` for a +By default, the command uses `~/.ssh/config`, the `gpu-ci` host alias, the +current Git checkout and its `HEAD`, and `~/abacus_gpu_ci` below the remote +user's home. It writes downloaded results to `./gpu-ci-artifacts`. Override any +value shown by `run --help`; for example, use `--target my-cluster` for a different local host alias. The command waits for Slurm completion and exits with a non-zero status when a case or infrastructure result is not fully successful. @@ -73,8 +79,8 @@ the same SHA sends no source data. For the public commands and all `run` options, use: ```bash -python3 .ci/sai/sai.py --help -python3 .ci/sai/sai.py run --help +python3 .ci/slurm/runner.py --help +python3 .ci/slurm/runner.py run --help ``` ## Matrix configuration @@ -100,10 +106,10 @@ shown as `N nodes / M GPUs`. Thus `gpu1`, `gpu2`, and `gpu4` display `1 GPU`, ## Results and retention -On SAI, a run is created below: +On the remote cluster, a run is created below: ``` -/runs//-/ +/runs//-/ ``` Its `results/` directory contains `result.json`, `summary.md`, build and case @@ -112,25 +118,25 @@ working data are alongside it while the run is active. After results are collected, CI archives `results/` and `jobs/` as: ``` -/archives//-.tar.gz +/archives//-.tar.gz ``` The client removes archived files older than 72 hours when preparing a later run, and removes the active run after archiving. On the GitHub runner, -`ARTIFACT_ROOT` is `${runner.temp}/sai-artifacts`; it contains the collected +`ARTIFACT_ROOT` is `${runner.temp}/gpu-ci-artifacts`; it contains the collected `results/`, `run.json`, and `client.log`. CI uploads that directory as -`sai-gpu--` and retains it for 30 days. A pull-request comment +`gpu-validation--` and retains it for 30 days. A pull-request comment links to the Actions run and the uploaded raw files. ## Troubleshooting -**SSH fails.** Check `SAI_SSH_HOST`, `SAI_SSH_PORT`, and `SAI_SSH_USER`, that +**SSH fails.** Check `REMOTE_SSH_HOST`, `REMOTE_SSH_PORT`, and `REMOTE_SSH_USER`, that the key matches the account, and that the target is reachable. CI uses the -committed `.ci/sai/known_hosts` with strict host-key checking. Test the same +committed `.ci/slurm/known_hosts` with strict host-key checking. Test the same target with the SSH config before retrying; do not disable host-key checking. **A module cannot be loaded.** `modules.sh` sources Lmod, purges modules, and -loads `cmake/3.31.6` and the configured ABACUS dependency module. Ask the SAI +loads `cmake/3.31.6` and the configured ABACUS dependency module. Ask the site administrator to provide or update those modules. Modules provide the compiler, CUDA, MPI, and library dependencies; do not add library paths to CI (`LD_LIBRARY_PATH`, `CPATH`, or `CMAKE_PREFIX_PATH`) or hard-code site paths. diff --git a/.ci/sai/build.sbatch.in b/.ci/slurm/build.sbatch.in similarity index 100% rename from .ci/sai/build.sbatch.in rename to .ci/slurm/build.sbatch.in diff --git a/.ci/sai/case.sbatch.in b/.ci/slurm/case.sbatch.in similarity index 84% rename from .ci/sai/case.sbatch.in rename to .ci/slurm/case.sbatch.in index b1a0735897c..91e22fb51d0 100644 --- a/.ci/sai/case.sbatch.in +++ b/.ci/slurm/case.sbatch.in @@ -23,22 +23,22 @@ export OMP_NUM_THREADS=1 # shellcheck source=/dev/null source @MAPPING_ROOT@/${SLURM_JOB_PARTITION}.bash source @CONTROL@/modules.sh -SAI_SYSTEM_MPIRUN=$(command -v mpirun) -export SAI_SYSTEM_MPIRUN +CI_SYSTEM_MPIRUN=$(command -v mpirun) +export CI_SYSTEM_MPIRUN export MAP_OPT export SLURM_EXPORT_ENV=ALL export OMPI_MCA_plm_slurm_args=--external-launcher export PRTE_MCA_plm_slurm_args=--external-launcher nvidia-smi dmon -s pucvmte -o T > @RESULTS@/dmon-${SLURM_JOB_ID}.log & -SAI_DISABLE_NCCL_IB=@DISABLE_NCCL_IB@ -if [[ $SAI_DISABLE_NCCL_IB == true ]]; then +CI_DISABLE_NCCL_IB=@DISABLE_NCCL_IB@ +if [[ $CI_DISABLE_NCCL_IB == true ]]; then export NCCL_IB_DISABLE=1 else unset NCCL_IB_DISABLE fi -python3 @CONTROL@/sai.py worker \ +python3 @CONTROL@/runner.py worker \ @SOURCE@ @CONTROL@ @INSTALL@ @RESULTS@ @MANIFEST@ exit $? diff --git a/.ci/sai/gpu-matrix.ini b/.ci/slurm/gpu-matrix.ini similarity index 100% rename from .ci/sai/gpu-matrix.ini rename to .ci/slurm/gpu-matrix.ini diff --git a/.ci/sai/known_hosts b/.ci/slurm/known_hosts similarity index 100% rename from .ci/sai/known_hosts rename to .ci/slurm/known_hosts diff --git a/.ci/sai/modules.sh b/.ci/slurm/modules.sh similarity index 100% rename from .ci/sai/modules.sh rename to .ci/slurm/modules.sh diff --git a/.ci/slurm/mpirun_with_mapping.sh b/.ci/slurm/mpirun_with_mapping.sh new file mode 100755 index 00000000000..d8f930b7e33 --- /dev/null +++ b/.ci/slurm/mpirun_with_mapping.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash + +set -euo pipefail +: "${CI_SYSTEM_MPIRUN:?}" +: "${MAP_OPT:?}" + +exec "$CI_SYSTEM_MPIRUN" --map-by "$MAP_OPT" "$@" diff --git a/.ci/sai/sai.py b/.ci/slurm/runner.py similarity index 97% rename from .ci/sai/sai.py rename to .ci/slurm/runner.py index a8096036386..a025dc916f8 100644 --- a/.ci/sai/sai.py +++ b/.ci/slurm/runner.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Run the ABACUS GPU matrix on SAI from GitHub or a workstation.""" +"""Run the ABACUS GPU matrix on a remote Slurm cluster.""" import argparse import configparser @@ -231,7 +231,7 @@ def _save_result(run: Path, components: Sequence[Mapping[str, Any]], rows: Seque def _result_markdown(result: Mapping[str, Any]) -> str: lines = [ - "# SAI GPU result", "", + "# GPU validation result", "", "Passed: **{}**; failed: **{}**; infrastructure: **{}**".format( result["passed"], result["failed"], result["infrastructure"] ), "", "| Component | State | Slurm job |", "| --- | --- | --- |", @@ -561,7 +561,7 @@ def remote_receive(project: Path, run: Path, source_sha: str, bundle_checksum: s _command(("git", "--git-dir", str(cache), "bundle", "verify", str(bundle))) _command(( "git", "--git-dir", str(cache), "fetch", str(bundle), - "HEAD:refs/sai/{}".format(source_sha), + "HEAD:refs/ci/{}".format(source_sha), )) else: if any(input_root.iterdir()): @@ -754,7 +754,7 @@ def run(args: argparse.Namespace) -> int: "{}:{}/".format(args.target, remote_control), )) prepared = _ssh(args.ssh_config, args.target, ( - "python3", remote_control + "/sai.py", "prepare", args.project_root, str(run), + "python3", remote_control + "/runner.py", "prepare", args.project_root, str(run), )) data = json.loads(prepared.stdout) base = data.get("cache_sha", "") @@ -778,11 +778,11 @@ def upload(item: Tuple[int, Path]) -> None: with ThreadPoolExecutor(max_workers=BUNDLE_PARTS) as pool: list(pool.map(upload, enumerate(parts))) _ssh(args.ssh_config, args.target, ( - "python3", remote_control + "/sai.py", "receive", + "python3", remote_control + "/runner.py", "receive", args.project_root, str(run), args.source_sha, checksum, )) launch = "nohup python3 {} remote-run {} > {}/results/coordinator.log 2>&1 < /dev/null &".format( - shlex.quote(remote_control + "/sai.py"), shlex.quote(str(run)), shlex.quote(str(run)) + shlex.quote(remote_control + "/runner.py"), shlex.quote(str(run)), shlex.quote(str(run)) ) _ssh(args.ssh_config, args.target, ("bash", "-lc", launch)) failures = 0 @@ -793,7 +793,7 @@ def upload(item: Tuple[int, Path]) -> None: if status.returncode: failures += 1 if failures == 10: - raise RuntimeError("lost contact with SAI") + raise RuntimeError("lost contact with the remote cluster") elif status.stdout.strip(): done = json.loads(status.stdout) if type(done) is not dict or set(done) != {"returncode"} or \ @@ -805,7 +805,7 @@ def upload(item: Tuple[int, Path]) -> None: time.sleep(30) command = ( "ssh", "-F", str(args.ssh_config), args.target, - shlex.join(("python3", remote_control + "/sai.py", "collect", str(run))), + shlex.join(("python3", remote_control + "/runner.py", "collect", str(run))), ) archive = args.artifacts / ".results.tar.gz" _retry_download(command, archive) @@ -853,7 +853,7 @@ def github_admit() -> int: repository = os.environ["GITHUB_REPOSITORY"] values = {"accepted": "true", "pr_number": "", "check_id": ""} if event == "schedule": - upstream = os.environ.get("SAI_UPSTREAM_REPOSITORY", "deepmodeling/abacus-develop") + upstream = os.environ.get("UPSTREAM_REPOSITORY", "deepmodeling/abacus-develop") metadata = _gh("repos/{}".format(upstream)) commit = _gh("repos/{}/commits/{}".format(upstream, metadata["default_branch"])) values.update(source_repository=upstream, source_sha=commit["sha"], namespace="daily") @@ -878,7 +878,7 @@ def github_admit() -> int: pr_number=number, ) check = _gh("repos/{}/check-runs".format(repository), "POST", { - "name": "SAI GPU validation", "head_sha": pull["head"]["sha"], + "name": "GPU validation", "head_sha": pull["head"]["sha"], "status": "in_progress", }) values["check_id"] = str(check["id"]) @@ -892,7 +892,7 @@ def github_admit() -> int: def github_finish() -> int: repository = os.environ["GITHUB_REPOSITORY"] - result = os.environ["SAI_RESULT"] + result = os.environ["GPU_RESULT"] conclusion = "success" if result == "success" else "failure" check_id = os.environ["CHECK_ID"] if check_id: @@ -902,7 +902,7 @@ def github_finish() -> int: pr = os.environ["PR_NUMBER"] if pr: body = ( - "## SAI GPU validation: {}\n\n" + "## GPU validation: {}\n\n" "GPU cases: **{} passed, {} failed, {} infrastructure**.\n\n" "[Open the Actions run]({}) | [Download raw test files]({})\n\n" "Candidate: `{}`" @@ -921,8 +921,8 @@ def parser() -> argparse.ArgumentParser: title="commands", dest="command", required=True, metavar="{run,config}", ) client = commands.add_parser( - "run", help="build and run the GPU matrix on SAI", - description="Transfer one committed ABACUS revision, build it, and run the GPU matrix on SAI.", + "run", help="build and run the GPU matrix on a remote cluster", + description="Transfer one committed ABACUS revision, build it, and run the GPU matrix through Slurm.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) client.add_argument( @@ -930,12 +930,12 @@ def parser() -> argparse.ArgumentParser: help="SSH config file containing the target host alias", ) client.add_argument( - "--target", default="sai-ci", + "--target", default="gpu-ci", help="host alias in the SSH config", ) client.add_argument( - "--project-root", default="~/abacus_sai_gpu_ci", - help="SAI directory for caches, runs, and archives", + "--project-root", default="~/abacus_gpu_ci", + help="remote directory for caches, runs, and archives", ) client.add_argument( "--source-repository", type=Path, default=Path("."), @@ -958,7 +958,7 @@ def parser() -> argparse.ArgumentParser: help="attempt number within the run ID", ) client.add_argument( - "--artifacts", type=Path, default=Path("sai-artifacts"), + "--artifacts", type=Path, default=Path("gpu-ci-artifacts"), help="local directory for downloaded results and client logs", ) commands.add_parser( @@ -1031,5 +1031,5 @@ def main() -> int: try: sys.exit(main()) except (OSError, RuntimeError, ValueError, SlurmError) as error: - print("sai-ci: {}".format(error), file=sys.stderr) + print("gpu-ci: {}".format(error), file=sys.stderr) sys.exit(2) diff --git a/.ci/sai/slurm.py b/.ci/slurm/slurm.py similarity index 98% rename from .ci/sai/slurm.py rename to .ci/slurm/slurm.py index b38d900e159..071cd14fe5c 100644 --- a/.ci/sai/slurm.py +++ b/.ci/slurm/slurm.py @@ -1,4 +1,4 @@ -"""Small Slurm adapter for the SAI GPU matrix.""" +"""Small Slurm adapter for the GPU matrix.""" import re import subprocess diff --git a/.ci/sai/test_sai.py b/.ci/slurm/test_runner.py similarity index 76% rename from .ci/sai/test_sai.py rename to .ci/slurm/test_runner.py index c19aa9943e3..85ce19c2e3a 100644 --- a/.ci/sai/test_sai.py +++ b/.ci/slurm/test_runner.py @@ -13,20 +13,20 @@ ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT)) -SPEC = importlib.util.spec_from_file_location("sai", ROOT / "sai.py") +SPEC = importlib.util.spec_from_file_location("runner", ROOT / "runner.py") assert SPEC and SPEC.loader -sai = importlib.util.module_from_spec(SPEC) -sys.modules["sai"] = sai -SPEC.loader.exec_module(sai) +runner = importlib.util.module_from_spec(SPEC) +sys.modules["runner"] = runner +SPEC.loader.exec_module(runner) import slurm # noqa: E402 def valid_result(): - config = sai.load_config() - components = [sai._component("build", "Compile", "PASS")] - components.extend(sai._component(name, profile.label, "PASS") for name, profile in config.resources.items()) + config = runner.load_config() + components = [runner._component("build", "Compile", "PASS")] + components.extend(runner._component(name, profile.label, "PASS") for name, profile in config.resources.items()) rows = [ - sai._result_row(case, "PASS", exit_code=0) + runner._result_row(case, "PASS", exit_code=0) for resource in config.resources for case in config.cases if case.resource == resource ] @@ -39,7 +39,7 @@ def valid_result(): class ConfigTests(unittest.TestCase): def test_current_matrix_is_loaded_from_ini(self): - config = sai.load_config() + config = runner.load_config() self.assertEqual(len(config.cases), 49) self.assertEqual(list(config.resources), ["gpu1", "gpu2", "gpu4", "gpu8x2"]) self.assertEqual(config.resources["gpu4"].label, "4 GPUs") @@ -53,7 +53,7 @@ def test_resource_names_are_not_hardcoded(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "matrix.ini" path.write_text(text, encoding="utf-8") - self.assertIn("single", sai.load_config(path).resources) + self.assertIn("single", runner.load_config(path).resources) def test_unknown_resource_and_noncontiguous_case_fail(self): original = (ROOT / "gpu-matrix.ini").read_text(encoding="utf-8") @@ -65,23 +65,23 @@ def test_unknown_resource_and_noncontiguous_case_fail(self): path = Path(directory) / "matrix.ini" path.write_text(text, encoding="utf-8") with self.assertRaises(ValueError): - sai.load_config(path) + runner.load_config(path) class CliTests(unittest.TestCase): def test_top_level_help_only_exposes_public_commands(self): help_text = subprocess.run( - (sys.executable, str(ROOT / "sai.py"), "--help"), + (sys.executable, str(ROOT / "runner.py"), "--help"), check=True, text=True, capture_output=True, ).stdout self.assertIn("{run,config}", help_text) - self.assertIn("build and run the GPU matrix on SAI", help_text) + self.assertIn("build and run the GPU matrix on a remote cluster", help_text) self.assertNotIn("remote-run", help_text) self.assertNotIn("github-admit", help_text) def test_run_help_describes_every_option(self): help_text = subprocess.run( - (sys.executable, str(ROOT / "sai.py"), "run", "--help"), + (sys.executable, str(ROOT / "runner.py"), "run", "--help"), check=True, text=True, capture_output=True, ).stdout for option in ( @@ -91,20 +91,20 @@ def test_run_help_describes_every_option(self): self.assertIn(option, help_text) normalized = " ".join(help_text.split()) for default in ( - "~/.ssh/config", "sai-ci", "~/abacus_sai_gpu_ci", "HEAD", "manual", - "sai-artifacts", + "~/.ssh/config", "gpu-ci", "~/abacus_gpu_ci", "HEAD", "manual", + "gpu-ci-artifacts", ): self.assertIn("default: {}".format(default), normalized) def test_run_options_parse(self): arguments = [ - "run", "--ssh-config", "/tmp/ssh/config", "--target", "sai-ci", - "--project-root", "/home/user/abacus_sai_gpu_ci", + "run", "--ssh-config", "/tmp/ssh/config", "--target", "gpu-ci", + "--project-root", "/home/user/abacus_gpu_ci", "--source-repository", "/tmp/abacus", "--source-sha", "a" * 40, "--namespace", "manual", "--run-id", "42", "--run-attempt", "1", "--artifacts", "/tmp/results", ] - args = sai.parser().parse_args(arguments) + args = runner.parser().parse_args(arguments) self.assertEqual(args.command, "run") self.assertEqual(args.ssh_config, Path("/tmp/ssh/config")) self.assertEqual(args.source_repository, Path("/tmp/abacus")) @@ -112,26 +112,26 @@ def test_run_options_parse(self): self.assertEqual(args.source_sha, "a" * 40) def test_run_defaults_parse(self): - with mock.patch("sai.time.time", return_value=42): - args = sai.parser().parse_args(["run"]) + with mock.patch("runner.time.time", return_value=42): + args = runner.parser().parse_args(["run"]) self.assertEqual(args.ssh_config, Path("~/.ssh/config")) - self.assertEqual(args.target, "sai-ci") - self.assertEqual(args.project_root, "~/abacus_sai_gpu_ci") + self.assertEqual(args.target, "gpu-ci") + self.assertEqual(args.project_root, "~/abacus_gpu_ci") self.assertEqual(args.source_repository, Path(".")) self.assertEqual(args.source_sha, "HEAD") self.assertEqual(args.namespace, "manual") self.assertEqual(args.run_id, "42") self.assertEqual(args.run_attempt, "1") - self.assertEqual(args.artifacts, Path("sai-artifacts")) + self.assertEqual(args.artifacts, Path("gpu-ci-artifacts")) class TemplateTests(unittest.TestCase): def test_resource_template_is_complete_and_has_no_cpu_request(self): - config = sai.load_config() + config = runner.load_config() with tempfile.TemporaryDirectory() as directory: run = Path(directory) / "runs" / "manual" / "1-1" destination = run / "jobs" / "gpu4.sbatch" - values = sai._job_values( + values = runner._job_values( config.resources["gpu4"], config, run, "gpu4", run / "results" / "gpu4-%A_%a.out", ) @@ -140,7 +140,7 @@ def test_resource_template_is_complete_and_has_no_cpu_request(self): "MAPPING_ROOT": config.mapping_root, "DISABLE_NCCL_IB": "false", "MANIFEST": run / "jobs" / "gpu4.tsv", }) - sai._render(ROOT / "case.sbatch.in", destination, values) + runner._render(ROOT / "case.sbatch.in", destination, values) text = destination.read_text(encoding="utf-8") self.assertIn("#SBATCH --nodes=1", text) self.assertIn("#SBATCH --array=0-3%4", text) @@ -179,11 +179,11 @@ def test_submit_and_accounting_require_each_array_task(self): self.assertEqual(states["101_1"], ("FAILED", "1:0")) def test_pass_requires_successful_slurm_accounting(self): - config = sai.Config( - "gpu", Path("/opt/sai_config/mps_mapping.d"), False, 1, - sai.Resource("build", "flood-gpu", 1, 1, 1, 60), - {"one": sai.Resource("one", "flood-gpu", 1, 1, 1, 60)}, - (sai.Case("suite", "case", "one", "autotest"),), + config = runner.Config( + "gpu", Path("/opt/cluster/mps_mapping.d"), False, 1, + runner.Resource("build", "flood-gpu", 1, 1, 1, 60), + {"one": runner.Resource("one", "flood-gpu", 1, 1, 1, 60)}, + (runner.Case("suite", "case", "one", "autotest"),), ) client = mock.Mock() client.submit.side_effect = ["100", "101"] @@ -197,10 +197,10 @@ def test_pass_requires_successful_slurm_accounting(self): status = run / "results" / "status" / "suite__case.json" status.parent.mkdir(parents=True) status.write_text(json.dumps({"state": "PASS", "exit_code": 0}), encoding="utf-8") - with mock.patch("sai.load_config", return_value=config), \ - mock.patch("sai.Slurm", return_value=client), \ - mock.patch("sai._render"): - self.assertEqual(sai.remote_run(run), 1) + with mock.patch("runner.load_config", return_value=config), \ + mock.patch("runner.Slurm", return_value=client), \ + mock.patch("runner._render"): + self.assertEqual(runner.remote_run(run), 1) row = json.loads((run / "results" / "result.json").read_text())["cases"][0] self.assertEqual(row["state"], "INFRA") self.assertEqual(row["slurm_exit_code"], "1:0") @@ -214,7 +214,7 @@ def _git(root, *arguments): def _bundle(self, repository, run, revision): bundle = run / "source.bundle.local" self._git(repository, "bundle", "create", str(bundle), revision) - parts, checksum = sai._split_bundle(bundle, run / "input") + parts, checksum = runner._split_bundle(bundle, run / "input") bundle.unlink() self.assertEqual(len(parts), 8) return checksum @@ -235,10 +235,10 @@ def test_full_then_incremental_bundle(self): project = home / "project" run1 = project / "runs" / "manual" / "1-1" (run1 / "control").mkdir(parents=True) - with mock.patch("sai.Path.home", return_value=home): - sai.remote_prepare(project, run1) + with mock.patch("runner.Path.home", return_value=home): + runner.remote_prepare(project, run1) checksum = self._bundle(repository, run1, "HEAD") - sai.remote_receive(project, run1, first, checksum) + runner.remote_receive(project, run1, first, checksum) self.assertEqual((run1 / "source" / "value.txt").read_text(), "one\n") (repository / "value.txt").write_text("two\n", encoding="utf-8") @@ -246,18 +246,18 @@ def test_full_then_incremental_bundle(self): second = self._git(repository, "rev-parse", "HEAD").stdout.strip() run2 = project / "runs" / "manual" / "2-1" (run2 / "control").mkdir(parents=True) - with mock.patch("sai.Path.home", return_value=home): - sai.remote_prepare(project, run2) + with mock.patch("runner.Path.home", return_value=home): + runner.remote_prepare(project, run2) checksum = self._bundle(repository, run2, first + "..HEAD") - sai.remote_receive(project, run2, second, checksum) + runner.remote_receive(project, run2, second, checksum) self.assertEqual((run2 / "source" / "value.txt").read_text(), "two\n") run3 = project / "runs" / "manual" / "3-1" (run3 / "control").mkdir(parents=True) - with mock.patch("sai.Path.home", return_value=home): - sai.remote_prepare(project, run3) - self.assertIsNone(sai._bundle_revision(repository, second, second)) - sai.remote_receive(project, run3, second, "-") + with mock.patch("runner.Path.home", return_value=home): + runner.remote_prepare(project, run3) + self.assertIsNone(runner._bundle_revision(repository, second, second)) + runner.remote_receive(project, run3, second, "-") self.assertEqual((run3 / "source" / "value.txt").read_text(), "two\n") def test_bundle_merge_rejects_corruption(self): @@ -265,10 +265,10 @@ def test_bundle_merge_rejects_corruption(self): root = Path(directory) bundle = root / "bundle" bundle.write_bytes(bytes(range(256)) * 4) - parts, checksum = sai._split_bundle(bundle, root) + parts, checksum = runner._split_bundle(bundle, root) parts[0].write_bytes(b"corrupt") with self.assertRaisesRegex(ValueError, "checksum mismatch"): - sai._assemble_bundle(root, checksum) + runner._assemble_bundle(root, checksum) def test_prepare_rejects_reused_run(self): with tempfile.TemporaryDirectory() as directory: @@ -276,17 +276,17 @@ def test_prepare_rejects_reused_run(self): project = home / "project" run = project / "runs" / "manual" / "1-1" (run / "control").mkdir(parents=True) - with mock.patch("sai.Path.home", return_value=home): - sai.remote_prepare(project, run) + with mock.patch("runner.Path.home", return_value=home): + runner.remote_prepare(project, run) with self.assertRaisesRegex(ValueError, "already exists"): - sai.remote_prepare(project, run) + runner.remote_prepare(project, run) def test_transfer_retry_is_bounded(self): completed = mock.Mock(returncode=0, stdout="done", stderr="") failed = mock.Mock(returncode=1, stdout="", stderr="disconnected") - with mock.patch("sai.subprocess.run", side_effect=[failed, failed, completed]) as run, \ - mock.patch("sai.time.sleep"): - self.assertEqual(sai._retry(("rsync", "source", "target")).stdout, "done") + with mock.patch("runner.subprocess.run", side_effect=[failed, failed, completed]) as run, \ + mock.patch("runner.time.sleep"): + self.assertEqual(runner._retry(("rsync", "source", "target")).stdout, "done") self.assertEqual(run.call_count, 3) def test_download_retry_replaces_partial_file(self): @@ -300,8 +300,8 @@ def download(_command, _cwd=None, stdout=None): raise RuntimeError("connection closed") download.calls = 0 - with mock.patch("sai._command", side_effect=download), mock.patch("sai.time.sleep"): - sai._retry_download(("ssh", "collect"), path) + with mock.patch("runner._command", side_effect=download), mock.patch("runner.time.sleep"): + runner._retry_download(("ssh", "collect"), path) self.assertEqual(path.read_bytes(), b"complete") def test_run_records_cleanup_metadata_before_remote_work(self): @@ -311,7 +311,7 @@ def test_run_records_cleanup_metadata_before_remote_work(self): args = argparse.Namespace( source_repository=root, project_root="/project", - target="sai-ci", + target="gpu-ci", namespace="manual", run_id="1", run_attempt="1", @@ -320,10 +320,10 @@ def test_run_records_cleanup_metadata_before_remote_work(self): ssh_config=root / "ssh-config", ) revision = mock.Mock(stdout=source_sha + "\n") - with mock.patch("sai._command", return_value=revision), \ - mock.patch("sai._retry", side_effect=RuntimeError("disconnected")), \ + with mock.patch("runner._command", return_value=revision), \ + mock.patch("runner._retry", side_effect=RuntimeError("disconnected")), \ self.assertRaisesRegex(RuntimeError, "disconnected"): - sai.run(args) + runner.run(args) metadata = json.loads((args.artifacts / "run.json").read_text()) self.assertEqual(metadata, { "project_root": "/project", @@ -337,8 +337,8 @@ def test_run_resolves_head_and_remote_home_defaults(self): source_sha = "a" * 40 args = argparse.Namespace( source_repository=root, - project_root="~/abacus_sai_gpu_ci", - target="sai-ci", + project_root="~/abacus_gpu_ci", + target="gpu-ci", namespace="manual", run_id="1", run_attempt="1", @@ -348,20 +348,20 @@ def test_run_resolves_head_and_remote_home_defaults(self): ) revision = mock.Mock(stdout=source_sha + "\n") home = mock.Mock(stdout="/home/user\n") - with mock.patch("sai.Path.home", return_value=Path("/home/local")), \ - mock.patch("sai._command", return_value=revision), \ - mock.patch("sai._retry", side_effect=[home, RuntimeError("disconnected")]), \ + with mock.patch("runner.Path.home", return_value=Path("/home/local")), \ + mock.patch("runner._command", return_value=revision), \ + mock.patch("runner._retry", side_effect=[home, RuntimeError("disconnected")]), \ self.assertRaisesRegex(RuntimeError, "disconnected"): - sai.run(args) + runner.run(args) metadata = json.loads((args.artifacts / "run.json").read_text()) self.assertEqual(metadata, { - "project_root": "/home/user/abacus_sai_gpu_ci", - "run_root": "/home/user/abacus_sai_gpu_ci/runs/manual/1-1", + "project_root": "/home/user/abacus_gpu_ci", + "run_root": "/home/user/abacus_gpu_ci/runs/manual/1-1", "source_sha": source_sha, }) def test_result_archive_rejects_traversal_and_links(self): - for name, link in (("../sai-ssh/id_ed25519", None), ("results/key", "../key")): + for name, link in (("../runner-ssh/id_ed25519", None), ("results/key", "../key")): payload = io.BytesIO() with tarfile.open(fileobj=payload, mode="w:gz") as archive: member = tarfile.TarInfo(name) @@ -374,7 +374,7 @@ def test_result_archive_rejects_traversal_and_links(self): archive.addfile(member) payload.seek(0) with tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError): - sai._extract_results(payload, Path(directory)) + runner._extract_results(payload, Path(directory)) def test_result_archive_extracts_regular_results(self): payload = io.BytesIO() @@ -384,7 +384,7 @@ def test_result_archive_extracts_regular_results(self): archive.addfile(member, io.BytesIO(b"{}\n")) payload.seek(0) with tempfile.TemporaryDirectory() as directory: - sai._extract_results(payload, Path(directory)) + runner._extract_results(payload, Path(directory)) self.assertEqual((Path(directory) / "results" / "result.json").read_text(), "{}\n") def test_result_archive_rejects_oversized_content(self): @@ -395,9 +395,9 @@ def test_result_archive_rejects_oversized_content(self): archive.addfile(member, io.BytesIO(b"abc")) payload.seek(0) with tempfile.TemporaryDirectory() as directory, \ - mock.patch("sai.MAX_RESULT_BYTES", 2), \ + mock.patch("runner.MAX_RESULT_BYTES", 2), \ self.assertRaisesRegex(ValueError, "too large"): - sai._extract_results(payload, Path(directory)) + runner._extract_results(payload, Path(directory)) def test_result_archive_limits_directory_members(self): payload = io.BytesIO() @@ -408,9 +408,9 @@ def test_result_archive_limits_directory_members(self): archive.addfile(member) payload.seek(0) with tempfile.TemporaryDirectory() as directory, \ - mock.patch("sai.MAX_RESULT_MEMBERS", 1), \ + mock.patch("runner.MAX_RESULT_MEMBERS", 1), \ self.assertRaisesRegex(ValueError, "too large"): - sai._extract_results(payload, Path(directory)) + runner._extract_results(payload, Path(directory)) class ResultTests(unittest.TestCase): @@ -421,11 +421,11 @@ def test_report_publishes_dynamic_components(self): path = root / "result.json" path.write_text(json.dumps(result), encoding="utf-8") args = argparse.Namespace(result=path, output=root / "output", summary=root / "summary") - sai.report(args) + runner.report(args) output = args.output.read_text(encoding="utf-8") self.assertIn("available=true", output) self.assertIn('"name":"gpu8x2"', output) - self.assertTrue(args.summary.read_text().startswith("# SAI GPU result\n")) + self.assertTrue(args.summary.read_text().startswith("# GPU validation result\n")) def test_report_rejects_untrusted_counts(self): invalid = ( @@ -443,7 +443,7 @@ def test_report_rejects_untrusted_counts(self): path.write_text(json.dumps(result), encoding="utf-8") args = argparse.Namespace(result=path, output=root / "output", summary=None) with self.assertRaisesRegex(ValueError, "result counts"): - sai.report(args) + runner.report(args) self.assertFalse(args.output.exists()) def test_report_rejects_wrong_matrix_identity(self): @@ -455,29 +455,29 @@ def test_report_rejects_wrong_matrix_identity(self): path.write_text(json.dumps(result), encoding="utf-8") args = argparse.Namespace(result=path, output=root / "output", summary=None) with self.assertRaisesRegex(ValueError, "result case"): - sai.report(args) + runner.report(args) def test_mpi_startup_failure_requires_complete_signature(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "log" path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE MPI_Init_thread PMIx_Init failed") - self.assertTrue(sai._mpi_startup_failure(path)) + self.assertTrue(runner._mpi_startup_failure(path)) path.write_bytes(b"PMIX_ERR_FILE_OPEN_FAILURE") - self.assertFalse(sai._mpi_startup_failure(path)) + self.assertFalse(runner._mpi_startup_failure(path)) path.write_bytes(b"srun returned non-zero exit status (512) from launching the per-node daemon") - self.assertTrue(sai._mpi_startup_failure(path)) + self.assertTrue(runner._mpi_startup_failure(path)) path.write_bytes(b"srun returned non-zero exit status (512)") - self.assertFalse(sai._mpi_startup_failure(path)) + self.assertFalse(runner._mpi_startup_failure(path)) class PolicyTests(unittest.TestCase): def test_workflow_uses_trusted_control_and_protected_environment(self): - text = (ROOT.parents[1] / ".github" / "workflows" / "sai-gpu.yml").read_text(encoding="utf-8") + text = (ROOT.parents[1] / ".github" / "workflows" / "gpu-validation.yml").read_text(encoding="utf-8") self.assertIn("ref: ${{ env.CONTROL_SHA }}", text) - self.assertIn("sai-ssh-scheduled", text) - self.assertIn("sai-ssh-manual", text) - self.assertIn("/abacus-ci sai-gpu", text) - self.assertIn("sai.py run", text) + self.assertIn("gpu-ci-scheduled", text) + self.assertIn("gpu-ci-manual", text) + self.assertIn("/abacus-ci gpu", text) + self.assertIn("runner.py run", text) self.assertIn("pull-requests: write", text) self.assertNotIn("cpus-per-task", text) diff --git a/.github/workflows/sai-gpu.yml b/.github/workflows/gpu-validation.yml similarity index 78% rename from .github/workflows/sai-gpu.yml rename to .github/workflows/gpu-validation.yml index 13f4df46ee2..04f41338d62 100644 --- a/.github/workflows/sai-gpu.yml +++ b/.github/workflows/gpu-validation.yml @@ -1,4 +1,4 @@ -name: SAI GPU validation +name: GPU validation on: issue_comment: @@ -6,11 +6,11 @@ on: workflow_dispatch: inputs: source_sha: - description: "Approved commit to execute on SAI" + description: "Approved commit to execute on the remote GPU cluster" required: true type: string project_root: - description: "Optional SAI project root" + description: "Optional remote project root" required: false type: string schedule: @@ -19,7 +19,7 @@ on: permissions: {} concurrency: - group: sai-gpu-${{ github.event_name == 'schedule' && 'daily' || github.run_id }} + group: gpu-validation-${{ github.event_name == 'schedule' && 'daily' || github.run_id }} cancel-in-progress: false jobs: @@ -28,7 +28,7 @@ jobs: if: >- github.event_name != 'issue_comment' || (github.event.issue.pull_request && - github.event.comment.body == '/abacus-ci sai-gpu') + github.event.comment.body == '/abacus-ci gpu') runs-on: ubuntu-24.04 permissions: checks: write @@ -65,17 +65,17 @@ jobs: env: GH_TOKEN: ${{ github.token }} MANUAL_SOURCE_SHA: ${{ inputs.source_sha || '' }} - SAI_UPSTREAM_REPOSITORY: deepmodeling/abacus-develop - run: python3 control/.ci/sai/sai.py github-admit + UPSTREAM_REPOSITORY: deepmodeling/abacus-develop + run: python3 control/.ci/slurm/runner.py github-admit rebuild-and-test: - name: Build and run on SAI + name: Build and run on GPU cluster needs: admit if: needs.admit.outputs.accepted == 'true' runs-on: ubuntu-24.04 timeout-minutes: 240 environment: - name: ${{ github.event_name == 'schedule' && 'sai-ssh-scheduled' || 'sai-ssh-manual' }} + name: ${{ github.event_name == 'schedule' && 'gpu-ci-scheduled' || 'gpu-ci-manual' }} permissions: contents: read outputs: @@ -90,10 +90,10 @@ jobs: CONTROL_SHA: ${{ needs.admit.outputs.control_sha }} PROJECT_ROOT_INPUT: ${{ inputs.project_root || '' }} RUN_NAMESPACE: ${{ needs.admit.outputs.namespace }} - SAI_PROJECT_ROOT: ${{ vars.SAI_PROJECT_ROOT }} - SAI_SSH_HOST: ${{ vars.SAI_SSH_HOST }} - SAI_SSH_PORT: ${{ vars.SAI_SSH_PORT }} - SAI_SSH_USER: ${{ vars.SAI_SSH_USER }} + REMOTE_PROJECT_ROOT: ${{ vars.REMOTE_PROJECT_ROOT }} + REMOTE_SSH_HOST: ${{ vars.REMOTE_SSH_HOST }} + REMOTE_SSH_PORT: ${{ vars.REMOTE_SSH_PORT }} + REMOTE_SSH_USER: ${{ vars.REMOTE_SSH_USER }} SOURCE_REPOSITORY: ${{ needs.admit.outputs.source_repository }} SOURCE_SHA: ${{ needs.admit.outputs.source_sha }} steps: @@ -115,20 +115,20 @@ jobs: - name: Configure SSH env: - SAI_SSH_PRIVATE_KEY: ${{ secrets.SAI_SSH_PRIVATE_KEY }} + REMOTE_SSH_PRIVATE_KEY: ${{ secrets.REMOTE_SSH_PRIVATE_KEY }} run: | set -euo pipefail - root="$RUNNER_TEMP/sai-ssh" + root="$RUNNER_TEMP/remote-ssh" mkdir -m 700 "$root" - printf '%s\n' "$SAI_SSH_PRIVATE_KEY" > "$root/id_ed25519" + printf '%s\n' "$REMOTE_SSH_PRIVATE_KEY" > "$root/id_ed25519" chmod 600 "$root/id_ed25519" - cp control/.ci/sai/known_hosts "$root/known_hosts" + cp control/.ci/slurm/known_hosts "$root/known_hosts" chmod 600 "$root/known_hosts" cat > "$root/config" </dev/null - echo "SAI_SSH_CONFIG=$root/config" >> "$GITHUB_ENV" - echo "ARTIFACT_ROOT=$RUNNER_TEMP/sai-artifacts" >> "$GITHUB_ENV" + echo "REMOTE_SSH_CONFIG=$root/config" >> "$GITHUB_ENV" + echo "ARTIFACT_ROOT=$RUNNER_TEMP/gpu-ci-artifacts" >> "$GITHUB_ENV" - - name: Run shared SAI client + - name: Run shared GPU client id: client continue-on-error: true run: | set -euo pipefail - project_root=${PROJECT_ROOT_INPUT:-$SAI_PROJECT_ROOT} - : "${project_root:?Set SAI_PROJECT_ROOT in the GitHub Environment}" + project_root=${PROJECT_ROOT_INPUT:-$REMOTE_PROJECT_ROOT} + : "${project_root:?Set REMOTE_PROJECT_ROOT in the GitHub Environment}" mkdir -p "$ARTIFACT_ROOT" set +e - python3 control/.ci/sai/sai.py run \ - --ssh-config "$SAI_SSH_CONFIG" \ - --target sai-ci \ + python3 control/.ci/slurm/runner.py run \ + --ssh-config "$REMOTE_SSH_CONFIG" \ + --target gpu-ci \ --project-root "$project_root" \ --source-repository "$GITHUB_WORKSPACE/source" \ --source-sha "$SOURCE_SHA" \ @@ -171,7 +171,7 @@ jobs: id: summary if: always() run: | - python3 control/.ci/sai/sai.py report \ + python3 control/.ci/slurm/runner.py report \ --result "$ARTIFACT_ROOT/results/result.json" \ --output "$GITHUB_OUTPUT" \ --summary "$GITHUB_STEP_SUMMARY" @@ -181,12 +181,12 @@ jobs: if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 with: - name: sai-gpu-${{ github.run_id }}-${{ github.run_attempt }} + name: gpu-validation-${{ github.run_id }}-${{ github.run_attempt }} path: ${{ env.ARTIFACT_ROOT }} if-no-files-found: warn retention-days: 30 - - name: Archive uploaded run on SAI + - name: Archive uploaded remote run if: always() && steps.upload.outcome == 'success' run: | set -euo pipefail @@ -194,15 +194,15 @@ jobs: project=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["project_root"])' "$ARTIFACT_ROOT/run.json") run=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["run_root"])' "$ARTIFACT_ROOT/run.json") for attempt in 1 2 3; do - ssh -F "$SAI_SSH_CONFIG" sai-ci \ - "test ! -d '$run' || python3 '$run/control/sai.py' archive '$project' '$run'" && exit 0 + ssh -F "$REMOTE_SSH_CONFIG" gpu-ci \ + "test ! -d '$run' || python3 '$run/control/runner.py' archive '$project' '$run'" && exit 0 test "$attempt" = 3 || sleep 5 done exit 1 - name: Remove SSH credentials if: always() - run: rm -rf "$RUNNER_TEMP/sai-ssh" + run: rm -rf "$RUNNER_TEMP/remote-ssh" - name: Validate result protocol if: always() @@ -213,7 +213,7 @@ jobs: test "$AVAILABLE" = true component-status: - name: SAI / ${{ matrix.component.label }} + name: GPU / ${{ matrix.component.label }} needs: [admit, rebuild-and-test] if: always() && needs.admit.outputs.accepted == 'true' runs-on: ubuntu-24.04 @@ -255,6 +255,6 @@ jobs: GPU_PASSED: ${{ needs.rebuild-and-test.outputs.passed }} PR_NUMBER: ${{ needs.admit.outputs.pr_number }} RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} - SAI_RESULT: ${{ needs.rebuild-and-test.result == 'success' && needs.component-status.result == 'success' && 'success' || 'failure' }} + GPU_RESULT: ${{ needs.rebuild-and-test.result == 'success' && needs.component-status.result == 'success' && 'success' || 'failure' }} SOURCE_SHA: ${{ needs.admit.outputs.source_sha }} - run: python3 control/.ci/sai/sai.py github-finish + run: python3 control/.ci/slurm/runner.py github-finish From d6e5f914519927e9117bb7f40dd5bb81215d07a1 Mon Sep 17 00:00:00 2001 From: Jiacheng Xu <13862180016@163.com> Date: Tue, 28 Jul 2026 02:34:17 +0800 Subject: [PATCH 14/14] ci: keep remote connection settings in config --- .ci/slurm/README.md | 34 ++++++++-------- .ci/slurm/{gpu-matrix.ini => config.ini} | 6 +++ .ci/slurm/runner.py | 51 ++++++++++++++++++++---- .ci/slurm/test_runner.py | 30 ++++++++++++-- .github/workflows/gpu-validation.yml | 16 +++++--- 5 files changed, 104 insertions(+), 33 deletions(-) rename .ci/slurm/{gpu-matrix.ini => config.ini} (98%) diff --git a/.ci/slurm/README.md b/.ci/slurm/README.md index e014cf1e303..59554a9bda4 100644 --- a/.ci/slurm/README.md +++ b/.ci/slurm/README.md @@ -3,8 +3,9 @@ ## Purpose and scope This workflow rebuilds ABACUS at one exact commit on a remote GPU cluster and -runs the matrix in `gpu-matrix.ini` through Slurm. The current matrix has 49 -cases. It reports build, resource, and case states and saves the raw logs. +runs the matrix in `config.ini` through Slurm. The configuration contains the +remote connection, Slurm resources, and 49 test cases. The workflow reports +build, resource, and case states and saves the raw logs. This is validation, not a benchmark or a general remote-shell service. The client requires a committed checkout: it bundles the checked-out `HEAD` and @@ -28,15 +29,14 @@ Create these GitHub Environments for the `GPU validation` workflow: - `gpu-ci-manual`: require the reviewers who should approve a remote run. Manual dispatches and the pull-request comment trigger use this environment. -In each environment, configure the SSH secret and connection variables: +In each environment, configure one secret: - Secret: `REMOTE_SSH_PRIVATE_KEY` (the private key used by the runner). -- Variables: `REMOTE_PROJECT_ROOT`, `REMOTE_SSH_HOST`, `REMOTE_SSH_PORT`, and - `REMOTE_SSH_USER`. -`REMOTE_PROJECT_ROOT` should be an absolute path below the remote user's home, for -example `/home///abacus_gpu_ci`. A manual dispatch may provide -the optional `project_root` input to override this variable for that run. +The two environments have different approval rules, so the same key must be +added to each one. Host, port, user, and the default project root come from the +trusted `[remote]` section in `config.ini`. A manual dispatch may provide the +optional `project_root` input to override the configured root for that run. ## Triggers @@ -83,10 +83,12 @@ python3 .ci/slurm/runner.py --help python3 .ci/slurm/runner.py run --help ``` -## Matrix configuration +## Configuration -`gpu-matrix.ini` is validated before jobs are submitted. +`config.ini` is validated before jobs are submitted. +- `[remote]`: SSH `host`, `port`, `user`, and `project_root`. The project root + must be an absolute path or use `~/` relative to the remote user's home. - `[cluster]`: Slurm `partition`, absolute `mapping_root` for the MPI mapping script, `disable_nccl_ib` (`true` or `false`), and `poll_seconds` (1-300). - `[build]`: build-job `qos`, `nodes`, `tasks_per_node`, `gpus_per_node`, and @@ -109,7 +111,7 @@ shown as `N nodes / M GPUs`. Thus `gpu1`, `gpu2`, and `gpu4` display `1 GPU`, On the remote cluster, a run is created below: ``` -/runs//-/ +/runs//-/ ``` Its `results/` directory contains `result.json`, `summary.md`, build and case @@ -118,7 +120,7 @@ working data are alongside it while the run is active. After results are collected, CI archives `results/` and `jobs/` as: ``` -/archives//-.tar.gz +/archives//-.tar.gz ``` The client removes archived files older than 72 hours when preparing a later @@ -130,10 +132,10 @@ links to the Actions run and the uploaded raw files. ## Troubleshooting -**SSH fails.** Check `REMOTE_SSH_HOST`, `REMOTE_SSH_PORT`, and `REMOTE_SSH_USER`, that -the key matches the account, and that the target is reachable. CI uses the -committed `.ci/slurm/known_hosts` with strict host-key checking. Test the same -target with the SSH config before retrying; do not disable host-key checking. +**SSH fails.** Check the `[remote]` values in `config.ini`, that the key matches +the configured account, and that the target is reachable. CI uses the committed +`.ci/slurm/known_hosts` with strict host-key checking. Test the same target with +the SSH config before retrying; do not disable host-key checking. **A module cannot be loaded.** `modules.sh` sources Lmod, purges modules, and loads `cmake/3.31.6` and the configured ABACUS dependency module. Ask the site diff --git a/.ci/slurm/gpu-matrix.ini b/.ci/slurm/config.ini similarity index 98% rename from .ci/slurm/gpu-matrix.ini rename to .ci/slurm/config.ini index 80d654f11b2..3453eefff97 100644 --- a/.ci/slurm/gpu-matrix.ini +++ b/.ci/slurm/config.ini @@ -1,3 +1,9 @@ +[remote] +host = c0.sai.ai-4s.com +port = 12022 +user = abacususer01 +project_root = ~/abacus_gpu_ci + [cluster] partition = 16V100 mapping_root = /opt/sai_config/mps_mapping.d diff --git a/.ci/slurm/runner.py b/.ci/slurm/runner.py index a025dc916f8..338e5d26c1c 100644 --- a/.ci/slurm/runner.py +++ b/.ci/slurm/runner.py @@ -69,8 +69,17 @@ def case_id(self) -> str: return self.suite + "/" + self.name +@dataclass(frozen=True) +class Remote: + host: str + port: int + user: str + project_root: str + + @dataclass(frozen=True) class Config: + remote: Remote partition: str mapping_root: Path disable_nccl_ib: bool @@ -111,7 +120,7 @@ def _resource(name: str, section: Mapping[str, str], array: bool) -> Resource: return profile -def load_config(path: Path = ROOT / "gpu-matrix.ini") -> Config: +def load_config(path: Path = ROOT / "config.ini") -> Config: parser = configparser.ConfigParser(interpolation=None, strict=True) with Path(path).open(encoding="utf-8") as stream: parser.read_file(stream) @@ -120,11 +129,27 @@ def load_config(path: Path = ROOT / "gpu-matrix.ini") -> Config: expected_cases = ["case.{:03d}".format(index) for index in range(1, len(cases) + 1)] if not resources or cases != expected_cases: raise ValueError("resources and contiguous case sections are required") - known = {"cluster", "build", *resources, *cases} - if set(parser.sections()) != known or set(parser["cluster"]) != { + known = {"remote", "cluster", "build", *resources, *cases} + if set(parser.sections()) != known or set(parser["remote"]) != { + "host", "port", "user", "project_root", + } or set(parser["cluster"]) != { "partition", "mapping_root", "disable_nccl_ib", "poll_seconds", }: raise ValueError("unexpected configuration section or key") + remote = parser["remote"] + project_root = remote["project_root"] + project = PurePosixPath(project_root) + if project_root == "~": + project_parts: Tuple[str, ...] = () + elif project_root.startswith("~/"): + project_parts = project.parts[1:] + elif project.is_absolute(): + project_parts = project.parts[1:] + else: + raise ValueError("invalid remote project root") + if not NAME.fullmatch(remote["host"]) or not NAME.fullmatch(remote["user"]) or \ + any(not NAME.fullmatch(part) for part in project_parts): + raise ValueError("invalid remote configuration") partition = parser["cluster"]["partition"] mapping = Path(parser["cluster"]["mapping_root"]) disable = parser["cluster"]["disable_nccl_ib"].lower() @@ -150,6 +175,10 @@ def load_config(path: Path = ROOT / "gpu-matrix.ini") -> Config: if any(not any(case.resource == name for case in matrix) for name in profiles): raise ValueError("every resource needs a case") return Config( + Remote( + remote["host"], _integer(remote, "port", 1, 65535), + remote["user"], project_root, + ), partition, mapping, disable == "true", _integer(parser["cluster"], "poll_seconds", 1, 300), _resource("build", parser["build"], False), profiles, tuple(matrix), @@ -244,7 +273,7 @@ def _result_markdown(result: Mapping[str, Any]) -> str: def remote_run(run: Path) -> int: run = run.resolve() - config = load_config(run / "control" / "gpu-matrix.ini") + config = load_config(run / "control" / "config.ini") results = run / "results" scripts = run / "jobs" results.mkdir(parents=True, exist_ok=True) @@ -934,7 +963,7 @@ def parser() -> argparse.ArgumentParser: help="host alias in the SSH config", ) client.add_argument( - "--project-root", default="~/abacus_gpu_ci", + "--project-root", default=load_config().remote.project_root, help="remote directory for caches, runs, and archives", ) client.add_argument( @@ -962,8 +991,8 @@ def parser() -> argparse.ArgumentParser: help="local directory for downloaded results and client logs", ) commands.add_parser( - "config", help="validate and summarize gpu-matrix.ini", - description="Validate gpu-matrix.ini and print its resource and case counts.", + "config", help="validate and summarize config.ini", + description="Validate config.ini and print its remote, resource, and case configuration.", ) def internal(name: str) -> argparse.ArgumentParser: @@ -1002,7 +1031,13 @@ def main() -> int: args = parser().parse_args() if args.command == "config": config = load_config() - print(json.dumps({"resources": list(config.resources), "cases": len(config.cases)})) + print(json.dumps({ + "remote": { + "host": config.remote.host, "port": config.remote.port, + "user": config.remote.user, "project_root": config.remote.project_root, + }, + "resources": list(config.resources), "cases": len(config.cases), + })) return 0 if args.command == "remote-run": return remote_run(args.run) diff --git a/.ci/slurm/test_runner.py b/.ci/slurm/test_runner.py index 85ce19c2e3a..adb13dad286 100644 --- a/.ci/slurm/test_runner.py +++ b/.ci/slurm/test_runner.py @@ -45,24 +45,42 @@ def test_current_matrix_is_loaded_from_ini(self): self.assertEqual(config.resources["gpu4"].label, "4 GPUs") self.assertEqual(config.resources["gpu8x2"].label, "2 nodes / 16 GPUs") self.assertEqual(config.cases[-1].runner, "cusolvermp") + self.assertEqual(config.remote.host, "c0.sai.ai-4s.com") + self.assertEqual(config.remote.port, 12022) + self.assertEqual(config.remote.user, "abacususer01") + self.assertEqual(config.remote.project_root, "~/abacus_gpu_ci") + known_hosts = (ROOT / "known_hosts").read_text(encoding="utf-8") + self.assertIn("[{}]:{}".format(config.remote.host, config.remote.port), known_hosts) def test_resource_names_are_not_hardcoded(self): - text = (ROOT / "gpu-matrix.ini").read_text(encoding="utf-8") + text = (ROOT / "config.ini").read_text(encoding="utf-8") text = text.replace("resource.gpu1", "resource.single", 1) text = text.replace("resource = gpu1", "resource = single", 1) with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / "matrix.ini" + path = Path(directory) / "config.ini" path.write_text(text, encoding="utf-8") self.assertIn("single", runner.load_config(path).resources) def test_unknown_resource_and_noncontiguous_case_fail(self): - original = (ROOT / "gpu-matrix.ini").read_text(encoding="utf-8") + original = (ROOT / "config.ini").read_text(encoding="utf-8") for text in ( original.replace("resource = gpu1", "resource = missing", 1), original.replace("[case.002]", "[case.999]", 1), ): with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / "matrix.ini" + path = Path(directory) / "config.ini" + path.write_text(text, encoding="utf-8") + with self.assertRaises(ValueError): + runner.load_config(path) + + def test_invalid_remote_configuration_fails(self): + original = (ROOT / "config.ini").read_text(encoding="utf-8") + for text in ( + original.replace("port = 12022", "port = 0", 1), + original.replace("project_root = ~/", "project_root = relative/", 1), + ): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "config.ini" path.write_text(text, encoding="utf-8") with self.assertRaises(ValueError): runner.load_config(path) @@ -180,6 +198,7 @@ def test_submit_and_accounting_require_each_array_task(self): def test_pass_requires_successful_slurm_accounting(self): config = runner.Config( + runner.Remote("cluster.example", 22, "user", "~/gpu-ci"), "gpu", Path("/opt/cluster/mps_mapping.d"), False, 1, runner.Resource("build", "flood-gpu", 1, 1, 1, 60), {"one": runner.Resource("one", "flood-gpu", 1, 1, 1, 60)}, @@ -477,8 +496,11 @@ def test_workflow_uses_trusted_control_and_protected_environment(self): self.assertIn("gpu-ci-scheduled", text) self.assertIn("gpu-ci-manual", text) self.assertIn("/abacus-ci gpu", text) + self.assertIn("runner.py config", text) + self.assertIn("fromJSON(steps.cluster.outputs.config).remote.host", text) self.assertIn("runner.py run", text) self.assertIn("pull-requests: write", text) + self.assertNotIn("${{ vars.", text) self.assertNotIn("cpus-per-task", text) diff --git a/.github/workflows/gpu-validation.yml b/.github/workflows/gpu-validation.yml index 04f41338d62..43b5dc3a95f 100644 --- a/.github/workflows/gpu-validation.yml +++ b/.github/workflows/gpu-validation.yml @@ -90,10 +90,6 @@ jobs: CONTROL_SHA: ${{ needs.admit.outputs.control_sha }} PROJECT_ROOT_INPUT: ${{ inputs.project_root || '' }} RUN_NAMESPACE: ${{ needs.admit.outputs.namespace }} - REMOTE_PROJECT_ROOT: ${{ vars.REMOTE_PROJECT_ROOT }} - REMOTE_SSH_HOST: ${{ vars.REMOTE_SSH_HOST }} - REMOTE_SSH_PORT: ${{ vars.REMOTE_SSH_PORT }} - REMOTE_SSH_USER: ${{ vars.REMOTE_SSH_USER }} SOURCE_REPOSITORY: ${{ needs.admit.outputs.source_repository }} SOURCE_SHA: ${{ needs.admit.outputs.source_sha }} steps: @@ -113,9 +109,18 @@ jobs: path: source persist-credentials: false + - name: Read trusted cluster configuration + id: cluster + run: | + config=$(python3 control/.ci/slurm/runner.py config) + printf 'config=%s\n' "$config" >> "$GITHUB_OUTPUT" + - name: Configure SSH env: + REMOTE_SSH_HOST: ${{ fromJSON(steps.cluster.outputs.config).remote.host }} + REMOTE_SSH_PORT: ${{ fromJSON(steps.cluster.outputs.config).remote.port }} REMOTE_SSH_PRIVATE_KEY: ${{ secrets.REMOTE_SSH_PRIVATE_KEY }} + REMOTE_SSH_USER: ${{ fromJSON(steps.cluster.outputs.config).remote.user }} run: | set -euo pipefail root="$RUNNER_TEMP/remote-ssh" @@ -145,10 +150,11 @@ jobs: - name: Run shared GPU client id: client continue-on-error: true + env: + REMOTE_PROJECT_ROOT: ${{ fromJSON(steps.cluster.outputs.config).remote.project_root }} run: | set -euo pipefail project_root=${PROJECT_ROOT_INPUT:-$REMOTE_PROJECT_ROOT} - : "${project_root:?Set REMOTE_PROJECT_ROOT in the GitHub Environment}" mkdir -p "$ARTIFACT_ROOT" set +e python3 control/.ci/slurm/runner.py run \