exec-dashboard / prewarm.py
DShah-11's picture
Executive translation-quality dashboard
14b7904
Raw
History Blame Contribute Delete
3.55 kB
#!/usr/bin/env python3
"""
Fetch the (private) eval data into a local dir before Streamlit starts serving.
The Space repo is PUBLIC, so the eval data is deliberately NOT committed to it.
Instead the data lives in a private HF dataset and is pulled at container
startup using the `HF_TOKEN` Space secret. The download is targeted — only the
two file types the dashboard reads, for the one subfolder — so it is a small
(~tens of MB) fetch that completes in seconds, well before the first visitor.
Because this runs in the Docker CMD *before* `streamlit run`, no visitor ever
sees a partial download: by the time the app serves, DATA_DIR is fully on disk
and `_resolve_eval_dir()` reads it as a local path (no network on the request
path). On a cold restart this re-fetches, but the targeted snapshot is fast.
Env / secrets:
HF_TOKEN — required to read the private dataset
HF_REPO_ID — dataset repo (default below)
HF_DEFAULT_BRANCH — branch/revision to pull
HF_DEFAULT_SUBFOLDER — subfolder within the branch to pull
DATA_DIR — where to write the data (default /app/data)
"""
import os
import shutil
import sys
def main() -> None:
repo_id = os.environ.get("HF_REPO_ID", "srmdtranslations/model_gu_en_evaluations")
branch = os.environ.get("HF_DEFAULT_BRANCH", "main")
subfolder = os.environ.get("HF_DEFAULT_SUBFOLDER", "")
token = os.environ.get("HF_TOKEN") or None
data_dir = os.environ.get("DATA_DIR", "/app/data")
if not token:
print("[prewarm] HF_TOKEN not set — skipping fetch (app will try at runtime).",
flush=True)
return
if not subfolder:
print("[prewarm] HF_DEFAULT_SUBFOLDER not set — skipping fetch.", flush=True)
return
try:
from huggingface_hub import snapshot_download
print(f"[prewarm] fetching {repo_id}@{branch}/{subfolder}{data_dir} …",
flush=True)
# Only the two file types the dashboard reads — keeps the fetch small/fast.
local_root = snapshot_download(
repo_id=repo_id,
repo_type="dataset",
revision=branch,
token=token,
allow_patterns=[
f"{subfolder}/**/*.holistic.jsonl",
f"{subfolder}/**/*.intent_entity.jsonl",
],
)
src = os.path.join(local_root, subfolder)
if not os.path.isdir(src):
print(f"[prewarm] subfolder missing after download: {src}", flush=True)
return
# Copy the matching files into DATA_DIR (resolve symlinks so the dir is
# self-contained and readable without the HF cache layout).
if os.path.exists(data_dir):
shutil.rmtree(data_dir)
n = 0
for root, _, files in os.walk(src):
for fname in files:
if not (fname.endswith(".holistic.jsonl")
or fname.endswith(".intent_entity.jsonl")):
continue
rel = os.path.relpath(os.path.join(root, fname), src)
out = os.path.join(data_dir, rel)
os.makedirs(os.path.dirname(out), exist_ok=True)
shutil.copy(os.path.realpath(os.path.join(root, fname)), out)
n += 1
print(f"[prewarm] staged {n} files into {data_dir}.", flush=True)
except Exception as exc: # noqa: BLE001 — prewarm must never block startup
print(f"[prewarm] skipped ({exc})", flush=True)
if __name__ == "__main__":
main()