Upload hf_backend/hub_store.py with huggingface_hub

hf_backend/hub_store.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import io
+import json
+import time
+from pathlib import Path
+
+from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
+
+from hf_backend.config import AppConfig
+from hf_backend.manifest import build_download_url, empty_manifest, normalize_manifest
+
+
+class HubStore:
+    _retryable_status_codes = {429, 500, 502, 503, 504}
+    _commit_retry_attempts = 5
+    _commit_retry_delay_seconds = 1.0
+
+    def __init__(self, config: AppConfig) -> None:
+        self.config = config
+        self.api = HfApi(token=config.hf_token)
+
+    def _is_retryable_hub_error(self, exc: HfHubHTTPError) -> bool:
+        response = getattr(exc, "response", None)
+        if response is None:
+            return False
+        return response.status_code in self._retryable_status_codes
+
+    def _run_with_commit_retry(self, operation):
+        attempt = 0
+        while True:
+            try:
+                return operation()
+            except HfHubHTTPError as exc:
+                attempt += 1
+                if attempt >= self._commit_retry_attempts or not self._is_retryable_hub_error(exc):
+                    raise
+                time.sleep(self._commit_retry_delay_seconds * attempt)
+
+    def ensure_dataset_repo(self) -> None:
+        try:
+            self.api.repo_info(
+                repo_id=self.config.dataset_repo_id,
+                repo_type="dataset",
+                token=self.config.hf_token,
+            )
+            return
+        except EntryNotFoundError:
+            pass
+        except HfHubHTTPError as exc:
+            if exc.response is None or exc.response.status_code != 404:
+                raise
+
+        self.api.create_repo(
+            self.config.dataset_repo_id,
+            repo_type="dataset",
+            private=False,
+            exist_ok=True,
+        )
+
+    def load_manifest(self) -> dict:
+        try:
+            manifest_path = hf_hub_download(
+                repo_id=self.config.dataset_repo_id,
+                repo_type="dataset",
+                filename=self.config.manifest_path,
+                token=self.config.hf_token,
+                force_download=True,
+            )
+            with open(manifest_path, encoding="utf-8") as handle:
+                return normalize_manifest(json.load(handle))
+        except (EntryNotFoundError, FileNotFoundError):
+            return empty_manifest()
+        except HfHubHTTPError as exc:
+            if exc.response is not None and exc.response.status_code == 404:
+                return empty_manifest()
+            raise
+
+    def save_manifest(self, manifest: dict, *, message: str = "Update job manifest") -> dict:
+        normalized = normalize_manifest(manifest)
+        data = json.dumps(normalized, ensure_ascii=False, indent=2).encode("utf-8")
+        self.upload_bytes(data, self.config.manifest_path, commit_message=message)
+        return normalized
+
+    def upload_bytes(self, content: bytes, repo_path: str, *, commit_message: str) -> str:
+        self._run_with_commit_retry(
+            lambda: self.api.upload_file(
+                path_or_fileobj=io.BytesIO(content),
+                path_in_repo=repo_path,
+                repo_id=self.config.dataset_repo_id,
+                repo_type="dataset",
+                token=self.config.hf_token,
+                commit_message=commit_message,
+            )
+        )
+        return build_download_url(self.config.dataset_repo_id, repo_path)
+
+    def upload_local_file(self, local_path: Path, repo_path: str, *, commit_message: str) -> str:
+        self._run_with_commit_retry(
+            lambda: self.api.upload_file(
+                path_or_fileobj=str(local_path),
+                path_in_repo=repo_path,
+                repo_id=self.config.dataset_repo_id,
+                repo_type="dataset",
+                token=self.config.hf_token,
+                commit_message=commit_message,
+            )
+        )
+        return build_download_url(self.config.dataset_repo_id, repo_path)
+
+    def delete_repo_file(self, repo_path: str, *, commit_message: str) -> None:
+        self._run_with_commit_retry(
+            lambda: self.api.delete_file(
+                path_in_repo=repo_path,
+                repo_id=self.config.dataset_repo_id,
+                repo_type="dataset",
+                token=self.config.hf_token,
+                commit_message=commit_message,
+            )
+        )
+
+    def download_to_directory(self, repo_path: str, directory: Path) -> Path:
+        directory.mkdir(parents=True, exist_ok=True)
+        downloaded = hf_hub_download(
+            repo_id=self.config.dataset_repo_id,
+            repo_type="dataset",
+            filename=repo_path,
+            token=self.config.hf_token,
+            local_dir=directory,
+            force_download=True,
+        )
+        return Path(downloaded)