code/official.diff

diff --git a/pyproject.toml b/pyproject.toml
index 5e68c70..068a90c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 # Miner-side SFT/RL training.
-train = ["trl", "accelerate", "deepspeed"]
+train = ["trl", "peft", "accelerate", "deepspeed"]
 
 [project.scripts]
 albedo-eval-api = "albedo_eval_service.control.api:main"
@@ -76,7 +76,7 @@ packages = [
 ]
 
 [tool.pytest.ini_options]
-pythonpath = ["src"]
+pythonpath = ["src", "."]
 testpaths = ["tests"]
 markers = [
     "integration: requires ALBEDO_TEST_DATABASE_URL and a Postgres database initialized from schema.sql",
diff --git a/scripts/prepare_datasets.py b/scripts/prepare_datasets.py
index a498824..8d6fd84 100644
--- a/scripts/prepare_datasets.py
+++ b/scripts/prepare_datasets.py
@@ -40,7 +40,7 @@ SOURCES: dict[str, dict] = {
     "open-swe-traces": {
         "repos": ["nvidia/Open-SWE-Traces"],
         "shard_glob": "data/train-*.parquet",
-        "raw_glob": "data/*/train-*.parquet",
+        "raw_glob": "data/**/train-*.parquet",
         "render": True,
         "family": "pr",
         "exclude_ids": _OPEN_SWE_LEAKS,
@@ -86,7 +86,20 @@ def _expected_parquet_shards(repo_id: str, shard_glob: str) -> set[str]:
 
 
 def _local_parquet_shards(dest: Path, shard_glob: str) -> set[str]:
-    return {p.relative_to(dest).as_posix() for p in dest.glob(shard_glob)}
+    """Match the same way as HuggingFace ``fnmatch`` (``*`` crosses ``/``).
+
+    ``Path.glob('data/*/train-*.parquet')`` is one directory deep, so nested
+    Open-SWE-Traces shards (``data/<agent>/<model>/<bench>/train-*.parquet``)
+    look missing after a successful download.
+    """
+    if not dest.is_dir():
+        return set()
+    found: set[str] = set()
+    for path in dest.rglob("*.parquet"):
+        rel = path.relative_to(dest).as_posix()
+        if fnmatch.fnmatch(rel, shard_glob):
+            found.add(rel)
+    return found
 
 
 def download_source(
diff --git a/scripts/render_trajectories.py b/scripts/render_trajectories.py
index 0f80a68..9806079 100644
--- a/scripts/render_trajectories.py
+++ b/scripts/render_trajectories.py
@@ -14,7 +14,7 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
-from prepare_datasets import SOURCES
+from prepare_datasets import SOURCES, _local_parquet_shards
 
 from albedo_eval_service.simulator.prompt_simulator import COMPLETE_MARKER
 
@@ -201,9 +201,10 @@ def _keep(row: dict, instance_id: str, spec: dict, seen_repos: Counter) -> str |
 
 def _raw_shards(raw_root: Path, spec: dict) -> list[Path]:
     files: list[Path] = []
+    glob = spec.get("raw_glob", "data/train-*.parquet")
     for repo in spec["repos"]:
         base = raw_root / repo.split("/")[-1]
-        files.extend(sorted(base.glob(spec.get("raw_glob", "data/train-*.parquet"))))
+        files.extend(sorted(base / rel for rel in _local_parquet_shards(base, glob)))
     return files
 
 
diff --git a/src/albedo_eval_service/remote/generation.py b/src/albedo_eval_service/remote/generation.py
index e384f70..06e49b0 100644
--- a/src/albedo_eval_service/remote/generation.py
+++ b/src/albedo_eval_service/remote/generation.py
@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import glob
 import multiprocessing as mp
 import os
 import queue as queue_module
+import signal
+import subprocess
+import sys
 import time
 from dataclasses import dataclass
 from typing import Any, Protocol
@@ -13,6 +17,73 @@ from .dataset import EvalSample
 from .prompt_remote import QWEN3_IM_END_TOKEN_ID
 
 
+def _bootstrap_cuda_env() -> None:
+    """Make nvcc visible to spawned vLLM/flashinfer workers.
+
+    Offline boxes often have CUDA only as the pip ``nvidia/cu*`` wheel.
+    ``CUDA_HOME`` set in the CLI parent is not always inherited by
+    EngineCore / Worker_TP processes, and flashinfer then falls back to
+    missing ``/usr/local/cuda``.
+    """
+    cache = os.path.join("/workspace/data/triton-cache", f"pid-{os.getpid()}")
+    os.makedirs(cache, exist_ok=True)
+    os.environ["TRITON_CACHE_DIR"] = cache
+    os.environ.setdefault("TRITON_HOME", "/workspace/data/triton-cache/home")
+    existing = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
+    if existing and os.path.isfile(os.path.join(existing, "bin", "nvcc")):
+        os.environ["CUDA_HOME"] = existing
+        os.environ["CUDA_PATH"] = existing
+        bin_dir = os.path.join(existing, "bin")
+        path = os.environ.get("PATH", "")
+        if bin_dir not in path.split(os.pathsep):
+            os.environ["PATH"] = f"{bin_dir}{os.pathsep}{path}"
+        return
+    # sys.executable may be a symlink to /usr/bin/python — use sys.prefix.
+    roots = [sys.prefix, getattr(sys, "base_prefix", sys.prefix), os.path.dirname(os.path.dirname(sys.executable))]
+    matches: list[str] = []
+    for venv_root in roots:
+        matches.extend(
+            glob.glob(
+                os.path.join(
+                    venv_root, "lib", "python*", "site-packages", "nvidia", "cu*", "bin", "nvcc"
+                )
+            )
+        )
+    matches = sorted(set(matches))
+    if not matches:
+        return
+    home = os.path.dirname(os.path.dirname(matches[-1]))
+    os.environ["CUDA_HOME"] = home
+    os.environ["CUDA_PATH"] = home
+    os.environ["PATH"] = f"{os.path.join(home, 'bin')}{os.pathsep}{os.environ.get('PATH', '')}"
+
+
+def _kill_process_tree(pid: int | None) -> None:
+    """SIGKILL a spawn worker and leftover EngineCore / Worker_TP children.
+
+    A generate() timeout used to return an error while the vLLM tree kept
+    the GPUs allocated. The next chain then failed with
+    ``Free memory ... less than desired GPU memory utilization``.
+    """
+    if pid is None:
+        return
+    try:
+        children = subprocess.check_output(
+            ["pgrep", "-P", str(pid)], text=True, stderr=subprocess.DEVNULL
+        ).split()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        children = []
+    for child in children:
+        try:
+            _kill_process_tree(int(child))
+        except ValueError:
+            continue
+    try:
+        os.kill(pid, signal.SIGKILL)
+    except ProcessLookupError:
+        pass
+
+
 @dataclass(frozen=True)
 class GenerationResult:
     sample_id: str
@@ -72,6 +143,7 @@ class VllmProcessGenerator:
         gpu_memory_utilization: float = 0.95,
         kv_cache_dtype: str = "auto",
         result_timeout_seconds: float = 900.0,
+        gdn_prefill_backend: str | None = None,
     ):
         self.model = model
         self.gpu_ids = gpu_ids
@@ -85,6 +157,7 @@ class VllmProcessGenerator:
         self.gpu_memory_utilization = gpu_memory_utilization
         self.kv_cache_dtype = kv_cache_dtype
         self.result_timeout_seconds = result_timeout_seconds
+        self.gdn_prefill_backend = gdn_prefill_backend
         self._ctx = mp.get_context("spawn")
         self._request_queue = None
         self._result_queue = None
@@ -116,12 +189,18 @@ class VllmProcessGenerator:
     def close(self) -> None:
         if self._process is None:
             return
+        pid = self._process.pid
         if self._process.is_alive() and self._request_queue is not None:
-            self._request_queue.put(None)
-            self._process.join(timeout=30)
+            try:
+                self._request_queue.put(None)
+            except Exception:
+                pass
+            self._process.join(timeout=8)
         if self._process.is_alive():
             self._process.terminate()
-            self._process.join(timeout=10)
+            self._process.join(timeout=5)
+        if self._process.is_alive() or pid:
+            _kill_process_tree(pid)
         self._process = None
         self._request_queue = None
         self._result_queue = None
@@ -147,6 +226,7 @@ class VllmProcessGenerator:
                 "compile_cache_dir": self.compile_cache_dir,
                 "gpu_memory_utilization": self.gpu_memory_utilization,
                 "kv_cache_dtype": self.kv_cache_dtype,
+                "gdn_prefill_backend": self.gdn_prefill_backend,
                 "queue": self._result_queue,
                 "request_queue": self._request_queue,
             },
@@ -164,6 +244,8 @@ class VllmProcessGenerator:
                         f"{self.result_timeout_seconds:g}s"
                     )
                 }
+                if self._process is not None:
+                    _kill_process_tree(self._process.pid)
                 break
             try:
                 candidate = self._result_queue.get(timeout=1)
@@ -203,11 +285,14 @@ def _vllm_worker(
     compile_cache_dir: str = "",
     gpu_memory_utilization: float = 0.95,
     kv_cache_dtype: str = "auto",
+    gdn_prefill_backend: str | None = None,
     queue=None,
     request_queue=None,
 ) -> None:
     try:
         os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(gpu_ids)
+        _bootstrap_cuda_env()
+        backend = gdn_prefill_backend or os.environ.get("ALBEDO_GDN_PREFILL_BACKEND") or None
 
         from vllm import LLM, SamplingParams
 
@@ -222,6 +307,8 @@ def _vllm_worker(
             "kv_cache_dtype": kv_cache_dtype,
             "limit_mm_per_prompt": {"image": 0, "video": 0},
         }
+        if backend:
+            llm_kwargs["gdn_prefill_backend"] = backend
         if max_model_len is not None:
             llm_kwargs["max_model_len"] = max_model_len
         if enforce_eager: