Archive private KB bundles for fast cold starts

app/config.py

@@ -152,11 +152,10 @@ def _snapshot_bundle(repo_id: str, *, token: str | bool | None) -> None:
 
     Mutes httpx's per-file flood during the cold-start download only. The
     GraphRAG experiment index (~150 MB) lives in the private repo for
-    reproducibility but prod does not use it; skip it on cold start. Pull it
-    explicitly to run that eval. ``README.md`` is the public repo's dataset
-    card, not runtime data. (The private bundle has neither a kb archive nor
-    a README, and the public one has no graphrag/ tree, so unmatched patterns
-    are no-ops.)
+    reproducibility but prod does not use it; skip it on cold start. The
+    contextual-nodes pickle is another rebuild-only artifact. Pull either
+    explicitly for data/eval work. ``README.md`` is the public repo's dataset
+    card, not runtime data. Both runtime bundles ship ``kb.tar.gz``.
     """
     from huggingface_hub import snapshot_download
 
@@ -169,7 +168,12 @@ def _snapshot_bundle(repo_id: str, *, token: str | bool | None) -> None:
             local_dir="data",
             repo_type="dataset",
             token=token,
-            ignore_patterns=["graphrag/**", "README.md"],
+            ignore_patterns=[
+                "graphrag/**",
+                "all_sources_contextual_nodes.pkl",
+                ".gitattributes",
+                "README.md",
+            ],
         )
     finally:
         httpx_logger.setLevel(previous_level)
@@ -179,10 +183,9 @@ def _snapshot_bundle(repo_id: str, *, token: str | bool | None) -> None:
 def _extract_kb_archive(base_dir: str = "data") -> None:
     """Extract ``kb.tar.gz`` into ``data/kb`` and delete the archive.
 
-    The public bundle ships the KB as one archive instead of ~3,000 files so
-    an anonymous cold start is not throttled by HF's per-request rate limits
-    (see build_public_docs_bundle.archive_kb). A no-op when no archive was
-    downloaded (the private bundle ships the unpacked tree).
+    Both private and public runtime bundles ship the KB as one archive instead
+    of ~3,000 files, avoiding per-file Hub overhead during cold starts. A no-op
+    for legacy bundles that still contain the unpacked tree.
     """
     import tarfile
 

data/scraping_scripts/README.md

@@ -212,8 +212,10 @@ place on startup either way).
 
 - **Built** by `data.scraping_scripts.build_kb_artifacts` (from
   `data/all_sources_data.jsonl`) and `data.scraping_scripts.update_kb_wiki`.
-- **Uploaded** to `towardsai-tutors/ai-tutor-vector-db` by
-  `data.scraping_scripts.upload_dbs_to_hf` (already includes `kb/**`).
+- **Archived and uploaded** to `towardsai-tutors/ai-tutor-vector-db` by
+  `data.scraping_scripts.upload_dbs_to_hf` as `kb.tar.gz`. The local unpacked
+  tree is preserved; legacy remote `kb/**` files are pruned only after the
+  archive upload is verified.
 - **Downloaded** by `app.config.ensure_local_vector_db` on the first
   chatbot start (or any start where the local KB is missing).
 
@@ -323,8 +325,12 @@ After the audit, the KB tree is packed into a single **`kb.tar.gz`** and the
 unpacked tree is dropped from staging: the public repo ships ~6 files instead
 of ~3,000, which keeps a token-free cold start clear of HF's anonymous
 per-request rate limits (~15 min unpacked vs a few min archived). The runtime
-extracts the archive after download (`app.config._extract_kb_archive`); the
-private bundle still ships the unpacked KB tree and is unaffected.
+extracts the archive after download (`app.config._extract_kb_archive`).
+
+The private production publisher now uses the same archive layout: it creates
+`data/kb.tar.gz` transiently, uploads and verifies it, then prunes the legacy
+remote `kb/**` tree. It never removes or rewrites the canonical local
+`data/kb/` directory.
 
 The staged tree lives in `data/public_docs_bundle/` (gitignored). Uploading
 needs an `HF_TOKEN` with write access to the `towardsai-tutors` org; the script
@@ -368,4 +374,6 @@ download small).
    - `update_kb_wiki.py` refreshes wiki navigation pages under `data/kb/wiki/`, preserving maintainer prose outside `<!-- AUTO-GENERATED -->` markers
    - `lint_kb_wiki.py` then verifies every KB path and relative link the wiki pages reference (existence on disk plus `corpus_manifest.jsonl` membership for `raw/` files) and fails the workflow before upload on any broken reference; run it standalone with `uv run -m data.scraping_scripts.lint_kb_wiki [--kb-dir data/kb]`
    - The docs and course workflows run these steps automatically unless `--skip-kb` is passed
-   - `upload_dbs_to_hf.py` uploads `data/kb/**` with the vector database so the runtime can download one KB bundle
+   - `upload_dbs_to_hf.py` packs `data/kb/` into `kb.tar.gz`, uploads it with
+     the vector database, verifies the remote archive, and only then prunes
+     the old unpacked remote KB files

data/scraping_scripts/upload_dbs_to_hf.py

@@ -1,7 +1,9 @@
-"""Upload the Chroma vector database to a Hugging Face dataset repository."""
+"""Publish the private vector database and archived KB to Hugging Face."""
 
 import argparse
+import gzip
 import os
+import tarfile
 from pathlib import Path
 
 from dotenv import load_dotenv
@@ -18,30 +20,93 @@ load_dotenv()
 DEFAULT_REPO_ID = "towardsai-tutors/ai-tutor-vector-db"
 
 FOLDER_PATH = "data"
+KB_DIR_NAME = "kb"
+KB_ARCHIVE_NAME = "kb.tar.gz"
 ALLOW_PATTERNS = [
     "chroma-db-all_sources/**",
     "all_sources_contextual_nodes.pkl",
-    "kb/**",
+    KB_ARCHIVE_NAME,
 ]
+# The upload contains only the archive, but the prune pass also owns the
+# unpacked tree left by older publishes. Keeping these scopes separate lets us
+# delete kb/** without re-uploading it from the local build artifacts.
+PRUNE_ALLOW_PATTERNS = [*ALLOW_PATTERNS, f"{KB_DIR_NAME}/**"]
 IGNORE_PATTERNS = ["*.py", "*.ipynb", "*.pyc", "__pycache__/**"]
 
 
+def _normalized_tar_info(info: tarfile.TarInfo) -> tarfile.TarInfo:
+    """Remove host-specific metadata so unchanged KBs produce one stable blob."""
+    info.uid = 0
+    info.gid = 0
+    info.uname = ""
+    info.gname = ""
+    info.mtime = 0
+    info.pax_headers = {}
+    if info.isdir():
+        info.mode = 0o755
+    elif info.isfile():
+        info.mode = 0o644
+    return info
+
+
+def create_kb_archive(folder_path: str = FOLDER_PATH) -> Path:
+    """Build ``kb.tar.gz`` atomically without modifying the local KB tree."""
+    root = Path(folder_path)
+    kb_dir = root / KB_DIR_NAME
+    required = (
+        kb_dir / "generated" / "corpus_manifest.jsonl",
+        kb_dir / "wiki" / "index.md",
+    )
+    missing = [str(path) for path in required if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(
+            "Cannot publish an incomplete KB; missing: " + ", ".join(missing)
+        )
+
+    archive_path = root / KB_ARCHIVE_NAME
+    temp_path = archive_path.with_name(f"{archive_path.name}.tmp")
+    temp_path.unlink(missing_ok=True)
+    print(f"Archiving {kb_dir} -> {archive_path} ...")
+    try:
+        with temp_path.open("wb") as raw:
+            # Fix the gzip timestamp/name and tar ownership/timestamps so an
+            # unchanged KB hashes identically across machines and publishes.
+            with gzip.GzipFile(
+                filename="", mode="wb", fileobj=raw, mtime=0
+            ) as compressed:
+                with tarfile.open(fileobj=compressed, mode="w") as tar:
+                    tar.add(
+                        kb_dir,
+                        arcname=KB_DIR_NAME,
+                        filter=_normalized_tar_info,
+                    )
+        os.replace(temp_path, archive_path)
+    except BaseException:
+        temp_path.unlink(missing_ok=True)
+        raise
+
+    print(
+        f"KB archive ready: {archive_path} ({archive_path.stat().st_size / 1e6:.1f} MB)"
+    )
+    return archive_path
+
+
 def _prune_stale_remote_files(
     api: HfApi,
     repo_id: str,
     *,
     folder_path: str,
     allow_patterns: list[str],
+    prune_allow_patterns: list[str] | None,
     ignore_patterns: list[str],
 ) -> None:
-    """Restore the auto-prune behavior we lost when switching from
-    `upload_folder(..., delete_patterns=["*"])` to `upload_large_folder`.
-
-    Compute the set of remote files that match `allow_patterns` (and don't
-    match `ignore_patterns`) but are no longer present locally, and delete
-    them in a single commit before the upload runs. We deliberately do this
-    in its own commit so the subsequent `upload_large_folder` stays
-    additive and resumable.
+    """Verify the upload, then prune files owned by this publisher.
+
+    ``allow_patterns`` describes files expected from the current upload.
+    ``prune_allow_patterns`` may be wider, which is how the archive migration
+    deletes legacy remote ``kb/**`` files without treating the local unpacked
+    tree as publishable. Refuse to delete anything unless every expected local
+    upload is visible remotely.
     """
     local_root = Path(folder_path)
     local_files_all = [
@@ -57,11 +122,22 @@ def _prune_stale_remote_files(
         )
     )
 
-    remote_files_all = api.list_repo_files(repo_id, repo_type="dataset")
+    remote_files_all = set(api.list_repo_files(repo_id, repo_type="dataset"))
+    missing_remote = sorted(local_kept - remote_files_all)
+    if missing_remote:
+        preview = ", ".join(missing_remote[:10])
+        suffix = "" if len(missing_remote) <= 10 else ", ..."
+        raise RuntimeError(
+            "Upload verification failed; refusing to prune because "
+            f"{len(missing_remote)} expected remote file(s) are missing: "
+            f"{preview}{suffix}"
+        )
+
+    prune_allow = prune_allow_patterns or allow_patterns
     remote_kept = list(
         filter_repo_objects(
             remote_files_all,
-            allow_patterns=allow_patterns,
+            allow_patterns=prune_allow,
             ignore_patterns=ignore_patterns,
         )
     )
@@ -71,9 +147,11 @@ def _prune_stale_remote_files(
         print("Prune step: no stale remote files to delete.")
         return
 
-    print(f"Prune step: deleting {len(stale)} stale remote file(s):")
-    for f in stale:
+    print(f"Prune step: deleting {len(stale)} stale remote file(s).")
+    for f in stale[:20]:
         print(f"  - {f}")
+    if len(stale) > 20:
+        print(f"  ... and {len(stale) - 20} more")
 
     api.create_commit(
         repo_id=repo_id,
@@ -88,10 +166,11 @@ def upload_bundle(
     *,
     folder_path: str = FOLDER_PATH,
     allow_patterns: list[str] | None = None,
+    prune_allow_patterns: list[str] | None = None,
     ignore_patterns: list[str] | None = None,
     create_public: bool = False,
 ) -> None:
-    """Prune stale remote files, then upload ``folder_path`` to ``repo_id``.
+    """Upload ``folder_path``, verify it, then prune stale remote files.
 
     The default arguments reproduce the production upload (the private
     all-sources bundle). ``build_public_docs_bundle`` reuses this with a
@@ -114,32 +193,42 @@ def upload_bundle(
             print(exc)
             raise SystemExit(1) from exc
 
-    # Step 1: prune. Restore the auto-cleanup we lost when moving off
-    # `upload_folder(..., delete_patterns=["*"])`. See helper docstring.
-    _prune_stale_remote_files(
-        api,
-        repo_id,
+    # Step 1: upload additively. Stale files remain available until the
+    # verification and explicit prune below, so a failed archive migration
+    # cannot strand the repo without either KB representation.
+    api.upload_folder(
         folder_path=folder_path,
+        repo_id=repo_id,
+        repo_type="dataset",
         allow_patterns=allow,
         ignore_patterns=ignore,
     )
 
-    # Step 2: upload. `upload_large_folder` is the recommended path once the
-    # total payload crosses ~1 GB — it commits in many small batches instead
-    # of one big atomic commit, which sidesteps the `commit/main` 500s
-    # `upload_folder` tends to hit at this size. It is resumable: re-running
-    # after a failure skips already-uploaded blobs.
-    api.upload_large_folder(
+    # Step 2: verify the expected files reached the repo, then prune. Uploading
+    # first is load-bearing for archive migrations: the old unpacked KB stays
+    # usable if the archive upload fails.
+    _prune_stale_remote_files(
+        api,
+        repo_id,
         folder_path=folder_path,
-        repo_id=repo_id,
-        repo_type="dataset",
         allow_patterns=allow,
+        prune_allow_patterns=prune_allow_patterns,
         ignore_patterns=ignore,
     )
 
 
 def upload_vector_db(repo_id: str = DEFAULT_REPO_ID) -> None:
-    upload_bundle(repo_id)
+    archive_path = create_kb_archive(FOLDER_PATH)
+    try:
+        upload_bundle(
+            repo_id,
+            allow_patterns=ALLOW_PATTERNS,
+            prune_allow_patterns=PRUNE_ALLOW_PATTERNS,
+        )
+    finally:
+        # The archive is a transient publishing artifact. The canonical local
+        # KB stays unpacked for development and maintenance workflows.
+        archive_path.unlink(missing_ok=True)
 
 
 GRAPHRAG_LOCAL_DIR = "data/graphrag/output"

tests/test_config.py

@@ -187,6 +187,28 @@ def test_download_bundle_falls_back_to_public_on_no_access() -> None:
     ]
 
 
+def test_snapshot_bundle_skips_rebuild_only_artifacts() -> None:
+    with (
+        patch("huggingface_hub.snapshot_download") as snapshot_download,
+        patch("app.config._extract_kb_archive") as extract,
+    ):
+        config._snapshot_bundle("org/bundle", token="tok")
+
+    snapshot_download.assert_called_once_with(
+        repo_id="org/bundle",
+        local_dir="data",
+        repo_type="dataset",
+        token="tok",
+        ignore_patterns=[
+            "graphrag/**",
+            "all_sources_contextual_nodes.pkl",
+            ".gitattributes",
+            "README.md",
+        ],
+    )
+    extract.assert_called_once_with()
+
+
 def test_extract_kb_archive_unpacks_and_removes_archive(tmp_path: Path) -> None:
     import tarfile
 

tests/test_upload_dbs_to_hf.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import hashlib
+import tarfile
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from data.scraping_scripts import upload_dbs_to_hf as uploader
+
+
+def _write_minimal_kb(root: Path) -> None:
+    kb = root / "kb"
+    (kb / "generated").mkdir(parents=True)
+    (kb / "generated" / "corpus_manifest.jsonl").write_text(
+        '{"doc_id":"1"}\n', encoding="utf-8"
+    )
+    (kb / "wiki").mkdir()
+    (kb / "wiki" / "index.md").write_text("# Index\n", encoding="utf-8")
+
+
+def test_create_kb_archive_is_deterministic_and_preserves_tree(
+    tmp_path: Path,
+) -> None:
+    _write_minimal_kb(tmp_path)
+
+    archive = uploader.create_kb_archive(str(tmp_path))
+    first_digest = hashlib.sha256(archive.read_bytes()).hexdigest()
+    archive = uploader.create_kb_archive(str(tmp_path))
+    second_digest = hashlib.sha256(archive.read_bytes()).hexdigest()
+
+    assert first_digest == second_digest
+    assert (tmp_path / "kb" / "wiki" / "index.md").is_file()
+    with tarfile.open(archive, "r:gz") as tar:
+        assert "kb/generated/corpus_manifest.jsonl" in tar.getnames()
+        assert "kb/wiki/index.md" in tar.getnames()
+
+
+def test_create_kb_archive_refuses_incomplete_kb(tmp_path: Path) -> None:
+    (tmp_path / "kb").mkdir()
+
+    with pytest.raises(FileNotFoundError, match="incomplete KB"):
+        uploader.create_kb_archive(str(tmp_path))
+
+    assert not (tmp_path / "kb.tar.gz").exists()
+
+
+class _FakeApi:
+    def __init__(self, remote_files: list[str]) -> None:
+        self.remote_files = remote_files
+        self.events: list[str] = []
+        self.deleted: list[str] = []
+
+    def upload_folder(self, **_kwargs) -> None:
+        self.events.append("upload")
+
+    def list_repo_files(self, _repo_id: str, *, repo_type: str) -> list[str]:
+        assert repo_type == "dataset"
+        self.events.append("list")
+        return self.remote_files
+
+    def create_commit(self, **kwargs) -> None:
+        self.events.append("prune")
+        self.deleted = [operation.path_in_repo for operation in kwargs["operations"]]
+
+
+def test_upload_verifies_archive_then_prunes_legacy_kb(tmp_path: Path) -> None:
+    (tmp_path / "chroma-db-all_sources").mkdir()
+    (tmp_path / "chroma-db-all_sources" / "chroma.sqlite3").write_bytes(b"db")
+    (tmp_path / "kb.tar.gz").write_bytes(b"archive")
+    api = _FakeApi(
+        [
+            "chroma-db-all_sources/chroma.sqlite3",
+            "kb.tar.gz",
+            "kb/wiki/index.md",
+            "kb/generated/corpus_manifest.jsonl",
+        ]
+    )
+
+    with patch(
+        "data.scraping_scripts.upload_dbs_to_hf.validate_hf_access", return_value=api
+    ):
+        uploader.upload_bundle(
+            "org/private-bundle",
+            folder_path=str(tmp_path),
+            allow_patterns=["chroma-db-all_sources/**", "kb.tar.gz"],
+            prune_allow_patterns=[
+                "chroma-db-all_sources/**",
+                "kb.tar.gz",
+                "kb/**",
+            ],
+        )
+
+    assert api.events == ["upload", "list", "prune"]
+    assert api.deleted == [
+        "kb/generated/corpus_manifest.jsonl",
+        "kb/wiki/index.md",
+    ]
+
+
+def test_upload_refuses_to_prune_when_archive_is_missing_remotely(
+    tmp_path: Path,
+) -> None:
+    (tmp_path / "kb.tar.gz").write_bytes(b"archive")
+    api = _FakeApi(["kb/wiki/index.md"])
+
+    with patch(
+        "data.scraping_scripts.upload_dbs_to_hf.validate_hf_access", return_value=api
+    ):
+        with pytest.raises(RuntimeError, match="refusing to prune"):
+            uploader.upload_bundle(
+                "org/private-bundle",
+                folder_path=str(tmp_path),
+                allow_patterns=["kb.tar.gz"],
+                prune_allow_patterns=["kb.tar.gz", "kb/**"],
+            )
+
+    assert api.events == ["upload", "list"]
+    assert api.deleted == []