| """Task packs: a directory with tasks.jsonl and manifest.json. Each task carries its references (the exact optimum and |
| the baselines); `public()` strips them, with the record's releases and flows, before anything reaches an agent.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
|
|
| HIDDEN = ("reference",) |
| HIDDEN_DAM = ("record_release_cfs",) |
| HIDDEN_POINT = ("record_flow_cfs", "local_cfs") |
|
|
|
|
| def tasks_root() -> Path: |
| here = Path(__file__).resolve() |
| for p in here.parents: |
| if (p / "tasks").is_dir() and any((p / "tasks").glob("*/tasks.jsonl")): |
| return p / "tasks" |
| return here.parents[2] / "tasks" |
|
|
|
|
| class TaskPack: |
| def __init__(self, tasks: list[dict]): |
| self.tasks = tasks |
| self._by_id = {t["task_id"]: t for t in tasks} |
| self._by_split: dict[str, list[dict]] = {} |
| for t in tasks: |
| self._by_split.setdefault(t["split"], []).append(t) |
|
|
| def splits(self) -> list[str]: |
| return list(self._by_split) |
|
|
| def count(self, split: str) -> int: |
| return len(self._by_split.get(split, [])) |
|
|
| def at(self, split: str, i: int) -> dict: |
| return self._by_split[split][i] |
|
|
| def get(self, task_id: str) -> dict: |
| return self._by_id[task_id] |
|
|
| @staticmethod |
| def public(task: dict) -> dict: |
| t = {k: v for k, v in task.items() if k not in HIDDEN} |
| t["dams"] = [{k: v for k, v in d.items() if k not in HIDDEN_DAM} for d in task["dams"]] |
| t["points"] = [{k: v for k, v in p.items() if k not in HIDDEN_POINT} for p in task["points"]] |
| return t |
|
|
|
|
| def read_pack(d: Path) -> list[dict]: |
| return [json.loads(line) for line in (Path(d) / "tasks.jsonl").read_text().splitlines() if line.strip()] |
|
|
|
|
| def write_pack(d: Path, tasks: list[dict], meta: dict): |
| d = Path(d) |
| d.mkdir(parents=True, exist_ok=True) |
| body = "".join(json.dumps(t, separators=(",", ":")) + "\n" for t in tasks) |
| (d / "tasks.jsonl").write_text(body) |
| manifest = {**meta, "n": len(tasks), "sha256": hashlib.sha256(body.encode()).hexdigest(), |
| "tasks": [{"task_id": t["task_id"], "split": t["split"], "tier": t["tier"], |
| "opt_cost": t["reference"]["opt_cost"], "lazy_cost": t["reference"]["lazy_cost"]} |
| for t in tasks]} |
| (d / "manifest.json").write_text(json.dumps(manifest, indent=1)) |
|
|
|
|
| def load_pack(dirs: str | None = None) -> TaskPack: |
| dirs = dirs or os.environ.get("FLOOD_TASKS_DIR") |
| paths = [Path(p) for p in dirs.split(":")] if dirs else sorted(p.parent for p in tasks_root().glob("*/tasks.jsonl")) |
| tasks = [] |
| for p in paths: |
| tasks += read_pack(p) |
| return TaskPack(tasks) |
|
|
