{ "cells": [ { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "7b233c4c-fb5b-4519-bc03-bcdd57c2ca57", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 1" }, "codeCollapsed": true, "collapsed": false }, "source": [ "# Kumo Engine task examples with TPC-H on Databricks\n", "\n", "Explore binary classification, regression, and link prediction with Kumo Engine on TPC-H data. \n", "\n", "This walkthrough reads the built-in `samples.tpch` dataset, materializes a local graph snapshot, and stores model bundles in temporary local scratch. No external credentials are required.\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" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "4df50a12-efa0-4ca4-8772-c62c29b41e67", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 2" }, "codeCollapsed": true, "collapsed": false }, "source": [ "## Prepare your environment\n", "\n", "**Runtime choices:** the default path is Serverless GPU; select a Python 3.12 AI Runtime GPU environment. You could also choose the classic all-purpose GPU cluster with dedicated access mode on Databricks Runtime 17.3 LTS ML (Python 3.12.3). \n", "\n", "**Compute and storage:** choose a CUDA-capable GPU and sufficient CPU RAM/GPU memory for the graph and model. Use the [GPU Sizing guide](https://docs.nvidia.com/sdgm/kumo-engine/get-started-with-kumo-engine/gpu-sizing) to size local scratch from the total source Parquet size. Allow additional space for runs and archives. Use `/local_disk0/kumo-scratch` for ephemeral scratch on classic dedicated compute, or `/tmp/kumoai-local` path on Serverless GPU; `/local_disk0` is not writable on serverless or standard access mode. `/tmp` is temporary and should not hold durable or sensitive exports.\n", "\n", "**Permissions and access:** attach/use the chosen GPU compute; install `kumo-engine[databricks]` from PyPI; use the SQL warehouse (`CAN USE`); read `samples.tpch` (`USE CATALOG`, `USE SCHEMA`, `SELECT`); create/replace views in the destination catalog/schema (`USE CATALOG`, `USE SCHEMA`, `CREATE TABLE`, plus `CREATE SCHEMA` on the catalog if this notebook creates that schema). Enabling prediction output also requires permission to create or overwrite its dedicated table. Existing views must be owned or manageable by the running identity.\n", "\n", "See the [Databricks Runtime 17.3 LTS ML release notes](https://docs.databricks.com/aws/en/release-notes/runtime/17.3lts-ml) and [Databricks local-file guidance](https://docs.databricks.com/aws/en/files/write-data) for current platform details.\n", "\n", "The TPC-H dataset contains `customer`, `orders`, `lineitem`, and `part` with the columns used below. The notebook materializes local graph data and adds a synthetic primary key for `lineitem`." ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "9d5c1f63-d0d5-4fda-815d-82c8fbd98f52", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 3" }, "codeCollapsed": true }, "source": [ "## 1. Enter your configuration\n", "\n", "Set the destination Unity Catalog catalog and schema for the views and output. The sample source remains `samples.tpch`. Kumo uses the configured local scratch directory to materialize graph and training data. Model bundles remain in temporary local scratch. Set the prediction date within the data's date range (1992-01-01 to 1998-08-02).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538189064, "inputWidgets": {}, "nuid": "a44ced76-08e9-403b-a9ee-12593b527e63", "showTitle": true, "startTime": 1792538188852, "submitTime": 1792538188738, "tableResultSettingsMap": {}, "title": "Configuration" }, "language": "python" }, "outputs": [], "source": [ "from datetime import datetime\n", "from pathlib import Path\n", "\n", "# ── TPC-H source (Databricks sample data) ──\n", "TPCH_CATALOG = \"samples\"\n", "TPCH_SCHEMA = \"tpch\"\n", "\n", "# ── Databricks connector ──\n", "CONNECTOR_CATALOG = \"\"\n", "CONNECTOR_SCHEMA = \"\"\n", "SQL_WAREHOUSE_ID = \"\"\n", "\n", "# MODEL_NAME = \"TPCH_CUSTOMER_ORDER_MODEL\"\n", "PREDICTION_TIME = datetime(1998, 1, 1) # Orders range 1992-01-01 to 1998-08-02.\n", "PREDICTION_OUTPUT_TABLE = \"TPCH_ORDER_PREDICTIONS\"\n", "\n", "# Durable prediction output.\n", "WRITE_PREDICTIONS = True\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "09d6d3ec-c322-4493-a9ad-90ba7034ade6", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 9" }, "codeCollapsed": true }, "source": [ "## 2. Install Kumo Engine from PyPI\n", "\n", "Run the installation cell once, then restart Python. After the restart, rerun configuration and continue at imports.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538169955, "inputWidgets": {}, "nuid": "c14ac6f7-4127-4ea9-a30f-efb043fef7a2", "showTitle": true, "startTime": 1792538165619, "submitTime": 1792538165560, "tableResultSettingsMap": {}, "title": "Install Kumo Engine" }, "language": "python" }, "outputs": [], "source": [ "%pip install kumo-engine[databricks]\n", "\n", "dbutils.library.restartPython()\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "615f150a-4304-41cf-aa80-5b2f757bd6de", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "\n", "## Import and configure scratch\n", "\n", "After each kernel restart, run configuration and verify that the runtime imports and a CUDA operation succeed. They do not depend on variables created during installation. " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538202476, "inputWidgets": {}, "nuid": "bfdcdc2a-78a5-4011-8019-f92f78b2a15b", "showTitle": true, "startTime": 1792538192463, "submitTime": 1792538192428, "tableResultSettingsMap": {}, "title": "Cell 12" }, "language": "python" }, "outputs": [], "source": [ "from importlib.metadata import version\n", "import shutil\n", "import tempfile\n", "\n", "from IPython.display import display\n", "import torch\n", "import torch_geometric\n", "\n", "import kumo.engine as kumo\n", "from kumo.engine import PredictionOutputConfig\n", "from kumoapi.graph import Edge\n", "from kumoapi.jobs import PredictionArtifactType, WriteMode\n", "from kumoapi.task import TaskType\n", "from kumo_databricks.output import SEPARATOR as _SEP\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538206478, "inputWidgets": {}, "nuid": "14357384-029d-4225-88ba-94cf181f3855", "showTitle": true, "startTime": 1792538206159, "submitTime": 1792538206127, "tableResultSettingsMap": {}, "title": "Cell 9" }, "language": "python" }, "outputs": [], "source": [ "if not torch.cuda.is_available():\n", " raise RuntimeError(\"No GPU is available to this notebook.\")\n", "\n", "values = torch.tensor([1.0, 2.0], device=\"cuda\")\n", "groups = torch.tensor([0, 0], device=\"cuda\")\n", "total = torch.zeros(1, device=\"cuda\").scatter_add_(0, groups, values)\n", "assert total.item() == 3.0\n", "\n", "print(\"Kumo:\", kumo.__version__)\n", "print(\"PyTorch:\", torch.__version__, \"CUDA:\", torch.version.cuda)\n", "print(\"PyG:\", torch_geometric.__version__)\n", "print(\"GPU:\", torch.cuda.get_device_name(0))" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "94b449eb-6621-4639-a9ad-812357539ec9", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "Configure scratch before constructing any Kumo table, graph, or trainer." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538210482, "inputWidgets": {}, "nuid": "685cd007-872c-4279-8b78-84cc1a3d5794", "showTitle": false, "startTime": 1792538210293, "submitTime": 1792538210263, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "SCRATCH_DIR = Path(\"/tmp/kumoai-local\")\n", "SCRATCH_DIR.mkdir(parents=True, exist_ok=True)\n", "with tempfile.NamedTemporaryFile(dir=SCRATCH_DIR) as probe:\n", " probe.write(b\"scratch check\")\n", " probe.flush()\n", "kumo.configure(scratch_dir=str(SCRATCH_DIR))\n", "print(\"Free scratch space (GiB):\", round(shutil.disk_usage(SCRATCH_DIR).free / 2**30, 1))\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "31c2ae98-1d93-4956-bdd2-50ee05152d46", "showTitle": true, "tableResultSettingsMap": {}, "title": "Connect section" }, "codeCollapsed": true }, "source": [ "## Connect to TPC-H data\n", "\n", "The notebook creates Unity Catalog views of the four TPC-H sample tables in the configured schema. The `lineitem` view adds a synthetic single-column primary key (`lineitem_id`) since TPC-H uses a composite key. A `DatabricksConnector` uses the notebook's default auth chain — no separate credentials are needed.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538283842, "inputWidgets": {}, "nuid": "f40ed503-c937-44a5-9402-bb535cbd7919", "showTitle": true, "startTime": 1792538216431, "submitTime": 1792538216390, "tableResultSettingsMap": {}, "title": "Connector (SF≈1 subset)" }, "language": "python" }, "outputs": [], "source": [ "# ── Create TPC-H views in the connector schema (SF≈1 subset) ──\n", "# SF=1 scale: ~150K customers, ~1.5M orders, ~6M line items, ~200K parts\n", "SF1_CUST = 150000\n", "SF1_PART = 200000\n", "\n", "spark.sql(f\"CREATE SCHEMA IF NOT EXISTS {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}\")\n", "\n", "spark.sql(f\"\"\"\n", " CREATE OR REPLACE VIEW {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.customer AS\n", " SELECT * FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.customer\n", " WHERE c_custkey <= {SF1_CUST}\n", "\"\"\")\n", "print(f\"View: {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.customer (SF≈1, ≤{SF1_CUST})\")\n", "\n", "spark.sql(f\"\"\"\n", " CREATE OR REPLACE VIEW {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.orders AS\n", " SELECT * FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.orders\n", " WHERE o_custkey <= {SF1_CUST}\n", "\"\"\")\n", "print(f\"View: {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.orders (SF≈1)\")\n", "\n", "spark.sql(f\"\"\"\n", " CREATE OR REPLACE VIEW {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.part AS\n", " SELECT * FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.part\n", " WHERE p_partkey <= {SF1_PART}\n", "\"\"\")\n", "print(f\"View: {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.part (SF≈1, ≤{SF1_PART})\")\n", "\n", "# Lineitem: synthetic PK + filter to SF≈1 customers and parts\n", "spark.sql(f\"\"\"\n", " CREATE OR REPLACE VIEW {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.lineitem AS\n", " SELECT *, CAST(l_orderkey * 8 + l_linenumber AS BIGINT) AS lineitem_id\n", " FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.lineitem\n", " WHERE l_orderkey IN (\n", " SELECT o_orderkey FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.orders\n", " WHERE o_custkey <= {SF1_CUST}\n", " )\n", " AND l_partkey <= {SF1_PART}\n", "\"\"\")\n", "print(f\"View: {CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.lineitem (SF≈1, with lineitem_id)\")\n", "\n", "# ── Create Kumo Databricks connector (uses notebook's default auth) ──\n", "connector = kumo.DatabricksConnector.from_notebook(\n", " catalog=CONNECTOR_CATALOG,\n", " schema_name=CONNECTOR_SCHEMA,\n", " warehouse_id=SQL_WAREHOUSE_ID,\n", ")\n", "print(\"Tables:\", connector.table_names())\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "54ab28de-982c-4c81-a507-c613836e3d5a", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 17" }, "codeCollapsed": true }, "source": [ "## Describe the four tables\n", "\n", "A primary key identifies a row. A physical type describes its stored representation, such as a date; a semantic type describes its modeling role, such as a category or numerical value. \n", "\n", "The customer and part definitions select the columns used in this example. " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538294549, "inputWidgets": {}, "nuid": "65d20853-119d-40fb-87da-eb4d94188fc7", "showTitle": true, "startTime": 1792538287784, "submitTime": 1792538287752, "tableResultSettingsMap": {}, "title": "Table definitions" }, "language": "python" }, "outputs": [], "source": [ "customer = (\n", " kumo.Table.from_source_table(\n", " connector[\"customer\"],\n", " [\"c_custkey\", \"c_nationkey\", \"c_mktsegment\", \"c_acctbal\"],\n", " primary_key=\"c_custkey\",\n", " )\n", ")\n", "\n", "part = (\n", " kumo.Table.from_source_table(\n", " connector[\"part\"],\n", " [\"p_partkey\", \"p_name\", \"p_brand\", \"p_type\", \"p_size\",\n", " \"p_container\", \"p_retailprice\"],\n", " primary_key=\"p_partkey\",\n", " )\n", ")" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "63c348b1-ea7b-47f1-b9c9-82481bfe4d15", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 19" }, "codeCollapsed": true }, "source": [ "Event timestamps let Kumo separate information available before a prediction boundary from future outcomes. Using future information as an input causes temporal leakage.\n", "\n", "This sample uses `o_orderdate` for orders and `l_shipdate` for line items. The ship date represents when the item became observable in the fulfillment process and is a reasonable availability timestamp for this dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538300554, "inputWidgets": {}, "nuid": "55a54e2b-d19a-4278-9041-13e761a23419", "showTitle": true, "startTime": 1792538298031, "submitTime": 1792538298002, "tableResultSettingsMap": {}, "title": "Orders table" }, "language": "python" }, "outputs": [], "source": [ "orders = (\n", " kumo.Table.from_source_table(\n", " connector[\"orders\"],\n", " [\"o_orderkey\", \"o_custkey\", \"o_orderdate\", \"o_orderstatus\",\n", " \"o_totalprice\", \"o_orderpriority\", \"o_shippriority\"],\n", " primary_key=\"o_orderkey\",\n", " time_column=\"o_orderdate\",\n", " )\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538307261, "inputWidgets": {}, "nuid": "66fb4065-950f-4b59-8c11-bd3234eac37f", "showTitle": true, "startTime": 1792538303853, "submitTime": 1792538303823, "tableResultSettingsMap": {}, "title": "Lineitem table" }, "language": "python" }, "outputs": [], "source": [ "lineitem = (\n", " kumo.Table.from_source_table(\n", " connector[\"lineitem\"],\n", " [\"lineitem_id\", \"l_orderkey\", \"l_partkey\", \"l_shipdate\",\n", " \"l_quantity\", \"l_extendedprice\", \"l_discount\", \"l_tax\",\n", " \"l_returnflag\", \"l_linestatus\", \"l_shipmode\"],\n", " primary_key=\"lineitem_id\",\n", " time_column=\"l_shipdate\",\n", " )\n", ")" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "968c5e60-af40-4b17-a632-7e878add4bfa", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 22" }, "codeCollapsed": true }, "source": [ "## Connect the tables and materialize a snapshot\n", "\n", "A **graph** describes the tables and their relationships. Each arrow below goes from a foreign-key column to the destination table's primary key:\n", "\n", "```text\n", "lineitem.l_partkey -> part.p_partkey\n", "lineitem.l_orderkey -> orders.o_orderkey\n", "orders.o_custkey -> customer.c_custkey\n", "```\n", "\n", "Validation checks the graph definition; it is not a blanket guarantee of source-data quality. A **snapshot** materializes source data into the local graph index used for training and prediction. Snapshot creation reads the tables and consumes local disk. Reuse requires an unchanged, valid retained snapshot; loading a model does not restore a lost index." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538311063, "inputWidgets": {}, "nuid": "53b31571-46bf-4919-aee4-67ff839d5bd7", "showTitle": true, "startTime": 1792538310907, "submitTime": 1792538310864, "tableResultSettingsMap": {}, "title": "Cell 23" }, "language": "python" }, "outputs": [], "source": [ "graph = kumo.Graph(\n", " tables=[customer, part, orders, lineitem],\n", " edges=[\n", " Edge(src_table=\"orders\", fkey=\"o_custkey\", dst_table=\"customer\"),\n", " Edge(src_table=\"lineitem\", fkey=\"l_orderkey\", dst_table=\"orders\"),\n", " Edge(src_table=\"lineitem\", fkey=\"l_partkey\", dst_table=\"part\"),\n", " ],\n", ")\n", "\n", "graph.validate()\n", "print(\"graph id:\", graph.id)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538332081, "inputWidgets": {}, "nuid": "c3cf8b86-9572-4207-adbe-571881f2ea7c", "showTitle": false, "startTime": 1792538316366, "submitTime": 1792538316337, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "snapshot = graph.snapshot()\n", "print(\"Snapshot:\", snapshot.id)\n", "print(\"Local index:\", snapshot.index_path)" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "573106d7-41ba-42d8-bd5b-1ac23c9c74a4", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 25" }, "codeCollapsed": true }, "source": [ "## Define the prediction question\n", "\n", "Predictive Query Language (PQL) is Kumo Engine's declarative syntax for defining a predictive modeling task: it specifies the target, the entity to score, and the data context. This default query predicts how many orders a customer active in the past year will place in the next 90 days." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538335483, "inputWidgets": {}, "nuid": "9231e678-701a-4c30-b501-af7e38c00265", "showTitle": true, "startTime": 1792538335254, "submitTime": 1792538335225, "tableResultSettingsMap": {}, "title": "Cell 26" }, "language": "python" }, "outputs": [], "source": [ "# Regression: predict how many orders a customer active in the past year will place in the next 90 days.\n", "pquery = kumo.PredictiveQuery(\n", " graph=graph,\n", " query=(\n", " \"PREDICT COUNT(orders.*, 0, 90, days) \"\n", " \"FOR EACH customer.c_custkey \"\n", " \"WHERE COUNT(orders.*, -365, 0, days) > 0\"\n", " ),\n", ")\n", "pquery.validate()\n", "print(\"Task type:\", pquery.get_task_type())" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "52b13150-dc8c-4b7b-9153-c025673caa3e", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Optional task alternatives\n", "\n", "The next two code cells are commented out. To try one, skip or comment out the regression query cell above, uncomment exactly one alternative, and run it before training. Both define `pquery` and use the existing `graph`, `snapshot`, and downstream variables. Train a new model for the selected task. The link example ranks parts associated with orders, not customer-level recommendations." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538339687, "inputWidgets": {}, "nuid": "97eb1030-cf7e-4a67-b2c6-e79d1195f750", "showTitle": false, "startTime": 1792538339560, "submitTime": 1792538339532, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "# Binary classification: predict whether a customer active in the past year will order again in the next 90 days.\n", "# pquery = kumo.PredictiveQuery(\n", "# graph=graph,\n", "# query=\"\"\"\n", "# PREDICT COUNT(orders.*, 0, 90, days) > 0\n", "# FOR EACH customer.c_custkey\n", "# WHERE COUNT(orders.*, -365, 0, days) > 0\n", "# \"\"\",\n", "# )\n", "\n", "# pquery.validate()\n", "# print(\"Task type:\", pquery.get_task_type())" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538347693, "inputWidgets": {}, "nuid": "2b6e7bc2-dcfc-4a4a-ad6c-64bac716ba2a", "showTitle": false, "startTime": 1792538347573, "submitTime": 1792538347543, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "# Link prediction: rank up to five parts that may ship in the next 90 days for an order with recent line items.\n", "# pquery = kumo.PredictiveQuery(\n", "# graph=graph,\n", "# query=(\n", "# \"PREDICT LIST_DISTINCT(lineitem.l_partkey, 0, 90, days) \"\n", "# \"RANK TOP 5 FOR EACH orders.o_orderkey \"\n", "# \"WHERE COUNT(lineitem.*, -365, 0, days) > 0\"\n", "# ),\n", "# )\n", "# pquery.validate()\n", "# print(\"Task type:\", pquery.get_task_type())" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "9dfe2453-e09a-4048-908c-0e8ba9c89d2c", "showTitle": true, "tableResultSettingsMap": {}, "title": "Cell 27" } }, "source": [ "### Set task-specific prediction options\n", "\n", "Regression needs no extra options, so this cell is commented out. For binary classification or link prediction, uncomment it and the `**prediction_options` lines in the scoring cells below, including when reloading a model. The binary example uses a 0.5 threshold; the link example requests five candidates.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538351797, "inputWidgets": {}, "nuid": "4e8418a5-ad6a-4227-9505-db592d0d635d", "showTitle": true, "startTime": 1792538351593, "submitTime": 1792538351562, "tableResultSettingsMap": {}, "title": "Prediction options" } }, "outputs": [], "source": [ "# prediction_options = {}\n", "# task_type = pquery.get_task_type()\n", "# if task_type == TaskType.BINARY_CLASSIFICATION:\n", "# prediction_options[\"binary_classification_threshold\"] = 0.5\n", "# elif task_type.is_link_pred:\n", "# prediction_options[\"num_classes_to_return\"] = 5\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "831d3575-b531-48db-86bc-09d84f87ae65", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "### Generate a Training Table\n", "\n", "After your query is validated, generate a training table to use for model fitting." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538358803, "inputWidgets": {}, "nuid": "67a01d89-8b12-4f03-9fb3-fabd6f2d3060", "showTitle": false, "startTime": 1792538355731, "submitTime": 1792538355703, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "train_table = pquery.generate_training_table(snapshot)\n", "print(\"Training table:\", train_table.train_table_path)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "cd201dce-9df3-4643-9ca0-c9aa8f326d33", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "## Create a Trainer\n", "\n", "The suggested model plan controls architecture, sampling, and optimization. The model plan is fully customizable: inspect and adjust its settings for your task, then create `trainer = kumo.Trainer(model_plan)` before fitting." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792538361305, "inputWidgets": {}, "nuid": "d4aeb266-f28b-431c-b8fb-5a5da57eca99", "showTitle": false, "startTime": 1792538361132, "submitTime": 1792538361102, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "model_plan = pquery.suggest_model_plan(snapshot)\n", "# Create a Trainer with your model plan.\n", "trainer = kumo.Trainer(model_plan)\n", "print(model_plan)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "05b15eee-3685-4e1b-bef2-a5ec00897d97", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "### Train a Model\n", "\n", "After creating a Trainer from the model plan, call `fit()` with the training table and compatible graph snapshot.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539229805, "inputWidgets": {}, "nuid": "303ee1ab-add8-4d24-a9da-184f95058cda", "showTitle": true, "startTime": 1792538365237, "submitTime": 1792538365205, "tableResultSettingsMap": {}, "title": "Fit (Serverless GPU)" }, "language": "python" }, "outputs": [], "source": [ "# Serverless GPU workaround for Optuna child process IPC authentication errors.\n", "# import multiprocessing\n", "# try:\n", "# multiprocessing.set_start_method(\"spawn\", force=True)\n", "# except RuntimeError:\n", "# pass\n", "\n", "training_result = trainer.fit(train_table, snapshot=snapshot)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "9fc0890f-7252-4d08-bfb0-cc3861713325", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "### Inspect the selected result\n", "\n", "Review the reported metrics and their stage names. Available metrics vary by task and build. The local model bundle preserves model state; it does not archive the complete experiment history. The next cell prints its path. The bundle remains in temporary local scratch. If you want to reuse it, copy the complete bundle to durable storage before the session ends.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539252523, "inputWidgets": {}, "nuid": "31c7ccc6-5ff6-4c6a-a8b6-e95187795906", "showTitle": false, "startTime": 1792539252377, "submitTime": 1792539252336, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "print(\"Available metrics:\", training_result.metrics)\n", "print(\"Reusable local bundle:\", training_result.bundle_dir)\n", "for name, path in training_result.artifact_paths.items():\n", " print(name, path)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539257527, "inputWidgets": {}, "nuid": "ab6a48c8-0c3f-49dd-a3a3-7a1a47b0f4a1", "showTitle": true, "startTime": 1792539256389, "submitTime": 1792539256364, "tableResultSettingsMap": {}, "title": "Cell 56" }, "language": "python" }, "outputs": [], "source": [ "date_range = spark.sql(f\"\"\"\n", " SELECT MIN(o_orderdate) AS first_order, MAX(o_orderdate) AS last_order\n", " FROM {TPCH_CATALOG}.{TPCH_SCHEMA}.orders\n", "\"\"\")\n", "display(date_range.toPandas())\n", "print(\"Configured prediction boundary:\", PREDICTION_TIME)" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "96b310f8-0f1c-4045-9111-1e5ec2f04a27", "showTitle": false, "tableResultSettingsMap": {}, "title": "" }, "codeCollapsed": true }, "source": [ "### Generate the scoring population\n", "\n", "For these temporal queries, use a timezone-naive timestamp interpreted as UTC and reuse the fitted trainer's saved training-table generation plan. Binary classification and regression score customers with an order in the past year; the link example scores orders with recent line items. This timestamp example does not apply to static queries." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539263232, "inputWidgets": {}, "nuid": "7abbe611-bb5d-4947-b0e8-9c5dab497ce1", "showTitle": false, "startTime": 1792539260482, "submitTime": 1792539260456, "tableResultSettingsMap": {}, "title": "" }, "language": "python" }, "outputs": [], "source": [ "prediction_table = pquery.generate_prediction_table(\n", " snapshot,\n", " prediction_time=PREDICTION_TIME,\n", " plan=trainer.training_table_generation_plan,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "81e84dc7-7de4-40e8-b0bd-ec76b838cb94", "showTitle": true, "tableResultSettingsMap": {}, "title": "Score section" }, "codeCollapsed": true }, "source": [ "\n", "## Generate Batch Predictions\n", "\n", "After training a model, generate batch predictions and write them to a Unity Catalog table. `WRITE_PREDICTIONS = True` enables that write by default; set it to `False` for local-only results.\n", "\n", "Check the TPC-H order-date range against `PREDICTION_TIME` before scoring. A chosen date alone does not guarantee the following 90 days of observed outcomes for historical evaluation." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539299661, "inputWidgets": {}, "nuid": "0d1f4f80-b489-44a6-8d2b-612a69766cea", "showTitle": true, "startTime": 1792539268677, "submitTime": 1792539268650, "tableResultSettingsMap": {}, "title": "Prediction output" }, "language": "python" }, "outputs": [], "source": [ "output_name = f\"{CONNECTOR_SCHEMA}{_SEP}{PREDICTION_OUTPUT_TABLE}\"\n", "\n", "output_config = None\n", "if WRITE_PREDICTIONS:\n", " output_config = PredictionOutputConfig(\n", " connector=connector,\n", " artifact_type=PredictionArtifactType.PREDICTIONS,\n", " table_name=output_name,\n", " write_mode=WriteMode.OVERWRITE,\n", " )\n", "result = trainer.predict(\n", " prediction_table,\n", " snapshot=snapshot,\n", " output_config=output_config,\n", " #**prediction_options,\n", ")\n", "\n", "print(\"Output rows:\", result.num_rows)\n", "if WRITE_PREDICTIONS:\n", " print(\"Output table:\", f\"{CONNECTOR_CATALOG}.{CONNECTOR_SCHEMA}.{PREDICTION_OUTPUT_TABLE}\")\n", "else:\n", " print(\"Local prediction files:\", result.files)" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "70676f00-4498-4809-884d-074a9efd06d1", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "## Optional: inspect training diagnostics\n", "\n", "These cells read the completed fit's local files; they do not retrain or recalculate model metrics. They require pandas, NumPy, and Matplotlib. History and explanations are separate from the model bundle. Missing explanation files can reflect the task, model, or `training_job.disable_explain` setting.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539303364, "inputWidgets": {}, "nuid": "343c516c-87ab-45bb-817b-8ecea91987b6", "showTitle": false, "startTime": 1792539303163, "submitTime": 1792539303135, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "import json\n", "import numpy as np\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "from IPython.display import display\n", "\n", "run_dir = Path(training_result.bundle_dir).parent\n", "automl_dirs = sorted({p.parent for p in run_dir.rglob(\"auto_train_progress.json\")})\n", "if len(automl_dirs) != 1:\n", " raise RuntimeError(\"Expected one AutoML directory for this fit. Inspect the run directory.\")\n", "automl_dir = automl_dirs[0]\n", "selection = json.loads((automl_dir / \"best_trainer.json\").read_text())\n", "selected_trial = automl_dir / Path(selection[\"best_trainer\"]).name\n", "if not selected_trial.is_dir():\n", " raise FileNotFoundError(selected_trial)\n", "refit_trial = automl_dir / f\"{selected_trial.name}_refit_trainval\"\n", "explanation_dir = refit_trial if refit_trial.is_dir() else selected_trial\n", "print(\"Selected trial:\", selected_trial)\n", "print(\"Explanation trial:\", explanation_dir)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "3bd2f34c-05b4-4742-9180-1361d420369d", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Compare trials and inspect the selected plan\n", "\n", "The summary and learning curves describe all search trials, not a later refit. Compare validation results using the selected plan's tuning metric. Test results are for final evaluation, not tuning. Missing values are not zeros.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539307369, "inputWidgets": {}, "nuid": "8044635b-4270-4bcb-aacb-f63a87b7cd9a", "showTitle": false, "startTime": 1792539307180, "submitTime": 1792539307150, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "progress = json.loads((automl_dir / \"auto_train_progress.json\").read_text())\n", "print(\"Completed / total trials:\", progress.get(\"completed_trials\"), \"/\", progress.get(\"total_trials\"))\n", "summaries, rows = [], []\n", "for trial, details in (progress.get(\"trial_progress\") or {}).items():\n", " details = details or {}\n", " summaries.append({\"trial\": trial, \"epochs\": details.get(\"num_epochs_completed\"),\n", " \"seconds\": details.get(\"elapsed_time\"), \"accelerator\": details.get(\"accelerator\")})\n", " for epoch, metrics in (details.get(\"metrics\") or {}).items():\n", " metrics = metrics or {}\n", " rows.append({\"trial\": trial, \"epoch\": int(epoch),\n", " **(metrics.get(\"train_metrics\") or {}), **(metrics.get(\"validation_metrics\") or {})})\n", "display(pd.DataFrame(summaries))\n", "learning_curves = pd.DataFrame(rows)\n", "if learning_curves.empty:\n", " raise ValueError(\"No epoch metrics were recorded. Check the training logs.\")\n", "learning_curves = learning_curves.sort_values([\"trial\", \"epoch\"])\n", "metric_columns = learning_curves.columns.difference([\"trial\", \"epoch\"])\n", "learning_curves[metric_columns] = learning_curves[metric_columns].apply(pd.to_numeric, errors=\"coerce\").replace([np.inf, -np.inf], np.nan)\n", "display(learning_curves)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539313773, "inputWidgets": {}, "nuid": "faaa5f8c-eb20-4c61-b1e7-f247fbf51a4c", "showTitle": false, "startTime": 1792539313543, "submitTime": 1792539313514, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "validation_files = [p for p in sorted(automl_dir.glob(\"version_*/val_result.csv\"))\n", " if p.parent.name.removeprefix(\"version_\").isdigit()]\n", "validation_results = [pd.read_csv(p).assign(trial=p.parent.name) for p in validation_files]\n", "display(pd.concat(validation_results, ignore_index=True) if validation_results else \"No validation CSVs were recorded.\")\n", "selected_plan = json.loads((selected_trial / \"model_plan.json\").read_text())\n", "tune_metric = selected_plan[\"training_job\"][\"tune_metric\"]\n", "print(\"Selected trial tuning metric:\", tune_metric)\n", "display(selected_plan)\n", "config_path = training_result.artifact_paths.get(\"trainer_config\")\n", "if config_path:\n", " print(\"Fitted trainer configuration:\")\n", " display(json.loads(Path(config_path).read_text()))\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "d4e9db4b-99f7-4861-9115-6c9953f397da", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Plot learning curves\n", "\n", "Plot loss and the tuning metric for each search trial. Missing or nonfinite values stay as gaps. A train/validation gap can suggest further investigation; it does not prove overfitting because sampling can differ.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539317977, "inputWidgets": {}, "nuid": "96d56ac2-534b-4acd-bd99-70e084b570c2", "showTitle": true, "startTime": 1792539317579, "submitTime": 1792539317549, "tableResultSettingsMap": {}, "title": "Plot learning curves" } }, "outputs": [], "source": [ "for metric in dict.fromkeys([\"loss\", tune_metric]):\n", " columns = [c for c in (f\"train_{metric}\", f\"val_{metric}\") if c in learning_curves]\n", " if not columns or not learning_curves[columns].notna().any().any():\n", " print(\"No finite values recorded for:\", metric)\n", " continue\n", " fig, ax = plt.subplots(figsize=(9, 5))\n", " for trial, values in learning_curves.groupby(\"trial\"):\n", " for column in columns:\n", " ax.plot(values[\"epoch\"], values[column], marker=\"o\", label=f\"{trial}: {column}\")\n", " ax.set(xlabel=\"Epoch\", ylabel=metric, title=f\"Search trials: {metric}\")\n", " ax.grid(alpha=0.25)\n", " ax.legend()\n", " display(fig)\n", " plt.close(fig)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "6deb4cfc-67d5-4c2e-971a-75aa5ad46a8c", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Read explanations by inspecting gradient scores\n", "\n", "Read explanations for the selected trial or its train/validation refit. Gradient scores describe input-column sensitivity on the explanation sample, not causal effects. Missing scores do not mean a column was excluded.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539322281, "inputWidgets": {}, "nuid": "8cb4b3ba-075e-4bb4-a482-7ba440adec5c", "showTitle": true, "startTime": 1792539321954, "submitTime": 1792539321929, "tableResultSettingsMap": {}, "title": "Gradient scores" } }, "outputs": [], "source": [ "scores_path = explanation_dir / \"explain_test/scores.csv\"\n", "if scores_path.is_file():\n", " scores = pd.read_csv(scores_path)\n", " scores[\"score\"] = pd.to_numeric(scores[\"score\"], errors=\"coerce\")\n", " if scores.empty or not np.isfinite(scores[\"score\"]).all():\n", " raise ValueError(\"Expected nonempty, finite gradient scores.\")\n", " display(scores.sort_values(\"score\", ascending=False))\n", " ranking = scores.nlargest(12, \"score\").sort_values(\"score\")\n", " fig, ax = plt.subplots(figsize=(8, 5))\n", " ax.barh(ranking[\"table\"] + \".\" + ranking[\"column\"], ranking[\"score\"])\n", " ax.set(xlabel=\"Gradient attribution score\", title=\"Input column sensitivity\")\n", " fig.tight_layout()\n", " display(fig)\n", " plt.close(fig)\n", "else:\n", " print(\"No gradient scores. Check task eligibility, disable_explain, and fit logs.\")\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "6ef800ef-af3e-42b4-b4be-d62bd0255cee", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Compare predictions and observed targets by cohort\n", "\n", "Select a `(table_name, hop, col_name)` from the displayed keys. For classification, choose an available class. Regression needs no class selection. The upper panel shows cohort averages; the lower panel shows population shares. Small cohorts and associations do not establish causal effects. Link prediction and forecasting do not support column cohorts.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539326483, "inputWidgets": {}, "nuid": "2ac3b5aa-bb74-41f7-9a92-33f77db432c8", "showTitle": true, "startTime": 1792539326089, "submitTime": 1792539326063, "tableResultSettingsMap": {}, "title": "Column cohorts" } }, "outputs": [], "source": [ "analysis_path = explanation_dir / \"column_analysis/test/column_analysis.json\"\n", "records = []\n", "if analysis_path.is_file():\n", " records = json.loads(analysis_path.read_text()).get(\"cohorts\") or []\n", "if records:\n", " keys = [(r[\"table_name\"], r[\"hop\"], r[\"col_name\"]) for r in records]\n", " print(\"Available columns:\", keys)\n", " column_key = keys[0] # Choose one of the keys above.\n", " matches = [r for key, r in zip(keys, records) if key == column_key]\n", " if len(matches) != 1:\n", " raise ValueError(\"Select one unique table, hop, and column.\")\n", " record = matches[0]\n", " predictions = record[\"predictions\"]\n", " class_name = None\n", " if isinstance(predictions, dict):\n", " if not predictions:\n", " raise ValueError(\"No prediction classes were recorded for this column.\")\n", " print(\"Available classes:\", list(predictions))\n", " class_name = next(iter(predictions)) # Choose an available class.\n", " print(\"Selected column and class:\", column_key, class_name)\n", " labels = record[\"cohorts\"]\n", " def cohort_values(field):\n", " values = record[field]\n", " values = values[class_name] if isinstance(values, dict) else values\n", " values = np.asarray(values, dtype=float)\n", " if values.ndim != 1 or len(values) != len(labels) or not np.isfinite(values).all():\n", " raise ValueError(f\"Invalid cohort values for {field}.\")\n", " return values\n", " population = cohort_values(\"populations\")\n", " if (population < 0).any() or (population > 1).any() or population.sum() <= 0:\n", " raise ValueError(\"Expected population shares between zero and one.\")\n", " fig, (ax, shares) = plt.subplots(2, 1, figsize=(8, 6), sharex=True)\n", " ax.plot(labels, cohort_values(\"predictions\"), \"o-\", label=\"Prediction\")\n", " if record.get(\"targets\") is not None:\n", " ax.plot(labels, cohort_values(\"targets\"), \"s--\", label=\"Observed target\")\n", " ax.set(title=f\"{column_key}, class={class_name}\", ylabel=\"Cohort average\")\n", " ax.legend()\n", " shares.bar(labels, population * 100)\n", " shares.set(xlabel=\"Feature value cohort\", ylabel=\"Population share (%)\")\n", " fig.tight_layout()\n", " display(fig)\n", " plt.close(fig)\n", "else:\n", " print(\"No column cohorts. Check task eligibility, disable_explain, and fit logs.\")\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "aa059a1e-d5e4-47f4-b225-18bd1f246687", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Locate other diagnostics\n", "\n", "The selected fit can also contain TensorBoard events, failed-trial details, optimization analysis, and sampled graph context. Their absence does not invalidate a completed fit.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539331088, "inputWidgets": {}, "nuid": "c032a9aa-97f1-4a12-aa01-1f7df7811189", "showTitle": false, "startTime": 1792539330920, "submitTime": 1792539330880, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "for pattern in [\"best_trainer.json\", \"version_*/train_progress.json\",\n", " \"**/events.out.tfevents.*\", \"failed_trials/*\", \"analyze/*\"]:\n", " paths = sorted(automl_dir.glob(pattern))\n", " print(pattern, \":\", len(paths), \"matches\")\n", "subgraph_paths = sorted((explanation_dir / \"subgraph_collector/test/subgraph_json\").glob(\"*.json\"))\n", "print(\"Sampled graph files:\", len(subgraph_paths))\n", "if subgraph_paths:\n", " print(\"First graph fields:\", list(json.loads(subgraph_paths[0].read_text())))\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "caa60a33-b7b3-4a75-b300-340d85382091", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "## Read a bounded prediction preview\n", "\n", "Read at most ten rows in Parquet batches. Accessing the result frame first would load every row, even if followed by `.head()`.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539335191, "inputWidgets": {}, "nuid": "8e0fc65c-c225-453b-9b73-f3af63a6b9b5", "showTitle": false, "startTime": 1792539334966, "submitTime": 1792539334940, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "import pandas as pd\n", "from IPython.display import display\n", "import pyarrow.parquet as pq\n", "\n", "def preview_parquet(files, limit=10):\n", " frames, remaining = [], limit\n", " for path in files:\n", " if remaining <= 0:\n", " break\n", " with pq.ParquetFile(path) as parquet:\n", " for batch in parquet.iter_batches(batch_size=remaining):\n", " frames.append(batch.slice(0, remaining).to_pandas())\n", " remaining -= len(frames[-1])\n", " if remaining <= 0:\n", " break\n", " return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()\n", "\n", "display(preview_parquet(result.files))\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "ab7335d8-bea5-43ee-b858-d5da9463b227", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "## Optional: preserve diagnostics\n", "\n", "The complete AutoML archive preserves history separately from the model bundle. It can include large checkpoints and sensitive source-derived values. Run the archive cells only when you want to retain this history, using a restricted destination. The ZIP is built on local scratch first.\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "a409fda1-f57c-4fc6-b817-a48f5f02394e", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Configure retained storage\n", "\n", "Set this to an existing Unity Catalog Volume with read/write access. Run this configuration cell for either diagnostic archiving or model retention; the two operations are independent.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539339595, "inputWidgets": {}, "nuid": "63bfb260-553d-42ea-a09e-94c96f454b5a", "showTitle": true, "startTime": 1792539338667, "submitTime": 1792539338642, "tableResultSettingsMap": {}, "title": "Configure retained storage" } }, "outputs": [], "source": [ "from uuid import uuid4\n", "from zipfile import ZipFile\n", "\n", "ARTIFACT_ROOT = Path(\"/Volumes////kumo-runs\")\n", "retained_run = ARTIFACT_ROOT / uuid4().hex\n", "retained_run.mkdir(parents=True, exist_ok=False)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539344499, "inputWidgets": {}, "nuid": "3d94496c-0c6e-4570-b841-93114c37f7de", "showTitle": false, "startTime": 1792539342713, "submitTime": 1792539342688, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "archive_dir = Path(tempfile.mkdtemp(dir=SCRATCH_DIR))\n", "local_zip = Path(shutil.make_archive(str(archive_dir / uuid4().hex), \"zip\", root_dir=automl_dir))\n", "archive_path = Path(shutil.copyfile(local_zip, retained_run / local_zip.name))\n", "print(\"Record archive path:\", archive_path)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "5a64249d-61fb-4962-8247-c1e27be7a3bf", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "### Retrieve and inspect an archive\n", "\n", "Copy the recorded Volume archive back to local scratch and list its contents without extracting it.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539348203, "inputWidgets": {}, "nuid": "65bc2601-dc07-4269-8de1-7b76cf4485ce", "showTitle": false, "startTime": 1792539347946, "submitTime": 1792539347923, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "downloaded_zip = Path(tempfile.mkdtemp(dir=SCRATCH_DIR)) / archive_path.name\n", "shutil.copyfile(archive_path, downloaded_zip)\n", "with ZipFile(downloaded_zip) as archive:\n", " print(\"Members:\", len(archive.infolist()))\n", " for member in archive.infolist()[:20]:\n", " print(member.filename, member.file_size, \"bytes\")\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "475b6e3e-68e5-4654-8b2e-d9e3645ff3a7", "showTitle": false, "tableResultSettingsMap": {}, "title": "" } }, "source": [ "## Optional: retain and reload the model\n", "\n", "The complete model bundle preserves fitted state, not source data, graph snapshots, detailed history, or predictions. Set the Volume destination in the first archive-section configuration cell before retaining a model; creating a diagnostic archive is optional. Keep the recorded model address. In a later session, rerun setup through the graph snapshot, predictive query, and task-specific prediction options, skip the original training and prediction cells, set the recorded address, then run the reload and scoring cells below. Load only a trusted model bundle.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539354208, "inputWidgets": {}, "nuid": "cc22175c-bddf-4f9e-a24c-20073842f37e", "showTitle": false, "startTime": 1792539352056, "submitTime": 1792539352034, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "retained_bundle = retained_run / \"model\"\n", "shutil.copytree(training_result.bundle_dir, retained_bundle)\n", "print(\"Record model path:\", retained_bundle)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539358311, "inputWidgets": {}, "nuid": "4f2c1297-5c15-4aa9-bd9c-de09203b99c0", "showTitle": false, "startTime": 1792539357473, "submitTime": 1792539357446, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "# In a later session, set retained_bundle to the recorded Volume path.\n", "local_bundle = Path(tempfile.mkdtemp(dir=SCRATCH_DIR)) / \"model\"\n", "shutil.copytree(retained_bundle, local_bundle)\n", "reloaded_trainer = kumo.Trainer.load(local_bundle)\n" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": {}, "inputWidgets": {}, "nuid": "829e3899-d336-443e-9237-c199504e1018", "showTitle": true, "tableResultSettingsMap": {}, "title": "Score with the reloaded model" } }, "source": [ "### Score with the reloaded model\n", "\n", "Regenerate the scoring population using the restored trainer's saved generation plan. Predictions stay local. For binary classification or link prediction, uncomment `**prediction_options` to apply the same task-specific options.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "finishTime": 1792539368520, "inputWidgets": {}, "nuid": "5cba78dc-aefc-440d-bae9-25539355c7c6", "showTitle": false, "startTime": 1792539361138, "submitTime": 1792539361113, "tableResultSettingsMap": {}, "title": "" } }, "outputs": [], "source": [ "reloaded_population = pquery.generate_prediction_table(\n", " snapshot, prediction_time=PREDICTION_TIME,\n", " plan=reloaded_trainer.training_table_generation_plan,\n", ")\n", "reloaded_result = reloaded_trainer.predict(\n", " reloaded_population, snapshot=snapshot, \n", " #**prediction_options,\n", ")\n", "print(\"Reloaded model rows:\", reloaded_result.num_rows)\n", "print(\"Local prediction files:\", reloaded_result.files)\n" ] } ], "metadata": { "application/vnd.databricks.v1+notebook": { "computePreferences": { "hardware": { "accelerator": "GPU_1xH100", "gpuPoolId": null, "memory": null }, "software": { "pinSparkToX86": null } }, "dashboards": [], "environmentMetadata": { "base_environment": "databricks_ai_v6", "environment_version": "6" }, "inputWidgetPreferences": null, "language": "python", "notebookMetadata": { "pythonIndentUnit": 4 }, "notebookName": "databricks_docs_example_notebook1", "widgets": {} }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }