worldmem-baseline-evals / deployment /repair_decmem_environment_berzelius.sh
BonanDing's picture
Declare VSA pytest dependency and add focused DecMem CPU repair job
49aa796 verified
Raw History Blame Contribute Delete
1.26 kB
#!/usr/bin/env bash
#SBATCH --job-name=decmem-env-repair
#SBATCH --partition=berzelius-hopper-cpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=2
#SBATCH --time=00:15:00
#SBATCH --account=berzelius-2026-196
#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
#SBATCH --output=slurm_logs/%x-%A_%a.out
#SBATCH --error=slurm_logs/%x-%A_%a.out
set -euo pipefail
module load buildenv-gcccuda/12.9.1-gcc11
ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
PYTHON=/proj/cvl/users/x_fahkh2/envs/decmem-eval/bin/python
OUTPUT=$ROOT/artifacts/decmem_environment_repair_$SLURM_JOB_ID
export TMPDIR=$ROOT/cache
export TRITON_CACHE_DIR=$ROOT/cache
export PIP_CACHE_DIR=$ROOT/cache/pip
export XDG_CACHE_HOME=$ROOT/cache
export PYTHONNOUSERSITE=1
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=2
export MKL_NUM_THREADS=2
mkdir -p "$TMPDIR" "$PIP_CACHE_DIR" "$OUTPUT"
# VSA imports pytest from its Triton module even when the H200 CUDA kernel is selected.
"$PYTHON" -m pip install pytest==9.1.1
"$PYTHON" -c 'import pytest; print("VSA dependency available: pytest", pytest.__version__)'
"$PYTHON" -m pip check
"$PYTHON" -m pip freeze > "$OUTPUT/decmem-eval_environment.txt"
printf 'DecMem environment repair completed: %s\n' "$OUTPUT"