From 3df1156614764e5981d32d4a2ba1e53cbdb274a5 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 21 Jul 2026 13:51:59 -0600 Subject: [PATCH 1/5] =?UTF-8?q?gw#585:=20consume=20tensorhub=20v4=20privat?= =?UTF-8?q?e-input=20manifests=20=E2=80=94=20v3->v4=20hard=20cut=20(PR=20#?= =?UTF-8?q?320=20port=20onto=20chaos)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proto adopted byte-for-byte from tensorhub chaos 07e5520b (protocol 4, RunJob.input_assets=15 — 14 taken by th#913 lane — InputAsset fields 1-6; CLI NDJSON stays v1); pb2 regenerated with grpcio-tools 1.82.1. Worker resolves the exact ordered manifest via one bounded resolve call under the attempt-scoped capability, verifies size/BLAKE3/MIME/kind/decode, preserves opaque refs, sets only local_path, attempt-owned temp cleanup; endpoint build rejects Asset-bearing set/frozenset and non-string maps; deterministic traversal (declaration order, list index, sorted string map keys). Full local suite 558 passed; mypy/ruff/ timeout-lint clean. uv.lock: mechanical 0.42.0 catch-up to pyproject. --- proto/CONTRACT.md | 48 +- proto/worker_scheduler.proto | 40 +- src/gen_worker/discovery/discover.py | 70 +- src/gen_worker/executor.py | 16 +- src/gen_worker/input_assets.py | 463 +++++++-- src/gen_worker/pb/worker_scheduler_pb2.py | 122 +-- src/gen_worker/pb/worker_scheduler_pb2.pyi | 22 +- tests/test_input_manifest_v4.py | 1031 ++++++++++++++++++++ 8 files changed, 1656 insertions(+), 156 deletions(-) create mode 100644 tests/test_input_manifest_v4.py diff --git a/proto/CONTRACT.md b/proto/CONTRACT.md index 993d536d..ab956fcb 100644 --- a/proto/CONTRACT.md +++ b/proto/CONTRACT.md @@ -3,7 +3,7 @@ Proto: `worker_scheduler.proto`, package `cozy.scheduler`, service `WorkerScheduler`, one RPC: `Connect(stream WorkerMessage) returns (stream SchedulerMessage)`. This REPLACES the old `scheduler.v1` proto entirely. No compat, no negotiation: -`ProtocolVersion.PROTOCOL_VERSION_CURRENT = 3` is the only accepted version. +`ProtocolVersion.PROTOCOL_VERSION_CURRENT = 4` is the only accepted version. Roles: **W** = gen-worker (Python), **O** = orchestrator (tensorhub Go). Every field below names its producer and consumer; anything without both was deleted. @@ -269,7 +269,7 @@ Reply to Hello; re-sent on config change (full replace). |---|---|---|---| | `protocol_version` | O constant | W version gate | symmetric validation | | `file_base_url` | O config (`FileAPIBaseURL`) | W upload stack / blob-ref PUT flow (§8) | tensorhub HTTP base for capability-token calls | -| `keep` | legacy v2 release config | none in v3 | retained field number only; v3 W ignores it | +| `keep` | legacy v2 release config | none in v4 | retained field number only; v4 W ignores it | | `resolutions` | O precision resolver | W endpoint bindings | full-replace declared-ref to worker-specific ref/cast picks | | `desired_residency` | O scheduler/controller | W model reconciler | full-replace per-worker disk and hot-instance goal (§7) | @@ -355,7 +355,7 @@ not inherit another alias's proof. | `request_id` | O | W executor | unique request id | | `attempt` | O fencing (§2) | W (echoed on all replies) | dispatch fence | | `function_name` | O from client submit | W registry lookup | dispatch key; unknown fn ⇒ immediate `JobResult{INVALID}` | -| `input_payload` | O assignment payload derived from `/invoke` | W deserializes, validates against the fn payload type | MessagePack bytes; O keeps canonical stored refs in durable request identity but replaces typed Asset refs only in this ephemeral payload with fresh authorized HTTP(S) URLs. W rejects opaque refs and caller `local_path`, downloads distinct URLs once in payload order, assigns that single fresh attempt-local path to every duplicate occurrence, and cleans it on every terminal path before model/handler work can observe a failure | +| `input_payload` | O copies the canonical `/invoke` payload byte-for-byte | W deserializes, materializes, then validates against the fn payload type | MessagePack bytes. Stored Asset refs remain opaque `source_ref` values; transport URLs never enter canonical bytes. Caller HTTP(S) refs remain in this payload and are absent from `input_assets`. W rejects caller `local_path` values, downloads each distinct input once in deterministic payload order, assigns one fresh attempt-local path to every duplicate occurrence, and cleans it on every terminal path | | `timeout_ms` | O from endpoint config | W deadline watchdog | 0 = none; on expiry W aborts and sends `JobResult{FATAL, safe_message:"deadline exceeded"}` | | `tenant` | O request state | W upload scoping + structured logs | invoking owner slug | | `invoker_id` | O request state | W `ctx.invoker_id` (read-only tenant surface) | optional | @@ -365,6 +365,46 @@ not inherit another alias's proof. | `models` | O binding resolver (endpoint defaults + `_models` envelope) | W model injection (`ensure_local` + typed path/pipeline) + per-request LoRA overlays | slot → ref (+ optional `loras`) | | `snapshots` | O resolver | W download stack | presigned snapshots for tensorhub-CAS refs in `models` (including LoRA overlay refs) that O doesn't know to be on this worker's disk; W ignores entries already local (digest match). hf/civitai refs need no snapshot | | `required_compile` | O unique target selection | W pre-setup + pre-execution fence | exact target/cell/contract required for this attempt; mandatory for W8A8 and fail-closed whenever present | +| `input_assets` | O submit-time stored-asset inspector | W private-input materializer | Unique by `(source_ref, kind)` in deterministic schema traversal order: endpoint field declaration, then list index, with `*` map keys sorted; the first traversal occurrence of a duplicate wins. Each entry is exactly `asset_id`, `source_ref`, lowercase BLAKE3 hex, `size_bytes`, `kind`, and lowercase `mime_type`; it contains no URL, path, auth, schema, or count field. | + +### Input-asset manifests (RunJob.input_assets) + +The endpoint contract must expose Asset-bearing collections through ordered +`list`/`tuple` shapes. Asset-bearing `set`/`frozenset` shapes and mappings with +non-string keys are rejected by the worker's schema discovery at build time: +they have no stable occurrence order, and O does not invent a sorting rule. +Generic `Asset`/`MediaAsset` fields use `kind = "media"`; typed +image/video/audio fields use their exact kind. Leading or trailing whitespace +in a typed Asset `ref` is invalid, so `input_payload` and +`input_assets[].source_ref` cannot disagree. + +W derives the ordered unique private sequence by `(ref, declared kind)` from +the decoded payload and requires exact count/order/ref/kind plus valid +immutable metadata against `input_assets` BEFORE any resolver call or GET; a +mismatch is a non-retryable failure. For non-empty `input_assets`, the +capability JWT carries `attempt` and the signed manifest digest, and W resolves +the whole list with exactly one bounded strict request per attempt: + +`POST {file_base_url}/api/v1/worker/input-assets/resolve` + +`Authorization: Bearer ` + +`{"request_id":"","attempt":}` + +The response is `{"assets":[...]}` with the same six manifest fields plus +`url` and `url_expires_at`. W rejects a missing, extra, reordered, changed, +unknown, or malformed entry, validates every URL before downloading, downloads +each unique input once, verifies exact byte count, streaming BLAKE3, media +kind, caller limits, and real image decode where applicable, keeps every +decoded Asset field (`ref` stays the opaque `source_ref`), and sets only the +worker-local `local_path`. URLs never replace `ref` and are never persisted. A +resolver `409` means the attempt is stale ⇒ W stops as CANCELED without +publishing; `503`/network failures and malformed platform responses are +RETRYABLE. Capabilities, URLs, opaque refs, and response bodies never appear in +errors or logs. Materialized inputs are attempt-scoped and transactional: +success, handler failure, cancellation, and supersession each clean exactly +their own attempt's directory — attempt N+1 never deletes N's scope and a +superseded attempt never publishes its result. ### RequiredCompileExecution (embedded in RunJob) @@ -572,7 +612,7 @@ There is no queued-only/item-level cancel. ### ModelOp (O → W) Producer: compile-cache adoption only. Model download/load/unload lifecycle is -declarative through `HelloAck.desired_residency`; v3 O MUST NOT send the legacy +declarative through `HelloAck.desired_residency`; v4 O MUST NOT send the legacy DOWNLOAD/LOAD/UNLOAD enum values. | field | producer | consumer | semantics | diff --git a/proto/worker_scheduler.proto b/proto/worker_scheduler.proto index c61a6271..4d1023b0 100644 --- a/proto/worker_scheduler.proto +++ b/proto/worker_scheduler.proto @@ -30,7 +30,7 @@ option go_package = "github.com/cozy-creator/tensorhub/internal/orchestrator/grp // FAILED_PRECONDITION. No minor versions, no feature gates, no negotiation. enum ProtocolVersion { PROTOCOL_VERSION_UNSPECIFIED = 0; - PROTOCOL_VERSION_CURRENT = 3; + PROTOCOL_VERSION_CURRENT = 4; } service WorkerScheduler { @@ -135,7 +135,7 @@ message ModelResidency { ResidencyTier tier = 2; // highest tier the model currently occupies int64 vram_bytes = 3; // measured footprint; set when tier = VRAM, else 0 // Exact immutable bytes and the DesiredResidency generation that caused - // their materialization. Both are additive protocol-v3 evidence: legacy + // their materialization. Both are protocol-v4 evidence: legacy // workers leave them zero/empty, which means identity is unknown rather // than inferred from the mutable ref. string snapshot_digest = 4; @@ -159,14 +159,14 @@ message InFlightJob { message HelloAck { ProtocolVersion protocol_version = 1; // orchestrator's version, for symmetric validation string file_base_url = 2; // tensorhub base URL for capability-token HTTP calls (uploads, blob refs) - repeated string keep = 3; // legacy v2 field; ignored by protocol v3 workers + repeated string keep = 3; // legacy v2 field; ignored by protocol v4 workers // th#697 precision ladder: per-model flavor resolutions for THIS worker's // card. The worker applies each entry to bindings whose wire ref equals // `ref` (a bare declared ref): download/load `resolved_ref` and apply the // `cast` transform at load. Full-replace with the ack. Additive: old // workers ignore it and load their declared bindings. repeated ModelResolution resolutions = 4; - // Full-replace desired model residency for this worker (protocol v3). + // Full-replace desired model residency for this worker (protocol v4). DesiredResidency desired_residency = 5; } @@ -338,7 +338,7 @@ message RunJob { string request_id = 1; int64 attempt = 2; // fencing token; bumped by the orchestrator on every (re)queue string function_name = 3; - bytes input_payload = 4; // MessagePack-encoded input struct (input file refs pre-materialized to URLs) + bytes input_payload = 4; // canonical MessagePack input; stored Asset refs remain opaque int64 timeout_ms = 5; // execution deadline; 0 = no explicit timeout string tenant = 6; // invoking owner (org/user slug) string invoker_id = 7; // surfaced read-only as ctx.invoker_id @@ -362,6 +362,18 @@ message RunJob { // instruction (worker's own policy). A worker that cannot serve the named // lane answers a typed refusal naming it — never a silent fallback. string lane = 14; + // Ordered, credential-free stored inputs authorized at submit. The worker + // resolves fresh transport URLs for the exact request+attempt capability. + repeated InputAsset input_assets = 15; +} + +message InputAsset { + string asset_id = 1; + string source_ref = 2; + string blake3 = 3; + int64 size_bytes = 4; + string kind = 5; + string mime_type = 6; } message RequiredCompileExecution { @@ -604,7 +616,7 @@ message ModelEvent { // Immutable bytes and DesiredResidency generation for this ref transition. // A receiver must treat an empty digest or zero generation as identity // unknown. Generation fences delayed events after a tag moves or a newer - // desired replacement is reconciled without changing protocol version 3. + // desired replacement is reconciled without changing protocol version 4. string snapshot_digest = 16; uint64 residency_generation = 17; // Echoes ModelOp.operation_id on ADOPTED/adopt-FAILED and at most one later @@ -615,14 +627,14 @@ message ModelEvent { // one later causal compiled-runtime FAILED; empty on boot-attached cells and // ordinary model events. string target_incarnation_id = 19; - // th#850 managed-tier fill hierarchy (gw#599): bytes actually fetched over - // the network (R2 origin) to produce this ON_DISK transition, as opposed - // to bytes served from a warm local-disk or endpoint-volume cache hit. Set - // on every fresh materialization's ON_DISK event, including an explicit 0 - // — pair with the DOWNLOADING events' bytes_total for this ref to read - // "warm boot ⇒ ~0 R2 bytes". Additive: old workers never set it, which a - // receiver must treat as identity unknown (same convention as - // snapshot_digest/residency_generation above), never as a verified zero. + // th#850: cumulative bytes actually fetched over the network (the public + // R2 source, never a local/volume CAS hit) for this DOWNLOADING operation. + // Monotonically non-decreasing within one operation, same cadence as + // bytes_done; set on progress events and the last one before completion is + // authoritative. This is the minimal worker-side signal the "volume- + // attached boot with warm blobs ⇒ R2 bytes ≈ 0" runtime assertion needs — + // distinct from bytes_done/bytes_total, which count ALL bytes materialized + // (including cache hits, both local disk and volume-backed CAS roots). int64 network_bytes = 20; } diff --git a/src/gen_worker/discovery/discover.py b/src/gen_worker/discovery/discover.py index f22a86d4..a9591134 100644 --- a/src/gen_worker/discovery/discover.py +++ b/src/gen_worker/discovery/discover.py @@ -48,9 +48,6 @@ def _is_msgspec_struct(t: Any) -> bool: return False -_MEDIA_ASSET_TYPES = (MediaAsset, ImageAsset, VideoAsset, AudioAsset) - - def _media_kind(t: type) -> str: if issubclass(t, ImageAsset): return "image" @@ -61,6 +58,46 @@ def _media_kind(t: type) -> str: return "media" +def _annotation_carries_asset(ann: Any, _seen: Optional[Set[type]] = None) -> bool: + """True when the annotation subtree can hold an input ``Asset``.""" + seen = _seen if _seen is not None else set() + origin = typing.get_origin(ann) + if origin is typing.Annotated: + args = typing.get_args(ann) + return bool(args) and _annotation_carries_asset(args[0], seen) + if origin in (typing.Union, py_types.UnionType): + return any( + _annotation_carries_asset(arg, seen) + for arg in typing.get_args(ann) + if arg is not type(None) + ) + if origin in (list, tuple, set, frozenset): + args = typing.get_args(ann) + return bool(args) and _annotation_carries_asset(args[0], seen) + if origin is dict: + args = typing.get_args(ann) + return len(args) == 2 and _annotation_carries_asset(args[1], seen) + if isinstance(ann, type): + if issubclass(ann, Asset): + return True + if issubclass(ann, Tensors): + return False + if _is_msgspec_struct(ann): + if ann in seen: + return False + seen.add(ann) + try: + hints = typing.get_type_hints(ann, include_extras=True) + except Exception: + hints = getattr(ann, "__annotations__", {}) + return any( + _annotation_carries_asset(hints[field], seen) + for field in getattr(ann, "__struct_fields__", ()) or () + if field in hints + ) + return False + + def _collect_payload_moderation_metadata(payload_type: type) -> Dict[str, Any]: out: Dict[str, list[Dict[str, str]]] = {"prompts": [], "media": []} seen_structs: set[type] = set() @@ -90,7 +127,20 @@ def walk(ann: Any, path: str) -> None: walk(arg, path) return - if origin in (list, tuple, set, frozenset): + if origin in (set, frozenset): + # th#886: input-asset manifests are ordered; unordered containers + # have no stable occurrence order, so an Asset here is a build error. + args = typing.get_args(ann) + if args and _annotation_carries_asset(args[0]): + raise ValueError( + f"{path}: Asset fields cannot ride unordered set/frozenset " + "containers; use list or tuple" + ) + if args: + walk(args[0], f"{path}[]") + return + + if origin in (list, tuple): args = typing.get_args(ann) if args: walk(args[0], f"{path}[]") @@ -99,14 +149,20 @@ def walk(ann: Any, path: str) -> None: if origin is dict: args = typing.get_args(ann) if len(args) == 2: + if args[0] is not str and _annotation_carries_asset(args[1]): + raise ValueError( + f"{path}: Asset-bearing mappings require string keys" + ) walk(args[1], f"{path}.*") return if isinstance(ann, type): - if issubclass(ann, _MEDIA_ASSET_TYPES): - out["media"].append({"field": path, "kind": _media_kind(ann)}) + if issubclass(ann, Tensors): return - if issubclass(ann, (Asset, Tensors)): + if issubclass(ann, Asset): + # Base Asset/MediaAsset = kind "media"; typed subclasses carry + # their exact kind (th#886 manifest vocabulary). + out["media"].append({"field": path, "kind": _media_kind(ann)}) return if _is_msgspec_struct(ann): if ann in seen_structs: diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 47d983f1..da2a6752 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -53,7 +53,7 @@ InsufficientDiskError, InsufficientHostRamError, ) -from .input_assets import cleanup_input_assets, materialize_input_assets +from .input_assets import cleanup_input_assets, manifest_from_run_job, materialize_input_assets from .models import cozy_snapshot from .models import disk_gc from .models import disk_telemetry @@ -6216,16 +6216,20 @@ async def _run_job_pinned( ) try: - # Tensorhub #871 owns canonical stored-ref authorization and - # presigning. RunJob must therefore contain only ephemeral - # HTTP(S) Asset transport refs. Validate/materialize that whole - # payload before any source/model acquisition or tenant handler - # work; opaque refs and caller local_path values fail closed. + # th#886 v4: canonical payload keeps the caller's opaque stored + # refs; RunJob.input_assets is the ordered credential-free + # manifest. Validate/materialize the whole payload (one resolver + # POST for private refs) before any source/model acquisition or + # tenant handler work; manifest drift and caller local_path + # values fail closed. await _to_thread_complete( materialize_input_assets, payload, run.request_id, attempt=run.attempt, + manifest=manifest_from_run_job(run.input_assets), + file_base_url=self.file_base_url or "", + capability_token=run.capability_token or "", cancel_check=lambda: ctx.cancelled, ) ctx.raise_if_cancelled("canceled") diff --git a/src/gen_worker/input_assets.py b/src/gen_worker/input_assets.py index 2fc86edc..4561451d 100644 --- a/src/gen_worker/input_assets.py +++ b/src/gen_worker/input_assets.py @@ -1,27 +1,33 @@ -"""Fail-closed materialization of RunJob input ``Asset`` values. - -Tensorhub owns stored-ref authorization and presigning. The worker accepts only -the resulting HTTP(S) transport refs, downloads each distinct ref once into a -fresh worker-owned directory, and writes that local path back to every ordered -occurrence. Opaque stored refs and caller-provided filesystem paths never cross -this trust boundary. +"""Fail-closed materialization of RunJob input ``Asset`` values (v4). + +Canonical MessagePack keeps the caller's opaque stored refs. Tensorhub ships an +ordered, credential-free ``RunJob.input_assets`` manifest; the assigned worker +validates the payload's private refs against it, resolves fresh transport URLs +once per attempt via the worker capability, downloads and verifies the exact +attested bytes, and sets only worker-local ``local_path`` on every occurrence. +Caller HTTP(S) refs stay public transports and never enter the manifest. +Capabilities, URLs, opaque refs, and resolver bodies never enter errors/logs. """ from __future__ import annotations import hashlib +import json import logging import os import shutil import tempfile import threading +import urllib.error import urllib.parse import urllib.request import warnings from dataclasses import dataclass +from datetime import datetime from pathlib import Path -from typing import Any, Callable, Iterator +from typing import Any, Callable, Iterator, Mapping, Sequence +import blake3 as blake3_mod import msgspec from .api.errors import CanceledError, RetryableError, ValidationError @@ -32,9 +38,14 @@ _DEFAULT_MAX_BYTES = 50 << 20 # matches tensorhub's default media cap _DOWNLOAD_TIMEOUT_S = 120 +_RESOLVE_TIMEOUT_S = 30 +_MAX_RESOLVE_BODY = 8 << 20 _MAX_WALK_DEPTH = 32 _CHUNK = 1 << 20 _INPUT_DIR_PREFIX = "gen-worker-inputs-" +_RESOLVE_PATH = "/api/v1/worker/input-assets/resolve" +_MANIFEST_KINDS = ("media", "image", "video", "audio") +_HEX = set("0123456789abcdef") _EXT_BY_MIME = { "image/png": ".png", @@ -46,6 +57,50 @@ "audio/wav": ".wav", } +# (url, headers, body) -> (http_status, response_body). Injectable for tests. +ResolveTransport = Callable[[str, Mapping[str, str], bytes], tuple[int, bytes]] + + +class InputManifestEntry(msgspec.Struct, frozen=True): + """One dispatched ``RunJob.input_assets`` entry (immutable identity).""" + + asset_id: str + source_ref: str + blake3: str + size_bytes: int + kind: str + mime_type: str + + +class _ResolvedInputAsset(msgspec.Struct, forbid_unknown_fields=True): + asset_id: str + source_ref: str + blake3: str + size_bytes: int + kind: str + mime_type: str + url: str + url_expires_at: str + + +class _ResolveResponse(msgspec.Struct, forbid_unknown_fields=True): + assets: list[_ResolvedInputAsset] + + +def manifest_from_run_job(entries: Any) -> tuple[InputManifestEntry, ...]: + """Convert ``RunJob.input_assets`` proto rows to typed manifest entries.""" + return tuple( + InputManifestEntry( + asset_id=str(e.asset_id), + source_ref=str(e.source_ref), + blake3=str(e.blake3), + size_bytes=int(e.size_bytes), + kind=str(e.kind), + mime_type=str(e.mime_type), + ) + for e in entries or () + ) + @dataclass(frozen=True) class _AssetOccurrence: @@ -53,6 +108,13 @@ class _AssetOccurrence: asset: Asset +@dataclass +class _DownloadUnit: + occurrences: list[_AssetOccurrence] + entry: InputManifestEntry | None = None # None = public caller transport + url: str = "" # transport ref (public) or resolver-minted URL (private) + + @dataclass class _ScopeState: path: Path @@ -81,6 +143,10 @@ def inputs_dir_for_request(request_id: str, attempt: int = 0) -> Path: def _iter_asset_occurrences( obj: Any, path: str = "$", depth: int = 0 ) -> Iterator[_AssetOccurrence]: + """Deterministic traversal: msgspec struct declaration order, list/tuple + index, lexicographically sorted string map keys. Unordered containers and + non-string map keys are defensively rejected (build discovery already + refuses such schemas).""" if depth > _MAX_WALK_DEPTH: raise ValidationError("input_asset_payload_too_deep: nested Asset walk exceeded its limit") if isinstance(obj, Asset): @@ -92,11 +158,27 @@ def _iter_asset_occurrences( elif isinstance(obj, (list, tuple)): for index, item in enumerate(obj): yield from _iter_asset_occurrences(item, f"{path}[{index}]", depth + 1) + elif isinstance(obj, (set, frozenset)): + for item in obj: + for _ in _iter_asset_occurrences(item, f"{path}[*]", depth + 1): + raise ValidationError( + f"input_asset_unordered_container: {path} carries an Asset " + "inside an unordered set/frozenset" + ) elif isinstance(obj, dict): - # Dict keys may contain caller data. Use stable insertion positions in - # causal errors while preserving the value traversal order exactly. - for index, item in enumerate(obj.values()): - yield from _iter_asset_occurrences(item, f"{path}[{index}]", depth + 1) + # Map keys may contain caller data: error paths use sorted positions, + # never the keys themselves. + for index, key in enumerate(_sorted_string_keys(obj, path)): + yield from _iter_asset_occurrences(obj[key], f"{path}[{index}]", depth + 1) + + +def _sorted_string_keys(obj: dict, path: str) -> list[str]: + for key in obj: + if not isinstance(key, str): + raise ValidationError( + f"input_asset_map_key_invalid: {path} mapping key is not a string" + ) + return sorted(obj) def _iter_assets(obj: Any, depth: int = 0) -> Iterator[Asset]: @@ -105,24 +187,53 @@ def _iter_assets(obj: Any, depth: int = 0) -> Iterator[Asset]: yield occurrence.asset -def _transport_url(ref: str, path: str) -> str: - value = str(ref or "").strip() +def _declared_kind(asset: Asset) -> str: + if isinstance(asset, ImageAsset): + return "image" + if isinstance(asset, VideoAsset): + return "video" + if isinstance(asset, AudioAsset): + return "audio" + return "media" + + +def _classify_ref(ref: str, path: str) -> bool: + """True = private opaque stored ref; False = public HTTP(S) transport. + + Mirrors tensorhub's submit-time classifier so ``input_payload`` and the + manifest can never disagree: padded/empty refs and every non-HTTP scheme + are invalid. + """ + value = str(ref or "") + if not value or value != value.strip(): + raise ValidationError(f"invalid_input_asset_ref: {path} ref is empty or padded") try: parsed = urllib.parse.urlsplit(value) except ValueError: - raise ValidationError( - f"invalid_input_asset_url: {path} has an invalid transport URL" - ) from None + raise ValidationError(f"invalid_input_asset_ref: {path} ref is not parseable") from None scheme = parsed.scheme.lower() - if not scheme: + if scheme in ("http", "https"): + if not parsed.netloc or parsed.hostname is None: + raise ValidationError(f"invalid_input_asset_url: {path} has an invalid transport URL") + return False + if scheme or value.startswith("//"): raise ValidationError( - f"unresolved_input_asset: {path} requires an authorized HTTP(S) transport URL" + f"unsupported_input_asset_scheme: {path} requires HTTP(S) or an opaque stored ref" ) - if scheme not in ("http", "https"): - raise ValidationError(f"unsupported_input_asset_scheme: {path} requires HTTP(S)") - if not parsed.netloc or parsed.hostname is None: - raise ValidationError(f"invalid_input_asset_url: {path} has an invalid transport URL") - return value + return True + + +def _manifest_entry_valid(entry: InputManifestEntry) -> bool: + return ( + bool(entry.asset_id.strip()) + and bool(entry.source_ref) + and entry.source_ref == entry.source_ref.strip() + and entry.size_bytes > 0 + and entry.kind in _MANIFEST_KINDS + and bool(entry.mime_type.strip()) + and len(entry.blake3) == 64 + and set(entry.blake3) <= _HEX + ) def _mime_allowed(mime: str, allowed: tuple[str, ...]) -> bool: @@ -151,6 +262,12 @@ def _declared_kind_allows(asset: Asset, mime: str) -> bool: return True +def _entry_kind_allows(kind: str, mime: str) -> bool: + if kind == "media": + return True + return str(mime or "").lower().startswith(f"{kind}/") + + def _requires_image_decode(occurrences: list[_AssetOccurrence]) -> bool: return any( isinstance(occurrence.asset, ImageAsset) @@ -225,16 +342,142 @@ def _effective_size_cap(occurrences: list[_AssetOccurrence]) -> int: return min(caps) if caps else _DEFAULT_MAX_BYTES +def _default_resolve_transport( + url: str, headers: Mapping[str, str], body: bytes +) -> tuple[int, bytes]: + req = urllib.request.Request(url, data=body, headers=dict(headers), method="POST") + try: + with urllib.request.urlopen(req, timeout=_RESOLVE_TIMEOUT_S) as response: + return int(response.status), response.read(_MAX_RESOLVE_BODY + 1) + except urllib.error.HTTPError as err: + try: + payload = err.read(_MAX_RESOLVE_BODY + 1) + except Exception: + payload = b"" + return int(err.code), payload + + +def _resolve_private_urls( + manifest: Sequence[InputManifestEntry], + *, + request_id: str, + attempt: int, + file_base_url: str, + capability_token: str, + resolve_transport: ResolveTransport | None, + cancel_check: Callable[[], bool] | None, +) -> list[str]: + """Exactly one bounded resolver POST per attempt. Returns URLs aligned to + ``manifest`` order. Response identity must byte-match the dispatched + manifest; anything else is a platform malfunction (retryable).""" + _check_cancel(cancel_check) + if not file_base_url or not capability_token: + raise RetryableError( + "input_asset_resolution_unavailable: worker has no resolver credentials" + ) + transport = resolve_transport or _default_resolve_transport + url = file_base_url.rstrip("/") + _RESOLVE_PATH + body = json.dumps({"request_id": request_id, "attempt": int(attempt)}).encode() + headers = { + "Authorization": f"Bearer {capability_token}", + "Content-Type": "application/json", + } + try: + status, raw = transport(url, headers, body) + except CanceledError: + raise + except Exception: + raise RetryableError( + "input_asset_resolution_unavailable: resolver request failed" + ) from None + _check_cancel(cancel_check) + if status == 409: + # This attempt is no longer the assigned one (or its manifest moved): + # stop without publishing anything; the hub owns the successor. + raise CanceledError("canceled") + if status != 200: + raise RetryableError( + "input_asset_resolution_unavailable: resolver refused the request" + ) + if len(raw) > _MAX_RESOLVE_BODY: + raise RetryableError( + "input_asset_response_mismatch: resolver response exceeds its size bound" + ) + try: + decoded = msgspec.json.decode(raw, type=_ResolveResponse, strict=True) + except (msgspec.DecodeError, msgspec.ValidationError): + raise RetryableError( + "input_asset_response_mismatch: resolver response is malformed" + ) from None + if len(decoded.assets) != len(manifest): + raise RetryableError( + "input_asset_response_mismatch: resolver response does not match the manifest" + ) + urls: list[str] = [] + for entry, resolved in zip(manifest, decoded.assets): + if ( + resolved.asset_id != entry.asset_id + or resolved.source_ref != entry.source_ref + or resolved.blake3 != entry.blake3 + or resolved.size_bytes != entry.size_bytes + or resolved.kind != entry.kind + or resolved.mime_type != entry.mime_type + ): + raise RetryableError( + "input_asset_response_mismatch: resolver response does not match the manifest" + ) + if not resolved.url or not _parse_rfc3339(resolved.url_expires_at): + raise RetryableError( + "input_asset_response_mismatch: resolver response is malformed" + ) + urls.append(resolved.url) + return urls + + +def _parse_rfc3339(value: str) -> bool: + try: + datetime.fromisoformat(str(value or "")) + except ValueError: + return False + return True + + +def _validate_transport_url(url: str, path: str) -> None: + try: + parsed = urllib.parse.urlsplit(str(url or "")) + except ValueError: + raise ValidationError( + f"invalid_input_asset_url: {path} has an invalid transport URL" + ) from None + if parsed.scheme.lower() not in ("http", "https") or parsed.hostname is None: + raise ValidationError(f"invalid_input_asset_url: {path} has an invalid transport URL") + try: + blocked = _url_is_blocked(url) + except Exception: + raise RetryableError( + "input_asset_validation_failed: transport policy check failed" + ) from None + if blocked: + raise ValidationError(f"input_asset_url_not_allowed: {path} transport is blocked") + + def _download( - url: str, - occurrences: list[_AssetOccurrence], + unit: _DownloadUnit, dest_dir: Path, index: int, cancel_check: Callable[[], bool] | None, ) -> tuple[Path, int, str]: _check_cancel(cancel_check) + occurrences = unit.occurrences + entry = unit.entry cap = _effective_size_cap(occurrences) - req = urllib.request.Request(url, headers={"User-Agent": "gen-worker"}) + if entry is not None: + if entry.size_bytes > cap: + raise ValidationError( + f"input_asset_too_large: {occurrences[0].path} exceeds its byte cap" + ) + cap = entry.size_bytes + req = urllib.request.Request(unit.url, headers={"User-Agent": "gen-worker"}) try: fd, tmp_name = tempfile.mkstemp(prefix=f"input-{index}-", suffix=".part", dir=dest_dir) except OSError: @@ -244,6 +487,7 @@ def _download( tmp = Path(tmp_name) total = 0 head = b"" + hasher = blake3_mod.blake3() if entry is not None else None try: try: with urllib.request.urlopen(req, timeout=_DOWNLOAD_TIMEOUT_S) as response: @@ -264,9 +508,16 @@ def _download( head = chunk[:64] total += len(chunk) if total > cap: + if entry is not None: + raise RetryableError( + "input_asset_integrity_failed: " + f"{occurrences[0].path} bytes exceed their attested size" + ) raise ValidationError( f"input_asset_too_large: {occurrences[0].path} exceeds its byte cap" ) + if hasher is not None: + hasher.update(chunk) output.write(chunk) if expected_length and total != expected_length: raise RetryableError( @@ -280,10 +531,27 @@ def _download( ) from None _check_cancel(cancel_check) - mime = _infer_mime_type(url, head) + if entry is not None: + if total != entry.size_bytes: + raise RetryableError( + f"input_asset_integrity_failed: {occurrences[0].path} bytes do not " + "match their attested size" + ) + assert hasher is not None + if hasher.hexdigest() != entry.blake3: + raise RetryableError( + f"input_asset_integrity_failed: {occurrences[0].path} bytes do not " + "match their attested BLAKE3" + ) + mime = _infer_mime_type(unit.url, head) if _requires_image_decode(occurrences): mime = _decode_image(tmp, occurrences) _check_cancel(cancel_check) + if entry is not None and not _entry_kind_allows(entry.kind, mime): + raise ValidationError( + f"input_asset_kind_mismatch: {occurrences[0].path} content is not its " + "manifest media kind" + ) for occurrence in occurrences: asset = occurrence.asset if not _declared_kind_allows(asset, mime): @@ -325,7 +593,11 @@ def _remove_scope_state(state: _ScopeState) -> None: def cleanup_input_assets(request_id: str, attempt: int = 0) -> None: - """Remove one attempt's materialized inputs and clear every assigned path.""" + """Remove one attempt's materialized inputs and clear every assigned path. + + Scopes are keyed by (request_id, attempt): cleaning attempt N+1 can never + remove attempt N's directory, and vice versa. + """ key = _scope_key(request_id, attempt) with _scope_lock: state = _scope_states.pop(key, None) @@ -354,19 +626,70 @@ def _start_scope(request_id: str, attempt: int, assets: tuple[Asset, ...]) -> _S return state +def _collect_units( + occurrences: list[_AssetOccurrence], + manifest: Sequence[InputManifestEntry], +) -> list[_DownloadUnit]: + """Classify occurrences, fence the private sequence against the dispatched + manifest, and return download units in payload first-occurrence order.""" + units: list[_DownloadUnit] = [] + public_by_url: dict[str, _DownloadUnit] = {} + private_by_key: dict[tuple[str, str], _DownloadUnit] = {} + for occurrence in occurrences: + ref = str(occurrence.asset.ref or "") + if _classify_ref(ref, occurrence.path): + key = (ref, _declared_kind(occurrence.asset)) + unit = private_by_key.get(key) + if unit is None: + unit = _DownloadUnit(occurrences=[]) + private_by_key[key] = unit + units.append(unit) + unit.occurrences.append(occurrence) + else: + unit = public_by_url.get(ref) + if unit is None: + unit = _DownloadUnit(occurrences=[], url=ref) + public_by_url[ref] = unit + units.append(unit) + unit.occurrences.append(occurrence) + + # The ordered unique private sequence must equal the dispatched manifest + # exactly (count, order, ref, kind) with valid immutable metadata BEFORE + # any resolver call or GET. A mismatch is a non-retryable contract breach. + private_sequence = list(private_by_key) + manifest_sequence = [(entry.source_ref, entry.kind) for entry in manifest] + if private_sequence != manifest_sequence: + raise ValidationError( + "input_asset_manifest_mismatch: dispatched manifest does not match " + "the payload's private inputs" + ) + for entry, key in zip(manifest, private_sequence): + if not _manifest_entry_valid(entry): + raise ValidationError( + "input_asset_manifest_invalid: dispatched manifest entry is malformed" + ) + private_by_key[key].entry = entry + return units + + def materialize_input_assets( payload: Any, request_id: str, *, attempt: int = 0, + manifest: Sequence[InputManifestEntry] = (), + file_base_url: str = "", + capability_token: str = "", cancel_check: Callable[[], bool] | None = None, + resolve_transport: ResolveTransport | None = None, ) -> int: """Materialize all ordered input assets transactionally. - The return value is the number of distinct HTTP(S) downloads. Duplicate - occurrences retain their payload positions and share one worker-owned - path. Any failure clears assigned paths and removes the whole attempt - directory before propagating a stable causal error. + Returns the number of distinct downloads. Duplicate occurrences retain + their payload positions and share one worker-owned path. Every decoded + Asset field is preserved — only ``local_path`` is assigned. Any failure + clears assigned paths and removes the whole attempt directory before + propagating a stable causal error. """ cleanup_input_assets(request_id, attempt) occurrences = list(_iter_asset_occurrences(payload)) @@ -375,28 +698,44 @@ def materialize_input_assets( # use it as evidence that the asset was materialized. occurrence.asset.local_path = None _check_cancel(cancel_check) + manifest = tuple(manifest) if not occurrences: + if manifest: + raise ValidationError( + "input_asset_manifest_mismatch: dispatched manifest does not match " + "the payload's private inputs" + ) return 0 - grouped: dict[str, list[_AssetOccurrence]] = {} - for occurrence in occurrences: - url = _transport_url(occurrence.asset.ref, occurrence.path) - grouped.setdefault(url, []).append(occurrence) + units = _collect_units(occurrences, manifest) + private_units = [unit for unit in units if unit.entry is not None] + if private_units and int(attempt) <= 0: + raise ValidationError( + "input_asset_attempt_invalid: private inputs require a positive attempt" + ) - # Validate the complete transport set before downloading the first item, - # so a later opaque/blocked ref cannot cause partial materialization work. - for occurrences_for_url in grouped.values(): - url = str(occurrences_for_url[0].asset.ref).strip() - try: - blocked = _url_is_blocked(url) - except Exception: - raise RetryableError( - "input_asset_validation_failed: transport policy check failed" - ) from None - if blocked: - raise ValidationError( - f"input_asset_url_not_allowed: {occurrences_for_url[0].path} transport is blocked" - ) + # Validate the complete public transport set before the resolver call and + # before downloading the first item, so a later blocked ref cannot cause + # partial materialization work. + for unit in units: + if unit.entry is None: + _validate_transport_url(unit.url, unit.occurrences[0].path) + + if private_units: + urls = _resolve_private_urls( + manifest, + request_id=request_id, + attempt=int(attempt), + file_base_url=file_base_url, + capability_token=capability_token, + resolve_transport=resolve_transport, + cancel_check=cancel_check, + ) + by_entry = dict(zip(manifest, urls)) + for unit in private_units: + assert unit.entry is not None + unit.url = by_entry[unit.entry] + _validate_transport_url(unit.url, unit.occurrences[0].path) state = _start_scope( request_id, @@ -404,30 +743,28 @@ def materialize_input_assets( tuple(occurrence.asset for occurrence in occurrences), ) try: - for index, (url, duplicates) in enumerate(grouped.items()): - local_path, size_bytes, mime_type = _download( - url, duplicates, state.path, index, cancel_check - ) - for occurrence in duplicates: - asset = occurrence.asset - asset.local_path = str(local_path) - asset.size_bytes = size_bytes - asset.mime_type = mime_type + for index, unit in enumerate(units): + local_path, size_bytes, mime_type = _download(unit, state.path, index, cancel_check) + for occurrence in unit.occurrences: + occurrence.asset.local_path = str(local_path) logger.info( - "materialized input asset index=%d aliases=%d bytes=%d mime=%s", + "materialized input asset index=%d aliases=%d bytes=%d mime=%s private=%s", index, - len(duplicates), + len(unit.occurrences), size_bytes, mime_type, + unit.entry is not None, ) - return len(grouped) + return len(units) except BaseException: cleanup_input_assets(request_id, attempt) raise __all__ = [ + "InputManifestEntry", "cleanup_input_assets", "inputs_dir_for_request", + "manifest_from_run_job", "materialize_input_assets", ] diff --git a/src/gen_worker/pb/worker_scheduler_pb2.py b/src/gen_worker/pb/worker_scheduler_pb2.py index 0e83ec1f..3677287c 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.py +++ b/src/gen_worker/pb/worker_scheduler_pb2.py @@ -24,7 +24,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xe6\x03\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x12\x39\n\x0f\x61\x63tivity_update\x18\t \x01(\x0b\x32\x1e.cozy.scheduler.ActivityUpdateH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xc7\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\x12\x1d\n\x15heartbeat_interval_ms\x18\x08 \x01(\x04\"\x9b\x02\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xdd\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\x12;\n\x11\x64\x65sired_residency\x18\x05 \x01(\x0b\x32 .cozy.scheduler.DesiredResidency\"\xf7\x01\n\x10\x44\x65siredResidency\x12\x12\n\ngeneration\x18\x01 \x01(\x04\x12\x11\n\tdisk_refs\x18\x02 \x03(\t\x12,\n\x03hot\x18\x03 \x03(\x0b\x32\x1f.cozy.scheduler.DesiredInstance\x12\x42\n\tsnapshots\x18\x04 \x03(\x0b\x32/.cozy.scheduler.DesiredResidency.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"V\n\x0f\x44\x65siredInstance\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\x06models\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\"P\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\x12\x0c\n\x04lane\x18\x04 \x01(\t\"\xe8\x02\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\x12%\n\x1dobserved_residency_generation\x18\x06 \x01(\x04\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\x12\x33\n\ndisk_usage\x18\t \x01(\x0b\x32\x1f.cozy.scheduler.DiskUsageReport\"\xa9\x01\n\x10StorageTierUsage\x12)\n\x04tier\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.StorageTier\x12\x12\n\nmount_path\x18\x02 \x01(\t\x12\x13\n\x0btotal_bytes\x18\x03 \x01(\x03\x12\x12\n\nfree_bytes\x18\x04 \x01(\x03\x12\x12\n\nused_bytes\x18\x05 \x01(\x03\x12\x19\n\x11reclaimable_bytes\x18\x06 \x01(\x03\"_\n\x0f\x44iskUsageReport\x12/\n\x05tiers\x18\x01 \x03(\x0b\x32 .cozy.scheduler.StorageTierUsage\x12\x1b\n\x13\x63\x61pacity_generation\x18\x02 \x01(\x04\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\x96\x04\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\x12\x0c\n\x04lane\x18\x0e \x01(\t\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"q\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xc2\x02\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\x12\x0c\n\x04lane\x18\r \x01(\t\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xa0\x01\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x9b\x04\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\x12\x15\n\rnetwork_bytes\x18\x14 \x01(\x03\"\xc6\x01\n\x0e\x41\x63tivityUpdate\x12\x0c\n\x04kind\x18\x01 \x01(\t\x12\r\n\x05phase\x18\x02 \x01(\t\x12\x0c\n\x04step\x18\x03 \x01(\x03\x12\x13\n\x0btotal_steps\x18\x04 \x01(\x03\x12\x0b\n\x03seq\x18\x05 \x01(\x04\x12,\n\x05state\x18\x06 \x01(\x0e\x32\x1d.cozy.scheduler.ActivityState\x12\r\n\x05\x65rror\x18\x07 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x08 \x01(\t\x12\x1a\n\x12updated_at_unix_ms\x18\t \x01(\x03\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x03*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*v\n\x0bStorageTier\x12\x1c\n\x18STORAGE_TIER_UNSPECIFIED\x10\x00\x12\x1a\n\x16STORAGE_TIER_CONTAINER\x10\x01\x12\x17\n\x13STORAGE_TIER_VOLUME\x10\x02\x12\x14\n\x10STORAGE_TIER_NFS\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa7\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04\"\x04\x08\x01\x10\x01\"\x04\x08\x02\x10\x02\"\x04\x08\x03\x10\x03*\x16MODEL_OP_KIND_DOWNLOAD*\x12MODEL_OP_KIND_LOAD*\x14MODEL_OP_KIND_UNLOAD*\x82\x02\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08*\x84\x01\n\rActivityState\x12\x1e\n\x1a\x41\x43TIVITY_STATE_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x41\x43TIVITY_STATE_RUNNING\x10\x01\x12\x1c\n\x18\x41\x43TIVITY_STATE_COMPLETED\x10\x02\x12\x19\n\x15\x41\x43TIVITY_STATE_FAILED\x10\x03\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xe6\x03\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x12\x39\n\x0f\x61\x63tivity_update\x18\t \x01(\x0b\x32\x1e.cozy.scheduler.ActivityUpdateH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xc7\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\x12\x1d\n\x15heartbeat_interval_ms\x18\x08 \x01(\x04\"\x9b\x02\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xdd\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\x12;\n\x11\x64\x65sired_residency\x18\x05 \x01(\x0b\x32 .cozy.scheduler.DesiredResidency\"\xf7\x01\n\x10\x44\x65siredResidency\x12\x12\n\ngeneration\x18\x01 \x01(\x04\x12\x11\n\tdisk_refs\x18\x02 \x03(\t\x12,\n\x03hot\x18\x03 \x03(\x0b\x32\x1f.cozy.scheduler.DesiredInstance\x12\x42\n\tsnapshots\x18\x04 \x03(\x0b\x32/.cozy.scheduler.DesiredResidency.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"V\n\x0f\x44\x65siredInstance\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\x06models\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\"P\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\x12\x0c\n\x04lane\x18\x04 \x01(\t\"\xe8\x02\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\x12%\n\x1dobserved_residency_generation\x18\x06 \x01(\x04\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\x12\x33\n\ndisk_usage\x18\t \x01(\x0b\x32\x1f.cozy.scheduler.DiskUsageReport\"\xa9\x01\n\x10StorageTierUsage\x12)\n\x04tier\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.StorageTier\x12\x12\n\nmount_path\x18\x02 \x01(\t\x12\x13\n\x0btotal_bytes\x18\x03 \x01(\x03\x12\x12\n\nfree_bytes\x18\x04 \x01(\x03\x12\x12\n\nused_bytes\x18\x05 \x01(\x03\x12\x19\n\x11reclaimable_bytes\x18\x06 \x01(\x03\"_\n\x0f\x44iskUsageReport\x12/\n\x05tiers\x18\x01 \x03(\x0b\x32 .cozy.scheduler.StorageTierUsage\x12\x1b\n\x13\x63\x61pacity_generation\x18\x02 \x01(\x04\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\xc8\x04\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\x12\x0c\n\x04lane\x18\x0e \x01(\t\x12\x30\n\x0cinput_assets\x18\x0f \x03(\x0b\x32\x1a.cozy.scheduler.InputAsset\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"w\n\nInputAsset\x12\x10\n\x08\x61sset_id\x18\x01 \x01(\t\x12\x12\n\nsource_ref\x18\x02 \x01(\t\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x12\n\nsize_bytes\x18\x04 \x01(\x03\x12\x0c\n\x04kind\x18\x05 \x01(\t\x12\x11\n\tmime_type\x18\x06 \x01(\t\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"q\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xc2\x02\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\x12\x0c\n\x04lane\x18\r \x01(\t\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xa0\x01\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x9b\x04\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\x12\x15\n\rnetwork_bytes\x18\x14 \x01(\x03\"\xc6\x01\n\x0e\x41\x63tivityUpdate\x12\x0c\n\x04kind\x18\x01 \x01(\t\x12\r\n\x05phase\x18\x02 \x01(\t\x12\x0c\n\x04step\x18\x03 \x01(\x03\x12\x13\n\x0btotal_steps\x18\x04 \x01(\x03\x12\x0b\n\x03seq\x18\x05 \x01(\x04\x12,\n\x05state\x18\x06 \x01(\x0e\x32\x1d.cozy.scheduler.ActivityState\x12\r\n\x05\x65rror\x18\x07 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x08 \x01(\t\x12\x1a\n\x12updated_at_unix_ms\x18\t \x01(\x03\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x04*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*v\n\x0bStorageTier\x12\x1c\n\x18STORAGE_TIER_UNSPECIFIED\x10\x00\x12\x1a\n\x16STORAGE_TIER_CONTAINER\x10\x01\x12\x17\n\x13STORAGE_TIER_VOLUME\x10\x02\x12\x14\n\x10STORAGE_TIER_NFS\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa7\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04\"\x04\x08\x01\x10\x01\"\x04\x08\x02\x10\x02\"\x04\x08\x03\x10\x03*\x16MODEL_OP_KIND_DOWNLOAD*\x12MODEL_OP_KIND_LOAD*\x14MODEL_OP_KIND_UNLOAD*\x82\x02\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08*\x84\x01\n\rActivityState\x12\x1e\n\x1a\x41\x43TIVITY_STATE_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x41\x43TIVITY_STATE_RUNNING\x10\x01\x12\x1c\n\x18\x41\x43TIVITY_STATE_COMPLETED\x10\x02\x12\x19\n\x15\x41\x43TIVITY_STATE_FAILED\x10\x03\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -40,24 +40,24 @@ _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_options = b'8\001' _globals['_FNUNAVAILABLE_AXESENTRY']._loaded_options = None _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_options = b'8\001' - _globals['_PROTOCOLVERSION']._serialized_start=6865 - _globals['_PROTOCOLVERSION']._serialized_end=6946 - _globals['_RESIDENCYTIER']._serialized_start=6948 - _globals['_RESIDENCYTIER']._serialized_end=7069 - _globals['_STORAGETIER']._serialized_start=7071 - _globals['_STORAGETIER']._serialized_end=7189 - _globals['_WORKERPHASE']._serialized_start=7192 - _globals['_WORKERPHASE']._serialized_end=7408 - _globals['_OUTPUTMODE']._serialized_start=7410 - _globals['_OUTPUTMODE']._serialized_end=7496 - _globals['_JOBSTATUS']._serialized_start=7499 - _globals['_JOBSTATUS']._serialized_end=7654 - _globals['_MODELOPKIND']._serialized_start=7657 - _globals['_MODELOPKIND']._serialized_end=7824 - _globals['_MODELSTATE']._serialized_start=7827 - _globals['_MODELSTATE']._serialized_end=8085 - _globals['_ACTIVITYSTATE']._serialized_start=8088 - _globals['_ACTIVITYSTATE']._serialized_end=8220 + _globals['_PROTOCOLVERSION']._serialized_start=7036 + _globals['_PROTOCOLVERSION']._serialized_end=7117 + _globals['_RESIDENCYTIER']._serialized_start=7119 + _globals['_RESIDENCYTIER']._serialized_end=7240 + _globals['_STORAGETIER']._serialized_start=7242 + _globals['_STORAGETIER']._serialized_end=7360 + _globals['_WORKERPHASE']._serialized_start=7363 + _globals['_WORKERPHASE']._serialized_end=7579 + _globals['_OUTPUTMODE']._serialized_start=7581 + _globals['_OUTPUTMODE']._serialized_end=7667 + _globals['_JOBSTATUS']._serialized_start=7670 + _globals['_JOBSTATUS']._serialized_end=7825 + _globals['_MODELOPKIND']._serialized_start=7828 + _globals['_MODELOPKIND']._serialized_end=7995 + _globals['_MODELSTATE']._serialized_start=7998 + _globals['_MODELSTATE']._serialized_end=8256 + _globals['_ACTIVITYSTATE']._serialized_start=8259 + _globals['_ACTIVITYSTATE']._serialized_end=8391 _globals['_WORKERMESSAGE']._serialized_start=43 _globals['_WORKERMESSAGE']._serialized_end=529 _globals['_SCHEDULERMESSAGE']._serialized_start=532 @@ -97,47 +97,49 @@ _globals['_COMPILETARGETBINDING']._serialized_start=3643 _globals['_COMPILETARGETBINDING']._serialized_end=3717 _globals['_RUNJOB']._serialized_start=3720 - _globals['_RUNJOB']._serialized_end=4254 + _globals['_RUNJOB']._serialized_end=4304 _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_start=2260 _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_end=2334 - _globals['_REQUIREDCOMPILEEXECUTION']._serialized_start=4257 - _globals['_REQUIREDCOMPILEEXECUTION']._serialized_end=4387 - _globals['_RESOLVEDCOMPUTE']._serialized_start=4389 - _globals['_RESOLVEDCOMPUTE']._serialized_end=4478 - _globals['_MODELBINDING']._serialized_start=4480 - _globals['_MODELBINDING']._serialized_end=4593 - _globals['_LORAOVERLAY']._serialized_start=4595 - _globals['_LORAOVERLAY']._serialized_end=4665 - _globals['_SNAPSHOT']._serialized_start=4667 - _globals['_SNAPSHOT']._serialized_end=4738 - _globals['_SNAPSHOTFILE']._serialized_start=4740 - _globals['_SNAPSHOTFILE']._serialized_end=4817 - _globals['_JOBACCEPTED']._serialized_start=4819 - _globals['_JOBACCEPTED']._serialized_end=4869 - _globals['_JOBRESULT']._serialized_start=4872 - _globals['_JOBRESULT']._serialized_end=5078 - _globals['_JOBMETRICS']._serialized_start=5081 - _globals['_JOBMETRICS']._serialized_end=5403 - _globals['_JOBPROGRESS']._serialized_start=5405 - _globals['_JOBPROGRESS']._serialized_end=5504 - _globals['_CANCELJOB']._serialized_start=5506 - _globals['_CANCELJOB']._serialized_end=5554 - _globals['_MODELOP']._serialized_start=5557 - _globals['_MODELOP']._serialized_end=5717 - _globals['_MODELEVENT']._serialized_start=5720 - _globals['_MODELEVENT']._serialized_end=6259 - _globals['_ACTIVITYUPDATE']._serialized_start=6262 - _globals['_ACTIVITYUPDATE']._serialized_end=6460 - _globals['_FNUNAVAILABLE']._serialized_start=6463 - _globals['_FNUNAVAILABLE']._serialized_end=6633 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=6590 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=6633 - _globals['_FNDEGRADED']._serialized_start=6636 - _globals['_FNDEGRADED']._serialized_end=6777 - _globals['_DRAIN']._serialized_start=6779 - _globals['_DRAIN']._serialized_end=6807 - _globals['_TOKENREFRESH']._serialized_start=6809 - _globals['_TOKENREFRESH']._serialized_end=6863 - _globals['_WORKERSCHEDULER']._serialized_start=8222 - _globals['_WORKERSCHEDULER']._serialized_end=8319 + _globals['_INPUTASSET']._serialized_start=4306 + _globals['_INPUTASSET']._serialized_end=4425 + _globals['_REQUIREDCOMPILEEXECUTION']._serialized_start=4428 + _globals['_REQUIREDCOMPILEEXECUTION']._serialized_end=4558 + _globals['_RESOLVEDCOMPUTE']._serialized_start=4560 + _globals['_RESOLVEDCOMPUTE']._serialized_end=4649 + _globals['_MODELBINDING']._serialized_start=4651 + _globals['_MODELBINDING']._serialized_end=4764 + _globals['_LORAOVERLAY']._serialized_start=4766 + _globals['_LORAOVERLAY']._serialized_end=4836 + _globals['_SNAPSHOT']._serialized_start=4838 + _globals['_SNAPSHOT']._serialized_end=4909 + _globals['_SNAPSHOTFILE']._serialized_start=4911 + _globals['_SNAPSHOTFILE']._serialized_end=4988 + _globals['_JOBACCEPTED']._serialized_start=4990 + _globals['_JOBACCEPTED']._serialized_end=5040 + _globals['_JOBRESULT']._serialized_start=5043 + _globals['_JOBRESULT']._serialized_end=5249 + _globals['_JOBMETRICS']._serialized_start=5252 + _globals['_JOBMETRICS']._serialized_end=5574 + _globals['_JOBPROGRESS']._serialized_start=5576 + _globals['_JOBPROGRESS']._serialized_end=5675 + _globals['_CANCELJOB']._serialized_start=5677 + _globals['_CANCELJOB']._serialized_end=5725 + _globals['_MODELOP']._serialized_start=5728 + _globals['_MODELOP']._serialized_end=5888 + _globals['_MODELEVENT']._serialized_start=5891 + _globals['_MODELEVENT']._serialized_end=6430 + _globals['_ACTIVITYUPDATE']._serialized_start=6433 + _globals['_ACTIVITYUPDATE']._serialized_end=6631 + _globals['_FNUNAVAILABLE']._serialized_start=6634 + _globals['_FNUNAVAILABLE']._serialized_end=6804 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=6761 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=6804 + _globals['_FNDEGRADED']._serialized_start=6807 + _globals['_FNDEGRADED']._serialized_end=6948 + _globals['_DRAIN']._serialized_start=6950 + _globals['_DRAIN']._serialized_end=6978 + _globals['_TOKENREFRESH']._serialized_start=6980 + _globals['_TOKENREFRESH']._serialized_end=7034 + _globals['_WORKERSCHEDULER']._serialized_start=8393 + _globals['_WORKERSCHEDULER']._serialized_end=8490 # @@protoc_insertion_point(module_scope) diff --git a/src/gen_worker/pb/worker_scheduler_pb2.pyi b/src/gen_worker/pb/worker_scheduler_pb2.pyi index 67fc0690..876cfb66 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.pyi +++ b/src/gen_worker/pb/worker_scheduler_pb2.pyi @@ -395,7 +395,7 @@ class CompileTargetBinding(_message.Message): def __init__(self, slot: _Optional[str] = ..., ref: _Optional[str] = ..., snapshot_digest: _Optional[str] = ...) -> None: ... class RunJob(_message.Message): - __slots__ = ("request_id", "attempt", "function_name", "input_payload", "timeout_ms", "tenant", "invoker_id", "capability_token", "output_mode", "compute", "models", "snapshots", "required_compile", "lane") + __slots__ = ("request_id", "attempt", "function_name", "input_payload", "timeout_ms", "tenant", "invoker_id", "capability_token", "output_mode", "compute", "models", "snapshots", "required_compile", "lane", "input_assets") class SnapshotsEntry(_message.Message): __slots__ = ("key", "value") KEY_FIELD_NUMBER: _ClassVar[int] @@ -417,6 +417,7 @@ class RunJob(_message.Message): SNAPSHOTS_FIELD_NUMBER: _ClassVar[int] REQUIRED_COMPILE_FIELD_NUMBER: _ClassVar[int] LANE_FIELD_NUMBER: _ClassVar[int] + INPUT_ASSETS_FIELD_NUMBER: _ClassVar[int] request_id: str attempt: int function_name: str @@ -431,7 +432,24 @@ class RunJob(_message.Message): snapshots: _containers.MessageMap[str, Snapshot] required_compile: RequiredCompileExecution lane: str - def __init__(self, request_id: _Optional[str] = ..., attempt: _Optional[int] = ..., function_name: _Optional[str] = ..., input_payload: _Optional[bytes] = ..., timeout_ms: _Optional[int] = ..., tenant: _Optional[str] = ..., invoker_id: _Optional[str] = ..., capability_token: _Optional[str] = ..., output_mode: _Optional[_Union[OutputMode, str]] = ..., compute: _Optional[_Union[ResolvedCompute, _Mapping]] = ..., models: _Optional[_Iterable[_Union[ModelBinding, _Mapping]]] = ..., snapshots: _Optional[_Mapping[str, Snapshot]] = ..., required_compile: _Optional[_Union[RequiredCompileExecution, _Mapping]] = ..., lane: _Optional[str] = ...) -> None: ... + input_assets: _containers.RepeatedCompositeFieldContainer[InputAsset] + def __init__(self, request_id: _Optional[str] = ..., attempt: _Optional[int] = ..., function_name: _Optional[str] = ..., input_payload: _Optional[bytes] = ..., timeout_ms: _Optional[int] = ..., tenant: _Optional[str] = ..., invoker_id: _Optional[str] = ..., capability_token: _Optional[str] = ..., output_mode: _Optional[_Union[OutputMode, str]] = ..., compute: _Optional[_Union[ResolvedCompute, _Mapping]] = ..., models: _Optional[_Iterable[_Union[ModelBinding, _Mapping]]] = ..., snapshots: _Optional[_Mapping[str, Snapshot]] = ..., required_compile: _Optional[_Union[RequiredCompileExecution, _Mapping]] = ..., lane: _Optional[str] = ..., input_assets: _Optional[_Iterable[_Union[InputAsset, _Mapping]]] = ...) -> None: ... + +class InputAsset(_message.Message): + __slots__ = ("asset_id", "source_ref", "blake3", "size_bytes", "kind", "mime_type") + ASSET_ID_FIELD_NUMBER: _ClassVar[int] + SOURCE_REF_FIELD_NUMBER: _ClassVar[int] + BLAKE3_FIELD_NUMBER: _ClassVar[int] + SIZE_BYTES_FIELD_NUMBER: _ClassVar[int] + KIND_FIELD_NUMBER: _ClassVar[int] + MIME_TYPE_FIELD_NUMBER: _ClassVar[int] + asset_id: str + source_ref: str + blake3: str + size_bytes: int + kind: str + mime_type: str + def __init__(self, asset_id: _Optional[str] = ..., source_ref: _Optional[str] = ..., blake3: _Optional[str] = ..., size_bytes: _Optional[int] = ..., kind: _Optional[str] = ..., mime_type: _Optional[str] = ...) -> None: ... class RequiredCompileExecution(_message.Message): __slots__ = ("target_incarnation_id", "cell_ref", "cell_snapshot_digest", "contract_digest") diff --git a/tests/test_input_manifest_v4.py b/tests/test_input_manifest_v4.py new file mode 100644 index 00000000..7b12d7bf --- /dev/null +++ b/tests/test_input_manifest_v4.py @@ -0,0 +1,1031 @@ +"""gw#585: Tensorhub v4 private-input manifests (th#886 hard cut). + +Hermetic: the resolver is a fake injected transport; downloads hit a local +loopback HTTP server only. Covers deterministic traversal, manifest fencing, +the single strict resolver POST, response strict-decoding, byte verification, +attempt-scoped cleanup/result fencing, discovery build rejection, and the v4 +proto descriptors. +""" + +from __future__ import annotations + +import asyncio +import base64 +import http.server +import json +import threading +import urllib.parse +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import blake3 +import msgspec +import pytest + +from gen_worker import input_assets +from gen_worker.api.errors import CanceledError, RetryableError, ValidationError +from gen_worker.api.types import Asset, AudioAsset, ImageAsset, MediaAsset, VideoAsset +from gen_worker.discovery.discover import _collect_payload_moderation_metadata +from gen_worker.executor import Executor +from gen_worker.input_assets import ( + InputManifestEntry, + cleanup_input_assets, + inputs_dir_for_request, + manifest_from_run_job, + materialize_input_assets, +) +from gen_worker.pb import worker_scheduler_pb2 as pb +from gen_worker.registry import EndpointSpec + +PNG_A = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAYAAAD0In+KAAAAEUlEQVR4nGP8z8Dwn4GBgQEADQUCAOAHawIAAAAASUVORK5CYII=" +) +PNG_B = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAYAAAD0In+KAAAAEUlEQVR4nGNkYPj/n4GBgQEACwcCAKXJIuMAAAAASUVORK5CYII=" +) +CORRUPT_PNG = b"\x89PNG\r\n\x1a\n" + b"not-a-decodable-image" +MP3 = b"ID3" + b"audio" * 16 +_SERVER_CHUNK = 1 << 20 +SLOW_BODY = PNG_A + b"x" * (2 * _SERVER_CHUNK) + +BASE_URL = "http://hub.internal.test" +CAPABILITY = "cap-token-SECRET" +EXPIRES = "2099-01-01T00:00:00+00:00" + +REF_A = "tenant/private/a" +REF_B = "tenant/private/b" +REF_C = "tenant/private/audio" + + +def b3(data: bytes) -> str: + return blake3.blake3(data).hexdigest() + + +def make_entry(ref: str, body: bytes, kind: str, mime: str, **overrides: Any) -> InputManifestEntry: + fields: dict[str, Any] = { + "asset_id": f"asset-{ref.rsplit('/', 1)[-1]}", + "source_ref": ref, + "blake3": b3(body), + "size_bytes": len(body), + "kind": kind, + "mime_type": mime, + } + fields.update(overrides) + return InputManifestEntry(**fields) + + +@dataclass +class FakeResolver: + """Injected resolver transport: records calls, serves a canned response.""" + + manifest: tuple[InputManifestEntry, ...] = () + urls: dict[str, str] = field(default_factory=dict) # source_ref -> url + status: int = 200 + raw_body: bytes | None = None + mutate: Any = None # callable(list[dict]) -> list[dict] + exc: Exception | None = None + on_call: Any = None + calls: list[tuple[str, dict[str, str], bytes]] = field(default_factory=list) + + def __call__(self, url: str, headers: Any, body: bytes) -> tuple[int, bytes]: + self.calls.append((url, dict(headers), bytes(body))) + if self.on_call is not None: + self.on_call() + if self.exc is not None: + raise self.exc + if self.raw_body is not None: + return self.status, self.raw_body + if self.status != 200: + return self.status, b'{"error":"resolver-error-body"}' + assets = [ + { + "asset_id": e.asset_id, + "source_ref": e.source_ref, + "blake3": e.blake3, + "size_bytes": e.size_bytes, + "kind": e.kind, + "mime_type": e.mime_type, + "url": self.urls[e.source_ref], + "url_expires_at": EXPIRES, + } + for e in self.manifest + ] + if self.mutate is not None: + assets = self.mutate(assets) + return 200, json.dumps({"assets": assets}).encode() + + +@dataclass +class HTTPRoot: + base_url: str + hits: dict[str, int] + slow_started: threading.Event + slow_release: threading.Event + + def url(self, name: str) -> str: + return f"{self.base_url}/{name}" + + +@pytest.fixture() +def http_root() -> Any: + hits: dict[str, int] = {} + slow_started = threading.Event() + slow_release = threading.Event() + bodies = { + "/a.png": ("image/png", PNG_A), + "/b.png": ("image/png", PNG_B), + "/audio.mp3": ("audio/mpeg", MP3), + "/corrupt.png": ("image/png", CORRUPT_PNG), + } + + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, *_args: Any) -> None: + pass + + def do_GET(self) -> None: # noqa: N802 - stdlib handler API + path = urllib.parse.urlsplit(self.path).path + hits[path] = hits.get(path, 0) + 1 + if path == "/slow.bin": + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(SLOW_BODY))) + self.end_headers() + try: + self.wfile.write(SLOW_BODY[:_SERVER_CHUNK]) + self.wfile.flush() + slow_started.set() + slow_release.wait(timeout=5) + self.wfile.write(SLOW_BODY[_SERVER_CHUNK:]) + except (BrokenPipeError, ConnectionResetError): + pass + return + content = bodies.get(path) + if content is None: + self.send_error(404) + return + mime, body = content + self.send_response(200) + self.send_header("Content-Type", mime) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + root = HTTPRoot( + base_url=f"http://127.0.0.1:{server.server_address[1]}", + hits=hits, + slow_started=slow_started, + slow_release=slow_release, + ) + yield root + slow_release.set() + server.shutdown() + thread.join(timeout=5) + + +@pytest.fixture(autouse=True) +def allow_localhost(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(input_assets, "_url_is_blocked", lambda _url: False) + + +class Payload(msgspec.Struct): + image: ImageAsset | None = None + prompt: str = "" + extra: list[Asset] = [] + audio: AudioAsset | None = None + mapped: dict[str, Asset] = {} + + +def materialize( + payload: Any, + request_id: str, + *, + attempt: int = 1, + manifest: tuple[InputManifestEntry, ...] = (), + resolver: FakeResolver | None = None, + cancel_check: Any = None, +) -> int: + return materialize_input_assets( + payload, + request_id, + attempt=attempt, + manifest=manifest, + file_base_url=BASE_URL, + capability_token=CAPABILITY, + cancel_check=cancel_check, + resolve_transport=resolver, + ) + + +def assert_no_leak(exc: BaseException, *secrets: str) -> None: + text = str(exc) + for secret in (CAPABILITY, BASE_URL, *secrets): + assert secret not in text + + +# --------------------------------------------------------------------------- +# Deterministic traversal + classification +# --------------------------------------------------------------------------- + + +def test_reverse_insertion_map_keys_traverse_sorted(http_root: HTTPRoot) -> None: + mapped: dict[str, Asset] = {} + mapped["zzz"] = Asset(ref=REF_B) + mapped["aaa"] = Asset(ref=REF_A) + payload = Payload(mapped=mapped) + assert [a.ref for a in input_assets._iter_assets(payload)] == [REF_A, REF_B] + + manifest = ( + make_entry(REF_A, PNG_A, "media", "image/png"), + make_entry(REF_B, PNG_B, "media", "image/png"), + ) + resolver = FakeResolver( + manifest=manifest, + urls={REF_A: http_root.url("a.png"), REF_B: http_root.url("b.png")}, + ) + assert materialize(payload, "req-map", manifest=manifest, resolver=resolver) == 2 + assert Path(mapped["aaa"].local_path or "").read_bytes() == PNG_A + assert Path(mapped["zzz"].local_path or "").read_bytes() == PNG_B + cleanup_input_assets("req-map", 1) + + # Insertion-order manifest (reversed) is a dispatch reorder: rejected. + reordered = (manifest[1], manifest[0]) + with pytest.raises(ValidationError, match="^input_asset_manifest_mismatch"): + materialize(Payload(mapped=dict(mapped)), "req-map-2", manifest=reordered) + + +def test_non_string_map_key_rejected_at_runtime() -> None: + class Loose(msgspec.Struct): + mapped: dict + + payload = Loose(mapped={1: Asset(ref=REF_A)}) + with pytest.raises(ValidationError, match="^input_asset_map_key_invalid"): + materialize_input_assets(payload, "req-badkey", attempt=1) + + +def test_asset_in_set_rejected_at_runtime() -> None: + # Assets are unhashable, so a real decoded payload can never place one in + # a set; the defensive runtime branch still refuses hostile lookalikes. + class HashableAsset(Asset): + def __hash__(self) -> int: # pragma: no cover - identity for the test + return 7 + + class Loose(msgspec.Struct): + items: Any + + for container in ({HashableAsset(ref=REF_A)}, frozenset({HashableAsset(ref=REF_A)})): + with pytest.raises(ValidationError, match="^input_asset_unordered_container"): + materialize_input_assets(Loose(items=container), "req-set", attempt=1) + + +def test_padded_or_empty_ref_rejected() -> None: + for ref in ("", " tenant/private/a", "tenant/private/a "): + payload = Payload(extra=[Asset(ref=ref)]) + with pytest.raises(ValidationError, match="^invalid_input_asset_ref"): + materialize_input_assets(payload, "req-pad", attempt=1) + + +@pytest.mark.parametrize( + "ref", ["file:///etc/passwd", "s3://bucket/key", "data:image/png,x", "//host/x"] +) +def test_non_http_schemes_rejected(ref: str) -> None: + payload = Payload(extra=[Asset(ref=ref)]) + with pytest.raises(ValidationError, match="^unsupported_input_asset_scheme") as caught: + materialize_input_assets(payload, "req-scheme", attempt=1) + assert ref not in str(caught.value) + + +# --------------------------------------------------------------------------- +# Manifest fencing (before any resolver call or GET) +# --------------------------------------------------------------------------- + + +def test_zero_assets_zero_manifest() -> None: + assert materialize(Payload(prompt="text"), "req-zero") == 0 + assert not inputs_dir_for_request("req-zero", 1).exists() + + +def test_zero_assets_with_manifest_rejected() -> None: + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + with pytest.raises(ValidationError, match="^input_asset_manifest_mismatch"): + materialize(Payload(prompt="text"), "req-zero-extra", manifest=manifest) + + +def test_public_only_never_calls_resolver(http_root: HTTPRoot) -> None: + resolver = FakeResolver() + payload = Payload(image=ImageAsset(ref=http_root.url("a.png"))) + assert materialize(payload, "req-public", resolver=resolver) == 1 + assert resolver.calls == [] + assert Path(payload.image.local_path or "").read_bytes() == PNG_A + cleanup_input_assets("req-public", 1) + + +@pytest.mark.parametrize( + "case,payload_refs,manifest_refs", + [ + ("missing", [REF_A], []), + ("extra", [REF_A], [REF_A, REF_B]), + ("reorder", [REF_A, REF_B], [REF_B, REF_A]), + ], +) +def test_dispatch_manifest_count_order_fenced( + case: str, payload_refs: list[str], manifest_refs: list[str] +) -> None: + bodies = {REF_A: PNG_A, REF_B: PNG_B} + payload = Payload(extra=[Asset(ref=r) for r in payload_refs]) + manifest = tuple(make_entry(r, bodies[r], "media", "image/png") for r in manifest_refs) + resolver = FakeResolver(manifest=manifest) + with pytest.raises(ValidationError, match="^input_asset_manifest_mismatch") as caught: + materialize(payload, f"req-{case}", manifest=manifest, resolver=resolver) + assert resolver.calls == [] + assert_no_leak(caught.value, REF_A, REF_B) + + +def test_dispatch_kind_mismatch_fenced() -> None: + payload = Payload(image=ImageAsset(ref=REF_A)) + manifest = (make_entry(REF_A, PNG_A, "video", "image/png"),) + resolver = FakeResolver(manifest=manifest) + with pytest.raises(ValidationError, match="^input_asset_manifest_mismatch"): + materialize(payload, "req-kindfence", manifest=manifest, resolver=resolver) + assert resolver.calls == [] + + +@pytest.mark.parametrize( + "overrides", + [ + {"asset_id": ""}, + {"blake3": "abc"}, + {"blake3": "Z" * 64}, + {"size_bytes": 0}, + {"mime_type": ""}, + ], +) +def test_invalid_manifest_metadata_fenced(overrides: dict[str, Any]) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png", **overrides),) + resolver = FakeResolver(manifest=manifest) + with pytest.raises(ValidationError, match="^input_asset_manifest_invalid"): + materialize(payload, "req-meta", manifest=manifest, resolver=resolver) + assert resolver.calls == [] + + +def test_private_inputs_require_positive_attempt() -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + with pytest.raises(ValidationError, match="^input_asset_attempt_invalid"): + materialize(payload, "req-attempt", attempt=0, manifest=manifest) + + +# --------------------------------------------------------------------------- +# Resolver call contract +# --------------------------------------------------------------------------- + + +def test_exactly_one_strict_resolver_post(http_root: HTTPRoot) -> None: + payload = Payload( + image=ImageAsset(ref=REF_A), + extra=[Asset(ref=REF_B)], + audio=AudioAsset(ref=REF_C), + ) + manifest = ( + make_entry(REF_A, PNG_A, "image", "image/png"), + make_entry(REF_B, PNG_B, "media", "image/png"), + make_entry(REF_C, MP3, "audio", "audio/mpeg"), + ) + resolver = FakeResolver( + manifest=manifest, + urls={ + REF_A: http_root.url("a.png"), + REF_B: http_root.url("b.png"), + REF_C: http_root.url("audio.mp3"), + }, + ) + assert materialize(payload, "req-one-post", attempt=7, manifest=manifest, resolver=resolver) == 3 + assert len(resolver.calls) == 1 + url, headers, body = resolver.calls[0] + assert url == f"{BASE_URL}/api/v1/worker/input-assets/resolve" + assert headers["Authorization"] == f"Bearer {CAPABILITY}" + assert json.loads(body) == {"request_id": "req-one-post", "attempt": 7} + cleanup_input_assets("req-one-post", 7) + + +def test_one_private_asset_preserves_all_fields(http_root: HTTPRoot) -> None: + asset = ImageAsset( + ref=REF_A, + local_path="/etc/passwd", + mime_type="image/png", + size_bytes=len(PNG_A), + blake3=b3(PNG_A), + sha256="caller-provided-sha", + ) + payload = Payload(image=asset) + manifest = (make_entry(REF_A, PNG_A, "image", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + assert materialize(payload, "req-priv-one", manifest=manifest, resolver=resolver) == 1 + assert asset.ref == REF_A # opaque ref NEVER rewritten + assert asset.local_path and Path(asset.local_path).read_bytes() == PNG_A + assert Path(asset.local_path).parent == inputs_dir_for_request("req-priv-one", 1) + assert asset.local_path != "/etc/passwd" + assert asset.mime_type == "image/png" + assert asset.size_bytes == len(PNG_A) + assert asset.sha256 == "caller-provided-sha" + + cleanup_input_assets("req-priv-one", 1) + assert asset.local_path is None + + +def test_duplicate_private_occurrences_share_one_download(http_root: HTTPRoot) -> None: + a1, b, a2 = Asset(ref=REF_A), Asset(ref=REF_B), Asset(ref=REF_A) + payload = Payload(extra=[a1, b, a2]) + manifest = ( + make_entry(REF_A, PNG_A, "media", "image/png"), + make_entry(REF_B, PNG_B, "media", "image/png"), + ) + resolver = FakeResolver( + manifest=manifest, + urls={REF_A: http_root.url("a.png"), REF_B: http_root.url("b.png")}, + ) + assert materialize(payload, "req-aba", manifest=manifest, resolver=resolver) == 2 + assert a1.local_path == a2.local_path and a1.local_path != b.local_path + assert http_root.hits == {"/a.png": 1, "/b.png": 1} + cleanup_input_assets("req-aba", 1) + + +def test_mixed_private_and_public(http_root: HTTPRoot) -> None: + private = ImageAsset(ref=REF_A) + public = ImageAsset(ref=http_root.url("b.png")) + payload = Payload(image=private, extra=[public]) + manifest = (make_entry(REF_A, PNG_A, "image", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + assert materialize(payload, "req-mixed", manifest=manifest, resolver=resolver) == 2 + assert len(resolver.calls) == 1 + assert Path(private.local_path or "").read_bytes() == PNG_A + assert Path(public.local_path or "").read_bytes() == PNG_B + assert public.ref == http_root.url("b.png") + cleanup_input_assets("req-mixed", 1) + + +def test_resolver_409_cancels_attempt() -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, status=409) + with pytest.raises(CanceledError): + materialize(payload, "req-409", manifest=manifest, resolver=resolver) + assert not inputs_dir_for_request("req-409", 1).exists() + + +@pytest.mark.parametrize("status", [400, 403, 500, 503]) +def test_resolver_failure_statuses_are_retryable(status: int) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, status=status) + with pytest.raises(RetryableError, match="^input_asset_resolution_unavailable") as caught: + materialize(payload, f"req-{status}", manifest=manifest, resolver=resolver) + assert "resolver-error-body" not in str(caught.value) + assert_no_leak(caught.value, REF_A) + + +def test_resolver_network_failure_is_retryable() -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, exc=ConnectionError("boom")) + with pytest.raises(RetryableError, match="^input_asset_resolution_unavailable"): + materialize(payload, "req-net", manifest=manifest, resolver=resolver) + + +def test_missing_resolver_credentials_is_retryable() -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + with pytest.raises(RetryableError, match="^input_asset_resolution_unavailable"): + materialize_input_assets( + payload, "req-nocreds", attempt=1, manifest=manifest, file_base_url="", + capability_token="", + ) + + +@pytest.mark.parametrize( + "raw", + [ + b"not json", + b"{}", + b'{"assets": [{"bogus": 1}]}', + b'{"assets": [], "extra_top": 1}', + ], +) +def test_malformed_resolver_response_is_retryable(raw: bytes) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, raw_body=raw) + with pytest.raises(RetryableError, match="^input_asset_response_mismatch"): + materialize(payload, "req-malformed", manifest=manifest, resolver=resolver) + + +def test_unknown_response_field_is_rejected(http_root: HTTPRoot) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + + def add_unknown(assets: list[dict]) -> list[dict]: + assets[0]["surprise"] = True + return assets + + resolver = FakeResolver( + manifest=manifest, urls={REF_A: http_root.url("a.png")}, mutate=add_unknown + ) + with pytest.raises(RetryableError, match="^input_asset_response_mismatch"): + materialize(payload, "req-unknown", manifest=manifest, resolver=resolver) + assert http_root.hits == {} + + +@pytest.mark.parametrize( + "case,mutate", + [ + ("missing", lambda assets: assets[:-1]), + ("extra", lambda assets: assets + [dict(assets[0])]), + ("reorder", lambda assets: list(reversed(assets))), + ], +) +def test_response_count_order_rejected( + http_root: HTTPRoot, case: str, mutate: Any +) -> None: + payload = Payload(extra=[Asset(ref=REF_A), Asset(ref=REF_B)]) + manifest = ( + make_entry(REF_A, PNG_A, "media", "image/png"), + make_entry(REF_B, PNG_B, "media", "image/png"), + ) + resolver = FakeResolver( + manifest=manifest, + urls={REF_A: http_root.url("a.png"), REF_B: http_root.url("b.png")}, + mutate=mutate, + ) + with pytest.raises(RetryableError, match="^input_asset_response_mismatch"): + materialize(payload, f"req-resp-{case}", manifest=manifest, resolver=resolver) + assert http_root.hits == {} + + +@pytest.mark.parametrize( + "field_name,value", + [ + ("asset_id", "asset-other"), + ("source_ref", "tenant/private/other"), + ("blake3", "0" * 64), + ("size_bytes", 999999), + ("kind", "video"), + ("mime_type", "image/jpeg"), + ("url", ""), + ("url_expires_at", "not-a-timestamp"), + ], +) +def test_changed_immutable_response_field_rejected( + http_root: HTTPRoot, field_name: str, value: Any +) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + + def change(assets: list[dict]) -> list[dict]: + assets[0][field_name] = value + return assets + + resolver = FakeResolver( + manifest=manifest, urls={REF_A: http_root.url("a.png")}, mutate=change + ) + with pytest.raises(RetryableError, match="^input_asset_response_mismatch"): + materialize(payload, f"req-chg-{field_name}", manifest=manifest, resolver=resolver) + assert http_root.hits == {} + + +def test_cancel_observed_at_resolve_boundary() -> None: + canceled = threading.Event() + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: "http://127.0.0.1:1/a"}) + resolver.on_call = canceled.set + with pytest.raises(CanceledError): + materialize( + payload, + "req-cancel-resolve", + manifest=manifest, + resolver=resolver, + cancel_check=canceled.is_set, + ) + assert len(resolver.calls) == 1 + assert not inputs_dir_for_request("req-cancel-resolve", 1).exists() + + +def test_blocked_resolved_url_rejected_without_download( + http_root: HTTPRoot, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(input_assets, "_url_is_blocked", lambda _url: True) + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(ValidationError, match="^input_asset_url_not_allowed") as caught: + materialize(payload, "req-blocked", manifest=manifest, resolver=resolver) + assert http_root.hits == {} + assert_no_leak(caught.value, REF_A, http_root.base_url) + + +# --------------------------------------------------------------------------- +# Download verification +# --------------------------------------------------------------------------- + + +def test_size_mismatch_is_retryable_integrity_failure(http_root: HTTPRoot) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = ( + make_entry(REF_A, PNG_A, "media", "image/png", size_bytes=len(PNG_A) + 1), + ) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(RetryableError, match="^input_asset_integrity_failed"): + materialize(payload, "req-size", manifest=manifest, resolver=resolver) + assert not inputs_dir_for_request("req-size", 1).exists() + + +def test_blake3_mismatch_is_retryable_integrity_failure(http_root: HTTPRoot) -> None: + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png", blake3="0" * 64),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(RetryableError, match="^input_asset_integrity_failed") as caught: + materialize(payload, "req-hash", manifest=manifest, resolver=resolver) + assert_no_leak(caught.value, REF_A, http_root.base_url) + + +def test_private_kind_and_mime_limits_enforced(http_root: HTTPRoot) -> None: + # Declared AudioAsset (manifest kind audio) but the stored bytes are a PNG. + payload = Payload(audio=AudioAsset(ref=REF_A)) + manifest = (make_entry(REF_A, PNG_A, "audio", "audio/mpeg"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(ValidationError, match="^input_asset_kind_mismatch"): + materialize(payload, "req-kind", manifest=manifest, resolver=resolver) + + # Caller MIME allowlist still applies to private inputs. + payload2 = Payload( + extra=[Asset(ref=REF_A, url_allowed_mime_types=("audio/*",))] + ) + manifest2 = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver2 = FakeResolver(manifest=manifest2, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(ValidationError, match="^input_asset_kind_mismatch"): + materialize(payload2, "req-allow", manifest=manifest2, resolver=resolver2) + + +def test_private_caller_byte_cap_precedes_download(http_root: HTTPRoot) -> None: + payload = Payload(extra=[Asset(ref=REF_A, url_max_bytes=8)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(ValidationError, match="^input_asset_too_large"): + materialize(payload, "req-cap", manifest=manifest, resolver=resolver) + assert http_root.hits == {} + + +def test_private_corrupt_image_fails_decode(http_root: HTTPRoot) -> None: + payload = Payload(image=ImageAsset(ref=REF_A)) + manifest = (make_entry(REF_A, CORRUPT_PNG, "image", "image/png"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("corrupt.png")}) + with pytest.raises(ValidationError, match="^input_asset_decode_failed"): + materialize(payload, "req-decode", manifest=manifest, resolver=resolver) + assert not inputs_dir_for_request("req-decode", 1).exists() + + +def test_private_download_cancel_cleans_partial_file(http_root: HTTPRoot) -> None: + asset = Asset(ref=REF_A) + payload = Payload(extra=[asset]) + manifest = (make_entry(REF_A, SLOW_BODY, "media", "application/octet-stream"),) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("slow.bin")}) + canceled = threading.Event() + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit( + materialize, + payload, + "req-dl-cancel", + manifest=manifest, + resolver=resolver, + cancel_check=canceled.is_set, + ) + assert http_root.slow_started.wait(timeout=5) + canceled.set() + http_root.slow_release.set() + with pytest.raises(CanceledError): + future.result(timeout=10) + + assert asset.local_path is None + assert not inputs_dir_for_request("req-dl-cancel", 1).exists() + + +def test_download_failure_cleans_prior_private_success(http_root: HTTPRoot) -> None: + ok, failing = Asset(ref=REF_A), Asset(ref=REF_B) + payload = Payload(extra=[ok, failing]) + manifest = ( + make_entry(REF_A, PNG_A, "media", "image/png"), + make_entry(REF_B, PNG_B, "media", "image/png"), + ) + resolver = FakeResolver( + manifest=manifest, + urls={REF_A: http_root.url("a.png"), REF_B: http_root.url("missing.bin")}, + ) + with pytest.raises(RetryableError, match="^input_asset_download_failed"): + materialize(payload, "req-partial", manifest=manifest, resolver=resolver) + assert ok.local_path is None and failing.local_path is None + assert not inputs_dir_for_request("req-partial", 1).exists() + + +# --------------------------------------------------------------------------- +# Attempt fencing: N and N+1 never touch each other's scope +# --------------------------------------------------------------------------- + + +def test_attempt_scopes_are_isolated(http_root: HTTPRoot) -> None: + first = Payload(image=ImageAsset(ref=http_root.url("a.png"))) + second = Payload(image=ImageAsset(ref=http_root.url("b.png"))) + materialize(first, "req-fence", attempt=1) + materialize(second, "req-fence", attempt=2) + dir1 = inputs_dir_for_request("req-fence", 1) + dir2 = inputs_dir_for_request("req-fence", 2) + try: + assert dir1.exists() and dir2.exists() and dir1 != dir2 + assert first.image.local_path and second.image.local_path + + cleanup_input_assets("req-fence", 2) + assert dir1.exists() and not dir2.exists() + assert first.image.local_path is not None + assert second.image.local_path is None + finally: + cleanup_input_assets("req-fence", 1) + cleanup_input_assets("req-fence", 2) + assert not dir1.exists() + + +class WorkerInput(msgspec.Struct): + images: list[ImageAsset] = [] + + +class WorkerOutput(msgspec.Struct): + count: int + + +def test_executor_passes_manifest_and_capability(monkeypatch: pytest.MonkeyPatch, http_root: HTTPRoot) -> None: + """The executor threads run.input_assets + file_base_url + capability + + (request_id, attempt) + the cancel callback into the materializer.""" + manifest_rows = [ + pb.InputAsset( + asset_id="asset-a", + source_ref=REF_A, + blake3=b3(PNG_A), + size_bytes=len(PNG_A), + kind="image", + mime_type="image/png", + ) + ] + resolver = FakeResolver( + manifest=manifest_from_run_job(manifest_rows), + urls={REF_A: http_root.url("a.png")}, + ) + monkeypatch.setattr(input_assets, "_default_resolve_transport", resolver) + + seen: dict[str, Any] = {} + + def handler(ctx: Any, payload: WorkerInput) -> WorkerOutput: + seen["ref"] = payload.images[0].ref + seen["bytes"] = Path(payload.images[0].local_path or "").read_bytes() + return WorkerOutput(count=len(payload.images)) + + async def scenario() -> pb.JobResult: + sent: list[pb.WorkerMessage] = [] + + async def send(message: pb.WorkerMessage) -> None: + sent.append(message) + + spec = EndpointSpec( + name="edit", + method=handler, + kind="inference", + payload_type=WorkerInput, + output_mode="single", + ) + executor = Executor([spec], send) + executor.file_base_url = BASE_URL + await executor.handle_run_job( + pb.RunJob( + request_id="exec-manifest", + attempt=3, + function_name="edit", + input_payload=msgspec.msgpack.encode( + WorkerInput(images=[ImageAsset(ref=REF_A)]) + ), + capability_token=CAPABILITY, + input_assets=manifest_rows, + ) + ) + job = executor.jobs[("exec-manifest", 3)] + assert job.task is not None + await job.task + results = [m.job_result for m in sent if m.WhichOneof("msg") == "job_result"] + assert len(results) == 1 + return results[0] + + result = asyncio.run(scenario()) + assert result.status == pb.JOB_STATUS_OK + assert seen["ref"] == REF_A and seen["bytes"] == PNG_A + assert len(resolver.calls) == 1 + _, headers, body = resolver.calls[0] + assert headers["Authorization"] == f"Bearer {CAPABILITY}" + assert json.loads(body) == {"request_id": "exec-manifest", "attempt": 3} + assert not inputs_dir_for_request("exec-manifest", 3).exists() + + +def test_executor_supersession_fences_result_and_scopes(http_root: HTTPRoot) -> None: + """Attempt N canceled by N+1 publishes nothing and never touches N+1's + inputs; N+1 completes normally.""" + handled: list[int] = [] + + def handler(ctx: Any, payload: WorkerInput) -> WorkerOutput: + handled.append(len(payload.images)) + return WorkerOutput(count=len(payload.images)) + + async def scenario() -> list[tuple[int, int]]: + sent: list[pb.WorkerMessage] = [] + + async def send(message: pb.WorkerMessage) -> None: + sent.append(message) + + spec = EndpointSpec( + name="edit", + method=handler, + kind="inference", + payload_type=WorkerInput, + output_mode="single", + ) + executor = Executor([spec], send) + await executor.handle_run_job( + pb.RunJob( + request_id="exec-fence", + attempt=1, + function_name="edit", + input_payload=msgspec.msgpack.encode( + WorkerInput(images=[ImageAsset(ref=http_root.url("slow.bin"))]) + ), + ) + ) + assert await asyncio.to_thread(http_root.slow_started.wait, 5) + await executor.handle_run_job( + pb.RunJob( + request_id="exec-fence", + attempt=2, + function_name="edit", + input_payload=msgspec.msgpack.encode( + WorkerInput(images=[ImageAsset(ref=http_root.url("a.png"))]) + ), + ) + ) + http_root.slow_release.set() + for key in (("exec-fence", 1), ("exec-fence", 2)): + job = executor.jobs[key] + assert job.task is not None + await asyncio.wait_for(job.task, timeout=10) + return [ + (m.job_result.attempt, m.job_result.status) + for m in sent + if m.WhichOneof("msg") == "job_result" + ] + + results = asyncio.run(scenario()) + # Attempt 1 was superseded: its terminal result is suppressed. Attempt 2 + # published exactly one OK result and both scopes are cleaned. + assert results == [(2, pb.JOB_STATUS_OK)] + assert handled == [1] + assert not inputs_dir_for_request("exec-fence", 1).exists() + assert not inputs_dir_for_request("exec-fence", 2).exists() + + +# --------------------------------------------------------------------------- +# Discovery: build-time schema contract +# --------------------------------------------------------------------------- + + +class DiscNested(msgspec.Struct): + clips: list[VideoAsset] = [] + + +class DiscPayload(msgspec.Struct): + generic: Asset | None = None + media: MediaAsset | None = None + image: ImageAsset | None = None + nested: DiscNested | None = None + mapped: dict[str, AudioAsset] = {} + + +def test_discovery_emits_media_kind_for_generic_assets() -> None: + meta = _collect_payload_moderation_metadata(DiscPayload) + assert meta["media"] == [ + {"field": "generic", "kind": "media"}, + {"field": "media", "kind": "media"}, + {"field": "image", "kind": "image"}, + {"field": "nested.clips[]", "kind": "video"}, + {"field": "mapped.*", "kind": "audio"}, + ] + + +def test_discovery_rejects_asset_bearing_sets() -> None: + class BadSet(msgspec.Struct): + items: set[ImageAsset] = set() + + class BadFrozen(msgspec.Struct): + items: frozenset[Asset] = frozenset() + + class NestedBad(msgspec.Struct): + wrapped: set[DiscNested] = set() + + for bad in (BadSet, BadFrozen, NestedBad): + with pytest.raises(ValueError, match="unordered set/frozenset"): + _collect_payload_moderation_metadata(bad) + + class OkSet(msgspec.Struct): + tags: set[str] = set() + + assert _collect_payload_moderation_metadata(OkSet) == {} + + +def test_discovery_rejects_non_string_key_asset_maps() -> None: + class BadMap(msgspec.Struct): + items: dict[int, ImageAsset] = {} + + with pytest.raises(ValueError, match="string keys"): + _collect_payload_moderation_metadata(BadMap) + + class OkMap(msgspec.Struct): + counts: dict[int, int] = {} + + assert _collect_payload_moderation_metadata(OkMap) == {} + + +# --------------------------------------------------------------------------- +# Proto descriptors: v4 wire contract +# --------------------------------------------------------------------------- + + +def test_protocol_version_is_4() -> None: + assert pb.PROTOCOL_VERSION_CURRENT == 4 + + +def test_run_job_input_assets_field_descriptor() -> None: + # 15, not 14: th#913 claimed RunJob.lane=14 on chaos before the v4 landing. + fd = pb.RunJob.DESCRIPTOR.fields_by_name["input_assets"] + assert fd.number == 15 + assert fd.is_repeated + assert fd.message_type is pb.InputAsset.DESCRIPTOR + + +def test_input_asset_field_numbers() -> None: + fields = [(f.name, f.number) for f in pb.InputAsset.DESCRIPTOR.fields] + assert fields == [ + ("asset_id", 1), + ("source_ref", 2), + ("blake3", 3), + ("size_bytes", 4), + ("kind", 5), + ("mime_type", 6), + ] + + +def test_manifest_from_run_job_round_trip() -> None: + rows = [ + pb.InputAsset( + asset_id="asset-a", + source_ref=REF_A, + blake3=b3(PNG_A), + size_bytes=len(PNG_A), + kind="image", + mime_type="image/png", + ), + pb.InputAsset( + asset_id="asset-b", + source_ref=REF_B, + blake3=b3(PNG_B), + size_bytes=len(PNG_B), + kind="media", + mime_type="image/png", + ), + ] + manifest = manifest_from_run_job(rows) + assert manifest == ( + InputManifestEntry( + asset_id="asset-a", + source_ref=REF_A, + blake3=b3(PNG_A), + size_bytes=len(PNG_A), + kind="image", + mime_type="image/png", + ), + InputManifestEntry( + asset_id="asset-b", + source_ref=REF_B, + blake3=b3(PNG_B), + size_bytes=len(PNG_B), + kind="media", + mime_type="image/png", + ), + ) + assert manifest_from_run_job(()) == () From 4048de6d6c8ec13198b215b9c6ef8dcb470fdaa4 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 21 Jul 2026 13:59:09 -0600 Subject: [PATCH 2/5] =?UTF-8?q?gw#585:=20marco-polo-attach=20=E2=80=94=20p?= =?UTF-8?q?rivate-input=20echo=20probe=20(blake3/size/local=5Fpath,=20opti?= =?UTF-8?q?onal=20in-flight=20delay)=20for=20the=20v4=20e2e=20lane?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/marco-polo/src/marco_polo/main.py | 40 ++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/examples/marco-polo/src/marco_polo/main.py b/examples/marco-polo/src/marco_polo/main.py index 3f0273bd..3fdbac2d 100644 --- a/examples/marco-polo/src/marco_polo/main.py +++ b/examples/marco-polo/src/marco_polo/main.py @@ -5,6 +5,7 @@ import msgspec from gen_worker import RequestContext, ValidationError, endpoint +from gen_worker.api.types import Asset class MarcoPoloInput(msgspec.Struct): @@ -15,6 +16,19 @@ class MarcoPoloOutput(msgspec.Struct): response: str +class MarcoAttachInput(msgspec.Struct): + image: Asset + text: str = "" + delay_s: float = 0.0 + + +class MarcoAttachOutput(msgspec.Struct): + response: str + blake3: str + size_bytes: int + local_path: str + + @endpoint class MarcoPolo: def marco_polo(self, ctx: RequestContext, data: MarcoPoloInput) -> MarcoPoloOutput: @@ -61,6 +75,32 @@ async def marco_polo_wedge( time.sleep(1800) return MarcoPoloOutput(response="unreachable") + async def marco_polo_attach( + self, ctx: RequestContext, data: MarcoAttachInput + ) -> MarcoAttachOutput: + """th#886 v4 private-input probe: echoes the materialized attachment's + BLAKE3/size/attempt-local path so the harness can verify exact bytes + and post-terminal temp cleanup. delay_s>0 holds the job in flight for + disconnect/retry chaos.""" + import blake3 + + if str(data.text or "").strip().lower() != "marco": + raise ValidationError(f"expected 'marco', got {data.text!r}") + path = data.image.path() + with open(path, "rb") as f: + raw = f.read() + waited = 0.0 + while waited < float(data.delay_s or 0.0): + ctx.raise_if_cancelled() + await asyncio.sleep(0.15) + waited += 0.15 + return MarcoAttachOutput( + response="polo", + blake3=blake3.blake3(raw).hexdigest(), + size_bytes=len(raw), + local_path=path, + ) + async def marco_polo_stream( self, ctx: RequestContext, data: MarcoPoloInput ) -> AsyncIterator[MarcoPoloOutput]: From 936e32dd20aa4c57fcaf6ddb495fc7a3accb4313 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 21 Jul 2026 14:06:35 -0600 Subject: [PATCH 3/5] =?UTF-8?q?gw#585:=20GEN=5FWORKER=5FINTERNAL=5FOBJECT?= =?UTF-8?q?=5FHOSTS=20=E2=80=94=20exact-host=20allowlist=20for=20resolver-?= =?UTF-8?q?minted=20private-input=20URLs=20(datacenter-internal=20object?= =?UTF-8?q?=20stores);=20caller=20public=20transports=20keep=20the=20full?= =?UTF-8?q?=20SSRF=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gen_worker/input_assets.py | 15 +++++++++++-- tests/test_input_manifest_v4.py | 37 +++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/src/gen_worker/input_assets.py b/src/gen_worker/input_assets.py index 4561451d..d960def0 100644 --- a/src/gen_worker/input_assets.py +++ b/src/gen_worker/input_assets.py @@ -442,7 +442,16 @@ def _parse_rfc3339(value: str) -> bool: return True -def _validate_transport_url(url: str, path: str) -> None: +def _internal_object_hosts() -> frozenset[str]: + """GEN_WORKER_INTERNAL_OBJECT_HOSTS: exact hostnames of a deployment's + internal object store (datacenter MinIO/NFS gateway). Honored ONLY for + resolver-minted private-input URLs; caller public transports always face + the full SSRF policy.""" + raw = os.environ.get("GEN_WORKER_INTERNAL_OBJECT_HOSTS", "") + return frozenset(host.strip().lower() for host in raw.split(",") if host.strip()) + + +def _validate_transport_url(url: str, path: str, *, resolver_minted: bool = False) -> None: try: parsed = urllib.parse.urlsplit(str(url or "")) except ValueError: @@ -451,6 +460,8 @@ def _validate_transport_url(url: str, path: str) -> None: ) from None if parsed.scheme.lower() not in ("http", "https") or parsed.hostname is None: raise ValidationError(f"invalid_input_asset_url: {path} has an invalid transport URL") + if resolver_minted and parsed.hostname.lower() in _internal_object_hosts(): + return try: blocked = _url_is_blocked(url) except Exception: @@ -735,7 +746,7 @@ def materialize_input_assets( for unit in private_units: assert unit.entry is not None unit.url = by_entry[unit.entry] - _validate_transport_url(unit.url, unit.occurrences[0].path) + _validate_transport_url(unit.url, unit.occurrences[0].path, resolver_minted=True) state = _start_scope( request_id, diff --git a/tests/test_input_manifest_v4.py b/tests/test_input_manifest_v4.py index 7b12d7bf..3e029ae0 100644 --- a/tests/test_input_manifest_v4.py +++ b/tests/test_input_manifest_v4.py @@ -628,6 +628,43 @@ def test_blocked_resolved_url_rejected_without_download( assert_no_leak(caught.value, REF_A, http_root.base_url) +def test_internal_object_host_allowlist_only_for_resolver_minted( + http_root: HTTPRoot, monkeypatch: pytest.MonkeyPatch +) -> None: + # Deployment-declared internal object store: resolver-minted URLs to the + # exact host pass; the same host stays blocked without the declaration. + monkeypatch.setattr(input_assets, "_url_is_blocked", lambda _url: True) + host = "127.0.0.1" + payload = Payload(extra=[Asset(ref=REF_A)]) + manifest = (make_entry(REF_A, PNG_A, "media", "image/png"),) + + monkeypatch.delenv("GEN_WORKER_INTERNAL_OBJECT_HOSTS", raising=False) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + with pytest.raises(ValidationError, match="^input_asset_url_not_allowed"): + materialize(payload, "req-internal-denied", manifest=manifest, resolver=resolver) + + monkeypatch.setenv("GEN_WORKER_INTERNAL_OBJECT_HOSTS", host) + resolver = FakeResolver(manifest=manifest, urls={REF_A: http_root.url("a.png")}) + count = materialize( + Payload(extra=[Asset(ref=REF_A)]), + "req-internal-allowed", + manifest=manifest, + resolver=resolver, + ) + assert count == 1 + + +def test_internal_object_host_never_unblocks_caller_public_urls( + http_root: HTTPRoot, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(input_assets, "_url_is_blocked", lambda _url: True) + monkeypatch.setenv("GEN_WORKER_INTERNAL_OBJECT_HOSTS", "127.0.0.1") + payload = Payload(extra=[Asset(ref=http_root.url("pub.png"))]) + with pytest.raises(ValidationError, match="^input_asset_url_not_allowed"): + materialize(payload, "req-public-still-blocked") + assert http_root.hits == {} + + # --------------------------------------------------------------------------- # Download verification # --------------------------------------------------------------------------- From b262b6701d2be430ca0e927664d06d945daad7b9 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 21 Jul 2026 14:08:44 -0600 Subject: [PATCH 4/5] gw#585: marco-polo-attach reads via os.fspath (Asset is os.PathLike; no .path() method) --- examples/marco-polo/src/marco_polo/main.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/marco-polo/src/marco_polo/main.py b/examples/marco-polo/src/marco_polo/main.py index 3fdbac2d..9ae78ed8 100644 --- a/examples/marco-polo/src/marco_polo/main.py +++ b/examples/marco-polo/src/marco_polo/main.py @@ -1,5 +1,6 @@ # source-hash bust: 2026-07-03 (#368: @endpoint API rewrite) import asyncio +import os import time from typing import AsyncIterator @@ -86,7 +87,7 @@ async def marco_polo_attach( if str(data.text or "").strip().lower() != "marco": raise ValidationError(f"expected 'marco', got {data.text!r}") - path = data.image.path() + path = os.fspath(data.image) with open(path, "rb") as f: raw = f.read() waited = 0.0 From 08dd86afcb3b75a19def53f76df07b6f85e9ae8f Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 21 Jul 2026 14:15:51 -0600 Subject: [PATCH 5/5] =?UTF-8?q?release:=20gen-worker=200.43.0=20=E2=80=94?= =?UTF-8?q?=20tensorhub=20v4=20private-input=20manifest=20hard=20cut=20(gw?= =?UTF-8?q?#585)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 20 ++++++++++++++++++++ pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f792054c..d9e296bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # Changelog +## 0.43.0 (2026-07-21) + +- **gw#585: tensorhub v4 private-input manifests — gRPC protocol v3 -> v4 + HARD CUT (th#886).** The hub no longer rewrites canonical payload bytes + with presigned URLs. `RunJob.input_assets` (field 15) carries the ordered + credential-free manifest; the worker resolves fresh transport URLs itself + with one strict `POST /worker/input-assets/resolve` under its + attempt-scoped capability, verifies exact size/BLAKE3/MIME/kind, preserves + opaque `Asset.ref`, sets only `local_path`, and cleans attempt-owned temp + files on every terminal path. Endpoint build now rejects Asset-bearing + `set`/`frozenset` and non-string-keyed mappings (unordered containers have + no manifest order); base `Asset`/`MediaAsset` discover as `kind=media`. + A v0.43 worker cannot serve a v3 hub and vice versa — deploy in lockstep + with the tensorhub v4 release. +- **`GEN_WORKER_INTERNAL_OBJECT_HOSTS`**: exact-host allowlist that exempts + resolver-minted private-input URLs from the private-IP SSRF gate for + deployments whose object store lives on an internal network. Caller public + transports always face the full SSRF policy. +- marco-polo example: `marco-polo-attach` private-input echo probe (e2e). + ## 0.42.0 (2026-07-21) - **gw#615: disk telemetry can no longer freeze the event loop (0.40.7 diff --git a/pyproject.toml b/pyproject.toml index 88cfc07b..d6927f78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "gen-worker" -version = "0.42.0" +version = "0.43.0" description = "A library used to build custom functions in Cozy Creator's serverless function platform." readme = "README.md" license = "MIT" diff --git a/uv.lock b/uv.lock index 179e72e3..0b11223b 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "gen-worker" -version = "0.42.0" +version = "0.43.0" source = { editable = "." } dependencies = [ { name = "blake3" },