Add reusable SpectraREML batch AI-REML engine

2026-08-09 14:13:38 +08:00
parent 777eeb288c
commit da971a2ca7
22 changed files with 4736 additions and 0 deletions
--- a/scripts/run_server.sh
+++ b/scripts/run_server.sh
@@ -0,0 +1,112 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+ACTION="${1:-}"
+if [[ ! "$ACTION" =~ ^(build|test|validate|run|dry-run|finalize|status|all)$ ]]; then
+  echo "Usage: bash $0 {build|test|validate|run|dry-run|finalize|status|all}" >&2
+  exit 2
+fi
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_DIR="${PROJECT_DIR:-$(cd -- "$SCRIPT_DIR/.." && pwd)}"
+PYTHON="${PYTHON:-python3}"
+BACKEND="${BACKEND:-MKL}"
+BACKEND="${BACKEND^^}"
+BUILD_DIR="${BUILD_DIR:-$PROJECT_DIR/build-${BACKEND,,}}"
+ENGINE="${ENGINE:-$BUILD_DIR/spectra_reml}"
+MANIFEST="${MANIFEST:-}"
+OUTPUT="${OUTPUT:-}"
+THREADS="${THREADS:-1}"
+BLAS_THREADS="${BLAS_THREADS:-1}"
+BLOCK_SIZE="${BLOCK_SIZE:-256}"
+RESUME="${RESUME:-1}"
+FORCE="${FORCE:-0}"
+CLI="$PROJECT_DIR/python/spectra_reml.py"
+
+die() {
+  echo "ERROR: $*" >&2
+  exit 1
+}
+
+[[ "$BACKEND" == "MKL" || "$BACKEND" == "OPENBLAS" ]] || die "BACKEND must be MKL or OPENBLAS"
+[[ "$THREADS" =~ ^[1-9][0-9]*$ ]] || die "THREADS must be positive"
+[[ "$BLAS_THREADS" =~ ^[1-9][0-9]*$ ]] || die "BLAS_THREADS must be positive"
+[[ "$BLOCK_SIZE" =~ ^[1-9][0-9]*$ ]] || die "BLOCK_SIZE must be positive"
+[[ "$RESUME" == "0" || "$RESUME" == "1" ]] || die "RESUME must be 0 or 1"
+[[ "$FORCE" == "0" || "$FORCE" == "1" ]] || die "FORCE must be 0 or 1"
+[[ "$RESUME" != "1" || "$FORCE" != "1" ]] || die "RESUME and FORCE are mutually exclusive"
+
+build_project() {
+  local args=(
+    -S "$PROJECT_DIR"
+    -B "$BUILD_DIR"
+    -DCMAKE_BUILD_TYPE=Release
+    -DREML_BLAS="$BACKEND"
+    -DREML_ENABLE_OPENMP=ON
+  )
+  if [[ "$BACKEND" == "MKL" ]]; then
+    args+=(
+      -DMKL_INTERFACE=lp64
+      -DMKL_LINK=dynamic
+      -DMKL_THREADING=sequential
+    )
+  fi
+  cmake "${args[@]}"
+  cmake --build "$BUILD_DIR" --parallel
+}
+
+if [[ "$ACTION" == "build" ]]; then
+  build_project
+  exit 0
+fi
+
+command -v "$PYTHON" >/dev/null 2>&1 || die "Python was not found: $PYTHON"
+[[ -s "$CLI" ]] || die "Generic CLI is missing: $CLI"
+
+if [[ "$ACTION" == "test" ]]; then
+  [[ -d "$BUILD_DIR" ]] || die "Build first: bash $0 build"
+  ctest --test-dir "$BUILD_DIR" --output-on-failure
+  "$PYTHON" "$PROJECT_DIR/tests/test_math_reference.py"
+  "$PYTHON" "$PROJECT_DIR/tests/python/test_spectra_reml_cli.py"
+  exit 0
+fi
+
+[[ -n "$MANIFEST" ]] || die "Set MANIFEST=/absolute/path/to/manifest.json"
+
+case "$ACTION" in
+  validate|status)
+    "$PYTHON" "$CLI" "$ACTION" --manifest "$MANIFEST"
+    ;;
+  finalize)
+    [[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
+    "$PYTHON" "$CLI" finalize --manifest "$MANIFEST" --output "$OUTPUT"
+    ;;
+  run|dry-run|all)
+    if [[ "$ACTION" != "dry-run" ]]; then
+      [[ -x "$ENGINE" ]] || die "Engine is missing/not executable: $ENGINE"
+    fi
+    mode=()
+    if [[ "$FORCE" == "1" ]]; then
+      mode+=(--force)
+    elif [[ "$RESUME" == "1" ]]; then
+      mode+=(--resume)
+    fi
+    common=(
+      --manifest "$MANIFEST"
+      --engine "$ENGINE"
+      --threads "$THREADS"
+      --blas-threads "$BLAS_THREADS"
+      --block-size "$BLOCK_SIZE"
+      "${mode[@]}"
+    )
+    if [[ "$ACTION" == "dry-run" ]]; then
+      "$PYTHON" "$CLI" run "${common[@]}" --dry-run
+    elif [[ "$ACTION" == "all" ]]; then
+      [[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
+      "$PYTHON" "$CLI" all "${common[@]}" --output "$OUTPUT"
+    else
+      "$PYTHON" "$CLI" run "${common[@]}"
+    fi
+    ;;
+esac