09 · A workflow in JAX: build it, use it, share it

You will learn

• Fit every parameter with jax.grad
• Draw the bootstrap with one jax.vmap call
• Save and restore jax.Array state, with jax in the env pin

You need

• Python 3.12 and uv (JAX installs for CPU)
• Section 3 is recorded; rerunning it needs the reference deployment

Run this notebook Offline, section 3 recorded

notebooks/09_workflow_jax.ipynb

Runs offline. Section 3 shows a recorded session against the reference deployment and cannot be run without it.

About 25 minutes

Access needed. The packages and notebooks are not yet publicly distributed. The steps below need a checkout of the UMP repository. Pages show recorded outputs, so you can follow along without running them.

cd notebooks && uv run jupyter lab

Runs offline; section 3 shows a recorded session against the reference deployment. Uses babelfin and JAX (CPU). The data is SYNTHETIC, generated below with the same known truth as notebook 06, so every answer can be checked.

Priya’s team prototypes media models in JAX: the whole model is one differentiable function, so jax.grad fits every parameter at once, and jax.vmap turns two hundred bootstrap fits into one call. The catch has always been the hand-off. A JAX model lives in a kernel full of jax.Arrays, and nobody else can rerun it next quarter, let alone ask it a question from another app.

This notebook follows the same arc as the scikit-learn and plain-Python ones, so the three can be read side by side:

1. Build: a seeded weekly world, an input contract, a JAX model module, named questions, a fit, answers checked against the truth.
2. Save and use: a bundle on disk, opened on a “new machine”, trusted after reading the code, scoring twelve new weeks, and four labelled ways a rerun can go differently.
3. Share and serve: Priya publishes, Sam approves, Li runs it with no manual trust, then serves it as an app.

In [1]:

import os
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import babelfin as fin
import httpx
import numpy as np
import pandas as pd
from IPython.display import display
from ump_sdk.ports import base_url

WF = Path("out/workflows")
WF.mkdir(parents=True, exist_ok=True)
CHANNELS = ["tv", "search", "social"]
SEED = 11
pd.set_option("display.max_colwidth", None)   # estimands read in full

1. Build

A synthetic world with a known truth

Three years of weekly sales with carryover (adstock) and diminishing returns (saturation), the truth from notebook 06. Twelve more weeks stand in for next quarter; their plan stays inside the spend range of the last three years, as next quarter’s plan usually does.

In [2]:

TRUE = {"decay": {"tv": 0.5, "search": 0.0, "social": 0.3}, "slope": {"tv": 1.8, "search": 3.0, "social": 2.4}}


def adstock(x, decay):
    out = np.zeros(len(x))
    for t, v in enumerate(x):
        out[t] = v + (decay * out[t - 1] if t else 0.0)
    return out


def media_effect(spend, k):
    """The true sales from media, per week, for a spend table (channels as columns)."""
    return {c: TRUE["slope"][c] * k[c] * np.log1p(adstock(spend[c].to_numpy(float), TRUE["decay"][c]) / k[c])
            for c in CHANNELS}


def make_world(seed=SEED, weeks=156, ahead=12):
    rng = np.random.default_rng(seed)
    n = weeks + ahead
    week = pd.date_range("2023-01-02", periods=n, freq="W-MON")
    spend = pd.DataFrame({"tv": rng.gamma(4, 9000, n), "search": rng.gamma(5, 2500, n),
                          "social": rng.gamma(3, 4000, n)}).round(2)
    for c in CHANNELS:   # next quarter's plan stays inside the range already spent
        spend.loc[weeks:, c] = spend.loc[weeks:, c].clip(spend[c][:weeks].min(), spend[c][:weeks].max())
    k = {c: adstock(spend[c][:weeks].to_numpy(), TRUE["decay"][c]).mean() for c in CHANNELS}   # half-saturation
    t, promo = np.arange(n), (np.arange(n) % 9 == 3).astype(int)
    sales = 200_000 + 30_000 * np.sin(2 * np.pi * t / 52) + 15_000 * promo + sum(media_effect(spend, k).values())
    sales = sales + rng.normal(0, 8_000, n)
    world = pd.DataFrame({"week": week, "sales": sales.round(2), **spend, "promo": promo})
    return world.iloc[:weeks].reset_index(drop=True), world.iloc[weeks:].reset_index(drop=True), k


train, next_quarter, TRUE_K = make_world()
train.tail(3)

Out [2]:

	week	sales	tv	search	social	promo
153	2025-12-08	351365.21	40334.67	11121.39	19767.60	0
154	2025-12-15	312421.59	15504.88	11518.47	14465.62	0
155	2025-12-22	351304.84	27929.58	28680.99	12070.82	0

The truth for a set of weeks answers the same question the model will: how much sales would those weeks lose if the channel’s spend in them were removed, per dollar removed. Carryover from earlier weeks stays in.

In [3]:

def true_roi(history, weeks, ch):
    full = pd.concat([history, weeks], ignore_index=True)
    cut = full.copy()
    cut.loc[len(history):, ch] = 0.0
    lost = media_effect(full, TRUE_K)[ch] - media_effect(cut, TRUE_K)[ch]
    return lost[len(history):].sum() / weeks[ch].sum()


TRUTH = {c: true_roi(train.iloc[:0], train, c) for c in CHANNELS}
{c: round(v, 3) for c, v in TRUTH.items()}

Out [3]:

{'tv': np.float64(2.457),
 'search': np.float64(2.017),
 'social': np.float64(2.289)}

The input contract

The contract says what “the same necessary shape” means for this workflow: which columns must be there, their types, the weekly time column, and the target. More weeks are always welcome.

In [4]:

contract = fin.Contract.from_frame(train, time="week", target="sales", freq="W-MON")
contract

Out [4]:

Contract: 4 required, 0 optional columns, time week, target sales

ump.input_contract/0.1

details

key	week
extra columns	drop
categories	none

column	dtype	required	missing
tv	float64	True	forbid
search	float64	True	forbid
social	float64	True	forbid
promo	int64	True	forbid

The model, as a module

The code goes in a file because a bundle saves source, never objects: a function defined in a notebook cell cannot be saved, hashed or read by the colleague who has to trust it.

The module has four parts. fit builds the design and runs two hundred block-bootstrap fits in one jax.vmap, each one Adam steps on jax.grad of the weighted squared error. to_state and from_state turn the fitted jax.Arrays into NumPy and back. One small function per business question returns one value per draw.

In [5]:

%%writefile out/workflows/northwind_jax.py
"""Northwind weekly sales in JAX: adstock, saturation and a seasonal baseline, fitted by gradient descent.

Draws come from one jax.vmap over block-bootstrap weights: each draw refits every parameter on its resample.
"""
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd

jax.config.update("jax_enable_x64", True)
CHANNELS = ["tv", "search", "social"]


def _baseline(df, start):
    t = (pd.to_datetime(df["week"]) - pd.Timestamp(start)).dt.days.to_numpy() / 7.0
    return jnp.asarray(np.column_stack([np.ones_like(t), np.sin(2 * np.pi * t / 52), np.cos(2 * np.pi * t / 52),
                                        t / 52, df["promo"].to_numpy(float)]))


def _adstock(x, decay):
    def step(carry, row):
        carry = row + decay * carry
        return carry, carry
    return jax.lax.scan(step, jnp.zeros_like(decay), x)[1]


def _predict(p, spend, base):
    """Scaled sales for the last len(base) weeks; spend may start earlier, so carryover is kept."""
    decay, k = jax.nn.sigmoid(p["decay_logit"]), jnp.exp(p["log_k"])
    media = jnp.exp(p["log_slope"]) * k * jnp.log1p(_adstock(spend, decay)[-base.shape[0]:] / k)
    return base @ p["base"] + media.sum(axis=1)


def _fit_one(weights, spend, base, y, *, steps, lr):
    loss = lambda p: jnp.sum(weights * (_predict(p, spend, base) - y) ** 2) / jnp.sum(weights)  # noqa: E731
    grad = jax.grad(loss)
    p0 = {"decay_logit": jnp.zeros(3), "log_k": jnp.full(3, jnp.log(2.0)), "log_slope": jnp.full(3, jnp.log(0.1)),
          "base": jnp.zeros(base.shape[1]).at[0].set(0.6)}

    def adam(state, i):
        p, m, v = state
        g = grad(p)
        m = jax.tree.map(lambda m, g: 0.9 * m + 0.1 * g, m, g)
        v = jax.tree.map(lambda v, g: 0.999 * v + 0.001 * g * g, v, g)
        p = jax.tree.map(lambda p, m, v: p - lr * (m / (1 - 0.9 ** (i + 1)))
                         / (jnp.sqrt(v / (1 - 0.999 ** (i + 1))) + 1e-8), p, m, v)
        return (p, m, v), None

    zero = jax.tree.map(jnp.zeros_like, p0)
    return jax.lax.scan(adam, (p0, zero, zero), jnp.arange(steps))[0][0]


def _block_weights(key, n, draws, block):
    starts = jax.random.randint(key, (draws, n // block + 1), 0, n - block + 1)
    rows = (starts[:, :, None] + jnp.arange(block)).reshape(draws, -1)[:, :n]
    return jax.vmap(lambda r: jnp.bincount(r, length=n))(rows).astype(float)


def fit(df, *, seed, params):
    start = pd.to_datetime(df["week"]).min().date().isoformat()
    scale, level = df[CHANNELS].mean().to_numpy(float), float(df["sales"].mean())
    spend = jnp.asarray(df[CHANNELS].to_numpy(float) / scale)
    base, y = _baseline(df, start), jnp.asarray(df["sales"].to_numpy(float) / level)
    weights = _block_weights(jax.random.key(seed), len(df), params["draws"], params["block"])
    fit_all = jax.jit(jax.vmap(partial(_fit_one, steps=params["steps"], lr=params["lr"]), in_axes=(0, None, None, None)))
    return {"params": fit_all(weights, spend, base, y), "scale": scale, "level": level, "start": start,
            "hist_days": _days(df), "hist_spend": df[CHANNELS].to_numpy(float)}


def _days(df):
    return pd.to_datetime(df["week"]).to_numpy("datetime64[D]").astype(np.int64)


def to_state(model):
    """jax.Array to NumPy: JSON-safe meta plus plain arrays, which the bundle writes as .npy (nothing pickled)."""
    arrays = {f"p_{k}": np.asarray(v) for k, v in model["params"].items()}
    arrays.update(hist_days=model["hist_days"], hist_spend=model["hist_spend"], scale=model["scale"])
    return {"level": model["level"], "start": model["start"]}, arrays


def from_state(meta, arrays):
    """NumPy back to jax.Array: the other half of the only JAX-specific step."""
    return {"params": {k[2:]: jnp.asarray(v) for k, v in arrays.items() if k.startswith("p_")},
            "scale": np.asarray(arrays["scale"]), "level": float(meta["level"]), "start": meta["start"],
            "hist_days": np.asarray(arrays["hist_days"]), "hist_spend": np.asarray(arrays["hist_spend"])}


def sales(model, df, **spend):
    """Weekly sales per draw (draws x weeks), some channels' spend replaced; earlier weeks' carryover kept."""
    df = df.assign(**spend)
    earlier = model["hist_days"] < _days(df).min()
    raw = np.vstack([model["hist_spend"][earlier], df[CHANNELS].to_numpy(float)])
    predict = jax.vmap(_predict, in_axes=(0, None, None))
    return predict(model["params"], jnp.asarray(raw / model["scale"]), _baseline(df, model["start"])) * model["level"]


def _roi(model, df, ch):
    return (sales(model, df) - sales(model, df, **{ch: 0.0})).sum(axis=1) / df[ch].sum()


def roi_tv(model, df):
    return _roi(model, df, "tv")


def roi_search(model, df):
    return _roi(model, df, "search")


def roi_social(model, df):
    return _roi(model, df, "social")


def mroi_tv(model, df):
    return (sales(model, df, tv=df["tv"] * 1.1) - sales(model, df)).sum(axis=1) / (0.1 * df["tv"].sum())


def cut_tv_20pct(model, df):
    return (sales(model, df, tv=df["tv"] * 0.8) - sales(model, df)).sum(axis=1)

Out [5]:

Overwriting out/workflows/northwind_jax.py

The workflow and its named questions

A question is a name, the question in words, the estimand that says exactly what the number means, and the name of the module function that answers it. By default the horizon follows the rows each run receives, so an answer on twelve new weeks says twelve weeks.

The draws are bootstrap refits, so each estimand says bootstrap_pct for its interval; that makes it a variant of the registered question, so the ROI and mROI questions name the summary number they feed (summary=).

The summary the bundle carries holds one interval per channel, so it declares per_channel dependence: anyone adding them up later gets the independence-to-comonotone envelope (G-DEP-01).

In [6]:

wf = fin.Workflow("Northwind weekly sales (JAX)", contract=contract, code=WF / "northwind_jax.py", seed=SEED,
                  params={"draws": 200, "block": 8, "steps": 1500, "lr": 0.03}, dependence="per_channel")
BOOT = {"kpi": "incremental_sales", "interval": "bootstrap_pct"}   # the draws are bootstrap fits, not a posterior
for ch in CHANNELS:
    wf.question(f"roi_{ch}", ask=f"What did each dollar of {ch} return over these weeks?",
                estimand=fin.estimand("roi", channels=[ch], **BOOT), summary=("roi", ch))
wf.question("mroi_tv", ask="What would the next dollar of TV return?",
            estimand=fin.estimand("mroi", channels=["tv"], **BOOT), summary=("mroi", "tv"))
wf.question("cut_tv_20pct", ask="What happens to sales if we cut TV by 20%?",
            estimand=fin.estimand("roi", channels=["tv"], intervention="scale", size=-0.2, baseline="current_plan",
                                  per="absolute", **BOOT))
wf

Out [6]:

Northwind weekly sales (JAX)

workflow, not fitted · 5 questions

details

module	northwind_jax.py
seed	11
prep recipe	none
contract	<Contract: 4 required, 0 optional columns, time week, target sales>

question	ask
roi_tv	What did each dollar of tv return over these weeks?
roi_search	What did each dollar of search return over these weeks?
roi_social	What did each dollar of social return over these weeks?
mroi_tv	What would the next dollar of TV return?
cut_tv_20pct	What happens to sales if we cut TV by 20%?

Fit

One call: babelfin checks the training frame against the contract, imports the module and runs its fit. The first run includes JAX’s compile time; the two hundred fits themselves run as one vectorised program.

In [7]:

start = time.perf_counter()
fitted = wf.fit(train)
print(f"fitted 200 bootstrap draws in {time.perf_counter() - start:.0f} s")
fitted

Out [7]:

fitted 200 bootstrap draws in 2 s

Northwind weekly sales (JAX)

fitted on 156 rows · 5 answers

details

module	northwind_jax.py
code sha256	cef367199d4e

question	value	lo	hi	level
roi_tv	2.29	2.00	2.61	0.9
roi_search	1.76	1.48	2.15	0.9
roi_social	2.10	1.69	2.59	0.9
mroi_tv	1.69	1.52	1.91	0.9
cut_tv_20pct	-2,069,131	-2,321,702	-1,852,707	0.9

Each answer is a fin.Number: the median of the draws, a 90% bootstrap interval (the level is proposed, the default when an estimand declares none), and the estimand in words. Three ROI questions, one next-dollar mROI and one scenario: they never get confused, because each carries its own estimand.

In [8]:

answers = fitted.answers()
pd.DataFrame([{"question": n, "value": a.value, "lo": a.lo, "hi": a.hi, "estimand": a.estimand.words}
              for n, a in answers.items()]).round(2)

Out [8]:

	question	value	lo	hi	estimand
0	roi_tv	2.29	2.00	2.61	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 156 weeks, posterior median, 90% bootstrap_pct
1	roi_search	1.76	1.48	2.15	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 156 weeks, posterior median, 90% bootstrap_pct
2	roi_social	2.10	1.69	2.59	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 156 weeks, posterior median, 90% bootstrap_pct
3	mroi_tv	1.69	1.52	1.91	incremental sales per dollar, net from a small increase in spend (+10%), geo_set, versus the current plan, over 156 weeks, posterior median, 90% bootstrap_pct
4	cut_tv_20pct	-2069131.20	-2321702.49	-1852707.28	incremental sales in total, net from scaling spend (-20%), geo_set, versus the current plan, over 156 weeks, posterior median, 90% bootstrap_pct

Did it find the truth?

The world is synthetic, so the true average ROI over the three years is known (estimand: incremental sales per dollar of the channel, all 156 weeks, removal of that channel’s spend). The intervals should cover it.

In [9]:

check = pd.DataFrame([{"channel": c, "estimate": answers[f"roi_{c}"].value, "lo": answers[f"roi_{c}"].lo,
                       "hi": answers[f"roi_{c}"].hi, "truth": TRUTH[c]} for c in CHANNELS])
check["covered"] = (check.lo <= check.truth) & (check.truth <= check.hi)
check.round(3)

Out [9]:

	channel	estimate	lo	hi	truth	covered
0	tv	2.293	2.002	2.606	2.457	True
1	search	1.757	1.478	2.147	2.017	True
2	social	2.103	1.688	2.588	2.289	True

The fitted decays are ordinary arrays on the model, one per draw. TV’s and social’s intervals cover the true carryover (0.5 and 0.3). Search’s decay is a sigmoid, which can approach zero but never reach it, so its interval sits just above the true 0.0; the spread is the part of the answer a single fit would hide.

In [10]:

decay = 1 / (1 + np.exp(-np.asarray(fitted.model["params"]["decay_logit"])))
pd.DataFrame({"truth": TRUE["decay"], "median": np.median(decay, axis=0),
              "5%": np.quantile(decay, 0.05, axis=0), "95%": np.quantile(decay, 0.95, axis=0)},
             index=CHANNELS).round(2)

Out [10]:

	truth	median	5%	95%
tv	0.5	0.50	0.44	0.56
search	0.0	0.06	0.03	0.12
social	0.3	0.24	0.15	0.33

2. Save and use

save writes a folder: the module’s source, the fitted state as JSON plus .npy arrays, the contract, the questions and a Tier D summary of the answers, every file hashed. Before writing, it checks that from_state(to_state(model)) gives the same draws. That conversion is the only JAX-specific step: to_state calls numpy.asarray on each jax.Array, from_state calls jax.numpy.asarray on the way back, and nothing is pickled.

In [11]:

bundle = WF / "northwind_jax.bundle"
if bundle.exists():
    import shutil
    shutil.rmtree(bundle)        # a bundle is written once; rerunning the notebook writes it again
path = fitted.save(bundle)
sorted(p.name for p in path.rglob("*") if p.is_file())

Out [11]:

['README.md',
 'SHA256SUMS',
 'contract.json',
 'hist_days.npy',
 'hist_spend.npy',
 'northwind_jax.py',
 'p_base.npy',
 'p_decay_logit.npy',
 'p_log_k.npy',
 'p_log_slope.npy',
 'questions.json',
 'scale.npy',
 'state.json',
 'summary.json',
 'workflow.json']

A new machine

Sam opens the folder on his laptop. Saving trusted the code for Priya’s own session only, so to play Sam in the same kernel we point finprint’s trust store at a fresh folder and forget this session’s trust.

In [12]:

def new_machine():
    from finprint import trust
    os.environ["FINPRINT_HOME"] = tempfile.mkdtemp(prefix="finprint-")
    trust._SESSION.clear()      # the in-memory trust of the author's own session


new_machine()
lw = fin.Workflow.load(bundle)
print("trusted:", lw.trusted)
print(lw.code[:400])

Out [12]:

trusted: False
"""Northwind weekly sales in JAX: adstock, saturation and a seasonal baseline, fitted by gradient descent.

Draws come from one jax.vmap over block-bootstrap weights: each draw refits every parameter on its resample.
"""
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd

jax.config.update("jax_enable_x64", True)
CHANNELS = ["tv", "search", "so

Loading checked every hash and parsed the data, but imported nothing. Until Sam trusts the code, a run answers from the bundle’s summary, and says so.

In [13]:

untrusted = lw.run(next_quarter)
print(untrusted.route, "/", untrusted.source)
untrusted.notes

Out [13]:

not_available / summary

[{'code': 'WORKFLOW_CODE_UNTRUSTED',
  'severity': 'warn',
  'column': None,
  'message': 'code code/northwind_jax.py is not trusted on this machine, so it was not run'},
 {'code': 'SUMMARY_FALLBACK',
  'severity': 'warn',
  'column': None,
  'message': "answers are the bundle's fit-time Tier D summary because the code is not trusted; no summary number for cut_tv_20pct (not available)"}]

Sam reads the module (above, and in full in lw.code), then trusts its hash on his machine.

In [14]:

lw.trust()
print("trusted:", lw.trusted, "· code sha256", lw.code_sha256[:12])

Out [14]:

trusted: True · code sha256 cef367199d4e

Score twelve new weeks

Next quarter’s twelve weeks follow straight on from the training weeks, so the saved fit scores them (score by default): no refit, the carryover from the last training weeks included.

In [15]:

run = lw.run(next_quarter)
print(run.route, "/", run.source, "·", run.rows, "rows")
new = run.answers()
pd.DataFrame([{"question": n, "value": a.value, "lo": a.lo, "hi": a.hi, "estimand": a.estimand.words}
              for n, a in new.items()]).round(2)

Out [15]:

score / score · 12 rows

	question	value	lo	hi	estimand
0	roi_tv	2.11	1.84	2.36	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 12 weeks, posterior median, 90% bootstrap_pct
1	roi_search	1.68	1.41	2.02	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 12 weeks, posterior median, 90% bootstrap_pct
2	roi_social	2.09	1.69	2.56	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 12 weeks, posterior median, 90% bootstrap_pct
3	mroi_tv	1.60	1.44	1.76	incremental sales per dollar, net from a small increase in spend (+10%), geo_set, versus the current plan, over 12 weeks, posterior median, 90% bootstrap_pct
4	cut_tv_20pct	-134230.50	-147098.72	-121393.65	incremental sales in total, net from scaling spend (-20%), geo_set, versus the current plan, over 12 weeks, posterior median, 90% bootstrap_pct

The truth for those twelve weeks (incremental sales per dollar of the channel over weeks 157 to 168, carryover from earlier weeks kept), beside the scored answers:

In [16]:

pd.DataFrame([{"channel": c, "estimate": new[f"roi_{c}"].value, "lo": new[f"roi_{c}"].lo, "hi": new[f"roi_{c}"].hi,
               "truth": true_roi(train, next_quarter, c)} for c in CHANNELS]).round(3)

Out [16]:

	channel	estimate	lo	hi	truth
0	tv	2.106	1.838	2.363	2.256
1	search	1.678	1.409	2.015	1.884
2	social	2.092	1.693	2.559	2.283

When the new data is different

A gap. Li exports the whole history plus the new quarter, and the feed loses two weeks in the middle of it. The export overlaps the saved history and has a hole, and scoring across a hole would give carryover the wrong past. So the run refits on the data it was given (with the saved seed and settings) and labels every answer as a fallback.

In [17]:

gappy = pd.concat([train, next_quarter.drop(index=[5, 6])], ignore_index=True)
gap_run = lw.run(gappy)
print(gap_run.route, "/", gap_run.source)
pd.DataFrame(gap_run.notes)[["code", "severity", "message"]]

Out [17]:

refit / refit

	code	severity	message
0	INPUT_IRREGULAR_TIME	warn	time column week is irregular (2 missing steps); scoring needs regular steps, so the rerun refits
1	INPUT_HISTORY_OVERLAP	warn	new input starts 2023-01-02, on or before the history end 2025-12-22; the rerun refits
2	HISTORY_ROUTED_TO_REFIT	warn	the input does not follow on from the saved history, so the saved state is not scored; the rerun refits
3	INPUT_ROWS_GREW	info	166 rows, 156 at training

In [18]:

gap_run.answers()["roi_tv"]

Out [18]:

roi_tv: 2.34

likely between 2.09 and 2.58 · medium confidence

details

means	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 166 weeks, posterior median, 90% bootstrap_pct
interval	90% bootstrap_pct 2.09 to 2.58
trust level	2
G-DATA-07 INPUT_IRREGULAR_TIME	time column week is irregular (2 missing steps); scoring needs regular steps, so the rerun refits; answer refitted on this input with the saved seed and params, not the saved state
G-DATA-07 INPUT_HISTORY_OVERLAP	new input starts 2023-01-02, on or before the history end 2025-12-22; the rerun refits; answer refitted on this input with the saved seed and params, not the saved state
G-DATA-07 HISTORY_ROUTED_TO_REFIT	the input does not follow on from the saved history, so the saved state is not scored; the rerun refits; answer refitted on this input with the saved seed and params, not the saved state

An extra column. Someone adds weather to the export. Extra columns never fail: it is dropped, with a note, and the run still scores.

In [19]:

extra = lw.run(next_quarter.assign(temperature_c=12.0))
print(extra.route, "/", extra.source)
[(n["code"], n["severity"]) for n in extra.notes]

Out [19]:

score / score

[('INPUT_EXTRA_DROPPED', 'info')]

A missing required column. The social feed is not in this export. The model cannot run without it, so the run falls back to the bundle’s labelled summary: the fit-time numbers, at trust level 1, with the reason on each.

In [20]:

missing = lw.run(next_quarter.drop(columns="social"))
print(missing.route, "/", missing.source)
missing.answers()["roi_tv"]

Out [20]:

not_available / summary

roi_tv: 2.29

likely between 2.00 and 2.61 · low confidence

details

means	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 156 weeks, posterior median, 90% bootstrap_pct
interval	90% bootstrap_pct 2.00 to 2.61
trust level	1
G-DATA-07 INPUT_REQUIRED_MISSING	required column social is missing; answer is the bundle's fit-time Tier D summary, not a run on this input
G-DATA-07 SUMMARY_FALLBACK	answers are the bundle's fit-time Tier D summary because the input does not meet the contract; no summary number for cut_tv_20pct (not available); answer is the bundle's fit-time Tier D summary, not a run on this input

inspect() shows what a bundle holds without trusting or running anything: its manifest, contract, questions and the versions it was saved with, JAX among them.

In [21]:

info = lw.inspect()
print(sorted(info))
info["env_pin"]

Out [21]:

['author', 'code_sha256', 'contract', 'created_at', 'default_mode', 'entry', 'env', 'env_pin', 'format', 'has_data', 'has_recipe', 'name', 'notes', 'questions', 'recipe_steps', 'tier', 'training', 'trusted', 'version']

{'python': '3.12.8',
 'numpy': '2.5.4',
 'pandas': '3.1.0',
 'finprint': '0.2.0',
 'ump_spec': '0.1.0',
 'jax': '0.12.0.1'}

3. Share and serve (recorded session)

A folder on Sam’s laptop helps Sam. To let anyone at Northwind run the workflow, Priya publishes it through UMP core as a Tier A model manifest; a second person must approve it before anyone else’s machine runs its code. These cells use the labelled test client northwind-e2e-live, never the demo client.

Recorded run

This section shows a recorded session against the reference deployment. It cannot be run without that deployment; the offline sections above run anywhere. The identity stub (as_=) and the localhost address belong to the reference deployment, not to the standard; a conforming core authenticates callers as section 19 requires.

In [22]:

CLIENT = "northwind-e2e-live"
NEEDS_POC = "Section 3 needs the running POC (scripts/dev.sh); skipped."


def core_up():
    try:
        return httpx.get(base_url("core") + "/healthz", timeout=2).status_code == 200
    except httpx.HTTPError:
        return False


LIVE = core_up()
print("core is up" if LIVE else NEEDS_POC)

Out [22]:

core is up

Priya publishes from the POC’s analyst login, maya.chen, and submits it for review.

In [23]:

if LIVE:
    with fin.connect(as_="maya.chen") as s:
        art = s.workflows.publish(fitted, client=CLIENT, title="Northwind weekly sales, JAX (notebook 09, synthetic)")
        art.submit(reason="notebook 09: a JAX workflow to share")
        print(art.id, art.state, "· tier", art.doc["tier"], "· engine", art.doc["engine"])
else:
    print(NEEDS_POC)

Out [23]:

mdl_704ac6aee7b34146b89f awaiting_approval · tier A · engine ump-workflow

She cannot approve her own work (G-INT-02). Sam can, and must name each fallback he accepts.

In [24]:

if LIVE:
    with fin.connect(as_="maya.chen") as s:
        own = s.artifact(art.id)
        try:
            own.approve(accept=own.fallbacks)
        except fin.Forbidden as exc:
            print("refused:", exc)
        with s.as_("sam.patel") as sam:
            review = sam.artifact(art.id)
            print("Sam accepts:", [f.code for f in review.fallbacks])
            review.approve(accept=review.fallbacks, comment="read northwind_jax.py; bootstrap draws, no pickles")
        print("state:", s.artifact(art.id).state)
else:
    print(NEEDS_POC)

Out [24]:

refused: SELF_APPROVAL_FORBIDDEN: an approver cannot approve their own work (owner or submitter)
Sam accepts: ['UNVALIDATED_METHOD']
state: validated

Li runs it with no manual trust

Li has never seen this code and her trust store is empty. get fetches every file, checks each hash against the approved manifest and trusts exactly what Sam’s approval vouches for.

In [25]:

if LIVE:
    new_machine()
    with fin.connect(as_="li.wong") as li:
        li_wf = li.workflows.get(art.id)
    li_run = li_wf.score(next_quarter)
    print("trusted:", li_wf.trusted, "· reason:", li_wf.approval.reason, "·", li_run.route, "/", li_run.source)
    display(li_run.answers()["roi_tv"])
else:
    print(NEEDS_POC)

Out [25]:

trusted: True · reason: None · score / score

roi_tv: 2.11

likely between 1.84 and 2.36 · medium confidence

details

means	incremental sales per dollar, net from removing all spend, geo_set, versus zero spend, over 12 weeks, posterior median, 90% bootstrap_pct
interval	90% bootstrap_pct 1.84 to 2.36
trust level	2

Serve it as an app

One served app per bundle: ump-workflow serve fetches the approved manifest, rebuilds the folder and answers questions over UMP’s adapter contract. The server has its own empty trust store, so the approval is again the only reason it runs the code. Any notebook or app can then ask it questions.

In [26]:

if LIVE:
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
    scratch = Path(tempfile.mkdtemp(prefix="nb09-serve-"))
    server = subprocess.Popen(
        [sys.executable, "-m", "ump_workflow", "serve", "--artifact", art.id, "--principal", "li.wong",
         "--port", str(port), "--dest", str(scratch / "served")],
        env={**os.environ, "FINPRINT_HOME": str(scratch / "server_home")},
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    url = f"http://127.0.0.1:{port}"
    for _ in range(120):
        try:
            if httpx.get(url + "/healthz", timeout=1).status_code == 200:
                break
        except httpx.HTTPError:
            time.sleep(0.5)
    print("serving on a free local port:", server.poll() is None)
else:
    print(NEEDS_POC)

Out [26]:

serving on a free local port: True

In [27]:

if LIVE:
    with fin.connect(as_="li.wong") as li:
        served = li.workflows.answer(url, next_quarter, questions=["roi_tv", "mroi_tv"])
    print(served.route, "/", served.source)
    display(pd.DataFrame([{"question": n, "value": a.value, "lo": a.lo, "hi": a.hi}
                          for n, a in served.answers().items()]).round(2))
else:
    print(NEEDS_POC)

Out [27]:

score / score

	question	value	lo	hi
0	roi_tv	2.11	1.84	2.36
1	mroi_tv	1.60	1.44	1.76

Stop the server, and move the published manifest to deprecated: a validated workflow is listed as a Tier A model in the models app, and a notebook demo should not linger there.

In [28]:

if LIVE:
    server.terminate()
    server.wait(timeout=15)
    with fin.connect(as_="maya.chen") as s:
        s.core.transition(art.id, "deprecated", reason="notebook 09 cleanup (synthetic shared workflow)")
        print("server stopped:", server.returncode is not None, "· manifest:", s.artifact(art.id).state)
else:
    print(NEEDS_POC)

Out [28]:

server stopped: True · manifest: deprecated

What changed for JAX

Very little. The module defines fit, to_state, from_state and one function per question, as in the scikit-learn and plain-Python notebooks. The one JAX-specific step is the state: numpy.asarray on save, jax.numpy.asarray on load, JSON and .npy on disk, nothing pickled. Everything else (the contract, score by default, the labelled fallbacks, the second-person approval, the served app) is the same for every toolkit.