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////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: