Spaces:
Running
Running
File size: 3,562 Bytes
a1e3907 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Task packs: directories of tasks (tasks.jsonl or tasks.jsonl.gz, with their reference plans) plus a manifest.
The server serves several packs at once; splits come from each task's `split` field. By default that is the dock-v1
eval and train packs (the hard tasks); the first pack, berth-v1, stays available by setting BERTH_TASKS_DIR.
BERTH_TASKS_DIR=/path/a:/path/b colon-separated pack directories
"""
from __future__ import annotations
import gzip
import hashlib
import json
import os
from pathlib import Path
from .model import Task
TASKS_ROOT = Path(__file__).resolve().parents[2] / "tasks"
DEFAULT_PACKS = ["dock-v1-eval", "dock-v1-train"]
LEGACY_PACK = "berth-v1"
def default_roots() -> list[Path]:
env = os.environ.get("BERTH_TASKS_DIR", "").strip()
if env:
return [Path(p) for p in env.split(":") if p]
roots = [TASKS_ROOT / p for p in DEFAULT_PACKS if _tasks_file(TASKS_ROOT / p)]
return roots or [TASKS_ROOT / LEGACY_PACK]
def _tasks_file(root: Path) -> Path | None:
for name in ("tasks.jsonl", "tasks.jsonl.gz"):
if (root / name).is_file():
return root / name
return None
def _read(root: Path) -> tuple[list[Task], dict]:
path = _tasks_file(root)
if path is None:
raise FileNotFoundError(f"no tasks.jsonl(.gz) in {root}")
raw = path.read_bytes()
body = gzip.decompress(raw) if path.suffix == ".gz" else raw
manifest = json.loads((root / "manifest.json").read_text()) if (root / "manifest.json").is_file() else {}
digest = hashlib.sha256(body).hexdigest()
if manifest.get("sha256") and manifest["sha256"] != digest:
raise ValueError(f"{path} does not match its manifest.json (sha256 {digest[:12]} != {manifest['sha256'][:12]})")
return [Task.from_dict(json.loads(line)) for line in body.decode().splitlines() if line.strip()], manifest
class TaskPack:
def __init__(self, roots: Path | str | list[Path | str] | None = None):
if roots is None:
roots = default_roots()
if isinstance(roots, (str, Path)):
roots = [roots]
self.roots = [Path(r) for r in roots]
self.root = self.roots[0]
self.tasks: list[Task] = []
self.manifests: dict[str, dict] = {}
for r in self.roots:
tasks, manifest = _read(r)
self.tasks += tasks
self.manifests[r.name] = manifest
self.manifest = self.manifests[self.root.name]
self._by_id = {}
for t in self.tasks:
if t.task_id in self._by_id:
raise ValueError(f"task id {t.task_id} appears in more than one pack")
self._by_id[t.task_id] = t
self._by_split: dict[str, list[Task]] = {}
for t in self.tasks:
self._by_split.setdefault(t.split, []).append(t)
def splits(self) -> list[str]:
return sorted(self._by_split)
def count(self, split: str) -> int:
return len(self._by_split.get(split, []))
def at(self, split: str, index: int) -> Task:
return self._by_split[split][index]
def get(self, task_id: str) -> Task:
return self._by_id[task_id]
def public(self, task: Task) -> dict:
"""What anyone may see: everything but the reference plans and costs."""
return task.to_dict(public=True)
_PACK: TaskPack | None = None
def load_pack(root: Path | str | list | None = None) -> TaskPack:
global _PACK
if root is not None:
return TaskPack(root)
if _PACK is None:
_PACK = TaskPack()
return _PACK
|