|
| 1 | +"""Bundles run straight from a Hugging Face Hub dataset repo, named ``hf://<org>/<repo>[@<revision>][/<path>]``. Each |
| 2 | +commit's bundle is downloaded once into a plain folder, since the Hub's own cache keeps files as links that would leave |
| 3 | +the bundle, and its ids are rooted in the ``@hf`` namespace, which only this path produces.""" |
| 4 | + |
| 5 | +from __future__ import annotations |
| 6 | + |
| 7 | +import json |
| 8 | +import os |
| 9 | +import re |
| 10 | +from collections.abc import Callable, Mapping, Sequence |
| 11 | +from dataclasses import dataclass, field |
| 12 | +from pathlib import Path |
| 13 | +from typing import Any |
| 14 | +from urllib.parse import unquote |
| 15 | + |
| 16 | +from huggingface_hub import HfApi |
| 17 | +from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError |
| 18 | +from huggingface_hub.repocard import metadata_load |
| 19 | + |
| 20 | +from agent_env.config.paths import cache_root |
| 21 | +from agent_env.plugins._bundles import ALIASES_FILE |
| 22 | +from agent_env.plugins._requirements import unmet_problems |
| 23 | +from agent_env.store.ids import HF_PREFIX, validate_local_id |
| 24 | +from agent_env.store.local_state import holding_locks |
| 25 | + |
| 26 | +from .parse import BUNDLE_TOML, BundleError, BundleKind, read_requires |
| 27 | +from .requires import UnmetRequirements |
| 28 | + |
| 29 | +SCHEME = "hf://" |
| 30 | +BUNDLES = "bundles" |
| 31 | +CARD = "README.md" |
| 32 | +SOURCE_KEY = "bundle_source" |
| 33 | + |
| 34 | +_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") |
| 35 | +_SHA = re.compile(r"[0-9a-f]{40}") |
| 36 | +_KINDS = frozenset(kind.value for kind in BundleKind) |
| 37 | +_FORM = "hf://<org>/<repo>[@<revision>][/<path>]" |
| 38 | + |
| 39 | + |
| 40 | +@dataclass(frozen=True) |
| 41 | +class HubReference: |
| 42 | + """An ``hf://`` reference's parts: ``<org>/<repo>``, the revision named, and the bundle's folder in the repo.""" |
| 43 | + |
| 44 | + repo: str |
| 45 | + revision: str | None |
| 46 | + path: str | None |
| 47 | + |
| 48 | + |
| 49 | +@dataclass(frozen=True) |
| 50 | +class HubSource: |
| 51 | + """Where a Hub bundle came from: its repo, the revision named (``main`` by default), its commit, and its folder in |
| 52 | + the repo, ``""`` for the root. Each run's context records it under ``SOURCE_KEY``.""" |
| 53 | + |
| 54 | + repo: str |
| 55 | + revision: str |
| 56 | + sha: str |
| 57 | + path: str |
| 58 | + |
| 59 | + @classmethod |
| 60 | + def from_metadata(cls, data: Mapping[str, str]) -> HubSource: |
| 61 | + return cls(data["repo"], data["revision"], data["sha"], data["path"]) |
| 62 | + |
| 63 | + def to_metadata(self) -> dict[str, str]: |
| 64 | + return {"hub": "hf", "repo": self.repo, "revision": self.revision, "sha": self.sha, "path": self.path} |
| 65 | + |
| 66 | + def describe(self) -> str: |
| 67 | + """As a run prints it: ``<org>/<repo> at <sha12>, bundle <path>``.""" |
| 68 | + return f"{self.repo} at {self.sha[:12]}, bundle {self.path or 'at its root'}" |
| 69 | + |
| 70 | + |
| 71 | +@dataclass(frozen=True) |
| 72 | +class HubBundle: |
| 73 | + """A bundle downloaded from the Hub: its folder, the id root it runs under, where it came from, and what the repo's |
| 74 | + card says it needs: the plugins it requires when it has no bundle.toml, and the command that sets it up.""" |
| 75 | + |
| 76 | + root: Path |
| 77 | + id_root: str |
| 78 | + source: HubSource |
| 79 | + plugins: tuple[str, ...] |
| 80 | + setup: str | None |
| 81 | + |
| 82 | + @property |
| 83 | + def setup_note(self) -> str | None: |
| 84 | + """What a run refused for something setting the bundle up provides ends with: the card's command for it.""" |
| 85 | + return f"{self.source.repo}'s card sets this bundle up with: {self.setup}" if self.setup else None |
| 86 | + |
| 87 | + |
| 88 | +@dataclass |
| 89 | +class _Commit: |
| 90 | + """What this machine holds of a commit, kept in its folder: the commit's file list, and the bundles downloaded.""" |
| 91 | + |
| 92 | + files: list[str] |
| 93 | + downloaded: list[str] = field(default_factory=list) |
| 94 | + |
| 95 | + @staticmethod |
| 96 | + def path(folder: Path) -> Path: |
| 97 | + return folder / ".cache" / "agentenv" / "commit.json" |
| 98 | + |
| 99 | + @classmethod |
| 100 | + def read(cls, folder: Path) -> _Commit | None: |
| 101 | + try: |
| 102 | + return cls(**json.loads(cls.path(folder).read_text())) |
| 103 | + except FileNotFoundError: |
| 104 | + return None |
| 105 | + |
| 106 | + def save(self, folder: Path) -> None: |
| 107 | + path = self.path(folder) |
| 108 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 109 | + partial = path.with_name(f"{path.name}.{os.getpid()}") |
| 110 | + partial.write_text(json.dumps({"files": self.files, "downloaded": self.downloaded})) |
| 111 | + os.replace(partial, path) |
| 112 | + |
| 113 | + |
| 114 | +def is_hub_reference(text: str) -> bool: |
| 115 | + return text.startswith(SCHEME) |
| 116 | + |
| 117 | + |
| 118 | +def parse_reference(text: str) -> HubReference: |
| 119 | + """``text``, ``hf://<org>/<repo>[@<revision>][/<path>]``, as its parts. The revision runs to the next ``/``, so one |
| 120 | + holding a ``/`` is written percent-encoded, as the Hub's own ``hf://`` paths take it (``@refs%2Fpr%2F3``). Raises |
| 121 | + BundleError.""" |
| 122 | + org, _, rest = text.removeprefix(SCHEME).partition("/") |
| 123 | + name, _, path = rest.partition("/") |
| 124 | + name, at, revision = name.partition("@") |
| 125 | + parts = [part for part in path.split("/") if part] |
| 126 | + if not (_NAME.fullmatch(org) and _NAME.fullmatch(name)) or (at and not revision): |
| 127 | + raise BundleError([f"{text}: name a Hub dataset's bundle as {_FORM}"]) |
| 128 | + if any(part in (".", "..") for part in parts): |
| 129 | + raise BundleError([f"{text}: the path in the repo can't hold '.' or '..'"]) |
| 130 | + return HubReference(f"{org}/{name}", unquote(revision) or None, "/".join(parts) or None) |
| 131 | + |
| 132 | + |
| 133 | +def fetch_bundle(text: str, *, on_progress: Callable[[str], None] | None = None) -> HubBundle: |
| 134 | + """The bundle ``text`` names, downloaded unless this machine has it. ``on_progress`` is given the line that says |
| 135 | + which repo, commit and bundle it is, before anything is downloaded. Raises BundleError.""" |
| 136 | + reference = parse_reference(text) |
| 137 | + try: |
| 138 | + return _fetch(reference, HfApi(), on_progress or (lambda line: None)) |
| 139 | + except HfHubHTTPError as e: |
| 140 | + error = BundleError([f"{text}: {str(e).strip() or type(e).__name__}"]) |
| 141 | + if isinstance(e, RepositoryNotFoundError): |
| 142 | + error.add_note("a private or gated repo needs a token with access: log in with `hf auth login`, or set " |
| 143 | + "HF_TOKEN") |
| 144 | + raise error from e |
| 145 | + |
| 146 | + |
| 147 | +def _fetch(reference: HubReference, api: HfApi, say: Callable[[str], None]) -> HubBundle: |
| 148 | + sha = reference.revision if reference.revision and _SHA.fullmatch(reference.revision) else api.repo_info( |
| 149 | + reference.repo, repo_type="dataset", revision=reference.revision).sha |
| 150 | + folder = cache_root() / "hf" / reference.repo / sha |
| 151 | + with holding_locks([f"{SCHEME}{reference.repo}@{sha}"]): |
| 152 | + commit, card = _commit(api, reference, sha, folder) |
| 153 | + hub = _select(reference, sha, commit.files, card, folder) |
| 154 | + say(hub.source.describe()) |
| 155 | + _download(api, hub, commit, folder) |
| 156 | + return hub |
| 157 | + |
| 158 | + |
| 159 | +def _commit(api: HfApi, reference: HubReference, sha: str, folder: Path) -> tuple[_Commit, dict[str, Any]]: |
| 160 | + """What this machine holds of the commit, and its card's ``agentenv`` table: listed, with the card downloaded, |
| 161 | + the first time.""" |
| 162 | + commit = _Commit.read(folder) |
| 163 | + if commit is None: |
| 164 | + commit = _Commit(api.list_repo_files(reference.repo, repo_type="dataset", revision=sha)) |
| 165 | + if CARD in commit.files: |
| 166 | + api.hf_hub_download(reference.repo, CARD, repo_type="dataset", revision=sha, local_dir=folder) |
| 167 | + commit.save(folder) |
| 168 | + return commit, _card(folder) |
| 169 | + |
| 170 | + |
| 171 | +def _select(reference: HubReference, sha: str, files: Sequence[str], card: Mapping[str, Any], folder: Path) -> HubBundle: |
| 172 | + """The bundle ``reference`` names in the commit of ``files``, kept under ``folder``: picked, with its id root, |
| 173 | + where it came from and what the card says it needs.""" |
| 174 | + path = _pick(reference, sha, files, card) |
| 175 | + plugins, setup = _needs(reference.repo, card, _bundle_name(reference, path)) |
| 176 | + source = HubSource(reference.repo, reference.revision or "main", sha, path) |
| 177 | + declared = _within(path, BUNDLE_TOML) in files |
| 178 | + return HubBundle(folder / path, _id_root(reference, path), source, plugins=() if declared else plugins, setup=setup) |
| 179 | + |
| 180 | + |
| 181 | +def _download(api: HfApi, hub: HubBundle, commit: _Commit, folder: Path) -> None: |
| 182 | + """Check what the bundle requires, its bundle.toml downloaded alone first when it has one, then download the |
| 183 | + bundle unless the commit's record has it.""" |
| 184 | + path = hub.source.path |
| 185 | + declared, downloaded = _within(path, BUNDLE_TOML) in commit.files, path in commit.downloaded |
| 186 | + if declared and not downloaded: |
| 187 | + api.hf_hub_download(hub.source.repo, _within(path, BUNDLE_TOML), repo_type="dataset", revision=hub.source.sha, |
| 188 | + local_dir=folder) |
| 189 | + _check_requirements(hub, declared) |
| 190 | + if not downloaded: |
| 191 | + api.snapshot_download(hub.source.repo, repo_type="dataset", revision=hub.source.sha, local_dir=folder, |
| 192 | + allow_patterns=_patterns(path)) |
| 193 | + commit.downloaded.append(path) |
| 194 | + commit.save(folder) |
| 195 | + |
| 196 | + |
| 197 | +def _check_requirements(hub: HubBundle, declared: bool) -> None: |
| 198 | + """Refuse a bundle whose bundle.toml (by a folder's rules) or, without one, card names a plugin not installed, |
| 199 | + the refusal ending with the card's setup command. Raises UnmetRequirements, or BundleError.""" |
| 200 | + if declared: |
| 201 | + requires, refused = read_requires(hub.root / BUNDLE_TOML) |
| 202 | + if refused: |
| 203 | + raise BundleError([f"{BUNDLE_TOML}: {problem}" for problem in refused]) |
| 204 | + problems, unmet = unmet_problems(requires, BUNDLE_TOML) |
| 205 | + else: |
| 206 | + problems, unmet = unmet_problems(hub.plugins, f"{hub.source.repo}'s card") |
| 207 | + if not problems: |
| 208 | + return |
| 209 | + if not unmet: |
| 210 | + raise BundleError(problems) |
| 211 | + error = UnmetRequirements(problems, unmet) |
| 212 | + if hub.setup_note: |
| 213 | + error.add_note(hub.setup_note) |
| 214 | + raise error |
| 215 | + |
| 216 | + |
| 217 | +def _card(folder: Path) -> dict[str, Any]: |
| 218 | + """The ``agentenv`` table of the card's metadata, ``{}`` without one.""" |
| 219 | + readme = folder / CARD |
| 220 | + table = (metadata_load(readme) or {}).get("agentenv") if readme.is_file() else None |
| 221 | + return table if isinstance(table, dict) else {} |
| 222 | + |
| 223 | + |
| 224 | +def _needs(repo: str, card: Mapping[str, Any], name: str) -> tuple[tuple[str, ...], str | None]: |
| 225 | + """What the card's ``agentenv: bundles: <name>`` table says the bundle needs: the plugins it requires, and the |
| 226 | + command that sets it up. Raises BundleError for a value of the wrong shape.""" |
| 227 | + where = f"{repo}'s card: agentenv: bundles" |
| 228 | + bundles = card.get("bundles") or {} |
| 229 | + if not isinstance(bundles, Mapping): |
| 230 | + raise BundleError([f"{where} must map each bundle's name to what it needs"]) |
| 231 | + needs = bundles.get(name) or {} |
| 232 | + if not isinstance(needs, Mapping): |
| 233 | + raise BundleError([f"{where}: {name} must be a table of plugins and setup"]) |
| 234 | + plugins, setup = needs.get("plugins") or [], needs.get("setup") |
| 235 | + if not isinstance(plugins, list) or not all(isinstance(plugin, str) for plugin in plugins): |
| 236 | + raise BundleError([f"{where}: {name}: plugins must be a list of pip requirements"]) |
| 237 | + if setup is not None and not isinstance(setup, str): |
| 238 | + raise BundleError([f"{where}: {name}: setup must be a command"]) |
| 239 | + return tuple(plugins), setup |
| 240 | + |
| 241 | + |
| 242 | +def _pick(reference: HubReference, sha: str, files: Sequence[str], card: Mapping[str, Any]) -> str: |
| 243 | + """The folder in the repo of the bundle ``reference`` names, ``""`` for the root: its path; else the only bundle |
| 244 | + under ``bundles/``, or the one the card names its default; else the root when it is a bundle. ``bundles/`` comes |
| 245 | + first, since a dataset can keep a ``tasks/`` of data files at its root.""" |
| 246 | + where = f"{reference.repo} at {sha[:12]}" |
| 247 | + if reference.path is not None: |
| 248 | + if not _is_bundle(files, reference.path): |
| 249 | + raise BundleError([f"{where}: {reference.path}/ isn't a bundle: it has none of {_kinds()}"]) |
| 250 | + return reference.path |
| 251 | + names = sorted({parts[1] for file in files if len(parts := file.split("/")) > 2 and parts[0] == BUNDLES}) |
| 252 | + if len(names) == 1: |
| 253 | + return _within(BUNDLES, names[0]) |
| 254 | + if card.get("default") in names: |
| 255 | + return _within(BUNDLES, card["default"]) |
| 256 | + if names: |
| 257 | + revision = f"@{reference.revision}" if reference.revision else "" |
| 258 | + runs = ", ".join(f"{SCHEME}{reference.repo}{revision}/{BUNDLES}/{name}" for name in names) |
| 259 | + raise BundleError([f"{where} holds several bundles under {BUNDLES}/ and its card names no default " |
| 260 | + f"(agentenv: default); run one of {runs}"]) |
| 261 | + if _is_bundle(files, ""): |
| 262 | + return "" |
| 263 | + raise BundleError([f"{where} holds no bundle: it has no {BUNDLES}/<name>/ folder, and none of {_kinds()} at its " |
| 264 | + "root"]) |
| 265 | + |
| 266 | + |
| 267 | +def _is_bundle(files: Sequence[str], path: str) -> bool: |
| 268 | + """Whether the folder at ``path`` in the repo holds a kind folder, as a bundle does.""" |
| 269 | + prefix = _within(path, "") |
| 270 | + return any(file.startswith(prefix) and file[len(prefix):].split("/", 1)[0] in _KINDS |
| 271 | + and "/" in file[len(prefix):] for file in files) |
| 272 | + |
| 273 | + |
| 274 | +def _patterns(path: str) -> list[str]: |
| 275 | + """What to download of the bundle at ``path``: its folder, or for the repo's root its kind folders and files.""" |
| 276 | + if path: |
| 277 | + return [f"{path}/*"] |
| 278 | + return [*(f"{kind}/*" for kind in sorted(_KINDS)), BUNDLE_TOML, ALIASES_FILE] |
| 279 | + |
| 280 | + |
| 281 | +def _id_root(reference: HubReference, path: str) -> str: |
| 282 | + """``@hf/<org>/<repo>``, then for a bundle below the root its name.""" |
| 283 | + id_root = f"{HF_PREFIX}{reference.repo}" + (f"/{_bundle_name(reference, path)}" if path else "") |
| 284 | + try: |
| 285 | + validate_local_id(id_root) |
| 286 | + except ValueError as e: |
| 287 | + raise BundleError([f"{SCHEME}{reference.repo}/{path}: can't be an id root: {e}"]) from None |
| 288 | + return id_root |
| 289 | + |
| 290 | + |
| 291 | +def _bundle_name(reference: HubReference, path: str) -> str: |
| 292 | + """The bundle's name, which keys its table in the card and ends its id root: its folder under ``bundles/``, its |
| 293 | + path for a folder elsewhere, or the repo's name for the root.""" |
| 294 | + parts = path.split("/") |
| 295 | + return parts[1] if len(parts) == 2 and parts[0] == BUNDLES else path or reference.repo.partition("/")[2] |
| 296 | + |
| 297 | + |
| 298 | +def _within(folder: str, name: str) -> str: |
| 299 | + return f"{folder}/{name}" if folder else name |
| 300 | + |
| 301 | + |
| 302 | +def _kinds() -> str: |
| 303 | + return ", ".join(f"{kind.value}/" for kind in BundleKind) |
0 commit comments