File size: 4,530 Bytes
a7525f2 c511323 a7525f2 c511323 a7525f2 c511323 a7525f2 c511323 a7525f2 c511323 | 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 104 105 106 107 108 109 110 111 112 113 114 | """Materialize the pinned 1,000 train / 250 test SmolDataEnvs tasks for HF Sandboxes.
Adapted from the FineEnvs env Space's prepare.py (same pinned revisions, same
instruction and grader hash checks), with one change: every task.toml gets
[environment] docker_image = "docker.io/savatar101/env-data-agent-train:base"
because Harbor's hf-sandbox backend needs a prebuilt image and will not build the
task's environment/Dockerfile. The Dockerfile is a one-line FROM of exactly this
image, so the sandbox runs bit-identical content.
Run once at job startup, then point OPENENV_DATASETS at the output:
python prepare_tasks.py --output /app/prepared
export OPENENV_DATASETS=/app/prepared/datasets/train,/app/prepared/datasets/test
"""
from __future__ import annotations
import argparse
import hashlib
import json
import shutil
import tomllib
from pathlib import Path
ROOT = Path(__file__).resolve().parent
# All tasks share this one base image (their Dockerfile is `FROM <this>` + WORKDIR).
DOCKER_IMAGE = "docker.io/savatar101/env-data-agent-train:base"
DATASETS = {
"train": (
"FineEnvs/SmolDataEnvs-harbor-train",
"4719635555de1666f374d847baebb500d368e493",
),
"test": (
"FineEnvs/SmolDataEnvs-harbor-test",
"b130595b579fc6026ad5762925a9a3b40e72f5bb",
),
}
GRADER_SHA256 = "16f906e9248fa8b7e1141939e425e88e59d52d65e98e360bbacef3720a43629f"
def task_rows(split):
return json.loads((ROOT / "data" / f"{split}.json").read_text())
def prepare(destination, splits=("train", "test")):
import tomli_w
from huggingface_hub import snapshot_download
destination = Path(destination).resolve()
for split in splits:
rows = task_rows(split)
repo, revision = DATASETS[split]
source = Path(
snapshot_download(
repo,
repo_type="dataset",
revision=revision,
allow_patterns=[f"tasks/{r['name']}/*" for r in rows],
max_workers=16,
)
)
for row in rows:
folder = source / "tasks" / row["name"]
instruction = (folder / "instruction.md").read_bytes()
if hashlib.sha256(instruction).hexdigest() != row["instruction_sha256"]:
raise ValueError(f"Historical task instruction changed: {split}/{row['name']}")
grader = (folder / "tests/grader.py").read_bytes()
if hashlib.sha256(grader).hexdigest() != GRADER_SHA256:
raise ValueError(f"Expected corrected numeric-only Harbor grader: {folder}")
target = destination / "datasets" / split / "tasks" / row["name"]
target.mkdir(parents=True, exist_ok=True)
for rel in (
"instruction.md",
"task.toml",
"environment/Dockerfile",
"environment/pull_bucket.py",
"tests/grader.py",
"tests/test.sh",
):
dst = target / rel
dst.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(folder / rel, dst)
spec = tomllib.loads((target / "task.toml").read_text())
spec["environment"]["memory_mb"] = 4096
spec["environment"]["cpus"] = 1
# hf-sandbox runs a prebuilt image instead of building environment/Dockerfile.
spec["environment"]["docker_image"] = DOCKER_IMAGE
# Cluster user IDs may not exist inside the sandbox user namespace.
spec["environment"].setdefault("env", {})["TAR_OPTIONS"] = "--no-same-owner"
(target / "task.toml").write_text(tomli_w.dumps(spec))
print(f"{split}: {len(rows)} tasks -> {destination / 'datasets' / split}")
# Overlap guard, as in the original.
if set(splits) == {"train", "test"}:
def names(split):
return {r["name"] for r in task_rows(split)}
notebooks = [
{r.get("notebook") for r in task_rows(s)} - {None} for s in ("train", "test")
]
if names("train") & names("test") or notebooks[0] & notebooks[1]:
raise ValueError("Train/test overlap")
print("no train/test overlap")
if __name__ == "__main__":
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--output", type=Path, default=ROOT / "prepared")
p.add_argument("--splits", nargs="+", default=["train", "test"], choices=["train", "test"])
a = p.parse_args()
prepare(a.output, tuple(a.splits))
|