From 38aa296de6bbe3b1c4e704c146bb506394fbd2b4 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 2 Jul 2026 16:50:58 -0700 Subject: [PATCH 01/10] prime images transfer-bulk --- .../prime/src/prime_cli/commands/images.py | 5 + .../src/prime_cli/commands/images_bulk.py | 91 +- .../commands/images_transfer_bulk.py | 864 ++++++++++++++++++ .../prime/tests/test_images_transfer_bulk.py | 492 ++++++++++ 4 files changed, 1432 insertions(+), 20 deletions(-) create mode 100644 packages/prime/src/prime_cli/commands/images_transfer_bulk.py create mode 100644 packages/prime/tests/test_images_transfer_bulk.py diff --git a/packages/prime/src/prime_cli/commands/images.py b/packages/prime/src/prime_cli/commands/images.py index f3210c1f..aa23469a 100644 --- a/packages/prime/src/prime_cli/commands/images.py +++ b/packages/prime/src/prime_cli/commands/images.py @@ -31,6 +31,7 @@ package_build_context, push_bulk, ) +from .images_transfer_bulk import transfer_bulk app = PlainTyper(help="Manage Docker images in Prime Intellect registry", no_args_is_help=True) console = get_console() @@ -740,6 +741,10 @@ def push_image( # Bulk push (JSONL manifest / Harbor task dirs) lives in images_bulk.py. app.command("push-bulk")(push_bulk) +# Bulk transfer (JSONL manifest / Harbor task dirs / Hugging Face datasets) +# lives in images_transfer_bulk.py. +app.command("transfer-bulk")(transfer_bulk) + @app.command("list", epilog=LIST_IMAGES_JSON_HELP) def list_images( diff --git a/packages/prime/src/prime_cli/commands/images_bulk.py b/packages/prime/src/prime_cli/commands/images_bulk.py index cbcf54af..8fa5cd12 100644 --- a/packages/prime/src/prime_cli/commands/images_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_bulk.py @@ -7,7 +7,7 @@ from collections import deque from dataclasses import dataclass from pathlib import Path -from typing import Any, Optional +from typing import Any, Callable, Optional if sys.version_info >= (3, 11): import tomllib @@ -68,6 +68,14 @@ class QuotaExceededError(APIError): """A 429 caused by the wallet's image count/storage quota, not the rate limiter.""" +class SubmitRateLimited(Exception): + """Raised by a submit callable to defer the spec instead of failing it.""" + + def __init__(self, message: str, retry_after: float): + super().__init__(message) + self.retry_after = retry_after + + @dataclass class BuildSpec: """One resolved build: where the context lives and what to call the image.""" @@ -97,13 +105,15 @@ def to_manifest_line(self) -> dict[str, str]: class BuildOutcome: """Terminal result for one spec. + ``spec`` is any object with ``image_ref``, ``source`` and + ``to_manifest_line()`` (BuildSpec here, TransferSpec for bulk transfers). ``status`` is a backend terminal status (COMPLETED/FAILED/CANCELLED) or a client-side one: SUBMIT_FAILED (initiate/upload/start failed), TIMEOUT (no terminal status within --build-timeout), SKIPPED (never submitted because the quota was exhausted mid-run). """ - spec: BuildSpec + spec: Any status: str build_id: Optional[str] = None full_image_path: Optional[str] = None @@ -163,7 +173,8 @@ def tar_filter(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]: # --------------------------------------------------------------------------- -def _duplicate_ref_problems(specs: list[BuildSpec], hint: str = "") -> list[str]: +def _duplicate_ref_problems(specs: list[Any], hint: str = "") -> list[str]: + """Report specs (anything with image_ref/source) sharing a destination ref.""" by_ref: dict[str, list[str]] = {} for spec in specs: by_ref.setdefault(spec.image_ref, []).append(spec.source) @@ -475,31 +486,37 @@ def _submit_build( @dataclass class _InFlightBuild: - spec: BuildSpec + spec: Any full_image_path: str deadline: float -def run_bulk_push( +def run_bulk_jobs( client: APIClient, - specs: list[BuildSpec], + specs: list[Any], *, - team_id: Optional[str], - visibility: Optional[ImageVisibility], + submit: Callable[[Any], tuple[str, str]], concurrency: int, build_timeout: int, + progress_description: str = "Pushing images", ) -> list[BuildOutcome]: - """Submit builds with a sliding window and poll them to terminal status. - - A finished build immediately frees a slot for the next queued build, so - one slow build never blocks the rest of a "batch". Once the wallet quota - is hit, submission stops (remaining specs become SKIPPED) but builds - already in flight are still polled to completion. + """Submit jobs with a sliding window and poll them to terminal status. + + ``submit(spec)`` starts one server-side build/transfer and returns + (build_id, full_image_path). A finished job immediately frees a slot for + the next queued spec, so one slow job never blocks the rest of a "batch". + Once the wallet quota is hit (QuotaExceededError), submission stops + (remaining specs become SKIPPED) but jobs already in flight are still + polled to completion. A submit that raises SubmitRateLimited requeues its + spec and pauses new submissions for ``retry_after`` seconds while polling + continues. """ - pending: deque[BuildSpec] = deque(specs) + pending: deque[Any] = deque(specs) in_flight: dict[str, _InFlightBuild] = {} outcomes: list[BuildOutcome] = [] quota_error: Optional[str] = None + submit_gate = 0.0 + rate_limit_notice_shown = False progress = Progress( TextColumn("[progress.description]{task.description}"), @@ -509,7 +526,7 @@ def run_bulk_push( console=console, ) with progress: - task_id = progress.add_task("Pushing images", total=len(specs)) + task_id = progress.add_task(progress_description, total=len(specs)) def record(outcome: BuildOutcome) -> None: outcomes.append(outcome) @@ -526,11 +543,21 @@ def record(outcome: BuildOutcome) -> None: while pending or in_flight: while pending and quota_error is None and len(in_flight) < concurrency: + if time.monotonic() < submit_gate: + break spec = pending.popleft() try: - build_id, full_image_path = _submit_build( - client, spec, team_id=team_id, visibility=visibility - ) + build_id, full_image_path = submit(spec) + except SubmitRateLimited as e: + pending.appendleft(spec) + submit_gate = time.monotonic() + e.retry_after + if not rate_limit_notice_shown: + rate_limit_notice_shown = True + progress.console.print( + "[dim]Server rate limit reached; pacing submissions " + "(jobs already submitted keep running)...[/dim]" + ) + break except QuotaExceededError as e: quota_error = str(e) record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) @@ -556,9 +583,14 @@ def record(outcome: BuildOutcome) -> None: ) pending.clear() - if not in_flight: + if not pending and not in_flight: break + if not in_flight: + # Everything left is waiting on the submit gate. + time.sleep(max(0.0, submit_gate - time.monotonic())) + continue + time.sleep(POLL_INTERVAL_SECONDS) for build_id, entry in list(in_flight.items()): @@ -603,6 +635,25 @@ def record(outcome: BuildOutcome) -> None: return outcomes +def run_bulk_push( + client: APIClient, + specs: list[BuildSpec], + *, + team_id: Optional[str], + visibility: Optional[ImageVisibility], + concurrency: int, + build_timeout: int, +) -> list[BuildOutcome]: + """Bulk build-and-push: submit build contexts through the shared engine.""" + return run_bulk_jobs( + client, + specs, + submit=lambda spec: _submit_build(client, spec, team_id=team_id, visibility=visibility), + concurrency=concurrency, + build_timeout=build_timeout, + ) + + def _write_failures_manifest(path: Path, failures: list[BuildOutcome]) -> None: lines = [json.dumps(outcome.spec.to_manifest_line()) for outcome in failures] path.write_text("\n".join(lines) + "\n") diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py new file mode 100644 index 00000000..12e35587 --- /dev/null +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -0,0 +1,864 @@ +import json +import os +import re +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + +import click +import httpx +import typer +from prime_sandboxes import ( + APIClient, + APIError, + Config, + ImageVisibility, + UnauthorizedError, +) +from rich.table import Table + +from ..utils import get_console +from .images_bulk import ( + _TAG_RE, + DEFAULT_BUILD_TIMEOUT_SECONDS, + DEFAULT_MAX_IN_FLIGHT, + FAILURE_TABLE_MAX_ROWS, + POLL_INTERVAL_SECONDS, + SUPPORTED_PLATFORMS, + BulkPushValidationError, + QuotaExceededError, + SubmitRateLimited, + _duplicate_ref_problems, + _is_http_429, + _is_quota_429, + _write_failures_manifest, + discover_harbor_tasks, + run_bulk_jobs, +) + +console = get_console() + +# How long to pause new submissions after the server's transfer rate limiter +# rejects one. +TRANSFER_RATE_LIMIT_PAUSE_SECONDS = 15.0 + +HF_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" +HF_ROWS_PAGE_LENGTH = 100 +HF_IMAGE_COLUMN_CANDIDATES = ("docker_image", "image_name", "container_image", "image") +HF_REQUEST_TIMEOUT_SECONDS = 30.0 +HF_REQUEST_MAX_ATTEMPTS = 6 +HF_REQUEST_BACKOFF_SECONDS = 2.0 +HF_REQUEST_BACKOFF_MAX_SECONDS = 60.0 + +_TRANSFER_MANIFEST_KEYS = {"source", "image", "platform"} + +_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") + + +@dataclass +class TransferSpec: + """One resolved transfer: a source registry image and its Prime destination. + + ``dest_name``/``dest_tag`` are what the transfer will be called. When + ``override`` is False they are derived from the source the same way the + server derives them and are only used for display and duplicate detection; + when True the user chose them and they are sent with the request. + """ + + source_image: str + dest_name: str + dest_tag: str + platform: str + source: str + override: bool = False + + @property + def image_ref(self) -> str: + return f"{self.dest_name}:{self.dest_tag}" + + def to_manifest_line(self) -> dict[str, str]: + line = {"source": self.source_image, "platform": self.platform} + if self.override: + line["image"] = self.image_ref + return line + + +def derive_transfer_destination(source_ref: str) -> tuple[str, str]: + """Derive the (name, tag) a transfer gets when no destination is given. + + Mirrors the server's parsing (last repository segment lowercased; tag kept, + or ``sha256-`` for digest refs). Used client-side for dry + runs and duplicate-destination detection; the server stays authoritative. + """ + ref = (source_ref or "").strip() + if not ref: + raise ValueError("empty image reference") + + digest: Optional[str] = None + if "@" in ref: + ref, _, digest = ref.rpartition("@") + if not _DIGEST_RE.match(digest): + raise ValueError("unsupported digest (only sha256 is supported)") + + first_segment = ref.split("/", 1)[0] + has_registry = "/" in ref and ( + "." in first_segment or ":" in first_segment or first_segment == "localhost" + ) + path_part = ref[len(first_segment) + 1 :] if has_registry else ref + if not path_part: + raise ValueError("missing repository") + + tag: Optional[str] = None + if ":" in path_part: + path_part, _, tag = path_part.rpartition(":") + if not tag: + raise ValueError("empty tag") + elif digest is None: + tag = "latest" + + repository = path_part.strip("/") + if not repository: + raise ValueError("missing repository") + + name = repository.split("/")[-1].lower() + if tag is None: + assert digest is not None + tag = "sha256-" + digest.split(":", 1)[1][:16] + return name, tag + + +_DUPLICATE_DEST_HINT = ( + ' — transfer them via a JSONL manifest with distinct "image" destination overrides' +) + + +# --------------------------------------------------------------------------- +# Transfer-list resolution: JSONL manifest +# --------------------------------------------------------------------------- + + +def load_transfer_manifest(manifest_path: Path, default_platform: str) -> list[TransferSpec]: + """Parse and fully validate a JSONL transfer manifest.""" + problems: list[str] = [] + specs: list[TransferSpec] = [] + + for lineno, raw_line in enumerate(manifest_path.read_text().splitlines(), start=1): + line = raw_line.strip() + if not line: + continue + where = f"{manifest_path.name}:{lineno}" + try: + entry = json.loads(line) + except json.JSONDecodeError as e: + problems.append(f"{where}: invalid JSON ({e})") + continue + if not isinstance(entry, dict): + problems.append(f"{where}: expected a JSON object") + continue + unknown = sorted(set(entry) - _TRANSFER_MANIFEST_KEYS) + if unknown: + problems.append( + f"{where}: unknown key(s) {', '.join(unknown)} " + f"(expected: {', '.join(sorted(_TRANSFER_MANIFEST_KEYS))})" + ) + continue + + source = entry.get("source") + if not source or not isinstance(source, str): + problems.append(f"{where}: 'source' is required") + continue + source = source.strip() + + platform = entry.get("platform") or default_platform + if platform not in SUPPORTED_PLATFORMS: + problems.append( + f"{where}: unsupported platform '{platform}' " + f"(supported: {', '.join(SUPPORTED_PLATFORMS)})" + ) + continue + + try: + derived_name, derived_tag = derive_transfer_destination(source) + except ValueError as e: + problems.append(f"{where}: invalid source '{source}' ({e})") + continue + + image = entry.get("image") + if image is not None and not isinstance(image, str): + problems.append(f"{where}: 'image' must be a string") + continue + if image: + if ":" in image: + dest_name, dest_tag = image.rsplit(":", 1) + else: + dest_name, dest_tag = image, "latest" + if not dest_name or "/" in dest_name: + problems.append( + f"{where}: invalid destination '{image}'; use simple names like 'myapp:v1'" + ) + continue + if not _TAG_RE.match(dest_tag): + problems.append(f"{where}: invalid destination tag '{dest_tag}'") + continue + override = True + else: + dest_name, dest_tag = derived_name, derived_tag + override = False + + specs.append( + TransferSpec( + source_image=source, + dest_name=dest_name, + dest_tag=dest_tag, + platform=platform, + source=where, + override=override, + ) + ) + + problems.extend(_duplicate_ref_problems(specs, hint=_DUPLICATE_DEST_HINT)) + if not specs and not problems: + problems.append(f"{manifest_path.name}: manifest contains no transfers") + if problems: + raise BulkPushValidationError(problems) + return specs + + +# --------------------------------------------------------------------------- +# Transfer-list resolution: Harbor task directories +# --------------------------------------------------------------------------- + + +def load_harbor_transfer_specs( + root: Path, *, platform: str +) -> tuple[list[TransferSpec], list[tuple[str, str]]]: + """Resolve transfer specs from a Harbor tasks directory. + + The complement of push-bulk's Harbor mode: it collects the prebuilt + ``[environment] docker_image`` references that push-bulk skips. Returns + (specs, skipped) where skipped lists tasks with nothing to transfer. + Tasks sharing the same prebuilt image collapse into one transfer. + """ + tasks = discover_harbor_tasks(root) + if not tasks: + raise BulkPushValidationError( + [ + f"no Harbor tasks found under {root} " + "(a task directory contains task.toml and environment/)" + ] + ) + + problems: list[str] = [] + specs: list[TransferSpec] = [] + skipped: list[tuple[str, str]] = [] + first_task_by_source: dict[str, str] = {} + for task_dir in tasks: + try: + with open(task_dir / "task.toml", "rb") as f: + config = tomllib.load(f) + except Exception as e: + problems.append(f"{task_dir.name}: failed to parse task.toml ({e})") + continue + + environment = config.get("environment") or {} + docker_image = environment.get("docker_image") if isinstance(environment, dict) else None + if not docker_image or not isinstance(docker_image, str): + skipped.append( + (task_dir.name, "no prebuilt docker_image (build it with push-bulk instead)") + ) + continue + docker_image = docker_image.strip() + + first_task = first_task_by_source.get(docker_image) + if first_task is not None: + skipped.append((task_dir.name, f"same image as {first_task} ({docker_image})")) + continue + first_task_by_source[docker_image] = task_dir.name + + try: + dest_name, dest_tag = derive_transfer_destination(docker_image) + except ValueError as e: + problems.append(f"{task_dir.name}: invalid docker_image '{docker_image}' ({e})") + continue + + specs.append( + TransferSpec( + source_image=docker_image, + dest_name=dest_name, + dest_tag=dest_tag, + platform=platform, + source=task_dir.name, + override=False, + ) + ) + + problems.extend(_duplicate_ref_problems(specs, hint=_DUPLICATE_DEST_HINT)) + if not specs and not problems: + problems.append( + f"no tasks with a prebuilt docker_image under {root} " + f"({len(skipped)} skipped: {', '.join(name for name, _ in skipped)})" + ) + if problems: + raise BulkPushValidationError(problems) + return specs, skipped + + +# --------------------------------------------------------------------------- +# Transfer-list resolution: Hugging Face datasets +# --------------------------------------------------------------------------- + + +def normalize_hf_dataset_id(raw: str) -> str: + """Accept 'org/name', hf://datasets/... and huggingface.co dataset URLs.""" + dataset = (raw or "").strip() + for prefix in ( + "hf://datasets/", + "https://huggingface.co/datasets/", + "http://huggingface.co/datasets/", + "huggingface.co/datasets/", + ): + if dataset.startswith(prefix): + dataset = dataset[len(prefix) :] + dataset = dataset.split("?", 1)[0].split("#", 1)[0] + # Drop trailing URL parts like /viewer/default/train. + parts = [part for part in dataset.split("/") if part] + dataset = "/".join(parts[:2]) + break + dataset = dataset.strip("/") + if not dataset or dataset.count("/") > 1: + raise BulkPushValidationError( + [f"invalid Hugging Face dataset '{raw}' (expected 'org/name' or a dataset URL)"] + ) + return dataset + + +def _hf_get(path: str, params: dict[str, Any]) -> dict[str, Any]: + """GET from the HF datasets-server with retries on transient failures. + + Uses HF_TOKEN from the environment for gated/private datasets. The + datasets-server rate limit is per IP and clears within about a minute, so + 429/5xx retries back off long enough to ride out a full window, honoring + Retry-After when the server sends one. + """ + token = os.environ.get("HF_TOKEN") + headers = {"Authorization": f"Bearer {token}"} if token else {} + url = f"{HF_DATASETS_SERVER_URL}{path}" + delay = HF_REQUEST_BACKOFF_SECONDS + for attempt in range(1, HF_REQUEST_MAX_ATTEMPTS + 1): + retry_delay = delay + delay = min(delay * 2, HF_REQUEST_BACKOFF_MAX_SECONDS) + try: + response = httpx.get( + url, params=params, headers=headers, timeout=HF_REQUEST_TIMEOUT_SECONDS + ) + except httpx.HTTPError as e: + if attempt == HF_REQUEST_MAX_ATTEMPTS: + raise BulkPushValidationError([f"Hugging Face request failed: {e}"]) + time.sleep(retry_delay) + continue + if response.status_code == 429 or response.status_code >= 500: + if attempt == HF_REQUEST_MAX_ATTEMPTS: + rate_limit_hint = ( + " (rate limited — wait a minute and re-run)" + if response.status_code == 429 + else "" + ) + raise BulkPushValidationError( + [ + f"Hugging Face request failed with HTTP " + f"{response.status_code}: {url}{rate_limit_hint}" + ] + ) + retry_after = response.headers.get("retry-after") + if retry_after: + try: + retry_delay = max(retry_delay, float(retry_after)) + except ValueError: + pass + time.sleep(retry_delay) + continue + if response.status_code != 200: + try: + detail = response.json().get("error") or response.text + except ValueError: + detail = response.text + raise BulkPushValidationError( + [ + f"Hugging Face request failed with HTTP {response.status_code} " + f"({detail}). The dataset may be private/gated (set HF_TOKEN) or have " + "the dataset viewer disabled — as a fallback, extract the image " + "references yourself and use --manifest" + ] + ) + return response.json() + raise AssertionError("unreachable") + + +def load_hf_specs( + dataset: str, + *, + config: Optional[str], + split: str, + column: Optional[str], + platform: str, +) -> tuple[list[TransferSpec], list[str]]: + """Resolve transfer specs from a Hugging Face dataset column. + + Pages the dataset through the datasets-server rows API (no local + `datasets` dependency), dedupes identical references preserving order, + and returns (specs, notes) where notes are informational messages. + """ + dataset_id = normalize_hf_dataset_id(dataset) + notes: list[str] = [] + + info = _hf_get("/info", {"dataset": dataset_id}) + configs = info.get("dataset_info") + if not isinstance(configs, dict) or not configs: + raise BulkPushValidationError([f"no dataset info available for '{dataset_id}'"]) + if config is None: + if len(configs) > 1: + raise BulkPushValidationError( + [ + f"dataset '{dataset_id}' has multiple configs " + f"({', '.join(sorted(configs))}); choose one with --hf-config" + ] + ) + config = next(iter(configs)) + elif config not in configs: + raise BulkPushValidationError( + [ + f"config '{config}' not found in dataset '{dataset_id}' " + f"(available: {', '.join(sorted(configs))})" + ] + ) + config_info = configs.get(config) or {} + + features = config_info.get("features") or {} + if column is None: + column = next((c for c in HF_IMAGE_COLUMN_CANDIDATES if c in features), None) + if column is None: + raise BulkPushValidationError( + [ + f"could not auto-detect an image column in '{dataset_id}' " + f"(looked for: {', '.join(HF_IMAGE_COLUMN_CANDIDATES)}; " + f"dataset columns: {', '.join(sorted(features)) or 'unknown'}); " + "pass --column" + ] + ) + notes.append(f"Using column '{column}' (auto-detected)") + elif features and column not in features: + raise BulkPushValidationError( + [ + f"column '{column}' not found in dataset '{dataset_id}' " + f"(columns: {', '.join(sorted(features))})" + ] + ) + + splits = config_info.get("splits") + if isinstance(splits, dict) and splits and split not in splits: + raise BulkPushValidationError( + [ + f"split '{split}' not found in dataset '{dataset_id}' " + f"(available: {', '.join(sorted(splits))})" + ] + ) + if info.get("partial"): + notes.append( + "Warning: the dataset viewer only covers part of this dataset; some rows may be missing" + ) + + sources: list[tuple[str, int]] = [] # (image ref, first row index), order-preserving + seen: set[str] = set() + scanned = 0 + empty_rows = 0 + offset = 0 + total: Optional[int] = None + while True: + page = _hf_get( + "/rows", + { + "dataset": dataset_id, + "config": config, + "split": split, + "offset": offset, + "length": HF_ROWS_PAGE_LENGTH, + }, + ) + rows = page.get("rows") or [] + if isinstance(page.get("num_rows_total"), int): + total = page["num_rows_total"] + for item in rows: + row = item.get("row") or {} + row_idx = item.get("row_idx", offset) + value = row.get(column) + if not isinstance(value, str) or not value.strip(): + empty_rows += 1 + continue + scanned += 1 + value = value.strip() + if value in seen: + continue + seen.add(value) + sources.append((value, row_idx)) + offset += len(rows) + if not rows or (total is not None and offset >= total): + break + + if empty_rows: + notes.append(f"Skipped {empty_rows} row(s) with an empty '{column}' value") + duplicates = scanned - len(sources) + if duplicates: + notes.append(f"Collapsed {duplicates} duplicate image reference(s)") + + problems: list[str] = [] + specs: list[TransferSpec] = [] + for ref, row_idx in sources: + try: + dest_name, dest_tag = derive_transfer_destination(ref) + except ValueError as e: + problems.append(f"{dataset_id} row {row_idx}: invalid image reference '{ref}' ({e})") + continue + specs.append( + TransferSpec( + source_image=ref, + dest_name=dest_name, + dest_tag=dest_tag, + platform=platform, + source=f"row {row_idx}", + override=False, + ) + ) + + problems.extend(_duplicate_ref_problems(specs, hint=_DUPLICATE_DEST_HINT)) + if not specs and not problems: + problems.append(f"no image references found in '{dataset_id}' column '{column}'") + if problems: + raise BulkPushValidationError(problems) + return specs, notes + + +# --------------------------------------------------------------------------- +# Submission +# --------------------------------------------------------------------------- + + +def _submit_transfer( + client: APIClient, + spec: TransferSpec, + *, + team_id: Optional[str], + visibility: Optional[ImageVisibility], +) -> tuple[str, str]: + """Queue one transfer. Returns (build_id, full_image_path). + + Transfers are counted per image by the server's transfer rate limiter + (a rolling window), so a plain 429 means "later", not "failed": it is + surfaced as SubmitRateLimited so the engine requeues the spec. Wallet + quota 429s become QuotaExceededError and stop the run. + """ + payload: dict[str, Any] = { + "source_image": spec.source_image, + "platform": spec.platform, + } + if spec.override: + payload["image_name"] = spec.dest_name + payload["image_tag"] = spec.dest_tag + if team_id: + payload["team_id"] = team_id + if visibility is not None: + payload["visibility"] = visibility.value + + try: + response = client.request("POST", "/images/build", json=payload) + except UnauthorizedError: + raise + except APIError as e: + if _is_quota_429(e): + raise QuotaExceededError(str(e)) from e + if _is_http_429(e): + raise SubmitRateLimited(str(e), retry_after=TRANSFER_RATE_LIMIT_PAUSE_SECONDS) from e + raise + + build_id = response.get("build_id") + if not build_id: + raise APIError("invalid response from server (missing build_id)") + return build_id, response.get("fullImagePath") or spec.image_ref + + +# --------------------------------------------------------------------------- +# Command +# --------------------------------------------------------------------------- + + +def transfer_bulk( + manifest: Optional[Path] = typer.Option( + None, + "--manifest", + "-m", + help=( + "JSONL manifest of transfers; each line is " + '{"source": "registry/repo:tag", "image"?: "name:tag", "platform"?: "..."}' + ), + ), + harbor: Optional[Path] = typer.Option( + None, + "--harbor", + help=( + "Harbor task directory (or directory of tasks); transfers each task's " + "prebuilt [environment] docker_image" + ), + ), + hf_dataset: Optional[str] = typer.Option( + None, + "--hf", + help=( + "Hugging Face dataset id or URL (e.g. 'R2E-Gym/R2E-Gym-Subset'); " + "transfers every image referenced in the dataset" + ), + ), + hf_split: str = typer.Option("train", "--hf-split", help="Dataset split for --hf mode"), + hf_config: Optional[str] = typer.Option( + None, + "--hf-config", + help="Dataset config for --hf mode", + show_default="the dataset's only config", + ), + column: Optional[str] = typer.Option( + None, + "--column", + help="Dataset column holding image references for --hf mode", + show_default=f"auto-detect: {', '.join(HF_IMAGE_COLUMN_CANDIDATES)}", + ), + platform: str = typer.Option( + "linux/amd64", + "--platform", + click_type=click.Choice(list(SUPPORTED_PLATFORMS)), + help="Required platform for the transferred images (manifest lines may override)", + ), + public: bool = typer.Option( + False, "--public", help="Make the images public when the transfers complete" + ), + private: bool = typer.Option( + False, "--private", help="Make the images private when the transfers complete" + ), + concurrency: int = typer.Option( + DEFAULT_MAX_IN_FLIGHT, "--concurrency", help="Maximum transfers in flight at once" + ), + build_timeout: int = typer.Option( + DEFAULT_BUILD_TIMEOUT_SECONDS, + "--build-timeout", + help="Seconds to wait for a single transfer before giving up on it", + ), + dry_run: bool = typer.Option( + False, "--dry-run", help="Resolve and print the transfer list without transferring" + ), + failures_out: Path = typer.Option( + Path("transfer-bulk-failures.jsonl"), + "--failures-out", + help="Where to write a re-runnable manifest of failed transfers", + ), +): + """ + Copy many existing registry images into Prime in one command. + + Reads image references from a JSONL manifest (--manifest), a Harbor tasks + directory (--harbor), or a Hugging Face dataset (--hf), validates + everything up front, then queues transfers server-side, pacing submissions + to the server's transfer rate limit and polling until every transfer + finishes. Failed transfers are written to a manifest you can re-run. + + \b + Manifest format (one JSON object per line; "image" overrides the + destination, which otherwise derives from the source reference): + {"source": "ghcr.io/org/app:v1"} + {"source": "docker.io/other/app:v1", "image": "other-app:v1"} + + \b + Harbor mode transfers the prebuilt docker_image of each task — exactly the + tasks push-bulk skips, so the two commands together cover a task set. + + \b + Hugging Face mode reads image references straight from a dataset column + (auto-detected, e.g. docker_image) using the dataset viewer API — no local + dataset download. Set HF_TOKEN for private or gated datasets. + + \b + Examples: + prime images transfer-bulk --manifest transfers.jsonl + prime images transfer-bulk --harbor ./tasks + prime images transfer-bulk --hf R2E-Gym/R2E-Gym-Subset --dry-run + prime images transfer-bulk --hf org/dataset --hf-split test --column image_name + prime images transfer-bulk --manifest transfer-bulk-failures.jsonl + """ + try: + modes = [m for m in (manifest, harbor, hf_dataset) if m is not None] + if len(modes) != 1: + console.print("[red]Error: Provide exactly one of --manifest, --harbor or --hf[/red]") + raise typer.Exit(1) + if hf_dataset is None and (hf_split != "train" or hf_config or column): + console.print( + "[red]Error: --hf-split, --hf-config and --column only apply to --hf mode[/red]" + ) + raise typer.Exit(1) + if public and private: + console.print("[red]Error: --public and --private cannot be used together[/red]") + raise typer.Exit(1) + if concurrency < 1: + console.print("[red]Error: --concurrency must be at least 1[/red]") + raise typer.Exit(1) + if build_timeout < 1: + console.print("[red]Error: --build-timeout must be at least 1[/red]") + raise typer.Exit(1) + + skipped: list[tuple[str, str]] = [] + notes: list[str] = [] + try: + if manifest is not None: + manifest_path = manifest.resolve() + if not manifest_path.is_file(): + raise BulkPushValidationError([f"manifest not found: {manifest_path}"]) + source_desc = f"manifest {manifest_path}" + specs = load_transfer_manifest(manifest_path, default_platform=platform) + elif harbor is not None: + harbor_root = harbor.resolve() + if not harbor_root.is_dir(): + raise BulkPushValidationError( + [f"Harbor task directory not found: {harbor_root}"] + ) + source_desc = f"Harbor tasks in {harbor_root}" + specs, skipped = load_harbor_transfer_specs(harbor_root, platform=platform) + else: + assert hf_dataset is not None + source_desc = f"Hugging Face dataset {normalize_hf_dataset_id(hf_dataset)}" + console.print(f"[cyan]Reading image references from {source_desc}...[/cyan]") + specs, notes = load_hf_specs( + hf_dataset, + config=hf_config, + split=hf_split, + column=column, + platform=platform, + ) + except BulkPushValidationError as e: + console.print( + f"[red]Error: cannot start bulk transfer ({len(e.problems)} problem(s)):[/red]" + ) + for problem in e.problems: + console.print(f"[red] - {problem}[/red]") + raise typer.Exit(1) + + for note in notes: + console.print(f"[dim]{note}[/dim]") + for task_name, reason in skipped: + console.print(f"[yellow]Skipping {task_name}: {reason}[/yellow]") + + if dry_run: + table = Table(title=f"Resolved {len(specs)} transfer(s)") + table.add_column("Source", style="cyan", overflow="fold") + table.add_column("Destination", overflow="fold") + table.add_column("Platform", no_wrap=True) + table.add_column("From", style="dim") + for spec in specs: + table.add_row(spec.source_image, spec.image_ref, spec.platform, spec.source) + console.print(table) + console.print("[dim]Dry run only — re-run without --dry-run to transfer.[/dim]") + return + + config = Config() + visibility: Optional[ImageVisibility] = None + if public: + visibility = ImageVisibility.PUBLIC + elif private: + visibility = ImageVisibility.PRIVATE + + console.print( + f"[bold blue]Bulk transferring {len(specs)} image(s)[/bold blue] " + f"[dim]({source_desc})[/dim]" + ) + if config.team_id: + console.print(f"[dim]Team: {config.team_id}[/dim]") + if visibility is not None: + console.print(f"[dim]Visibility: {visibility.value}[/dim]") + else: + console.print( + "[dim]Visibility: PRIVATE for new images " + "(existing tags keep their current visibility)[/dim]" + ) + console.print( + f"[dim]Up to {concurrency} transfers in flight; " + f"polling every {int(POLL_INTERVAL_SECONDS)}s. Transfers are rate-limited " + "per account server-side, so large batches take a while to submit.[/dim]" + ) + console.print() + + client = APIClient() + outcomes = run_bulk_jobs( + client, + specs, + submit=lambda spec: _submit_transfer( + client, spec, team_id=config.team_id or None, visibility=visibility + ), + concurrency=concurrency, + build_timeout=build_timeout, + progress_description="Transferring images", + ) + + failures = [o for o in outcomes if o.status != "COMPLETED"] + completed_count = len(outcomes) - len(failures) + console.print() + console.print(f"[bold]{completed_count}/{len(outcomes)} transfers completed[/bold]") + if not failures: + console.print("[bold green]All images transferred successfully![/bold green]") + console.print() + console.print("[dim]Use them with: prime sandbox create [/dim]") + return + + failure_table = Table(title=f"{len(failures)} transfer(s) did not complete") + failure_table.add_column("Source", style="cyan", overflow="fold") + failure_table.add_column("Destination", overflow="fold") + failure_table.add_column("Status", no_wrap=True) + failure_table.add_column("From", style="dim") + failure_table.add_column("Error") + for outcome in failures[:FAILURE_TABLE_MAX_ROWS]: + failure_table.add_row( + outcome.spec.source_image, + outcome.spec.image_ref, + outcome.status, + outcome.spec.source, + outcome.error or "", + ) + console.print(failure_table) + if len(failures) > FAILURE_TABLE_MAX_ROWS: + console.print( + f"[dim]... and {len(failures) - FAILURE_TABLE_MAX_ROWS} more, " + f"all included in {failures_out}[/dim]" + ) + if any("limit exceeded" in (o.error or "").lower() for o in failures): + console.print( + "[red]Image quota reached — delete unused images (prime images delete) " + "or request a higher limit, then retry.[/red]" + ) + + _write_failures_manifest(failures_out, failures) + console.print() + console.print(f"Wrote {len(failures)} failed transfer(s) to [bold]{failures_out}[/bold]") + console.print( + f"[dim]Retry with: prime images transfer-bulk --manifest {failures_out}[/dim]" + ) + raise typer.Exit(1) + + except UnauthorizedError: + console.print("[red]Error: Not authenticated. Please run 'prime login' first.[/red]") + raise typer.Exit(1) + except KeyboardInterrupt: + console.print( + "\n[yellow]Cancelled. Transfers already queued keep running server-side; " + "check them with 'prime images list'.[/yellow]" + ) + raise typer.Exit(1) diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py new file mode 100644 index 00000000..833b1c8d --- /dev/null +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -0,0 +1,492 @@ +import json + +import prime_cli.commands.images_bulk as images_bulk +import prime_cli.commands.images_transfer_bulk as images_transfer_bulk +import pytest +from prime_cli.commands.images_transfer_bulk import derive_transfer_destination +from prime_cli.main import app +from prime_sandboxes import APIError +from typer.testing import CliRunner + +runner = CliRunner() + +TEST_ENV = { + "COLUMNS": "200", + "LINES": "50", + "PRIME_DISABLE_VERSION_CHECK": "1", + "PRIME_TEAM_ID": "", + "HF_TOKEN": "", +} + + +def _write_manifest(path, entries): + path.write_text("\n".join(json.dumps(e) for e in entries) + "\n") + + +def _make_harbor_task(root, dirname, *, docker_image=None, with_dockerfile=False): + task = root / dirname + env = task / "environment" + env.mkdir(parents=True) + if with_dockerfile: + (env / "Dockerfile").write_text("FROM busybox\n") + toml_lines = ['version = "1.0"', "", "[task]", f'name = "{dirname}"'] + if docker_image: + toml_lines += ["", "[environment]", f'docker_image = "{docker_image}"'] + (task / "task.toml").write_text("\n".join(toml_lines) + "\n") + return task + + +class FakeTransferAPI: + """Scriptable stand-in for APIClient, shared across one test invocation.""" + + def __init__(self): + self.calls = [] + self.payloads = [] + self.build_counter = 0 + # build_id -> list of statuses returned by successive GET polls; + # the last status repeats once the list is exhausted. + self.poll_scripts = {} + self.default_poll = ["COMPLETED"] + # Exceptions raised (FIFO) by POST /images/build before any transfer is queued. + self.build_error_queue = [] + + def request(self, method, path, json=None, params=None): + self.calls.append((method, path)) + if method == "POST" and path == "/images/build": + if self.build_error_queue: + raise self.build_error_queue.pop(0) + assert "source_image" in json + self.payloads.append(json) + self.build_counter += 1 + build_id = f"build-{self.build_counter}" + return { + "build_id": build_id, + "fullImagePath": f"user/{json.get('image_name') or 'derived'}", + } + if method == "GET" and path.startswith("/images/build/"): + build_id = path.rsplit("/", 1)[1] + script = self.poll_scripts.setdefault(build_id, list(self.default_poll)) + return {"status": script.pop(0) if len(script) > 1 else script[0]} + raise AssertionError(f"Unexpected request: {method} {path}") + + def post_build_count(self): + return len([c for c in self.calls if c == ("POST", "/images/build")]) + + +@pytest.fixture +def fake_api(monkeypatch): + api = FakeTransferAPI() + monkeypatch.setattr("prime_cli.main.check_for_update", lambda: (False, None)) + monkeypatch.setattr(images_transfer_bulk, "APIClient", lambda: api) + monkeypatch.setattr(images_bulk, "POLL_INTERVAL_SECONDS", 0.0) + monkeypatch.setattr(images_transfer_bulk, "TRANSFER_RATE_LIMIT_PAUSE_SECONDS", 0.0) + return api + + +def _fake_hf(monkeypatch, *, info, pages, total, rate_limit_first_rows=False): + """Serve datasets-server /info and /rows from canned data. + + ``pages`` maps a row offset to the list of row values (the image column) + served at that offset; None values become rows missing the column. With + ``rate_limit_first_rows`` the first /rows request gets a 429 response. + """ + calls = [] + state = {"rate_limited": rate_limit_first_rows} + + class FakeResponse: + def __init__(self, payload, status_code=200, headers=None): + self.status_code = status_code + self._payload = payload + self.text = json.dumps(payload) + self.headers = headers or {} + + def json(self): + return self._payload + + def fake_get(url, params=None, headers=None, timeout=None): + params = dict(params or {}) + calls.append((url, params)) + if url.endswith("/info"): + return FakeResponse(info) + if url.endswith("/rows"): + if state["rate_limited"]: + state["rate_limited"] = False + return FakeResponse({}, status_code=429, headers={"retry-after": "0"}) + offset = params["offset"] + values = pages.get(offset, []) + rows = [ + { + "row_idx": offset + i, + "row": {} if value is None else {"docker_image": value}, + } + for i, value in enumerate(values) + ] + return FakeResponse({"rows": rows, "num_rows_total": total}) + raise AssertionError(f"Unexpected HF request: {url}") + + monkeypatch.setattr(images_transfer_bulk.httpx, "get", fake_get) + return calls + + +HF_INFO_ONE_CONFIG = { + "dataset_info": { + "default": { + "features": {"docker_image": {"dtype": "string"}, "prompt": {"dtype": "string"}}, + "splits": {"train": {"name": "train"}}, + } + }, + "partial": False, +} + + +# --------------------------------------------------------------------------- +# Destination derivation +# --------------------------------------------------------------------------- + + +def test_derive_destination_matches_server_rules(): + assert derive_transfer_destination("ubuntu") == ("ubuntu", "latest") + assert derive_transfer_destination("docker.io/library/ubuntu:22.04") == ("ubuntu", "22.04") + assert derive_transfer_destination("ghcr.io/Org/My-App:v1") == ("my-app", "v1") + assert derive_transfer_destination("quay.io/org/app@sha256:" + "a" * 64) == ( + "app", + "sha256-" + "a" * 16, + ) + with pytest.raises(ValueError): + derive_transfer_destination("") + with pytest.raises(ValueError): + derive_transfer_destination("app@sha512:abc") + + +# --------------------------------------------------------------------------- +# Manifest mode +# --------------------------------------------------------------------------- + + +def test_manifest_happy_path(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest( + manifest, + [ + {"source": "docker.io/org/app:v1"}, + {"source": "ghcr.io/org/other:v2", "image": "renamed:v9"}, + ], + ) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 0, result.output + assert "All images transferred successfully" in result.output + assert fake_api.payloads[0]["source_image"] == "docker.io/org/app:v1" + assert "image_name" not in fake_api.payloads[0] + assert fake_api.payloads[1]["image_name"] == "renamed" + assert fake_api.payloads[1]["image_tag"] == "v9" + assert all(p["platform"] == "linux/amd64" for p in fake_api.payloads) + + +def test_manifest_validation_fails_before_any_submission(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + manifest.write_text( + "\n".join( + [ + "not json", + json.dumps({"source": "org/app:v1", "context": "./x"}), + json.dumps({"image": "no-source:v1"}), + # Same derived destination app:v2 from two different namespaces. + json.dumps({"source": "a/app:v2"}), + json.dumps({"source": "b/app:v2"}), + ] + ) + + "\n" + ) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 1 + assert "invalid JSON" in result.output + assert "unknown key(s) context" in result.output + assert "'source' is required" in result.output + assert "duplicate image reference 'app:v2'" in result.output + assert fake_api.post_build_count() == 0 + + +def test_manifest_destination_override_resolves_duplicates(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest( + manifest, + [ + {"source": "a/app:v2"}, + {"source": "b/app:v2", "image": "app-b:v2"}, + ], + ) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 0, result.output + assert fake_api.post_build_count() == 2 + + +def test_dry_run_prints_derived_destinations_without_api_calls(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "ghcr.io/org/My-App:v1"}]) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--dry-run"], + env=TEST_ENV, + ) + assert result.exit_code == 0, result.output + assert "my-app:v1" in result.output + assert "Dry run only" in result.output + assert fake_api.calls == [] + + +# --------------------------------------------------------------------------- +# Harbor mode +# --------------------------------------------------------------------------- + + +def test_harbor_transfers_prebuilt_images_and_skips_buildable_tasks(tmp_path, fake_api): + root = tmp_path / "tasks" + _make_harbor_task(root, "task-a", docker_image="ghcr.io/org/img-a:v1") + _make_harbor_task(root, "task-b", with_dockerfile=True) + _make_harbor_task(root, "task-c", docker_image="ghcr.io/org/img-a:v1") + result = runner.invoke(app, ["images", "transfer-bulk", "--harbor", str(root)], env=TEST_ENV) + assert result.exit_code == 0, result.output + assert "Skipping task-b" in result.output + assert "Skipping task-c: same image as task-a" in result.output + assert fake_api.post_build_count() == 1 + assert fake_api.payloads[0]["source_image"] == "ghcr.io/org/img-a:v1" + + +def test_harbor_with_no_transferable_tasks_fails(tmp_path, fake_api): + root = tmp_path / "tasks" + _make_harbor_task(root, "task-a", with_dockerfile=True) + result = runner.invoke(app, ["images", "transfer-bulk", "--harbor", str(root)], env=TEST_ENV) + assert result.exit_code == 1 + assert "no tasks with a prebuilt docker_image" in result.output + assert fake_api.post_build_count() == 0 + + +# --------------------------------------------------------------------------- +# Hugging Face mode +# --------------------------------------------------------------------------- + + +def test_hf_pages_rows_dedupes_and_autodetects_column(tmp_path, fake_api, monkeypatch): + hf_calls = _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={ + 0: ["org/img-a:v1", "org/img-b:v1", "org/img-a:v1"], + 3: ["org/img-c:v1", None], + }, + total=5, + ) + result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "Org/DataSet"], env=TEST_ENV) + assert result.exit_code == 0, result.output + assert "Using column 'docker_image' (auto-detected)" in result.output + assert "Collapsed 1 duplicate image reference(s)" in result.output + assert "Skipped 1 row(s) with an empty 'docker_image' value" in result.output + assert fake_api.post_build_count() == 3 + sources = [p["source_image"] for p in fake_api.payloads] + assert sources == ["org/img-a:v1", "org/img-b:v1", "org/img-c:v1"] + row_requests = [params for url, params in hf_calls if url.endswith("/rows")] + assert [p["offset"] for p in row_requests] == [0, 3] + assert all(p["dataset"] == "Org/DataSet" for p in row_requests) + + +def test_hf_rate_limited_request_retries(tmp_path, fake_api, monkeypatch): + monkeypatch.setattr(images_transfer_bulk, "HF_REQUEST_BACKOFF_SECONDS", 0.0) + hf_calls = _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={0: ["org/img-a:v1"]}, + total=1, + rate_limit_first_rows=True, + ) + result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + assert result.exit_code == 0, result.output + assert fake_api.post_build_count() == 1 + # The rate-limited /rows request is retried: info + 429 + 200. + assert len(hf_calls) == 3 + + +def test_hf_accepts_dataset_url(tmp_path, fake_api, monkeypatch): + hf_calls = _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: ["org/img-a:v1"]}, total=1) + result = runner.invoke( + app, + [ + "images", + "transfer-bulk", + "--hf", + "https://huggingface.co/datasets/R2E-Gym/R2E-Gym-Subset/viewer/default/train", + ], + env=TEST_ENV, + ) + assert result.exit_code == 0, result.output + assert all(params["dataset"] == "R2E-Gym/R2E-Gym-Subset" for _, params in hf_calls) + + +def test_hf_multiple_configs_requires_flag(tmp_path, fake_api, monkeypatch): + info = { + "dataset_info": { + "cfg-a": HF_INFO_ONE_CONFIG["dataset_info"]["default"], + "cfg-b": HF_INFO_ONE_CONFIG["dataset_info"]["default"], + } + } + _fake_hf(monkeypatch, info=info, pages={0: ["org/img-a:v1"]}, total=1) + result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + assert result.exit_code == 1 + assert "multiple configs" in result.output + + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--hf-config", "cfg-a"], + env=TEST_ENV, + ) + assert result.exit_code == 0, result.output + + +def test_hf_unknown_column_and_split_fail(tmp_path, fake_api, monkeypatch): + _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--column", "nope"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "column 'nope' not found" in result.output + + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--hf-split", "test"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "split 'test' not found" in result.output + + +def test_hf_no_detectable_column_asks_for_flag(tmp_path, fake_api, monkeypatch): + info = { + "dataset_info": { + "default": { + "features": {"prompt": {"dtype": "string"}}, + "splits": {"train": {"name": "train"}}, + } + } + } + _fake_hf(monkeypatch, info=info, pages={0: []}, total=0) + result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + assert result.exit_code == 1 + assert "pass --column" in result.output + + +# --------------------------------------------------------------------------- +# Submission behavior +# --------------------------------------------------------------------------- + + +def test_rate_limited_submit_defers_and_retries(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app-a:v1"}, {"source": "b/app-b:v1"}]) + fake_api.build_error_queue.append( + APIError("HTTP 429: Image transfer rate limit exceeded: 10/10 images") + ) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 0, result.output + assert "pacing submissions" in result.output + assert "2/2 transfers completed" in result.output + # First POST is rejected, the spec is requeued, then both succeed. + assert fake_api.post_build_count() == 3 + + +def test_quota_429_aborts_and_skips_remaining(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest( + manifest, + [{"source": "a/app-a:v1"}, {"source": "b/app-b:v1"}, {"source": "c/app-c:v1"}], + ) + fake_api.build_error_queue.append( + APIError("HTTP 429: Image limit exceeded. Personal has 100 images") + ) + failures_out = tmp_path / "failures.jsonl" + result = runner.invoke( + app, + [ + "images", + "transfer-bulk", + "--manifest", + str(manifest), + "--failures-out", + str(failures_out), + ], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "0/3 transfers completed" in result.output + assert "Image quota reached" in result.output + lines = [json.loads(line) for line in failures_out.read_text().splitlines()] + assert [entry["source"] for entry in lines] == ["a/app-a:v1", "b/app-b:v1", "c/app-c:v1"] + assert fake_api.post_build_count() == 1 + + +def test_failed_transfer_writes_rerunnable_failures_manifest(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest( + manifest, + [{"source": "a/app-a:v1"}, {"source": "b/app-b:v1", "image": "renamed:v1"}], + ) + fake_api.poll_scripts["build-2"] = ["PENDING", "FAILED"] + failures_out = tmp_path / "failures.jsonl" + result = runner.invoke( + app, + [ + "images", + "transfer-bulk", + "--manifest", + str(manifest), + "--failures-out", + str(failures_out), + ], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "1/2 transfers completed" in result.output + lines = [json.loads(line) for line in failures_out.read_text().splitlines()] + assert lines == [{"source": "b/app-b:v1", "platform": "linux/amd64", "image": "renamed:v1"}] + assert f"transfer-bulk --manifest {failures_out}" in result.output + + +# --------------------------------------------------------------------------- +# Mode validation +# --------------------------------------------------------------------------- + + +def test_requires_exactly_one_mode(tmp_path, fake_api): + result = runner.invoke(app, ["images", "transfer-bulk"], env=TEST_ENV) + assert result.exit_code == 1 + assert "exactly one of --manifest, --harbor or --hf" in result.output + + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app:v1"}]) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--hf", "org/ds"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "exactly one of --manifest, --harbor or --hf" in result.output + + +def test_hf_flags_rejected_outside_hf_mode(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app:v1"}]) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--column", "docker_image"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "only apply to --hf mode" in result.output From cb1fc1cdeaff5f3942747c927738b123ff652840 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 2 Jul 2026 21:31:18 -0700 Subject: [PATCH 02/10] 429s --- .../src/prime_cli/commands/images_bulk.py | 42 ++++++++-- .../commands/images_transfer_bulk.py | 29 ++++++- .../prime/tests/test_images_transfer_bulk.py | 81 +++++++++++++++++++ 3 files changed, 144 insertions(+), 8 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_bulk.py b/packages/prime/src/prime_cli/commands/images_bulk.py index 8fa5cd12..daf7697f 100644 --- a/packages/prime/src/prime_cli/commands/images_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_bulk.py @@ -47,6 +47,14 @@ RATE_LIMIT_BACKOFF_INITIAL_SECONDS = 2.0 RATE_LIMIT_BACKOFF_MAX_SECONDS = 60.0 +# Liveness bound for SubmitRateLimited deferrals: rate-limit pauses are normal +# while pacing a large batch, but this many consecutive deferrals without a +# single successful submission (~5 minutes at the transfer pause) means the +# server is persistently rejecting us, so the run gives up instead of pacing +# forever. +MAX_CONSECUTIVE_SUBMIT_DEFERRALS = 20 +RATE_LIMIT_SKIP_REASON = "not submitted: the server kept rate-limiting submissions" + TERMINAL_BUILD_STATUSES = {"COMPLETED", "FAILED", "CANCELLED"} FAILURE_TABLE_MAX_ROWS = 20 @@ -509,13 +517,19 @@ def run_bulk_jobs( (remaining specs become SKIPPED) but jobs already in flight are still polled to completion. A submit that raises SubmitRateLimited requeues its spec and pauses new submissions for ``retry_after`` seconds while polling - continues. + continues; after MAX_CONSECUTIVE_SUBMIT_DEFERRALS deferrals with no + successful submission in between, the run stops submitting the same way + the quota path does, so a persistently 429ing server cannot stall it + forever. """ pending: deque[Any] = deque(specs) in_flight: dict[str, _InFlightBuild] = {} outcomes: list[BuildOutcome] = [] - quota_error: Optional[str] = None + # When set, no more submissions happen and the remaining specs are + # SKIPPED with this message (quota exhausted or persistent rate limiting). + stop_skip_reason: Optional[str] = None submit_gate = 0.0 + consecutive_deferrals = 0 rate_limit_notice_shown = False progress = Progress( @@ -542,13 +556,27 @@ def record(outcome: BuildOutcome) -> None: ) while pending or in_flight: - while pending and quota_error is None and len(in_flight) < concurrency: + while pending and stop_skip_reason is None and len(in_flight) < concurrency: if time.monotonic() < submit_gate: break spec = pending.popleft() try: build_id, full_image_path = submit(spec) except SubmitRateLimited as e: + consecutive_deferrals += 1 + if consecutive_deferrals > MAX_CONSECUTIVE_SUBMIT_DEFERRALS: + record( + BuildOutcome( + spec=spec, + status="SUBMIT_FAILED", + error=( + f"{e} (gave up after {MAX_CONSECUTIVE_SUBMIT_DEFERRALS} " + "rate-limit deferrals with no successful submission)" + ), + ) + ) + stop_skip_reason = RATE_LIMIT_SKIP_REASON + break pending.appendleft(spec) submit_gate = time.monotonic() + e.retry_after if not rate_limit_notice_shown: @@ -559,26 +587,28 @@ def record(outcome: BuildOutcome) -> None: ) break except QuotaExceededError as e: - quota_error = str(e) + stop_skip_reason = "not submitted: image quota exceeded" record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) except UnauthorizedError: raise except (APIError, httpx.HTTPError, OSError) as e: + consecutive_deferrals = 0 record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) else: + consecutive_deferrals = 0 in_flight[build_id] = _InFlightBuild( spec=spec, full_image_path=full_image_path, deadline=time.monotonic() + build_timeout, ) - if quota_error is not None and pending: + if stop_skip_reason is not None and pending: for spec in pending: record( BuildOutcome( spec=spec, status="SKIPPED", - error="not submitted: image quota exceeded", + error=stop_skip_reason, ) ) pending.clear() diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 12e35587..2fda427f 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -26,11 +26,13 @@ from ..utils import get_console from .images_bulk import ( + _QUOTA_DETAIL_MARKERS, _TAG_RE, DEFAULT_BUILD_TIMEOUT_SECONDS, DEFAULT_MAX_IN_FLIGHT, FAILURE_TABLE_MAX_ROWS, POLL_INTERVAL_SECONDS, + RATE_LIMIT_SKIP_REASON, SUPPORTED_PLATFORMS, BulkPushValidationError, QuotaExceededError, @@ -586,7 +588,21 @@ def _submit_transfer( raise SubmitRateLimited(str(e), retry_after=TRANSFER_RATE_LIMIT_PAUSE_SECONDS) from e raise - build_id = response.get("build_id") + # Single-source transfers return a top-level build_id today; the endpoint + # also has a per-source results shape (used for comma-separated sources), + # so unwrap it like `prime images push --source-image` does in case the + # contract shifts. + results = response.get("results") + if isinstance(results, list): + entry = results[0] if results and isinstance(results[0], dict) else {} + if not entry.get("success") or not entry.get("buildId"): + error = entry.get("error") or "invalid response from server (transfer not queued)" + if any(marker in error.lower() for marker in _QUOTA_DETAIL_MARKERS): + raise QuotaExceededError(error) + raise APIError(error) + return entry["buildId"], entry.get("fullImagePath") or spec.image_ref + + build_id = response.get("build_id") or response.get("buildId") if not build_id: raise APIError("invalid response from server (missing build_id)") return build_id, response.get("fullImagePath") or spec.image_ref @@ -839,11 +855,20 @@ def transfer_bulk( f"[dim]... and {len(failures) - FAILURE_TABLE_MAX_ROWS} more, " f"all included in {failures_out}[/dim]" ) - if any("limit exceeded" in (o.error or "").lower() for o in failures): + failure_errors = [(o.error or "").lower() for o in failures] + if any(marker in error for error in failure_errors for marker in _QUOTA_DETAIL_MARKERS): console.print( "[red]Image quota reached — delete unused images (prime images delete) " "or request a higher limit, then retry.[/red]" ) + if any( + (o.error or "") == RATE_LIMIT_SKIP_REASON or "rate-limit deferrals" in (o.error or "") + for o in failures + ): + console.print( + "[yellow]The server kept rate-limiting submissions — wait a few minutes, " + "then retry with the failures manifest.[/yellow]" + ) _write_failures_manifest(failures_out, failures) console.print() diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index 833b1c8d..28a0f57a 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -49,6 +49,9 @@ def __init__(self): self.default_poll = ["COMPLETED"] # Exceptions raised (FIFO) by POST /images/build before any transfer is queued. self.build_error_queue = [] + # Respond with the per-source results shape instead of a top-level build_id. + self.respond_bulk_shape = False + self.bulk_entry_error = None def request(self, method, path, json=None, params=None): self.calls.append((method, path)) @@ -56,9 +59,29 @@ def request(self, method, path, json=None, params=None): if self.build_error_queue: raise self.build_error_queue.pop(0) assert "source_image" in json + if self.respond_bulk_shape and self.bulk_entry_error is not None: + entry = { + "sourceImage": json["source_image"], + "success": False, + "error": self.bulk_entry_error, + "retryable": False, + } + return {"results": [entry], "failed": [entry]} self.payloads.append(json) self.build_counter += 1 build_id = f"build-{self.build_counter}" + if self.respond_bulk_shape: + return { + "results": [ + { + "sourceImage": json["source_image"], + "success": True, + "buildId": build_id, + "fullImagePath": f"user/{json.get('image_name') or 'derived'}", + } + ], + "failed": [], + } return { "build_id": build_id, "fullImagePath": f"user/{json.get('image_name') or 'derived'}", @@ -402,6 +425,64 @@ def test_rate_limited_submit_defers_and_retries(tmp_path, fake_api): assert fake_api.post_build_count() == 3 +def test_bulk_shape_response_is_unwrapped(tmp_path, fake_api): + fake_api.respond_bulk_shape = True + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app-a:v1"}, {"source": "b/app-b:v1"}]) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 0, result.output + assert "All images transferred successfully" in result.output + # The build ids from the unwrapped entries are what gets polled. + assert ("GET", "/images/build/build-1") in fake_api.calls + assert ("GET", "/images/build/build-2") in fake_api.calls + + +def test_bulk_shape_failed_entry_records_submit_failed(tmp_path, fake_api): + fake_api.respond_bulk_shape = True + fake_api.bulk_entry_error = "Invalid transfer source: nope" + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app-a:v1"}]) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 1 + assert "0/1 transfers completed" in result.output + assert "Invalid transfer source: nope" in result.output + + +def test_persistent_rate_limit_gives_up_instead_of_pacing_forever(tmp_path, fake_api, monkeypatch): + monkeypatch.setattr(images_bulk, "MAX_CONSECUTIVE_SUBMIT_DEFERRALS", 2) + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app-a:v1"}, {"source": "b/app-b:v1"}]) + fake_api.build_error_queue.extend( + APIError("HTTP 429: Image transfer rate limit exceeded: 10/10 images") for _ in range(10) + ) + failures_out = tmp_path / "failures.jsonl" + result = runner.invoke( + app, + [ + "images", + "transfer-bulk", + "--manifest", + str(manifest), + "--failures-out", + str(failures_out), + ], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "0/2 transfers completed" in result.output + assert "kept rate-limiting" in result.output + # Rate-limit failures must not trigger the quota guidance. + assert "Image quota reached" not in result.output + # cap+1 attempts for the first spec, then the run stops without touching the second. + assert fake_api.post_build_count() == 3 + lines = [json.loads(line) for line in failures_out.read_text().splitlines()] + assert [entry["source"] for entry in lines] == ["a/app-a:v1", "b/app-b:v1"] + + def test_quota_429_aborts_and_skips_remaining(tmp_path, fake_api): manifest = tmp_path / "transfers.jsonl" _write_manifest( From 16c4977749693f7fbf870f8f72485f61016c4b3a Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 2 Jul 2026 21:47:30 -0700 Subject: [PATCH 03/10] fix derive_transfer_destination --- .../commands/images_transfer_bulk.py | 13 ++++- .../prime/tests/test_images_transfer_bulk.py | 55 +++++++++++++++---- 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 2fda427f..6008187d 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -102,6 +102,11 @@ def derive_transfer_destination(source_ref: str) -> tuple[str, str]: ref = (source_ref or "").strip() if not ref: raise ValueError("empty image reference") + # A comma is never valid inside an image reference. Reject it here rather + # than letting the server split it into multiple transfers, which this + # command tracks one-per-spec. + if "," in ref: + raise ValueError("commas are not allowed; use one entry per image reference") digest: Optional[str] = None if "@" in ref: @@ -594,7 +599,13 @@ def _submit_transfer( # contract shifts. results = response.get("results") if isinstance(results, list): - entry = results[0] if results and isinstance(results[0], dict) else {} + # Specs are validated to hold exactly one image reference, so anything + # other than one entry means we would silently drop transfers. + if len(results) != 1 or not isinstance(results[0], dict): + raise APIError( + f"invalid response from server (expected one transfer result, got {len(results)})" + ) + entry = results[0] if not entry.get("success") or not entry.get("buildId"): error = entry.get("error") or "invalid response from server (transfer not queued)" if any(marker in error.lower() for marker in _QUOTA_DETAIL_MARKERS): diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index 28a0f57a..45a37347 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -52,6 +52,7 @@ def __init__(self): # Respond with the per-source results shape instead of a top-level build_id. self.respond_bulk_shape = False self.bulk_entry_error = None + self.bulk_results_count = 1 def request(self, method, path, json=None, params=None): self.calls.append((method, path)) @@ -71,17 +72,13 @@ def request(self, method, path, json=None, params=None): self.build_counter += 1 build_id = f"build-{self.build_counter}" if self.respond_bulk_shape: - return { - "results": [ - { - "sourceImage": json["source_image"], - "success": True, - "buildId": build_id, - "fullImagePath": f"user/{json.get('image_name') or 'derived'}", - } - ], - "failed": [], + entry = { + "sourceImage": json["source_image"], + "success": True, + "buildId": build_id, + "fullImagePath": f"user/{json.get('image_name') or 'derived'}", } + return {"results": [entry] * self.bulk_results_count, "failed": []} return { "build_id": build_id, "fullImagePath": f"user/{json.get('image_name') or 'derived'}", @@ -179,6 +176,9 @@ def test_derive_destination_matches_server_rules(): derive_transfer_destination("") with pytest.raises(ValueError): derive_transfer_destination("app@sha512:abc") + # Comma-separated refs (the push --source-image ad-hoc form) are one-per-entry here. + with pytest.raises(ValueError, match="commas"): + derive_transfer_destination("a/app:v1,b/app:v2") # --------------------------------------------------------------------------- @@ -439,6 +439,41 @@ def test_bulk_shape_response_is_unwrapped(tmp_path, fake_api): assert ("GET", "/images/build/build-2") in fake_api.calls +def test_comma_separated_sources_rejected_in_manifest_and_hf(tmp_path, fake_api, monkeypatch): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app:v1,b/app:v2"}]) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 1 + assert "commas are not allowed" in result.output + assert fake_api.post_build_count() == 0 + + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={0: ["org/img-a:v1,org/img-b:v1"]}, + total=1, + ) + result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + assert result.exit_code == 1 + assert "row 0" in result.output + assert "commas are not allowed" in result.output + assert fake_api.post_build_count() == 0 + + +def test_bulk_shape_multi_entry_response_fails_loudly(tmp_path, fake_api): + fake_api.respond_bulk_shape = True + fake_api.bulk_results_count = 2 + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app-a:v1"}]) + result = runner.invoke( + app, ["images", "transfer-bulk", "--manifest", str(manifest)], env=TEST_ENV + ) + assert result.exit_code == 1 + assert "expected one transfer result, got 2" in result.output + + def test_bulk_shape_failed_entry_records_submit_failed(tmp_path, fake_api): fake_api.respond_bulk_shape = True fake_api.bulk_entry_error = "Invalid transfer source: nope" From dba8685d53f6132dabc01bd8dba5fd0a46bd90e2 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 3 Jul 2026 13:15:21 -0700 Subject: [PATCH 04/10] rm HF_IMAGE_COLUMN_CANDIDATES --- .../src/prime_cli/commands/images_bulk.py | 18 ++--- .../commands/images_transfer_bulk.py | 63 ++++++++------- .../prime/tests/test_images_transfer_bulk.py | 80 ++++++++++++++----- 3 files changed, 99 insertions(+), 62 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_bulk.py b/packages/prime/src/prime_cli/commands/images_bulk.py index daf7697f..175f972e 100644 --- a/packages/prime/src/prime_cli/commands/images_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_bulk.py @@ -47,20 +47,12 @@ RATE_LIMIT_BACKOFF_INITIAL_SECONDS = 2.0 RATE_LIMIT_BACKOFF_MAX_SECONDS = 60.0 -# Liveness bound for SubmitRateLimited deferrals: rate-limit pauses are normal -# while pacing a large batch, but this many consecutive deferrals without a -# single successful submission (~5 minutes at the transfer pause) means the -# server is persistently rejecting us, so the run gives up instead of pacing -# forever. +# Max consecutive deferrals without a single successful submission MAX_CONSECUTIVE_SUBMIT_DEFERRALS = 20 -RATE_LIMIT_SKIP_REASON = "not submitted: the server kept rate-limiting submissions" - -TERMINAL_BUILD_STATUSES = {"COMPLETED", "FAILED", "CANCELLED"} FAILURE_TABLE_MAX_ROWS = 20 _TAG_RE = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9._-]{0,127}$") - _MANIFEST_KEYS = {"image", "context", "dockerfile", "platform"} @@ -525,8 +517,6 @@ def run_bulk_jobs( pending: deque[Any] = deque(specs) in_flight: dict[str, _InFlightBuild] = {} outcomes: list[BuildOutcome] = [] - # When set, no more submissions happen and the remaining specs are - # SKIPPED with this message (quota exhausted or persistent rate limiting). stop_skip_reason: Optional[str] = None submit_gate = 0.0 consecutive_deferrals = 0 @@ -575,7 +565,9 @@ def record(outcome: BuildOutcome) -> None: ), ) ) - stop_skip_reason = RATE_LIMIT_SKIP_REASON + stop_skip_reason = ( + "not submitted: the server kept rate-limiting submissions" + ) break pending.appendleft(spec) submit_gate = time.monotonic() + e.retry_after @@ -632,7 +624,7 @@ def record(outcome: BuildOutcome) -> None: except APIError: build_status = "" # transient poll failure; the deadline still applies - if build_status in TERMINAL_BUILD_STATUSES: + if build_status in {"COMPLETED", "FAILED", "CANCELLED"}: del in_flight[build_id] record( BuildOutcome( diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 6008187d..26a033e7 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -32,7 +32,6 @@ DEFAULT_MAX_IN_FLIGHT, FAILURE_TABLE_MAX_ROWS, POLL_INTERVAL_SECONDS, - RATE_LIMIT_SKIP_REASON, SUPPORTED_PLATFORMS, BulkPushValidationError, QuotaExceededError, @@ -53,7 +52,6 @@ HF_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" HF_ROWS_PAGE_LENGTH = 100 -HF_IMAGE_COLUMN_CANDIDATES = ("docker_image", "image_name", "container_image", "image") HF_REQUEST_TIMEOUT_SECONDS = 30.0 HF_REQUEST_MAX_ATTEMPTS = 6 HF_REQUEST_BACKOFF_SECONDS = 2.0 @@ -93,18 +91,11 @@ def to_manifest_line(self) -> dict[str, str]: def derive_transfer_destination(source_ref: str) -> tuple[str, str]: - """Derive the (name, tag) a transfer gets when no destination is given. - - Mirrors the server's parsing (last repository segment lowercased; tag kept, - or ``sha256-`` for digest refs). Used client-side for dry - runs and duplicate-destination detection; the server stays authoritative. - """ + """Derive the (name, tag) a transfer gets when no destination is given.""" ref = (source_ref or "").strip() if not ref: raise ValueError("empty image reference") - # A comma is never valid inside an image reference. Reject it here rather - # than letting the server split it into multiple transfers, which this - # command tracks one-per-spec. + # A comma is never valid inside an image reference. if "," in ref: raise ValueError("commas are not allowed; use one entry per image reference") @@ -449,24 +440,33 @@ def load_hf_specs( features = config_info.get("features") or {} if column is None: - column = next((c for c in HF_IMAGE_COLUMN_CANDIDATES if c in features), None) - if column is None: - raise BulkPushValidationError( - [ - f"could not auto-detect an image column in '{dataset_id}' " - f"(looked for: {', '.join(HF_IMAGE_COLUMN_CANDIDATES)}; " - f"dataset columns: {', '.join(sorted(features)) or 'unknown'}); " - "pass --column" - ] - ) - notes.append(f"Using column '{column}' (auto-detected)") - elif features and column not in features: + raise BulkPushValidationError( + [ + "--column is required with --hf (columns in " + f"'{dataset_id}': {', '.join(sorted(features)) or 'unknown'})" + ] + ) + if features and column not in features: raise BulkPushValidationError( [ f"column '{column}' not found in dataset '{dataset_id}' " f"(columns: {', '.join(sorted(features))})" ] ) + # Reject columns whose feature type cannot hold image reference strings + # (e.g. actual pictures, nested lists) before submitting anything. + feature_spec = features.get(column) if isinstance(features, dict) else None + if feature_spec is not None and not ( + isinstance(feature_spec, dict) + and feature_spec.get("_type") == "Value" + and feature_spec.get("dtype") in ("string", "large_string") + ): + raise BulkPushValidationError( + [ + f"column '{column}' in '{dataset_id}' is not a string column, " + "so it cannot hold image references" + ] + ) splits = config_info.get("splits") if isinstance(splits, dict) and splits and split not in splits: @@ -660,8 +660,9 @@ def transfer_bulk( column: Optional[str] = typer.Option( None, "--column", - help="Dataset column holding image references for --hf mode", - show_default=f"auto-detect: {', '.join(HF_IMAGE_COLUMN_CANDIDATES)}", + help=( + "Dataset column holding image references (required for --hf mode; e.g. 'docker_image')" + ), ), platform: str = typer.Option( "linux/amd64", @@ -712,15 +713,15 @@ def transfer_bulk( tasks push-bulk skips, so the two commands together cover a task set. \b - Hugging Face mode reads image references straight from a dataset column - (auto-detected, e.g. docker_image) using the dataset viewer API — no local - dataset download. Set HF_TOKEN for private or gated datasets. + Hugging Face mode reads image references straight from the dataset column + named by --column (e.g. docker_image) using the dataset viewer API — no + local dataset download. Set HF_TOKEN for private or gated datasets. \b Examples: prime images transfer-bulk --manifest transfers.jsonl prime images transfer-bulk --harbor ./tasks - prime images transfer-bulk --hf R2E-Gym/R2E-Gym-Subset --dry-run + prime images transfer-bulk --hf R2E-Gym/R2E-Gym-Subset --column docker_image --dry-run prime images transfer-bulk --hf org/dataset --hf-split test --column image_name prime images transfer-bulk --manifest transfer-bulk-failures.jsonl """ @@ -873,7 +874,9 @@ def transfer_bulk( "or request a higher limit, then retry.[/red]" ) if any( - (o.error or "") == RATE_LIMIT_SKIP_REASON or "rate-limit deferrals" in (o.error or "") + # Matches the engine's skip reason and its give-up SUBMIT_FAILED error. + "kept rate-limiting submissions" in (o.error or "") + or "rate-limit deferrals" in (o.error or "") for o in failures ): console.print( diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index 45a37347..c48c307b 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -151,7 +151,11 @@ def fake_get(url, params=None, headers=None, timeout=None): HF_INFO_ONE_CONFIG = { "dataset_info": { "default": { - "features": {"docker_image": {"dtype": "string"}, "prompt": {"dtype": "string"}}, + "features": { + "docker_image": {"dtype": "string", "_type": "Value"}, + "prompt": {"dtype": "string", "_type": "Value"}, + "picture": {"_type": "Image"}, + }, "splits": {"train": {"name": "train"}}, } }, @@ -295,7 +299,7 @@ def test_harbor_with_no_transferable_tasks_fails(tmp_path, fake_api): # --------------------------------------------------------------------------- -def test_hf_pages_rows_dedupes_and_autodetects_column(tmp_path, fake_api, monkeypatch): +def test_hf_pages_rows_and_dedupes(tmp_path, fake_api, monkeypatch): hf_calls = _fake_hf( monkeypatch, info=HF_INFO_ONE_CONFIG, @@ -305,9 +309,12 @@ def test_hf_pages_rows_dedupes_and_autodetects_column(tmp_path, fake_api, monkey }, total=5, ) - result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "Org/DataSet"], env=TEST_ENV) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "Org/DataSet", "--column", "docker_image"], + env=TEST_ENV, + ) assert result.exit_code == 0, result.output - assert "Using column 'docker_image' (auto-detected)" in result.output assert "Collapsed 1 duplicate image reference(s)" in result.output assert "Skipped 1 row(s) with an empty 'docker_image' value" in result.output assert fake_api.post_build_count() == 3 @@ -327,7 +334,11 @@ def test_hf_rate_limited_request_retries(tmp_path, fake_api, monkeypatch): total=1, rate_limit_first_rows=True, ) - result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--column", "docker_image"], + env=TEST_ENV, + ) assert result.exit_code == 0, result.output assert fake_api.post_build_count() == 1 # The rate-limited /rows request is retried: info + 429 + 200. @@ -343,6 +354,8 @@ def test_hf_accepts_dataset_url(tmp_path, fake_api, monkeypatch): "transfer-bulk", "--hf", "https://huggingface.co/datasets/R2E-Gym/R2E-Gym-Subset/viewer/default/train", + "--column", + "docker_image", ], env=TEST_ENV, ) @@ -364,7 +377,16 @@ def test_hf_multiple_configs_requires_flag(tmp_path, fake_api, monkeypatch): result = runner.invoke( app, - ["images", "transfer-bulk", "--hf", "org/ds", "--hf-config", "cfg-a"], + [ + "images", + "transfer-bulk", + "--hf", + "org/ds", + "--hf-config", + "cfg-a", + "--column", + "docker_image", + ], env=TEST_ENV, ) assert result.exit_code == 0, result.output @@ -382,26 +404,42 @@ def test_hf_unknown_column_and_split_fail(tmp_path, fake_api, monkeypatch): result = runner.invoke( app, - ["images", "transfer-bulk", "--hf", "org/ds", "--hf-split", "test"], + [ + "images", + "transfer-bulk", + "--hf", + "org/ds", + "--column", + "docker_image", + "--hf-split", + "test", + ], env=TEST_ENV, ) assert result.exit_code == 1 assert "split 'test' not found" in result.output -def test_hf_no_detectable_column_asks_for_flag(tmp_path, fake_api, monkeypatch): - info = { - "dataset_info": { - "default": { - "features": {"prompt": {"dtype": "string"}}, - "splits": {"train": {"name": "train"}}, - } - } - } - _fake_hf(monkeypatch, info=info, pages={0: []}, total=0) +def test_hf_column_is_required(tmp_path, fake_api, monkeypatch): + _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) assert result.exit_code == 1 - assert "pass --column" in result.output + assert "--column is required" in result.output + # The error lists the dataset's columns so the user can pick one. + assert "docker_image" in result.output + assert fake_api.post_build_count() == 0 + + +def test_hf_non_string_column_rejected(tmp_path, fake_api, monkeypatch): + _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--column", "picture"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "not a string column" in result.output + assert fake_api.post_build_count() == 0 # --------------------------------------------------------------------------- @@ -455,7 +493,11 @@ def test_comma_separated_sources_rejected_in_manifest_and_hf(tmp_path, fake_api, pages={0: ["org/img-a:v1,org/img-b:v1"]}, total=1, ) - result = runner.invoke(app, ["images", "transfer-bulk", "--hf", "org/ds"], env=TEST_ENV) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--column", "docker_image"], + env=TEST_ENV, + ) assert result.exit_code == 1 assert "row 0" in result.output assert "commas are not allowed" in result.output From 31db1618f99fe4fa29ea25ebade41a6edba0a31b Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 3 Jul 2026 13:32:39 -0700 Subject: [PATCH 05/10] org/dataset --- packages/prime/src/prime_cli/commands/images_transfer_bulk.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 26a033e7..a2587dbc 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -646,7 +646,7 @@ def transfer_bulk( None, "--hf", help=( - "Hugging Face dataset id or URL (e.g. 'R2E-Gym/R2E-Gym-Subset'); " + "Hugging Face dataset id or URL (e.g. 'org/dataset'); " "transfers every image referenced in the dataset" ), ), @@ -721,7 +721,7 @@ def transfer_bulk( Examples: prime images transfer-bulk --manifest transfers.jsonl prime images transfer-bulk --harbor ./tasks - prime images transfer-bulk --hf R2E-Gym/R2E-Gym-Subset --column docker_image --dry-run + prime images transfer-bulk --hf org/dataset --column docker_image --dry-run prime images transfer-bulk --hf org/dataset --hf-split test --column image_name prime images transfer-bulk --manifest transfer-bulk-failures.jsonl """ From a162e16bdc044230574cb798a8c9cf936952916e Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 3 Jul 2026 13:47:50 -0700 Subject: [PATCH 06/10] progress bar --- .../commands/images_transfer_bulk.py | 86 ++++++++++++------- .../prime/tests/test_images_transfer_bulk.py | 13 +++ 2 files changed, 69 insertions(+), 30 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index a2587dbc..8e20db46 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -22,6 +22,7 @@ ImageVisibility, UnauthorizedError, ) +from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn from rich.table import Table from ..utils import get_console @@ -52,6 +53,7 @@ HF_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" HF_ROWS_PAGE_LENGTH = 100 +LARGE_DATASET_WARN_BYTES = 500_000_000 HF_REQUEST_TIMEOUT_SECONDS = 30.0 HF_REQUEST_MAX_ATTEMPTS = 6 HF_REQUEST_BACKOFF_SECONDS = 2.0 @@ -380,6 +382,10 @@ def _hf_get(path: str, params: dict[str, Any]) -> dict[str, Any]: retry_delay = max(retry_delay, float(retry_after)) except ValueError: pass + console.print( + f"[dim]Hugging Face returned HTTP {response.status_code}; " + f"retrying in {int(retry_delay)}s...[/dim]" + ) time.sleep(retry_delay) continue if response.status_code != 200: @@ -481,42 +487,62 @@ def load_hf_specs( "Warning: the dataset viewer only covers part of this dataset; some rows may be missing" ) + # The rows API returns full rows (a columns filter is not honored), so + # reading time scales with total dataset size, not just row count. + dataset_size = config_info.get("dataset_size") + if isinstance(dataset_size, (int, float)) and dataset_size > LARGE_DATASET_WARN_BYTES: + console.print( + f"[dim]Dataset rows total ~{dataset_size / 1e9:.1f}GB and the viewer API " + "streams full rows, so reading may take a few minutes.[/dim]" + ) + sources: list[tuple[str, int]] = [] # (image ref, first row index), order-preserving seen: set[str] = set() scanned = 0 empty_rows = 0 offset = 0 total: Optional[int] = None - while True: - page = _hf_get( - "/rows", - { - "dataset": dataset_id, - "config": config, - "split": split, - "offset": offset, - "length": HF_ROWS_PAGE_LENGTH, - }, - ) - rows = page.get("rows") or [] - if isinstance(page.get("num_rows_total"), int): - total = page["num_rows_total"] - for item in rows: - row = item.get("row") or {} - row_idx = item.get("row_idx", offset) - value = row.get(column) - if not isinstance(value, str) or not value.strip(): - empty_rows += 1 - continue - scanned += 1 - value = value.strip() - if value in seen: - continue - seen.add(value) - sources.append((value, row_idx)) - offset += len(rows) - if not rows or (total is not None and offset >= total): - break + progress = Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TextColumn("{task.completed}/{task.total}"), + TimeElapsedColumn(), + console=console, + ) + with progress: + task_id = progress.add_task("Reading dataset rows", total=None) + while True: + page = _hf_get( + "/rows", + { + "dataset": dataset_id, + "config": config, + "split": split, + "offset": offset, + "length": HF_ROWS_PAGE_LENGTH, + }, + ) + rows = page.get("rows") or [] + if isinstance(page.get("num_rows_total"), int): + total = page["num_rows_total"] + progress.update(task_id, total=total) + for item in rows: + row = item.get("row") or {} + row_idx = item.get("row_idx", offset) + value = row.get(column) + if not isinstance(value, str) or not value.strip(): + empty_rows += 1 + continue + scanned += 1 + value = value.strip() + if value in seen: + continue + seen.add(value) + sources.append((value, row_idx)) + offset += len(rows) + progress.update(task_id, completed=offset) + if not rows or (total is not None and offset >= total): + break if empty_rows: notes.append(f"Skipped {empty_rows} row(s) with an empty '{column}' value") diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index c48c307b..ce34f482 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -430,6 +430,19 @@ def test_hf_column_is_required(tmp_path, fake_api, monkeypatch): assert fake_api.post_build_count() == 0 +def test_hf_large_dataset_prints_size_note(tmp_path, fake_api, monkeypatch): + info = json.loads(json.dumps(HF_INFO_ONE_CONFIG)) + info["dataset_info"]["default"]["dataset_size"] = 3_700_000_000 + _fake_hf(monkeypatch, info=info, pages={0: ["org/img-a:v1"]}, total=1) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--hf", "org/ds", "--column", "docker_image"], + env=TEST_ENV, + ) + assert result.exit_code == 0, result.output + assert "may take a few minutes" in result.output + + def test_hf_non_string_column_rejected(tmp_path, fake_api, monkeypatch): _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) result = runner.invoke( From 4ceac8fe2b6cd26c93adf78ad8f67c9a52fd57f6 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Mon, 6 Jul 2026 14:28:01 -0700 Subject: [PATCH 07/10] bump HF_REQUEST_MAX_ATTEMPTS --- packages/prime/src/prime_cli/commands/images_transfer_bulk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 8e20db46..d1d5473b 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -55,7 +55,7 @@ HF_ROWS_PAGE_LENGTH = 100 LARGE_DATASET_WARN_BYTES = 500_000_000 HF_REQUEST_TIMEOUT_SECONDS = 30.0 -HF_REQUEST_MAX_ATTEMPTS = 6 +HF_REQUEST_MAX_ATTEMPTS = 9 HF_REQUEST_BACKOFF_SECONDS = 2.0 HF_REQUEST_BACKOFF_MAX_SECONDS = 60.0 From 737bd4043c09bcb318c49756bdc8029ab5b1d148 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 7 Jul 2026 11:56:42 -0700 Subject: [PATCH 08/10] fix team context for platform images --- .../prime/src/prime_cli/commands/images.py | 16 +++--- .../commands/images_transfer_bulk.py | 42 +++++++++++++-- packages/prime/tests/test_images_push.py | 20 +++++-- .../prime/tests/test_images_transfer_bulk.py | 52 +++++++++++++++++++ 4 files changed, 114 insertions(+), 16 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images.py b/packages/prime/src/prime_cli/commands/images.py index aa23469a..0c60a099 100644 --- a/packages/prime/src/prime_cli/commands/images.py +++ b/packages/prime/src/prime_cli/commands/images.py @@ -429,9 +429,6 @@ def push_image( if platform_image and private: console.print("[red]Error: Platform images must be public[/red]") raise typer.Exit(1) - if platform_image and config.team_id: - console.print("[red]Error: Platform images cannot be pushed in a team context[/red]") - raise typer.Exit(1) if not is_transfer and image_reference is None: console.print( "[red]Error: Image reference is required unless --source-image is used[/red]" @@ -484,7 +481,9 @@ def push_image( console.print(f"[bold]Destination:[/bold] {destination_display}") if platform_image: console.print("[bold]Owner:[/bold] Platform") - if config.team_id: + if config.team_id: + console.print("[dim]Team context ignored: platform images are org-less[/dim]") + elif config.team_id: console.print(f"[dim]Team: {config.team_id}[/dim]") console.print() @@ -503,7 +502,7 @@ def push_image( image_name=image_name, image_tag=image_tag, platform=platform, - team_id=config.team_id or None, + team_id=None if platform_image else (config.team_id or None), visibility=visibility, owner_scope="platform" if platform_image else None, ) @@ -582,7 +581,10 @@ def push_image( console.print( f"[bold blue]Building and pushing image:[/bold blue] {image_name}:{image_tag}" ) - if config.team_id: + if platform_image: + if config.team_id: + console.print("[dim]Team context ignored: platform images are org-less[/dim]") + elif config.team_id: console.print(f"[dim]Team: {config.team_id}[/dim]") console.print() @@ -626,7 +628,7 @@ def push_image( "dockerfile_path": PACKAGED_DOCKERFILE_PATH, "platform": platform, } - if config.team_id: + if config.team_id and not platform_image: build_payload["team_id"] = config.team_id if platform_image: build_payload["owner_scope"] = "platform" diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index d1d5473b..207e8aec 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -588,6 +588,7 @@ def _submit_transfer( *, team_id: Optional[str], visibility: Optional[ImageVisibility], + owner_scope: Optional[str] = None, ) -> tuple[str, str]: """Queue one transfer. Returns (build_id, full_image_path). @@ -607,6 +608,8 @@ def _submit_transfer( payload["team_id"] = team_id if visibility is not None: payload["visibility"] = visibility.value + if owner_scope is not None: + payload["owner_scope"] = owner_scope try: response = client.request("POST", "/images/build", json=payload) @@ -702,6 +705,11 @@ def transfer_bulk( private: bool = typer.Option( False, "--private", help="Make the images private when the transfers complete" ), + platform_image: bool = typer.Option( + False, + "--platform-image", + help="Transfer org-less platform images (admins only; implies --public)", + ), concurrency: int = typer.Option( DEFAULT_MAX_IN_FLIGHT, "--concurrency", help="Maximum transfers in flight at once" ), @@ -743,12 +751,17 @@ def transfer_bulk( named by --column (e.g. docker_image) using the dataset viewer API — no local dataset download. Set HF_TOKEN for private or gated datasets. + \b + Admins can pass --platform-image to transfer org-less public platform + images, exactly like 'prime images push --source-image --platform-image'. + \b Examples: prime images transfer-bulk --manifest transfers.jsonl prime images transfer-bulk --harbor ./tasks prime images transfer-bulk --hf org/dataset --column docker_image --dry-run prime images transfer-bulk --hf org/dataset --hf-split test --column image_name + prime images transfer-bulk --hf org/dataset --column docker_image --platform-image prime images transfer-bulk --manifest transfer-bulk-failures.jsonl """ try: @@ -764,6 +777,9 @@ def transfer_bulk( if public and private: console.print("[red]Error: --public and --private cannot be used together[/red]") raise typer.Exit(1) + if platform_image and private: + console.print("[red]Error: Platform images must be public[/red]") + raise typer.Exit(1) if concurrency < 1: console.print("[red]Error: --concurrency must be at least 1[/red]") raise typer.Exit(1) @@ -830,12 +846,18 @@ def transfer_bulk( visibility = ImageVisibility.PUBLIC elif private: visibility = ImageVisibility.PRIVATE + if platform_image: + visibility = ImageVisibility.PUBLIC console.print( f"[bold blue]Bulk transferring {len(specs)} image(s)[/bold blue] " f"[dim]({source_desc})[/dim]" ) - if config.team_id: + if platform_image: + console.print("[dim]Owner: Platform[/dim]") + if config.team_id: + console.print("[dim]Team context ignored: platform images are org-less[/dim]") + elif config.team_id: console.print(f"[dim]Team: {config.team_id}[/dim]") if visibility is not None: console.print(f"[dim]Visibility: {visibility.value}[/dim]") @@ -844,10 +866,18 @@ def transfer_bulk( "[dim]Visibility: PRIVATE for new images " "(existing tags keep their current visibility)[/dim]" ) + # Platform transfers skip the server's per-account transfer rate limiter. + pacing_note = ( + "" + if platform_image + else ( + " Transfers are rate-limited per account server-side, " + "so large batches take a while to submit." + ) + ) console.print( f"[dim]Up to {concurrency} transfers in flight; " - f"polling every {int(POLL_INTERVAL_SECONDS)}s. Transfers are rate-limited " - "per account server-side, so large batches take a while to submit.[/dim]" + f"polling every {int(POLL_INTERVAL_SECONDS)}s.{pacing_note}[/dim]" ) console.print() @@ -856,7 +886,11 @@ def transfer_bulk( client, specs, submit=lambda spec: _submit_transfer( - client, spec, team_id=config.team_id or None, visibility=visibility + client, + spec, + team_id=None if platform_image else (config.team_id or None), + visibility=visibility, + owner_scope="platform" if platform_image else None, ), concurrency=concurrency, build_timeout=build_timeout, diff --git a/packages/prime/tests/test_images_push.py b/packages/prime/tests/test_images_push.py index 9618e6ae..a28d5a40 100644 --- a/packages/prime/tests/test_images_push.py +++ b/packages/prime/tests/test_images_push.py @@ -301,23 +301,33 @@ def request(self, method, path, json=None, params=None): assert "Platform images must be public" in result.output -def test_push_platform_image_rejects_team_context(monkeypatch): +def test_push_platform_image_ignores_team_context(monkeypatch): monkeypatch.setattr("prime_cli.main.check_for_update", lambda: (False, None)) + captured = {} class DummyAPIClient: def request(self, method, path, json=None, params=None): - raise AssertionError(f"Unexpected request: {method} {path}") + captured["json"] = json + return { + "build_id": "build-123", + "buildIds": ["build-123"], + "upload_url": None, + "fullImagePath": "ubuntu:22.04", + "visibility": "PUBLIC", + } monkeypatch.setattr("prime_cli.commands.images.APIClient", DummyAPIClient) result = runner.invoke( app, - ["images", "push", "ubuntu:22.04", "--platform-image"], + ["images", "push", "--source-image", "ubuntu:22.04", "--platform-image"], env={**TEST_ENV, "PRIME_TEAM_ID": "team-123"}, ) - assert result.exit_code == 1 - assert "Platform images cannot be pushed in a team context" in result.output + assert result.exit_code == 0, result.output + assert "Team context ignored" in result.output + assert captured["json"]["owner_scope"] == "platform" + assert "team_id" not in captured["json"] def test_push_image_source_image_queues_transfer_without_upload(monkeypatch): diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index ce34f482..9087f02f 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -661,3 +661,55 @@ def test_hf_flags_rejected_outside_hf_mode(tmp_path, fake_api): ) assert result.exit_code == 1 assert "only apply to --hf mode" in result.output + + +# --------------------------------------------------------------------------- +# Platform images +# --------------------------------------------------------------------------- + + +def test_platform_image_sends_owner_scope_and_public_visibility(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest( + manifest, + [{"source": "docker.io/org/app:v1"}, {"source": "ghcr.io/org/other:v2"}], + ) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--platform-image"], + env=TEST_ENV, + ) + assert result.exit_code == 0, result.output + assert "Owner: Platform" in result.output + assert "Visibility: PUBLIC" in result.output + for payload in fake_api.payloads: + assert payload["owner_scope"] == "platform" + assert payload["visibility"] == "PUBLIC" + assert "team_id" not in payload + + +def test_platform_image_rejects_private(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app:v1"}]) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--platform-image", "--private"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "Platform images must be public" in result.output + assert fake_api.post_build_count() == 0 + + +def test_platform_image_ignores_team_context(tmp_path, fake_api): + manifest = tmp_path / "transfers.jsonl" + _write_manifest(manifest, [{"source": "a/app:v1"}]) + result = runner.invoke( + app, + ["images", "transfer-bulk", "--manifest", str(manifest), "--platform-image"], + env={**TEST_ENV, "PRIME_TEAM_ID": "team-123"}, + ) + assert result.exit_code == 0, result.output + assert "Team context ignored" in result.output + assert fake_api.payloads[0]["owner_scope"] == "platform" + assert "team_id" not in fake_api.payloads[0] From 4fc63bcea72f92eab22a2a9fca650f9866c4d60b Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 7 Jul 2026 13:44:08 -0700 Subject: [PATCH 09/10] hf support for prime images push-bulk --- .../src/prime_cli/commands/images_bulk.py | 114 +++++- .../prime/src/prime_cli/commands/images_hf.py | 366 ++++++++++++++++++ .../commands/images_transfer_bulk.py | 249 ++---------- packages/prime/tests/test_images_push_bulk.py | 288 +++++++++++++- .../prime/tests/test_images_transfer_bulk.py | 5 +- 5 files changed, 785 insertions(+), 237 deletions(-) create mode 100644 packages/prime/src/prime_cli/commands/images_hf.py diff --git a/packages/prime/src/prime_cli/commands/images_bulk.py b/packages/prime/src/prime_cli/commands/images_bulk.py index 175f972e..6312e45b 100644 --- a/packages/prime/src/prime_cli/commands/images_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_bulk.py @@ -1,5 +1,6 @@ import json import re +import shutil import sys import tarfile import tempfile @@ -705,14 +706,44 @@ def push_bulk( "environment/ folder is the build context" ), ), - tag: str = typer.Option("latest", "--tag", help="Image tag for Harbor mode"), + hf_dataset: Optional[str] = typer.Option( + None, + "--hf", + help=( + "Hugging Face dataset id or URL (e.g. 'org/dataset'); " + "builds every Dockerfile stored in the dataset" + ), + ), + hf_split: str = typer.Option("train", "--hf-split", help="Dataset split for --hf mode"), + hf_config: Optional[str] = typer.Option( + None, + "--hf-config", + help="Dataset config (called 'subset' in the HF viewer) for --hf mode", + show_default="the dataset's only config", + ), + dockerfile_column: Optional[str] = typer.Option( + None, + "--dockerfile-column", + help=( + "Dataset column holding Dockerfile contents (required for --hf mode; e.g. 'dockerfile')" + ), + ), + name_column: Optional[str] = typer.Option( + None, + "--name-column", + help=( + "Dataset column naming each image (required for --hf mode; " + "e.g. 'instance_id'; values are sanitized to valid image names)" + ), + ), + tag: str = typer.Option("latest", "--tag", help="Image tag for Harbor and --hf modes"), name_template: str = typer.Option( "{dir}", "--name-template", help=( - "Image name template for Harbor mode; placeholders: {dir} (task directory " - "name) and {name} (task.toml [task].name). The result is sanitized to a " - "valid image name" + "Image name template for Harbor and --hf modes; placeholders: {dir} (task " + "directory name) and {name} (task.toml [task].name; in --hf mode both are " + "the --name-column value). The result is sanitized to a valid image name" ), ), platform: str = typer.Option( @@ -747,10 +778,11 @@ def push_bulk( """ Build and push many Docker images in one command. - Reads builds from a JSONL manifest (--manifest) or a Harbor tasks - directory (--harbor), validates everything up front, then keeps up to - --concurrency builds running server-side, starting the next build as soon - as one finishes. Failed builds are written to a manifest you can re-run. + Reads builds from a JSONL manifest (--manifest), a Harbor tasks directory + (--harbor), or a Hugging Face dataset of Dockerfiles (--hf), validates + everything up front, then keeps up to --concurrency builds running + server-side, starting the next build as soon as one finishes. Failed + builds are written to a manifest you can re-run. \b Manifest format (one JSON object per line; paths relative to the manifest): @@ -764,16 +796,37 @@ def push_bulk( with --tag. Tasks with a prebuilt docker_image or no environment/Dockerfile are skipped. + \b + Hugging Face mode builds one image per dataset row using the dataset + viewer API — no local dataset download. --dockerfile-column holds each + row's Dockerfile contents and --name-column names the image (sanitized, + tagged with --tag). Set HF_TOKEN for private or gated datasets. Datasets + whose rows reference prebuilt registry images are transfer-bulk's job: + prime images transfer-bulk --hf. + \b Examples: prime images push-bulk --manifest builds.jsonl prime images push-bulk --harbor ./tasks --tag v1 prime images push-bulk --harbor ./tasks --name-template "swe-{dir}" --dry-run + prime images push-bulk --hf org/dataset --dockerfile-column dockerfile \\ + --name-column instance_id prime images push-bulk --manifest push-bulk-failures.jsonl """ + hf_context_root: Optional[Path] = None + keep_hf_context = False try: - if (manifest is None) == (harbor is None): - console.print("[red]Error: Provide exactly one of --manifest or --harbor[/red]") + modes = [m for m in (manifest, harbor, hf_dataset) if m is not None] + if len(modes) != 1: + console.print("[red]Error: Provide exactly one of --manifest, --harbor or --hf[/red]") + raise typer.Exit(1) + if hf_dataset is None and ( + hf_split != "train" or hf_config or dockerfile_column or name_column + ): + console.print( + "[red]Error: --hf-split, --hf-config, --dockerfile-column and " + "--name-column only apply to --hf mode[/red]" + ) raise typer.Exit(1) if public and private: console.print("[red]Error: --public and --private cannot be used together[/red]") @@ -785,10 +838,13 @@ def push_bulk( console.print("[red]Error: --build-timeout must be at least 1[/red]") raise typer.Exit(1) if manifest is not None and (tag != "latest" or name_template != "{dir}"): - console.print("[red]Error: --tag and --name-template only apply to --harbor mode[/red]") + console.print( + "[red]Error: --tag and --name-template only apply to --harbor and --hf modes[/red]" + ) raise typer.Exit(1) skipped: list[tuple[str, str]] = [] + notes: list[str] = [] try: if manifest is not None: manifest_path = manifest.resolve() @@ -796,8 +852,7 @@ def push_bulk( raise BulkPushValidationError([f"manifest not found: {manifest_path}"]) source_desc = f"manifest {manifest_path}" specs = load_manifest(manifest_path, default_platform=platform) - else: - assert harbor is not None + elif harbor is not None: harbor_root = harbor.resolve() if not harbor_root.is_dir(): raise BulkPushValidationError( @@ -809,6 +864,27 @@ def push_bulk( specs, skipped = load_harbor_specs( harbor_root, tag=tag, name_template=name_template, platform=platform ) + else: + assert hf_dataset is not None + if not _TAG_RE.match(tag): + raise BulkPushValidationError([f"invalid image tag '{tag}'"]) + # Imported here: images_hf imports BuildSpec helpers from this module. + from .images_hf import load_hf_build_specs, normalize_hf_dataset_id + + source_desc = f"Hugging Face dataset {normalize_hf_dataset_id(hf_dataset)}" + console.print(f"[cyan]Reading Dockerfiles from {source_desc}...[/cyan]") + hf_context_root = Path(tempfile.mkdtemp(prefix="prime-push-bulk-hf-")) + specs, notes = load_hf_build_specs( + hf_dataset, + config=hf_config, + split=hf_split, + dockerfile_column=dockerfile_column, + name_column=name_column, + tag=tag, + name_template=name_template, + platform=platform, + context_root=hf_context_root, + ) except BulkPushValidationError as e: console.print( f"[red]Error: cannot start bulk push ({len(e.problems)} problem(s)):[/red]" @@ -817,6 +893,8 @@ def push_bulk( console.print(f"[red] - {problem}[/red]") raise typer.Exit(1) + for note in notes: + console.print(f"[dim]{note}[/dim]") for task_name, reason in skipped: console.print(f"[yellow]Skipping {task_name}: {reason}[/yellow]") @@ -909,6 +987,13 @@ def push_bulk( console.print() console.print(f"Wrote {len(failures)} failed build(s) to [bold]{failures_out}[/bold]") console.print(f"[dim]Retry with: prime images push-bulk --manifest {failures_out}[/dim]") + if hf_context_root is not None: + # The failures manifest points into the generated build contexts. + keep_hf_context = True + console.print( + f"[dim]Keeping generated Dockerfiles in {hf_context_root} " + "so the failures manifest can be re-run.[/dim]" + ) raise typer.Exit(1) except UnauthorizedError: @@ -920,3 +1005,6 @@ def push_bulk( "check them with 'prime images list'.[/yellow]" ) raise typer.Exit(1) + finally: + if hf_context_root is not None and not keep_hf_context: + shutil.rmtree(hf_context_root, ignore_errors=True) diff --git a/packages/prime/src/prime_cli/commands/images_hf.py b/packages/prime/src/prime_cli/commands/images_hf.py new file mode 100644 index 00000000..916f3d36 --- /dev/null +++ b/packages/prime/src/prime_cli/commands/images_hf.py @@ -0,0 +1,366 @@ +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterator, Optional + +import httpx +from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn + +from ..utils import get_console +from .images_bulk import ( + BuildSpec, + BulkPushValidationError, + _duplicate_ref_problems, + _render_name_template, +) + +console = get_console() + +HF_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" +HF_ROWS_PAGE_LENGTH = 100 +LARGE_DATASET_WARN_BYTES = 500_000_000 +HF_REQUEST_TIMEOUT_SECONDS = 30.0 +HF_REQUEST_MAX_ATTEMPTS = 9 +HF_REQUEST_BACKOFF_SECONDS = 2.0 +HF_REQUEST_BACKOFF_MAX_SECONDS = 60.0 + +PARTIAL_DATASET_NOTE = ( + "Warning: the dataset viewer only covers part of this dataset; some rows may be missing" +) + + +def normalize_hf_dataset_id(raw: str) -> str: + """Accept 'org/name', hf://datasets/... and huggingface.co dataset URLs.""" + dataset = (raw or "").strip() + for prefix in ( + "hf://datasets/", + "https://huggingface.co/datasets/", + "http://huggingface.co/datasets/", + "huggingface.co/datasets/", + ): + if dataset.startswith(prefix): + dataset = dataset[len(prefix) :] + dataset = dataset.split("?", 1)[0].split("#", 1)[0] + # Drop trailing URL parts like /viewer/default/train. + parts = [part for part in dataset.split("/") if part] + dataset = "/".join(parts[:2]) + break + dataset = dataset.strip("/") + if not dataset or dataset.count("/") > 1: + raise BulkPushValidationError( + [f"invalid Hugging Face dataset '{raw}' (expected 'org/name' or a dataset URL)"] + ) + return dataset + + +def hf_get(path: str, params: dict[str, Any]) -> dict[str, Any]: + """GET from the HF datasets-server with retries on transient failures. + + Uses HF_TOKEN from the environment for gated/private datasets. The + datasets-server rate limit is per IP and clears within about a minute, so + 429/5xx retries back off long enough to ride out a full window, honoring + Retry-After when the server sends one. + """ + token = os.environ.get("HF_TOKEN") + headers = {"Authorization": f"Bearer {token}"} if token else {} + url = f"{HF_DATASETS_SERVER_URL}{path}" + delay = HF_REQUEST_BACKOFF_SECONDS + for attempt in range(1, HF_REQUEST_MAX_ATTEMPTS + 1): + retry_delay = delay + delay = min(delay * 2, HF_REQUEST_BACKOFF_MAX_SECONDS) + try: + response = httpx.get( + url, params=params, headers=headers, timeout=HF_REQUEST_TIMEOUT_SECONDS + ) + except httpx.HTTPError as e: + if attempt == HF_REQUEST_MAX_ATTEMPTS: + raise BulkPushValidationError([f"Hugging Face request failed: {e}"]) + time.sleep(retry_delay) + continue + if response.status_code == 429 or response.status_code >= 500: + if attempt == HF_REQUEST_MAX_ATTEMPTS: + rate_limit_hint = ( + " (rate limited — wait a minute and re-run)" + if response.status_code == 429 + else "" + ) + raise BulkPushValidationError( + [ + f"Hugging Face request failed with HTTP " + f"{response.status_code}: {url}{rate_limit_hint}" + ] + ) + retry_after = response.headers.get("retry-after") + if retry_after: + try: + retry_delay = max(retry_delay, float(retry_after)) + except ValueError: + pass + console.print( + f"[dim]Hugging Face returned HTTP {response.status_code}; " + f"retrying in {int(retry_delay)}s...[/dim]" + ) + time.sleep(retry_delay) + continue + if response.status_code != 200: + try: + detail = response.json().get("error") or response.text + except ValueError: + detail = response.text + raise BulkPushValidationError( + [ + f"Hugging Face request failed with HTTP {response.status_code} " + f"({detail}). The dataset may be private/gated (set HF_TOKEN) or have " + "the dataset viewer disabled — as a fallback, extract the image " + "references yourself and use --manifest" + ] + ) + return response.json() + raise AssertionError("unreachable") + + +@dataclass +class HFDataset: + """One selected (dataset, config) pair with its /info metadata.""" + + dataset_id: str + config: str + config_info: dict[str, Any] + partial: bool + + @property + def features(self) -> dict[str, Any]: + features = self.config_info.get("features") + return features if isinstance(features, dict) else {} + + +def select_hf_config(dataset: str, config: Optional[str]) -> HFDataset: + """Fetch /info for the dataset and resolve which config to read.""" + dataset_id = normalize_hf_dataset_id(dataset) + info = hf_get("/info", {"dataset": dataset_id}) + configs = info.get("dataset_info") + if not isinstance(configs, dict) or not configs: + raise BulkPushValidationError([f"no dataset info available for '{dataset_id}'"]) + if config is None: + if len(configs) > 1: + raise BulkPushValidationError( + [ + f"dataset '{dataset_id}' has multiple configs " + f"({', '.join(sorted(configs))}); choose one with --hf-config" + ] + ) + config = next(iter(configs)) + elif config not in configs: + raise BulkPushValidationError( + [ + f"config '{config}' not found in dataset '{dataset_id}' " + f"(available: {', '.join(sorted(configs))})" + ] + ) + return HFDataset( + dataset_id=dataset_id, + config=config, + config_info=configs.get(config) or {}, + partial=bool(info.get("partial")), + ) + + +def missing_column_error(flag: str, ds: HFDataset) -> BulkPushValidationError: + """Error for a required column flag that was not given, listing the columns.""" + return BulkPushValidationError( + [ + f"{flag} is required with --hf (columns in " + f"'{ds.dataset_id}': {', '.join(sorted(ds.features)) or 'unknown'})" + ] + ) + + +def require_string_column(ds: HFDataset, column: str, *, reason: str) -> None: + """Reject columns that are missing or whose feature type cannot hold the + needed strings (e.g. actual pictures, nested lists) before submitting + anything.""" + features = ds.features + if features and column not in features: + raise BulkPushValidationError( + [ + f"column '{column}' not found in dataset '{ds.dataset_id}' " + f"(columns: {', '.join(sorted(features))})" + ] + ) + feature_spec = features.get(column) + if feature_spec is not None and not ( + isinstance(feature_spec, dict) + and feature_spec.get("_type") == "Value" + and feature_spec.get("dtype") in ("string", "large_string") + ): + raise BulkPushValidationError( + [f"column '{column}' in '{ds.dataset_id}' is not a string column, so {reason}"] + ) + + +def check_split(ds: HFDataset, split: str) -> None: + splits = ds.config_info.get("splits") + if isinstance(splits, dict) and splits and split not in splits: + raise BulkPushValidationError( + [ + f"split '{split}' not found in dataset '{ds.dataset_id}' " + f"(available: {', '.join(sorted(splits))})" + ] + ) + + +def warn_if_large(ds: HFDataset) -> None: + # The rows API returns full rows (a columns filter is not honored), so + # reading time scales with total dataset size, not just row count. + dataset_size = ds.config_info.get("dataset_size") + if isinstance(dataset_size, (int, float)) and dataset_size > LARGE_DATASET_WARN_BYTES: + console.print( + f"[dim]Dataset rows total ~{dataset_size / 1e9:.1f}GB and the viewer API " + "streams full rows, so reading may take a few minutes.[/dim]" + ) + + +def iter_hf_rows(ds: HFDataset, split: str) -> Iterator[tuple[int, dict[str, Any]]]: + """Yield (row_idx, row) for every row of the split, with a progress bar.""" + offset = 0 + total: Optional[int] = None + progress = Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TextColumn("{task.completed}/{task.total}"), + TimeElapsedColumn(), + console=console, + ) + with progress: + task_id = progress.add_task("Reading dataset rows", total=None) + while True: + page = hf_get( + "/rows", + { + "dataset": ds.dataset_id, + "config": ds.config, + "split": split, + "offset": offset, + "length": HF_ROWS_PAGE_LENGTH, + }, + ) + rows = page.get("rows") or [] + if isinstance(page.get("num_rows_total"), int): + total = page["num_rows_total"] + progress.update(task_id, total=total) + for item in rows: + yield item.get("row_idx", offset), item.get("row") or {} + offset += len(rows) + progress.update(task_id, completed=offset) + if not rows or (total is not None and offset >= total): + break + + +# --------------------------------------------------------------------------- +# Build-list resolution for push-bulk: Hugging Face datasets of Dockerfiles +# --------------------------------------------------------------------------- + + +def load_hf_build_specs( + dataset: str, + *, + config: Optional[str], + split: str, + dockerfile_column: Optional[str], + name_column: Optional[str], + tag: str, + name_template: str, + platform: str, + context_root: Path, +) -> tuple[list[BuildSpec], list[str]]: + """Resolve build specs from a Hugging Face dataset of Dockerfiles. + + The push-bulk complement of transfer-bulk's --hf mode: instead of image + references, each row carries Dockerfile *contents* (--dockerfile-column) + plus a name for the resulting image (--name-column). Each usable row's + Dockerfile is written under ``context_root`` as a one-file build context. + Rows with identical (name, Dockerfile) values collapse into one build. + Returns (specs, notes) where notes are informational messages. + """ + ds = select_hf_config(dataset, config) + notes: list[str] = [] + if dockerfile_column is None: + raise missing_column_error("--dockerfile-column", ds) + if name_column is None: + raise missing_column_error("--name-column", ds) + require_string_column(ds, dockerfile_column, reason="it cannot hold Dockerfile contents") + require_string_column(ds, name_column, reason="it cannot hold image names") + check_split(ds, split) + if ds.partial: + notes.append(PARTIAL_DATASET_NOTE) + warn_if_large(ds) + + rows: list[tuple[int, str, str]] = [] # (row idx, image name value, Dockerfile) + seen: set[tuple[str, str]] = set() + scanned = 0 + incomplete_rows = 0 + for row_idx, row in iter_hf_rows(ds, split): + name_value = row.get(name_column) + dockerfile = row.get(dockerfile_column) + if ( + not isinstance(name_value, str) + or not name_value.strip() + or not isinstance(dockerfile, str) + or not dockerfile.strip() + ): + incomplete_rows += 1 + continue + scanned += 1 + key = (name_value.strip(), dockerfile) + if key in seen: + continue + seen.add(key) + rows.append((row_idx, name_value.strip(), dockerfile)) + + if incomplete_rows: + notes.append( + f"Skipped {incomplete_rows} row(s) with an empty " + f"'{name_column}' or '{dockerfile_column}' value" + ) + duplicates = scanned - len(rows) + if duplicates: + notes.append(f"Collapsed {duplicates} duplicate row(s)") + + problems: list[str] = [] + specs: list[BuildSpec] = [] + for row_idx, name_value, dockerfile in rows: + # In HF mode both {dir} and {name} render as the row's name value. + image_name = _render_name_template( + name_template, task_dir_name=name_value, toml_name=name_value + ) + if not image_name: + problems.append( + f"{ds.dataset_id} row {row_idx}: image name is empty " + f"after sanitizing '{name_value}'" + ) + continue + context_dir = context_root / f"row-{row_idx}" + context_dir.mkdir(parents=True, exist_ok=True) + dockerfile_path = context_dir / "Dockerfile" + dockerfile_path.write_text(dockerfile) + specs.append( + BuildSpec( + image_name=image_name, + image_tag=tag, + context=context_dir, + dockerfile=dockerfile_path, + platform=platform, + source=f"row {row_idx}", + ) + ) + + problems.extend(_duplicate_ref_problems(specs, hint=" — use --name-template to disambiguate")) + if not specs and not problems: + problems.append( + f"no builds found in '{ds.dataset_id}' (no rows with both " + f"'{name_column}' and '{dockerfile_column}')" + ) + if problems: + raise BulkPushValidationError(problems) + return specs, notes diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 207e8aec..6270464e 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -1,8 +1,6 @@ import json -import os import re import sys -import time from dataclasses import dataclass from pathlib import Path from typing import Any, Optional @@ -13,7 +11,6 @@ import tomli as tomllib import click -import httpx import typer from prime_sandboxes import ( APIClient, @@ -22,7 +19,6 @@ ImageVisibility, UnauthorizedError, ) -from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn from rich.table import Table from ..utils import get_console @@ -44,6 +40,16 @@ discover_harbor_tasks, run_bulk_jobs, ) +from .images_hf import ( + PARTIAL_DATASET_NOTE, + check_split, + iter_hf_rows, + missing_column_error, + normalize_hf_dataset_id, + require_string_column, + select_hf_config, + warn_if_large, +) console = get_console() @@ -51,14 +57,6 @@ # rejects one. TRANSFER_RATE_LIMIT_PAUSE_SECONDS = 15.0 -HF_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" -HF_ROWS_PAGE_LENGTH = 100 -LARGE_DATASET_WARN_BYTES = 500_000_000 -HF_REQUEST_TIMEOUT_SECONDS = 30.0 -HF_REQUEST_MAX_ATTEMPTS = 9 -HF_REQUEST_BACKOFF_SECONDS = 2.0 -HF_REQUEST_BACKOFF_MAX_SECONDS = 60.0 - _TRANSFER_MANIFEST_KEYS = {"source", "image", "platform"} _DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") @@ -315,96 +313,6 @@ def load_harbor_transfer_specs( # --------------------------------------------------------------------------- -def normalize_hf_dataset_id(raw: str) -> str: - """Accept 'org/name', hf://datasets/... and huggingface.co dataset URLs.""" - dataset = (raw or "").strip() - for prefix in ( - "hf://datasets/", - "https://huggingface.co/datasets/", - "http://huggingface.co/datasets/", - "huggingface.co/datasets/", - ): - if dataset.startswith(prefix): - dataset = dataset[len(prefix) :] - dataset = dataset.split("?", 1)[0].split("#", 1)[0] - # Drop trailing URL parts like /viewer/default/train. - parts = [part for part in dataset.split("/") if part] - dataset = "/".join(parts[:2]) - break - dataset = dataset.strip("/") - if not dataset or dataset.count("/") > 1: - raise BulkPushValidationError( - [f"invalid Hugging Face dataset '{raw}' (expected 'org/name' or a dataset URL)"] - ) - return dataset - - -def _hf_get(path: str, params: dict[str, Any]) -> dict[str, Any]: - """GET from the HF datasets-server with retries on transient failures. - - Uses HF_TOKEN from the environment for gated/private datasets. The - datasets-server rate limit is per IP and clears within about a minute, so - 429/5xx retries back off long enough to ride out a full window, honoring - Retry-After when the server sends one. - """ - token = os.environ.get("HF_TOKEN") - headers = {"Authorization": f"Bearer {token}"} if token else {} - url = f"{HF_DATASETS_SERVER_URL}{path}" - delay = HF_REQUEST_BACKOFF_SECONDS - for attempt in range(1, HF_REQUEST_MAX_ATTEMPTS + 1): - retry_delay = delay - delay = min(delay * 2, HF_REQUEST_BACKOFF_MAX_SECONDS) - try: - response = httpx.get( - url, params=params, headers=headers, timeout=HF_REQUEST_TIMEOUT_SECONDS - ) - except httpx.HTTPError as e: - if attempt == HF_REQUEST_MAX_ATTEMPTS: - raise BulkPushValidationError([f"Hugging Face request failed: {e}"]) - time.sleep(retry_delay) - continue - if response.status_code == 429 or response.status_code >= 500: - if attempt == HF_REQUEST_MAX_ATTEMPTS: - rate_limit_hint = ( - " (rate limited — wait a minute and re-run)" - if response.status_code == 429 - else "" - ) - raise BulkPushValidationError( - [ - f"Hugging Face request failed with HTTP " - f"{response.status_code}: {url}{rate_limit_hint}" - ] - ) - retry_after = response.headers.get("retry-after") - if retry_after: - try: - retry_delay = max(retry_delay, float(retry_after)) - except ValueError: - pass - console.print( - f"[dim]Hugging Face returned HTTP {response.status_code}; " - f"retrying in {int(retry_delay)}s...[/dim]" - ) - time.sleep(retry_delay) - continue - if response.status_code != 200: - try: - detail = response.json().get("error") or response.text - except ValueError: - detail = response.text - raise BulkPushValidationError( - [ - f"Hugging Face request failed with HTTP {response.status_code} " - f"({detail}). The dataset may be private/gated (set HF_TOKEN) or have " - "the dataset viewer disabled — as a fallback, extract the image " - "references yourself and use --manifest" - ] - ) - return response.json() - raise AssertionError("unreachable") - - def load_hf_specs( dataset: str, *, @@ -419,130 +327,33 @@ def load_hf_specs( `datasets` dependency), dedupes identical references preserving order, and returns (specs, notes) where notes are informational messages. """ - dataset_id = normalize_hf_dataset_id(dataset) + ds = select_hf_config(dataset, config) + dataset_id = ds.dataset_id notes: list[str] = [] - info = _hf_get("/info", {"dataset": dataset_id}) - configs = info.get("dataset_info") - if not isinstance(configs, dict) or not configs: - raise BulkPushValidationError([f"no dataset info available for '{dataset_id}'"]) - if config is None: - if len(configs) > 1: - raise BulkPushValidationError( - [ - f"dataset '{dataset_id}' has multiple configs " - f"({', '.join(sorted(configs))}); choose one with --hf-config" - ] - ) - config = next(iter(configs)) - elif config not in configs: - raise BulkPushValidationError( - [ - f"config '{config}' not found in dataset '{dataset_id}' " - f"(available: {', '.join(sorted(configs))})" - ] - ) - config_info = configs.get(config) or {} - - features = config_info.get("features") or {} if column is None: - raise BulkPushValidationError( - [ - "--column is required with --hf (columns in " - f"'{dataset_id}': {', '.join(sorted(features)) or 'unknown'})" - ] - ) - if features and column not in features: - raise BulkPushValidationError( - [ - f"column '{column}' not found in dataset '{dataset_id}' " - f"(columns: {', '.join(sorted(features))})" - ] - ) - # Reject columns whose feature type cannot hold image reference strings - # (e.g. actual pictures, nested lists) before submitting anything. - feature_spec = features.get(column) if isinstance(features, dict) else None - if feature_spec is not None and not ( - isinstance(feature_spec, dict) - and feature_spec.get("_type") == "Value" - and feature_spec.get("dtype") in ("string", "large_string") - ): - raise BulkPushValidationError( - [ - f"column '{column}' in '{dataset_id}' is not a string column, " - "so it cannot hold image references" - ] - ) - - splits = config_info.get("splits") - if isinstance(splits, dict) and splits and split not in splits: - raise BulkPushValidationError( - [ - f"split '{split}' not found in dataset '{dataset_id}' " - f"(available: {', '.join(sorted(splits))})" - ] - ) - if info.get("partial"): - notes.append( - "Warning: the dataset viewer only covers part of this dataset; some rows may be missing" - ) - - # The rows API returns full rows (a columns filter is not honored), so - # reading time scales with total dataset size, not just row count. - dataset_size = config_info.get("dataset_size") - if isinstance(dataset_size, (int, float)) and dataset_size > LARGE_DATASET_WARN_BYTES: - console.print( - f"[dim]Dataset rows total ~{dataset_size / 1e9:.1f}GB and the viewer API " - "streams full rows, so reading may take a few minutes.[/dim]" - ) + raise missing_column_error("--column", ds) + require_string_column(ds, column, reason="it cannot hold image references") + check_split(ds, split) + if ds.partial: + notes.append(PARTIAL_DATASET_NOTE) + warn_if_large(ds) sources: list[tuple[str, int]] = [] # (image ref, first row index), order-preserving seen: set[str] = set() scanned = 0 empty_rows = 0 - offset = 0 - total: Optional[int] = None - progress = Progress( - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("{task.completed}/{task.total}"), - TimeElapsedColumn(), - console=console, - ) - with progress: - task_id = progress.add_task("Reading dataset rows", total=None) - while True: - page = _hf_get( - "/rows", - { - "dataset": dataset_id, - "config": config, - "split": split, - "offset": offset, - "length": HF_ROWS_PAGE_LENGTH, - }, - ) - rows = page.get("rows") or [] - if isinstance(page.get("num_rows_total"), int): - total = page["num_rows_total"] - progress.update(task_id, total=total) - for item in rows: - row = item.get("row") or {} - row_idx = item.get("row_idx", offset) - value = row.get(column) - if not isinstance(value, str) or not value.strip(): - empty_rows += 1 - continue - scanned += 1 - value = value.strip() - if value in seen: - continue - seen.add(value) - sources.append((value, row_idx)) - offset += len(rows) - progress.update(task_id, completed=offset) - if not rows or (total is not None and offset >= total): - break + for row_idx, row in iter_hf_rows(ds, split): + value = row.get(column) + if not isinstance(value, str) or not value.strip(): + empty_rows += 1 + continue + scanned += 1 + value = value.strip() + if value in seen: + continue + seen.add(value) + sources.append((value, row_idx)) if empty_rows: notes.append(f"Skipped {empty_rows} row(s) with an empty '{column}' value") @@ -683,7 +494,7 @@ def transfer_bulk( hf_config: Optional[str] = typer.Option( None, "--hf-config", - help="Dataset config for --hf mode", + help="Dataset config (called 'subset' in the HF viewer) for --hf mode", show_default="the dataset's only config", ), column: Optional[str] = typer.Option( diff --git a/packages/prime/tests/test_images_push_bulk.py b/packages/prime/tests/test_images_push_bulk.py index f6b61ff6..b90e757b 100644 --- a/packages/prime/tests/test_images_push_bulk.py +++ b/packages/prime/tests/test_images_push_bulk.py @@ -1,6 +1,9 @@ +import io import json +import tarfile import prime_cli.commands.images_bulk as images_bulk +import prime_cli.commands.images_hf as images_hf import pytest from prime_cli.main import app from prime_sandboxes import APIError @@ -45,6 +48,8 @@ class FakeAPI: def __init__(self): self.calls = [] + self.payloads = [] + self.uploads = [] # raw tar.gz bytes PUT to each upload URL, in order self.build_counter = 0 # build_id -> list of statuses returned by successive GET polls; # the last status repeats once the list is exhausted. @@ -58,6 +63,7 @@ def request(self, method, path, json=None, params=None): if method == "POST" and path == "/images/build": if self.build_error_queue: raise self.build_error_queue.pop(0) + self.payloads.append(json) self.build_counter += 1 build_id = f"build-{self.build_counter}" return { @@ -87,7 +93,7 @@ def raise_for_status(self): def fake_put(url, content, headers=None, timeout=None): api.calls.append(("PUT", url)) - content.read() + api.uploads.append(content.read()) return DummyUploadResponse() monkeypatch.setattr("prime_cli.main.check_for_update", lambda: (False, None)) @@ -259,7 +265,7 @@ def test_requires_exactly_one_mode(tmp_path, fake_api, monkeypatch): monkeypatch.chdir(tmp_path) result = runner.invoke(app, ["images", "push-bulk"], env=TEST_ENV) assert result.exit_code == 1 - assert "exactly one of --manifest or --harbor" in result.output + assert "exactly one of --manifest, --harbor or --hf" in result.output manifest = tmp_path / "builds.jsonl" manifest.write_text("") @@ -269,7 +275,15 @@ def test_requires_exactly_one_mode(tmp_path, fake_api, monkeypatch): env=TEST_ENV, ) assert result.exit_code == 1 - assert "exactly one of --manifest or --harbor" in result.output + assert "exactly one of --manifest, --harbor or --hf" in result.output + + result = runner.invoke( + app, + ["images", "push-bulk", "--manifest", str(manifest), "--hf", "org/ds"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "exactly one of --manifest, --harbor or --hf" in result.output def test_tag_flag_rejected_in_manifest_mode(tmp_path, fake_api, monkeypatch): @@ -374,3 +388,271 @@ def test_harbor_root_is_single_task(tmp_path, fake_api, monkeypatch): assert result.exit_code == 0, result.output assert len(fake_api.post_build_indices()) == 1 assert "1/1 builds completed" in result.output + + +# --------------------------------------------------------------------------- +# Hugging Face mode +# --------------------------------------------------------------------------- + + +def _fake_hf(monkeypatch, *, info, pages, total): + """Serve datasets-server /info and /rows from canned data. + + ``pages`` maps a row offset to the list of row dicts served at that offset. + """ + calls = [] + + class FakeResponse: + def __init__(self, payload, status_code=200, headers=None): + self.status_code = status_code + self._payload = payload + self.text = json.dumps(payload) + self.headers = headers or {} + + def json(self): + return self._payload + + def fake_get(url, params=None, headers=None, timeout=None): + params = dict(params or {}) + calls.append((url, params)) + if url.endswith("/info"): + return FakeResponse(info) + if url.endswith("/rows"): + offset = params["offset"] + rows = [ + {"row_idx": offset + i, "row": row} for i, row in enumerate(pages.get(offset, [])) + ] + return FakeResponse({"rows": rows, "num_rows_total": total}) + raise AssertionError(f"Unexpected HF request: {url}") + + monkeypatch.setattr(images_hf.httpx, "get", fake_get) + return calls + + +HF_INFO_ONE_CONFIG = { + "dataset_info": { + "default": { + "features": { + "dockerfile": {"dtype": "string", "_type": "Value"}, + "instance_id": {"dtype": "string", "_type": "Value"}, + "picture": {"_type": "Image"}, + }, + "splits": {"train": {"name": "train"}}, + } + }, + "partial": False, +} + +HF_ARGS = ["--dockerfile-column", "dockerfile", "--name-column", "instance_id"] + + +def _packaged_dockerfile(tar_bytes): + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r:gz") as tar: + member = tar.extractfile(images_bulk.PACKAGED_DOCKERFILE_PATH) + assert member is not None + return member.read().decode() + + +def _use_local_hf_context_root(monkeypatch, tmp_path): + """Route the generated HF build contexts to a known directory.""" + root = tmp_path / "hf-contexts" + + def fake_mkdtemp(prefix): + root.mkdir() + return str(root) + + monkeypatch.setattr(images_bulk.tempfile, "mkdtemp", fake_mkdtemp) + return root + + +def test_hf_builds_each_row_dockerfile(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + context_root = _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={ + 0: [ + {"instance_id": "Org__Repo-1", "dockerfile": "FROM busybox\nRUN echo a\n"}, + {"instance_id": "org__repo-2", "dockerfile": "FROM busybox\nRUN echo b\n"}, + ] + }, + total=2, + ) + + result = runner.invoke( + app, ["images", "push-bulk", "--hf", "org/ds", "--tag", "v1", *HF_ARGS], env=TEST_ENV + ) + + assert result.exit_code == 0, result.output + assert [(p["image_name"], p["image_tag"]) for p in fake_api.payloads] == [ + ("org__repo-1", "v1"), + ("org__repo-2", "v1"), + ] + # Each build context carries that row's Dockerfile contents. + assert [_packaged_dockerfile(upload) for upload in fake_api.uploads] == [ + "FROM busybox\nRUN echo a\n", + "FROM busybox\nRUN echo b\n", + ] + assert "2/2 builds completed" in result.output + # Generated contexts are cleaned up when nothing failed. + assert not context_root.exists() + + +def test_hf_name_template_and_dedupe_and_empty_rows(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _use_local_hf_context_root(monkeypatch, tmp_path) + row = {"instance_id": "repo-1", "dockerfile": "FROM busybox\n"} + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={0: [row, dict(row), {"instance_id": "no-dockerfile"}]}, + total=3, + ) + + result = runner.invoke( + app, + ["images", "push-bulk", "--hf", "org/ds", "--name-template", "swe-{name}", *HF_ARGS], + env=TEST_ENV, + ) + + assert result.exit_code == 0, result.output + assert "Collapsed 1 duplicate row(s)" in result.output + assert "Skipped 1 row(s) with an empty 'instance_id' or 'dockerfile' value" in result.output + assert [p["image_name"] for p in fake_api.payloads] == ["swe-repo-1"] + + +def test_hf_same_name_different_dockerfiles_fail(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={ + 0: [ + {"instance_id": "repo-1", "dockerfile": "FROM busybox\nRUN echo a\n"}, + {"instance_id": "repo-1", "dockerfile": "FROM busybox\nRUN echo b\n"}, + ] + }, + total=2, + ) + + result = runner.invoke(app, ["images", "push-bulk", "--hf", "org/ds", *HF_ARGS], env=TEST_ENV) + + assert result.exit_code == 1 + assert "duplicate image reference 'repo-1:latest'" in result.output + assert "--name-template" in result.output + assert fake_api.calls == [] + + +def test_hf_column_flags_are_required(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) + + result = runner.invoke(app, ["images", "push-bulk", "--hf", "org/ds"], env=TEST_ENV) + assert result.exit_code == 1 + assert "--dockerfile-column is required" in result.output + # The error lists the dataset's columns so the user can pick one. + assert "instance_id" in result.output + + result = runner.invoke( + app, + ["images", "push-bulk", "--hf", "org/ds", "--dockerfile-column", "dockerfile"], + env=TEST_ENV, + ) + assert result.exit_code == 1 + assert "--name-column is required" in result.output + assert fake_api.calls == [] + + +def test_hf_non_string_column_rejected(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf(monkeypatch, info=HF_INFO_ONE_CONFIG, pages={0: []}, total=0) + + result = runner.invoke( + app, + [ + "images", + "push-bulk", + "--hf", + "org/ds", + "--dockerfile-column", + "picture", + "--name-column", + "instance_id", + ], + env=TEST_ENV, + ) + + assert result.exit_code == 1 + assert "not a string column" in result.output + assert fake_api.calls == [] + + +def test_hf_flags_rejected_outside_hf_mode(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _make_context(tmp_path, "a") + manifest = tmp_path / "builds.jsonl" + _write_manifest(manifest, [{"image": "app-a:v1", "context": "a"}]) + + result = runner.invoke( + app, + ["images", "push-bulk", "--manifest", str(manifest), "--name-column", "instance_id"], + env=TEST_ENV, + ) + + assert result.exit_code == 1 + assert "only apply to --hf mode" in result.output + assert fake_api.calls == [] + + +def test_hf_dry_run_resolves_without_api_calls(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + context_root = _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={0: [{"instance_id": "repo-1", "dockerfile": "FROM busybox\n"}]}, + total=1, + ) + + result = runner.invoke( + app, ["images", "push-bulk", "--hf", "org/ds", "--dry-run", *HF_ARGS], env=TEST_ENV + ) + + assert result.exit_code == 0, result.output + assert "repo-1:latest" in result.output + assert "Dry run only" in result.output + assert fake_api.calls == [] + assert not context_root.exists() + + +def test_hf_failures_keep_contexts_and_manifest_is_rerunnable(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + context_root = _use_local_hf_context_root(monkeypatch, tmp_path) + _fake_hf( + monkeypatch, + info=HF_INFO_ONE_CONFIG, + pages={0: [{"instance_id": "repo-1", "dockerfile": "FROM busybox\n"}]}, + total=1, + ) + fake_api.default_poll = ["FAILED"] + + result = runner.invoke(app, ["images", "push-bulk", "--hf", "org/ds", *HF_ARGS], env=TEST_ENV) + + assert result.exit_code == 1 + failures_file = tmp_path / "push-bulk-failures.jsonl" + lines = [json.loads(line) for line in failures_file.read_text().splitlines()] + assert [line["image"] for line in lines] == ["repo-1:latest"] + # The generated contexts stay behind so the failures manifest re-runs. + assert str(context_root) in result.output + assert all((tmp_path / line["dockerfile"]).is_file() for line in lines) + + fake_api.default_poll = ["COMPLETED"] + result = runner.invoke( + app, ["images", "push-bulk", "--manifest", str(failures_file)], env=TEST_ENV + ) + assert result.exit_code == 0, result.output + assert "1/1 builds completed" in result.output diff --git a/packages/prime/tests/test_images_transfer_bulk.py b/packages/prime/tests/test_images_transfer_bulk.py index 9087f02f..a8648584 100644 --- a/packages/prime/tests/test_images_transfer_bulk.py +++ b/packages/prime/tests/test_images_transfer_bulk.py @@ -1,6 +1,7 @@ import json import prime_cli.commands.images_bulk as images_bulk +import prime_cli.commands.images_hf as images_hf import prime_cli.commands.images_transfer_bulk as images_transfer_bulk import pytest from prime_cli.commands.images_transfer_bulk import derive_transfer_destination @@ -144,7 +145,7 @@ def fake_get(url, params=None, headers=None, timeout=None): return FakeResponse({"rows": rows, "num_rows_total": total}) raise AssertionError(f"Unexpected HF request: {url}") - monkeypatch.setattr(images_transfer_bulk.httpx, "get", fake_get) + monkeypatch.setattr(images_hf.httpx, "get", fake_get) return calls @@ -326,7 +327,7 @@ def test_hf_pages_rows_and_dedupes(tmp_path, fake_api, monkeypatch): def test_hf_rate_limited_request_retries(tmp_path, fake_api, monkeypatch): - monkeypatch.setattr(images_transfer_bulk, "HF_REQUEST_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr(images_hf, "HF_REQUEST_BACKOFF_SECONDS", 0.0) hf_calls = _fake_hf( monkeypatch, info=HF_INFO_ONE_CONFIG, From 903fefaaad5b03e03033a958ec37a815028535ac Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 7 Jul 2026 14:15:25 -0700 Subject: [PATCH 10/10] interruptions --- .../src/prime_cli/commands/images_bulk.py | 248 +++++++++++------- .../commands/images_transfer_bulk.py | 16 ++ packages/prime/tests/test_images_push_bulk.py | 59 ++++- 3 files changed, 225 insertions(+), 98 deletions(-) diff --git a/packages/prime/src/prime_cli/commands/images_bulk.py b/packages/prime/src/prime_cli/commands/images_bulk.py index 6312e45b..9576e949 100644 --- a/packages/prime/src/prime_cli/commands/images_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_bulk.py @@ -77,6 +77,19 @@ def __init__(self, message: str, retry_after: float): self.retry_after = retry_after +class BulkRunInterrupted(Exception): + """Ctrl+C during a bulk run. + + Carries the outcomes recorded so far plus every never-submitted spec + (marked SKIPPED), so the command can write a resume manifest. Jobs already + in flight are not included — they keep running server-side. + """ + + def __init__(self, outcomes: list["BuildOutcome"]): + super().__init__("bulk run interrupted") + self.outcomes = outcomes + + @dataclass class BuildSpec: """One resolved build: where the context lives and what to call the image.""" @@ -546,114 +559,134 @@ def record(outcome: BuildOutcome) -> None: f"[dim]({outcome.status.lower()}: {outcome.error})[/dim]" ) - while pending or in_flight: - while pending and stop_skip_reason is None and len(in_flight) < concurrency: - if time.monotonic() < submit_gate: - break - spec = pending.popleft() - try: - build_id, full_image_path = submit(spec) - except SubmitRateLimited as e: - consecutive_deferrals += 1 - if consecutive_deferrals > MAX_CONSECUTIVE_SUBMIT_DEFERRALS: + try: + while pending or in_flight: + while pending and stop_skip_reason is None and len(in_flight) < concurrency: + if time.monotonic() < submit_gate: + break + spec = pending.popleft() + try: + build_id, full_image_path = submit(spec) + except KeyboardInterrupt: + # The submit was cut short and may have half-queued the + # job; treating it as never submitted keeps it in the + # resume manifest. + pending.appendleft(spec) + raise + except SubmitRateLimited as e: + consecutive_deferrals += 1 + if consecutive_deferrals > MAX_CONSECUTIVE_SUBMIT_DEFERRALS: + record( + BuildOutcome( + spec=spec, + status="SUBMIT_FAILED", + error=( + f"{e} (gave up after {MAX_CONSECUTIVE_SUBMIT_DEFERRALS} " + "rate-limit deferrals with no successful submission)" + ), + ) + ) + stop_skip_reason = ( + "not submitted: the server kept rate-limiting submissions" + ) + break + pending.appendleft(spec) + submit_gate = time.monotonic() + e.retry_after + if not rate_limit_notice_shown: + rate_limit_notice_shown = True + progress.console.print( + "[dim]Server rate limit reached; pacing submissions " + "(jobs already submitted keep running)...[/dim]" + ) + break + except QuotaExceededError as e: + stop_skip_reason = "not submitted: image quota exceeded" + record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) + except UnauthorizedError: + raise + except (APIError, httpx.HTTPError, OSError) as e: + consecutive_deferrals = 0 + record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) + else: + consecutive_deferrals = 0 + in_flight[build_id] = _InFlightBuild( + spec=spec, + full_image_path=full_image_path, + deadline=time.monotonic() + build_timeout, + ) + + if stop_skip_reason is not None and pending: + for spec in pending: record( BuildOutcome( spec=spec, - status="SUBMIT_FAILED", - error=( - f"{e} (gave up after {MAX_CONSECUTIVE_SUBMIT_DEFERRALS} " - "rate-limit deferrals with no successful submission)" - ), + status="SKIPPED", + error=stop_skip_reason, ) ) - stop_skip_reason = ( - "not submitted: the server kept rate-limiting submissions" - ) - break - pending.appendleft(spec) - submit_gate = time.monotonic() + e.retry_after - if not rate_limit_notice_shown: - rate_limit_notice_shown = True - progress.console.print( - "[dim]Server rate limit reached; pacing submissions " - "(jobs already submitted keep running)...[/dim]" - ) + pending.clear() + + if not pending and not in_flight: break - except QuotaExceededError as e: - stop_skip_reason = "not submitted: image quota exceeded" - record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) - except UnauthorizedError: - raise - except (APIError, httpx.HTTPError, OSError) as e: - consecutive_deferrals = 0 - record(BuildOutcome(spec=spec, status="SUBMIT_FAILED", error=str(e))) - else: - consecutive_deferrals = 0 - in_flight[build_id] = _InFlightBuild( - spec=spec, - full_image_path=full_image_path, - deadline=time.monotonic() + build_timeout, - ) - if stop_skip_reason is not None and pending: - for spec in pending: - record( - BuildOutcome( - spec=spec, - status="SKIPPED", - error=stop_skip_reason, + if not in_flight: + # Everything left is waiting on the submit gate. + time.sleep(max(0.0, submit_gate - time.monotonic())) + continue + + time.sleep(POLL_INTERVAL_SECONDS) + + for build_id, entry in list(in_flight.items()): + build_error: Optional[str] = None + try: + status_response = client.request("GET", f"/images/build/{build_id}") + build_status = str(status_response.get("status") or "") + error_message = status_response.get("errorMessage") or status_response.get( + "error_message" ) - ) - pending.clear() - - if not pending and not in_flight: - break - - if not in_flight: - # Everything left is waiting on the submit gate. - time.sleep(max(0.0, submit_gate - time.monotonic())) - continue - - time.sleep(POLL_INTERVAL_SECONDS) - - for build_id, entry in list(in_flight.items()): - try: - status_response = client.request("GET", f"/images/build/{build_id}") - build_status = str(status_response.get("status") or "") - except UnauthorizedError: - raise - except APIError: - build_status = "" # transient poll failure; the deadline still applies - - if build_status in {"COMPLETED", "FAILED", "CANCELLED"}: - del in_flight[build_id] - record( - BuildOutcome( - spec=entry.spec, - status=build_status, - build_id=build_id, - full_image_path=entry.full_image_path, - error=( - None - if build_status == "COMPLETED" - else f"build ended as {build_status}" - ), + if isinstance(error_message, str) and error_message.strip(): + build_error = error_message.strip() + except UnauthorizedError: + raise + except APIError: + build_status = "" # transient poll failure; the deadline still applies + + if build_status in {"COMPLETED", "FAILED", "CANCELLED"}: + del in_flight[build_id] + failure_reason = f"build ended as {build_status}" + if build_error: + failure_reason += f": {build_error}" + record( + BuildOutcome( + spec=entry.spec, + status=build_status, + build_id=build_id, + full_image_path=entry.full_image_path, + error=None if build_status == "COMPLETED" else failure_reason, + ) ) - ) - elif time.monotonic() > entry.deadline: - del in_flight[build_id] - record( - BuildOutcome( - spec=entry.spec, - status="TIMEOUT", - build_id=build_id, - full_image_path=entry.full_image_path, - error=( - f"no terminal status after {build_timeout}s " - "(the build may still finish server-side)" - ), + elif time.monotonic() > entry.deadline: + del in_flight[build_id] + record( + BuildOutcome( + spec=entry.spec, + status="TIMEOUT", + build_id=build_id, + full_image_path=entry.full_image_path, + error=( + f"no terminal status after {build_timeout}s " + "(the build may still finish server-side)" + ), + ) ) + except KeyboardInterrupt: + for spec in pending: + outcomes.append( + BuildOutcome( + spec=spec, status="SKIPPED", error="not submitted: run interrupted" ) + ) + raise BulkRunInterrupted(outcomes) from None return outcomes @@ -999,6 +1032,27 @@ def push_bulk( except UnauthorizedError: console.print("[red]Error: Not authenticated. Please run 'prime login' first.[/red]") raise typer.Exit(1) + except BulkRunInterrupted as e: + console.print( + "\n[yellow]Cancelled. Builds already started keep running server-side; " + "check them with 'prime images list'.[/yellow]" + ) + unfinished = [o for o in e.outcomes if o.status != "COMPLETED"] + if unfinished: + _write_failures_manifest(failures_out, unfinished) + console.print( + f"Wrote {len(unfinished)} unfinished build(s) to [bold]{failures_out}[/bold]" + ) + console.print( + f"[dim]Resume with: prime images push-bulk --manifest {failures_out}[/dim]" + ) + if hf_context_root is not None: + keep_hf_context = True + console.print( + f"[dim]Keeping generated Dockerfiles in {hf_context_root} " + "so the manifest can be re-run.[/dim]" + ) + raise typer.Exit(1) except KeyboardInterrupt: console.print( "\n[yellow]Cancelled. Builds already started keep running server-side; " diff --git a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py index 6270464e..48cddce8 100644 --- a/packages/prime/src/prime_cli/commands/images_transfer_bulk.py +++ b/packages/prime/src/prime_cli/commands/images_transfer_bulk.py @@ -31,6 +31,7 @@ POLL_INTERVAL_SECONDS, SUPPORTED_PLATFORMS, BulkPushValidationError, + BulkRunInterrupted, QuotaExceededError, SubmitRateLimited, _duplicate_ref_problems, @@ -766,6 +767,21 @@ def transfer_bulk( except UnauthorizedError: console.print("[red]Error: Not authenticated. Please run 'prime login' first.[/red]") raise typer.Exit(1) + except BulkRunInterrupted as e: + console.print( + "\n[yellow]Cancelled. Transfers already queued keep running server-side; " + "check them with 'prime images list'.[/yellow]" + ) + unfinished = [o for o in e.outcomes if o.status != "COMPLETED"] + if unfinished: + _write_failures_manifest(failures_out, unfinished) + console.print( + f"Wrote {len(unfinished)} unfinished transfer(s) to [bold]{failures_out}[/bold]" + ) + console.print( + f"[dim]Resume with: prime images transfer-bulk --manifest {failures_out}[/dim]" + ) + raise typer.Exit(1) except KeyboardInterrupt: console.print( "\n[yellow]Cancelled. Transfers already queued keep running server-side; " diff --git a/packages/prime/tests/test_images_push_bulk.py b/packages/prime/tests/test_images_push_bulk.py index b90e757b..7fede006 100644 --- a/packages/prime/tests/test_images_push_bulk.py +++ b/packages/prime/tests/test_images_push_bulk.py @@ -57,6 +57,8 @@ def __init__(self): self.default_poll = ["COMPLETED"] # Exceptions raised (FIFO) by POST /images/build before any build is created. self.build_error_queue = [] + # errorMessage returned alongside FAILED poll statuses. + self.poll_error_message = None def request(self, method, path, json=None, params=None): self.calls.append((method, path)) @@ -76,7 +78,11 @@ def request(self, method, path, json=None, params=None): if method == "GET" and path.startswith("/images/build/"): build_id = path.rsplit("/", 1)[1] script = self.poll_scripts.setdefault(build_id, list(self.default_poll)) - return {"status": script.pop(0) if len(script) > 1 else script[0]} + status_value = script.pop(0) if len(script) > 1 else script[0] + response = {"status": status_value} + if status_value == "FAILED" and self.poll_error_message: + response["errorMessage"] = self.poll_error_message + return response raise AssertionError(f"Unexpected request: {method} {path}") def post_build_indices(self): @@ -656,3 +662,54 @@ def test_hf_failures_keep_contexts_and_manifest_is_rerunnable(tmp_path, fake_api ) assert result.exit_code == 0, result.output assert "1/1 builds completed" in result.output + + +# --------------------------------------------------------------------------- +# Failure detail and interruption +# --------------------------------------------------------------------------- + + +def test_failed_build_surfaces_server_error_message(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + _make_context(tmp_path, "a") + manifest = tmp_path / "builds.jsonl" + _write_manifest(manifest, [{"image": "app-a:v1", "context": "a"}]) + fake_api.default_poll = ["FAILED"] + fake_api.poll_error_message = "failed to pull base image openswe-python-3.12: not found" + + result = runner.invoke(app, ["images", "push-bulk", "--manifest", str(manifest)], env=TEST_ENV) + + assert result.exit_code == 1 + assert "openswe-python-3.12" in result.output + + +def test_interrupt_writes_resume_manifest(tmp_path, fake_api, monkeypatch): + monkeypatch.chdir(tmp_path) + for name in ("a", "b", "c"): + _make_context(tmp_path, name) + manifest = tmp_path / "builds.jsonl" + _write_manifest(manifest, [{"image": f"app-{n}:v1", "context": n} for n in ("a", "b", "c")]) + + original_request = fake_api.request + + def interrupt_on_poll(method, path, json=None, params=None): + if method == "GET" and path.startswith("/images/build/"): + raise KeyboardInterrupt() + return original_request(method, path, json=json, params=params) + + monkeypatch.setattr(fake_api, "request", interrupt_on_poll) + + result = runner.invoke( + app, + ["images", "push-bulk", "--manifest", str(manifest), "--concurrency", "1"], + env=TEST_ENV, + ) + + assert result.exit_code == 1 + assert "Cancelled" in result.output + assert "Resume with" in result.output + # app-a was submitted (still running server-side, not in the manifest); + # app-b and app-c were never submitted and land in the resume manifest. + failures_file = tmp_path / "push-bulk-failures.jsonl" + lines = [json.loads(line) for line in failures_file.read_text().splitlines()] + assert [line["image"] for line in lines] == ["app-b:v1", "app-c:v1"]