"""HF Dataset queue: async Mac → HF prediction pipeline. Architecture (pull-based, Mac never opens a port): HF Space → writes "pending" to HF Dataset queue Mac worker → polls queue every 5 min → runs prediction → writes "done" HF Space → reads fresh result from queue on next request Queue file (queue.json in HF Dataset repo): { "": { "status": "pending|done|error", "requested_at": "ISO", "completed_at": "ISO|null", "source": "hf_space|mac_worker", "result": {...}|null } } """ import json import logging import os import tempfile import time from datetime import datetime, timedelta, timezone from pathlib import Path logger = logging.getLogger(__name__) DATASET_REPO = os.getenv("HF_QUEUE_DATASET", "DennisChan0909/stock-predictor-queue") QUEUE_FILE = "queue.json" RESULT_FRESH_HOURS = 24 # "done" results older than this are re-queued STALE_REFRESH_HOURS = 2 # Mac refreshes results older than this _QUEUE_READ_TTL = 90 # seconds — cache queue.json in-memory to avoid repeated HF Hub downloads _queue_mem_cache: dict = {} # {"data": {...}, "ts": float} def _get_token() -> str | None: token = os.getenv("HF_TOKEN") if token: return token token_file = Path(__file__).resolve().parent.parent / ".hf_token" if token_file.exists(): t = token_file.read_text().strip() if t and t != "PASTE_YOUR_HF_TOKEN_HERE": return t return None def _hf_available() -> bool: try: import huggingface_hub # noqa: F401 return True except ImportError: return False def is_on_hf_space() -> bool: return bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID") or os.getenv("SPACE_HOST")) def read_queue(force: bool = False) -> dict: if not _hf_available(): return {} token = _get_token() if not token: return {} # In-memory cache — avoids repeated HF Hub downloads on every predict request if not force: cached = _queue_mem_cache.get("data") if cached is not None and time.time() - _queue_mem_cache.get("ts", 0) < _QUEUE_READ_TTL: return cached try: from huggingface_hub import hf_hub_download path = hf_hub_download( repo_id=DATASET_REPO, filename=QUEUE_FILE, repo_type="dataset", token=token, force_download=True, ) with open(path) as f: data = json.load(f) _queue_mem_cache["data"] = data _queue_mem_cache["ts"] = time.time() return data except Exception as e: logger.debug("HF queue read: %s", e) return {} def write_queue(queue: dict) -> bool: if not _hf_available(): return False token = _get_token() if not token: return False try: from huggingface_hub import HfApi api = HfApi() with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(queue, f, indent=2, default=str) tmp = f.name api.upload_file( path_or_fileobj=tmp, path_in_repo=QUEUE_FILE, repo_id=DATASET_REPO, repo_type="dataset", token=token, commit_message=f"worker {datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M')}", ) os.unlink(tmp) # Invalidate in-memory cache so next read gets fresh data _queue_mem_cache.clear() return True except Exception as e: logger.warning("HF queue write: %s", e) return False def is_publishable_result(result: object) -> bool: """Return True when a queue result is a completed ML-style prediction.""" if not isinstance(result, dict): return False if result.get("source") == "quick_rules": return False if str(result.get("disclaimer", "")).startswith("Quick estimate"): return False return True def get_cached_result(stock_no: str, max_age_hours: float | None = RESULT_FRESH_HOURS) -> dict | None: """Return a fresh 'done' prediction from the queue, or None.""" queue = read_queue() item = queue.get(stock_no) if not item or item.get("status") != "done": return None result = item.get("result") if not is_publishable_result(result): return None completed = item.get("completed_at") if completed and max_age_hours is not None: try: ts = datetime.fromisoformat(completed.replace("Z", "+00:00")) if datetime.now(timezone.utc) - ts > timedelta(hours=max_age_hours): return None except Exception: pass return result def enqueue(stock_no: str) -> bool: """Add stock_no to queue as pending. Idempotent — skips if already pending or freshly done.""" queue = read_queue() existing = queue.get(stock_no, {}) status = existing.get("status") if status == "pending": return True if status == "done": completed = existing.get("completed_at") if completed: try: ts = datetime.fromisoformat(completed.replace("Z", "+00:00")) if datetime.now(timezone.utc) - ts < timedelta(hours=STALE_REFRESH_HOURS): return True # still fresh except Exception: pass queue[stock_no] = { "status": "pending", "requested_at": datetime.now(timezone.utc).isoformat(), "completed_at": None, "source": "hf_space", "result": None, } return write_queue(queue) def queue_status(stock_no: str) -> dict: """Return the current queue entry for a stock, or {"status": "not_queued"}.""" queue = read_queue() return queue.get(stock_no, {"status": "not_queued"})