uv-scripts/embeddings

1# /// script2# requires-python = ">=3.10"3# dependencies = [4#     "datasets",5#     "huggingface-hub>=1.12",6# ]7# ///8"""9Fan one embedding run out across N Hugging Face Jobs, then consolidate.10 11Runs generate-embeddings.py once per shard (each worker gets RANK / NUM_SHARDS /12RUN_ID / OUTPUT_BUCKET via env), workers write parquet shards + progress heartbeats13to the run bucket (object PUTs — no repo-commit contention), and when all workers14finish a consolidation Job merges the shards into the final Hub dataset with one15commit. Runs locally on your laptop; only the workers/consolidator run on Jobs.16 17Examples:18    # Smoke: 2 small-GPU jobs over a 20k-row slice19    uv run launch-embedding-fleet.py stanfordnlp/imdb your-name/imdb-emb \\20        --max-samples 20000 --num-shards 2 --flavor t4-small --timeout 20m21 22    # Mid-size: 8 L4s over a few million rows23    uv run launch-embedding-fleet.py your-name/corpus your-name/corpus-emb \\24        --num-shards 8 --flavor l4x1 --timeout 1h25 26    # Re-run one failed shard, then consolidate an existing run27    uv run launch-embedding-fleet.py ... --run-id 20260709-1200-abc123 --retry-rank 328    uv run launch-embedding-fleet.py ... --run-id 20260709-1200-abc123 --consolidate-only29 30Resume model: every worker writes only runs/<run-id>/data/<rank>.parquet and its own31status file, so re-running a rank is idempotent. The launcher exiting early never32orphans a run — --retry-rank / --consolidate-only pick it back up.33"""34import argparse35import json36import logging37import os38import secrets as pysecrets39import sys40import time41 42logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")43logger = logging.getLogger("launch-embedding-fleet")44 45SCRIPT_BASE = "https://huggingface.co/datasets/uv-scripts/embeddings/raw/main"46 47 48def put_json(bucket, path, obj, token=None):49    from huggingface_hub import batch_bucket_files50    batch_bucket_files(bucket, add=[(json.dumps(obj, indent=2).encode(), path)], token=token)51 52 53def rows_in_split(input_dataset, config, split):54    """Row count WITHOUT downloading: builder metadata, else the dataset-viewer size API."""55    try:56        from datasets import load_dataset_builder57        b = load_dataset_builder(input_dataset, config) if config else load_dataset_builder(input_dataset)58        n = b.info.splits[split].num_examples59        if n:60            return n61    except Exception as e:62        logger.info(f"builder metadata unavailable ({e}); trying dataset-viewer size API")63    try:64        import huggingface_hub65        r = huggingface_hub.get_session().get(66            "https://datasets-server.huggingface.co/size", params={"dataset": input_dataset}67        )68        r.raise_for_status()69        for s in r.json()["size"]["splits"]:70            if s["split"] == split and (config is None or s["config"] == config):71                return s["num_rows"]72    except Exception as e:73        logger.info(f"size API unavailable ({e})")74    return None75 76 77def main():78    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)79    p.add_argument("input_dataset")80    p.add_argument("output_dataset")81    p.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2")82    p.add_argument("--column", default="text")83    p.add_argument("--split", default="train")84    p.add_argument("--config", default=None)85    p.add_argument("--max-samples", type=int, default=None)86    p.add_argument("--num-shards", type=int, default=8)87    p.add_argument("--flavor", default="l4x1")88    p.add_argument("--timeout", default="1h", help="Per-worker timeout — also the hard cost ceiling")89    p.add_argument("--bucket", default=None,90                   help="Run bucket (default: <namespace>/embedding-runs)")91    p.add_argument("--script", default=f"{SCRIPT_BASE}/generate-embeddings.py",92                   help="Worker script: raw URL (default) or a local .py path for dev")93    p.add_argument("--consolidate-script", default=f"{SCRIPT_BASE}/consolidate-shards.py",94                   help="Consolidator script: raw URL (default) or local path for dev")95    p.add_argument("--consolidate-flavor", default="cpu-upgrade",96                   help="Consolidation job flavor. cpu-upgrade has 50 GB disk — use cpu-xl "97                        "(1 TB) when total shard size approaches ~20 GB.")98    p.add_argument("--rows-total", type=int, default=None,99                   help="Override when the dataset has no split metadata")100    p.add_argument("--streaming", action="store_true",101                   help="Workers stream + shard at the FILE level (no full-split download per "102                        "rank) — for very big datasets. Text only; incompatible with --max-samples.")103    p.add_argument("--private", action="store_true", help="Final output dataset is private")104    p.add_argument("--embed-args", nargs=argparse.REMAINDER, default=[],105                   help="Everything after --embed-args is passed through to generate-embeddings.py "106                        "verbatim — put it LAST (any launcher flags after it are swallowed too).")107    p.add_argument("--run-id", default=None, help="Attach to an existing run (with --resume/--retry-rank/--consolidate-only)")108    p.add_argument("--resume", action="store_true",109                   help="Converge an existing --run-id: find ranks without a 'done' status, re-run "110                        "exactly those, then consolidate. Idempotent — safe to run repeatedly.")111    p.add_argument("--retry-rank", type=int, default=None, help="Re-spawn a single failed shard of --run-id")112    p.add_argument("--consolidate-only", action="store_true", help="Just run consolidation for --run-id")113    p.add_argument("--no-wait", action="store_true",114                   help="Spawn workers and exit (run --consolidate-only later)")115    args = p.parse_args()116 117    from huggingface_hub import (JobStage, create_bucket, download_bucket_files, get_token,118                                 run_uv_job, wait_for_job, whoami)119 120    # Workers need a real token as a secret (bucket writes + output push); env may be empty121    # when the user authenticated via `hf auth login` (keyring/hub cache).122    token = os.environ.get("HF_TOKEN") or get_token()123    if not token:124        p.error("No HF token found — set HF_TOKEN or run `hf auth login`.")125    namespace = whoami(token=token)["name"]126    bucket = args.bucket or f"{namespace}/embedding-runs"127 128    if (args.retry_rank is not None or args.consolidate_only or args.resume) and not args.run_id:129        p.error("--resume / --retry-rank / --consolidate-only need --run-id")130    if args.num_shards < 1:131        p.error(f"--num-shards must be >= 1 (got {args.num_shards})")132    if args.streaming and args.max_samples:133        p.error("--streaming is incompatible with --max-samples (use row mode for capped tests)")134 135    def read_manifest(run_id):136        import tempfile137        from pathlib import Path138        with tempfile.TemporaryDirectory() as td:139            dst = Path(td) / "run.json"140            download_bucket_files(bucket, [(f"runs/{run_id}/run.json", dst)],141                                  raise_on_missing_files=True, token=token)142            return json.loads(dst.read_text())143 144    # Workers are always spawned from the run manifest, never from current CLI flags —145    # a --retry-rank months later must reproduce the original slice/model/args exactly.146    def spawn_worker(rank, manifest):147        script_args = [manifest["input_dataset"], manifest["output_dataset"],148                       "--model", manifest["model"], "--column", manifest["column"],149                       "--split", manifest["split"]]150        if manifest.get("config"):151            script_args += ["--config", manifest["config"]]152        if manifest.get("max_samples"):153            script_args += ["--max-samples", str(manifest["max_samples"])]154        script_args += list(manifest.get("embed_args") or [])155        env = {156            "RANK": str(rank),157            "NUM_SHARDS": str(manifest["num_shards"]),158            "RUN_ID": manifest["run_id"],159            "OUTPUT_BUCKET": bucket,160        }161        if manifest.get("revision"):162            env["REVISION"] = manifest["revision"]163        if manifest.get("streaming"):164            env["STREAMING"] = "1"165        job = run_uv_job(166            args.script,167            script_args=script_args,168            flavor=manifest["flavor"],169            timeout=manifest["timeout"],170            env=env,171            secrets={"HF_TOKEN": token},172            labels={"embedding-fleet-run": manifest["run_id"], "rank": str(rank)},173            token=token,174        )175        logger.info(f"  rank {rank} → job {job.id} ({manifest['flavor']})")176        return job177 178    def spawn_consolidator(run_id):179        job = run_uv_job(180            args.consolidate_script,181            script_args=[182                "--bucket", bucket, "--run-id", run_id,183            ] + (["--private"] if args.private else []),184            flavor=args.consolidate_flavor,185            timeout="2h",186            secrets={"HF_TOKEN": token},187            labels={"embedding-fleet-run": run_id, "role": "consolidate"},188            token=token,189        )190        logger.info(f"Consolidation job {job.id} ({args.consolidate_flavor}) — "191                    f"merges shards → {args.output_dataset}")192        return job193 194    def execute_workers(ranks, manifest):195        """Spawn the given ranks, wait, auto-retry failures ONCE, return still-failed ranks."""196        for attempt in (1, 2):197            jobs = {rank: spawn_worker(rank, manifest) for rank in ranks}198            infos = wait_for_job([j.id for j in jobs.values()], token=token)199            ranks = [rank for (rank, job), info in zip(jobs.items(), infos)200                     if info.status.stage != JobStage.COMPLETED]201            if not ranks:202                return []203            if attempt == 1:204                logger.warning(f"{len(ranks)} worker(s) failed; auto-retrying once: {ranks}")205        return ranks206 207    def done_ranks(run_id, n):208        """Ranks whose bucket status reports state == 'done' (works for both shard modes)."""209        import tempfile210        from pathlib import Path211        done = set()212        with tempfile.TemporaryDirectory() as td:213            pairs = [(f"runs/{run_id}/status/{i:05d}.json", Path(td) / f"{i}.json") for i in range(n)]214            download_bucket_files(bucket, pairs, token=token)215            for i, (_, dst) in enumerate(pairs):216                if dst.exists() and json.loads(dst.read_text()).get("state") == "done":217                    done.add(i)218        return done219 220    # --- attach-to-existing-run paths ---221    if args.resume:222        manifest = read_manifest(args.run_id)223        n = manifest["num_shards"]224        # Don't double-spawn ranks that are still running — wait for them, then diff.225        from huggingface_hub import list_jobs226        try:227            in_flight = [j for j in list_jobs(labels={"embedding-fleet-run": args.run_id},228                                              namespace=namespace, token=token)229                         if j.status.stage in (JobStage.RUNNING, JobStage.SCHEDULING)230                         and (j.labels or {}).get("rank") is not None]231        except Exception as e:232            logger.warning(f"in-flight check skipped ({e})")233            in_flight = []234        if in_flight:235            ranks = sorted({j.labels["rank"] for j in in_flight})236            logger.info(f"{len(in_flight)} worker(s) still in flight (ranks {ranks}) — waiting before resuming.")237            wait_for_job([j.id for j in in_flight], token=token)238        todo = sorted(set(range(n)) - done_ranks(args.run_id, n))239        if todo:240            logger.info(f"Resume {args.run_id}: {n - len(todo)}/{n} shards done; re-running {todo}")241            still_failed = execute_workers(todo, manifest)242            if still_failed:243                logger.error(f"Ranks still failing after retry: {still_failed} — investigate, then --resume again.")244                sys.exit(1)245        else:246            logger.info(f"Resume {args.run_id}: all {n} shards already done — consolidating.")247        job = spawn_consolidator(args.run_id)248        info = wait_for_job(job.id, token=token)249        if info.status.stage != JobStage.COMPLETED:250            logger.error(f"Consolidation failed (job {job.id}); run --resume again.")251            sys.exit(1)252        logger.info(f"✅ https://huggingface.co/datasets/{manifest['output_dataset']}")253        return254 255    if args.retry_rank is not None:256        manifest = read_manifest(args.run_id)257        if not 0 <= args.retry_rank < manifest["num_shards"]:258            p.error(f"--retry-rank must be in [0, {manifest['num_shards']}) for run {args.run_id}")259        logger.info(f"Re-spawning rank {args.retry_rank} of run {args.run_id} (config from manifest)")260        job = spawn_worker(args.retry_rank, manifest)261        info = wait_for_job(job.id, token=token)262        if info.status.stage != JobStage.COMPLETED:263            logger.error(f"Retry of rank {args.retry_rank} did not complete "264                         f"(stage={info.status.stage}, job {job.id}).")265            sys.exit(1)266        logger.info("Retry completed; run --consolidate-only when all shards are done.")267        return268 269    if args.consolidate_only:270        job = spawn_consolidator(args.run_id)271        info = wait_for_job(job.id, token=token)272        sys.exit(0 if info.status.stage == JobStage.COMPLETED else 1)273 274    # --- fresh run ---275    rows_total = args.rows_total or rows_in_split(args.input_dataset, args.config, args.split)276    if rows_total is None and not args.streaming:277        p.error("Couldn't determine the split's row count — pass --rows-total.")278    if rows_total is not None:279        if args.max_samples:280            rows_total = min(rows_total, args.max_samples)281        if not args.streaming and args.num_shards > rows_total:282            p.error(f"--num-shards {args.num_shards} exceeds the row count ({rows_total}) — "283                    f"some shards would be empty.")284 285    # Pin the input snapshot: every rank (and any later --retry-rank) must slice the IDENTICAL286    # revision, or a mid-run commit to the input dataset silently breaks the exact partition.287    from huggingface_hub import dataset_info288    revision = dataset_info(args.input_dataset, token=token).sha289    logger.info(f"Pinned input revision: {revision[:12]}")290 291    run_id = time.strftime("%Y%m%d-%H%M%S") + "-" + pysecrets.token_hex(3)292    create_bucket(bucket, private=True, exist_ok=True, token=token)293 294    manifest = {295        "run_id": run_id,296        "input_dataset": args.input_dataset,297        "output_dataset": args.output_dataset,298        "model": args.model,299        "column": args.column,300        "split": args.split,301        "config": args.config,302        "max_samples": args.max_samples,303        "num_shards": args.num_shards,304        "rows_total": rows_total,305        "flavor": args.flavor,306        "timeout": args.timeout,307        "private": args.private,308        "embed_args": list(args.embed_args),309        "streaming": args.streaming,310        "revision": revision,311        "started_at": time.time(),312        "job_ids": [],313    }314    put_json(bucket, f"runs/{run_id}/run.json", manifest, token=token)315    if rows_total is not None:316        logger.info(f"Run {run_id}: {rows_total:,} rows → {args.num_shards} shards "317                    f"(~{rows_total // args.num_shards:,} rows each) on {args.flavor}")318    else:319        logger.info(f"Run {run_id}: streaming file-shards × {args.num_shards} on {args.flavor} "320                    f"(row count unknown upfront)")321 322    jobs = [spawn_worker(rank, manifest) for rank in range(args.num_shards)]323    manifest["job_ids"] = [j.id for j in jobs]324    put_json(bucket, f"runs/{run_id}/run.json", manifest, token=token)325 326    logger.info(f"Manifest: hf://buckets/{bucket}/runs/{run_id}/run.json")327    logger.info(f"Dashboard: https://huggingface.co/spaces/davanstrien/embedding-fleet-dashboard?run={run_id}")328 329    if args.no_wait:330        logger.info(f"--no-wait: consolidate later with --run-id {run_id} --consolidate-only")331        return332 333    logger.info("Waiting for workers…")334    infos = wait_for_job([j.id for j in jobs], token=token)335    failed = [rank for rank, (j, i) in enumerate(zip(jobs, infos))336              if i.status.stage != JobStage.COMPLETED]337    if failed:338        logger.warning(f"{len(failed)} worker(s) did not complete; auto-retrying: {failed}")339        still_failed = execute_workers(failed, manifest)340        if still_failed:341            logger.error(f"Ranks still failing after retries: {still_failed}")342            logger.error(f"Investigate (job logs / heartbeat age), then converge with: "343                         f"--run-id {run_id} --resume")344            sys.exit(1)345 346    job = spawn_consolidator(run_id)347    info = wait_for_job(job.id, token=token)348    if info.status.stage != JobStage.COMPLETED:349        logger.error(f"Consolidation failed (job {job.id}); retry with --consolidate-only --run-id {run_id}")350        sys.exit(1)351    logger.info(f"✅ https://huggingface.co/datasets/{args.output_dataset}")352 353 354if __name__ == "__main__":355    main()356