{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Jue7eQwunYn"
      },
      "source": [
        "# AdventureWorks customer repeat-order prediction with Kumo in Google Colab\n",
        "\n",
        "This follows the **prepare → connect → define tables/graph → query → train/reload → inspect → predict** workflow. It uses the Linux `FileConnector` and four AdventureWorks tables: `Customer`, `Product`, `SalesOrderHeader`, and `SalesOrderDetail`. There is no Snowflake, Databricks, S3, or AWS-credentials step.\n",
        "\n",
        "The four source files are **headerless tab-separated exports despite their `.csv` extension**. Set the data path below to a location available inside your Colab runtime. If the files are in Google Drive, mount Drive in Colab before running the relevant cells. The preparation step creates separate, labeled CSVs and does not modify the source files.\n",
        "\n",
        "Choose a **Linux x86-64 Colab runtime with Python 3.12 (2026.07) and a GPU**. Install Kumo Engine, restart the runtime, then resume at Configuration. Colab local storage is ephemeral; retain any model and prediction artifacts you need elsewhere before the runtime ends.\n"
      ],
      "id": "1Jue7eQwunYn"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQrmoH-WunYr"
      },
      "source": [
        "## Install Kumo Engine (once per new Colab runtime)\n",
        "\n",
        "Run the installation cell, then restart the Colab runtime and continue at Configuration.\n"
      ],
      "id": "pQrmoH-WunYr"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "85_ucT8DunYs"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "%pip install kumo-engine\n"
      ],
      "id": "85_ucT8DunYs"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d2mygEQunYs"
      },
      "source": [
        "## Configuration\n",
        "\n",
        "Keep `BASE` on the local runtime disk under `/content`, outside mounted Drive. Raw copies, prepared data, and graph/training scratch stay there. `SOURCE_DATA_DIR` can point to Drive; set `ARTIFACT_ROOT` separately to Drive if you want retained copies there, and wait for syncing before ending the runtime.\n"
      ],
      "id": "7d2mygEQunYs"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OuPxmTOAunYs"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "from pathlib import Path\n",
        "\n",
        "MODE = \"train\"  # \"train\" or \"reload\" from a saved local bundle.\n",
        "MODEL_BUNDLE_PATH = None  # For reload, set to Path(\"<path/to/model_bundle>\").\n",
        "SOURCE_DATA_DIR = Path(\"<path/to/adventureworks_csvs>\")\n",
        "BASE = Path(\"/content/kumo-workspace\")\n",
        "RAW_DATA_DIR = BASE / \"raw\"\n",
        "PREPARED_DATA_DIR = BASE / \"prepared\"\n",
        "SCRATCH_DIR = BASE / \"scratch\"\n",
        "ARTIFACT_ROOT = BASE / \"artifacts\"\n",
        "PREDICTION_TIME = datetime(2025, 6, 30)  # After the latest supplied OrderDate.\n",
        "\n",
        "INSPECT_DIAGNOSTICS = True\n",
        "PLOT_LEARNING_CURVES = True\n",
        "ARCHIVE_DIAGNOSTICS = False\n",
        "SAVE_PREDICTIONS = True\n",
        "RUN_INDEX = None  # Select this fit's unique new run when possible.\n",
        "\n",
        "if MODE not in {\"train\", \"reload\"}:\n",
        "    raise ValueError(\"MODE must be 'train' or 'reload'.\")\n",
        "for name, path in {\"SOURCE_DATA_DIR\": SOURCE_DATA_DIR, \"BASE\": BASE}.items():\n",
        "    if \"<\" in str(path):\n",
        "        raise ValueError(f\"Set {name} to a path available inside the Colab runtime.\")\n",
        "\n",
        "runtime_root = Path(\"/content\").resolve()\n",
        "for name, path in {\"BASE\": BASE, \"SCRATCH_DIR\": SCRATCH_DIR}.items():\n",
        "    resolved = path.expanduser().resolve()\n",
        "    if not resolved.is_relative_to(runtime_root):\n",
        "        raise ValueError(f\"{name} must be on the local runtime disk under /content.\")\n",
        "    existing_parent = resolved\n",
        "    while not existing_parent.exists():\n",
        "        existing_parent = existing_parent.parent\n",
        "    if existing_parent.stat().st_dev != runtime_root.stat().st_dev:\n",
        "        raise ValueError(f\"{name} must use local runtime storage, not mounted Drive or another filesystem.\")\n"
      ],
      "id": "OuPxmTOAunYs"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NBJEpiHbunYs"
      },
      "source": [
        "## Initialize the local Kumo runtime"
      ],
      "id": "NBJEpiHbunYs"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VvHqVORPunYs"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "import csv\n",
        "import json\n",
        "import shutil\n",
        "import tempfile\n",
        "from uuid import uuid4\n",
        "\n",
        "import pandas as pd\n",
        "import torch\n",
        "from IPython.display import display\n",
        "\n",
        "import kumo.engine as kumo\n",
        "from kumoapi.typing import Dtype, Stype\n",
        "from kumoapi.graph import Edge\n",
        "\n",
        "scratch_root = SCRATCH_DIR.expanduser().resolve()\n",
        "scratch_root.mkdir(parents=True, exist_ok=True)\n",
        "ARTIFACT_ROOT.mkdir(parents=True, exist_ok=True)\n",
        "with tempfile.TemporaryFile(dir=scratch_root) as probe:\n",
        "    probe.write(b\"kumo scratch check\")\n",
        "    probe.flush()\n",
        "kumo.configure(scratch_dir=str(scratch_root))\n",
        "print(\"Kumo:\", kumo.__version__)\n",
        "print(\"Scratch:\", scratch_root)\n",
        "print(\"Free local disk GiB:\", round(shutil.disk_usage(scratch_root).free / 2**30, 1))\n",
        "if not torch.cuda.is_available():\n",
        "    raise RuntimeError(\"Choose a GPU-backed Colab runtime before materializing the graph.\")\n",
        "for index in range(torch.cuda.device_count()):\n",
        "    print(\"GPU:\", torch.cuda.get_device_name(index))\n",
        "\n",
        "def disk_checkpoint(label):\n",
        "    free_gib = round(shutil.disk_usage(scratch_root).free / 2**30, 2)\n",
        "    print(f\"Disk after {label}: {free_gib} GiB free\")\n",
        "    if free_gib < 5:\n",
        "        raise RuntimeError(\"Less than 5 GiB remains in local Colab storage; stop before the next step.\")\n"
      ],
      "id": "VvHqVORPunYs"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6Xq_H-VunYt"
      },
      "source": [
        "## Verify access to the four AdventureWorks files\n",
        "\n",
        "SHA-256 checks are optional for separately obtained exports. The notebook still checks that the source and local copy have identical bytes; later cells validate the expected table shape.\n"
      ],
      "id": "Z6Xq_H-VunYt"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rh6C7vdVunYt"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "import hashlib\n",
        "import shutil\n",
        "\n",
        "# Optional: set a SHA-256 value for any file whose exact bytes must be pinned.\n",
        "# Leave as None when using a separately obtained AdventureWorks export.\n",
        "EXPECTED_DATA_SHA256 = {\n",
        "    \"Customer.csv\": None,\n",
        "    \"Product.csv\": None,\n",
        "    \"SalesOrderHeader.csv\": None,\n",
        "    \"SalesOrderDetail.csv\": None,\n",
        "}\n",
        "\n",
        "def sha256(path):\n",
        "    digest = hashlib.sha256()\n",
        "    with Path(path).open(\"rb\") as stream:\n",
        "        for chunk in iter(lambda: stream.read(1024 * 1024), b\"\"):\n",
        "            digest.update(chunk)\n",
        "    return digest.hexdigest()\n",
        "\n",
        "if not SOURCE_DATA_DIR.is_dir():\n",
        "    raise FileNotFoundError(f\"Source directory is not visible: {SOURCE_DATA_DIR}\")\n",
        "RAW_DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
        "for name, expected_hash in EXPECTED_DATA_SHA256.items():\n",
        "    source = SOURCE_DATA_DIR / name\n",
        "    if not source.is_file():\n",
        "        raise FileNotFoundError(f\"Missing AdventureWorks source file: {source}\")\n",
        "    source_hash = sha256(source)\n",
        "    if expected_hash is not None and source_hash != expected_hash:\n",
        "        raise RuntimeError(f\"{name} does not match its configured SHA-256 checksum.\")\n",
        "    destination = RAW_DATA_DIR / name\n",
        "    if destination.exists():\n",
        "        if sha256(destination) != source_hash:\n",
        "            raise RuntimeError(f\"Local file differs from source at {destination}; use a clean workspace.\")\n",
        "    elif source.resolve() != destination.resolve():\n",
        "        shutil.copy2(source, destination)\n",
        "        if sha256(destination) != source_hash:\n",
        "            raise RuntimeError(f\"Copied file differs from source at {destination}.\")\n",
        "    print(\"Verified AdventureWorks file:\", destination)\n"
      ],
      "id": "Rh6C7vdVunYt"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L_Wzgv47unYu"
      },
      "source": [
        "## Prepare the four headerless TSV exports as local CSV tables"
      ],
      "id": "L_Wzgv47unYu"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zwmllfsJunYu"
      },
      "source": [
        "The source files stay unchanged. The prepared directory is separate from Kumo scratch; only the columns needed for this AdventureWorks example are copied. Each input row is checked for the expected field count."
      ],
      "id": "zwmllfsJunYu"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9S9McgONunYu"
      },
      "outputs": [],
      "source": [
        "# (expected raw field count, selected column name -> zero-based raw field index)\n",
        "TABLE_LAYOUTS = {\n",
        "    'Customer': (7, {\n",
        "        'CustomerID': 0, 'StoreID': 2, 'TerritoryID': 3,\n",
        "    }),\n",
        "    'Product': (25, {\n",
        "        'ProductID': 0, 'Name': 1, 'Color': 5, 'ProductLine': 15,\n",
        "        'Class': 16, 'StandardCost': 8, 'ListPrice': 9,\n",
        "    }),\n",
        "    'SalesOrderHeader': (26, {\n",
        "        'SalesOrderID': 0, 'CustomerID': 10, 'OrderDate': 2,\n",
        "        'OnlineOrderFlag': 6, 'TerritoryID': 12,\n",
        "        'SubTotal': 19, 'TotalDue': 22,\n",
        "    }),\n",
        "    'SalesOrderDetail': (11, {\n",
        "        'SalesOrderDetailID': 1, 'SalesOrderID': 0, 'ProductID': 4,\n",
        "        'ModifiedDate': 10, 'OrderQty': 3, 'UnitPrice': 6,\n",
        "        'LineTotal': 8,\n",
        "    }),\n",
        "}\n",
        "\n",
        "PREPARED_DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
        "row_counts = {}\n",
        "date_bounds = [None, None]\n",
        "for table_name, (expected_fields, columns) in TABLE_LAYOUTS.items():\n",
        "    raw_path = RAW_DATA_DIR / f'{table_name}.csv'\n",
        "    prepared_path = PREPARED_DATA_DIR / f'{table_name}.csv'\n",
        "    if not raw_path.is_file():\n",
        "        raise FileNotFoundError(f'Missing uploaded table: {raw_path}')\n",
        "    temporary_path = prepared_path.with_suffix('.csv.tmp')\n",
        "    row_count = 0\n",
        "    with raw_path.open(newline='', encoding='utf-8-sig') as raw_file, \\\n",
        "         temporary_path.open('w', newline='', encoding='utf-8') as prepared_file:\n",
        "        reader = csv.reader(raw_file, delimiter='\\t')\n",
        "        writer = csv.writer(prepared_file)\n",
        "        writer.writerow(columns)\n",
        "        for row_count, row in enumerate(reader, start=1):\n",
        "            if len(row) != expected_fields:\n",
        "                raise ValueError(f'{raw_path.name} row {row_count}: expected '\n",
        "                                 f'{expected_fields} fields, found {len(row)}')\n",
        "            selected = [row[index] for index in columns.values()]\n",
        "            if table_name == 'SalesOrderHeader':\n",
        "                order_date = row[2][:10]\n",
        "                date_bounds[0] = min(date_bounds[0], order_date) if date_bounds[0] else order_date\n",
        "                date_bounds[1] = max(date_bounds[1], order_date) if date_bounds[1] else order_date\n",
        "            for date_column in ('OrderDate', 'ModifiedDate'):\n",
        "                if date_column in columns:\n",
        "                    selected[list(columns).index(date_column)] = row[columns[date_column]][:10]\n",
        "            writer.writerow(selected)\n",
        "    if row_count == 0:\n",
        "        raise ValueError(f'{raw_path.name} is empty')\n",
        "    temporary_path.replace(prepared_path)\n",
        "    row_counts[table_name] = row_count\n",
        "print('Prepared rows:', row_counts)\n",
        "print('OrderDate range:', *date_bounds)\n",
        "if PREDICTION_TIME.date().isoformat() <= date_bounds[1]:\n",
        "    raise ValueError('PREDICTION_TIME must be after the latest OrderDate for this scoring demo.')\n"
      ],
      "id": "9S9McgONunYu"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-3tXs3funYu"
      },
      "source": [
        "## Connect to the local files"
      ],
      "id": "y-3tXs3funYu"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HZf0yT2xunYu"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "connector = kumo.FileConnector(root_dir=str(PREPARED_DATA_DIR))\n",
        "print(\"Local tables:\", connector.table_names())\n",
        "display(pd.DataFrame([\n",
        "    {\"column\": column.name, \"physical_type\": column.dtype, \"semantic_type\": column.stype}\n",
        "    for column in connector[\"SalesOrderHeader\"].columns\n",
        "]))\n"
      ],
      "id": "HZf0yT2xunYu"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yg904viyunYu"
      },
      "source": [
        "## Define the four AdventureWorks tables"
      ],
      "id": "yg904viyunYu"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bVqa_qjrunYv"
      },
      "outputs": [],
      "source": [
        "customer = (\n",
        "    kumo.Table.from_source_table(\n",
        "        connector['Customer'],\n",
        "        ['CustomerID', 'StoreID', 'TerritoryID'],\n",
        "        primary_key='CustomerID',\n",
        "    )\n",
        "    .with_column('StoreID', stype=Stype.categorical)\n",
        "    .with_column('TerritoryID', stype=Stype.categorical)\n",
        ")\n",
        "\n",
        "product = (\n",
        "    kumo.Table.from_source_table(\n",
        "        connector['Product'],\n",
        "        ['ProductID', 'Name', 'Color', 'ProductLine', 'Class',\n",
        "         'StandardCost', 'ListPrice'],\n",
        "        primary_key='ProductID',\n",
        "    )\n",
        "    .with_column('Name', stype=Stype.text)\n",
        "    .with_column('Color', stype=Stype.categorical)\n",
        "    .with_column('ProductLine', stype=Stype.categorical)\n",
        "    .with_column('Class', stype=Stype.categorical)\n",
        "    .with_column('StandardCost', stype=Stype.numerical)\n",
        "    .with_column('ListPrice', stype=Stype.numerical)\n",
        ")\n",
        "\n",
        "sales_order_header = (\n",
        "    kumo.Table.from_source_table(\n",
        "        connector['SalesOrderHeader'],\n",
        "        ['SalesOrderID', 'CustomerID', 'OrderDate', 'OnlineOrderFlag',\n",
        "         'TerritoryID', 'SubTotal', 'TotalDue'],\n",
        "        primary_key='SalesOrderID', time_column='OrderDate',\n",
        "    )\n",
        "    .with_column('CustomerID', stype=Stype.ID)\n",
        "    .with_column('OrderDate', dtype=Dtype.date)\n",
        "    .with_column('OnlineOrderFlag', stype=Stype.categorical)\n",
        "    .with_column('TerritoryID', stype=Stype.categorical)\n",
        "    .with_column('SubTotal', stype=Stype.numerical)\n",
        "    .with_column('TotalDue', stype=Stype.numerical)\n",
        ")"
      ],
      "id": "bVqa_qjrunYv"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-gcM47EaunYy"
      },
      "source": [
        "`SalesOrderDetail` has its own `ModifiedDate`, which the original walkthrough used as the line-item time column. Verify this is the right availability time for your intended use before treating model metrics as a benchmark."
      ],
      "id": "-gcM47EaunYy"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txoVERmwunYy"
      },
      "outputs": [],
      "source": [
        "sales_order_detail = (\n",
        "    kumo.Table.from_source_table(\n",
        "        connector['SalesOrderDetail'],\n",
        "        ['SalesOrderDetailID', 'SalesOrderID', 'ProductID', 'ModifiedDate',\n",
        "         'OrderQty', 'UnitPrice', 'LineTotal'],\n",
        "        primary_key='SalesOrderDetailID', time_column='ModifiedDate',\n",
        "    )\n",
        "    .with_column('SalesOrderID', stype=Stype.ID)\n",
        "    .with_column('ProductID', stype=Stype.ID)\n",
        "    .with_column('ModifiedDate', dtype=Dtype.date)\n",
        "    .with_column('OrderQty', stype=Stype.numerical)\n",
        "    .with_column('UnitPrice', stype=Stype.numerical)\n",
        "    .with_column('LineTotal', stype=Stype.numerical)\n",
        ")"
      ],
      "id": "txoVERmwunYy"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tbfJTsr5unYy"
      },
      "source": [
        "## Validate the graph and materialize a snapshot"
      ],
      "id": "tbfJTsr5unYy"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8LN3NsoWunYy"
      },
      "outputs": [],
      "source": [
        "graph = kumo.Graph(\n",
        "    tables=[customer, product, sales_order_header, sales_order_detail],\n",
        "    edges=[\n",
        "        Edge(src_table='SalesOrderHeader', fkey='CustomerID', dst_table='Customer'),\n",
        "        Edge(src_table='SalesOrderDetail', fkey='SalesOrderID', dst_table='SalesOrderHeader'),\n",
        "        Edge(src_table='SalesOrderDetail', fkey='ProductID', dst_table='Product'),\n",
        "    ],\n",
        ")\n",
        "graph.validate()\n",
        "print('Graph ID:', graph.id)\n",
        "print('Tables:', list(graph._named_tables))\n",
        "print('Edges:', len(graph.edges))\n",
        "disk_checkpoint('graph validation')"
      ],
      "id": "8LN3NsoWunYy"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "veHi9fB-unYz"
      },
      "outputs": [],
      "source": [
        "snapshot = graph.snapshot()\n",
        "print(\"Snapshot:\", snapshot.id)\n",
        "print(\"Local index:\", snapshot.index_path)\n",
        "disk_checkpoint(\"snapshot\")"
      ],
      "id": "veHi9fB-unYz"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9QrMIakTunYz"
      },
      "source": [
        "## Define the predictive query"
      ],
      "id": "9QrMIakTunYz"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MwTzyT7eunYz"
      },
      "source": [
        "## Predict whether an active customer orders in the next 90 days\n",
        "\n",
        "This temporal binary-classification query is the same task used by the original AdventureWorks walkthrough. The population is customers who ordered in the preceding 365 days."
      ],
      "id": "MwTzyT7eunYz"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hofNDrNxunYz"
      },
      "outputs": [],
      "source": [
        "pquery = kumo.PredictiveQuery(\n",
        "    graph=graph,\n",
        "    query=(\n",
        "        'PREDICT COUNT(SalesOrderHeader.*, 0, 90, days) > 0 '\n",
        "        'FOR EACH Customer.CustomerID '\n",
        "        'WHERE COUNT(SalesOrderHeader.*, -365, 0, days) > 0'\n",
        "    ),\n",
        ")\n",
        "pquery.validate()\n",
        "print('Task type:', pquery.get_task_type())\n",
        "print('Prediction time:', PREDICTION_TIME)"
      ],
      "id": "hofNDrNxunYz"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eArGwggYunYz"
      },
      "source": [
        "## Set up training or reload a retained model"
      ],
      "id": "eArGwggYunYz"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W5FO4WZOunYz"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "if MODE == \"train\":\n",
        "    model_plan = pquery.suggest_model_plan(snapshot)\n",
        "    model_plan.optimization.max_epochs = 5\n",
        "    print(model_plan)\n",
        "    print(\"Validation selection metric:\", getattr(\n",
        "        model_plan.training_job, \"tune_metric\", \"not exposed by this build\"))\n",
        "else:\n",
        "    if not MODEL_BUNDLE_PATH or not Path(MODEL_BUNDLE_PATH).is_dir():\n",
        "        raise ValueError(\"Set MODEL_BUNDLE_PATH to a retained local model bundle directory.\")\n",
        "    trainer = kumo.Trainer.load(Path(MODEL_BUNDLE_PATH))\n",
        "    print(\"Loaded:\", MODEL_BUNDLE_PATH)\n"
      ],
      "id": "W5FO4WZOunYz"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1XAIqhNunYz"
      },
      "source": [
        "### Generate historical examples"
      ],
      "id": "d1XAIqhNunYz"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "txcLlSjGunY0"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "train_table = None\n",
        "if MODE == \"train\":\n",
        "    train_table = pquery.generate_training_table(snapshot)\n",
        "    print(\"Training table:\", train_table.train_table_path)\n",
        "    disk_checkpoint(\"training-table generation\")\n",
        "else:\n",
        "    print(\"Reload route: training-table generation skipped.\")\n"
      ],
      "id": "txcLlSjGunY0"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YuGn0SQAunY0"
      },
      "source": [
        "### Fit and retain the model bundle"
      ],
      "id": "YuGn0SQAunY0"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rR3xsyREunY0"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "runs_before_fit = {path.parent for path in scratch_root.rglob(\"auto_train_progress.json\")}\n",
        "training_result = None\n",
        "retained_bundle = Path(MODEL_BUNDLE_PATH) if MODE == \"reload\" else None\n",
        "if MODE == \"train\":\n",
        "    trainer = kumo.Trainer(model_plan)\n",
        "    training_result = trainer.fit(train_table, snapshot=snapshot)\n",
        "    disk_checkpoint(\"model fit\")\n",
        "    retained_bundle = ARTIFACT_ROOT / f\"model-{uuid4().hex}\"\n",
        "    shutil.copytree(Path(training_result.bundle_dir), retained_bundle)\n",
        "    print(\"Retained local bundle:\", retained_bundle)\n",
        "else:\n",
        "    print(\"Reload route: fit skipped.\")\n"
      ],
      "id": "rR3xsyREunY0"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zcc_cYhuunY0"
      },
      "source": [
        "### Inspect model metrics and artifacts"
      ],
      "id": "Zcc_cYhuunY0"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BR-E157NunY1"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "if training_result is None:\n",
        "    print(\"Reloaded from:\", retained_bundle)\n",
        "    print(\"Saved training-table plan:\", trainer.training_table_generation_plan)\n",
        "else:\n",
        "    print(\"Available metrics:\", training_result.metrics)\n",
        "    print(\"Reusable local bundle:\", retained_bundle)\n",
        "    for name, path in training_result.artifact_paths.items():\n",
        "        print(name, path)\n",
        "print(\"Download the retained bundle before Colab recycles the runtime.\")\n"
      ],
      "id": "BR-E157NunY1"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JSOlu2V0unY1"
      },
      "source": [
        "## Optional diagnostics\n",
        "\n",
        "Diagnostic inspection and learning-curve plotting are enabled by default. Set `INSPECT_DIAGNOSTICS` or `PLOT_LEARNING_CURVES` to `False` to disable them. Diagnostic archiving is disabled unless `ARCHIVE_DIAGNOSTICS = True`. Training metrics and the model bundle are retained independently."
      ],
      "id": "JSOlu2V0unY1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZL10RW2xunY1"
      },
      "outputs": [],
      "source": [
        "automl_dirs = sorted({\n",
        "    path.parent for path in scratch_root.rglob(\"auto_train_progress.json\")\n",
        "    if path.is_file()\n",
        "})\n",
        "for index, path in enumerate(automl_dirs):\n",
        "    print(index, path)\n",
        "if not automl_dirs:\n",
        "    print(\"No retained AutoML runs found under configured scratch.\")"
      ],
      "id": "ZL10RW2xunY1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2zB368munY1"
      },
      "outputs": [],
      "source": [
        "automl_dir = None\n",
        "if INSPECT_DIAGNOSTICS:\n",
        "    if RUN_INDEX is not None:\n",
        "        if not isinstance(RUN_INDEX, int) or not 0 <= RUN_INDEX < len(automl_dirs):\n",
        "            raise ValueError(\"Choose one of the displayed run indices.\")\n",
        "        automl_dir = automl_dirs[RUN_INDEX]\n",
        "    elif \"runs_before_fit\" in globals():\n",
        "        new_runs = [path for path in automl_dirs if path not in runs_before_fit]\n",
        "        if len(new_runs) == 1:\n",
        "            automl_dir = new_runs[0]\n",
        "    if automl_dir is None:\n",
        "        print(\"No unambiguous run selected. Set RUN_INDEX from the displayed list.\")\n",
        "    else:\n",
        "        print(\"Selected:\", automl_dir)\n",
        "else:\n",
        "    print(\"Diagnostic inspection disabled.\")"
      ],
      "id": "r2zB368munY1"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UycqpZjhunY1"
      },
      "source": [
        "### Read progress and compare trials\n",
        "\n",
        "These small readers report missing or unreadable files without blocking prediction. The progress summary shows completed/total trials and each trial's recorded epochs, elapsed seconds, and accelerator."
      ],
      "id": "UycqpZjhunY1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "27qGorgounY1"
      },
      "outputs": [],
      "source": [
        "def read_json(path):\n",
        "    try:\n",
        "        return json.loads(path.read_text(encoding=\"utf-8\"))\n",
        "    except (OSError, ValueError) as error:\n",
        "        print(f\"Cannot read {path.name}: {error}\")\n",
        "        return {}\n",
        "\n",
        "def read_csv(path):\n",
        "    try:\n",
        "        return pd.read_csv(path)\n",
        "    except (OSError, pd.errors.EmptyDataError, pd.errors.ParserError) as error:\n",
        "        print(f\"Cannot read {path.name}: {error}\")\n",
        "        return pd.DataFrame()\n",
        "\n",
        "def history_frames(progress):\n",
        "    summaries, curves = [], []\n",
        "    for trial, details in (progress.get(\"trial_progress\") or {}).items():\n",
        "        details = details or {}\n",
        "        summaries.append({\"trial\": str(trial),\n",
        "                          \"epochs_completed\": details.get(\"num_epochs_completed\"),\n",
        "                          \"elapsed_seconds\": details.get(\"elapsed_time\"),\n",
        "                          \"accelerator\": details.get(\"accelerator\")})\n",
        "        for epoch, metrics in (details.get(\"metrics\") or {}).items():\n",
        "            try:\n",
        "                epoch_number = int(epoch)\n",
        "            except (TypeError, ValueError):\n",
        "                print(\"Skipping unrecognized epoch:\", epoch)\n",
        "                continue\n",
        "            metrics = metrics or {}\n",
        "            curves.append({\"trial\": str(trial), \"epoch\": epoch_number,\n",
        "                           **(metrics.get(\"train_metrics\") or {}),\n",
        "                           **(metrics.get(\"validation_metrics\") or {})})\n",
        "    summary = pd.DataFrame(summaries)\n",
        "    history = pd.DataFrame(curves)\n",
        "    if not history.empty:\n",
        "        history = history.sort_values([\"trial\", \"epoch\"])\n",
        "    return summary, history"
      ],
      "id": "27qGorgounY1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S_ojurwvunY1"
      },
      "outputs": [],
      "source": [
        "progress = read_json(automl_dir / \"auto_train_progress.json\") if automl_dir else {}\n",
        "trial_summary, learning_curves = history_frames(progress)\n",
        "print(\"Completed / total trials:\", progress.get(\"completed_trials\", \"not recorded\"),\n",
        "      \"/\", progress.get(\"total_trials\", \"not recorded\"))\n",
        "if not trial_summary.empty:\n",
        "    display(trial_summary)\n",
        "else:\n",
        "    print(\"No trial progress recorded or no run selected.\")\n",
        "if not learning_curves.empty:\n",
        "    display(learning_curves)\n",
        "print(\"Available epoch metrics:\", [\n",
        "    c for c in learning_curves.columns if c not in {\"trial\", \"epoch\"}\n",
        "])"
      ],
      "id": "S_ojurwvunY1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hF2O5WJVunY2"
      },
      "outputs": [],
      "source": [
        "trial_frames = []\n",
        "for path in sorted(automl_dir.glob(\"version_*/val_result.csv\")) if automl_dir else []:\n",
        "    frame = read_csv(path)\n",
        "    if not frame.empty:\n",
        "        frame[\"trial_directory\"] = path.parent.name\n",
        "        trial_frames.append(frame)\n",
        "if trial_frames:\n",
        "    display(pd.concat(trial_frames, ignore_index=True))\n",
        "else:\n",
        "    print(\"No completed trial validation results found.\")\n"
      ],
      "id": "hF2O5WJVunY2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OwNu6di7unY2"
      },
      "outputs": [],
      "source": [
        "selected_trial = explanation_dir = None\n",
        "selection = read_json(automl_dir / \"best_trainer.json\") if automl_dir else {}\n",
        "if selection.get(\"best_trainer\"):\n",
        "    selected_trial = automl_dir / Path(selection[\"best_trainer\"]).name\n",
        "    if selected_trial.is_dir():\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",
        "        display(read_json(selected_trial / \"model_plan.json\"))\n",
        "    else:\n",
        "        print(\"Selected trial directory is unavailable:\", selected_trial)\n",
        "else:\n",
        "    print(\"No selected trial recorded or no run selected.\")\n"
      ],
      "id": "OwNu6di7unY2"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qw8YTreWunY2"
      },
      "source": [
        "### Plot available learning curves\n",
        "\n",
        "Select exact metric names printed above. The default tries training and validation loss. Compare how validation changes while training improves. Divergence can suggest overfitting."
      ],
      "id": "Qw8YTreWunY2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iz20yKJounY2"
      },
      "outputs": [],
      "source": [
        "METRIC_COLUMNS = [\"train_loss\", \"val_loss\"]\n",
        "columns = [name for name in METRIC_COLUMNS if name in learning_curves.columns]\n",
        "if not PLOT_LEARNING_CURVES:\n",
        "    print(\"Learning-curve plotting disabled.\")\n",
        "elif learning_curves.empty or not columns:\n",
        "    print(\"No matching completed epoch metrics. Choose from the available names.\")\n",
        "else:\n",
        "    try:\n",
        "        import matplotlib.pyplot as plt\n",
        "    except ImportError:\n",
        "        print(\"Matplotlib is unavailable. Use the displayed epoch table instead.\")\n",
        "    else:\n",
        "        figure, axis = plt.subplots(figsize=(9, 5))\n",
        "        for trial, rows in learning_curves.groupby(\"trial\", sort=True):\n",
        "            for column in columns:\n",
        "                values = pd.to_numeric(rows[column], errors=\"coerce\")\n",
        "                if values.notna().any():\n",
        "                    axis.plot(rows[\"epoch\"], values, marker=\"o\", label=f\"{trial}: {column}\")\n",
        "        axis.set(xlabel=\"Epoch\", ylabel=\"Metric\", title=\"Learning curves\")\n",
        "        axis.grid(alpha=0.25)\n",
        "        if axis.lines:\n",
        "            axis.legend()\n",
        "        plt.show()"
      ],
      "id": "iz20yKJounY2"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xdpc3tbxunY2"
      },
      "source": [
        "### Inspect gradient scores and column analysis\n",
        "\n",
        "Gradient scores describe model sensitivity to input columns, not causal effects. Column analysis summarizes cohort-level patterns. Availability depends on task, model, and build. Inspect the selected trial, using its matching training/validation refit when present."
      ],
      "id": "Xdpc3tbxunY2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lvSNfoxUunY2"
      },
      "outputs": [],
      "source": [
        "score_path = explanation_dir / \"explain_test/scores.csv\" if explanation_dir else None\n",
        "score_paths = [score_path] if score_path and score_path.is_file() else []\n",
        "if not score_paths:\n",
        "    print(\"No gradient scores. Check selected run, build, explanation setting, and logs.\")\n",
        "for path in score_paths:\n",
        "    print(\"Scores:\", path.relative_to(automl_dir))\n",
        "    scores = read_csv(path)\n",
        "    if \"score\" in scores.columns:\n",
        "        scores = scores.sort_values(\"score\", ascending=False)\n",
        "    display(scores)"
      ],
      "id": "lvSNfoxUunY2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VAcmyeLKunY2"
      },
      "outputs": [],
      "source": [
        "analysis_path = explanation_dir / \"column_analysis/test/column_analysis.json\" if explanation_dir else None\n",
        "analysis_paths = [analysis_path] if analysis_path and analysis_path.is_file() else []\n",
        "if not analysis_paths:\n",
        "    print(\"No column analysis found for the selected run.\")\n",
        "for path in analysis_paths:\n",
        "    print(\"Column analysis:\", path.relative_to(automl_dir))\n",
        "    analysis = read_json(path)\n",
        "    display(pd.json_normalize(analysis.get(\"cohorts\") or []))"
      ],
      "id": "VAcmyeLKunY2"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wK8KGb72unY2"
      },
      "source": [
        "### Locate other diagnostics\n",
        "\n",
        "When present, `best_trainer.json` identifies the selected trial; `train_progress.json` records individual trial progress; TensorBoard event files support deeper curve inspection."
      ],
      "id": "wK8KGb72unY2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_JAXozqfunY2"
      },
      "outputs": [],
      "source": [
        "if automl_dir:\n",
        "    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",
        "        for path in paths[:20]:\n",
        "            print(\" \", path.relative_to(automl_dir))"
      ],
      "id": "_JAXozqfunY2"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r9OAtMbKunY2"
      },
      "source": [
        "## Optional: archive diagnostics locally"
      ],
      "id": "r9OAtMbKunY2"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m5T6gAp0unY3"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "archived_diagnostics = None\n",
        "if ARCHIVE_DIAGNOSTICS:\n",
        "    if automl_dir is None:\n",
        "        raise ValueError(\"Select a diagnostic run first.\")\n",
        "    archived_diagnostics = Path(shutil.make_archive(\n",
        "        str(ARTIFACT_ROOT / f\"diagnostics-{uuid4().hex}\"), \"zip\", root_dir=automl_dir))\n",
        "    print(\"Local archive:\", archived_diagnostics)\n",
        "else:\n",
        "    print(\"Diagnostic archive skipped.\")\n"
      ],
      "id": "m5T6gAp0unY3"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yk9ZnYQtunY3"
      },
      "source": [
        "## Generate predictions"
      ],
      "id": "Yk9ZnYQtunY3"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQzCYy5sunY3"
      },
      "source": [
        "This is future scoring: the files have no observed labels after the configured boundary, so the scores are not an accuracy evaluation. The population is customers with an order during the preceding 365 days."
      ],
      "id": "bQzCYy5sunY3"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l5WGij-1unY6"
      },
      "outputs": [],
      "source": [
        "print('OrderDate range:', *date_bounds)\n",
        "print('Configured prediction boundary:', PREDICTION_TIME)\n",
        "print('Population: customers with an order in the preceding 365 days')"
      ],
      "id": "l5WGij-1unY6"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hlC0bx0qunY6"
      },
      "source": [
        "### Generate the scoring population\n",
        "\n",
        "For this temporal query, pass a timezone-naive UTC timestamp and reuse the fitted trainer's saved training-table generation plan."
      ],
      "id": "hlC0bx0qunY6"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TIBKB4V-unY6"
      },
      "outputs": [],
      "source": [
        "prediction_table = pquery.generate_prediction_table(\n",
        "    snapshot,\n",
        "    prediction_time=PREDICTION_TIME,\n",
        "    plan=trainer.training_table_generation_plan,\n",
        ")\n",
        "print('Scoring population rows:', prediction_table.row_count)\n",
        "disk_checkpoint('prediction-table generation')"
      ],
      "id": "TIBKB4V-unY6"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pjN7KhcrunY6"
      },
      "source": [
        "### Score locally"
      ],
      "id": "pjN7KhcrunY6"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qZNcfOZIunY7"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "prediction_result = trainer.predict(\n",
        "    prediction_table, snapshot=snapshot,\n",
        "    binary_classification_threshold=0.5, output_config=None,\n",
        ")\n",
        "print(\"Rows:\", prediction_result.num_rows)\n",
        "print(\"Local Parquet files:\", prediction_result.files)\n",
        "disk_checkpoint(\"prediction\")\n"
      ],
      "id": "qZNcfOZIunY7"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFUhQab_unY7"
      },
      "source": [
        "### Read a bounded preview\n",
        "\n",
        "Read at most ten rows across the local Parquet files in batches. This avoids collecting the entire result."
      ],
      "id": "QFUhQab_unY7"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ECPFeB-WunY7"
      },
      "outputs": [],
      "source": [
        "import pyarrow.parquet as pq\n",
        "\n",
        "def preview_parquet(files, limit=10):\n",
        "    frames = []\n",
        "    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(prediction_result.files))"
      ],
      "id": "ECPFeB-WunY7"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vfvFov05unY7"
      },
      "execution_count": null,
      "outputs": [],
      "source": [
        "if SAVE_PREDICTIONS:\n",
        "    prediction_dir = ARTIFACT_ROOT / f\"predictions-{uuid4().hex}\"\n",
        "    prediction_dir.mkdir(parents=True)\n",
        "    saved_files = []\n",
        "    for index, source in enumerate(prediction_result.files):\n",
        "        target = prediction_dir / f\"part-{index}.parquet\"\n",
        "        shutil.copy2(source, target)\n",
        "        saved_files.append(target)\n",
        "    manifest = {\n",
        "        \"kumo_engine_version\": kumo.__version__,\n",
        "        \"prediction_time\": PREDICTION_TIME.isoformat(),\n",
        "        \"prediction_rows\": prediction_result.num_rows,\n",
        "        \"model_bundle\": str(retained_bundle),\n",
        "        \"metrics\": training_result.metrics if training_result else None,\n",
        "        \"prediction_files\": [str(path) for path in saved_files],\n",
        "    }\n",
        "    (prediction_dir / \"run.json\").write_text(json.dumps(manifest, indent=2, default=str))\n",
        "    print(\"Saved prediction directory:\", prediction_dir)\n",
        "    print(\"Download predictions and the model bundle before the Colab runtime ends.\")\n",
        "else:\n",
        "    print(\"SDK prediction Parquet remains in local scratch for this runtime.\")\n"
      ],
      "id": "vfvFov05unY7"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F38G5JTJunY7"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "For another scoring run, retain/download the **complete** model bundle, place it back on the Colab filesystem, set `MODE=\"reload\"` and `MODEL_BUNDLE_PATH`, and rerun the source, graph, snapshot, and prediction sections. The source CSVs, prepared CSVs, snapshot, model bundle, diagnostics, and predictions are separate artifacts. Changing source columns, graph edges, query horizon, or training split requires renewed validation and may require retraining.\n"
      ],
      "id": "F38G5JTJunY7"
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
