hf_store ¶

def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
    """Return whether *record* may contribute to rolling baselines.

    Legacy records predate ``baseline_eligible`` and ``run_source``. They were
    uploaded only by the old successful main/full-suite path, so keep them
    eligible until the HF history naturally rolls forward.
    """
    if record.get("baseline_eligible") is True:
        return True
    return "baseline_eligible" not in record and "run_source" not in record

def load_as_dataframe(
    local_dir: str,
    *,
    days: int | None = None,
    successful_only: bool = False,
) -> pd.DataFrame:
    """Load and normalize records from *local_dir* into a pandas DataFrame.

    Combines :func:`load_records` + :func:`normalize_dataframe` into a single
    call for consumers (e.g. the dashboard) that work exclusively with
    DataFrames.

    Args:
        local_dir: Root directory previously populated by :func:`sync_from_hf`.
        days: Passed through to :func:`load_records`.
        successful_only: Passed through to :func:`load_records`.

    Returns:
        Normalized DataFrame, or an empty DataFrame if no records were found.
    """
    records = load_records(local_dir, days=days, successful_only=successful_only)
    if not records:
        return pd.DataFrame()

    df = pd.DataFrame(records)
    return normalize_dataframe(df)

def load_records(
    local_dir: str,
    *,
    days: int | None = None,
    successful_only: bool = False,
    baseline_eligible_only: bool = False,
) -> list[dict[str, Any]]:
    """Return raw JSON dicts from *local_dir*.

    Args:
        local_dir: Root directory previously populated by :func:`sync_from_hf`.
        days: When set, discard records whose ``timestamp`` is older than this
            many days. Records with a missing/unparsable timestamp are kept.
        successful_only: When True, only records with ``success=True`` are
            returned. Useful when building a regression baseline.
        baseline_eligible_only: When True, only baseline-eligible records are
            returned. Legacy records missing both ``baseline_eligible`` and
            ``run_source`` are treated as eligible.

    Returns:
        List of raw dicts sorted by ``timestamp`` ascending (records that could
        not be parsed are silently skipped).
    """
    cutoff: datetime | None = None
    if days is not None:
        cutoff = datetime.now(timezone.utc) - timedelta(days=days)

    records: list[tuple[datetime, str, dict[str, Any]]] = []

    for path in sorted(glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True)):
        try:
            with open(path, encoding="utf-8") as fh:
                data: dict[str, Any] = json.load(fh)
        except (OSError, json.JSONDecodeError):
            continue

        if successful_only and not data.get("success", True):
            continue

        if baseline_eligible_only and not is_baseline_eligible_record(data):
            continue

        ts = _parse_record_timestamp(data)
        if cutoff is not None and ts is not None and ts < cutoff:
            continue

        records.append((
            ts or datetime.min.replace(tzinfo=timezone.utc),
            path,
            data,
        ))

    return [data for _ts, _path, data in sorted(records)]

def load_records_for_identity(
    local_dir: str,
    identity_filters: dict[str, str],
    *,
    last_n: int | None = None,
    successful_only: bool = True,
    baseline_eligible_only: bool = False,
) -> list[dict[str, Any]]:
    """Return records matching comparable identity fields.

    V2 performance comparison is intentionally independent from the legacy
    ``model_id`` directory and ``gpu_type`` display string. The comparable
    identity filters usually contain the full v2 comparison key, and may also
    contain a subset when looking for same-cohort recipe mismatches.
    """
    records = load_records(
        local_dir,
        successful_only=successful_only,
        baseline_eligible_only=baseline_eligible_only,
    )

    for key, expected in identity_filters.items():
        records = [r for r in records if str(r.get(key)) == str(expected)]

    if last_n is not None:
        records = records[-last_n:]

    return records

def load_records_for_model(
    local_dir: str,
    model_id: str,
    gpu_type: str | None = None,
    *,
    workload_id: str | None = None,
    variant_id: str | None = None,
    benchmark_version: str | None = None,
    recipe_fingerprint: str | None = None,
    hardware_profile_id: str | None = None,
    software_profile_id: str | None = None,
    last_n: int | None = None,
    successful_only: bool = True,
    baseline_eligible_only: bool = False,
) -> list[dict[str, Any]]:
    """Return records for a specific *model_id*, optionally filtered by cohort.

    Args:
        local_dir: Root directory previously populated by :func:`sync_from_hf`.
        model_id: Matches the ``model_id`` field inside each JSON record.
        gpu_type: When set, only records whose ``gpu_type`` matches are returned.
        workload_id: When set, only records from the same workload are returned.
        variant_id: When set, only records from the same workload variant are returned.
        benchmark_version: When set, only records from the same benchmark version are returned.
        recipe_fingerprint: When set, only records from the same benchmark
            recipe are returned.
        hardware_profile_id: When set, only records from the same hardware
            cohort are returned.
        software_profile_id: When set, only records from the same software
            cohort are returned.
        last_n: When set, return only the most recent *n* records (after all
            other filters). Useful for sliding-window baseline calculations.
        successful_only: Passed through to :func:`load_records`.
        baseline_eligible_only: Passed through to :func:`load_records`.

    Returns:
        List of matching dicts sorted by timestamp ascending.
    """
    model_dir = os.path.join(local_dir, sanitize(model_id))
    if not os.path.isdir(model_dir):
        return []

    records = load_records(
        model_dir,
        successful_only=successful_only,
        baseline_eligible_only=baseline_eligible_only,
    )

    if gpu_type is not None:
        records = [r for r in records if r.get("gpu_type") == gpu_type]

    identity_filters = {
        "workload_id": workload_id,
        "variant_id": variant_id,
        "benchmark_version": benchmark_version,
        "recipe_fingerprint": recipe_fingerprint,
        "hardware_profile_id": hardware_profile_id,
        "software_profile_id": software_profile_id,
    }
    for key, expected in identity_filters.items():
        if expected is not None:
            records = [r for r in records if str(r.get(key)) == str(expected)]

    if last_n is not None:
        records = records[-last_n:]

    return records

def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Apply standard type coercions to a raw records DataFrame.

    - Parses ``timestamp`` to UTC-aware datetime.
    - Coerces ``latency``, ``throughput``, ``memory``, ``text_encoder_time_s``,
      ``dit_time_s``, ``vae_decode_time_s`` to float.
    - Adds a ``config_id`` column (first 7 chars of ``commit_sha``).

    Returns the mutated DataFrame (also modifies in place for efficiency).
    """
    if df.empty:
        return df

    df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True, errors="coerce")
    df["config_id"] = df.get("commit_sha", pd.Series(dtype=str)).fillna("unknown").str[:7]

    for col in _NUMERIC_COLS:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    return df

def resolve_hf_token() -> str | None:
    """Return the first configured Hugging Face token env var."""
    for env_var in HF_TOKEN_ENV_VARS:
        token = os.environ.get(env_var)
        if token:
            return token
    return None

def safe_float(value: Any) -> float | None:
    """Coerce *value* to float, returning None on failure."""
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None

def sanitize(value: str) -> str:
    """Return a filesystem- and HF-path-safe version of *value*."""
    sanitized = re.sub(r"[^A-Za-z0-9._-]", "_", value)
    return f"_{sanitized}" if not sanitized or sanitized.startswith(".") else sanitized

def sync_from_hf(
    local_dir: str,
    *,
    strict: bool = False,
    reuse_existing: bool = False,
    revision: str | None = None,
) -> str:
    """Download the HF dataset repo snapshot to *local_dir*.

    Returns *local_dir* so callers can chain: ``load_records(sync_from_hf(...))``.

    By default (``strict=False``) failures are logged and *local_dir* is
    returned unchanged, so dashboard / PR consumers stay resilient when HF is
    unavailable. Callers that depend on the sync for correctness (e.g. the
    main-branch baseline writer) must pass ``strict=True`` so that misconfig
    or transient HF errors fail loud rather than silently reset the baseline.

    When ``reuse_existing=True``, a previous successful sync in ``local_dir``
    is reused only while its marker is fresh. This avoids duplicate HF
    snapshot checks when compare and dashboard scripts run sequentially in the
    same CI job, without silently reusing stale data in persistent local or
    long-lived runner environments.

    Pass ``revision`` to pin the snapshot to a previously read Hub commit. This
    is used by conditional writers that must validate one exact remote state
    before committing with that revision as their parent.
    """
    marker_path = _sync_marker_path(local_dir)
    if reuse_existing and os.path.exists(marker_path):
        if _sync_marker_is_fresh(marker_path) and _sync_marker_matches_request(marker_path, revision):
            print(f"hf_store: reusing existing sync at {local_dir}")
            return local_dir
        os.remove(marker_path)
        print(f"hf_store: existing sync at {local_dir} is stale or mismatched; refreshing")

    if not reuse_existing and os.path.exists(marker_path):
        os.remove(marker_path)

    if not HF_REPO_ID:
        msg = "hf_store: HF_REPO_ID not set"
        if strict:
            raise RuntimeError(f"{msg}; cannot sync.")
        print(f"{msg}, skipping sync.")
        return local_dir

    print(f"hf_store: syncing from {HF_REPO_ID} → {local_dir}")
    try:
        snapshot_download(
            repo_id=HF_REPO_ID,
            repo_type="dataset",
            local_dir=local_dir,
            token=resolve_hf_token(),
            allow_patterns="*.json",
            revision=revision,
        )
        os.makedirs(local_dir, exist_ok=True)
        with open(marker_path, "w", encoding="utf-8") as marker:
            json.dump(
                {
                    "endpoint": ENDPOINT,
                    "repo_id": HF_REPO_ID,
                    "synced_at": datetime.now(timezone.utc).isoformat(),
                    "revision": revision,
                }, marker)
    except Exception as exc:
        if strict:
            raise
        print(f"hf_store: sync skipped — {exc}")

    return local_dir

def upload_record(
    local_path: str,
    record: dict[str, Any],
    *,
    strict: bool = False,
) -> None:
    """Upload *local_path* to the HF repo under ``<model_id>/<filename>``.

    By default failures (missing token, network errors) are logged and
    swallowed. Pass ``strict=True`` when the upload is part of a write-path
    that must not silently lose records — otherwise the rolling baseline can
    stop advancing without any signal in the build log.
    """
    token = resolve_hf_token()
    if not token:
        msg = f"hf_store: none of {', '.join(HF_TOKEN_ENV_VARS)} set"
        if strict:
            raise RuntimeError(f"{msg}; cannot upload.")
        print(f"{msg}, skipping upload.")
        return

    model_id = record.get("model_id", "unknown")
    path_in_repo = f"{sanitize(model_id)}/{os.path.basename(local_path)}"
    commit_sha = (record.get("commit_sha") or "unknown")[:7]

    api = HfApi(token=token)
    try:
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=path_in_repo,
            repo_id=HF_REPO_ID,
            repo_type="dataset",
            commit_message=f"Perf: {model_id} at {commit_sha}",
        )
        print(f"hf_store: uploaded → {HF_REPO_ID}/{path_in_repo}")
    except Exception as exc:
        if strict:
            raise
        print(f"hf_store: upload failed — {exc}")