File size: 2,318 Bytes
4d68493 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Optional persistence: round-trip the OHLCV cache to a private HF Dataset.
The HF Space container's ``/tmp`` is wiped on every restart. If the user
sets ``HF_TOKEN`` (or ``HUGGING_FACE_HUB_TOKEN``) and ``FSCANNER_CACHE_REPO``,
we push the parquet cache to a private dataset repo and pull it back on
startup. This makes cold-start scans near-instant for a previously-scanned
universe.
"""
from __future__ import annotations
import os
from typing import Optional
import pandas as pd
from . import paths
CACHE_REPO = os.environ.get("FSCANNER_CACHE_REPO", "").strip()
CACHE_FILE = "ohlcv_cache.parquet"
def _is_enabled() -> bool:
return bool(CACHE_REPO) and bool(
os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
)
def pull_remote_cache() -> int:
"""Try to download the remote cache and write to CACHE_PATH. Returns the
number of tickers loaded, 0 on any failure or if not enabled.
Pulls whenever remote persistence is configured; the remote copy
overrides whatever (possibly stale or missing) file exists locally.
"""
if not _is_enabled():
return 0
try:
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id=CACHE_REPO,
repo_type="dataset",
filename=CACHE_FILE,
force_download=False,
)
# Replace the local cache with the downloaded one
df = pd.read_parquet(path)
os.makedirs(os.path.dirname(paths.CACHE_PATH) or ".", exist_ok=True)
df.to_parquet(paths.CACHE_PATH, index=False)
tickers = df["Ticker"].nunique() if "Ticker" in df.columns else 0
return int(tickers)
except Exception:
return 0
def push_remote_cache() -> bool:
"""Upload the local cache to the remote dataset repo. Returns True on
success.
"""
if not _is_enabled() or not os.path.exists(paths.CACHE_PATH):
return False
try:
from huggingface_hub import HfApi
api = HfApi()
api.upload_file(
path_or_fileobj=paths.CACHE_PATH,
path_in_repo=CACHE_FILE,
repo_id=CACHE_REPO,
repo_type="dataset",
commit_message="Update OHLCV cache",
)
return True
except Exception:
return False
|