fix(hf): fix download request to hf download

Author: kshitijrajsharma

infra/run-artifact-refresh.sh

@@ -11,7 +11,7 @@ dsn="${OSMSG_PSQL_DSN:-postgresql://osmsg:osmsg@db:5432/osmsg}"
 repo="${OSMSG_HISTORY_REPO:-kshitijrajsharma/osmsg-history}"
 
 echo "[artifact-refresh] advancing ${artifact_host} from ${repo}"
-docker compose run --rm -v "${artifact_host}:/artifact" --entrypoint osmsg worker \
+docker compose run --rm -e HF_XET_HIGH_PERFORMANCE=1 -v "${artifact_host}:/artifact" --entrypoint osmsg worker \
   maintain refresh --artifact-dir /artifact --repo "${repo}"
 
 echo "[artifact-refresh] reloading the API onto the advanced frontier"

osmsg/maintain/refresh.py

@@ -9,49 +9,27 @@
 import shutil
 
 import duckdb
-import requests
+import huggingface_hub
 
 from ..exceptions import OsmsgError
 from ..history import fetch_manifest
-from ..ui import info, warn
+from ..ui import info
 
 UTC = dt.UTC
 VERIFY_TOLERANCE = dt.timedelta(days=2)
-_HF_DATASETS = "https://huggingface.co/datasets"
-_DOWNLOAD_CHUNK = 1 << 20
-_DOWNLOAD_RETRIES = 5
-_INTERRUPTED = (requests.ConnectionError, requests.Timeout, requests.exceptions.ChunkedEncodingError)
 
 _ROLLUP_FILES = {
     "hashtag_changeset.parquet": "rollup/hashtag_changeset/data.parquet",
     "users.parquet": "rollup/users/data.parquet",
 }
 
 
-def _download_file(repo: str, remote: str, dest: pathlib.Path) -> pathlib.Path:
-    """Stream a published file from the dataset's resolve/main URL to dest, preserving its exact bytes
-    (and so its row-group layout). Resumes with a Range request when the connection drops mid-stream, so
-    the multi-GB rollup survives a flaky link; raises once the retry budget is spent, before the swap."""
-    url = f"{_HF_DATASETS}/{repo}/resolve/main/{remote}"
-    stalls = 0
-    while stalls < _DOWNLOAD_RETRIES:
-        have = dest.stat().st_size if dest.exists() else 0
-        headers = {"Range": f"bytes={have}-"} if have else {}
-        try:
-            with requests.get(url, headers=headers, stream=True, timeout=60) as response:
-                response.raise_for_status()
-                resuming = have > 0 and response.status_code == 206
-                with open(dest, "ab" if resuming else "wb") as handle:
-                    for chunk in response.iter_content(_DOWNLOAD_CHUNK):
-                        handle.write(chunk)
-            return dest
-        except _INTERRUPTED as exc:
-            now = dest.stat().st_size if dest.exists() else 0
-            stalls = 0 if now > have else stalls + 1
-            if stalls >= _DOWNLOAD_RETRIES:
-                raise OsmsgError(f"{remote}: download stalled after {stalls} attempts with no progress: {exc}") from exc
-            warn(f"{remote}: interrupted ({type(exc).__name__}); resuming from {now:,} bytes")
-    raise OsmsgError(f"{remote}: download failed")
+def _download(repo: str, remote: str, into_dir: pathlib.Path) -> pathlib.Path:
+    """Download a published dataset file into into_dir. huggingface_hub handles resume, retry, and the
+    parallel chunked transfer that a single stream cannot sustain over a long-haul link."""
+    return pathlib.Path(
+        huggingface_hub.hf_hub_download(repo_id=repo, filename=remote, repo_type="dataset", local_dir=str(into_dir))
+    )
 
 
 def _rollup_bounds(parquet: pathlib.Path) -> tuple[int, dt.datetime | None]:
@@ -99,10 +77,8 @@ def refresh_artifact(repo: str, artifact_dir: pathlib.Path) -> bool:
         shutil.rmtree(scratch)
     scratch.mkdir()
 
-    downloaded = {
-        name: _download_file(repo, remote_path, scratch / name) for name, remote_path in _ROLLUP_FILES.items()
-    }
-    manifest = _download_file(repo, "manifest.json", scratch / "manifest.json")
+    downloaded = {name: _download(repo, remote_path, scratch) for name, remote_path in _ROLLUP_FILES.items()}
+    manifest = _download(repo, "manifest.json", scratch)
     _verify(downloaded["hashtag_changeset.parquet"], remote.frontier, artifact_dir / "hashtag_changeset.parquet")
     for name, path in downloaded.items():
         os.replace(path, artifact_dir / name)

pyproject.toml

@@ -11,6 +11,7 @@ license-files = ["LICENSE"]
 requires-python = ">=3.11"
 dependencies = [
     "duckdb>=1.5.2",
+    "huggingface-hub>=1.0",
     "osmium>=4.3.1",
     "platformdirs>=4.5.1",
     "pyarrow>=24.0.0",

tests/test_maintain_refresh.py

@@ -6,7 +6,6 @@
 
 import duckdb
 import pytest
-import requests
 
 from osmsg.exceptions import OsmsgError
 from osmsg.history import Manifest
@@ -15,55 +14,6 @@
 UTC = dt.UTC
 
 
-class _FakeResponse:
-    def __init__(self, status_code, chunks, break_after=None):
-        self.status_code = status_code
-        self._chunks = chunks
-        self._break_after = break_after
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *args):
-        return False
-
-    def raise_for_status(self):
-        pass
-
-    def iter_content(self, chunk_size):
-        for index, chunk in enumerate(self._chunks):
-            if self._break_after is not None and index == self._break_after:
-                raise requests.exceptions.ChunkedEncodingError("connection broken")
-            yield chunk
-
-
-def test_download_resumes_after_interruption(tmp_path, monkeypatch):
-    dest = tmp_path / "rollup.parquet"
-    calls = []
-
-    def fake_get(url, headers=None, stream=True, timeout=60):
-        calls.append(headers or {})
-        if not headers:
-            return _FakeResponse(200, [b"aa", b"bb", b"cc"], break_after=2)
-        return _FakeResponse(206, [b"cc", b"dd"])
-
-    monkeypatch.setattr(refresh.requests, "get", fake_get)
-    refresh._download_file("test/repo", "rollup/x/data.parquet", dest)
-    assert dest.read_bytes() == b"aabbccdd"
-    assert calls[1]["Range"] == "bytes=4-"
-
-
-def test_download_raises_after_retry_budget(tmp_path, monkeypatch):
-    dest = tmp_path / "rollup.parquet"
-
-    def always_break(url, headers=None, stream=True, timeout=60):
-        return _FakeResponse(200, [b"aa", b"bb"], break_after=0)
-
-    monkeypatch.setattr(refresh.requests, "get", always_break)
-    with pytest.raises(OsmsgError, match="stalled after"):
-        refresh._download_file("test/repo", "rollup/x/data.parquet", dest)
-
-
 def _month_start(year, month):
     return dt.datetime(year, month, 1, tzinfo=UTC)
 
@@ -77,9 +27,11 @@ def _write_rollup(path, latest, rows):
 
 
 def _fake_download(latest, manifest_dict):
-    """Stand in for the HTTP download: materialize each requested file at its scratch destination."""
+    """Stand in for huggingface_hub.hf_hub_download: materialize each requested file under into_dir."""
 
-    def _download(repo, remote, dest):
+    def _download(repo, remote, into_dir):
+        dest = into_dir / remote
+        dest.parent.mkdir(parents=True, exist_ok=True)
         if remote.endswith("manifest.json"):
             dest.write_text(json.dumps(manifest_dict))
         elif "hashtag_changeset" in remote:
@@ -108,7 +60,7 @@ def test_noop_when_already_current(tmp_path, monkeypatch):
     def _boom(*args, **kwargs):
         raise AssertionError("must not download when already current")
 
-    monkeypatch.setattr(refresh, "_download_file", _boom)
+    monkeypatch.setattr(refresh, "_download", _boom)
     assert refresh.refresh_artifact("test/repo", artifact) is False
 
 
@@ -122,7 +74,7 @@ def test_happy_path_advances_and_swaps(tmp_path, monkeypatch):
     new_manifest = {"schema_version": 1, "min_month": "2005-04", "max_month": "2026-07"}
     _patch_manifests(monkeypatch, artifact, remote_frontier=_month_start(2026, 8), local_frontier=_month_start(2026, 7))
     monkeypatch.setattr(
-        refresh, "_download_file", _fake_download(dt.datetime(2026, 7, 31, 23, 59, tzinfo=UTC), new_manifest)
+        refresh, "_download", _fake_download(dt.datetime(2026, 7, 31, 23, 59, tzinfo=UTC), new_manifest)
     )
 
     assert refresh.refresh_artifact("test/repo", artifact) is True
@@ -140,7 +92,7 @@ def test_short_rollup_is_rejected_and_live_files_untouched(tmp_path, monkeypatch
     _write_rollup(artifact / "hashtag_changeset.parquet", _month_start(2026, 6), rows=5)
     _patch_manifests(monkeypatch, artifact, remote_frontier=_month_start(2026, 8), local_frontier=_month_start(2026, 7))
     monkeypatch.setattr(
-        refresh, "_download_file", _fake_download(dt.datetime(2026, 6, 15, tzinfo=UTC), {"max_month": "2026-07"})
+        refresh, "_download", _fake_download(dt.datetime(2026, 6, 15, tzinfo=UTC), {"max_month": "2026-07"})
     )
 
     with pytest.raises(OsmsgError, match="short of frontier"):
@@ -154,7 +106,7 @@ def test_shrunk_rollup_is_rejected(tmp_path, monkeypatch):
     _write_rollup(artifact / "hashtag_changeset.parquet", _month_start(2026, 6), rows=50)
     _patch_manifests(monkeypatch, artifact, remote_frontier=_month_start(2026, 8), local_frontier=_month_start(2026, 7))
     monkeypatch.setattr(
-        refresh, "_download_file", _fake_download(dt.datetime(2026, 7, 31, tzinfo=UTC), {"max_month": "2026-07"})
+        refresh, "_download", _fake_download(dt.datetime(2026, 7, 31, tzinfo=UTC), {"max_month": "2026-07"})
     )
 
     with pytest.raises(OsmsgError, match="fewer rows"):