scripts/slurm/example_agentic_retrieval.slurm

#!/bin/bash
# Example SLURM array script for the Plan-Act-Read agentic retrieval pipeline
# over hierarchical memory.
#
#SBATCH -J agentic_hier
#SBATCH -A ${SLURM_ACCOUNT:-your-account}
#SBATCH -p ${PARTITIONS:-cpu}
#SBATCH --nodes=1
#SBATCH --time=04:00:00
#SBATCH --array=0-7
#SBATCH --output=logs/agentic_hier_%A_%a.log
#SBATCH --export=ALL,NV_API_KEY

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
cd "$PROJECT_DIR"

MODEL_NAME="${MODEL_NAME:-gpt-5.5}"
TOP_K="${TOP_K:-20}"
SHARD_ROOT="${SHARD_ROOT:-output/shards/v5_${MODEL_NAME//./_}_nchunks10}"

shard_id=$(printf "%02d" "$SLURM_ARRAY_TASK_ID")

export ret_cache="$SHARD_ROOT/ret_cache/shard_${shard_id}.jsonl"
export plan_cache="response_cache/qa/${MODEL_NAME//./_}_plan_cache_shard_${shard_id}"
export reading_cache="response_cache/qa/${MODEL_NAME//./_}_reading_cache_shard_${shard_id}"

python main.py \
    --in_file  "$SHARD_ROOT/dataset/shard_${shard_id}.json" \
    --out_file "$SHARD_ROOT/agentic_hier/part_${shard_id}.jsonl" \
    --model_name "$MODEL_NAME" \
    --top_k "$TOP_K" \
    --n_chunks 10 \
    --nvidia \
    --all_sessions_file dataset/all_sessions.json \
    --hier_v2 \
    --hier_union \
    --mode agent