feat(Day 7): ShardRouter multi-repo support — admin shard via dataset repo

Free-tier 1GB LFS quota on the Space repo is exhausted at 10 shards.
Workaround without paying for Pro: spill the remaining shards into
separate HF Dataset repo(s), each with its own 1GB quota.

ShardRouter changes (~120 lines):
- Accepts `extra_repos` arg or TAU_RAG_EXTRA_SHARDS_REPOS env (CSV)
- At init: for each extra repo, snapshot_download to a local cache
(TAU_RAG_EXTRA_SHARDS_CACHE, default /tmp/legal_eye_extra_shards)
- Scans all dirs (primary + extras); primary wins on conflict
- _shard_dir maps shard_name → root, _load_shard uses that mapping
- New status field 'shards_dirs' shows all sources

Dockerfile:
- TAU_RAG_EXTRA_SHARDS_REPOS=Legal-i/legal-eye-shards-extra
- TAU_RAG_EXTRA_SHARDS_CACHE=/tmp/legal_eye_extra_shards

Verified locally: download of legal-eye-shards-extra (169MB / 5 files)
in 7.9s, router scans both dirs, reports '11 shards across 2 dir(s)'.

Next: criminal also goes to legal-eye-shards-extra once it finishes
building (~15min).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Dockerfile

@@ -54,6 +54,8 @@ ENV PYTHONPATH=/app \
     TAU_RAG_AUTOLOAD_CORPUS=1 \
     TAU_RAG_AUTOLOAD_CORPUS_PATH=/app/tau_rag/runtime/uploads/corpus_paid.jsonl \
     TAU_RAG_TIER=paid \
+    TAU_RAG_EXTRA_SHARDS_REPOS=Legal-i/legal-eye-shards-extra \
+    TAU_RAG_EXTRA_SHARDS_CACHE=/tmp/legal_eye_extra_shards \
     TAU_RAG_CLUSTER_AUGMENT=1 \
     TAU_RAG_AUTH_REQUIRED=true \
     TAU_RAG_RATE_LIMIT_QPS=3 \

tau_rag/retrieve/shard_router.py

@@ -51,6 +51,23 @@ from typing import Any, Dict, List, Optional, Tuple
 _DEFAULT_SHARDS_DIR = (Path(__file__).resolve().parent.parent
                        / "runtime" / "shards")
 
+# v2.96.8 (Day 7) — Multi-repo support. When the free-tier 1GB LFS
+# quota on the main HF Space repo is exhausted, additional shards
+# live in separate HF Dataset repos (e.g. Legal-i/legal-eye-shards-
+# extra). The router scans multiple shard dirs at startup. The
+# secondary dataset(s) are downloaded once at boot via
+# huggingface_hub.snapshot_download and cached locally.
+#
+# Configured via env:
+#   TAU_RAG_EXTRA_SHARDS_REPOS  comma-separated list of HF dataset
+#                                repo IDs to download (e.g.
+#                                "Legal-i/legal-eye-shards-extra")
+#   TAU_RAG_EXTRA_SHARDS_CACHE  local dir to cache downloads (default
+#                                /tmp/extra_shards)
+_EXTRA_REPOS_ENV = "TAU_RAG_EXTRA_SHARDS_REPOS"
+_EXTRA_CACHE_ENV = "TAU_RAG_EXTRA_SHARDS_CACHE"
+_EXTRA_CACHE_DEFAULT = "/tmp/legal_eye_extra_shards"
+
 # LRU cache size — how many shard pipelines to keep in memory at once.
 # Each shard ≈ 30-900MB depending on size. Default 4 = ~1-2GB typical.
 _DEFAULT_MAX_IN_MEMORY = int(os.environ.get(
@@ -71,7 +88,9 @@ class ShardRouter:
 
     def __init__(self,
                  shards_dir: Optional[Path] = None,
-                 max_in_memory: int = _DEFAULT_MAX_IN_MEMORY):
+                 max_in_memory: int = _DEFAULT_MAX_IN_MEMORY,
+                 extra_repos: Optional[List[str]] = None):
+        # Primary shards dir (typically the Space's bundled shards)
         self.shards_dir = Path(shards_dir or _DEFAULT_SHARDS_DIR)
         self.max_in_memory = max_in_memory
         self._lock = threading.RLock()
@@ -83,21 +102,102 @@ class ShardRouter:
         # Stats — for /v1/system/shards endpoint
         self._stats: Dict[str, Dict[str, int]] = {}
 
-        # Load top-level manifest
+        # Multi-dir support: list of (dir, source_label) pairs scanned
+        # in order. The first dir is the primary (bundled shards); any
+        # additional dirs come from HF Dataset repos downloaded at boot.
+        self._shards_dirs: List[tuple] = [(self.shards_dir, "primary")]
+
+        # Bootstrap secondary repos before scanning, so their shards
+        # are visible in available_shards from the get-go.
+        repos_env = (os.environ.get(_EXTRA_REPOS_ENV) or "").strip()
+        repos_list = extra_repos if extra_repos else (
+            [r.strip() for r in repos_env.split(",") if r.strip()]
+            if repos_env else []
+        )
+        if repos_list:
+            cache_root = Path(os.environ.get(_EXTRA_CACHE_ENV,
+                                              _EXTRA_CACHE_DEFAULT))
+            for repo_id in repos_list:
+                downloaded = self._download_extra_repo(repo_id, cache_root)
+                if downloaded:
+                    self._shards_dirs.append((downloaded, repo_id))
+
+        # Aggregate manifest across all dirs
         self.manifest: Optional[dict] = None
         self.available_shards: List[str] = []
-        mp = self.shards_dir / "manifest.json"
-        if mp.exists():
-            try:
-                self.manifest = json.loads(mp.read_text(encoding="utf-8"))
-                # Only consider shards that have a built retriever_state
-                self.available_shards = [
-                    s for s in self.manifest.get("shards", [])
-                    if (self.shards_dir / s / "retriever_state"
-                        / "manifest.json").exists()
-                ]
-            except Exception as e:
-                print(f"[shard_router] manifest load failed: {e}")
+        # Track which dir each shard lives in (for fast lookup at load)
+        self._shard_dir: Dict[str, Path] = {}
+
+        for dir_path, label in self._shards_dirs:
+            mp = dir_path / "manifest.json"
+            top_manifest = None
+            if mp.exists():
+                try:
+                    top_manifest = json.loads(mp.read_text(encoding="utf-8"))
+                except Exception as e:
+                    print(f"[shard_router] manifest load failed for "
+                          f"{label}: {e}")
+
+            # If no top-level manifest, scan subdirs directly
+            if top_manifest is None:
+                shard_names = [d.name for d in dir_path.iterdir()
+                               if d.is_dir()
+                               and (d / "retriever_state"
+                                    / "manifest.json").exists()] \
+                              if dir_path.exists() else []
+            else:
+                shard_names = top_manifest.get("shards", [])
+
+            for s in shard_names:
+                if (dir_path / s / "retriever_state"
+                        / "manifest.json").exists():
+                    if s not in self._shard_dir:  # primary wins on conflict
+                        self.available_shards.append(s)
+                        self._shard_dir[s] = dir_path
+
+            # Keep the primary manifest as the "main" one
+            if label == "primary" and top_manifest:
+                self.manifest = top_manifest
+
+        if self._shards_dirs:
+            n_extra = sum(1 for _, lbl in self._shards_dirs if lbl != "primary")
+            print(f"[shard_router] ready · {len(self.available_shards)} "
+                  f"shards across {len(self._shards_dirs)} dir(s) "
+                  f"({n_extra} extra repo(s))")
+
+    def _download_extra_repo(self, repo_id: str,
+                             cache_root: Path) -> Optional[Path]:
+        """Download a HF Dataset repo's shards dir to a local cache.
+        Returns the local path to the shards root inside the cache,
+        or None on failure. Cached: subsequent boots reuse the
+        download if `etag` hasn't changed."""
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            print(f"[shard_router] huggingface_hub not installed; "
+                  f"skipping extra repo {repo_id}")
+            return None
+        target = cache_root / repo_id.replace("/", "__")
+        target.mkdir(parents=True, exist_ok=True)
+        try:
+            print(f"[shard_router] downloading extra repo {repo_id}...")
+            t0 = time.time()
+            local = snapshot_download(
+                repo_id=repo_id,
+                repo_type="dataset",
+                local_dir=str(target),
+            )
+            elapsed = time.time() - t0
+            print(f"[shard_router] downloaded {repo_id} in "
+                  f"{elapsed:.1f}s → {local}")
+            # Expect a `shards/` subdir in the dataset repo
+            shards_root = Path(local) / "shards"
+            if shards_root.exists():
+                return shards_root
+            return Path(local)
+        except Exception as e:
+            print(f"[shard_router] failed to download {repo_id}: {e}")
+            return None
 
     def is_available(self) -> bool:
         """Does the router have any usable shards?"""
@@ -187,16 +287,41 @@ class ShardRouter:
         cfg.retrieval.enabled = list(_SHARD_RETRIEVERS)
         return Pipeline.from_config(cfg)
 
+    def _read_state_bytes(self, state_dir: Path,
+                          base_name: str) -> Optional[bytes]:
+        """Read state file bytes. Handles two layouts:
+        1. Single file: `<base_name>.pkl.gz`  → read directly
+        2. Chunked:     `<base_name>.pkl.gz.part{N}` → concatenate
+
+        Chunked form is used when the file exceeded HF Hub's
+        non-LFS file-size threshold (~50MB) and was split by
+        tau_rag/scripts/chunk_pkl.py to stay in regular git (free
+        1GB LFS quota is precious).
+
+        Returns the raw gzipped bytes, or None if neither layout
+        is present.
+        """
+        single = state_dir / f"{base_name}.pkl.gz"
+        if single.exists():
+            return single.read_bytes()
+        # Chunked form
+        parts = sorted(state_dir.glob(f"{base_name}.pkl.gz.part*"),
+                       key=lambda p: int(p.suffix.removeprefix(".part")))
+        if parts:
+            return b"".join(p.read_bytes() for p in parts)
+        return None
+
     def _load_state_into(self, pipe, state_dir: Path) -> None:
         """Pickle.load each retriever state file into the pipeline."""
+        import io
         named = getattr(pipe.retrievers, "_retrievers", None) or {}
         if not named:
             for r in getattr(pipe.retrievers, "retrievers", []):
                 named[getattr(r, "name", type(r).__name__)] = r
 
         for name, retriever in named.items():
-            pkl = state_dir / f"{name}.pkl.gz"
-            if not pkl.exists():
+            state_bytes = self._read_state_bytes(state_dir, name)
+            if state_bytes is None:
                 continue
             # Unwrap RefAware-style wrappers to reach the inner retriever
             inner = retriever
@@ -209,7 +334,7 @@ class ShardRouter:
             if not hasattr(inner, "load_state_dict"):
                 continue
             try:
-                with gzip.open(pkl, "rb") as f:
+                with gzip.open(io.BytesIO(state_bytes), "rb") as f:
                     state = pickle.load(f)
                 inner.load_state_dict(state)
             except Exception as e:
@@ -217,7 +342,10 @@ class ShardRouter:
 
     def _load_shard(self, shard_name: str):
         """Load a shard's pipeline (or return cached). Evicts LRU on
-        budget overflow. Thread-safe."""
+        budget overflow. Thread-safe.
+
+        Multi-dir aware: looks up which dir holds this shard via
+        `_shard_dir`, falls back to `shards_dir` for backwards-compat."""
         with self._lock:
             if shard_name in self._cache:
                 # Move to MRU end
@@ -227,12 +355,14 @@ class ShardRouter:
                     self._stats[shard_name].get("hits", 0) + 1
                 return self._cache[shard_name]
 
-            # Cold load
+            # Cold load — find the right dir for this shard
             t0 = time.time()
-            state_dir = self.shards_dir / shard_name / "retriever_state"
+            shard_root = self._shard_dir.get(shard_name, self.shards_dir)
+            state_dir = shard_root / shard_name / "retriever_state"
             if not state_dir.exists():
                 raise FileNotFoundError(
-                    f"shard {shard_name} has no retriever_state")
+                    f"shard {shard_name} has no retriever_state "
+                    f"in {shard_root}")
             pipe = self._build_pipeline()
             self._load_state_into(pipe, state_dir)
             load_ms = int((time.time() - t0) * 1000)