evalstate/diffusers-pr-api

1from __future__ import annotations2 3from pathlib import Path4 5from slop_farmer.config import SnapshotAdoptOptions6from slop_farmer.data.parquet_io import read_json, read_parquet_rows, write_json7 8__all__ = ["adopt_snapshot_for_pipeline"]9 10 11def adopt_snapshot_for_pipeline(options: SnapshotAdoptOptions) -> Path:12    snapshot_dir = options.snapshot_dir.resolve()13    manifest_path = snapshot_dir / "manifest.json"14    manifest = read_json(manifest_path)15    repo = str(manifest.get("repo") or "")16    if not repo:17        raise ValueError(f"Snapshot manifest has no repo: {manifest_path}")18    snapshot_id = str(manifest.get("snapshot_id") or snapshot_dir.name)19    manifest_watermark = (20        manifest.get("watermark") if isinstance(manifest.get("watermark"), dict) else {}21    )22    next_since = (23        options.next_since24        or manifest_watermark.get("next_since")25        or manifest.get("crawl_started_at")26        or _infer_snapshot_next_since(snapshot_dir, manifest)27        or manifest.get("extracted_at")28    )29    if not isinstance(next_since, str) or not next_since:30        raise ValueError(f"Could not determine next_since from {manifest_path}")31 32    output_dir = options.output_dir.resolve()33    latest_path = output_dir / "snapshots" / "latest.json"34    watermark_path = output_dir / "state" / "watermark.json"35    write_json(36        {37            "repo": repo,38            "latest_snapshot_id": snapshot_id,39            "snapshot_dir": str(snapshot_dir),40            "manifest_path": str(manifest_path),41            "next_since": next_since,42        },43        latest_path,44    )45    write_json(46        {47            "repo": repo,48            "last_successful_snapshot_id": snapshot_id,49            "snapshot_dir": str(snapshot_dir),50            "effective_since": manifest_watermark.get("effective_since"),51            "next_since": next_since,52            "updated_at": manifest.get("imported_at") or manifest.get("extracted_at") or next_since,53        },54        watermark_path,55    )56    return latest_path57 58 59def _infer_snapshot_next_since(snapshot_dir: Path, manifest: dict[str, object]) -> str | None:60    if manifest.get("source_type") != "hf_checkpoint_import":61        return None62    table_specs = (63        ("pull_requests.parquet", ("updated_at", "created_at")),64        ("issues.parquet", ("updated_at", "created_at")),65        ("comments.parquet", ("updated_at", "created_at")),66        ("reviews.parquet", ("submitted_at",)),67        ("review_comments.parquet", ("updated_at", "created_at")),68        ("events.parquet", ("created_at",)),69    )70    for filename, fields in table_specs:71        candidates: list[str] = []72        for row in read_parquet_rows(snapshot_dir / filename):73            for field in fields:74                value = row.get(field)75                if isinstance(value, str) and value:76                    candidates.append(value)77                    break78        if candidates:79            return max(candidates)80    return None81