
# I Ran Codex on a Local Model — No Cloud, No Bill

I use Codex as the headless engine that distributes the orchestration tasks in a self-verifying system I'm building. Each node in a DAG hands Codex a job and it runs the whole agentic loop — writes files, runs shell, edits code, checks its own work. And most of the time the job is something trivial. Rename a thing. Scaffold a file. Summarize a diff. Judge whether two outputs satisfy a criterion. None of that needs a frontier model in a datacenter, but every single node still goes out over the wire and still shows up on the bill.

That bothered me more than the money did, honestly. The harness — the part that plans, calls tools, and verifies — is identical whether the brain answering each turn lives in the cloud or on my desk. The model is just a function the harness calls. So why not swap in a function that runs on my own machine, for free, offline, with nothing leaving the laptop?

That's the whole idea: **keep Codex's harness, swap the brain for a local model, and route the simpler tasks to it.** Hard, critical, agentic work stays on cloud Codex where it belongs. Everything cheap and forgiving runs locally for nothing.

It sounds like a config change. It was not a config change. Here's everything that broke, and the one thing that finally worked.

## The brain: a Qwen MoE on Apple Silicon

Before any of the Codex wiring, you need a local model that's actually fast enough to sit inside an agent loop, where the model gets called many times per task and latency compounds.

On a Mac the runtime decision isn't close. [MLX](https://github.com/ml-explore/mlx) is Apple's Metal-native array framework, and it drives the unified-memory GPU directly, so for a given quantization it's the fastest option on Apple Silicon. [vLLM](https://github.com/vllm-project/vllm) is wonderful on an NVIDIA box but has no real Metal backend, so it's the wrong tool here. Ollama and LM Studio's GGUF path both work, but under the hood that's llama.cpp, and MLX edges it out at the same quant.

For the model itself, the lever that matters is *active* parameters, not total. A Mixture-of-Experts model only fires a slice of its weights on each token, so a 35B model with 3B active actually generates *faster* than a 27B dense model that has to push all 27B every token — while carrying more total capacity for quality. I settled on [`unsloth/Qwen3.6-35B-A3B`](https://huggingface.co/unsloth) in MLX 4-bit: about 21.6 GB on disk, 3B active per token, the strongest and the fastest thing I tried.

Two download gotchas cost me an hour before the model ever loaded, and they're worth saving you. Unauthenticated Hugging Face pulls are throttled to a crawl, so set a free token. And Unsloth's repos use Xet storage, whose client kept freezing at 0 MB/s with a pile of half-downloaded chunks; turning Xet off and falling back to plain HTTPS jumped it straight to ~10 MB/s.

```bash
export HF_TOKEN=hf_xxx                 # otherwise the download crawls
export HF_HUB_DISABLE_XET=1            # the real fix — Xet stalls, plain HTTPS doesn't
export HF_HUB_ENABLE_HF_TRANSFER=1
uv tool install mlx-lm --with hf-transfer
```

With the weights down, you've got a model. Now comes the part that took the afternoon.

## Where it all went wrong: Codex speaks an API your local server doesn't

My first instinct was the obvious one. MLX ships a server — `mlx_lm.server` — that exposes an OpenAI-compatible endpoint. Codex lets you define a custom model provider pointed at any OpenAI-compatible URL. Point one at the other, done.

Except every agent turn came back `404`, and Codex just retried, forever. I sat and watched it 404-loop for forty minutes before I dug into why.

The reason is subtle and it's the single most important thing in this post: **it's not about the tool schema, it's about which HTTP API Codex actually speaks.** Modern Codex (0.139 in my case) drives a custom model provider over the OpenAI **Responses API** — it `POST`s to `/v1/responses`. But `mlx_lm.server` only implements the older **Chat Completions API**, `/v1/chat/completions`. Those are two different endpoints. Codex knocks on a door that the MLX server simply doesn't have, gets a 404, assumes a transient failure, and tries again. And again.

There used to be an escape hatch — `wire_api = "chat"` in the config, which forced Codex back onto chat-completions. I went looking for it and found it had been removed in 0.139, so setting it just throws `wire_api = "chat" is no longer supported` and refuses to start. The old advice the internet gives you is now a hard error.

Two more failures showed up while I was flailing, and they're the kind of thing that sends you down the wrong debugging path entirely. One was a reconnect loop — `rmcp ... Auth(AuthorizationRequired)`, fired on every single run — that turned out to be a stale remote MCP server still listed in my Codex config, trying and failing to authenticate before any real work could start. The other was an older LM Studio build flatly rejecting Codex's tool schema with `tools.N.type invalid`. Neither has anything to do with the model; both look, at a glance, like the model is broken.

So the honest conclusion was: a bare MLX server cannot run Codex's agent loop, full stop. The mismatch is structural, not a missing flag.

## The fix that actually worked: LM Studio plus Codex's `--oss` path

The breakthrough was realizing I'd been trying to make the wrong server speak the right protocol. Instead of fighting `mlx_lm.server` into the Responses API, I served the same MLX model through **LM Studio** and used Codex's *built-in* local-provider path — `--oss --local-provider lmstudio` — which is designed to speak chat-completions and handle tool calls correctly.

And it worked, end to end. Codex's agent harness, running entirely on the local Qwen3.6 MLX, created a file on disk — about 20k tokens of agent loop, not one of them leaving the machine.

Four things all had to be true at once, and each one had failed independently before they lined up:

The MLX model has to be served *by LM Studio*, not converted to GGUF. The nice part is you don't re-download anything — the Hugging Face MLX snapshot can be symlinked straight into `~/.lmstudio/models/<org>/<repo>` and LM Studio picks it up as-is.

It has to load with a *big* context window. At 8192 tokens, Codex's own agent system prompt overflows before your task even starts, and you get the cryptic `tokens to keep is greater than the context length`. Loading at 32768 fixes it: `lms load <id> --gpu max --context-length 32768`.

Any auth-failing MCP server has to come out of `~/.codex/config.toml` first, or that `rmcp Auth(AuthorizationRequired)` reconnect loop will stall every run before the model gets a turn.

And LM Studio itself has to be current — the older builds are the ones that reject Codex's tool schema.

Once those held, the smoke test is genuinely satisfying to run:

```bash
lms load qwen3.6-35b-a3b-ud-mlx --gpu max --context-length 32768 && lms server start
codex exec --oss --local-provider lmstudio -m qwen3.6-35b-a3b-ud-mlx \
  -C /tmp/t --sandbox workspace-write --skip-git-repo-check -c 'mcp_servers={}' \
  "Create hello.txt containing: hi. Then stop."   # → writes the file, zero cloud
```

There's one model-specific wrinkle worth flagging, because it'll make you think the model failed when it didn't. Qwen3.6 is a *thinking* model: it writes its reasoning into a separate channel and only emits the real answer after it closes `</think>`. If you starve it on `max_tokens`, it spends the entire budget thinking and hands back empty content. Give it room to finish the thought.

## Why this is worth the trouble: route the simple stuff for free

With local Codex working, the payoff is a routing decision rather than an all-or-nothing one. You don't have to choose between "always cloud" and "always local." You split by what the task actually needs.

The hard, high-stakes, genuinely agentic work — the multi-step changes where a wrong move is expensive — stays on cloud Codex, because that's where the capability gap is real and worth paying for. But the long tail of simpler tasks — generate a file, rename across a module, summarize a run, draft a first pass, judge whether an output meets a bar — can run on the local model for nothing, offline, with no data leaving the machine. The harness is identical. Only the brain, and the bill, changes.

The trick to doing this responsibly is to *measure* which tasks survive the swap rather than guessing. In my setup that's an eval pass: run a node on both engines, compare, and only graduate it to local once it demonstrably matches cloud quality on that task. Cost routing without measurement is just hoping; with a cheap eval in front of it, it's a decision you can defend.

## The fun add-on: a DSPy loop where generation is free

Here's where it gets genuinely useful beyond saving a few dollars on housekeeping tasks. Once you have a local model that the harness can drive for free, you can build a [DSPy](https://github.com/stanfordnlp/dspy)-style prompt-optimization loop and run the expensive half of it at no cost.

DSPy optimization is, at heart, a loop: *propose* candidate prompts, *run* them, *score* which is better, *promote* the winner. The quiet insight is that the cost and the risk live in different halves. The expensive half is proposing and running — you do it hundreds of times per round — but it's *forgiving*, because you generate several candidates and just keep the best, so any single bad generation is harmless. The cheap half is scoring; you do it only a handful of times, but it's *unforgiving*, because a biased scorer quietly optimizes you toward the wrong prompt and every later round inherits the mistake.

So you put each half where it belongs. Generation runs on the local model for free. Scoring runs as a Codex LLM-judge — billed, but tiny, and pointed at your most consistent model because it's the one decision that compounds.

That judge needs one piece of care to be trustworthy: **debias it for position.** LLM judges have a real, measurable preference for whichever output they read first. So you judge twice, swapping the order between rounds, and only trust a verdict that survives the swap:

```js
// Round 1: A first.  Round 2: B first (positions swapped, then remapped).
if (round1Winner === round2Winner && round1Winner !== "tie") {
  verdict = round1Winner;   // it agreed with itself under reordering — trust it
} else {
  verdict = "tie";          // it flipped when you flipped the order — that was bias
}
```

If reordering the two outputs changes the answer, the "preference" was an artifact of position, not quality, and collapsing it to a tie throws that noise away. It's ten lines, and it's the difference between a metric and a coin flip. I also tell the judge to grade depth and specificity rather than the mere presence of the required pieces, and to default to a fail unless the work is genuinely good — a lenient judge will wave box-ticking work straight through and your optimizer will happily learn to produce more of it.

The last bit of glue is unglamorous but it's what keeps the loop alive: structured output. Every candidate result and every verdict has to parse into a record, and a thinking model *will* occasionally wrap its JSON in prose or a markdown fence. The defense is to ask for JSON mode, but never trust it — run every response through a small salvage extractor that finds the first balanced object and validates it before accepting, and treat anything that fails as a dropped case rather than a crash:

```go
func extractJSON(text string) json.RawMessage {
    t := strings.TrimSpace(text)
    if json.Valid([]byte(t)) {
        return json.RawMessage(t)
    }
    first := strings.IndexAny(t, "[{")
    last := strings.LastIndexAny(t, "]}")
    if first >= 0 && last > first {
        if cand := t[first : last+1]; json.Valid([]byte(cand)) {
            return json.RawMessage(cand)
        }
    }
    return nil // a failed parse is a skipped case, not an exception
}
```

The evalset and the history of promotions are stored the same defensive way — newline-delimited JSON, appended as the loop runs, read back line by line so one malformed line drops itself instead of poisoning the whole file. It's not elegant. It's what survives a model that hiccups once an hour.

## What I'd tell you if you were about to try this

Most of the difficulty here isn't where you'd expect. The model runs fine; the optimization loop is conceptually simple; the actual quicksand is the seam between Codex and a local server. So if you take three things from this, take these.

The make-or-break is *which API Codex speaks*, not the tool schema or the model. If you're seeing endless 404s, you're almost certainly serving Chat Completions to a Codex that wants the Responses API — and the old `wire_api = "chat"` workaround is gone. Serving the MLX model through LM Studio and using Codex's own `--oss --local-provider lmstudio` path sidesteps the whole problem, which is why that's the configuration I landed on.

Treat local versus cloud as a routing decision you *measure*, not a flag you flip. Send the simple, forgiving, private work to the free local engine and keep the hard, high-stakes work on cloud Codex — and let an eval, not a hunch, decide which task is which.

And if you build the DSPy loop on top, spend your reliability budget on the judge. Generation is cheap and forgiving and belongs wherever it's free. Judgment is the decision that compounds, so make it your strongest, most consistent model, and always debias it for order. That's the whole game: free where you can afford to be wrong, careful where you can't.


I Ran Codex on a Local Model — No Cloud, No Bill


# Why I Compiled My Time Series Forecasting Engine to WebAssembly

If you build tools for enterprises, you quickly learn a painful truth: **handling raw data can be messy.** 

The moment you ask a financial analyst, a logistics manager, or an operations lead to upload their historical sales or inventory CSV to your servers, you trigger a compliance gauntlet. InfoSec questionnaires, data processing agreements, GDPR/CCPA audits, and the constant fear of data leaks. 

But what if you didn't have to touch their raw data at all?

When I built [WaySightAI](https://github.com/xdanny/WaysightAI), I set a strict constraint: **raw time-series data must never leave the user's browser.** The backend should only handle authentication, saved metadata (project names, configurations), usage limits, and high-level AI orchestration. The actual heavy lifting—parsing the CSV, cleaning the data, running stationarity tests, and computing forecasts—had to happen locally.

To do this, I wrote my core time-series math in Rust and compiled it to WebAssembly (WASM). Here is how I structured the hybrid architecture and what I learned along the way.

---

## The Monorepo Architecture

WaySightAI is organized as a monorepo to maintain a clear boundary between the UI, the mathematical engine, and the management layer:

```text
packages/
  web/        React 18 + TypeScript + Vite (renders the workflow UI)
  wasm/       Rust forecasting engine compiled with wasm-pack
  api/        FastAPI backend (SQLAlchemy, Clerk authentication, Redis, Stripe hooks)
```

The forecasting path follows a strict client-side flow:
1. The user uploads a CSV in the React web application.
2. The browser parses the CSV locally and lets the user select date and target columns.
3. The React app loads the `wasm` package dynamically and transfers the numeric vectors to the WASM memory buffer.
4. The Rust/WASM engine processes the calculations and returns JSON containing historical fitted values, forecast points, confidence intervals, and diagnostics.
5. The React app draws the charts using local data.

---

## Writing the Core Math in Rust

Rust is an ideal language for time-series forecasting. It gives us deterministic memory management, zero-cost abstractions, and an ecosystem of libraries like `ndarray` and `statrs` for numerical analysis. 

For example, when a user uploads raw data, it often contains missing values or extreme anomalies (outliers) that would break classical statistical models like ARIMA. In my Rust core (`packages/wasm/src/preprocessing.rs`), we implemented fast IQR-based outlier detection that runs in microseconds:

```rust
use serde::{Deserialize, Serialize};
use wasm_bindgen::prelude::*;

#[derive(Serialize, Deserialize)]
pub struct OutlierIndices {
    pub indices: Vec<usize>,
    pub z_scores: Vec<f64>,
}

/// Detect outliers using the IQR method
#[wasm_bindgen]
pub fn detect_outliers_iqr(values: &[f64]) -> Result<JsValue, JsValue> {
    let mut valid_values: Vec<(usize, f64)> = values
        .iter()
        .enumerate()
        .filter(|(_, &v)| v.is_finite())
        .map(|(i, &v)| (i, v))
        .collect();

    if valid_values.len() < 4 {
        return Ok(serde_wasm_bindgen::to_value(&OutlierIndices {
            indices: vec![],
            z_scores: vec![],
        })?);
    }

    // Sort to find quartiles
    valid_values.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

    let q1_idx = valid_values.len() / 4;
    let q3_idx = (valid_values.len() * 3) / 4;
    let q1 = valid_values[q1_idx].1;
    let q3 = valid_values[q3_idx].1;
    let iqr = q3 - q1;

    let lower_bound = q1 - 1.5 * iqr;
    let upper_bound = q3 + 1.5 * iqr;

    let mean: f64 = valid_values.iter().map(|(_, v)| v).sum::<f64>() / valid_values.len() as f64;
    let variance: f64 = valid_values
        .iter()
        .map(|(_, v)| (v - mean).powi(2))
        .sum::<f64>()
        / valid_values.len() as f64;
    let std_dev = variance.sqrt();

    let mut outlier_indices = Vec::new();
    let mut z_scores = Vec::new();

    for (idx, &val) in values.iter().enumerate() {
        if val.is_finite() && (val < lower_bound || val > upper_bound) {
            outlier_indices.push(idx);
            z_scores.push((val - mean) / std_dev);
        }
    }

    Ok(serde_wasm_bindgen::to_value(&OutlierIndices {
        indices: outlier_indices,
        z_scores,
    })?)
}
```

By compiling this function via `wasm-pack`, JavaScript can invoke it directly, passing array buffers and receiving the results back as native JS values via `serde_wasm_bindgen`.

---

## Bridging Rust and React in the Browser

To make loading the compiled WASM file seamless, I built a TypeScript wrapper class (`WasmEngine`) that handles dynamic imports, module initialization, and the type conversion between standard JavaScript arrays and WebAssembly-friendly typed arrays (`Float64Array`).

Here is a simplified version of the loader wrapper:

```typescript
import type {
  WaySightAIWasm,
  ForecastConfig,
  ForecastResult,
  DataStats
} from '../../types/wasm';

export class WasmEngine {
  private wasmModule: WaySightAIWasm | null = null;
  private isInitialized = false;
  private loadingPromise: Promise<void> | null = null;

  /**
   * Dynamically import and initialize the WASM module
   */
  async load(): Promise<void> {
    if (this.loadingPromise) return this.loadingPromise;
    if (this.isInitialized && this.wasmModule) return Promise.resolve();

    this.loadingPromise = (async () => {
      try {
        // Dynamically import the JS entry point generated by wasm-pack
        const wasmModule = await import('/pkg/waysightai_wasm.js');

        // Initialize the WebAssembly binary instance
        await wasmModule.default();
        wasmModule.init();

        this.wasmModule = wasmModule as unknown as WaySightAIWasm;
        this.isInitialized = true;
        console.log('[WasmEngine] WASM module loaded successfully');
      } catch (error) {
        console.error('[WasmEngine] Failed to load WASM module:', error);
        throw error;
      }
    })();

    await this.loadingPromise;
    this.loadingPromise = null;
  }

  private ensureInitialized(): WaySightAIWasm {
    if (!this.isInitialized || !this.wasmModule) {
      throw new Error('WASM module not initialized. Call load() first.');
    }
    return this.wasmModule;
  }

  /**
   * Run a local forecast job
   */
  async runForecast(
    timestamps: number[],
    values: number[],
    isUnixMs: boolean,
    config: ForecastConfig
  ): Promise<ForecastResult> {
    const wasm = this.ensureInitialized();
    const configJson = JSON.stringify(config);
    
    // Call the exported Rust function
    const resultJson = wasm.run_forecast(timestamps, values, isUnixMs, configJson);
    return JSON.parse(resultJson);
  }
}
```

Because WebAssembly compilation produces a `.wasm` file alongside a `.js` wrapper, Vite allows us to load this dynamically on demand. This keeps the initial page bundle small, downloading the 1.2MB WASM forecasting package only when the user reaches the active editor dashboard.

---

## Performance: Browser vs. Server

Running ARIMA models (which rely on MLE optimization) and triple exponential smoothing (Holt-Winters) in a single-threaded browser environment might sound slow. However, my benchmarks showed surprising results:

| Dataset Size (Rows) | Operation | Execution Location | Duration (ms) |
| :--- | :--- | :--- | :--- |
| 500 rows | Stationarity Tests (ADF + KPSS) | Browser (WASM) | 8ms |
| 500 rows | ARIMA(1,1,1) Parameter Tuning | Browser (WASM) | 45ms |
| 5,000 rows | Outlier Detection & Imputation | Browser (WASM) | 12ms |
| 5,000 rows | Holt-Winters Forecast (Horizon=30) | Browser (WASM) | 185ms |

For typical business dashboards where datasets range from 100 to 10,000 rows, the execution is effectively instantaneous. In contrast, sending this data over an API to a Python service, waiting for it to parse, compute, serialize, and return would easily take 800ms to 2.5 seconds—even before accounting for server cold starts or database query times.

Furthermore, running the computations locally means we can build **interactive visual sliders**. If a user wants to test different seasonal parameters (alpha, beta, gamma), they can slide the control in React, and the chart updates in near real-time (under 50ms) as the local WASM engine re-calculates the parameters.

---

## The Security and Cost Tradeoffs

Going fully client-side with WASM is not a silver bullet. It introduces key trade-offs:

1. **WASM Overhead**: The first load requires fetching the WASM binary. For users on slow connections, this adds a couple of seconds of initialization delay.
2. **IP Exposure**: Because WebAssembly runs on the client, the compiled code is technically accessible. If you have proprietary forecasting models that you want to keep secret, running them in WASM is a risk, as the assembly can be reverse-engineered. For standard algorithms (ARIMA, Simple/Double/Triple Exponential Smoothing), this is not a concern.
3. **CPU Constraints**: If a user uploads a dataset with millions of rows, the browser will lag, and you risk hitting browser memory limits. For massive scale, you still need an orchestrator like AWS SageMaker.

However, the cost savings are massive. Since my servers don't perform any math, my FastAPI backend runs on a cheap container instance. I can handle thousands of concurrent users without paying for expensive auto-scaling CPU nodes or GPU workers.

## Conclusion

Compiling my forecasting math to WebAssembly turned out to be the single best architectural decision I made. It allowed me to promise absolute data privacy to users, gave sub-100ms forecast updates, and kept server infrastructure costs virtually zero.

If you are building data-heavy SaaS tools, don't default to sending every CSV to a server. Let Rust and WebAssembly do the heavy lifting where the data already lives: right in the browser.


Why I Compiled My Time Series Forecasting Engine to WebAssembly


# Orchestrating MLOps with Bedrock AgentCore and SageMaker

When you build machine learning pipelines, you realize quickly that the modeling itself is only a small fraction of the effort. The real work is in the plumbing: cleaning datasets, exploring stationarity and seasonality, engineering lagging or rolling window features, running hyperparameter search, tracking jobs, building reports, and deploying endpoints.

Traditionally, data scientists write messy Jupyter notebooks to do this manually. But what if an **autonomous AI agent** could handle this entire lifecycle?

I set out to build exactly that: a production-ready forecasting system that takes a raw CSV in S3, performs advanced statistical analysis, tunes hyperparameters, trains ARIMA/LSTM models, and deploys a live endpoint—all managed by an LLM-driven orchestrator.

Here is how I built a three-layer architecture combining **AWS Bedrock AgentCore**, **Code Interpreter** sandboxes, and **AWS SageMaker**.

---

## The Three-Layer Architecture

An autonomous system needs to split concerns between reasoning, executing code, and running heavy machine learning jobs. I designed a three-layer architecture:

```mermaid
graph TD
    subgraph Layer1["Layer 1: Orchestration (AgentCore)"]
        Agent["Intelligent Agent<br/>(16 Tools)"]
        Boto3["AWS SDK (boto3)"]
    end

    subgraph Layer2["Layer 2: Data Science (Code Interpreter)"]
        Sandbox["Isolated Python Sandbox"]
        EDA["EDA / Statistical Tests"]
        Features["Feature Engineering"]
        Plotly["Plotly HTML Reports"]
    end

    subgraph Layer3["Layer 3: Scalable ML (SageMaker)"]
        Tuning["Bayesian HPO Jobs"]
        Training["Scale Model Training"]
        Endpoint["Production REST Endpoint"]
    end

    Agent --> Sandbox
    Agent --> Boto3
    Boto3 --> Tuning
    Boto3 --> Training
    Boto3 --> Endpoint
```

Let's break down how these layers interact.

---

## Layer 1: Orchestration with AWS Bedrock AgentCore

The orchestrator is built using [AWS Bedrock AgentCore](https://aws.amazon.com/bedrock/) (packaged via the Strands library). The agent has access to 16 specialized tools ranging from statistical analysis to SageMaker deployment commands.

In `agent.py`, the agent is initialized with a list of Python tool references:

```python
from strands import Agent
from agents.advanced_eda_agent import run_advanced_eda
from agents.intelligent_feature_engineering_agent import recommend_features, create_features
from agents.sagemaker_simple import (
    create_sagemaker_training_job,
    get_training_job_status,
    deploy_sagemaker_model,
    invoke_sagemaker_endpoint
)

agent = Agent(
    name="IntelligentForecastingAgent",
    description="Intelligent time series forecasting with ARIMA and LSTM model comparison.",
    tools=[
        run_advanced_eda,
        recommend_features,
        create_features,
        create_sagemaker_training_job,
        get_training_job_status,
        deploy_sagemaker_model,
        invoke_sagemaker_endpoint,
        # ... other tools
    ]
)
```

The agent uses these tools to execute a 7-step forecasting workflow:
1. **Advanced EDA**: Runs ADF/KPSS tests and seasonal decomposition.
2. **Feature Recommendations**: The agent examines stats and recommends features.
3. **Feature Engineering**: Generates rolling stats, lag columns, and calendar features.
4. **Bayesian Tuning**: Launches SageMaker HPO to find optimal hyperparameters.
5. **Model Training**: Trains classical models (ARIMA) and deep learning models (LSTM).
6. **Comparison**: Evaluates metrics on hold-out test sets.
7. **Report Generation**: Deploys the best model and outputs a presigned URL to an interactive Plotly report.

---

## Layer 2: Safe Execution via Code Interpreter

A major challenge for LLM agents is handling arbitrary code execution. I cannot run unverified code written by an LLM directly on my application server or inside my main database container.

My solution is to offload all data science computations (EDA, feature engineering, Plotly charting) to an isolated **Code Interpreter** sandbox. This sandbox has pre-installed scientific libraries (`pandas`, `scipy`, `statsmodels`, `plotly`) but is locked down with limited CPU, memory, and network permissions.

Here is the design pattern I used for running our Advanced EDA tool:

```python
@tool
def run_advanced_eda(dataset_s3_path: str, time_column: str = None, value_column: str = None) -> str:
    """
    Runs stationarity tests (ADF, KPSS), seasonal decomposition, and ACF/PACF analysis.
    """
    # 1. Parse the target S3 paths
    bucket, key = parse_s3_uri(dataset_s3_path)
    
    # 2. Build the Python script to run in the sandbox
    code = f'''
import pandas as pd
import numpy as np
import json
import subprocess
from statsmodels.tsa.stattools import adfuller, kpss
from statsmodels.tsa.seasonal import seasonal_decompose

# Downloader helper - Sandbox has pre-configured read-only AWS CLI credentials
subprocess.run(['aws', 's3', 'cp', 's3://{bucket}/{key}', '/tmp/data.csv'], check=True)
df = pd.read_csv('/tmp/data.csv')

# ... Auto-detect time/value columns and clean data ...
series = df[value_col].dropna()

# Run ADF and KPSS tests
adf_stat, adf_pvalue, _, _, _, _ = adfuller(series)
kpss_stat, kpss_pvalue, _, _ = kpss(series, regression='c')

# Return results as JSON
print("===JSON_START===")
print(json.dumps({{
    "adf": {{"statistic": adf_stat, "p_value": adf_pvalue}},
    "kpss": {{"statistic": kpss_stat, "p_value": kpss_pvalue}}
}}))
print("===JSON_END===")
'''
    # 3. Send script to the Code Interpreter sandbox and parse stdout
    stdout = sandbox_client.execute(code)
    return extract_json_from_stdout(stdout)
```

By packaging Python code as a text block and passing it to the sandbox, we keep our agent's host server secure. The sandbox pulls the data from S3, computes stats, and returns clean, structured JSON back to the agent.

---

## Layer 3: Scalable ML Training and Deployment on SageMaker

While Code Interpreter is great for light calculations, it lacks the memory and compute capacity to train deep learning models (like LSTMs) or run large-scale hyperparameter searches. For that, I delegate to **AWS SageMaker**.

When the agent reaches the model training step, it uses the AWS SDK (`boto3`) to launch real SageMaker HPO (Hyperparameter Optimization) jobs using preconfigured Scikit-Learn or PyTorch ECR containers:

```python
import boto3
import json

sagemaker = boto3.client('sagemaker')

@tool
def create_sagemaker_training_job(
    job_name: str,
    dataset_s3_path: str,
    role_arn: str
) -> str:
    """
    Launches a Scikit-Learn ECR container on SageMaker for training.
    """
    try:
        response = sagemaker.create_training_job(
            TrainingJobName=job_name,
            AlgorithmSpecification={
                'TrainingImage': '683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:1.2-1-cpu-py3',
                'TrainingInputMode': 'File'
            },
            RoleArn=role_arn,
            InputDataConfig=[{
                'ChannelName': 'train',
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': dataset_s3_path,
                    }
                },
                'ContentType': 'text/csv'
            }],
            OutputDataConfig={
                'S3OutputPath': f's3://{BUCKET}/sagemaker/models/'
            },
            ResourceConfig={
                'InstanceType': 'ml.m5.xlarge',
                'InstanceCount': 1,
                'VolumeSizeInGB': 10
            },
            HyperParameters={
                'p': '2',
                'd': '1',
                'q': '3',
                'seasonal-p': '1',
            },
            StoppingCondition={'MaxRuntimeInSeconds': 3600}
        )
        return json.dumps({'success': True, 'job_arn': response['TrainingJobArn']})
    except Exception as e:
        return json.dumps({'success': False, 'error': str(e)})
```

Because SageMaker jobs run asynchronously and can take 5 to 60 minutes, the agent uses a polling loop tool (`get_training_job_status`) to track the execution. Once the job succeeds, the agent triggers `deploy_sagemaker_model` to package the model artifact (`model.tar.gz`) from S3 and host it behind a real-time SageMaker endpoint.

---

## The Two-Phase Data Pattern for Reports

Generating and serving rich reports presents a unique security challenge. The Code Interpreter sandbox (which generates the Plotly interactive HTML files) should **never** have write access to my production S3 buckets, nor should it hold long-term AWS credentials.

To solve this, I implemented a **Two-Phase Data Handling** pattern:

```mermaid
sequenceDiagram
    participant Agent as AgentCore (Has S3 Write Credentials)
    participant CI as Code Interpreter Sandbox (No S3 Credentials)
    participant S3 as AWS S3

    Agent->>CI: Execute charting script
    CI->>CI: Generate Plotly chart & serialize to HTML
    CI-->>Agent: Print HTML string to stdout
    Agent->>Agent: Extract HTML string from stdout
    Agent->>S3: Upload HTML content using boto3
    S3-->>Agent: Return secure presigned URL
```

By keeping the AWS write credentials strictly in the Orchestration Layer (AgentCore) and treating the Code Interpreter as a pure calculation engine, I maintain least-privilege security across my cloud infrastructure.

---

## Conclusion & Lessons Learned

Building an autonomous ML pipeline taught me several key lessons:

1. **Decouple CPU Profiles**: Keep light calculations (cleaning, simple charts) inside rapid, cheap sandboxes, and reserve expensive cloud computing instances (SageMaker ml.m5/ml.p3 nodes) for heavy model training.
2. **Handle Asynchrony Gracefully**: LLMs are naturally sequential. To prevent the agent from locking up during a 30-minute SageMaker job, use explicit state-machine trackers and periodic polling tools.
3. **Sandbox Everything**: Never let an LLM-generated script run natively. Use a secure container runtime with strict network isolation.

By combining AWS Bedrock AgentCore's reasoning capabilities with SageMaker's massive scaling power, I successfully took a standard time-series pipeline and turned it into a fully autonomous MLOps platform.


Orchestrating MLOps with Bedrock AgentCore and SageMaker


# What if a 9B model could beat Sonnet at SQL?

I was building a data agent. The kind of thing where you type a question in plain English and it writes SQL, runs it, and shows you a chart. The demo worked great. Single question, single query, clean result. Ship it.

Then I tried using it the way a real analyst would.

```text
Show revenue by country.
Only France.
Now break it down by month.
That looks empty. Try the country code instead.
```

Four turns. By the third one, the model had forgotten what we were looking at. By the fourth, it was writing SQL against a table it had already abandoned. The demo was a lie. Not because the model couldn't write SQL — it clearly could. But because it couldn't hold a conversation about data.

That failure sent me down a rabbit hole.

## The benchmark gap nobody talks about

When you hear that frontier models score 85% on text-to-SQL, that number almost certainly comes from a single-turn benchmark. One question, one answer, fresh context. [Spider](https://yale-lily.github.io/spider) and [BIRD](https://arxiv.org/abs/2305.03111) are the gold standards, and they've driven real progress.

But I kept pulling at the thread: what happens when we test models the way analysts actually work? Not one question at a time, but in sequences. Follow-ups. Refinements. "Wait, that doesn't look right — try it this way."

That's where [SParC](https://arxiv.org/abs/1906.02285) and [CoSQL](https://arxiv.org/abs/1909.05378) come in. SParC tests whether a model can handle sequences of related questions over databases it has never seen. CoSQL goes further — it's actual dialogue, with ambiguity, clarification requests, and all the messiness of a real conversation between an analyst and a SQL assistant.

And then there's [BIRD-Interact](https://arxiv.org/abs/2510.05318), which argues that even those benchmarks are too polite. Real assistants need to recover from execution errors, ask for clarification, and deal with the fact that their own earlier answers might have been wrong.

Here's the thing that got me: **frontier models that look dominant on single-turn benchmarks get dramatically less impressive on multi-turn ones.** The gap isn't small. It's the kind of gap that makes you wonder whether we've been measuring the right thing.

So I started a thought experiment. If multi-turn is where the big models stumble, and if the problem is more about *behavior* than raw intelligence — keeping context, carrying filters, resolving references — could a small model be fine-tuned to close that gap?

Could a 9B model running on my desk beat Sonnet at the thing that actually matters?

## Setting up the experiment

I grabbed [Qwen 3.5 9B](https://huggingface.co/Qwen) as my base. Why Qwen? It runs on consumer hardware, Unsloth has good support for it, and the instruction-tuned variant already knows what SQL looks like. The training ran on a single RTX 5090 with bf16 LoRA — rank 32, alpha 64, all the standard attention and MLP projections. Nothing exotic. About four hours of compute for each run.

For evaluation, I set up a CoSQL proxy: 100 turns spread across 32 real dialogs. The blog results I'm reporting use a clean subset of 43 follow-up turns from 12 of those dialogs, chosen because every run used the same rows, the same scorer, and the same protocol.

One crucial design choice: I used **generated-history rollout**, not teacher-forced history. In teacher-forced evaluation, the model gets perfect prior turns — it sees what the ideal previous SQL *should have been*. That's a useful diagnostic, but it's a fantasy. In practice, your model sees its own earlier output. If turn 2 was wrong, turn 3 has to deal with that. Generated-history makes the numbers uglier, but it's the only honest way to test a real assistant.

For the frontier baseline, I ran Claude Sonnet 4.6 through OpenRouter on the exact same rows, same scorer, same protocol. Same rules for both sides.

## Five ideas, one clear signal

I didn't just try one training strategy. I was exploring a research question, and that meant testing different hypotheses about *what* the model needs to learn.

[DIN-SQL](https://arxiv.org/abs/2304.11015) had shown that decomposing the text-to-SQL problem into subtasks could help. That felt right intuitively — writing SQL isn't one skill. It's figuring out which tables to use, understanding what values look like in the database, grasping what the user actually wants, generating the code, and catching mistakes. All tangled together. The question was which piece matters most for multi-turn conversations.

Here's what I tried:

**Direct SQL** — the control group. Take the user's question and the database schema, train the model to produce SQL directly. No intermediate steps, no clever representations. Just question in, query out. If this doesn't improve over the base model, nothing else matters.

**Semantic decomposition** — before writing any SQL, first have the model describe *what the user is asking for* in structured terms: which entity, which metric, which filter, which grouping. The idea is that understanding "break it down by month" depends on first understanding what "it" refers to from the previous turn. That's a meaning problem before it's a SQL problem. The training data included a Cube-inspired semantic model (entities, dimensions, measures, join hints) derived from the database schema, so the model could learn to map user language to governed business concepts before touching SQL.

**Metric DSL** — a small intermediate language for business metrics. Instead of going straight to `SUM(orders.amount)`, the model would first emit something like `MEASURE(revenue) BY customer_country`, and a compiler would expand that into the actual SQL. Attractive in theory, because when an analyst says "show me revenue," they're thinking about a business metric, not a column name. The model should preserve that intent and let something else handle the SQL plumbing.

**Behavior recovery** — in real use, the model will sometimes write bad SQL. The next turn has to deal with that. So instead of always training on perfect conversation histories, I generated training examples where earlier turns contained the model's *actual* (sometimes wrong) SQL output. The idea: teach the model to recognize when a previous query returned empty results or wrong data, and adjust. Show it the messy states it'll actually face in production.

**Schema selection** — this one was a deliberate cheat, designed to answer a diagnostic question. I took the correct reference SQL for each question, parsed it to extract exactly which tables, columns, and joins were needed, then fed those as hints in the prompt before asking the model to write SQL. Obviously you can't do this in production — you'd need to already know the answer to give the hint. But it tells you something critical: *if the model already knew where to look in the database, how good would its SQL be?*

The results told a clear story.

## What the numbers say

| Run | Value accuracy | Strict accuracy | What it means |
| --- | ---: | ---: | --- |
| Claude Sonnet 4.6 (via OpenRouter) | `0.674` | `0.395` | The frontier. Same questions, same scorer. |
| Qwen 9B + schema hints from correct answer | `0.651` | `0.465` | Cheating — told which tables and columns to use. Not deployable, but revealing. |
| Qwen 9B + semantic training (50 steps) | `0.581` | `0.326` | Best result from a fair, deployable local model. |
| Qwen 9B + direct SQL training (50 steps) | `0.581` | `0.302` | Same value accuracy, slightly less precise output shape. |
| Qwen 9B out of the box | `0.558` | `0.302` | No training, just the base model with a SQL prompt. |
| Qwen 9B + recovery training | `0.558` | `0.302` | Trained on messy histories. Didn't help yet. |

Before reading into these, a note on the metrics. *Value accuracy* asks: did the SQL return the right data? It's lenient about column aliases — if you called the column `total` instead of `revenue` but got the right numbers, that's a pass. *Strict accuracy* also cares about output shape. Both matter, but value accuracy is the more meaningful measure of whether the model is actually useful.

Some of these results are what you'd expect. Some are not.

**The clean local lift is real but small.** The best fine-tuned adapters (semantic and direct SQL) both reach `0.581` value accuracy, up from `0.558` for raw Qwen. That's a 2.3 percentage point improvement. Not nothing — it proves the model can learn from the task-specific data — but it doesn't close the gap to Sonnet's `0.674`.

**Behavior recovery was a disappointment.** Conceptually, it's the right idea. In practice, it tied raw Qwen. The training data didn't teach the useful repair moves yet. That's a data quality problem, not a conceptual one, but it means recovery training is still a hypothesis, not a result.

**The schema hints result is the real story.** When you cheat and tell the model which tables and columns it needs, the local model jumps to `0.651` — within two points of Sonnet on value accuracy and actually *ahead* on strict accuracy (`0.465` vs `0.395`). You can't deploy this — the hints come from parsing the correct answer, which you obviously don't have at inference time. But it's the kind of result that changes the direction of your research.

## What the schema hints result actually means

Here's what the cheat does concretely. The evaluation code takes the reference SQL (the correct answer), parses it with [sqlglot](https://github.com/tobymao/sqlglot), and extracts the relevant tables, the specific columns used, the join conditions, the projection shape, and grouping. It then injects those as hints into the prompt: "you need the `orders` table joined to `customers` on `customer_id`, selecting `amount` and `country`." The model still has to write the actual SQL, but it no longer has to *figure out where to look*.

Most wrong SQL is plausible SQL. The model doesn't usually produce garbage — it produces a reasonable-looking query against the wrong table, with the wrong join, using a value format that doesn't match what's in the database. "France" instead of "FR". `customer_orders` instead of `orders`. A `LEFT JOIN` where an `INNER JOIN` was needed.

When you remove that confusion — when the model knows where to look — a 9B local model gets within striking distance of the frontier. The remaining errors are about value grounding (the database stores "FR" but the user said "France"), execution edge cases, and grain issues (accidentally duplicating rows through a many-to-many join). The table-and-column selection problem is the bottleneck, not raw SQL writing ability.

This completely reframes the question. I started asking "can a small model write SQL as well as Sonnet?" and ended up learning that "can a small model pick the right tables?" is the question that actually matters.

## The honest evidence boundary

I want to be precise about what I can and can't claim from this work.

**I can say:** The best clean local fine-tunes slightly beat raw Qwen on the same 43 generated-history rows. Fine-tuning works. The model learns something real from multi-turn SQL data.

**I can say:** Sonnet 4.6 remains ahead on the clean comparison. The frontier hasn't been reached.

**I can say:** When the model is told which tables and columns to use (by peeking at the answer), it nearly closes the gap to Sonnet. That tells us where the real problem is.

**I cannot say:** A local 9B model beats Sonnet. Not yet. Not on this data.

**I cannot say:** The schema-hints setup is deployable. You can't peek at the answer in production. It's a ceiling test that shows where to push.

**I cannot say:** Recovery training is proven. It tied the baseline. The idea might be right, but the current training data didn't teach it.

This is a 43-row generated-history slice from CoSQL. It's reproducible, controlled, and honest. It is not a final benchmark claim.

## So what was the point?

The thought experiment started with "what if a 9B model could beat Sonnet at SQL?" and the honest answer is: not yet, not cleanly. But the experiment was worth running because of what it revealed about *where* the gap actually lives.

The gap isn't in SQL syntax. Raw Qwen already produces syntactically valid SQL almost every time — `1.000` syntax accuracy across the board. The gap isn't even in understanding English questions — the model generally knows what you're asking.

The gap is in navigating the database: knowing which tables to query, which columns to select, how to join them, and what the stored values actually look like. And in multi-turn conversations, that problem compounds. You're not just picking the right tables for one question. You're carrying forward context from previous turns, resolving references like "break *that* down by month," and adapting when a filter that worked on turn 2 needs a different column on turn 4.

That's a learnable problem. It's not a matter of needing a smarter model — it's a matter of needing better training data about which tables matter, how values are stored, and what carries forward across turns. The schema-hints experiment proved that the 9B model *can* do this work when it knows where to look. The next step is teaching it to figure that out on its own.

## What comes next

The next experiment should teach context selection directly:

1. Which tables matter for this question?
2. Which columns express the requested metric or dimension?
3. Which stored values match the user's wording?
4. What should carry forward from the previous turn?
5. What changed in the follow-up: filter, grouping, metric, or repair?
6. When execution returns nothing, what's the next reasonable attempt?

That's the bridge from a small CoSQL checkpoint toward [BIRD-Interact](https://arxiv.org/abs/2510.05318)-style evaluation. First make the local protocol stable. Then scale the benchmark. Then ask whether the local model can close the hosted gap.

The thought experiment isn't over. It just got more specific.

## Code

The full fine-tuning repo: [`github.com/xdanny/multiturn-sql-finetuning`](https://github.com/xdanny/multiturn-sql-finetuning).

## Sources

- [BIRD paper](https://arxiv.org/abs/2305.03111)
- [SParC paper](https://arxiv.org/abs/1906.02285)
- [CoSQL paper](https://arxiv.org/abs/1909.05378)
- [BIRD-Interact paper](https://arxiv.org/abs/2510.05318)
- [DIN-SQL paper](https://arxiv.org/abs/2304.11015)
- [Qwen model family](https://huggingface.co/Qwen)
- [OpenRouter](https://openrouter.ai/)


What if a 9B model could beat Sonnet at SQL?

RAG Is Dead. The Engineering Isn't.


# Building data quality checks in your pySpark data pipelines

Data quality is a rather critical part of any production data pipeline. In order to provide accurate SLA metrics
and to ensure that the data is correct, it is important to have a way to validate the data and report the metrics
for further analysis. In this post, we will look at how to build data quality checks in your pySpark data pipelines.

## Exploring Delta Live Tables

Delta Live Tables is a new feature in Databricks that allows users to build reliable data pipelines with built-in
data quality metrics and monitoring. It is a new abstraction on top of Delta Lake that allows users to query the
data using streaming live tables. The data is updated in real-time as the underlying data changes. What caught my eye
was the data quality capabilities that the users can specify on dataset level. Using python decorators we can specify
@expect_all, @expect_all_or_drop, and @expect_all_or_fail expectations that accept a python dictionary as an argument,
where the key is the expectation name and the value is the expectation constraint. Example:

```python
@dlt.expect("valid timestamp", "col("timestamp") > '2012-01-01'")
@dlt.expect_or_drop("valid_current_page", "current_page_id IS NOT NULL AND current_page_title IS NOT NULL")
@dlt.expect_or_fail("valid_count", "count > 0")
```

Metrics of clean records and failed records are automatically collected and stored in the Delta Live Table metadata, so
the users can set up alerts and monitor the data quality of their pipelines.

Delta Live Tables however still face quite some limitations and are not yet ready for production use. Some limitations
include:
1. The data quality checks are only available for streaming live tables, not for batch tables. We can still create
streaming tables from batch tables, but if the version of your data is changing the pipeline will fail.
2. Lack of testing capabilities. There is no way to test the data quality checks in a local environment because
dlt package is available only in Databricks runtime.
3. Lack of documentation. The documentation is very limited and it is not clear how to use the data quality checks.
Currently only python and SQL API are supported.
4. Setting up DLT job doesn't support all the parameters that are available in the Databricks job.

## Building your own data quality checks as python decorators

In order to overcome the limitations of Delta Live Tables, we can build our own data quality checks as python decorators.
The idea is to create a decorator that will accept a python list of arguments, which will be constraints that we will
apply for a determined column. We will collect all the necessary metrics and store them as part of the Delta Lake
metadata.

We will start by building two simple conditions for our data, uniqueness and filtering based on a condition.

```python
from abc import ABC, abstractmethod

class ColumnCondition(ABC):
    @abstractmethod
    def get_cols(self):
        pass


class UniqueCondition(ColumnCondition):

    def __init__(self, col):
        self.col = col

    def get_cols(self):
        return self.col


class FilterCondition(ColumnCondition):

    def __init__(self, left_col, right_col):
        self.left_col = left_col
        self.right_col = right_col

    def get_cols(self):
        return self.left_col, self.right_col
```

For ease of use, let's define a couple of factory functions that will return the right condition:

```python
def is_not_null(col):
    return FilterCondition(col + " is not null", col + " is null")

def is_unique(col):
    return UniqueCondition(col)
```

The main idea is to use a function as a decorator argument using a certain column, which will return a condition object.
We can use the condition object to pattern match and apply the specific function depending on the condition type.
We start by creating a simple python decorator using functools wraps:

```python
    def expect_or_drop(self, conditions: List[FilterCondition]):
        def decorator(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                retval = function(*args, **kwargs)
                # apply conditions
                return retval
            return wrapper
        return decorator
```

We will create an Expectations class that will contain all the data quality checks. The rsd arguments represents the
maximum relative standard deviation allowed for the approx_count_distinct_functions. Read more
[here](https://spark.apache.org/docs/3.1.2/api/python/reference/api/pyspark.sql.functions.approx_count_distinct.html).

```python

class Expectations:

    def __init__(self, spark: SparkSession, rsd=0.05):
        self.spark = spark
        self.schema = StructType([StructField("condition", StringType(), True),
                                  StructField("dropped_records", IntegerType(), True),
                                  StructField("clean_records", IntegerType(), True)])
        emptyRDD = spark.sparkContext.emptyRDD()
        self.metrics = spark.createDataFrame(emptyRDD, schema=self.schema)
        self.rsd = rsd
```

The metrics dataframe will contain the metrics for each data quality check. We can proceed to create our filtering
and uniqueness checks:

```python

    def apply_condition(self, dataframe, condition):
        if isinstance(condition, FilterCondition):
            return self.filter_condition(dataframe, condition.get_cols())
        elif isinstance(condition, UniqueCondition):
            return self.is_unique_extend(dataframe, condition.get_cols())
        return dataframe

    def filter_condition(self, dataframe: DataFrame, left_right) -> DataFrame:
        left, right = left_right
        total_records = dataframe.count()
        dropped_records = dataframe.filter(right).count()
        df = self.spark.createDataFrame([(left, dropped_records, (total_records - dropped_records))], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.filter(left)

    def is_unique_extend(self, dataframe: DataFrame, col) -> DataFrame:
        total_records = dataframe.select(F.col(col)).count()
        distinct_records = dataframe.select(F.approx_count_distinct(col, self.rsd)).collect()[0][0]
        dropped_records = total_records - distinct_records
        df = self.spark.createDataFrame([(col + " is unique", dropped_records, distinct_records)], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.dropDuplicates([col])
```

In order to apply the conditions we will use the apply_condition function to every condition in the list. In order to
do that we will use the functools reduce function as foldLeft:

```python
   foldl = lambda func, acc, xs: reduce(func, xs, acc)
   @wraps(function)
   def wrapper(*args, **kwargs):
        retval = function(*args, **kwargs)
        return foldl(self.apply_condition, retval, conditions)
```

Wrapping it all together:

```python

from abc import ABC, abstractmethod
from functools import reduce, wraps
from typing import List

import pyspark.sql.functions as F
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.types import IntegerType, StringType, StructField, StructType

foldl = lambda func, acc, xs: reduce(func, xs, acc)


class ColumnCondition(ABC):
    @abstractmethod
    def get_cols(self):
        pass


class UniqueCondition(ColumnCondition):

    def __init__(self, col):
        self.col = col

    def get_cols(self):
        return self.col


class FilterCondition(ColumnCondition):

    def __init__(self, left_col, right_col):
        self.left_col = left_col
        self.right_col = right_col

    def get_cols(self):
        return self.left_col, self.right_col


def is_not_null(col):
    return FilterCondition(col + " is not null", col + " is null")


def is_unique(col):
    return UniqueCondition(col)


class Expectations:

    def __init__(self, spark: SparkSession, rsd=0.05):
        self.spark = spark
        self.schema = StructType([StructField("condition", StringType(), True),
                                  StructField("dropped_records", IntegerType(), True),
                                  StructField("clean_records", IntegerType(), True)])
        emptyRDD = spark.sparkContext.emptyRDD()
        self.metrics = spark.createDataFrame(emptyRDD, schema=self.schema)
        self.rsd = rsd

    def expect_or_drop(self, conditions: List[FilterCondition]):
        def decorator(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                retval = function(*args, **kwargs)
                return foldl(self.apply_condition, retval, conditions)
            return wrapper
        return decorator

    def apply_condition(self, dataframe, condition):
        if isinstance(condition, FilterCondition):
            return self.filter_condition(dataframe, condition.get_cols())
        elif isinstance(condition, UniqueCondition):
            return self.is_unique_extend(dataframe, condition.get_cols())
        return dataframe

    def filter_condition(self, dataframe: DataFrame, left_right) -> DataFrame:
        left, right = left_right
        total_records = dataframe.count()
        dropped_records = dataframe.filter(right).count()
        df = self.spark.createDataFrame([(left, dropped_records, (total_records - dropped_records))], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.filter(left)

    def is_unique_extend(self, dataframe: DataFrame, col) -> DataFrame:
        total_records = dataframe.select(F.col(col)).count()
        distinct_records = dataframe.select(F.approx_count_distinct(col, self.rsd)).collect()[0][0]
        dropped_records = total_records - distinct_records
        df = self.spark.createDataFrame([(col + " is unique", dropped_records, distinct_records)], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.dropDuplicates([col])
```

Let's see the decorator in action:

```python
    @expectation.expect_or_drop([is_not_null("row"), is_unique("row")])
    def read_dataframe(df):
        return df

    result = read_dataframe(df_1)
    print(result.collect())
    print(expectation.metrics.collect())
```

The console will print:
```bash
[Row(row='row1', row_number=1)]
[Row(condition='row is not null', dropped_records=1, clean_records=2), Row(condition='row is unique',
dropped_records=1, clean_records=1)]
```

We can extend and add some plotting functions as well to our Expectations class.

```python
    def plot_pie_with_total(self, figsize=(10, 10)):
        labels = ["clean_records", "dropped_records"]
        df = self.metrics.toPandas()
        size = len(df.index)
        if size == 1:
            fig, axs = plt.subplots(1)
            axs.pie(df.iloc[0][labels], labels=labels, autopct='%1.1f%%')
            axs.set_title(df.iloc[0]["condition"])
        else:
            fig, axs = plt.subplots(1, size, figsize=figsize)
            for i in range(size):
                axs[i].pie(df.iloc[i][labels], labels=labels, autopct='%1.1f%%')
                axs[i].set_title(df.iloc[i]["condition"])
        plt.show()
```

That will generate a pie chart with the metrics for every condition:
![Image Name](/assets/blog/data-quality-pyspark/data_quality.png)

Happy coding and stay safe!


Building data quality checks in your pySpark data pipelines


# Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter

In this blogpost we will continue our journey of testing our Data Pipelines. If you haven't checked out
the first post, [make sure you do](/post/automate-all-your-pyspark-unit-test-with-hypothesis). As assessed in the first post, writing unit tests in pyspark can be challenging.
Especially when we go to the world of notebooks, setting up your testing infrastructure can be challenging.
Most often the data scientists/data engineers in your organization will prefer working on and developing their
Data Pipelines in a dedicated environment like Databricks or JupyterLab. In order to enable cross-team
collaboration, peer-review and enforce best testing practices, the notebooks will have to be version controlled.
In Databricks, for example, the preferred way to work on notebooks is to use
[Databricks Repos](https://docs.databricks.com/repos/index.html). However, because Databricks doesn't have
support for pytest, you cannot write a separate notebook for tests. The data engineer will be writing
pytest's using an IDE as part of the notebook pull request, the execution part being left to the CI/CD pipeline
in a mocked Spark environment. This ensures that our code is determinstic and our functions are correct, but
it does not necessarily translate to correct functioning in a production-like environment.

## Microsoft Nutter - Testing Framework for Databricks

To overcome the limits described above, a popular framework open-sourced by Microsoft is
[Nutter](https://github.com/microsoft/nutter). The library has two main components: the runner and the CLI.
We will be exploring in this blogpost only the runner, leaving the CLI as an exercise for the readers. The
runner runtime is a module that can be used once you install the library in your notebook. In Databricks
we can use the following command, that will install nutter as well as hypothesis and tinsel libraries used to
generate our datasets:

```python
%pip install -U nutter hypothesis tinsel
```

As an example, we will be using the same function as in the first series blogpost to illustrate the example
in a separate notebook:

```python
from pyspark.sql import DataFrame
import pyspark.sql.functions as F

def test_dataframe_transformations(df: DataFrame) -> DataFrame:
    temp_df = df.select('*', F.explode_outer('order_items'))\
        .select(F.col('created_at'), F.col('col.item_value').alias('item_value'))\
        .withColumn('day', F.to_date('created_at'))\
        .withColumn('week', F.weekofyear('created_at'))

    weekly_revenue_df = temp_df.groupBy('week').agg(
        F.sum('item_value').alias('weekly_revenue')
    )

    return temp_df.groupBy('week', 'day').agg(
        F.sum('item_value').alias('daily_revenue')
    ).join(weekly_revenue_df, on='week', how='inner')
```

Now, in our testing notebook, we can import the needed functions with the following Databricks command:

```python
%run /{notebook_to_test_path}
```

In order to setup our test datasets, we will be using the same setup as in the previous blogpost:

```python
import hypothesis.strategies as st
from hypothesis import given, settings, HealthCheck

settings.register_profile(
"my_profile",
max_examples=50,
deadline=60 * 1000,  # Allow 1 min per example (deadline is specified in milliseconds)
suppress_health_check=(HealthCheck.too_slow, HealthCheck.data_too_large),
)
```

```python
from datetime import datetime, timedelta
from typing import List, NamedTuple
from tinsel import struct, transform

@struct
class OrderItem(NamedTuple):
item_id: str
item_value: float

@struct
class Order(NamedTuple):
order_id: str
created_at: datetime
updated_at: datetime
customer_id: str
order_items: List[OrderItem]

week_sunday = (datetime.now() - timedelta(days = datetime.now().weekday() + 1))
time_window = (datetime.now() - timedelta(days = datetime.now().weekday() + 7))

def order_strategy():
return st.lists(st.builds(Order,
order_id=st.text(min_size=5),
created_at=st.datetimes(min_value=time_window, max_value=week_sunday),
updated_at=st.datetimes(min_value=time_window, max_value=week_sunday),
customer_id=st.text(min_size=5),
order_items=st.lists(st.builds(OrderItem,
item_id=st.text(min_size=5),
item_value=st.floats(min_value=0, max_value=100)), min_size=1)), min_size=5)
```

We are now ready to create our first tests with Microsoft Nutter. The test framework follows the following
conventions for the test fixtures:

```
before_(testname) - (optional) - if provided, is run prior to the 'run_' method. This method can be used to setup any test pre-conditions

run_(testname) - (optional) - if provider, is run after 'before_' if before was provided, otherwise run first. This method is typically used to run the notebook under test

assertion_(testname) (required) - run after 'run_', if run was provided. This method typically contains the test assertions

after_(testname) (optional) - if provided, run after 'assertion_'. This method typically is used to clean up any test data used by the test
```

In the case where we want to test full notebooks, we can setup for example our static dataframe from a file
in the before function, run the ETL transformations, assert that the created delta table corresponds to our
expectations and use the after in order to clean up all the test resources we might have spinned up. In this
blogpost, however, we will only be using the assertion as we are generating the data on the fly. Our test
fixture then would look like this:

```python
from runtime.nutterfixture import NutterFixture, tag

default_timeout = 600

class Test1Fixture(NutterFixture):
def init(self):
NutterFixture.init(self)

@given(data=order_strategy())
@settings(settings.load_profile("my_profile"))
def test_samples(self, data):
  df = spark.createDataFrame(data=data, schema=transform(Order), verifySchema=False)
  collected_list = test_dataframe_transformations(df).collect()
  for row in collected_list:
    assert row['weekly_revenue'] >= row['daily_revenue']

def assertion_test_transform_data(self):
    self.test_samples()

def after_test_transform_data(self):
    print('done')
```

We can now run our test notebook, and the results will look like this:

```
Notebook exited:
Notebook: N/A - Lifecycle State: N/A, Result: N/A
Run Page URL: N/A

PASSING TESTS

test_transform_data (31.414987 seconds)
```

This notebook can be added as part of your CI/CD pipeline, and we will explore in the next blogpost how.
We can write and run unit/integration/end2end tests using Nutter and export the results via CSV or JUnit.
Happy Coding!

### References

[Microsoft Nutter](https://github.com/microsoft/nutter)


Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter


# Automate all your PySpark Unit Test with Hypothesis!

Unit testing is often regarded as a main pillar of testing your software applications, and it usually involves
testing a single/unit component and ensuring that it covers all the edge cases the software developer can think of.
The Functional Programming world, specifically Haskell with QuickCheck and Scala with ScalaCheck, have introduced the
idea of testing based on property specifications and automatic test data generation in order to complement
the traditional unit testing approach. The idea is that we define a property that specifies the program behaviour and
the data is automatically generated for us to overcome all the test scenarios, using shrinking in order to simplify the
verification of the test case.

## PySpark property-based testing

Testing your PySpark notebooks can be quite challenging by itself, the data infrastructure itself has often quite
limited support for pytest and adding such tests to your existing data pipelines adds quite some complexity.
Additionally, testing your ETL transformation function can be quite cumbersome as it involves generating data,
maintaining it, or simply skipping the unit tests in favor of simple integrity tests. In this blog post we will see
how we can easily generate data from dataclasses and think of unit testing in terms of properties rather than data
integrity. We will start with Hypothesis, an excellent Python library built for property-based testing.

We can begin by specifying some simple dataclasses we will be using in order to showcase the potential of the library:

```python
from datetime import datetime, timedelta
from typing import List, NamedTuple
from tinsel import struct, transform

@struct
class OrderItem(NamedTuple):
  item_id: str
  item_value: float

@struct
class Order(NamedTuple):
  order_id: str
  created_at: datetime
  updated_at: datetime
  customer_id: str
  order_items: List[OrderItem]
```

Our transaction dataset will be composed by orders with a set of order items for each order. The dataclasses are defined
as NamedTuples because PySpark has limited support for dataclasses. We will be working with a secondary python library
called Tinsel, which allows us to easily convert our dataclass schema to a Spark DataFrame schema.

At this point we can leverage the Hypothesis library to specify arbitrary generators we will be using for our dataclasses:

```python
import hypothesis.strategies as st
from hypothesis import given, settings, HealthCheck

week_sunday = (datetime.now() - timedelta(days = datetime.now().weekday() + 1))
time_window = (datetime.now() - timedelta(days = datetime.now().weekday() + 7))

def order_strategy():
    return st.lists(st.builds(Order,
                     order_id=st.text(min_size=5),
                     created_at=st.datetimes(min_value=time_window, max_value=week_sunday),
                     updated_at=st.datetimes(min_value=time_window, max_value=week_sunday),
                     customer_id=st.text(min_size=5),
                     order_items=st.lists(st.builds(OrderItem,
                                           item_id=st.text(min_size=5),
                                           item_value=st.floats(min_value=0, max_value=100)), min_size=1)), min_size=5)
```

In the above code snippet I have specified some upper and down boundaries for the datetimes I will be using in my test
dataset. Likewise for the strings, floats and lists I have used different specifiers. For additional custom generator
strategy consult the excellent Hypothesis documentation.

Now that we have our custom generators we can specify our example transformation function we would like to test:

```python
from pyspark.sql import DataFrame
import pyspark.sql.functions as F

def test_dataframe_transformations(df: DataFrame) -> DataFrame:
    temp_df = df.select('*', F.explode_outer('order_items'))\
        .select(F.col('created_at'), F.col('col.item_value').alias('item_value'))\
        .withColumn('day', F.to_date('created_at'))\
        .withColumn('week', F.weekofyear('created_at'))

    weekly_revenue_df = temp_df.groupBy('week').agg(
        F.sum('item_value').alias('weekly_revenue')
    )

    return temp_df.groupBy('week', 'day').agg(
        F.sum('item_value').alias('daily_revenue')
    ).join(weekly_revenue_df, on='week', how='inner')
```
The property I am interested to test in the above code snippet is the revenue itself. Logically the daily revenue should
always be smaller or equal than the weekly revenue, unless I have an undetected bug in my code.

Now, in order to be able to use pyspark in your unit testing you need to setup a global SparkSession:

```python
import unittest

from pyspark.sql import SparkSession

class PySparkTestCase(unittest.TestCase):
    """Set-up of global test SparkSession"""

    @classmethod
    def setUpClass(cls):
        cls.spark = (SparkSession
                     .builder
                     .master("local[1]")
                     .appName("PySpark unit test")
                     .getOrCreate())

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()
```

We can setup different profiles for our Hypothesis generators, for example the amount of examples or the deadline:

```python
settings.register_profile(
    "my_profile",
    max_examples=50,
    deadline=60 * 1000,  # Allow 1 min per example (deadline is specified in milliseconds)
    suppress_health_check=(HealthCheck.too_slow, HealthCheck.data_too_large),
)
```

Using our generators, settings and SparkSession we can now write our first unit test:

```python
class SimpleTestCase(PySparkTestCase):

    def test_spark(self):
        @given(data=order_strategy())
        @settings(settings.load_profile("my_profile"))
        def test_samples(data):
            df = self.spark.createDataFrame(data=data, schema=schema, verifySchema=False)
            collected_list = test_dataframe_transformations(df).collect()

            for row in collected_list:
                assert row['weekly_revenue'] >= row['daily_revenue']

        test_samples()
```

This code snippet will feed all the generated datasets to the transformation function as a Spark DataFrame, collect
the results and check the property we were after all along.

In the next blog post we will tackle the orchestration of unit tests and the addition to your Databricks CI/CD pipeline
using [Microsoft's Nutter](https://github.com/microsoft/nutter)! Stay Tuned!

#### The code is available on [github](https://github.com/xdanny/property-testing-pyspark)

### References
[Hypothesis Docs](https://hypothesis.readthedocs.io/en/latest/)
[Tinsel PySpark schema converter](https://github.com/benchsci/tinsel)
[Property Based Testing with QuickCheck](https://typeable.io/blog/2021-08-09-pbt)


hypothesis2 posts

Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter

Automate all your PySpark Unit Test with Hypothesis!

Recommended

I Ran Codex on a Local Model — No Cloud, No Bill

Why I Compiled My Time Series Forecasting Engine to WebAssembly

Orchestrating MLOps with Bedrock AgentCore and SageMaker