fix(labeler-pro): pre-load dataset once, add memory spec and thin extraction script

Pre-load ScaleAI/SWE-bench_Pro dataset once in main() instead of once per
worker thread to avoid N concurrent load_dataset() calls causing OOM.
Add --mem=128G and 48h wall time to swebench_pro_label.sh.
Add extract_swebench_thin.sh (4 GPUs, no fat constraint) for Laguna extraction.
Add laguna_xs2_swebench_pro_full_labeler.yaml config.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: andre15silva

configs/labeling/laguna_xs2_swebench_pro_full_labeler.yaml

@@ -0,0 +1,13 @@
+# Label full Laguna-XS.2 run on SWE-bench Pro (all available trajectories, resume-safe)
+# Usage:
+#   sbatch slurm/swebench_pro_label.sh \
+#     --config configs/labeling/laguna_xs2_swebench_pro_full_labeler.yaml
+
+trajectory_dir: generations/swebench_pro/laguna_xs2_full
+output_dir: labels/swebench_pro/laguna_xs2_full
+scripts_dir: SWE-bench_Pro-os
+modal_app_name: program-probes-labeler-pro
+sandbox_timeout: 3600
+eval_timeout: 600
+resume: true
+n_workers: 100

run_labeler_pro.py

@@ -18,12 +18,12 @@
 from pathlib import Path
 
 from src.configs import SwebenchProLabelerConfig, load_config
-from src.labeling.swebench_pro_labeler import label_pro_trajectory, load_pro_instance
+from src.labeling.swebench_pro_labeler import label_pro_trajectory
 
 
 def _process(
     tf: Path,
-    iid: str,
+    instance: dict,
     cfg: SwebenchProLabelerConfig,
     out_dir: Path,
     scripts_dir: Path,
@@ -34,7 +34,6 @@ def _process(
         return f"{tf.name}: skipped (labels exist)"
     if jitter > 0:
         time.sleep(random.uniform(0, jitter))
-    instance = load_pro_instance(iid)
     label_pro_trajectory(
         tf,
         instance=instance,
@@ -82,17 +81,28 @@ def _iid(tf: Path) -> str:
 
     print(f"Labeling {len(traj_files)} trajectories from {traj_dir} → {out_dir}", flush=True)
 
+    # Pre-load dataset once to avoid N concurrent load_dataset() calls in threads
+    print("Loading SWE-bench Pro dataset ...", flush=True)
+    from datasets import load_dataset as _hf_load
+    _data = _hf_load("ScaleAI/SWE-bench_Pro", split="test")
+    instances: dict[str, dict] = {inst["instance_id"]: dict(inst) for inst in _data}
+    print(f"Loaded {len(instances)} instances.", flush=True)
+
+    def _get_instance(iid: str) -> dict:
+        if iid not in instances:
+            raise ValueError(f"Instance {iid!r} not found in SWE-bench Pro")
+        return instances[iid]
+
     jitter = 2.0 if cfg.n_workers > 1 else 0.0
 
     if cfg.n_workers == 1:
         for tf in traj_files:
-            iid = _iid(tf)
-            result = _process(tf, iid, cfg, out_dir, scripts_dir, jitter=0.0)
+            result = _process(tf, _get_instance(_iid(tf)), cfg, out_dir, scripts_dir, jitter=0.0)
             print(result, flush=True)
     else:
         with ThreadPoolExecutor(max_workers=cfg.n_workers) as pool:
             futures = {
-                pool.submit(_process, tf, _iid(tf), cfg, out_dir, scripts_dir, jitter): tf.name
+                pool.submit(_process, tf, _get_instance(_iid(tf)), cfg, out_dir, scripts_dir, jitter): tf.name
                 for tf in traj_files
             }
             for fut in as_completed(futures):

slurm/extract_swebench_thin.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Extract hidden states from SWE-bench trajectories — thin A100 nodes (4 GPUs, no fat constraint).
+#
+# Array job (8 shards):
+#   sbatch --array=0-7 slurm/extract_swebench_thin.sh \
+#     --model-config configs/models/laguna_xs2.yaml \
+#     --generation-config configs/generation_laguna_xs2.yaml \
+#     --traj-dir generations/swebench_pro/laguna_xs2_full \
+#     --output-dir outputs/swebench_pro/laguna_xs2_full
+#
+#SBATCH -J pp-extract-swebench
+#SBATCH -p berzelius
+#SBATCH --gpus=4
+#SBATCH -t 12:00:00
+#SBATCH -o logs/extract_swebench_%A_%a.out
+#SBATCH -e logs/extract_swebench_%A_%a.err
+
+set -euo pipefail
+mkdir -p logs
+
+module load buildenv-gcccuda/12.4.1-gcc13.3.0
+unset CPATH
+export LIBRARY_PATH="/usr/local/cuda/lib64:${LIBRARY_PATH:-}"
+export CUDA_HOME=/usr/local/cuda
+export PATH="/usr/local/cuda/bin:$PATH"
+
+RANK=${SLURM_ARRAY_TASK_ID:-0}
+NUM_SHARDS=${SLURM_ARRAY_TASK_COUNT:-1}
+
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+uv run python run_extract_swebench.py \
+  --shard-rank "$RANK" \
+  --num-shards "$NUM_SHARDS" \
+  --extraction-batch-size 1 \
+  "$@"

slurm/swebench_pro_label.sh

@@ -6,7 +6,8 @@
 #
 #SBATCH -J pp-pro-label
 #SBATCH -p berzelius-cpu
-#SBATCH -t 12:00:00
+#SBATCH --mem=128G
+#SBATCH -t 48:00:00
 #SBATCH -o logs/pro_labeler_%A_%a.out
 #SBATCH -e logs/pro_labeler_%A_%a.err