"""Background worker for benchmark card generation. Detects new benchmark folders in EEE_datastore, generates cards via run_eee_pipeline(), and uploads them to evaleval/auto-benchmarkcards. Uses Jenny's Entity Registry for canonical ID resolution and dedup. """ import json import logging import os import tempfile import time from datetime import datetime, timezone from functools import wraps from pathlib import Path from typing import Any, Optional import requests from huggingface_hub import HfApi, snapshot_download logger = logging.getLogger("worker") EEE_REPO = "evaleval/EEE_datastore" CARDS_REPO = "evaleval/auto-benchmarkcards" ENTITY_REGISTRY_URL = "https://evaleval-entity-registry.hf.space/api/v1" # Persistent storage on HF Spaces (mounted volume). # Falls back to local /tmp for development. PERSISTENT_DIR = Path(os.environ.get("PERSISTENT_DIR", "/data")) STATE_FILE = PERSISTENT_DIR / "state.json" FORCE_REGENERATE = os.environ.get("FORCE_REGENERATE", "").lower() in ("1", "true", "yes") # -- Retry decorator for transient failures -- def retry(max_attempts=3, delay=5, backoff=2): """Retry decorator with exponential backoff for transient failures.""" def decorator(func): @wraps(func) def wrapper(*args, **kwargs): for attempt in range(max_attempts): try: return func(*args, **kwargs) except Exception as e: if attempt == max_attempts - 1: raise wait = delay * (backoff ** attempt) logger.warning( "%s failed (attempt %d/%d), retrying in %ds: %s", func.__name__, attempt + 1, max_attempts, wait, e, ) time.sleep(wait) return wrapper return decorator # -- State management (atomic writes) -- def load_state() -> dict: """Load persistent state (known folders, job history, pending queue).""" if STATE_FILE.exists(): try: return json.loads(STATE_FILE.read_text()) except Exception: logger.exception("Failed to read state file, starting fresh") return {"known_folders": [], "jobs": [], "pending_folders": []} def save_state(state: dict) -> None: """Save persistent state atomically (write-then-rename).""" PERSISTENT_DIR.mkdir(parents=True, exist_ok=True) tmp = STATE_FILE.with_suffix(".tmp") tmp.write_text(json.dumps(state, indent=2)) tmp.rename(STATE_FILE) def save_pending(folders: list[str]) -> None: """Add folders to the pending queue.""" state = load_state() pending = state.get("pending_folders", []) for f in folders: if f not in pending: pending.append(f) state["pending_folders"] = pending save_state(state) def pop_pending() -> list[str]: """Pop all pending folders from the queue.""" state = load_state() pending = state.pop("pending_folders", []) state["pending_folders"] = [] save_state(state) return pending # -- Entity Registry -- _canonical_cache: dict[str, Optional[str]] = {} def resolve_canonical_id(benchmark_name: str) -> Optional[str]: """Resolve benchmark name to canonical_id via Entity Registry. Returns canonical_id string (e.g. "math") or None if not found. Uses an in-memory cache to avoid repeated API calls within a job. """ if benchmark_name in _canonical_cache: return _canonical_cache[benchmark_name] try: resp = requests.post( f"{ENTITY_REGISTRY_URL}/resolve", json={"raw_value": benchmark_name, "entity_type": "benchmark"}, timeout=10, ) resp.raise_for_status() data = resp.json() canonical_id = data.get("canonical_id") _canonical_cache[benchmark_name] = canonical_id if canonical_id: logger.info("Entity Registry: '%s' -> '%s'", benchmark_name, canonical_id) return canonical_id except Exception: logger.debug("Entity Registry lookup failed for '%s'", benchmark_name) _canonical_cache[benchmark_name] = None return None def resolve_canonical_ids_batch(names: list[str]) -> dict[str, Optional[str]]: """Batch-resolve benchmark names to canonical_ids.""" # Check cache first, only query uncached names uncached = [n for n in names if n not in _canonical_cache] if not uncached: return {n: _canonical_cache[n] for n in names} try: payload = [{"raw_value": n, "entity_type": "benchmark"} for n in uncached] resp = requests.post( f"{ENTITY_REGISTRY_URL}/resolve/batch", json=payload, timeout=30, ) resp.raise_for_status() results = resp.json() for name, result in zip(uncached, results): canonical_id = result.get("canonical_id") _canonical_cache[name] = canonical_id if canonical_id: logger.info("Entity Registry: '%s' -> '%s'", name, canonical_id) except Exception: logger.warning("Entity Registry batch resolve failed, using fallback") for name in uncached: _canonical_cache[name] = None return {n: _canonical_cache.get(n) for n in names} def _get_card_filename(benchmark_name: str) -> str: """Get the canonical filename for a benchmark card. Uses Entity Registry canonical_id when available, falls back to sanitize_benchmark_name from the main package. """ canonical = resolve_canonical_id(benchmark_name) if canonical: return canonical from auto_benchmarkcard.output import sanitize_benchmark_name return sanitize_benchmark_name(benchmark_name).lower() # -- EEE folder detection -- def _extract_folders(file_list: list[str]) -> set[str]: """Extract unique top-level folder names under data/.""" folders = set() for path in file_list: parts = path.split("/") if len(parts) >= 2 and parts[0] == "data": folders.add(parts[1]) return folders @retry(max_attempts=3, delay=5) def detect_new_benchmarks() -> list[str]: """Compare current EEE_datastore file listing against known state.""" api = HfApi() all_files = api.list_repo_files(EEE_REPO, repo_type="dataset") current_folders = _extract_folders(all_files) state = load_state() known = set(state.get("known_folders", [])) new_folders = sorted(current_folders - known) if new_folders: logger.info("Detected %d new folders: %s", len(new_folders), new_folders) else: logger.info("No new folders (known: %d, current: %d)", len(known), len(current_folders)) return new_folders # -- Download & upload -- @retry(max_attempts=3, delay=10) def _download_folders(folder_names: list[str], target_dir: str) -> Path: """Download EEE folders into a shared temp directory.""" patterns = [f"data/{f}/**/*.json" for f in folder_names] logger.info("Downloading %d EEE folders to %s", len(folder_names), target_dir) snapshot_download( repo_id=EEE_REPO, repo_type="dataset", local_dir=target_dir, allow_patterns=patterns, ) return Path(target_dir) / "data" @retry(max_attempts=3, delay=5) def _upload_card(card: dict, benchmark_name: str, canonical_id: Optional[str] = None) -> bool: """Upload a generated card to evaleval/auto-benchmarkcards.""" api = HfApi() filename = canonical_id or _get_card_filename(benchmark_name) remote_path = f"cards/{filename}.json" with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(card, f, indent=2) tmp_path = f.name try: api.upload_file( path_or_fileobj=tmp_path, path_in_repo=remote_path, repo_id=CARDS_REPO, repo_type="dataset", commit_message=f"Auto-generated card: {benchmark_name}", ) logger.info("Uploaded card to %s/%s", CARDS_REPO, remote_path) return True finally: try: os.unlink(tmp_path) except OSError: pass @retry(max_attempts=2, delay=5) def _list_existing_cards() -> set[str]: """List all card filenames (without extension) in the cards repo.""" api = HfApi() all_files = api.list_repo_files(CARDS_REPO, repo_type="dataset") cards = set() for path in all_files: if path.startswith("cards/") and path.endswith(".json"): name = path[len("cards/"):-len(".json")] cards.add(name) return cards # -- Main processing -- def _build_dedup_filter( benchmark_names: list[str], existing_cards: set[str], ) -> list[str]: """Return list of benchmark names that don't already have cards. Checks in order: Entity Registry canonical_id, exact fallback name, and parent prefix match (for 'Parent - Child' pattern). """ if FORCE_REGENERATE: logger.info("FORCE_REGENERATE=true, skipping dedup") return benchmark_names # Batch-resolve all names canonical_map = resolve_canonical_ids_batch(benchmark_names) from auto_benchmarkcard.output import sanitize_benchmark_name new_benchmarks = [] for name in benchmark_names: canonical = canonical_map.get(name) fallback = sanitize_benchmark_name(name).lower() # 1. Entity Registry canonical_id match if canonical and canonical in existing_cards: logger.info("Skipping '%s' (card exists as '%s')", name, canonical) continue # 2. Exact fallback name match if fallback in existing_cards: logger.info("Skipping '%s' (card exists as '%s')", name, fallback) continue # 3. Prefix match: 'MGSM - Bengali' -> check if 'mgsm' card exists if " - " in name: parent = name.split(" - ", 1)[0].strip() parent_lower = sanitize_benchmark_name(parent).lower() if parent_lower in existing_cards: logger.info("Skipping '%s' (parent card exists as '%s')", name, parent_lower) continue new_benchmarks.append(name) logger.info("Dedup: %d total, %d new, %d existing", len(benchmark_names), len(new_benchmarks), len(benchmark_names) - len(new_benchmarks)) return new_benchmarks def process_new_benchmarks(new_folders: list[str]) -> None: """Generate and upload cards for benchmarks in new folders. Delegates to run_eee_pipeline() for the actual generation, using a callback to upload each card as it's generated. """ from auto_benchmarkcard.eee_workflow import run_eee_pipeline from auto_benchmarkcard.tools.eee.eee_tool import scan_eee_folder from auto_benchmarkcard.workflow import setup_logging_suppression setup_logging_suppression(debug_mode=False) state = load_state() job_record: dict[str, Any] = { "started_at": datetime.now(timezone.utc).isoformat(), "folders": new_folders, "results": [], } # Pre-fetch existing cards for dedup try: existing_cards = _list_existing_cards() logger.info("Found %d existing cards in %s", len(existing_cards), CARDS_REPO) except Exception: logger.warning("Failed to list existing cards, dedup disabled for this job") existing_cards = set() # Download all folders into one shared temp dir with tempfile.TemporaryDirectory(prefix="eee_batch_") as tmpdir: try: data_path = _download_folders(new_folders, tmpdir) except Exception: logger.exception("Failed to download EEE folders") job_record["results"].append({ "folders": new_folders, "status": "download_failed", }) job_record["completed_at"] = datetime.now(timezone.utc).isoformat() state["jobs"].append(job_record) state["jobs"] = state["jobs"][-50:] save_state(state) return # Scan to discover benchmark names for dedup try: scan_result = scan_eee_folder(str(data_path)) except Exception: logger.exception("Failed to scan EEE data") job_record["results"].append({ "folders": new_folders, "status": "scan_failed", }) job_record["completed_at"] = datetime.now(timezone.utc).isoformat() state["jobs"].append(job_record) state["jobs"] = state["jobs"][-50:] save_state(state) return all_names = ( list(scan_result.benchmarks.keys()) + list(scan_result.composites.keys()) ) benchmarks_to_generate = _build_dedup_filter(all_names, existing_cards) if not benchmarks_to_generate: logger.info("All benchmarks already have cards, nothing to generate") job_record["results"].append({"status": "all_existing"}) else: # Upload callback: called by run_eee_pipeline for each generated card def _on_card_generated(name: str, card: dict) -> None: canonical = resolve_canonical_id(name) # Enrich card metadata inner = card.get("benchmark_card", card) info = inner.get("card_info", {}) info["source"] = "webhook" if canonical: info["canonical_id"] = canonical inner["card_info"] = info try: _upload_card(card, name, canonical_id=canonical) job_record["results"].append({ "benchmark": name, "canonical_id": canonical, "status": "uploaded", }) except Exception: logger.exception("Failed to upload card for %s", name) job_record["results"].append({ "benchmark": name, "status": "upload_failed", }) # Run the unified pipeline summary = run_eee_pipeline( eee_path=str(data_path), output_path=str(PERSISTENT_DIR / "output"), benchmarks_filter=benchmarks_to_generate, on_card_generated=_on_card_generated, ) # Record skipped/failed from pipeline summary for item in summary.get("skipped", []): job_record["results"].append({ "benchmark": item.get("benchmark", "unknown"), "status": f"skipped:{item.get('reason', 'unknown')}", }) for name in summary.get("failed", []): # Only add if not already recorded by callback existing = {r.get("benchmark") for r in job_record["results"]} if name not in existing: job_record["results"].append({ "benchmark": name, "status": "generation_failed", }) # Mark folders as known for folder_name in new_folders: if folder_name not in state["known_folders"]: state["known_folders"].append(folder_name) job_record["completed_at"] = datetime.now(timezone.utc).isoformat() results = job_record["results"] uploaded = sum(1 for r in results if r.get("status") == "uploaded") failed = sum(1 for r in results if "failed" in r.get("status", "")) skipped = sum(1 for r in results if r.get("status", "").startswith("skipped")) logger.info("Job complete: %d uploaded, %d failed, %d skipped", uploaded, failed, skipped) state["jobs"].append(job_record) state["jobs"] = state["jobs"][-50:] save_state(state)