From 04f9230eda39b032633eaad1875310c6983612e7 Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Wed, 21 Oct 2026 22:36:02 -0700 Subject: [PATCH 1/6] A Kumo Tabular quickstart notebook for Databricks Deploys the tabular model to a one-GPU Model Serving endpoint that scales to zero, then asks the AdventureWorks products table a color classification, a list price regression, and the color question in PQL. Every Unity Catalog name part is quoted with backticks. Signed-off-by: Adi Wadaskar --- .../databricks/kumo_tabular_quickstart.ipynb | 649 ++++++++++++++++++ 1 file changed, 649 insertions(+) create mode 100644 examples/databricks/kumo_tabular_quickstart.ipynb diff --git a/examples/databricks/kumo_tabular_quickstart.ipynb b/examples/databricks/kumo_tabular_quickstart.ipynb new file mode 100644 index 000000000..da2ef50e6 --- /dev/null +++ b/examples/databricks/kumo_tabular_quickstart.ipynb @@ -0,0 +1,649 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "18beb70fdc90345950ccb127fd06ff75", + "metadata": {}, + "source": [ + "# Get started with Kumo Tabular\n", + "\n", + "Kumo Tabular is a foundation model for a single table: it fills in a column from the other columns, with no feature engineering and no model trained for each task.\n", + "You give it rows where the answer is known and rows where it is not, and it predicts the missing answers from the known ones.\n", + "\n", + "This notebook deploys the model in your workspace, then asks three questions of Microsoft's AdventureWorks products table: which color a product has, what its list price is, and the color question again as a predictive query.\n", + "Each prediction is checked against the value the table records, and compared with a simple rule, such as always guessing the most common color." + ] + }, + { + "cell_type": "markdown", + "id": "94b63d84cfc7b541f6144d1d49a85e45", + "metadata": {}, + "source": [ + "## How Kumo Tabular works\n", + "\n", + "> **1. The input is one table.**\n", + "\n", + "The table holds numerical and categorical columns, such as a product's weight and its size.\n", + "The model reads no related tables and no time windows, so it needs no graph: rows in, answers out.\n", + "\n", + "> **2. Rows where the answer is known are the context.**\n", + "\n", + "The rows where the target column is filled in are the context examples.\n", + "The rows you want answered are the rows to score, and have the target column left out.\n", + "\n", + "> **3. A prediction is classification or regression.**\n", + "\n", + "A categorical target, such as a color, is a classification and returns a predicted class with a probability for each class.\n", + "A numerical target, such as a price, is a regression and returns a number.\n", + "\n", + "There is no training step: for each request, the client sends the context rows with the rows to score, and the model predicts from them." + ] + }, + { + "cell_type": "markdown", + "id": "1d2d1413ac076a95031bac81a0c15417", + "metadata": {}, + "source": [ + "## Before you start\n", + "\n", + "You need:\n", + "\n", + "- the Kumo Tabular model, in Unity Catalog or installed from Databricks Marketplace;\n", + "- permission to create a Model Serving endpoint, in a region with GPU Model Serving;\n", + "- a catalog and schema where you can create tables;\n", + "- serverless compute with environment 4 or later, or Databricks Runtime 17.3 LTS or later.\n", + "\n", + "The first deployment of the endpoint takes 20 to 40 minutes, then the rest of the notebook runs in a few minutes.\n", + "The endpoint bills for GPU time while it runs; it scales to zero when idle, and the last step shows how to delete it.\n", + "\n", + "Run each cell in order with Shift+Enter.\n", + "\n", + "**Let's get started!**" + ] + }, + { + "cell_type": "markdown", + "id": "7a0c878d48c1b6d0bebbe0a2d7bf40a4", + "metadata": {}, + "source": [ + "## Step 1. Install the client\n", + "\n", + "Install the client, the `structured-data-models-client` package, with its Databricks extras, then restart Python.\n", + "pip may print a red message about `protobuf` versions; it does not affect this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5445169ae07af0aa529f51f1cb835515", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install --quiet \"structured-data-models-client[databricks,databricks-serving]\"\n", + "dbutils.library.restartPython()" + ] + }, + { + "cell_type": "markdown", + "id": "c8a993bc41f4ddcb078426d5e83070f2", + "metadata": {}, + "source": [ + "## Step 2. Set your parameters\n", + "\n", + "Set these names, then run the cell.\n", + "`MODEL_NAME` is the Kumo Tabular model, in the catalog you installed it into, and `CATALOG` and `SCHEMA` are where the sample table goes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e8e09e66cfe3045b176857b3dd28c3a", + "metadata": {}, + "outputs": [], + "source": [ + "from datetime import timedelta\n", + "\n", + "import pandas as pd\n", + "from databricks.sdk import WorkspaceClient\n", + "from databricks.sdk.errors import NotFound\n", + "from databricks.sdk.service.serving import (\n", + " EndpointCoreConfigInput,\n", + " ServedEntityInput,\n", + " ServingModelWorkloadType,\n", + ")\n", + "from sdm_client import (\n", + " KumoTabular,\n", + " PredictiveQuery,\n", + " Request,\n", + " Schema,\n", + " SDMClient,\n", + ")\n", + "from sdm_connectors import quote_ident\n", + "\n", + "# The model, as .., and its version.\n", + "MODEL_NAME = 'kumo_tabular.models.kumo_tabular'\n", + "MODEL_VERSION = '1'\n", + "ENDPOINT_NAME = 'kumo-tabular'\n", + "# Where the sample table goes: a catalog and schema you can create tables in.\n", + "CATALOG = 'main'\n", + "SCHEMA = 'kumo_tabular_quickstart'\n", + "TABLE = 'adventureworks_products'\n", + "\n", + "\n", + "def quoted(*parts):\n", + " # A Unity Catalog name, each part quoted for SQL.\n", + " return '.'.join(quote_ident(part, char='`') for part in parts)\n", + "\n", + "\n", + "PRODUCTS = quoted(CATALOG, SCHEMA, TABLE)" + ] + }, + { + "cell_type": "markdown", + "id": "6c222fe7761b0460ac651bf4d49baea1", + "metadata": {}, + "source": [ + "## Step 3. Deploy the model\n", + "\n", + "The next cell creates a Model Serving endpoint for the model in your workspace, on one GPU (`GPU_MEDIUM`, size Small) that scales to zero when idle.\n", + "It does not wait for the deployment; Step 6 does.\n", + "If the endpoint already exists, the cell checks that it serves this model.\n", + "You can also create it with **Serve this model** on the model's page in Catalog Explorer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a133dc9b0e07efd33ac4da64235bcbdb", + "metadata": {}, + "outputs": [], + "source": [ + "workspace = WorkspaceClient()\n", + "try:\n", + " endpoint = workspace.serving_endpoints.get(ENDPOINT_NAME)\n", + "except NotFound:\n", + " workspace.serving_endpoints.create(\n", + " name=ENDPOINT_NAME,\n", + " config=EndpointCoreConfigInput(\n", + " name=ENDPOINT_NAME,\n", + " served_entities=[\n", + " ServedEntityInput(\n", + " entity_name=MODEL_NAME,\n", + " entity_version=MODEL_VERSION,\n", + " workload_type=ServingModelWorkloadType.GPU_MEDIUM,\n", + " workload_size='Small',\n", + " scale_to_zero_enabled=True,\n", + " )\n", + " ],\n", + " ),\n", + " )\n", + " print(\n", + " f'Creating the endpoint {ENDPOINT_NAME}; this takes 20 to 40 minutes.'\n", + " )\n", + "else:\n", + " config = endpoint.config or endpoint.pending_config\n", + " served = {\n", + " (entity.entity_name.lower(), entity.entity_version)\n", + " for entity in config.served_entities\n", + " }\n", + " if served != {(MODEL_NAME.lower(), MODEL_VERSION)}:\n", + " raise ValueError(\n", + " f'The endpoint {ENDPOINT_NAME} already exists and serves {sorted(served)}, '\n", + " f'not {MODEL_NAME} version {MODEL_VERSION}. Set ENDPOINT_NAME to another name.'\n", + " )\n", + " print(f'The endpoint {ENDPOINT_NAME} already exists and serves this model.')" + ] + }, + { + "cell_type": "markdown", + "id": "df62876d09335df77ab439e3935f6141", + "metadata": {}, + "source": [ + "## Step 4. Load the sample data\n", + "\n", + "The next cell downloads the AdventureWorks products from the [Microsoft SQL Server samples repository](https://github.com/microsoft/sql-server-samples/tree/master/samples/databases/adventure-works/oltp-install-script) at a fixed commit and saves them as a table in your schema.\n", + "It needs internet access, and skips the download if the table already exists with the expected columns.\n", + "\n", + "> Each user is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n", + "> The repository's samples are provided under the [MIT License](https://github.com/microsoft/sql-server-samples/blob/master/license.txt)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1c9a61987b2e78e004b7da1d055a40b2", + "metadata": {}, + "outputs": [], + "source": [ + "URL = (\n", + " 'https://raw.githubusercontent.com/microsoft/sql-server-samples/'\n", + " 'b47eadc85277f72843647223cde53aa1b85ffcb8/samples/databases/'\n", + " 'adventure-works/oltp-install-script/Product.csv'\n", + ")\n", + "COLUMNS = (\n", + " 'ProductID Name ProductNumber MakeFlag FinishedGoodsFlag Color '\n", + " 'SafetyStockLevel ReorderPoint StandardCost ListPrice Size '\n", + " 'SizeUnitMeasureCode WeightUnitMeasureCode Weight DaysToManufacture '\n", + " 'ProductLine Class Style ProductSubcategoryID ProductModelID '\n", + " 'SellStartDate SellEndDate DiscontinuedDate rowguid ModifiedDate'\n", + ").split()\n", + "\n", + "if spark.catalog.tableExists(PRODUCTS):\n", + " if spark.table(PRODUCTS).columns != COLUMNS:\n", + " raise ValueError(\n", + " f'{PRODUCTS} already exists with other columns. '\n", + " 'Drop it, or set SCHEMA to another schema.'\n", + " )\n", + " print(f'The sample table is already in {CATALOG}.{SCHEMA}.')\n", + "else:\n", + " frame = pd.read_csv(\n", + " URL,\n", + " sep='\\t',\n", + " header=None,\n", + " names=COLUMNS,\n", + " dtype_backend='numpy_nullable',\n", + " )\n", + " for column in frame.columns:\n", + " if column.endswith('Date'):\n", + " frame[column] = pd.to_datetime(frame[column])\n", + " spark.sql(f'CREATE SCHEMA IF NOT EXISTS {quoted(CATALOG, SCHEMA)}')\n", + " spark.createDataFrame(frame).write.mode('ignore').saveAsTable(PRODUCTS)\n", + " print(f'Saved the sample table to {CATALOG}.{SCHEMA}.')" + ] + }, + { + "cell_type": "markdown", + "id": "fb44d873a5f8030a3e185da590015cb5", + "metadata": {}, + "source": [ + "## Step 5. Read the table and choose the features\n", + "\n", + "The model reads one table, so the table is read into the notebook: unlike the relational model, which samples linked tables where they are, there is nothing to sample here.\n", + "A table small enough to score is simply read.\n", + "\n", + "The model reads numerical and categorical columns.\n", + "The cell leaves out the rest: the product name, product number and `rowguid`, which name a row rather than describe it, and the dates, which are not numbers or categories.\n", + "It also leaves out `ProductID`, a row identifier, which Example 3 adds back as the key of a predictive query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7015e2d235fde236028420efe0c29b31", + "metadata": {}, + "outputs": [], + "source": [ + "products = spark.table(PRODUCTS).toPandas()\n", + "LEFT_OUT = [\n", + " 'ProductID',\n", + " 'rowguid',\n", + " 'ModifiedDate',\n", + " 'Name',\n", + " 'ProductNumber',\n", + " 'SellStartDate',\n", + " 'SellEndDate',\n", + " 'DiscontinuedDate',\n", + "]\n", + "features = products.drop(columns=LEFT_OUT)\n", + "print(f'{len(products)} products; features: {\", \".join(features.columns)}')\n", + "products.head()" + ] + }, + { + "cell_type": "markdown", + "id": "7bef8d59a20ee6a31c5b8f3e705a62a3", + "metadata": {}, + "source": [ + "## Step 6. Connect to the model\n", + "\n", + "The next cell waits until the endpoint is ready, then connects to it.\n", + "If the endpoint has scaled to zero, the first prediction takes a few extra minutes while it starts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d8ef001398acc6e75952a6b3152b862", + "metadata": {}, + "outputs": [], + "source": [ + "endpoint = workspace.serving_endpoints.wait_get_serving_endpoint_not_updating(\n", + " ENDPOINT_NAME, timeout=timedelta(hours=1)\n", + ")\n", + "ready = endpoint.state.ready.value if endpoint.state.ready else 'unknown'\n", + "if ready != 'READY':\n", + " raise RuntimeError(\n", + " f'The endpoint {ENDPOINT_NAME} is {ready}. '\n", + " 'Open it under Serving to see why, then run this cell again.'\n", + " )\n", + "# Allow 15 minutes per request, for an endpoint that has to scale up from zero.\n", + "client = SDMClient.for_databricks_serving(ENDPOINT_NAME, timeout=900)\n", + "model = KumoTabular()\n", + "print(f'Connected to {ENDPOINT_NAME}.')" + ] + }, + { + "cell_type": "markdown", + "id": "cdcff7ffa76899086f96e1d9b857ac11", + "metadata": {}, + "source": [ + "## Step 7. Make predictions\n", + "\n", + "Each prediction takes two calls:\n", + "\n", + "1. `Request(context=..., predict=..., target=..., task=...)` holds the labelled rows, the rows to score, the column to predict, and the kind of question.\n", + "2. `client.predict(request, model=model)` sends it, and `to_pandas()` returns the result.\n", + "\n", + "`task` is required: it is `'classification'` or `'regression'`, and a column of whole numbers could be either.\n", + "\n", + "**Note:** AdventureWorks is sample data, and the results below are for demonstration.\n", + "Benchmark Kumo Tabular on data representative of your own use case before drawing conclusions about its performance." + ] + }, + { + "cell_type": "markdown", + "id": "e9b3e0cdbb9ae38d0faed7d135aa8b85", + "metadata": {}, + "source": [ + "### Example 1: Predict a product's color\n", + "\n", + "Which color does a product have?\n", + "\n", + "Some products have no color recorded, but only the products that do can be checked.\n", + "The cell sets aside a random 60 of those as the rows to score, with their color hidden, and gives the rest as context.\n", + "The result is one row per product to score, with the predicted class and a probability for each class." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1304e32f1411c7445bf8c0f3098cac4d", + "metadata": {}, + "outputs": [], + "source": [ + "TARGET = 'Color'\n", + "known = features[features[TARGET].notna()].sample(frac=1, random_state=0)\n", + "held_out = known.head(60).reset_index(drop=True)\n", + "context = known.iloc[60:]\n", + "# Which rows are which, by their place in the table, for Example 3.\n", + "held_out_rows, context_rows = known.index[:60], known.index[60:]\n", + "\n", + "request = Request(\n", + " context=context,\n", + " predict=held_out.drop(columns=[TARGET]),\n", + " target=TARGET,\n", + " task='classification',\n", + ")\n", + "colors = client.predict(request, model=model).to_pandas()\n", + "colors.head()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6151e733e52d70acbff2706114823da7", + "metadata": {}, + "outputs": [], + "source": [ + "held_out['predicted'] = held_out.index.map(\n", + " colors.set_index('row_index')['prediction']\n", + ")\n", + "most_common = context[TARGET].mode()[0]\n", + "correct = (held_out['predicted'] == held_out[TARGET]).sum()\n", + "guessed = (held_out[TARGET] == most_common).sum()\n", + "print(\n", + " f'{correct} of {len(held_out)} colors predicted correctly; always '\n", + " f'guessing the most common color, {most_common}, gets {guessed} right.'\n", + ")\n", + "display(\n", + " held_out.loc[\n", + " held_out['predicted'] != held_out[TARGET], [TARGET, 'predicted']\n", + " ]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "626bd7b7d84310cb239b73a01b94dd1f", + "metadata": {}, + "source": [ + "How to interpret the result:\n", + "\n", + "1. The first table has one row per product, keyed by `row_index`, its position among the rows to score, with the predicted color in `prediction` and one probability column per color.\n", + "2. The printed line counts the correct predictions, and the right guesses of a rule that always picks the most common color.\n", + "3. The second table lists the products the model got wrong.\n", + "\n", + "The probabilities are the model's, not calibrated scores: read them as an ordering over the colors rather than as frequencies." + ] + }, + { + "cell_type": "markdown", + "id": "4e8f15fa3c864abe0944c5bd7eecf152", + "metadata": {}, + "source": [ + "### Example 2: Predict a product's list price\n", + "\n", + "The same table, a number instead of a label.\n", + "Changing the question means changing the target and the task; everything else is the same call.\n", + "Here the answer is a price, so the task is a regression, and the result has one `prediction` column and no probabilities.\n", + "\n", + "Only products with a list price above zero are used, because a price of zero marks a part that is never sold.\n", + "The cell sets aside a random 60 of them as the rows to score, with their price hidden, and compares the predictions with the recorded prices and with a rule that always guesses the median price." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a4a6dcc8d052704213c857fe345cebd", + "metadata": {}, + "outputs": [], + "source": [ + "priced = features[features['ListPrice'] > 0].sample(frac=1, random_state=0)\n", + "to_price = priced.head(60).reset_index(drop=True)\n", + "\n", + "request = Request(\n", + " context=priced.iloc[60:],\n", + " predict=to_price.drop(columns=['ListPrice']),\n", + " target='ListPrice',\n", + " task='regression',\n", + ")\n", + "prices = client.predict(request, model=model).to_pandas()\n", + "\n", + "to_price['predicted'] = to_price.index.map(\n", + " prices.set_index('row_index')['prediction']\n", + ")\n", + "error = (to_price['predicted'] - to_price['ListPrice']).abs().mean()\n", + "median = priced.iloc[60:]['ListPrice'].median()\n", + "rule_error = (median - to_price['ListPrice']).abs().mean()\n", + "print(\n", + " f'Average error of the predicted price: ${error:,.2f}; always guessing '\n", + " f'the median price, ${median:,.2f}: ${rule_error:,.2f}.'\n", + ")\n", + "to_price[['ListPrice', 'predicted']].head(10)" + ] + }, + { + "cell_type": "markdown", + "id": "018d3ad2f0e002fe4b16eaf285e0e2e0", + "metadata": {}, + "source": [ + "How to interpret the result:\n", + "\n", + "1. The printed line compares the average difference between the predicted and recorded price with that of a rule that always guesses the median price.\n", + "2. The table shows the first ten products: the recorded price and the predicted one." + ] + }, + { + "cell_type": "markdown", + "id": "5bf68bea28f4377838271a6fff96ba84", + "metadata": {}, + "source": [ + "### Example 3: Ask the color question in PQL\n", + "\n", + "Example 1 built a `Request` by hand.\n", + "The Predictive Query Language says what to predict and which entity it is about, and gives the same question in one line:\n", + "\n", + "```\n", + "PREDICT products.Color ← the column to fill in\n", + "FOR EACH products.ProductID ← one prediction per product\n", + "```\n", + "\n", + "The schema comes from the rows already in the notebook, so this asks nothing more of the catalog.\n", + "`materialize(predict=...)` takes the rows to score, the same 60 products as Example 1, and builds the request.\n", + "`ProductID` is added back here because the schema needs a key.\n", + "Almost every other clause of the language, such as aggregations, time windows and `RANK`, is refused on one table, and the refusal says so." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69de2c28ce1e8d8f058dc216a7b2a44b", + "metadata": {}, + "outputs": [], + "source": [ + "with_id = products.drop(columns=[c for c in LEFT_OUT if c != 'ProductID'])\n", + "ask = with_id.loc[held_out_rows].reset_index(drop=True)\n", + "schema = Schema.from_data(\n", + " {'products': with_id.loc[context_rows]},\n", + " primary_keys={'products': 'ProductID'},\n", + ")\n", + "request = PredictiveQuery(\n", + " f'PREDICT products.{TARGET} FOR EACH products.ProductID', schema\n", + ").materialize(predict=ask.drop(columns=[TARGET]), model=model)\n", + "pql_colors = client.predict(request, model=model).to_pandas()\n", + "\n", + "ask['predicted'] = ask['ProductID'].map(\n", + " pql_colors.set_index('id')['prediction']\n", + ")\n", + "print(\n", + " f'{(ask[\"predicted\"] == ask[TARGET]).sum()} of {len(ask)} colors predicted '\n", + " 'correctly.'\n", + ")\n", + "pql_colors.head()" + ] + }, + { + "cell_type": "markdown", + "id": "83812a9f0e71d8e85a42a7865b937a13", + "metadata": {}, + "source": [ + "How to interpret the result:\n", + "\n", + "1. The result is keyed by `ProductID`, in the `id` column, as well as by `row_index`.\n", + "2. The printed line counts the correct predictions for the same 60 products as Example 1, from the same context rows, so it matches Example 1's count." + ] + }, + { + "cell_type": "markdown", + "id": "35191b5835272b7849ed573c1de6294e", + "metadata": {}, + "source": [ + "## Step 8. Try it yourself\n", + "\n", + "The next cell asks the same price question as Example 2, with a number of context rows, `CONTEXT_ROWS`, that you can change.\n", + "It prints the average error of the predicted prices.\n", + "Change `CONTEXT_ROWS`, or uncomment the challenge line, and run it again." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3cfc087b258d3276f5ac18bbe7c1fb2", + "metadata": {}, + "outputs": [], + "source": [ + "CONTEXT_ROWS = 200\n", + "# Challenge: uncomment the line, run the cell again, and compare the error.\n", + "# CONTEXT_ROWS = 20 # how few examples can the model learn from?\n", + "\n", + "request = Request(\n", + " context=priced.iloc[60:].head(CONTEXT_ROWS),\n", + " predict=to_price.drop(columns=['ListPrice', 'predicted']),\n", + " target='ListPrice',\n", + " task='regression',\n", + ")\n", + "result = client.predict(request, model=model).to_pandas()\n", + "predicted = to_price.index.map(result.set_index('row_index')['prediction'])\n", + "print(\n", + " f'Average error with {CONTEXT_ROWS} context rows: '\n", + " f'${(predicted - to_price[\"ListPrice\"]).abs().mean():,.2f}.'\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "6e837325e79e8c45fc589d24752670b6", + "metadata": {}, + "source": [ + "## Step 9. Use your own data\n", + "\n", + "For your own table, name the column to fill in and say which rows to score:\n", + "\n", + "```python\n", + "known = df[df[''].notna()]\n", + "unknown = df[df[''].isna()].drop(columns=[''])\n", + "request = Request(\n", + " context=known,\n", + " predict=unknown,\n", + " target='',\n", + " task='classification', # or 'regression'\n", + ")\n", + "result = client.predict(request, model=model).to_pandas()\n", + "```\n", + "\n", + "| Question | Target | Task |\n", + "|---|---|---|\n", + "| Which category does this item belong to? | a categorical column | `classification` |\n", + "| Will this customer churn? | a yes-or-no column | `classification` |\n", + "| What will this order be worth? | a numerical column | `regression` |\n", + "| What is this item's missing weight? | a numerical column | `regression` |\n", + "\n", + "For linked tables and questions about the future, such as who will stop ordering in the next 90 days, use Kumo Relational instead: it reads the links between tables and the time of each event.\n", + "\n", + "Tips:\n", + "\n", + "- Leave out random row IDs, free text, and dates: the model reads numerical and categorical columns.\n", + "- A request carries its rows in one call, up to 10,000 context rows and a limit of 16 MB on the request body; send fewer rows if the endpoint refuses it as too large.\n", + "- To keep the predictions, save them as a table: `spark.createDataFrame(result).write.saveAsTable('')`.\n", + "\n", + "Learn more:\n", + "\n", + "- [One table](https://github.com/NVIDIA/structured-data-models-client/blob/main/docs/get-started/quickstart.md#one-table) in the client's quickstart\n", + "- [Limits](https://github.com/NVIDIA/structured-data-models-client/blob/main/docs/reference/limits.md)" + ] + }, + { + "cell_type": "markdown", + "id": "3b725598852163af887b0ac980a1ed6c", + "metadata": {}, + "source": [ + "## Step 10. Clean up\n", + "\n", + "The endpoint bills for GPU time only while it runs.\n", + "To delete it, uncomment the last line and run the cell.\n", + "The sample table stays in your schema until you drop it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5948acae96bf3d2d978e9a2402f176fc", + "metadata": {}, + "outputs": [], + "source": [ + "client.close()\n", + "# workspace.serving_endpoints.delete(ENDPOINT_NAME)" + ] + } + ], + "metadata": { + "application/vnd.databricks.v1+notebook": { + "language": "python" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 2e69ce1d76628d7d259c4ce21dab075884ae6957 Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Wed, 21 Oct 2026 22:37:07 -0700 Subject: [PATCH 2/6] Say fewer rows, not fewer entities, when serving refuses a request as too large A tabular request has no entities and no num_context_samples, so the 16 MB refusal now asks for fewer context rows or fewer rows to score. The limits table and the quickstart name the tabular model too, and the notebook test expects the new quickstart. Signed-off-by: Adi Wadaskar --- CHANGELOG.md | 6 ++++++ docs/get-started/quickstart.md | 7 +++++++ docs/reference/limits.md | 2 +- .../src/sdm_client/core/databricks_serving.py | 5 +++-- .../tests/test_databricks_notebooks.py | 1 + .../tests/test_databricks_serving.py | 4 ++++ 6 files changed, 22 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 156de1221..47e276828 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,9 @@ these packages. All five distributions share one version and one entry. `connections.toml` profile, `SF_PARTNER`, or the connector's own detection of Streamlit and notebooks. A connection or Snowpark session you open yourself is used as it is. +- A Databricks get-started notebook for the tabular model, + `examples/databricks/kumo_tabular_quickstart.ipynb`, which deploys it to a + serving endpoint and fills in a column of the AdventureWorks products table. - The relational engine publishes wheels for macOS (Apple silicon and Intel) and 64-bit Windows, and for CPython 3.14, alongside the `manylinux_2_28` x86-64 and aarch64 wheels it already built, so `structured-data-models-client[relational]` @@ -32,6 +35,9 @@ these packages. All five distributions share one version and one entry. ### Changed +- The refusal of a request over a Databricks serving endpoint's 16 MB limit + no longer names `num_context_samples` and entities alone, which a tabular + request has none of: it asks for fewer context rows or fewer rows to score. - Every engine wheel is built on a machine of its own, native rather than emulated, and each is smoke-tested on the architecture it targets. The release no longer builds aarch64 under QEMU, which also takes a privileged diff --git a/docs/get-started/quickstart.md b/docs/get-started/quickstart.md index 750fffacd..3db45e64c 100644 --- a/docs/get-started/quickstart.md +++ b/docs/get-started/quickstart.md @@ -148,6 +148,13 @@ with SDMClient.for_databricks_serving('kumo-relational') as client: result = client.predict(request, model=KumoRelational()) ``` +A `kumo-tabular` endpoint is reached the same way, with `model=KumoTabular()`. +The Databricks notebooks +[`kumo_relational_quickstart`](https://github.com/NVIDIA/structured-data-models-client/blob/main/examples/databricks/kumo_relational_quickstart.ipynb) and +[`kumo_tabular_quickstart`](https://github.com/NVIDIA/structured-data-models-client/blob/main/examples/databricks/kumo_tabular_quickstart.ipynb) +deploy each model to a serving endpoint and ask it questions about the +AdventureWorks sample data. + This needs `pip install 'structured-data-models-client[databricks-serving]'`. A serving endpoint answers one request at a time, so `fit()` is refused there; `predict()` sends the context with every call. diff --git a/docs/reference/limits.md b/docs/reference/limits.md index 1700037b3..463067622 100644 --- a/docs/reference/limits.md +++ b/docs/reference/limits.md @@ -26,7 +26,7 @@ Both endpoints enforce sizes, and most of these are the relational one's. | entities one question gathers | relational | 100,000, in at most 100 calls | at `materialize()`, and this client's own: the endpoint is never asked for more than one call's worth | | calls a question asked in parts makes | relational | 200, one per candidate or step | at `materialize()`, as `sdm_client.decompose.MOST_PARTS` | | request body | relational | 30 MiB per call | at `materialize()` | -| request body | relational, through Databricks model serving | 16 MB per call, the gateway's own decimal limit | before sending, on the wrapped envelope | +| request body | relational and tabular, through Databricks model serving | 16 MB per call, the gateway's own decimal limit | before sending, on the wrapped envelope | | an explanation | relational | exactly one entity | before sending | | an explanation | relational | only within the 1,000-row context budget | before sending | | hops | relational | 6 | at `materialize()` | diff --git a/packages/structured-data-models-client/src/sdm_client/core/databricks_serving.py b/packages/structured-data-models-client/src/sdm_client/core/databricks_serving.py index c2d4d2896..78b60bdbf 100644 --- a/packages/structured-data-models-client/src/sdm_client/core/databricks_serving.py +++ b/packages/structured-data-models-client/src/sdm_client/core/databricks_serving.py @@ -213,8 +213,9 @@ def predict(self, payload: dict[str, Any]) -> dict[str, Any]: f'this request is {size:,} bytes once wrapped for serving ' f'endpoint {self._endpoint!r}, above the ' f'{self._max_request_bytes:,} byte limit it accepts. Send ' - f'fewer rows per call: a smaller num_context_samples, or ' - f'fewer entities to score', + f'fewer rows per call: fewer context rows (a smaller ' + f'num_context_samples, where the model takes one), or fewer ' + f'rows to score', code='INVALID_REQUEST', details={'bytes': size, 'limit': self._max_request_bytes}, ) diff --git a/packages/structured-data-models-client/tests/test_databricks_notebooks.py b/packages/structured-data-models-client/tests/test_databricks_notebooks.py index d2ef76445..053a2abca 100644 --- a/packages/structured-data-models-client/tests/test_databricks_notebooks.py +++ b/packages/structured-data-models-client/tests/test_databricks_notebooks.py @@ -37,6 +37,7 @@ def _python(source: str) -> str: def test_there_is_a_notebook_for_each_way_in() -> None: assert {path.name for path in NOTEBOOKS} >= { 'kumo_relational_quickstart.ipynb', + 'kumo_tabular_quickstart.ipynb', 'lakebase_relational.ipynb', } diff --git a/packages/structured-data-models-client/tests/test_databricks_serving.py b/packages/structured-data-models-client/tests/test_databricks_serving.py index 511995efa..56c45f01b 100644 --- a/packages/structured-data-models-client/tests/test_databricks_serving.py +++ b/packages/structured-data-models-client/tests/test_databricks_serving.py @@ -439,6 +439,10 @@ def test_an_oversized_request_is_refused_before_it_is_sent(): assert error.value.details['limit'] == 200 assert error.value.details['bytes'] > 200 assert workspace.serving_endpoints.calls == [] + # The advice holds for any model: a tabular request has no entities and + # no num_context_samples, so it names rows rather than either. + assert 'context rows' in str(error.value) + assert 'rows to score' in str(error.value) def test_the_size_counted_is_the_wrapped_envelope_not_the_payload(): From 1b3718350814a0860fd1e2a68dea44183f3491e9 Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Wed, 21 Oct 2026 23:13:08 -0700 Subject: [PATCH 3/6] Say the price question is an easy one in the Kumo Tabular quickstart The products table carries each product's standard cost, which tracks its list price closely, so the regression error is small. Say so, so nobody reads it as the model's typical error on a hard column. Signed-off-by: Adi Wadaskar --- examples/databricks/kumo_tabular_quickstart.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/databricks/kumo_tabular_quickstart.ipynb b/examples/databricks/kumo_tabular_quickstart.ipynb index da2ef50e6..7ee89ddde 100644 --- a/examples/databricks/kumo_tabular_quickstart.ipynb +++ b/examples/databricks/kumo_tabular_quickstart.ipynb @@ -428,7 +428,8 @@ "Here the answer is a price, so the task is a regression, and the result has one `prediction` column and no probabilities.\n", "\n", "Only products with a list price above zero are used, because a price of zero marks a part that is never sold.\n", - "The cell sets aside a random 60 of them as the rows to score, with their price hidden, and compares the predictions with the recorded prices and with a rule that always guesses the median price." + "The cell sets aside a random 60 of them as the rows to score, with their price hidden, and compares the predictions with the recorded prices and with a rule that always guesses the median price.\n", + "The table also holds each product's standard cost, which tracks its list price closely, so this is an easy question; a column with less to give away is a harder one." ] }, { From 344001efb17056712546a655fdfd69c2ae9ddadc Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Wed, 21 Oct 2026 23:13:57 -0700 Subject: [PATCH 4/6] Save the outputs of a run on Databricks into the Kumo Tabular quickstart The notebook was run end to end on serverless environment 5 against a GPU_MEDIUM endpoint serving Kumo Tabular. The catalog, schema and endpoint names in the outputs are the notebook's defaults. The pip output is left out, because it only repeats the protobuf notice the text already explains. Signed-off-by: Adi Wadaskar --- .../databricks/kumo_tabular_quickstart.ipynb | 747 +++++++++++++++++- 1 file changed, 729 insertions(+), 18 deletions(-) diff --git a/examples/databricks/kumo_tabular_quickstart.ipynb b/examples/databricks/kumo_tabular_quickstart.ipynb index 7ee89ddde..e51bdabe5 100644 --- a/examples/databricks/kumo_tabular_quickstart.ipynb +++ b/examples/databricks/kumo_tabular_quickstart.ipynb @@ -153,10 +153,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "a133dc9b0e07efd33ac4da64235bcbdb", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Creating the endpoint kumo-tabular; this takes 20 to 40 minutes.\n" + ] + } + ], "source": [ "workspace = WorkspaceClient()\n", "try:\n", @@ -210,10 +218,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "1c9a61987b2e78e004b7da1d055a40b2", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Saved the sample table to main.kumo_tabular_quickstart.\n" + ] + } + ], "source": [ "URL = (\n", " 'https://raw.githubusercontent.com/microsoft/sql-server-samples/'\n", @@ -268,10 +284,226 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "7015e2d235fde236028420efe0c29b31", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "504 products; features: MakeFlag, FinishedGoodsFlag, Color, SafetyStockLevel, ReorderPoint, StandardCost, ListPrice, Size, SizeUnitMeasureCode, WeightUnitMeasureCode, Weight, DaysToManufacture, ProductLine, Class, Style, ProductSubcategoryID, ProductModelID\n" + ] + }, + { + "output_type": "execute_result", + "execution_count": 3, + "metadata": {}, + "data": { + "text/plain": [ + " ProductID ... ModifiedDate\n", + "0 1 ... 2025-02-07 10:01:36.827\n", + "1 2 ... 2025-02-07 10:01:36.827\n", + "2 3 ... 2025-02-07 10:01:36.827\n", + "3 4 ... 2025-02-07 10:01:36.827\n", + "4 316 ... 2025-02-07 10:01:36.827\n", + "\n", + "[5 rows x 25 columns]" + ], + "text/html": [ + "
\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
ProductIDNameProductNumberMakeFlagFinishedGoodsFlagColorSafetyStockLevelReorderPointStandardCostListPriceSizeSizeUnitMeasureCodeWeightUnitMeasureCodeWeightDaysToManufactureProductLineClassStyleProductSubcategoryIDProductModelIDSellStartDateSellEndDateDiscontinuedDaterowguidModifiedDate
01Adjustable RaceAR-538100None10007500.00.0NoneNoneNoneNaN0NoneNoneNoneNaNNaN2019-04-30NaTNaT694215B7-08F7-4C0D-ACB1-D734BA44C0C82025-02-07 10:01:36.827
12Bearing BallBA-832700None10007500.00.0NoneNoneNoneNaN0NoneNoneNoneNaNNaN2019-04-30NaTNaT58AE3C20-4F3A-4749-A7D4-D568806CC5372025-02-07 10:01:36.827
23BB Ball BearingBE-234910None8006000.00.0NoneNoneNoneNaN1NoneNoneNoneNaNNaN2019-04-30NaTNaT9C21AED2-5BFA-4F18-BCB8-F11638DC2E4E2025-02-07 10:01:36.827
34Headset Ball BearingsBE-290800None8006000.00.0NoneNoneNoneNaN0NoneNoneNoneNaNNaN2019-04-30NaTNaTECFED6CB-51FF-49B5-B06C-7D8AC834DB8B2025-02-07 10:01:36.827
4316BladeBL-203610None8006000.00.0NoneNoneNoneNaN1NoneNoneNoneNaNNaN2019-04-30NaTNaTE73E9750-603B-4131-89F5-3DD15ED5FF802025-02-07 10:01:36.827
\n", + "" + ] + } + } + ], "source": [ "products = spark.table(PRODUCTS).toPandas()\n", "LEFT_OUT = [\n", @@ -302,10 +534,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "3d8ef001398acc6e75952a6b3152b862", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Connected to kumo-tabular.\n" + ] + } + ], "source": [ "endpoint = workspace.serving_endpoints.wait_get_serving_endpoint_not_updating(\n", " ENDPOINT_NAME, timeout=timedelta(hours=1)\n", @@ -356,10 +596,135 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "1304e32f1411c7445bf8c0f3098cac4d", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "execute_result", + "execution_count": 5, + "metadata": {}, + "data": { + "text/plain": [ + " row_index prediction ... probability_Grey probability_Silver/Black\n", + "0 0 White ... 0.052111 0.014584\n", + "1 1 Black ... 0.000534 0.000115\n", + "2 2 Blue ... 0.000522 0.000128\n", + "3 3 Black ... 0.000458 0.000123\n", + "4 4 Black ... 0.006497 0.003588\n", + "\n", + "[5 rows x 11 columns]" + ], + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
row_indexpredictionprobability_Blueprobability_Blackprobability_Silverprobability_Redprobability_Yellowprobability_Whiteprobability_Multiprobability_Greyprobability_Silver/Black
00White0.0178690.0415470.0185810.0111820.0143580.8025200.0272470.0521110.014584
11Black0.0000840.9962270.0020470.0005870.0001230.0002030.0000790.0005340.000115
22Blue0.5911750.0002630.0001750.0004470.4063080.0006110.0003710.0005220.000128
33Black0.0001850.6659460.0001850.3322530.0003460.0003050.0001970.0004580.000123
44Black0.0102210.9345660.0086070.0080860.0061030.0124260.0099070.0064970.003588
\n", + "
" + ] + } + } + ], "source": [ "TARGET = 'Color'\n", "known = features[features[TARGET].notna()].sample(frac=1, random_state=0)\n", @@ -380,10 +745,104 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "id": "6151e733e52d70acbff2706114823da7", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "47 of 60 colors predicted correctly; always guessing the most common color, Black, gets 21 right.\n" + ] + }, + { + "output_type": "display_data", + "metadata": {}, + "data": { + "text/html": [ + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
Colorpredicted
BlackRed
YellowBlue
YellowBlue
BlackSilver
RedBlack
YellowBlue
YellowBlue
YellowBlue
BlackRed
BlackRed
YellowBlue
MultiBlack
YellowBlue
" + ], + "text/plain": [ + " Color predicted\n", + " Black Red\n", + "Yellow Blue\n", + "Yellow Blue\n", + " Black Silver\n", + " Red Black\n", + "Yellow Blue\n", + "Yellow Blue\n", + "Yellow Blue\n", + " Black Red\n", + " Black Red\n", + "Yellow Blue\n", + " Multi Black\n", + "Yellow Blue" + ] + } + } + ], "source": [ "held_out['predicted'] = held_out.index.map(\n", " colors.set_index('row_index')['prediction']\n", @@ -434,10 +893,116 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "id": "5a4a6dcc8d052704213c857fe345cebd", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Average error of the predicted price: $2.63; always guessing the median price, $337.22: $544.46.\n" + ] + }, + { + "output_type": "execute_result", + "execution_count": 7, + "metadata": {}, + "data": { + "text/plain": [ + " ListPrice predicted\n", + "0 782.99 780.330139\n", + "1 74.99 73.376358\n", + "2 249.79 251.080139\n", + "3 742.35 745.315979\n", + "4 59.99 58.066715\n", + "5 348.76 350.437347\n", + "6 3374.99 3370.106934\n", + "7 594.83 597.447815\n", + "8 124.73 125.464111\n", + "9 69.99 67.738365" + ], + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
ListPricepredicted
0782.99780.330139
174.9973.376358
2249.79251.080139
3742.35745.315979
459.9958.066715
5348.76350.437347
63374.993370.106934
7594.83597.447815
8124.73125.464111
969.9967.738365
\n", + "
" + ] + } + } + ], "source": [ "priced = features[features['ListPrice'] > 0].sample(frac=1, random_state=0)\n", "to_price = priced.head(60).reset_index(drop=True)\n", @@ -497,10 +1062,148 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, "id": "69de2c28ce1e8d8f058dc216a7b2a44b", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "47 of 60 colors predicted correctly.\n" + ] + }, + { + "output_type": "execute_result", + "execution_count": 8, + "metadata": {}, + "data": { + "text/plain": [ + " row_index id ... probability_Grey probability_Silver/Black\n", + "0 0 874 ... 0.052111 0.014584\n", + "1 1 777 ... 0.000534 0.000115\n", + "2 2 890 ... 0.000522 0.000128\n", + "3 3 795 ... 0.000458 0.000123\n", + "4 4 862 ... 0.006497 0.003588\n", + "\n", + "[5 rows x 12 columns]" + ], + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
row_indexidpredictionprobability_Blueprobability_Blackprobability_Silverprobability_Redprobability_Yellowprobability_Whiteprobability_Multiprobability_Greyprobability_Silver/Black
00874White0.0178690.0415470.0185810.0111820.0143580.8025200.0272470.0521110.014584
11777Black0.0000840.9962270.0020470.0005870.0001230.0002030.0000790.0005340.000115
22890Blue0.5911750.0002630.0001750.0004470.4063080.0006110.0003710.0005220.000128
33795Black0.0001850.6659460.0001850.3322530.0003460.0003050.0001970.0004580.000123
44862Black0.0102210.9345660.0086070.0080860.0061030.0124260.0099070.0064970.003588
\n", + "
" + ] + } + } + ], "source": [ "with_id = products.drop(columns=[c for c in LEFT_OUT if c != 'ProductID'])\n", "ask = with_id.loc[held_out_rows].reset_index(drop=True)\n", @@ -548,10 +1251,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "id": "b3cfc087b258d3276f5ac18bbe7c1fb2", "metadata": {}, - "outputs": [], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Average error with 200 context rows: $2.58.\n" + ] + } + ], "source": [ "CONTEXT_ROWS = 200\n", "# Challenge: uncomment the line, run the cell again, and compare the error.\n", From b3f41a5d9e19ad586dd381ec3cc1f3cba73d0b95 Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Thu, 29 Oct 2026 11:16:22 -0700 Subject: [PATCH 5/6] Check the served Tabular model after endpoint updates finish Signed-off-by: Adi Wadaskar --- .../databricks/kumo_tabular_quickstart.ipynb | 11 ++++ .../tests/test_databricks_notebooks.py | 65 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/examples/databricks/kumo_tabular_quickstart.ipynb b/examples/databricks/kumo_tabular_quickstart.ipynb index e51bdabe5..551af6c73 100644 --- a/examples/databricks/kumo_tabular_quickstart.ipynb +++ b/examples/databricks/kumo_tabular_quickstart.ipynb @@ -556,6 +556,17 @@ " f'The endpoint {ENDPOINT_NAME} is {ready}. '\n", " 'Open it under Serving to see why, then run this cell again.'\n", " )\n", + "# A pending update may have changed the model while we waited.\n", + "served = {\n", + " (entity.entity_name.lower(), entity.entity_version)\n", + " for entity in endpoint.config.served_entities\n", + "}\n", + "if served != {(MODEL_NAME.lower(), MODEL_VERSION)}:\n", + " raise ValueError(\n", + " f'The endpoint {ENDPOINT_NAME} now serves {sorted(served)}, '\n", + " f'not {MODEL_NAME} version {MODEL_VERSION}. '\n", + " 'Set ENDPOINT_NAME to another name and run from Step 3 again.'\n", + " )\n", "# Allow 15 minutes per request, for an endpoint that has to scale up from zero.\n", "client = SDMClient.for_databricks_serving(ENDPOINT_NAME, timeout=900)\n", "model = KumoTabular()\n", diff --git a/packages/structured-data-models-client/tests/test_databricks_notebooks.py b/packages/structured-data-models-client/tests/test_databricks_notebooks.py index 053a2abca..b7e2a3250 100644 --- a/packages/structured-data-models-client/tests/test_databricks_notebooks.py +++ b/packages/structured-data-models-client/tests/test_databricks_notebooks.py @@ -10,7 +10,10 @@ import ast import importlib import json +from datetime import timedelta from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -71,3 +74,65 @@ def test_every_name_imported_from_this_project_is_public( f"{notebook.name} imports {private}, which are not in their module's " f'__all__' ) + + +@pytest.mark.parametrize( + ('ready_model', 'ready_version', 'accepted'), + [ + ('catalog.schema.tabular', '1', True), + ('CATALOG.SCHEMA.TABULAR', '1', True), + ('catalog.schema.other', '1', False), + ('catalog.schema.tabular', '2', False), + ], +) +def test_tabular_notebook_checks_the_model_after_a_pending_update( + ready_model: str, ready_version: str, accepted: bool +) -> None: + cells = _cells(EXAMPLES / 'kumo_tabular_quickstart.ipynb', 'code') + deploy = next( + cell for cell in cells if 'workspace = WorkspaceClient()' in cell + ) + connect = next( + cell + for cell in cells + if 'wait_get_serving_endpoint_not_updating(' in cell + ) + + def config(name: str, version: str) -> SimpleNamespace: + return SimpleNamespace( + served_entities=[ + SimpleNamespace(entity_name=name, entity_version=version) + ] + ) + + workspace = Mock() + workspace.serving_endpoints.get.return_value = SimpleNamespace( + config=config('catalog.schema.tabular', '1'), + pending_config=config(ready_model, ready_version), + ) + workspace.serving_endpoints.wait_get_serving_endpoint_not_updating.return_value = SimpleNamespace( + config=config(ready_model, ready_version), + state=SimpleNamespace(ready=SimpleNamespace(value='READY')), + ) + client_factory = Mock() + namespace = { + 'WorkspaceClient': lambda: workspace, + 'NotFound': LookupError, + 'MODEL_NAME': 'catalog.schema.tabular', + 'MODEL_VERSION': '1', + 'ENDPOINT_NAME': 'tabular-demo', + 'timedelta': timedelta, + 'SDMClient': SimpleNamespace(for_databricks_serving=client_factory), + 'KumoTabular': Mock(), + } + exec(deploy, namespace) + workspace.serving_endpoints.create.assert_not_called() + if accepted: + exec(connect, namespace) + client_factory.assert_called_once_with('tabular-demo', timeout=900) + else: + with pytest.raises( + ValueError, match='not catalog.schema.tabular version 1' + ): + exec(connect, namespace) + client_factory.assert_not_called() From 01887cb934b454b0234f342dc3e65ae52c974032 Mon Sep 17 00:00:00 2001 From: Adi Wadaskar Date: Thu, 29 Oct 2026 19:17:57 -0700 Subject: [PATCH 6/6] Accept an endpoint already updating to the requested Tabular model Signed-off-by: Adi Wadaskar --- .../databricks/kumo_tabular_quickstart.ipynb | 2 +- .../tests/test_databricks_notebooks.py | 17 +++++++++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/examples/databricks/kumo_tabular_quickstart.ipynb b/examples/databricks/kumo_tabular_quickstart.ipynb index 551af6c73..27bbb55fb 100644 --- a/examples/databricks/kumo_tabular_quickstart.ipynb +++ b/examples/databricks/kumo_tabular_quickstart.ipynb @@ -189,7 +189,7 @@ " f'Creating the endpoint {ENDPOINT_NAME}; this takes 20 to 40 minutes.'\n", " )\n", "else:\n", - " config = endpoint.config or endpoint.pending_config\n", + " config = endpoint.pending_config or endpoint.config\n", " served = {\n", " (entity.entity_name.lower(), entity.entity_version)\n", " for entity in config.served_entities\n", diff --git a/packages/structured-data-models-client/tests/test_databricks_notebooks.py b/packages/structured-data-models-client/tests/test_databricks_notebooks.py index b7e2a3250..04d92b6ed 100644 --- a/packages/structured-data-models-client/tests/test_databricks_notebooks.py +++ b/packages/structured-data-models-client/tests/test_databricks_notebooks.py @@ -77,16 +77,17 @@ def test_every_name_imported_from_this_project_is_public( @pytest.mark.parametrize( - ('ready_model', 'ready_version', 'accepted'), + ('active_model', 'ready_model', 'ready_version', 'accepted'), [ - ('catalog.schema.tabular', '1', True), - ('CATALOG.SCHEMA.TABULAR', '1', True), - ('catalog.schema.other', '1', False), - ('catalog.schema.tabular', '2', False), + ('catalog.schema.tabular', 'catalog.schema.tabular', '1', True), + ('catalog.schema.other', 'catalog.schema.tabular', '1', True), + ('catalog.schema.tabular', 'CATALOG.SCHEMA.TABULAR', '1', True), + ('catalog.schema.tabular', 'catalog.schema.other', '1', False), + ('catalog.schema.tabular', 'catalog.schema.tabular', '2', False), ], ) def test_tabular_notebook_checks_the_model_after_a_pending_update( - ready_model: str, ready_version: str, accepted: bool + active_model: str, ready_model: str, ready_version: str, accepted: bool ) -> None: cells = _cells(EXAMPLES / 'kumo_tabular_quickstart.ipynb', 'code') deploy = next( @@ -107,8 +108,8 @@ def config(name: str, version: str) -> SimpleNamespace: workspace = Mock() workspace.serving_endpoints.get.return_value = SimpleNamespace( - config=config('catalog.schema.tabular', '1'), - pending_config=config(ready_model, ready_version), + config=config(active_model, '1'), + pending_config=config('catalog.schema.tabular', '1'), ) workspace.serving_endpoints.wait_get_serving_endpoint_not_updating.return_value = SimpleNamespace( config=config(ready_model, ready_version),