evalstate/diffusers-pr-api

1from __future__ import annotations2 3import json4from collections.abc import Callable, Iterable5from dataclasses import dataclass6from datetime import UTC, datetime7from pathlib import Path8from typing import Any, Protocol, cast9 10from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download11 12from slop_farmer.app.save_cache import _save_analysis_cache_api13from slop_farmer.config import PublishAnalysisArtifactsOptions14from slop_farmer.data.parquet_io import read_json15from slop_farmer.data.snapshot_paths import (16    ANALYSIS_REPORT_FILENAME_BY_VARIANT,17    HYBRID_ANALYSIS_REVIEWS_FILENAME,18    ROOT_MANIFEST_FILENAME,19    analysis_run_artifact_path,20    analysis_run_manifest_path,21    archived_snapshot_manifest_path,22    build_archived_analysis_run_manifest,23    build_current_analysis_manifest,24    current_analysis_artifact_path,25    resolve_snapshot_dir_from_output,26)27 28 29class HubApiLike(Protocol):30    def create_repo(31        self,32        repo_id: str,33        *,34        repo_type: str,35        private: bool,36        exist_ok: bool,37    ) -> None: ...38 39    def create_commit(40        self,41        repo_id: str,42        operations: Iterable[CommitOperationAdd],43        *,44        commit_message: str,45        repo_type: str,46    ) -> Any: ...47 48    def upload_folder(49        self,50        *,51        repo_id: str,52        folder_path: Path,53        path_in_repo: str,54        repo_type: str,55        commit_message: str,56    ) -> None: ...57 58 59@dataclass(frozen=True, slots=True)60class PublishableAnalysisArtifacts:61    repo: str62    snapshot_id: str63    model: str | None64    report_path: Path65    reviews_path: Path | None66    report_payload: dict[str, Any]67 68 69def run_publish_analysis_artifacts(options: PublishAnalysisArtifactsOptions) -> dict[str, Any]:70    snapshot_dir = resolve_snapshot_dir_from_output(options.output_dir, options.snapshot_dir)71    return publish_analysis_artifacts(72        snapshot_dir=snapshot_dir,73        analysis_input=options.analysis_input,74        hf_repo_id=options.hf_repo_id,75        analysis_id=options.analysis_id,76        canonical=options.canonical,77        save_cache=options.save_cache,78        private=options.private_hf_repo,79    )80 81 82def publish_analysis_artifacts(83    *,84    snapshot_dir: Path,85    analysis_input: Path | None,86    hf_repo_id: str,87    analysis_id: str,88    canonical: bool,89    private: bool,90    save_cache: bool = False,91    log: Callable[[str], None] | None = None,92) -> dict[str, Any]:93    return _publish_analysis_artifacts_api(94        cast("HubApiLike", HfApi()),95        snapshot_dir=snapshot_dir,96        analysis_input=analysis_input,97        hf_repo_id=hf_repo_id,98        analysis_id=analysis_id,99        canonical=canonical,100        private=private,101        save_cache=save_cache,102        log=log,103    )104 105 106def _publish_analysis_artifacts_api(107    api: HubApiLike,108    *,109    snapshot_dir: Path,110    analysis_input: Path | None = None,111    hf_repo_id: str,112    analysis_id: str,113    canonical: bool,114    private: bool,115    save_cache: bool = False,116    log: Callable[[str], None] | None = None,117) -> dict[str, Any]:118    artifacts = _discover_publishable_analysis(snapshot_dir, analysis_input=analysis_input)119    published_at = _iso_now()120    channel = "canonical" if canonical else "comparison"121    archived_manifest = build_archived_analysis_run_manifest(122        repo=artifacts.repo,123        snapshot_id=artifacts.snapshot_id,124        analysis_id=analysis_id,125        variant="hybrid",126        channel=channel,127        model=artifacts.model,128        published_at=published_at,129        include_hybrid_reviews=artifacts.reviews_path is not None,130    )131    current_manifest = (132        build_current_analysis_manifest(133            repo=artifacts.repo,134            snapshot_id=artifacts.snapshot_id,135            analysis_id=analysis_id,136            variant="hybrid",137            channel=channel,138            model=artifacts.model,139            published_at=published_at,140            include_hybrid_reviews=artifacts.reviews_path is not None,141        )142        if canonical143        else None144    )145    snapshot_manifest = _updated_snapshot_manifest(146        snapshot_dir=snapshot_dir,147        hf_repo_id=hf_repo_id,148        snapshot_id=artifacts.snapshot_id,149        analysis_id=analysis_id,150        archived_manifest=archived_manifest,151        canonical=canonical,152    )153    operations = _commit_operations(154        artifacts=artifacts,155        analysis_id=analysis_id,156        archived_manifest=archived_manifest,157        current_manifest=current_manifest,158        snapshot_manifest=snapshot_manifest,159    )160 161    if log:162        log(f"Ensuring Hub dataset repo exists: {hf_repo_id}")163    api.create_repo(hf_repo_id, repo_type="dataset", private=private, exist_ok=True)164    if log:165        log(f"Publishing analysis {analysis_id} for snapshot {artifacts.snapshot_id}")166    api.create_commit(167        hf_repo_id,168        operations,169        commit_message=f"Publish analysis {analysis_id} for snapshot {artifacts.snapshot_id}",170        repo_type="dataset",171    )172    cache_result = (173        _save_analysis_cache_api(174            api,175            snapshot_dir=snapshot_dir,176            hf_repo_id=hf_repo_id,177            private=private,178            log=log,179        )180        if save_cache181        else None182    )183    result: dict[str, Any] = {184        "repo": artifacts.repo,185        "dataset_id": hf_repo_id,186        "snapshot_id": artifacts.snapshot_id,187        "analysis_id": analysis_id,188        "canonical": canonical,189        "save_cache": save_cache,190        "published_at": published_at,191        "artifact_paths": [operation.path_in_repo for operation in operations],192    }193    if cache_result is not None:194        result["cache"] = cache_result195    if log:196        log(f"Published analysis artifacts to {hf_repo_id}")197    return result198 199 200def _discover_publishable_analysis(201    snapshot_dir: Path, *, analysis_input: Path | None202) -> PublishableAnalysisArtifacts:203    manifest_path = snapshot_dir / ROOT_MANIFEST_FILENAME204    if not manifest_path.exists():205        raise FileNotFoundError(f"Snapshot manifest is missing: {manifest_path}")206    manifest = read_json(manifest_path)207    if not isinstance(manifest, dict):208        raise ValueError(f"Snapshot manifest at {manifest_path} must contain a JSON object.")209    snapshot_id = str(manifest.get("snapshot_id") or snapshot_dir.name).strip()210    repo = str(manifest.get("repo") or "").strip()211    if not repo:212        raise ValueError(f"Snapshot manifest at {manifest_path} does not define repo.")213 214    report_path = (215        analysis_input.resolve()216        if analysis_input is not None217        else snapshot_dir / ANALYSIS_REPORT_FILENAME_BY_VARIANT["hybrid"]218    )219    if not report_path.exists():220        raise FileNotFoundError(f"Hybrid analysis report is missing: {report_path}")221    report_payload = read_json(report_path)222    if not isinstance(report_payload, dict):223        raise ValueError(f"Hybrid analysis report at {report_path} must contain a JSON object.")224    report_snapshot_id = str(report_payload.get("snapshot_id") or snapshot_id).strip()225    if report_snapshot_id != snapshot_id:226        raise ValueError(227            f"Hybrid analysis report snapshot_id {report_snapshot_id!r} does not match manifest snapshot_id {snapshot_id!r}."228        )229    report_repo = str(report_payload.get("repo") or repo).strip()230    if report_repo != repo:231        raise ValueError(232            f"Hybrid analysis report repo {report_repo!r} does not match manifest repo {repo!r}."233        )234    model = report_payload.get("model")235    if model is not None:236        model = str(model)237 238    reviews_path = report_path.with_name(f"{report_path.stem}.llm-reviews.json")239    return PublishableAnalysisArtifacts(240        repo=repo,241        snapshot_id=snapshot_id,242        model=model,243        report_path=report_path,244        reviews_path=reviews_path if reviews_path.exists() else None,245        report_payload={str(key): value for key, value in report_payload.items()},246    )247 248 249def _updated_snapshot_manifest(250    *,251    snapshot_dir: Path,252    hf_repo_id: str,253    snapshot_id: str,254    analysis_id: str,255    archived_manifest: dict[str, Any],256    canonical: bool,257) -> dict[str, Any]:258    manifest = _load_remote_snapshot_manifest(hf_repo_id, snapshot_id) or read_json(259        snapshot_dir / ROOT_MANIFEST_FILENAME260    )261    if not isinstance(manifest, dict):262        raise ValueError("Archived snapshot manifest must contain a JSON object.")263    updated = {str(key): value for key, value in manifest.items()}264    published_analysis: dict[str, Any] | Any = updated.get("published_analysis")265    if not isinstance(published_analysis, dict):266        published_analysis = {"schema_version": 1, "runs": {}}267    runs: dict[str, Any] | Any = published_analysis.get("runs")268    if not isinstance(runs, dict):269        runs = {}270    runs[analysis_id] = {271        "analysis_id": analysis_id,272        "variant": archived_manifest["variant"],273        "channel": archived_manifest["channel"],274        "model": archived_manifest.get("model"),275        "published_at": archived_manifest["published_at"],276        "manifest_path": analysis_run_manifest_path(snapshot_id, analysis_id),277        "artifacts": archived_manifest["artifacts"],278    }279    published_analysis["schema_version"] = 1280    published_analysis["runs"] = runs281    if canonical:282        published_analysis["canonical_analysis_id"] = analysis_id283    updated["published_analysis"] = published_analysis284    return updated285 286 287def _load_remote_snapshot_manifest(hf_repo_id: str, snapshot_id: str) -> dict[str, Any] | None:288    try:289        downloaded = hf_hub_download(290            repo_id=hf_repo_id,291            repo_type="dataset",292            filename=archived_snapshot_manifest_path(snapshot_id),293        )294    except Exception:295        return None296    payload = json.loads(Path(downloaded).read_text(encoding="utf-8"))297    return payload if isinstance(payload, dict) else None298 299 300def _commit_operations(301    *,302    artifacts: PublishableAnalysisArtifacts,303    analysis_id: str,304    archived_manifest: dict[str, Any],305    current_manifest: dict[str, Any] | None,306    snapshot_manifest: dict[str, Any],307) -> list[CommitOperationAdd]:308    report_filename = ANALYSIS_REPORT_FILENAME_BY_VARIANT["hybrid"]309    operations = [310        CommitOperationAdd(311            path_in_repo=analysis_run_artifact_path(312                artifacts.snapshot_id,313                analysis_id,314                report_filename,315            ),316            path_or_fileobj=artifacts.report_path,317        ),318        CommitOperationAdd(319            path_in_repo=analysis_run_manifest_path(artifacts.snapshot_id, analysis_id),320            path_or_fileobj=_json_bytes(archived_manifest),321        ),322        CommitOperationAdd(323            path_in_repo=archived_snapshot_manifest_path(artifacts.snapshot_id),324            path_or_fileobj=_json_bytes(snapshot_manifest),325        ),326    ]327    if artifacts.reviews_path is not None:328        operations.append(329            CommitOperationAdd(330                path_in_repo=analysis_run_artifact_path(331                    artifacts.snapshot_id,332                    analysis_id,333                    HYBRID_ANALYSIS_REVIEWS_FILENAME,334                ),335                path_or_fileobj=artifacts.reviews_path,336            )337        )338    if current_manifest is not None:339        operations.extend(340            [341                CommitOperationAdd(342                    path_in_repo=current_analysis_artifact_path(report_filename),343                    path_or_fileobj=artifacts.report_path,344                ),345                CommitOperationAdd(346                    path_in_repo=current_analysis_artifact_path(ROOT_MANIFEST_FILENAME),347                    path_or_fileobj=_json_bytes(current_manifest),348                ),349            ]350        )351        if artifacts.reviews_path is not None:352            operations.append(353                CommitOperationAdd(354                    path_in_repo=current_analysis_artifact_path(HYBRID_ANALYSIS_REVIEWS_FILENAME),355                    path_or_fileobj=artifacts.reviews_path,356                )357            )358    return operations359 360 361def _json_bytes(payload: dict[str, Any]) -> bytes:362    return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8")363 364 365def _iso_now() -> str:366    return datetime.now(tz=UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")367