From 3869101b5aa9245a01c869467304f177f55a7501 Mon Sep 17 00:00:00 2001 From: lullu57 Date: Wed, 4 Mar 2026 15:31:17 -0500 Subject: [PATCH 1/2] Add configurable Slurm submission workflow with storage and email alerts --- README.md | 48 +++++ .../slurm_experiment_list.example.txt | 4 + scripts/slurm/run_kerneltuner_array.sbatch | 188 ++++++++++++++++++ scripts/slurm/submit_kerneltuner.sh | 170 ++++++++++++++++ 4 files changed, 410 insertions(+) create mode 100644 configs/experiments/slurm_experiment_list.example.txt create mode 100755 scripts/slurm/run_kerneltuner_array.sbatch create mode 100755 scripts/slurm/submit_kerneltuner.sh diff --git a/README.md b/README.md index 55ef940..3b6af62 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,54 @@ The initial workflow is document-driven: 3. Implement modules against the specs in `docs/specs/`. 4. Keep deviations explicit through ADRs rather than ad hoc code changes. +## Slurm Submission Helpers + +For cluster execution, the repo includes reusable Slurm scripts: + +- `scripts/slurm/run_kerneltuner_array.sbatch`: array worker that maps one array task to one experiment YAML. +- `scripts/slurm/submit_kerneltuner.sh`: submit wrapper that computes array size from a list file. + +Example list file: + +- `configs/experiments/slurm_experiment_list.example.txt` + +Basic usage: + +```bash +chmod +x scripts/slurm/submit_kerneltuner.sh +cp configs/experiments/slurm_experiment_list.example.txt configs/experiments/slurm_experiment_list.txt +scripts/slurm/submit_kerneltuner.sh \ + --list configs/experiments/slurm_experiment_list.txt \ + --partition gpunodes \ + --gpu-type rtx_a2000 \ + --gpus 1 \ + --time 0-04:00 \ + --cpus 4 \ + --mem 24GB \ + --log-dir /scratch/scratch-space/expires-xxxx/$USER/kerneltuner_logs \ + --scratch-root /scratch/scratch-space/expires-xxxx/$USER/kerneltuner \ + --artifact-root /scratch/scratch-space/expires-xxxx/$USER/kerneltuner_artifacts \ + --mail-user you@example.com \ + --mail-type BEGIN,END,FAIL \ + --alert-email you@example.com \ + --alert-on-start +``` + +Useful environment overrides: + +- `RUN_COMMAND_TEMPLATE` (default: `ktune run-experiment --experiment "{experiment}"`) +- `INSTALL_PACKAGES` (`0` to skip pip install in jobs) +- `EXTRA_PIP_PACKAGES` (space-separated extra pip packages) +- `SKIP_IF_ARTIFACTS_EXIST` (`0` to force reruns) +- `DRY_RUN=1` (worker-level dry run) +- `ALERT_EMAIL`, `ALERT_ON_START`, `ALERT_ON_END`, `ALERT_ON_FAIL` for worker-level alerts + +Notes: + +- `--artifact-root` overrides `artifact_root` in experiment YAML at submission time. +- If no scratch path exists on the node, jobs fall back to `/.scratch/$USER`. +- Slurm native email (`--mail-user/--mail-type`) and worker-level alerts can be used together. + ## What This Repo Is Not - It is not a Triton compiler redesign effort. diff --git a/configs/experiments/slurm_experiment_list.example.txt b/configs/experiments/slurm_experiment_list.example.txt new file mode 100644 index 0000000..d777083 --- /dev/null +++ b/configs/experiments/slurm_experiment_list.example.txt @@ -0,0 +1,4 @@ +# One experiment YAML path per line. +# Blank lines and lines starting with '#' are ignored. + +configs/experiments/gemm_smoke.example.yaml diff --git a/scripts/slurm/run_kerneltuner_array.sbatch b/scripts/slurm/run_kerneltuner_array.sbatch new file mode 100755 index 0000000..05fc86b --- /dev/null +++ b/scripts/slurm/run_kerneltuner_array.sbatch @@ -0,0 +1,188 @@ +#!/bin/bash --login +#SBATCH --job-name=kerneltuner +#SBATCH --partition=gpunodes +#SBATCH --time=0-02:00 +#SBATCH -c 4 +#SBATCH --mem=24GB +#SBATCH --gres=gpu:1 +#SBATCH --array=0-0 +#SBATCH --output=slurm_jobs/kerneltuner_%A_%a.out +#SBATCH --error=slurm_jobs/kerneltuner_%A_%a.err + +set -euo pipefail + +send_alert() { + local phase="$1" + local details="$2" + local email="${ALERT_EMAIL:-}" + if [[ -z "$email" ]]; then + return 0 + fi + local subject_prefix="${ALERT_SUBJECT_PREFIX:-[KernelTuner]}" + local subject="${subject_prefix} ${phase} job=${SLURM_JOB_ID:-local} task=${SLURM_ARRAY_TASK_ID:-0}" + local body + body=$( + cat </dev/null 2>&1; then + printf '%s\n' "$body" | mail -s "$subject" "$email" || true + elif command -v sendmail >/dev/null 2>&1; then + { + printf 'To: %s\n' "$email" + printf 'Subject: %s\n' "$subject" + printf '\n%s\n' "$body" + } | sendmail -t || true + fi +} + +on_exit() { + local exit_code=$? + if [[ $exit_code -ne 0 && "${ALERT_ON_FAIL:-1}" == "1" ]]; then + send_alert "FAIL" "exit_code=${exit_code}" + fi +} +trap on_exit EXIT + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [workspace_root] [venv_name]" >&2 + exit 2 +fi + +EXPERIMENT_LIST_FILE="$1" +WORKSPACE_ROOT="${2:-$SLURM_SUBMIT_DIR}" +VENV_NAME="${3:-kerneltuner_env}" + +if [[ ! -f "$EXPERIMENT_LIST_FILE" ]]; then + echo "ERROR: Experiment list file not found: $EXPERIMENT_LIST_FILE" >&2 + exit 1 +fi + +mkdir -p "$WORKSPACE_ROOT/slurm_jobs" + +mapfile -t EXPERIMENTS < <(awk 'NF && $1 !~ /^#/' "$EXPERIMENT_LIST_FILE") + +if [[ "${#EXPERIMENTS[@]}" -eq 0 ]]; then + echo "ERROR: No experiments found in $EXPERIMENT_LIST_FILE" >&2 + exit 1 +fi + +TASK_ID="${SLURM_ARRAY_TASK_ID:-0}" +if (( TASK_ID < 0 || TASK_ID >= ${#EXPERIMENTS[@]} )); then + echo "ERROR: SLURM_ARRAY_TASK_ID=$TASK_ID out of range (0..$((${#EXPERIMENTS[@]} - 1)))" >&2 + exit 1 +fi + +EXPERIMENT_CFG="${EXPERIMENTS[$TASK_ID]}" +if [[ ! -f "$EXPERIMENT_CFG" ]]; then + echo "ERROR: Experiment config not found: $EXPERIMENT_CFG" >&2 + exit 1 +fi + +resolve_artifact_root() { + local cfg_path="$1" + local raw_root + raw_root="$(awk -F': *' '$1=="artifact_root"{print $2; exit}' "$cfg_path" | tr -d "\"'[:space:]")" + if [[ -z "$raw_root" ]]; then + raw_root="artifacts" + fi + if [[ "$raw_root" = /* ]]; then + printf '%s\n' "$raw_root" + else + printf '%s/%s\n' "$WORKSPACE_ROOT" "$raw_root" + fi +} + +echo ">>> KernelTuner Slurm Task <<<" +echo "Working directory: $WORKSPACE_ROOT" +echo "Experiment list: $EXPERIMENT_LIST_FILE" +echo "Array task id: $TASK_ID" +echo "Experiment config: $EXPERIMENT_CFG" +if [[ "${ALERT_ON_START:-0}" == "1" ]]; then + send_alert "START" "experiment_cfg=${EXPERIMENT_CFG}" +fi + +EXPERIMENT_ID="$(awk -F': *' '$1=="experiment_id"{print $2; exit}' "$EXPERIMENT_CFG" | tr -d "\"'[:space:]")" +if [[ -z "${EXPERIMENT_ID}" ]]; then + EXPERIMENT_ID="$(basename "$EXPERIMENT_CFG" .yaml)" +fi + +EXPERIMENT_TO_RUN="$EXPERIMENT_CFG" + +SCRATCH_ROOT="${SCRATCH_ROOT:-}" +if [[ -z "$SCRATCH_ROOT" ]]; then + SCRATCH_ROOT="$(ls -d /scratch/scratch-space/expires-* 2>/dev/null | head -n 1 || true)" +fi +if [[ -z "$SCRATCH_ROOT" ]]; then + SCRATCH_ROOT="$WORKSPACE_ROOT/.scratch/$USER" +fi +mkdir -p "$SCRATCH_ROOT" +echo "Scratch root: $SCRATCH_ROOT" + +if [[ -n "${ARTIFACT_ROOT_OVERRIDE:-}" ]]; then + mkdir -p "$SCRATCH_ROOT/kerneltuner_tmp" + OVERRIDE_CFG="$SCRATCH_ROOT/kerneltuner_tmp/${EXPERIMENT_ID}_${SLURM_JOB_ID:-local}_${TASK_ID}.yaml" + awk -v override="$ARTIFACT_ROOT_OVERRIDE" ' + BEGIN { seen = 0 } + /^artifact_root:/ { print "artifact_root: " override; seen = 1; next } + { print } + END { if (seen == 0) print "artifact_root: " override } + ' "$EXPERIMENT_CFG" > "$OVERRIDE_CFG" + EXPERIMENT_TO_RUN="$OVERRIDE_CFG" + echo "Using artifact_root override: $ARTIFACT_ROOT_OVERRIDE" +fi + +if [[ "${SKIP_IF_ARTIFACTS_EXIST:-1}" == "1" ]]; then + ARTIFACT_ROOT_RESOLVED="$(resolve_artifact_root "$EXPERIMENT_TO_RUN")" + if compgen -G "$ARTIFACT_ROOT_RESOLVED/$EXPERIMENT_ID/*/summary.json" > /dev/null; then + echo ">>> Found existing summary for experiment_id=$EXPERIMENT_ID; skipping." + if [[ "${ALERT_ON_END:-0}" == "1" ]]; then + send_alert "SKIP" "existing_summary=1 experiment_id=${EXPERIMENT_ID}" + fi + exit 0 + fi +fi + +VENV_PATH="$SCRATCH_ROOT/$VENV_NAME" +if [[ ! -d "$VENV_PATH" ]]; then + echo "Creating virtual environment at $VENV_PATH" + python3 -m venv "$VENV_PATH" +fi + +source "$VENV_PATH/bin/activate" + +if [[ "${INSTALL_PACKAGES:-1}" == "1" ]]; then + echo "Installing dependencies into $VENV_PATH" + pip install --upgrade pip + pip install --no-cache-dir -e "$WORKSPACE_ROOT" pyyaml pandas pyarrow + if [[ -n "${EXTRA_PIP_PACKAGES:-}" ]]; then + pip install --no-cache-dir ${EXTRA_PIP_PACKAGES} + fi +fi + +export PYTHONPATH="$WORKSPACE_ROOT/src:${PYTHONPATH:-}" + +RUN_COMMAND_TEMPLATE="${RUN_COMMAND_TEMPLATE:-ktune run-experiment --experiment \"{experiment}\"}" +EXPERIMENT_TOKEN='{experiment}' +WORKSPACE_TOKEN='{workspace}' +RUN_COMMAND="${RUN_COMMAND_TEMPLATE//$EXPERIMENT_TOKEN/$EXPERIMENT_TO_RUN}" +RUN_COMMAND="${RUN_COMMAND//$WORKSPACE_TOKEN/$WORKSPACE_ROOT}" + +echo "Run command: $RUN_COMMAND" +if [[ "${DRY_RUN:-0}" == "1" ]]; then + echo "DRY_RUN=1 set; not executing." + exit 0 +fi + +bash -lc "$RUN_COMMAND" + +echo ">>> Task complete" +if [[ "${ALERT_ON_END:-1}" == "1" ]]; then + send_alert "SUCCESS" "experiment_cfg=${EXPERIMENT_TO_RUN}" +fi diff --git a/scripts/slurm/submit_kerneltuner.sh b/scripts/slurm/submit_kerneltuner.sh new file mode 100755 index 0000000..c68c69f --- /dev/null +++ b/scripts/slurm/submit_kerneltuner.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: + scripts/slurm/submit_kerneltuner.sh --list [options] + +Options: + --list Required. Text file with one experiment YAML path per line. + --job-name Slurm job name. Default: kerneltuner + --partition Slurm partition. Default: gpunodes + --time Slurm time limit. Default: 0-02:00 + --cpus CPUs per task. Default: 4 + --mem Memory per task. Default: 24GB + --gpu-type GPU type for --gres (e.g. rtx_a2000) + --gpus Number of GPUs. Default: 1 + --mail-user Email for Slurm notifications (optional) + --mail-type Slurm mail type string (optional) + --venv-name Virtualenv directory name. Default: kerneltuner_env + --workspace Workspace root. Default: current directory + --log-dir Directory for Slurm .out/.err files. Default: /slurm_jobs + --scratch-root Override SCRATCH_ROOT used by worker script + --artifact-root Override experiment artifact_root for this submission + --alert-email Send worker alerts to this email (mail/sendmail on node) + --alert-on-start Enable start alert from worker + --alert-on-end Enable success/skip alert from worker + --alert-on-fail Enable failure alert from worker (default on) + --dry-run Print sbatch command without submitting + +Environment vars passed through to worker script: + RUN_COMMAND_TEMPLATE e.g. 'ktune run-experiment --experiment "{experiment}"' + DRY_RUN Set to 1 for worker-level dry run + INSTALL_PACKAGES Set to 0 to skip pip install + EXTRA_PIP_PACKAGES Extra pip packages to install in job env + SKIP_IF_ARTIFACTS_EXIST Set to 0 to disable skip-on-existing-summary behavior +EOF +} + +LIST_FILE="" +JOB_NAME="kerneltuner" +PARTITION="gpunodes" +TIME_LIMIT="0-02:00" +CPUS="4" +MEMORY="24GB" +GPU_TYPE="" +GPU_COUNT="1" +MAIL_USER="" +MAIL_TYPE="" +VENV_NAME="kerneltuner_env" +WORKSPACE_ROOT="$(pwd)" +LOG_DIR="" +SCRATCH_ROOT_OVERRIDE="" +ARTIFACT_ROOT_OVERRIDE="" +ALERT_EMAIL="" +ALERT_ON_START="0" +ALERT_ON_END="1" +ALERT_ON_FAIL="1" +DRY_RUN_SUBMIT="0" + +while [[ $# -gt 0 ]]; do + case "$1" in + --list) LIST_FILE="${2:-}"; shift 2 ;; + --job-name) JOB_NAME="${2:-}"; shift 2 ;; + --partition) PARTITION="${2:-}"; shift 2 ;; + --time) TIME_LIMIT="${2:-}"; shift 2 ;; + --cpus) CPUS="${2:-}"; shift 2 ;; + --mem) MEMORY="${2:-}"; shift 2 ;; + --gpu-type) GPU_TYPE="${2:-}"; shift 2 ;; + --gpus) GPU_COUNT="${2:-}"; shift 2 ;; + --mail-user) MAIL_USER="${2:-}"; shift 2 ;; + --mail-type) MAIL_TYPE="${2:-}"; shift 2 ;; + --venv-name) VENV_NAME="${2:-}"; shift 2 ;; + --workspace) WORKSPACE_ROOT="${2:-}"; shift 2 ;; + --log-dir) LOG_DIR="${2:-}"; shift 2 ;; + --scratch-root) SCRATCH_ROOT_OVERRIDE="${2:-}"; shift 2 ;; + --artifact-root) ARTIFACT_ROOT_OVERRIDE="${2:-}"; shift 2 ;; + --alert-email) ALERT_EMAIL="${2:-}"; shift 2 ;; + --alert-on-start) ALERT_ON_START="1"; shift ;; + --alert-on-end) ALERT_ON_END="1"; shift ;; + --alert-on-fail) ALERT_ON_FAIL="1"; shift ;; + --dry-run) DRY_RUN_SUBMIT="1"; shift ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown option: $1" >&2; usage; exit 2 ;; + esac +done + +if [[ -z "$LIST_FILE" ]]; then + echo "ERROR: --list is required" >&2 + usage + exit 2 +fi + +if [[ ! -f "$LIST_FILE" ]]; then + echo "ERROR: list file not found: $LIST_FILE" >&2 + exit 1 +fi + +readarray -t EXP_LINES < <(awk 'NF && $1 !~ /^#/' "$LIST_FILE") +NUM_EXPERIMENTS="${#EXP_LINES[@]}" +if [[ "$NUM_EXPERIMENTS" -eq 0 ]]; then + echo "ERROR: no experiment entries found in $LIST_FILE" >&2 + exit 1 +fi + +ARRAY_SPEC="0-$((NUM_EXPERIMENTS - 1))" +if [[ -z "$LOG_DIR" ]]; then + LOG_DIR="$WORKSPACE_ROOT/slurm_jobs" +fi +mkdir -p "$LOG_DIR" + +if [[ -n "$GPU_TYPE" ]]; then + GRES="gpu:${GPU_TYPE}:${GPU_COUNT}" +else + GRES="gpu:${GPU_COUNT}" +fi + +SBATCH_CMD=( + sbatch + "--job-name=$JOB_NAME" + "--partition=$PARTITION" + "--time=$TIME_LIMIT" + "--cpus-per-task=$CPUS" + "--mem=$MEMORY" + "--gres=$GRES" + "--array=$ARRAY_SPEC" + "--output=$LOG_DIR/${JOB_NAME}_%A_%a.out" + "--error=$LOG_DIR/${JOB_NAME}_%A_%a.err" +) + +if [[ -n "$MAIL_USER" ]]; then + SBATCH_CMD+=("--mail-user=$MAIL_USER") +fi +if [[ -n "$MAIL_TYPE" ]]; then + SBATCH_CMD+=("--mail-type=$MAIL_TYPE") +fi + +EXPORT_VARS=("ALL") +if [[ -n "$SCRATCH_ROOT_OVERRIDE" ]]; then + EXPORT_VARS+=("SCRATCH_ROOT=$SCRATCH_ROOT_OVERRIDE") +fi +if [[ -n "$ARTIFACT_ROOT_OVERRIDE" ]]; then + EXPORT_VARS+=("ARTIFACT_ROOT_OVERRIDE=$ARTIFACT_ROOT_OVERRIDE") +fi +if [[ -n "$ALERT_EMAIL" ]]; then + EXPORT_VARS+=("ALERT_EMAIL=$ALERT_EMAIL") +fi +EXPORT_VARS+=("ALERT_ON_START=$ALERT_ON_START") +EXPORT_VARS+=("ALERT_ON_END=$ALERT_ON_END") +EXPORT_VARS+=("ALERT_ON_FAIL=$ALERT_ON_FAIL") +SBATCH_CMD+=("--export=$(IFS=,; echo "${EXPORT_VARS[*]}")") + +SBATCH_CMD+=( + "$WORKSPACE_ROOT/scripts/slurm/run_kerneltuner_array.sbatch" + "$LIST_FILE" + "$WORKSPACE_ROOT" + "$VENV_NAME" +) + +echo "Submitting $NUM_EXPERIMENTS experiments as array: $ARRAY_SPEC" +echo "sbatch command:" +printf ' %q' "${SBATCH_CMD[@]}" +echo + +if [[ "$DRY_RUN_SUBMIT" == "1" ]]; then + echo "Dry run enabled; not submitting." + exit 0 +fi + +"${SBATCH_CMD[@]}" From 3aaa352bfbe9e393bd9f7fea890925feaf7b01fb Mon Sep 17 00:00:00 2001 From: lullu57 Date: Fri, 6 Mar 2026 12:12:53 -0500 Subject: [PATCH 2/2] Fix Slurm worker path resolution and workspace fallback --- scripts/slurm/run_kerneltuner_array.sbatch | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/scripts/slurm/run_kerneltuner_array.sbatch b/scripts/slurm/run_kerneltuner_array.sbatch index 05fc86b..1fcc8c2 100755 --- a/scripts/slurm/run_kerneltuner_array.sbatch +++ b/scripts/slurm/run_kerneltuner_array.sbatch @@ -56,7 +56,7 @@ if [[ $# -lt 1 ]]; then fi EXPERIMENT_LIST_FILE="$1" -WORKSPACE_ROOT="${2:-$SLURM_SUBMIT_DIR}" +WORKSPACE_ROOT="${2:-${SLURM_SUBMIT_DIR:-$(pwd)}}" VENV_NAME="${3:-kerneltuner_env}" if [[ ! -f "$EXPERIMENT_LIST_FILE" ]]; then @@ -79,7 +79,12 @@ if (( TASK_ID < 0 || TASK_ID >= ${#EXPERIMENTS[@]} )); then exit 1 fi -EXPERIMENT_CFG="${EXPERIMENTS[$TASK_ID]}" +EXPERIMENT_CFG_RAW="${EXPERIMENTS[$TASK_ID]}" +if [[ "$EXPERIMENT_CFG_RAW" = /* ]]; then + EXPERIMENT_CFG="$EXPERIMENT_CFG_RAW" +else + EXPERIMENT_CFG="$WORKSPACE_ROOT/$EXPERIMENT_CFG_RAW" +fi if [[ ! -f "$EXPERIMENT_CFG" ]]; then echo "ERROR: Experiment config not found: $EXPERIMENT_CFG" >&2 exit 1