urun-cli 0.5.1__tar.gz → 0.5.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {urun_cli-0.5.1 → urun_cli-0.5.2}/PKG-INFO +1 -1
- {urun_cli-0.5.1 → urun_cli-0.5.2}/pyproject.toml +1 -1
- urun_cli-0.5.2/src/urun/catalog.py +652 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/cli.py +164 -45
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/deps.py +149 -7
- urun_cli-0.5.2/src/urun/serve.py +524 -0
- urun_cli-0.5.1/src/urun/catalog.py +0 -291
- urun_cli-0.5.1/src/urun/serve.py +0 -360
- {urun_cli-0.5.1 → urun_cli-0.5.2}/.gitignore +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/CHANGELOG.md +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/LICENSE +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/README.md +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/SECURITY.md +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/__init__.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/api.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/config.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/discovery.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/errors.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/manifest.py +0 -0
- {urun_cli-0.5.1 → urun_cli-0.5.2}/src/urun/retry.py +0 -0
|
@@ -0,0 +1,652 @@
|
|
|
1
|
+
"""Thin PostgREST client for the control-plane model catalog.
|
|
2
|
+
|
|
3
|
+
``public.model_catalog`` is org-AGNOSTIC, anon-readable reference data published
|
|
4
|
+
by a urun-infra migration (one row per ``(model_id, variant, gpu_spec)``; see
|
|
5
|
+
``urun-infra/catalog/CONTRACT.md``). ``urun serve catalog`` enumerates it and
|
|
6
|
+
``urun serve <id>[:<variant>]`` resolves a single row to deploy.
|
|
7
|
+
|
|
8
|
+
CANONICAL PATH (no silent fallbacks — see CLAUDE.md): the catalog is read over
|
|
9
|
+
PostgREST at the control-plane's ``/rest/v1/model_catalog`` endpoint. The base is
|
|
10
|
+
derived from the SAME control-plane API URL the rest of the CLI uses (the origin
|
|
11
|
+
of ``--api-url`` / ``URUN_API_URL`` / saved login, default ``api.urun.sh``), and
|
|
12
|
+
PostgREST is reached with the public Supabase anon key. There is exactly ONE
|
|
13
|
+
resolution for each input; a missing anon key FAILS HARD rather than guessing.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
import json
|
|
19
|
+
import urllib.error
|
|
20
|
+
import urllib.parse
|
|
21
|
+
import urllib.request
|
|
22
|
+
from collections.abc import Callable
|
|
23
|
+
from dataclasses import dataclass
|
|
24
|
+
from typing import Any
|
|
25
|
+
|
|
26
|
+
from . import __version__
|
|
27
|
+
from .api import require_safe_url
|
|
28
|
+
from .errors import ApiError, UrunError
|
|
29
|
+
from .retry import DEPLOY_BACKOFF, BackoffPolicy, retry_call
|
|
30
|
+
|
|
31
|
+
CATALOG_TABLE = "model_catalog"
|
|
32
|
+
|
|
33
|
+
# PostgREST is mounted under /rest/v1 on the Supabase control plane; the CLI's
|
|
34
|
+
# control-plane API URL (e.g. https://api.urun.sh/v1) shares the same origin.
|
|
35
|
+
POSTGREST_PREFIX = "/rest/v1"
|
|
36
|
+
|
|
37
|
+
# CLI-LOCAL SKU -> per-GPU VRAM (GB). DELIBERATELY DUPLICATED from urun-python's
|
|
38
|
+
# GPU_VRAM_GB map: the CLI is standalone and cannot import the (often unreleased)
|
|
39
|
+
# urun-python at runtime, so it carries its own small table. Keys are lowercase
|
|
40
|
+
# SKU "kind" strings matching the prefix of a `gpu_spec` (`<kind>:<count>`).
|
|
41
|
+
GPU_VRAM_GB: dict[str, int] = {
|
|
42
|
+
"l4": 24,
|
|
43
|
+
"a10g": 24,
|
|
44
|
+
"l40": 48,
|
|
45
|
+
"l40s": 48,
|
|
46
|
+
"a100": 80,
|
|
47
|
+
"h100": 80,
|
|
48
|
+
"h200": 141,
|
|
49
|
+
"b200": 192,
|
|
50
|
+
"5090": 32,
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
# CLI-LOCAL engine -> serving-deps mapping. DELIBERATELY DUPLICATED (like
|
|
55
|
+
# GPU_VRAM_GB above) from urun-python's SOURCE OF TRUTH
|
|
56
|
+
# ``urun/serve/engines/extras.py`` (``ENGINE_DEPS`` / ``serve_deps_for_engine``):
|
|
57
|
+
# the CLI is standalone and cannot import the (often unreleased) urun-python at
|
|
58
|
+
# deploy time, but it MUST know each engine's python deps so it can render them
|
|
59
|
+
# statically into the ``urun serve`` shim's ``@app.function(deps=...)``. Without
|
|
60
|
+
# this the builder's pure-AST deps extractor sees a shim with no declared deps,
|
|
61
|
+
# skips the deps venv, and the serve pod crashloops with no engine installed.
|
|
62
|
+
#
|
|
63
|
+
# KEEP IN SYNC with urun-python ``ENGINE_DEPS``. ``tests/test_serve_catalog.py``
|
|
64
|
+
# pins the exact contents (mirroring the GPU_VRAM_GB sync test); when urun-python
|
|
65
|
+
# bumps a pin, mirror it here and update that test.
|
|
66
|
+
_SERVE_COMMON_DEPS: tuple[str, ...] = (
|
|
67
|
+
"torch==2.6.0",
|
|
68
|
+
"transformers==4.51.3",
|
|
69
|
+
"tokenizers>=0.20.3",
|
|
70
|
+
"accelerate",
|
|
71
|
+
"safetensors",
|
|
72
|
+
"huggingface_hub[hf_transfer]>=0.30.0,<1.0.0",
|
|
73
|
+
# Cap below 2.0 to match the urun wheel's B200-safe floor (urun-python
|
|
74
|
+
# pyproject ``numpy>=1.23.5,<2``, #387). The engine deps and ``urun[deploy]``
|
|
75
|
+
# install into the SAME deps venv; a ``<2.3`` cap resolves numpy 2.2.x and
|
|
76
|
+
# makes the resolver unsatisfiable against the wheel's ``<2`` ("No solution
|
|
77
|
+
# found"). Stay <2 so the two coexist. KEEP IN SYNC with urun-python.
|
|
78
|
+
"numpy<2",
|
|
79
|
+
"pillow>=10.0.0",
|
|
80
|
+
"pycrdt>=0.13.0",
|
|
81
|
+
"orjson",
|
|
82
|
+
)
|
|
83
|
+
_LLAMACPP_CUDA_WHEEL_INDEX = "https://abetlen.github.io/llama-cpp-python/whl/cu124"
|
|
84
|
+
|
|
85
|
+
# The upstream PyTorch prebuilt CUDA 12.4 wheel index (matches the dev L4 sm_89 /
|
|
86
|
+
# CUDA 12.4 runtime base). Engines whose deps do NOT transitively pull a CUDA
|
|
87
|
+
# torch (vLLM/sglang bring their own CUDA torch; ``custom-pytorch`` / diffusers
|
|
88
|
+
# does NOT) MUST pass this so ``torch==2.6.0`` resolves to the ``+cu124`` wheel
|
|
89
|
+
# rather than the PyPI-default ``+cpu`` build (a CPU torch cannot touch the L4).
|
|
90
|
+
_TORCH_CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu124"
|
|
91
|
+
|
|
92
|
+
# Image-engine (custom-pytorch / diffusers) serving deps. Diverges from
|
|
93
|
+
# ``_SERVE_COMMON_DEPS`` on TWO axes the vLLM/text stack does not need:
|
|
94
|
+
# 1. transformers: Qwen-Image-Edit-2511's text encoder is
|
|
95
|
+
# ``Qwen2_5_VLForConditionalGeneration``; transformers 4.51.3 (the vLLM-0.8.5
|
|
96
|
+
# lockstep pin) crashes building its nested ``text_config`` in
|
|
97
|
+
# ``GenerationConfig.from_model_config`` (``'dict' object has no attribute
|
|
98
|
+
# 'to_dict'``). 4.53.x carries the Qwen2.5-VL config fix and still loads
|
|
99
|
+
# under torch 2.6.0.
|
|
100
|
+
# 2. torch: pulled with the cu124 index so it is GPU-capable (diffusers brings
|
|
101
|
+
# no CUDA torch of its own).
|
|
102
|
+
# KEEP IN SYNC with urun-python ``_SERVE_IMAGE_DEPS``.
|
|
103
|
+
_SERVE_IMAGE_DEPS: tuple[str, ...] = (
|
|
104
|
+
"--extra-index-url",
|
|
105
|
+
_TORCH_CUDA_WHEEL_INDEX,
|
|
106
|
+
"torch==2.6.0",
|
|
107
|
+
# torchvision (pinned in lockstep with torch 2.6.0 -> 0.21.0, cu124 build via
|
|
108
|
+
# the index above). Qwen-Image-Edit-2511's processor pulls transformers'
|
|
109
|
+
# AutoVideoProcessor, which hard-gates on torchvision at load — without it
|
|
110
|
+
# DiffusionPipeline.from_pretrained raises "AutoVideoProcessor requires the
|
|
111
|
+
# Torchvision library but it was not found".
|
|
112
|
+
"torchvision==0.21.0",
|
|
113
|
+
"transformers==4.53.3",
|
|
114
|
+
"tokenizers>=0.20.3",
|
|
115
|
+
"accelerate",
|
|
116
|
+
"safetensors",
|
|
117
|
+
"huggingface_hub[hf_transfer]>=0.30.0,<1.0.0",
|
|
118
|
+
# Cap below 2.0 to match the urun wheel's B200-safe floor (see _SERVE_COMMON_DEPS).
|
|
119
|
+
# The image deps + ``urun[deploy]`` install into the SAME venv; ``<2.3`` resolves
|
|
120
|
+
# numpy 2.2.x and conflicts with the wheel's ``<2``. KEEP IN SYNC with urun-python.
|
|
121
|
+
"numpy<2",
|
|
122
|
+
"pillow>=10.0.0",
|
|
123
|
+
"pycrdt>=0.13.0",
|
|
124
|
+
"orjson",
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
# Full-duplex AUDIO (PersonaPlex / Moshi) serving deps. DIVERGES HARD from
|
|
128
|
+
# ``_SERVE_COMMON_DEPS``: PersonaPlex ships its OWN VENDORED ``moshi`` fork
|
|
129
|
+
# (github.com/NVIDIA/personaplex ``moshi/``, package ``moshi-personaplex``,
|
|
130
|
+
# imports as ``moshi``) whose ``loaders.py`` loads the ``nvidia/personaplex-7b-v1``
|
|
131
|
+
# checkpoint with HARDCODED kwargs — the upstream PyPI ``moshi``'s
|
|
132
|
+
# ``CheckpointInfo.from_hf_repo`` crashes on the PersonaPlex stub config
|
|
133
|
+
# (``KeyError: 'dep_q'``). PyPI ``moshi>=0.2.0`` ALSO caps ``aiohttp<3.12`` which
|
|
134
|
+
# conflicts with the urun wheel's ``aiohttp>=3.12`` ("No solution found"). The
|
|
135
|
+
# vendored fork does NOT cap aiohttp. Its pyproject pins ``torch>=2.2,<2.5`` +
|
|
136
|
+
# ``huggingface-hub<0.25`` + ``numpy<2.2`` + ``sentencepiece==0.2`` + ``sphn``. We
|
|
137
|
+
# pin ``numpy<2`` here (NOT the fork's ``<2.2``): the deploy builder installs
|
|
138
|
+
# ``urun[deploy]`` (pyproject ``numpy>=1.23.5,<2`` — the B200-segfault floor) into
|
|
139
|
+
# the SAME deps venv, and a ``<2.2`` cap resolves numpy 2.1.3 → collides with the
|
|
140
|
+
# wheel's ``<2`` ("No solution found"). ``numpy<2`` (1.26.x) satisfies BOTH the
|
|
141
|
+
# fork's ``<2.2`` AND urun's ``<2``. Pinned here as a self-contained group (torch
|
|
142
|
+
# 2.4.1 +cu124 runs on the dev L4). KEEP IN SYNC with urun-python ``_SERVE_MOSHI_DEPS``.
|
|
143
|
+
_SERVE_MOSHI_DEPS: tuple[str, ...] = (
|
|
144
|
+
"--extra-index-url",
|
|
145
|
+
_TORCH_CUDA_WHEEL_INDEX,
|
|
146
|
+
# PersonaPlex's vendored moshi caps torch <2.5; 2.4.1 is the latest cu124 wheel.
|
|
147
|
+
"torch==2.4.1",
|
|
148
|
+
# <2 (not the fork's <2.2): coexist with urun[deploy]'s numpy>=1.23.5,<2 in the
|
|
149
|
+
# shared deps-venv. 1.26.x satisfies both the fork's <2.2 and urun's <2.
|
|
150
|
+
"numpy<2",
|
|
151
|
+
"safetensors>=0.4.0,<0.5",
|
|
152
|
+
"huggingface_hub[hf_transfer]>=0.24,<0.25",
|
|
153
|
+
"sentencepiece==0.2.0",
|
|
154
|
+
"sphn>=0.1.4,<0.2",
|
|
155
|
+
"soundfile>=0.12.1",
|
|
156
|
+
"einops==0.7",
|
|
157
|
+
# The vendored PersonaPlex moshi fork (imports as ``moshi``). Pinned to the repo
|
|
158
|
+
# so the loader + LMGen API match the personaplex-7b-v1 checkpoint exactly.
|
|
159
|
+
"moshi-personaplex @ git+https://github.com/NVIDIA/personaplex.git@main#subdirectory=moshi",
|
|
160
|
+
)
|
|
161
|
+
ENGINE_DEPS: dict[str, list[str]] = {
|
|
162
|
+
"echo": [],
|
|
163
|
+
# CPU echo doubles per modality (image / full-duplex audio). Like ``echo``
|
|
164
|
+
# they need NOTHING heavy: ``echo-image`` is pure Pillow (already in the
|
|
165
|
+
# runtime base), ``echo-audio`` is pure-python PCM16 loopback.
|
|
166
|
+
"echo-image": [],
|
|
167
|
+
"echo-audio": [],
|
|
168
|
+
"vllm": [*_SERVE_COMMON_DEPS, "vllm==0.8.5.post1"],
|
|
169
|
+
"sglang": [*_SERVE_COMMON_DEPS, "sglang[all]>=0.4.0"],
|
|
170
|
+
"llamacpp": [
|
|
171
|
+
*_SERVE_COMMON_DEPS,
|
|
172
|
+
"--extra-index-url",
|
|
173
|
+
_LLAMACPP_CUDA_WHEEL_INDEX,
|
|
174
|
+
"llama-cpp-python>=0.3.2",
|
|
175
|
+
],
|
|
176
|
+
"llamacpp-diffusion": [
|
|
177
|
+
*_SERVE_COMMON_DEPS,
|
|
178
|
+
"--extra-index-url",
|
|
179
|
+
_LLAMACPP_CUDA_WHEEL_INDEX,
|
|
180
|
+
"llama-cpp-python>=0.3.2",
|
|
181
|
+
],
|
|
182
|
+
# custom-pytorch — generic custom-pytorch / diffusers IMAGE shape
|
|
183
|
+
# (Qwen-Image-Edit-Lightning fp8). Needs the diffusion pipeline (``diffusers``)
|
|
184
|
+
# on top of the common torch / transformers / safetensors / accelerate /
|
|
185
|
+
# huggingface_hub / pillow stack. NO llama.cpp / vllm.
|
|
186
|
+
"custom-pytorch": [*_SERVE_IMAGE_DEPS, "diffusers>=0.31.0"],
|
|
187
|
+
"transformers": [*_SERVE_COMMON_DEPS],
|
|
188
|
+
"diffusers": [*_SERVE_COMMON_DEPS, "diffusers"],
|
|
189
|
+
# Full-duplex AUDIO (PersonaPlex / Moshi speech-to-speech). Installs the
|
|
190
|
+
# VENDORED PersonaPlex moshi fork (NOT PyPI ``moshi`` — see _SERVE_MOSHI_DEPS):
|
|
191
|
+
# its loaders.py loads the nvidia/personaplex-7b-v1 checkpoint with the right
|
|
192
|
+
# hardcoded LM/Mimi config, where upstream moshi's CheckpointInfo crashes, and
|
|
193
|
+
# the fork does not cap aiohttp (PyPI moshi caps <3.12 vs the urun wheel >=3.12).
|
|
194
|
+
"moshi": [*_SERVE_MOSHI_DEPS],
|
|
195
|
+
}
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def serve_deps_for_engine(engine: str | None, *, strict: bool = False) -> list[str]:
|
|
199
|
+
"""Return the FLAT python dep list for a catalog ``engine`` value.
|
|
200
|
+
|
|
201
|
+
Mirror of urun-python ``urun.serve.engines.extras.serve_deps_for_engine``:
|
|
202
|
+
normalize (strip + lowercase; empty/None -> ``echo``). By default an UNKNOWN
|
|
203
|
+
engine yields ``[]`` rather than raising — this lenient form exists only for
|
|
204
|
+
the cross-impl parity test.
|
|
205
|
+
|
|
206
|
+
``strict=True`` (the DEPLOY path) instead RAISES on an engine that is not in
|
|
207
|
+
``ENGINE_DEPS``. This closes the (d) non-determinism: the CLI's ``ENGINE_DEPS``
|
|
208
|
+
is a hand-maintained DUPLICATE of urun-python's source of truth, so when the
|
|
209
|
+
catalog row's ``engine`` is a real serving engine the CLI's copy has drifted on
|
|
210
|
+
(or simply does not know), the lenient ``[]`` silently shipped an EMPTY deps
|
|
211
|
+
venv -> ``deps_required=false`` -> urun-only venv -> the serve pod crashloops
|
|
212
|
+
with no vllm/sglang/llama-cpp installed (and INTERMITTENTLY, depending on which
|
|
213
|
+
model/engine was deployed). A known engine that DELIBERATELY needs no deps (e.g.
|
|
214
|
+
``echo``) is listed in ``ENGINE_DEPS`` with ``[]`` and stays a valid no-venv
|
|
215
|
+
deploy; only a genuinely-UNKNOWN engine is rejected.
|
|
216
|
+
"""
|
|
217
|
+
name = (engine or "echo").strip().lower()
|
|
218
|
+
if strict and name not in ENGINE_DEPS:
|
|
219
|
+
known = ", ".join(sorted(ENGINE_DEPS))
|
|
220
|
+
raise UrunError(
|
|
221
|
+
f"serve engine {name!r} is not in the CLI's known serving-engine set "
|
|
222
|
+
f"({known}); refusing to deploy a serve build that would silently ship an "
|
|
223
|
+
f"EMPTY deps venv (no engine installed). The CLI's ENGINE_DEPS map has "
|
|
224
|
+
f"likely drifted from urun-python's source of truth — add {name!r} to "
|
|
225
|
+
f"urun.catalog.ENGINE_DEPS (kept in sync with urun-python "
|
|
226
|
+
f"urun.serve.engines.extras.ENGINE_DEPS)."
|
|
227
|
+
)
|
|
228
|
+
return list(ENGINE_DEPS.get(name, []))
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def parse_gpu_spec(spec: str) -> tuple[str, int]:
|
|
232
|
+
"""Parse a ``<kind>:<count>`` SKU spec into ``(kind, count)``.
|
|
233
|
+
|
|
234
|
+
``count`` defaults to 1 when there is no ``:`` (or a non-integer count),
|
|
235
|
+
mirroring :pyattr:`CatalogRow.gpu_count`. ``kind`` is lowercased so it keys
|
|
236
|
+
into :data:`GPU_VRAM_GB`.
|
|
237
|
+
"""
|
|
238
|
+
kind, _, count = (spec or "").partition(":")
|
|
239
|
+
kind = kind.strip().lower()
|
|
240
|
+
try:
|
|
241
|
+
n = int(count)
|
|
242
|
+
except ValueError:
|
|
243
|
+
n = 1
|
|
244
|
+
return kind, (n if n > 0 else 1)
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
def gpu_vram_gb(spec: str) -> int | None:
|
|
248
|
+
"""Total VRAM (GB) a ``<kind>:<count>`` SKU provides, or None if unknown kind."""
|
|
249
|
+
kind, count = parse_gpu_spec(spec)
|
|
250
|
+
per = GPU_VRAM_GB.get(kind)
|
|
251
|
+
return per * count if per is not None else None
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
# CLI-LOCAL GPU kind -> uv ``--torch-backend`` token. DELIBERATELY DUPLICATED
|
|
255
|
+
# (like GPU_VRAM_GB / ENGINE_DEPS) from urun-python's SOURCE OF TRUTH — the runtime
|
|
256
|
+
# vanilla base image map (``urun.runtime_controller.render._VANILLA_BASE_IMAGE_DEFAULTS``,
|
|
257
|
+
# mirrored in ``urun.serve.deploy._GPU_CLASS_TORCH_BACKEND``). The base tag encodes
|
|
258
|
+
# the CUDA ABI (``cu124-…-sm89`` for L4/L40S, ``cu130-…-sm90`` for H100, ``cu130-…-sm100``
|
|
259
|
+
# for B200), so the app-deps torch index MUST match it or the resolved wheel won't run.
|
|
260
|
+
# `urun serve` bakes the resolved value into the manifest's ``deps.torch_backend`` so
|
|
261
|
+
# the build installs the engine's pinned torch from the GPU-correct index instead of
|
|
262
|
+
# the deploy-side cu130 default (which has NO wheel for the engine's torch on the L4).
|
|
263
|
+
# KEEP IN SYNC with urun-python; pinned by tests/test_serve_catalog.py.
|
|
264
|
+
GPU_TORCH_BACKEND: dict[str, str] = {
|
|
265
|
+
"l4": "cu124", # Ada / sm_89
|
|
266
|
+
"l40": "cu124",
|
|
267
|
+
"l40s": "cu124",
|
|
268
|
+
"a10": "cu124", # Ampere / sm_86 — cu124 wheels cover it
|
|
269
|
+
"a10g": "cu124",
|
|
270
|
+
"a100": "cu124", # Ampere / sm_80
|
|
271
|
+
"h100": "cu130", # Hopper / sm_90
|
|
272
|
+
"h200": "cu130",
|
|
273
|
+
"b200": "cu130", # Blackwell / sm_100
|
|
274
|
+
}
|
|
275
|
+
# Conservative default: the dev fleet is the L4, so an unknown kind must NOT fall
|
|
276
|
+
# back to the broken cu130 default.
|
|
277
|
+
DEFAULT_TORCH_BACKEND = "cu124"
|
|
278
|
+
|
|
279
|
+
|
|
280
|
+
def torch_backend_for_gpu(spec: str) -> str:
|
|
281
|
+
"""Resolve uv's ``--torch-backend`` token for a ``<kind>:<count>`` SKU spec."""
|
|
282
|
+
kind, _count = parse_gpu_spec(spec)
|
|
283
|
+
return GPU_TORCH_BACKEND.get(kind, DEFAULT_TORCH_BACKEND)
|
|
284
|
+
|
|
285
|
+
|
|
286
|
+
@dataclass(frozen=True)
|
|
287
|
+
class CatalogRow:
|
|
288
|
+
"""One flattened ``(model_id, variant, gpu_spec)`` catalog placement.
|
|
289
|
+
|
|
290
|
+
Mirrors ``public.model_catalog`` exactly (see CONTRACT.md). ``is_dev_testable``
|
|
291
|
+
is the contract rule: a single-GPU, non-disagg placement runs on dev.
|
|
292
|
+
"""
|
|
293
|
+
|
|
294
|
+
model_id: str
|
|
295
|
+
variant: str
|
|
296
|
+
gpu_spec: str
|
|
297
|
+
engine: str
|
|
298
|
+
family: str
|
|
299
|
+
task: str
|
|
300
|
+
hf_repo: str
|
|
301
|
+
description: str
|
|
302
|
+
size_gb: int
|
|
303
|
+
files: list[str] | None
|
|
304
|
+
disagg: bool
|
|
305
|
+
targets: dict[str, Any]
|
|
306
|
+
engine_args: dict[str, Any]
|
|
307
|
+
# Placement/serving floors (additive — older catalogs predate them, so they
|
|
308
|
+
# are optional and default to None; published by a separate urun-infra PR).
|
|
309
|
+
# ``default_gpu`` is model-level recommended placement (e.g. "l4:1");
|
|
310
|
+
# ``min_vram_gb`` is the GPU-resident VRAM floor for this placement;
|
|
311
|
+
# ``min_host_ram_gb`` is the host-RAM floor for tiered-offload (GGUF) models.
|
|
312
|
+
default_gpu: str | None = None
|
|
313
|
+
min_vram_gb: int | None = None
|
|
314
|
+
min_host_ram_gb: int | None = None
|
|
315
|
+
|
|
316
|
+
@property
|
|
317
|
+
def ref(self) -> str:
|
|
318
|
+
"""The stable ``<id>:<variant>`` address of this row's model+variant."""
|
|
319
|
+
return f"{self.model_id}:{self.variant}"
|
|
320
|
+
|
|
321
|
+
@property
|
|
322
|
+
def name(self) -> str:
|
|
323
|
+
"""The stable serve key ``<id>:<variant>@<gpu_spec>`` (per CONTRACT.md)."""
|
|
324
|
+
return f"{self.model_id}:{self.variant}@{self.gpu_spec}"
|
|
325
|
+
|
|
326
|
+
@property
|
|
327
|
+
def gpu_count(self) -> int:
|
|
328
|
+
"""GPU count parsed from ``<kind>:<n>`` (e.g. ``b200:8`` -> 8)."""
|
|
329
|
+
_, _, count = self.gpu_spec.partition(":")
|
|
330
|
+
try:
|
|
331
|
+
return int(count)
|
|
332
|
+
except ValueError:
|
|
333
|
+
return 1
|
|
334
|
+
|
|
335
|
+
@property
|
|
336
|
+
def is_dev_testable(self) -> bool:
|
|
337
|
+
"""Contract rule: 1-GPU, non-disagg placements are dev-testable."""
|
|
338
|
+
return self.gpu_count == 1 and not self.disagg
|
|
339
|
+
|
|
340
|
+
|
|
341
|
+
def parse_row(raw: dict[str, Any]) -> CatalogRow:
|
|
342
|
+
"""Build a :class:`CatalogRow` from one PostgREST JSON object.
|
|
343
|
+
|
|
344
|
+
Tolerant of extra columns (e.g. ``updated_at``) but strict about the
|
|
345
|
+
contract fields: a missing/invalid required field is a loud error so a
|
|
346
|
+
malformed catalog never silently produces a degraded serve.
|
|
347
|
+
"""
|
|
348
|
+
if not isinstance(raw, dict):
|
|
349
|
+
raise UrunError("catalog row is not a JSON object")
|
|
350
|
+
|
|
351
|
+
def req_str(key: str) -> str:
|
|
352
|
+
value = raw.get(key)
|
|
353
|
+
if not isinstance(value, str) or not value:
|
|
354
|
+
raise UrunError(f"catalog row missing required string field {key!r}")
|
|
355
|
+
return value
|
|
356
|
+
|
|
357
|
+
files = raw.get("files")
|
|
358
|
+
if files is not None and not isinstance(files, list):
|
|
359
|
+
raise UrunError("catalog row 'files' must be a JSON array or null")
|
|
360
|
+
size_gb = raw.get("size_gb")
|
|
361
|
+
if not isinstance(size_gb, int) or isinstance(size_gb, bool):
|
|
362
|
+
raise UrunError("catalog row 'size_gb' must be an integer")
|
|
363
|
+
targets = raw.get("targets") or {}
|
|
364
|
+
engine_args = raw.get("engine_args") or {}
|
|
365
|
+
if not isinstance(targets, dict) or not isinstance(engine_args, dict):
|
|
366
|
+
raise UrunError("catalog row 'targets'/'engine_args' must be JSON objects")
|
|
367
|
+
|
|
368
|
+
# Additive placement floors: tolerant, NOT required (older catalogs lack them).
|
|
369
|
+
default_gpu_raw = raw.get("default_gpu")
|
|
370
|
+
default_gpu = default_gpu_raw if isinstance(default_gpu_raw, str) and default_gpu_raw else None
|
|
371
|
+
|
|
372
|
+
def opt_int(key: str) -> int | None:
|
|
373
|
+
value = raw.get(key)
|
|
374
|
+
return value if isinstance(value, int) and not isinstance(value, bool) else None
|
|
375
|
+
|
|
376
|
+
return CatalogRow(
|
|
377
|
+
model_id=req_str("model_id"),
|
|
378
|
+
variant=req_str("variant"),
|
|
379
|
+
gpu_spec=req_str("gpu_spec"),
|
|
380
|
+
engine=req_str("engine"),
|
|
381
|
+
family=req_str("family"),
|
|
382
|
+
task=req_str("task"),
|
|
383
|
+
hf_repo=req_str("hf_repo"),
|
|
384
|
+
description=raw.get("description") if isinstance(raw.get("description"), str) else "",
|
|
385
|
+
size_gb=size_gb,
|
|
386
|
+
files=[str(f) for f in files] if files is not None else None,
|
|
387
|
+
disagg=bool(raw.get("disagg", False)),
|
|
388
|
+
targets=targets,
|
|
389
|
+
engine_args=engine_args,
|
|
390
|
+
default_gpu=default_gpu,
|
|
391
|
+
min_vram_gb=opt_int("min_vram_gb"),
|
|
392
|
+
min_host_ram_gb=opt_int("min_host_ram_gb"),
|
|
393
|
+
)
|
|
394
|
+
|
|
395
|
+
|
|
396
|
+
def split_ref(ref: str) -> tuple[str, str | None]:
|
|
397
|
+
"""Split ``<id>`` or ``<id>:<variant>`` into ``(model_id, variant|None)``.
|
|
398
|
+
|
|
399
|
+
Empty id or empty variant (a trailing/leading colon) is a loud error.
|
|
400
|
+
"""
|
|
401
|
+
text = (ref or "").strip()
|
|
402
|
+
if not text:
|
|
403
|
+
raise UrunError("missing model id; expected <id> or <id>:<variant>")
|
|
404
|
+
model_id, sep, variant = text.partition(":")
|
|
405
|
+
model_id = model_id.strip()
|
|
406
|
+
if not model_id:
|
|
407
|
+
raise UrunError(f"invalid model reference {ref!r}: empty model id")
|
|
408
|
+
if not sep:
|
|
409
|
+
return model_id, None
|
|
410
|
+
variant = variant.strip()
|
|
411
|
+
if not variant:
|
|
412
|
+
raise UrunError(f"invalid model reference {ref!r}: empty variant after ':'")
|
|
413
|
+
if ":" in variant:
|
|
414
|
+
raise UrunError(f"invalid model reference {ref!r}: at most one ':' separator")
|
|
415
|
+
return model_id, variant
|
|
416
|
+
|
|
417
|
+
|
|
418
|
+
def resolve_row(
|
|
419
|
+
rows: list[CatalogRow],
|
|
420
|
+
model_id: str,
|
|
421
|
+
variant: str | None,
|
|
422
|
+
gpu_spec: str | None,
|
|
423
|
+
) -> CatalogRow:
|
|
424
|
+
"""Resolve ``(model_id, variant?, gpu?)`` to exactly one catalog row.
|
|
425
|
+
|
|
426
|
+
Defaults (per CONTRACT.md): no variant -> the model's FIRST listed variant;
|
|
427
|
+
no ``--gpu`` -> that variant's FIRST listed placement. Unknown id, unknown
|
|
428
|
+
variant, or an unmatched ``--gpu`` each raise a clear, enumerated error.
|
|
429
|
+
"""
|
|
430
|
+
for_model = [r for r in rows if r.model_id == model_id]
|
|
431
|
+
if not for_model:
|
|
432
|
+
known = ", ".join(sorted({r.model_id for r in rows})) or "(catalog is empty)"
|
|
433
|
+
raise UrunError(f"unknown model id {model_id!r}. Known models: {known}")
|
|
434
|
+
|
|
435
|
+
if variant is None:
|
|
436
|
+
# First listed variant — PostgREST returns rows already PK-ordered.
|
|
437
|
+
variant = for_model[0].variant
|
|
438
|
+
for_variant = [r for r in for_model if r.variant == variant]
|
|
439
|
+
if not for_variant:
|
|
440
|
+
known = ", ".join(sorted({r.variant for r in for_model}))
|
|
441
|
+
raise UrunError(
|
|
442
|
+
f"unknown variant {variant!r} for model {model_id!r}. Known variants: {known}"
|
|
443
|
+
)
|
|
444
|
+
|
|
445
|
+
if gpu_spec is None:
|
|
446
|
+
return for_variant[0]
|
|
447
|
+
for_gpu = [r for r in for_variant if r.gpu_spec == gpu_spec]
|
|
448
|
+
if not for_gpu:
|
|
449
|
+
known = ", ".join(sorted({r.gpu_spec for r in for_variant}))
|
|
450
|
+
raise UrunError(
|
|
451
|
+
f"no placement on GPU {gpu_spec!r} for {model_id}:{variant}. "
|
|
452
|
+
f"Available placements: {known}"
|
|
453
|
+
)
|
|
454
|
+
return for_gpu[0]
|
|
455
|
+
|
|
456
|
+
|
|
457
|
+
def _host_ram_note(row: CatalogRow) -> str:
|
|
458
|
+
"""Informational host-RAM NOTE for tiered-offload models, or "" if unknown.
|
|
459
|
+
|
|
460
|
+
The CLI cannot know the target node's RAM, so ``min_host_ram_gb`` is surfaced
|
|
461
|
+
as advice (especially for GGUF/llamacpp tiered-offload placements).
|
|
462
|
+
"""
|
|
463
|
+
if not isinstance(row.min_host_ram_gb, int) or row.min_host_ram_gb <= 0:
|
|
464
|
+
return ""
|
|
465
|
+
return f" Note: needs >= {row.min_host_ram_gb}GB host RAM (tiered offload)."
|
|
466
|
+
|
|
467
|
+
|
|
468
|
+
def validate_gpu_against_row(
|
|
469
|
+
rows: list[CatalogRow],
|
|
470
|
+
model_id: str,
|
|
471
|
+
variant: str | None,
|
|
472
|
+
requested_gpu: str,
|
|
473
|
+
) -> tuple[str, str]:
|
|
474
|
+
"""Validate a requested ``--gpu`` SKU against a variant's placement floors.
|
|
475
|
+
|
|
476
|
+
Returns ``(level, message)`` where ``level`` is ``"error" | "warn" | "ok"``:
|
|
477
|
+
|
|
478
|
+
* unknown SKU kind -> error (lists known SKUs)
|
|
479
|
+
* available VRAM < known ``min_vram_gb`` -> error (numbers)
|
|
480
|
+
* clears floor but ``--gpu`` is NOT an enumerated
|
|
481
|
+
placement for the variant -> warn (untuned)
|
|
482
|
+
* ``--gpu`` IS an enumerated placement -> ok (clean)
|
|
483
|
+
|
|
484
|
+
A known ``min_host_ram_gb`` is appended to the message as an informational
|
|
485
|
+
NOTE (the CLI can't know the node's RAM). Floors that are null fall back to
|
|
486
|
+
the placement-membership decision only.
|
|
487
|
+
"""
|
|
488
|
+
# Resolve the variant (gpu_spec=None -> representative row); reuse resolve_row
|
|
489
|
+
# so unknown id/variant produce the same enumerated errors as the deploy path.
|
|
490
|
+
representative = resolve_row(rows, model_id, variant, None)
|
|
491
|
+
resolved_variant = representative.variant
|
|
492
|
+
for_variant = [r for r in rows if r.model_id == model_id and r.variant == resolved_variant]
|
|
493
|
+
enumerated = sorted({r.gpu_spec for r in for_variant})
|
|
494
|
+
ref = f"{model_id}:{resolved_variant}"
|
|
495
|
+
|
|
496
|
+
available = gpu_vram_gb(requested_gpu)
|
|
497
|
+
if available is None:
|
|
498
|
+
kind, _ = parse_gpu_spec(requested_gpu)
|
|
499
|
+
known = ", ".join(sorted(GPU_VRAM_GB))
|
|
500
|
+
return (
|
|
501
|
+
"error",
|
|
502
|
+
f"unknown GPU SKU {kind!r} in {requested_gpu!r}. Known SKUs: {known}.",
|
|
503
|
+
)
|
|
504
|
+
|
|
505
|
+
# The matching placement (exact --gpu) drives the floor; otherwise use the
|
|
506
|
+
# variant's minimum (representative) min_vram_gb across its placements.
|
|
507
|
+
exact = next((r for r in for_variant if r.gpu_spec == requested_gpu), None)
|
|
508
|
+
floor_row = exact
|
|
509
|
+
if floor_row is None or floor_row.min_vram_gb is None:
|
|
510
|
+
with_floor = [r for r in for_variant if isinstance(r.min_vram_gb, int)]
|
|
511
|
+
if with_floor:
|
|
512
|
+
floor_row = min(with_floor, key=lambda r: r.min_vram_gb) # type: ignore[arg-type,return-value]
|
|
513
|
+
min_vram = floor_row.min_vram_gb if floor_row is not None else None
|
|
514
|
+
host_note = _host_ram_note(exact or representative)
|
|
515
|
+
|
|
516
|
+
if isinstance(min_vram, int) and available < min_vram:
|
|
517
|
+
return (
|
|
518
|
+
"error",
|
|
519
|
+
f"{requested_gpu} provides {available}GB VRAM but {ref} needs "
|
|
520
|
+
f">= {min_vram}GB (min_vram_gb). Choose a larger GPU.{host_note}",
|
|
521
|
+
)
|
|
522
|
+
|
|
523
|
+
if exact is not None:
|
|
524
|
+
# Enumerated, tuned placement that clears (or has no) floor -> clean.
|
|
525
|
+
return ("ok", f"{requested_gpu} is a tuned placement for {ref}.{host_note}")
|
|
526
|
+
|
|
527
|
+
# Not enumerated, but clears the floor (or no floor known) -> untuned warn.
|
|
528
|
+
floor_clause = f"clears the {min_vram}GB floor but " if isinstance(min_vram, int) else ""
|
|
529
|
+
return (
|
|
530
|
+
"warn",
|
|
531
|
+
f"{requested_gpu} ({available}GB) {floor_clause}is not a tuned placement "
|
|
532
|
+
f"for {ref} (enumerated: {', '.join(enumerated)}). Proceeding untuned.{host_note}",
|
|
533
|
+
)
|
|
534
|
+
|
|
535
|
+
|
|
536
|
+
def valid_skus_for_row(row: CatalogRow) -> list[str]:
|
|
537
|
+
"""Single-device SKUs from :data:`GPU_VRAM_GB` that clear ``row.min_vram_gb``.
|
|
538
|
+
|
|
539
|
+
Reuses the SAME VRAM-floor logic as :func:`validate_gpu_against_row`: a SKU is
|
|
540
|
+
"valid on" this placement when one device's VRAM (``GPU_VRAM_GB[kind] * 1``)
|
|
541
|
+
is at least ``min_vram_gb``. With no known floor, every SKU is considered
|
|
542
|
+
valid (the placement imposes no VRAM constraint we can check). Returned as a
|
|
543
|
+
sorted ``<kind>:1`` list — the per-device floor a single card must clear.
|
|
544
|
+
"""
|
|
545
|
+
floor = row.min_vram_gb
|
|
546
|
+
if not isinstance(floor, int) or floor <= 0:
|
|
547
|
+
return [f"{kind}:1" for kind in sorted(GPU_VRAM_GB)]
|
|
548
|
+
return [f"{kind}:1" for kind, vram in sorted(GPU_VRAM_GB.items()) if vram >= floor]
|
|
549
|
+
|
|
550
|
+
|
|
551
|
+
class CatalogClient:
|
|
552
|
+
"""Reads ``public.model_catalog`` over PostgREST.
|
|
553
|
+
|
|
554
|
+
``base_url`` is the PostgREST base (``.../rest/v1``); ``anon_key`` is the
|
|
555
|
+
public Supabase anon key sent as the ``apikey`` header (the table grants
|
|
556
|
+
SELECT to ``anon``). Both are resolved by the CLI's single canonical path
|
|
557
|
+
(see :func:`urun.serve.resolve_catalog_client`).
|
|
558
|
+
"""
|
|
559
|
+
|
|
560
|
+
def __init__(
|
|
561
|
+
self,
|
|
562
|
+
base_url: str,
|
|
563
|
+
anon_key: str,
|
|
564
|
+
timeout: float = 30.0,
|
|
565
|
+
backoff: BackoffPolicy = DEPLOY_BACKOFF,
|
|
566
|
+
on_retry: Callable[[str], None] | None = None,
|
|
567
|
+
):
|
|
568
|
+
require_safe_url(base_url, "catalog URL")
|
|
569
|
+
if not anon_key:
|
|
570
|
+
raise UrunError("missing catalog anon key")
|
|
571
|
+
self.base_url = base_url.rstrip("/")
|
|
572
|
+
self.anon_key = anon_key
|
|
573
|
+
self.timeout = timeout
|
|
574
|
+
self.backoff = backoff
|
|
575
|
+
self.on_retry = on_retry
|
|
576
|
+
|
|
577
|
+
def select_rows(self) -> list[CatalogRow]:
|
|
578
|
+
"""Select the full catalog, PK-ordered, as :class:`CatalogRow` objects.
|
|
579
|
+
|
|
580
|
+
Raises a clear, actionable error when the table is not seeded yet (it
|
|
581
|
+
ships via a urun-infra migration) instead of a raw PostgREST 404.
|
|
582
|
+
"""
|
|
583
|
+
params = urllib.parse.urlencode(
|
|
584
|
+
{"select": "*", "order": "model_id,variant,gpu_spec"}, safe=",:"
|
|
585
|
+
)
|
|
586
|
+
payload = self._get(f"/{CATALOG_TABLE}?{params}")
|
|
587
|
+
if not isinstance(payload, list):
|
|
588
|
+
raise UrunError("catalog response was not a JSON array")
|
|
589
|
+
return [parse_row(item) for item in payload]
|
|
590
|
+
|
|
591
|
+
def _get(self, path: str) -> Any:
|
|
592
|
+
def attempt() -> bytes:
|
|
593
|
+
req = urllib.request.Request(
|
|
594
|
+
self.base_url + path,
|
|
595
|
+
method="GET",
|
|
596
|
+
headers={
|
|
597
|
+
"apikey": self.anon_key,
|
|
598
|
+
"Authorization": f"Bearer {self.anon_key}",
|
|
599
|
+
"Accept": "application/json",
|
|
600
|
+
"User-Agent": f"urun-cli/{__version__}",
|
|
601
|
+
},
|
|
602
|
+
)
|
|
603
|
+
try:
|
|
604
|
+
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
|
605
|
+
return resp.read()
|
|
606
|
+
except urllib.error.HTTPError as exc:
|
|
607
|
+
raise _catalog_http_error(exc) from exc
|
|
608
|
+
|
|
609
|
+
try:
|
|
610
|
+
raw = retry_call(
|
|
611
|
+
attempt,
|
|
612
|
+
policy=self.backoff,
|
|
613
|
+
on_retry=self.on_retry,
|
|
614
|
+
description=f"GET {path}",
|
|
615
|
+
)
|
|
616
|
+
except (urllib.error.URLError, TimeoutError) as exc:
|
|
617
|
+
reason = getattr(exc, "reason", exc)
|
|
618
|
+
raise UrunError(f"network error reading catalog: {reason}") from exc
|
|
619
|
+
|
|
620
|
+
if not raw:
|
|
621
|
+
return []
|
|
622
|
+
try:
|
|
623
|
+
return json.loads(raw.decode("utf-8"))
|
|
624
|
+
except json.JSONDecodeError as exc:
|
|
625
|
+
raise UrunError("invalid JSON response from catalog API") from exc
|
|
626
|
+
|
|
627
|
+
|
|
628
|
+
# PostgREST raises 404 with code PGRST205 ("could not find the table ...") when
|
|
629
|
+
# the relation is absent — i.e. the seed migration has not shipped to this
|
|
630
|
+
# control plane yet. Translate it to a clear, actionable message.
|
|
631
|
+
_UNSEEDED_MESSAGE = (
|
|
632
|
+
"the model catalog is not available on this control plane yet "
|
|
633
|
+
"(public.model_catalog is unseeded). It ships via a urun-infra migration; "
|
|
634
|
+
"try again once the catalog has been published."
|
|
635
|
+
)
|
|
636
|
+
|
|
637
|
+
|
|
638
|
+
def _catalog_http_error(exc: urllib.error.HTTPError) -> UrunError:
|
|
639
|
+
payload = exc.read()
|
|
640
|
+
code = "http_error"
|
|
641
|
+
message = payload.decode("utf-8", "replace") or exc.reason
|
|
642
|
+
try:
|
|
643
|
+
parsed = json.loads(payload.decode("utf-8"))
|
|
644
|
+
if isinstance(parsed, dict):
|
|
645
|
+
code = parsed.get("code") or code
|
|
646
|
+
message = parsed.get("message") or message
|
|
647
|
+
except Exception:
|
|
648
|
+
pass
|
|
649
|
+
# PGRST205 = missing relation; a bare 404 on the table path means the same.
|
|
650
|
+
if exc.code == 404 or (isinstance(code, str) and code == "PGRST205"):
|
|
651
|
+
return UrunError(_UNSEEDED_MESSAGE)
|
|
652
|
+
return ApiError(exc.code, str(code), str(message))
|