From 8b051c19ebbee2068fdcbbb6c2d7e14773a50c0b Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Wed, 12 Aug 2026 18:46:27 +0000 Subject: [PATCH] databricks: scope model-services listing to system.ai schema list_model_services walked the entire metastore (hundreds of services across ~23 sequential ~2s pages, ~50s on a busy workspace) only to keep the system.ai.* subset client-side. Passing parent=schemas/system.ai scopes it server-side: same 53 ids come back in a single ~1s page. This backs both ucode setup ("Fetching available models") and ucode apply (manifest validation), so both drop from ~50s to ~1s. Co-authored-by: Isaac --- src/ucode/databricks.py | 32 +++++++++++++++++++++++--------- tests/test_databricks.py | 3 +++ 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 705b5519..eb709236 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1328,10 +1328,18 @@ def build_auth_shell_command( # part after the prefix is exactly the model string agents send (no # `databricks-` infix — that only appears on the inner destination name). _MODEL_SERVICE_NAME_PREFIX = "model-services/" -# The metastore-scope listing returns services from EVERY schema (e.g. -# `main.user.foo`, `temp.*`, internal DLT schemas). We only want the -# Databricks-managed foundation models under `system.ai`. +# The listing can return services from EVERY schema (e.g. `main.user.foo`, +# `temp.*`, internal DLT schemas). We only want the Databricks-managed +# foundation models under `system.ai`. _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai." +# Scope the listing to the `system.ai` schema via the `parent` query param +# (`schemas/{catalog}.{schema}`). Without it the endpoint walks the ENTIRE +# metastore — hundreds of unrelated services across dozens of ~2s pages, then +# discards all but `system.ai.*` client-side (a ~50s walk on a busy workspace). +# Parent-scoped, the same set comes back in a single page (~1s). The endpoint +# ignores the other filters (`catalog_name`/`schema_name`/`filter`), so `parent` +# is the only server-side narrowing that works. +_MODEL_SERVICE_PARENT_SCHEMA = "schemas/system.ai" # Supported OSS chat families, matched by name substring. Add an entry to # support a new family. @@ -1473,11 +1481,14 @@ def list_model_services( ) -> tuple[list[str], str | None]: """List all `system.ai.*` model ids via the UC model-services API. - Pages through ``/api/2.1/unity-catalog/model-services`` (metastore scope) - with a bounded ``page_size`` (the endpoint 499s without one) and returns the - de-duplicated, sorted list of ``system.ai.`` ids. Returns - (ids, reason); reason is None on success, otherwise it describes why the - list is empty (HTTP/network error or no services). + Pages through ``/api/2.1/unity-catalog/model-services`` scoped to the + ``system.ai`` schema (``parent=schemas/system.ai``) with a bounded + ``page_size`` (the endpoint 499s without one) and returns the de-duplicated, + sorted list of ``system.ai.`` ids. Returns (ids, reason); reason + is None on success, otherwise it describes why the list is empty (HTTP/network + error or no services). Scoping matters: the unscoped metastore listing walks + every schema across dozens of ~2s pages (~50s on a busy workspace) only to + keep the same ``system.ai.*`` subset — see ``_MODEL_SERVICE_PARENT_SCHEMA``. A successful result is memoized per workspace for the life of the process; pass ``use_cache=False`` to force a fresh walk. @@ -1493,7 +1504,10 @@ def list_model_services( seen_tokens: set[str] = set() last_reason: str | None = None for _ in range(max_pages): - params: dict[str, str] = {"page_size": str(page_size)} + params: dict[str, str] = { + "parent": _MODEL_SERVICE_PARENT_SCHEMA, + "page_size": str(page_size), + } if page_token: params["page_token"] = page_token url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0b51f8bc..e1a7f67a 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -355,6 +355,9 @@ def fake_get(url, token, timeout=10): assert ids == ["system.ai.gpt-5"] assert reason is None assert all("page_size=" in u for u in urls) + # Scope to the `system.ai` schema so the endpoint returns just the + # foundation models rather than walking the whole metastore. + assert all("parent=schemas%2Fsystem.ai" in u for u in urls) def test_retries_page_before_giving_up(self, monkeypatch): payload = {"model_services": [_model_service("system.ai.gpt-5")]}