{ "cells": [ { "cell_type": "markdown", "id": "prerequisites", "metadata": { "codeCollapsed": false, "collapsed": false }, "source": "## Prepare your environment\n\n1. **Runtime and compute:** select a Snowflake Notebook on GPU Container Runtime 2.9 with Python 3.12, backed by a GPU compute pool. The query warehouse runs SQL; increasing its size does not increase GPU memory. Choose a pool with enough GPU memory and CPU RAM for the model and graph; there is no validated one-size-fits-all GPU size for this sample.\n2. **Local working disk:** 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. Confirm that `SCRATCH_DIR` uses the intended supported local disk; the notebook prints the available space at that path.\n3. **Access and permissions:** package installation needs an approved PyPI route, such as a Snowflake external access integration (EAI). The role needs the privileges below; a successful import does not prove them:\n\n| Action | Required access |\n| --- | --- |\n| Run notebook | `USAGE` on the GPU compute pool and query warehouse; access to the notebook |\n| Install | PyPI access through an approved EAI and `USAGE` on that EAI, where required |\n| Read data | `IMPORTED PRIVILEGES` on shared database `SNOWFLAKE_SAMPLE_DATA` (or equivalent inherited access) and query access to `TPCH_SF1` |\n| Create views | Permission to create the database/schema and create or replace these views; otherwise have an administrator prepare them and skip the view-setup cell |\n| Publish a model | `CREATE MODEL` on the configured schema (or schema ownership) |\n| Save predictions | `CREATE SCHEMA` on the database for `OUTPUT_SCHEMA` (or an existing schema), and permission to create or overwrite the output table in it |\n| Archive diagnostics (optional) | `CREATE STAGE` on the configured schema for `KUMO_DIAGNOSTICS` (or an existing stage), plus read and write on it |\n\nThis notebook uses four TPC-H tables: `CUSTOMER`, `PART`, `ORDERS`, and `LINEITEM`. These are accessed through views created in your configured database/schema over the read-only `SNOWFLAKE_SAMPLE_DATA.TPCH_SF1` data. The LINEITEM view adds a synthetic single-column primary key (`LINEITEM_ID`) since TPC-H uses a composite key (`L_ORDERKEY`, `L_LINENUMBER`).\n\n**Quickstart sample:** by default, the views select 2,000 customers from TPC-H SF1, their orders and line items, and the parts referenced by those line items. This preserves the graph's customer–order–line-item–part relationships. Increase `SAMPLE_CUSTOMERS` only after the small run succeeds.\n\nSee [Snowflake GPU runtime 2.9](https://docs.snowflake.com/en/developer-guide/snowflake-ml/container-runtime/releases/gpu/2_9) and [Notebook setup privileges](https://docs.snowflake.com/en/user-guide/ui-snowsight/notebooks-setup) for current platform details.\n\n**Scratch** is local working storage for snapshots, generated rows, model files, predictions, and diagnostics. It may survive a kernel restart if the filesystem remains, but container or machine loss can remove it. Model publication preserves the reusable model; other scratch files remain temporary.\n\n**Each user is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use**" }, { "cell_type": "markdown", "id": "configuration-intro", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## 1. Enter your configuration\n\nFill in your warehouse, database, and schema below. Use ordinary unquoted database and schema identifiers. The view-setup cell creates views in your configured database over the read-only `SNOWFLAKE_SAMPLE_DATA.TPCH_SF1`. Views in an owned database are required because Kumo's snapshot uses staged reads, which Snowflake does not allow on shared/imported databases.\n\nThis workflow trains and publishes a model, then generates predictions and writes them to the configured Snowflake output table.\n\nJanuary 1, 1998 is the configured prediction date. TPC-H SF1 data spans approximately 1992-01-02 to 1998-08-02. After a kernel restart, rerun this configuration and skip installation." }, { "cell_type": "code", "execution_count": null, "id": "configuration", "metadata": { "language": "python" }, "outputs": [], "source": "from datetime import datetime\n\nWAREHOUSE_NAME = \"\"\nDATABASE_NAME = \"\" # Created if it does not exist.\nSCHEMA_NAME = \"\" # Source views are created here.\n\nMODEL_NAME = \"\" # Model registry name; each fit adds a version.\nPREDICTION_TIME = datetime(1998, 1, 1) # TPC-H data spans ~1992-01-02 to 1998-08-02.\nOUTPUT_SCHEMA = \"\" # Created if it does not exist.\nPREDICTION_OUTPUT_TABLE = \"\"\nTPCH_SCHEMA = \"TPCH_SF1\"\nSAMPLE_CUSTOMERS = 2000 # Small, relationship-preserving quickstart slice." }, { "id": "e93c1a05-ed23-49d1-8657-090ef8244648", "cell_type": "markdown", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## 1b. Create views over TPC-H sample data\n\n`SNOWFLAKE_SAMPLE_DATA` is a shared database that does not allow staged exports. This cell creates views in your configured database over a bounded, relationship-preserving slice of four TPC-H SF1 tables so Kumo can unload them to local Parquet. The LINEITEM view adds a synthetic primary key (`LINEITEM_ID`) since TPC-H uses a composite key.\n\nRun this once per environment. Safe to rerun (uses `IF NOT EXISTS` and `OR REPLACE`). The printed row counts describe this run's sample; views do not copy source data." }, { "id": "08e70897-d6f1-4caa-8c75-29ed91283f60", "cell_type": "code", "metadata": { "language": "python", "title": "Create TPC-H views" }, "source": "from snowflake.snowpark.context import get_active_session\nfrom time import perf_counter\n\nview_setup_started = perf_counter()\nsession = get_active_session()\n\n# Create writable database with views over read-only SNOWFLAKE_SAMPLE_DATA.\n# Staged reads require an owned database; SNOWFLAKE_SAMPLE_DATA is shared and read-only.\nsession.sql(f\"CREATE DATABASE IF NOT EXISTS {DATABASE_NAME}\").collect()\nsession.sql(f\"CREATE SCHEMA IF NOT EXISTS {DATABASE_NAME}.{SCHEMA_NAME}\").collect()\n\nsession.sql(f\"\"\"\nCREATE OR REPLACE VIEW {DATABASE_NAME}.{SCHEMA_NAME}.CUSTOMER AS\nSELECT C_CUSTKEY, C_NAME, C_NATIONKEY, C_MKTSEGMENT, C_ACCTBAL\nFROM SNOWFLAKE_SAMPLE_DATA.{TPCH_SCHEMA}.CUSTOMER\nWHERE C_CUSTKEY <= {SAMPLE_CUSTOMERS}\n\"\"\").collect()\n\nsession.sql(f\"\"\"\nCREATE OR REPLACE VIEW {DATABASE_NAME}.{SCHEMA_NAME}.ORDERS AS\nSELECT O_ORDERKEY, O_CUSTKEY, O_ORDERDATE, O_ORDERSTATUS, O_ORDERPRIORITY, O_TOTALPRICE, O_SHIPPRIORITY\nFROM SNOWFLAKE_SAMPLE_DATA.{TPCH_SCHEMA}.ORDERS\nWHERE O_CUSTKEY <= {SAMPLE_CUSTOMERS}\n\"\"\").collect()\n\n# LINEITEM needs a synthetic single-column primary key (TPC-H uses composite L_ORDERKEY + L_LINENUMBER).\nsession.sql(f\"\"\"\nCREATE OR REPLACE VIEW {DATABASE_NAME}.{SCHEMA_NAME}.LINEITEM AS\nSELECT\n (L_ORDERKEY * 100 + L_LINENUMBER) AS LINEITEM_ID,\n L_ORDERKEY, L_PARTKEY, L_LINENUMBER, L_QUANTITY,\n L_EXTENDEDPRICE, L_DISCOUNT, L_TAX, L_SHIPDATE,\n L_RETURNFLAG, L_LINESTATUS, L_SHIPMODE\nFROM SNOWFLAKE_SAMPLE_DATA.{TPCH_SCHEMA}.LINEITEM\nWHERE L_ORDERKEY IN (\n SELECT O_ORDERKEY FROM {DATABASE_NAME}.{SCHEMA_NAME}.ORDERS\n)\n\"\"\").collect()\n\nsession.sql(f\"\"\"\nCREATE OR REPLACE VIEW {DATABASE_NAME}.{SCHEMA_NAME}.PART AS\nSELECT P_PARTKEY, P_NAME, P_MFGR, P_BRAND, P_TYPE, P_SIZE, P_CONTAINER, P_RETAILPRICE\nFROM SNOWFLAKE_SAMPLE_DATA.{TPCH_SCHEMA}.PART\nWHERE P_PARTKEY IN (\n SELECT L_PARTKEY FROM {DATABASE_NAME}.{SCHEMA_NAME}.LINEITEM\n)\n\"\"\").collect()\n\nfor table_name in (\"CUSTOMER\", \"ORDERS\", \"LINEITEM\", \"PART\"):\n count = session.sql(\n f\"SELECT COUNT(*) FROM {DATABASE_NAME}.{SCHEMA_NAME}.{table_name}\"\n ).collect()[0][0]\n print(f\"{table_name} rows: {count:,}\")\nprint(f\"View setup and row counts: {perf_counter() - view_setup_started:.1f} seconds\")", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "id": "install-intro", "metadata": { "codeCollapsed": true }, "source": "## 2. Install Kumo Engine from PyPI\n\nRun this installation cell once. **Then restart the kernel, rerun configuration, skip installation, and continue at imports.** Do not use Run All across the restart boundary. Installation uses the active interpreter." }, { "cell_type": "code", "execution_count": null, "id": "install", "metadata": { "language": "python" }, "outputs": [], "source": "%pip install kumo-engine[snowflake]" }, { "cell_type": "markdown", "id": "runtime", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "\n## Import and configure scratch\n\nAfter each kernel restart, run configuration and these cells. They do not depend on variables created during installation. Configure scratch before constructing any Kumo table, graph, or trainer." }, { "cell_type": "code", "execution_count": null, "id": "imports", "metadata": { "language": "python" }, "outputs": [], "source": "import sys\nimport shutil\nimport tempfile\nfrom importlib.metadata import version\nfrom pathlib import Path\nfrom snowflake.snowpark.context import get_active_session\n\nimport kumo.engine as kumo\nfrom kumo.engine import SnowflakeConnector\nfrom kumoapi.graph import Edge\nfrom kumoapi.jobs import PredictionArtifactType, WriteMode\nfrom kumoapi.task import TaskType\n\nsession = get_active_session()" }, { "cell_type": "code", "execution_count": null, "id": "scratch", "metadata": { "language": "python" }, "outputs": [], "source": "SCRATCH_DIR = \"/tmp/kumoai-local\"\nscratch = Path(SCRATCH_DIR)\nscratch.mkdir(parents=True, exist_ok=True)\nwith tempfile.TemporaryFile(dir=scratch) as probe:\n probe.write(b\"kumo scratch check\")\n probe.flush()\nkumo.configure(scratch_dir=SCRATCH_DIR)\nprint(\"Kumo version:\", kumo.__version__)\nprint(\"Free working space (GiB):\", round(shutil.disk_usage(scratch).free / 2**30, 1))" }, { "cell_type": "markdown", "id": "scratch-result", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "The printed version identifies the installed build. The probe checks write access and free space at this path; it does not establish that the path is on supported local storage. Confirm the backing disk and capacity before materializing data." }, { "cell_type": "markdown", "id": "connect", "metadata": { "codeCollapsed": true }, "source": "## Connect and inspect the source\n\n`from_spcs_env()` uses runtime-provided account and host settings and the mounted OAuth token. Never copy credentials into this notebook. This connector reads the source views and publishes the model in your configured database/schema; predictions are written later through a second connector for `OUTPUT_SCHEMA`. Confirm all four table views appear." }, { "cell_type": "code", "execution_count": null, "id": "connector", "metadata": { "language": "python" }, "outputs": [], "source": "connector = kumo.SnowflakeConnector.from_spcs_env(\n warehouse=WAREHOUSE_NAME, database=DATABASE_NAME, schema_name=SCHEMA_NAME,\n)\nprint(\"Tables:\", connector.table_names())" }, { "cell_type": "markdown", "id": "tables", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## Describe the four tables\n\nA primary key identifies a row, and a time column orders events. Kumo infers each column's physical type (its stored representation, such as a date) and semantic type (its modeling role, such as a category or numerical value)." }, { "cell_type": "code", "execution_count": null, "id": "dimension-tables", "metadata": { "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\npart = (\n kumo.Table.from_source_table(\n connector[\"PART\"],\n [\"P_PARTKEY\", \"P_NAME\", \"P_BRAND\", \"P_TYPE\",\n \"P_SIZE\", \"P_CONTAINER\", \"P_RETAILPRICE\"],\n primary_key=\"P_PARTKEY\",\n )\n)" }, { "cell_type": "markdown", "id": "timestamps", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "Event timestamps let Kumo separate information available before a prediction boundary from future outcomes. Using future information as an input causes temporal leakage." }, { "cell_type": "code", "execution_count": null, "id": "order-header", "metadata": { "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_ORDERPRIORITY\", \"O_TOTALPRICE\"],\n primary_key=\"O_ORDERKEY\",\n time_column=\"O_ORDERDATE\",\n )\n)" }, { "cell_type": "code", "execution_count": null, "id": "order-detail", "metadata": { "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\"],\n primary_key=\"LINEITEM_ID\",\n time_column=\"L_SHIPDATE\",\n )\n)" }, { "cell_type": "markdown", "id": "graph", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## Connect the tables and materialize a snapshot\n\nA **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\nLINEITEM.L_PARTKEY -> PART.P_PARTKEY\nLINEITEM.L_ORDERKEY -> ORDERS.O_ORDERKEY\nORDERS.O_CUSTKEY -> CUSTOMER.C_CUSTKEY\n```\n\nValidation 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, "id": "graph-definition", "metadata": { "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\ngraph.validate()\n\nsnapshot = graph.snapshot()\n\nprint(\"Snapshot:\", snapshot.id)" }, { "cell_type": "markdown", "id": "question", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## Define the prediction question\n\nPredictive 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, "id": "pquery", "metadata": { "language": "python" }, "outputs": [], "source": "# Regression: predict how many orders a customer active in the past year will place in the next 90 days.\npquery = 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)\npquery.validate()\nprint(\"Task type:\", pquery.get_task_type())" }, { "cell_type": "markdown", "id": "other-tasks", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "### Optional task alternatives\n\nThe 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. " }, { "cell_type": "code", "execution_count": null, "id": "binary-query", "metadata": { "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, "id": "link-query", "metadata": { "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": { "codeCollapsed": true }, "source": "### Set task-specific prediction options\n\nRun this after choosing your query, including when reloading a model. Regression needs no extra options. The link example requests five candidates.\n", "id": "prediction-options-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "execution_count": null, "outputs": [], "source": "prediction_options = {}\ntask_type = pquery.get_task_type()\nif task_type == TaskType.BINARY_CLASSIFICATION:\n prediction_options[\"binary_classification_threshold\"] = 0.5\nelif task_type.is_link_pred:\n prediction_options[\"num_classes_to_return\"] = 5\n", "id": "prediction-options" }, { "cell_type": "markdown", "id": "generate-intro", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "### Generate a Training Table\n\nAfter your query is validated, generate a training table to use for model fitting. " }, { "cell_type": "code", "execution_count": null, "id": "generate-training", "metadata": { "language": "python" }, "outputs": [], "source": "train_table = pquery.generate_training_table(snapshot)" }, { "cell_type": "markdown", "id": "model-route", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "## Create a Trainer\n\nThe suggested model plan controls architecture, sampling, and optimization. The model plan is fully customizable; inspect and adjust its settings and validation selection metric for your task before creating `trainer = kumo.Trainer(model_plan)`. Compare like-for-like validation metrics and reserve test metrics for evaluating the selected model." }, { "cell_type": "code", "execution_count": null, "id": "plan-or-reload", "metadata": { "language": "python" }, "outputs": [], "source": "model_plan = pquery.suggest_model_plan(snapshot)\n\ntrainer = kumo.Trainer(model_plan)" }, { "cell_type": "markdown", "id": "fit-intro", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "### Train a Model\n\nAfter creating a Trainer from the model plan, call `fit()` with the training table and compatible graph snapshot. The fit step also publishes the model to Snowflake.\n\nFor a deliberately local-only experiment, set `model_output = None` in this cell before running it." }, { "cell_type": "code", "execution_count": null, "id": "fit", "metadata": { "language": "python" }, "outputs": [], "source": "model_output = kumo.ModelOutputConfig(\n connector, location=kumo.ModelRegistryLocation(MODEL_NAME),\n)\n\ntraining_result = trainer.fit(\n train_table, snapshot=snapshot, export_config=model_output,\n)" }, { "cell_type": "markdown", "id": "result-intro", "metadata": { "codeCollapsed": true }, "source": "### Inspect the selected result\n\nReview the reported metrics and their stage names. Available metrics vary by task and build; this example does not promise calibration output or a fixed set. A successful fit alone does not establish business usefulness.\n\nThe reusable bundle and registry publication preserve model state. They do not archive the complete experiment history. Continue to batch prediction below." }, { "cell_type": "code", "execution_count": null, "id": "training-result", "metadata": { "language": "python" }, "outputs": [], "source": "print(\"Available metrics:\", training_result.metrics)\nprint(\"Reusable bundle:\", training_result.bundle_dir)" }, { "cell_type": "markdown", "id": "prediction-table-intro", "metadata": { "codeCollapsed": true }, "source": "### Generate the scoring population\n\nFor this temporal query, use a timezone-naive timestamp interpreted as UTC. Reuse the fitted trainer's saved training-table generation plan. The generated population contains customers meeting the past-year restriction at that date. This timestamp example does not apply to static queries." }, { "cell_type": "code", "execution_count": null, "id": "prediction-table", "metadata": { "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", "id": "prediction-write-warning", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "\n## Generate Batch Predictions\n\nAfter training a model, generate batch predictions and write them to the configured Snowflake table. This cell creates `OUTPUT_SCHEMA` if needed and writes through a separate connector for that schema.\n\n`PREDICTION_TIME` must fall within the source date range. A chosen date alone does not guarantee the following 90 days of observed outcomes for historical evaluation." }, { "cell_type": "code", "execution_count": null, "id": "predict", "metadata": { "language": "python" }, "outputs": [], "source": "# The output schema must exist before Kumo writes to it.\nsession.sql(f\"CREATE SCHEMA IF NOT EXISTS {DATABASE_NAME}.{OUTPUT_SCHEMA}\").collect()\n\n# Separate connector for writes, reusing the source connector's identity and warehouse.\noutput_connector = SnowflakeConnector(\n identity=connector.identity,\n warehouse=connector.warehouse,\n database=connector.database,\n role=connector.role,\n schema_name=OUTPUT_SCHEMA,\n)\n\nprediction_output = kumo.PredictionOutputConfig(\n connector=output_connector,\n artifact_type=PredictionArtifactType.PREDICTIONS,\n table_name=PREDICTION_OUTPUT_TABLE,\n write_mode=WriteMode.OVERWRITE,\n)\n\nprediction_result = trainer.predict(\n prediction_table, snapshot=snapshot, output_config=prediction_output, **prediction_options,\n)\nprint(\"Rows:\", prediction_result.num_rows)\nprint(\"Local Parquet files:\", prediction_result.files)\nprint(\"Snowflake output:\", f\"{DATABASE_NAME}.{OUTPUT_SCHEMA}.{PREDICTION_OUTPUT_TABLE}\")" }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "## Optional: inspect training diagnostics\n\nThese 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", "id": "diagnostics-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "import json\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom IPython.display import display\n\nrun_dir = Path(training_result.bundle_dir).parent\nautoml_dirs = sorted({p.parent for p in run_dir.rglob(\"auto_train_progress.json\")})\nif len(automl_dirs) != 1:\n raise RuntimeError(\"Expected one AutoML directory for this fit. Inspect the run directory.\")\nautoml_dir = automl_dirs[0]\nselection = json.loads((automl_dir / \"best_trainer.json\").read_text())\nselected_trial = automl_dir / Path(selection[\"best_trainer\"]).name\nif not selected_trial.is_dir():\n raise FileNotFoundError(selected_trial)\nrefit_trial = automl_dir / f\"{selected_trial.name}_refit_trainval\"\nexplanation_dir = refit_trial if refit_trial.is_dir() else selected_trial\nprint(\"Selected trial:\", selected_trial)\nprint(\"Explanation trial:\", explanation_dir)\n", "id": "select-experiment", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Compare trials and inspect the selected plan\n\nThe 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", "id": "trial-metrics-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "progress = json.loads((automl_dir / \"auto_train_progress.json\").read_text())\nprint(\"Completed / total trials:\", progress.get(\"completed_trials\"), \"/\", progress.get(\"total_trials\"))\nsummaries, rows = [], []\nfor 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 {})})\ndisplay(pd.DataFrame(summaries))\nlearning_curves = pd.DataFrame(rows)\nif learning_curves.empty:\n raise ValueError(\"No epoch metrics were recorded. Check the training logs.\")\nlearning_curves = learning_curves.sort_values([\"trial\", \"epoch\"])\nmetric_columns = learning_curves.columns.difference([\"trial\", \"epoch\"])\nlearning_curves[metric_columns] = learning_curves[metric_columns].apply(pd.to_numeric, errors=\"coerce\").replace([np.inf, -np.inf], np.nan)\ndisplay(learning_curves)\n", "id": "trial-metrics", "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "validation_files = [p for p in sorted(automl_dir.glob(\"version_*/val_result.csv\"))\n if p.parent.name.removeprefix(\"version_\").isdigit()]\nvalidation_results = [pd.read_csv(p).assign(trial=p.parent.name) for p in validation_files]\ndisplay(pd.concat(validation_results, ignore_index=True) if validation_results else \"No validation CSVs were recorded.\")\nselected_plan = json.loads((selected_trial / \"model_plan.json\").read_text())\ntune_metric = selected_plan[\"training_job\"][\"tune_metric\"]\nprint(\"Selected trial tuning metric:\", tune_metric)\ndisplay(selected_plan)\nconfig_path = training_result.artifact_paths.get(\"trainer_config\")\nif config_path:\n print(\"Fitted trainer configuration:\")\n display(json.loads(Path(config_path).read_text()))\n", "id": "selected-plan", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Plot learning curves\n\nPlot 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", "id": "curves-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "# Kumo's imports switch matplotlib to the non-interactive Agg backend; restore inline rendering.\n%matplotlib inline\n\nfor 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 plt.show()\n", "id": "learning-curves", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true, "collapsed": false }, "source": "### Read explanations by inspecting gradient scores\n\nRead 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", "id": "explanations-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "scores_path = explanation_dir / \"explain_test/scores.csv\"\nif 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 plt.show()\nelse:\n print(\"No gradient scores. Check task eligibility, disable_explain, and fit logs.\")\n", "id": "gradient-scores", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Compare predictions and observed targets by cohort\n\nSelect 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", "id": "cohorts-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "analysis_path = explanation_dir / \"column_analysis/test/column_analysis.json\"\nrecords = []\nif analysis_path.is_file():\n records = json.loads(analysis_path.read_text()).get(\"cohorts\") or []\nif 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 plt.show()\nelse:\n print(\"No column cohorts. Check task eligibility, disable_explain, and fit logs.\")\n", "id": "column-cohorts", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Locate other diagnostics\n\nThe 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", "id": "other-diagnostics-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "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\")\nsubgraph_paths = sorted((explanation_dir / \"subgraph_collector/test/subgraph_json\").glob(\"*.json\"))\nprint(\"Sampled graph files:\", len(subgraph_paths))\nif subgraph_paths:\n print(\"First graph fields:\", list(json.loads(subgraph_paths[0].read_text())))\n", "id": "other-diagnostics", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "## Read a bounded prediction preview\n\nRead at most ten rows in Parquet batches. Accessing the result frame first would load every row, even if followed by `.head()`.\n", "id": "preview-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "import pandas as pd\nfrom IPython.display import display\nimport pyarrow.parquet as pq\n\ndef 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\ndisplay(preview_parquet(prediction_result.files))\n", "id": "preview-predictions", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "## Optional: preserve diagnostics\n\nThe 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", "id": "archive-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "from uuid import uuid4\nfrom zipfile import ZipFile\n\nARCHIVE_STAGE = f\"@{DATABASE_NAME}.{SCHEMA_NAME}.KUMO_DIAGNOSTICS\"\nsession.sql(f\"CREATE STAGE IF NOT EXISTS {ARCHIVE_STAGE[1:]}\").collect()\narchive_dir = Path(tempfile.mkdtemp(dir=SCRATCH_DIR))\nlocal_zip = Path(shutil.make_archive(str(archive_dir / uuid4().hex), \"zip\", root_dir=automl_dir))\nuploads = session.file.put(f\"file://{local_zip}\", ARCHIVE_STAGE, overwrite=False, auto_compress=False)\nif not uploads or any(item.status != \"UPLOADED\" for item in uploads):\n raise RuntimeError(\"Archive upload did not complete. The local ZIP remains available.\")\narchive_path = f\"{ARCHIVE_STAGE}/{local_zip.name}\"\nprint(\"Record archive path:\", archive_path)\n", "id": "archive-diagnostics", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Retrieve and inspect an archive\n\nUse the recorded archive path to download and list its contents without extracting it. The notebook role needs access to the internal stage.\n", "id": "retrieve-archive-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "download_dir = Path(tempfile.mkdtemp(dir=SCRATCH_DIR))\nsession.file.get(archive_path, str(download_dir))\ndownloaded_zip = download_dir / archive_path.rsplit(\"/\", 1)[-1]\nwith 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", "id": "retrieve-archive", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "## Optional: retain and reload the model\n\nThe complete model bundle preserves fitted state, not source data, graph snapshots, detailed history, or predictions. Keep the recorded model name and version. In a later session, rerun setup through the graph snapshot, predictive query, and task-specific prediction options, and skip the original training, prediction, and record cells. Then set `model_name` and `model_version` to the recorded values and run the reload and scoring cells below. Load only a trusted model bundle.\n", "id": "reload-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "published = training_result.published_model\nif published is None:\n raise ValueError(\"This fit did not publish a model. Retain the complete local bundle instead.\")\nmodel_name, model_version = published.name, published.version\nprint(\"Record model name and version:\", model_name, model_version)\n", "id": "record-model", "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "# In a later session, set model_name and model_version to the recorded values.\nreloaded_trainer = kumo.Trainer.load_from_connector(\n connector, location=kumo.ModelRegistryLocation(model_name, version=model_version),\n)\n", "id": "reload-model", "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "codeCollapsed": true }, "source": "### Score with the reloaded model\n\nRegenerate the scoring population using the restored trainer's saved generation plan. Predictions stay local; the same task-specific options apply.\n", "id": "reload-prediction-intro" }, { "cell_type": "code", "metadata": { "language": "python" }, "source": "reloaded_population = pquery.generate_prediction_table(\n snapshot, prediction_time=PREDICTION_TIME,\n plan=reloaded_trainer.training_table_generation_plan,\n)\nreloaded_result = reloaded_trainer.predict(\n reloaded_population, snapshot=snapshot, **prediction_options,\n)\nprint(\"Reloaded model rows:\", reloaded_result.num_rows)\nprint(\"Local prediction files:\", reloaded_result.files)\n", "id": "reload-predictions", "execution_count": null, "outputs": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12.13" } }, "nbformat": 4, "nbformat_minor": 5 }