Beam decision models

These examples use the reusable decision primitives in apache_beam.ml.inference.decision. A model evaluates named, typed questions for each Beam element. EvaluateDecisions returns the original state, normalized answers, provider metadata, and request latency. Policy and sinks remain ordinary Beam transforms after evaluation.

Public API

The core module contains ChoiceQuestion, BooleanQuestion, and ScoreQuestion, their typed answers (ChoiceAnswer, BooleanAnswer, and ScoreAnswer), the DecisionModel ABC, and the EvaluateDecisions PTransform.

JevDecisionModel maps BooleanQuestion to Jev's Noul, returning BooleanAnswer.probability; a missing confidence is None.

DecisionModel is the adapter ABC. A subclass implements evaluate(state, questions) and returns one matching typed answer per named question in a DecisionResponse. Override __enter__ and __exit__ to reuse a client, connection pool, or model weights. The transform enters the model during DoFn setup and closes it during teardown for each DoFn instance.

EvaluateDecisions(model, questions) is a PTransform from input states to DecisionResult. Evaluation is synchronous, timestamps and windows are preserved, exceptions propagate to the runner, and request, failure, and latency metrics are recorded.

Before returning a result, JevDecisionModel checks that its response contains the requested questions with matching answer types. It rejects unknown choice labels and non-finite or out-of-range scores, probabilities, and confidence values with ValueError. Custom DecisionModel implementations must enforce the same output contract. These errors fail the pipeline by default.

Architecture

The optional Jev adapter lives in apache_beam.ml.inference.typesafe_inference; the deterministic LocalDecisionModel stays in apache_beam.examples.inference.decision_models.local_model. Both implement the same core contract.

Integration

Questions and events are regular Python values. The model is the only provider-specific argument to EvaluateDecisions:

import apache_beam as beam

from apache_beam.ml.inference.decision import ChoiceQuestion
from apache_beam.ml.inference.decision import EvaluateDecisions
from apache_beam.examples.inference.decision_models.local_model import (
    LocalDecisionModel,
)
from apache_beam.ml.inference.typesafe_inference import JevDecisionModel

questions = {
    'team': ChoiceQuestion('Which team should handle this message?', {
        'billing': 'Invoices', 'technical': 'Technical problems'}),
}

def select_team(result):
  return result.response.answers['team'].choice

model = LocalDecisionModel()
# model = JevDecisionModel()

with beam.Pipeline() as pipeline:
  events = pipeline | 'Events' >> beam.Create(
      [{'message': 'The API returns 503 after I rotated my key.'}])
  decisions = events | 'Evaluate decisions' >> EvaluateDecisions(
      model, questions)
  rows = decisions | 'Apply downstream policy' >> beam.Map(select_team)

The downstream policy stays unchanged when swapping adapters. The Jev client is optional; install the shared SDK requirements listed below and set TYPESAFE_API_KEY when running it.

Install

From a Beam checkout, activate a virtual environment and install the core SDK:

cd sdks/python
python3 -m pip install -e .

Install the optional Jev SDK used by the adapter and examples:

cd sdks/python
python3 -m pip install -r \
  apache_beam/ml/inference/typesafe_tests_requirements.txt

Install Beam's GCP extra for Pub/Sub or BigQuery sinks:

cd sdks/python
python3 -m pip install -e '.[gcp]'

The example requirements file delegates to that same bounded dependency file.

Run examples

The main demo uses a three-element TestStream and writes one JSONL file per destination:

cd sdks/python
output_dir="$(mktemp -d /tmp/beam-decision-model-output.XXXXXX)"
python3 -m apache_beam.examples.inference.decision_models.message_router \
  --model=local \
  --output-dir="$output_dir" \
  --graph="$output_dir/message_router.svg"
find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;

Render the graph without running the job. The sink stays in the graph:

cd sdks/python
python3 -m apache_beam.examples.inference.decision_models.message_router \
  --model=jev \
  --graph-only \
  --bq-dataset=PROJECT:DATASET \
  --graph=/tmp/beam-decision-model.svg

The .dot export works without Graphviz; SVG and PNG require dot on PATH.

Run the Jev path against the same stream with TYPESAFE_API_KEY set:

cd sdks/python
export TYPESAFE_API_KEY='replace-with-your-key'
jev_output_dir="$(mktemp -d /tmp/beam-decision-model-jev-output.XXXXXX)"
python3 -m apache_beam.examples.inference.decision_models.message_router \
  --model=jev \
  --min-confidence=0.65 \
  --output-dir="$jev_output_dir" \
  --graph="$jev_output_dir/message_router.svg"

For a streaming source, provide a Pub/Sub subscription containing JSON objects with event_id and message fields:

cd sdks/python
python3 -m apache_beam.examples.inference.decision_models.message_router \
  --model=jev \
  --pubsub-subscription=projects/PROJECT/subscriptions/SUBSCRIPTION \
  --bq-dataset=PROJECT:DATASET

With --bq-dataset=PROJECT:DATASET, rows go to messages_billing, messages_technical, messages_sales, or messages_review; --output-dir writes local JSONL files. A normal run accepts one sink option.

The fraud example prints one JSON row per sample message. Noul supplies the boolean probability; values at or above 0.7 are flagged:

cd sdks/python
python3 -m apache_beam.examples.inference.decision_models.fraud_review \
  --model=local --primitive=noul
python3 -m apache_beam.examples.inference.decision_models.fraud_review \
  --model=local --primitive=both
python3 -m apache_beam.examples.inference.decision_models.fraud_review \
  --model=jev --primitive=both

Pass --dead-letter-dir to write per-element ValueError failures from decision evaluation as JSONL records containing the original input and the exception representation. This opt-in catches malformed-response validation errors; without the option, they fail the pipeline. The Jev SDK retries transient transport and service failures with exponential backoff and jitter, and honors Retry-After when supplied. Those retries run independently on each Beam worker, which does not impose a global request-rate limit. Failures that remain after retries still fail the pipeline instead of going to this sink. Errors while the model context is created during worker setup also fail the pipeline and are not routed to this sink.

For example, enable the local dead-letter sink with the Jev adapter:

cd sdks/python
export TYPESAFE_API_KEY='replace-with-your-key'
dead_letter_dir="$(mktemp -d /tmp/beam-decision-model-dead-letter.XXXXXX)"
python3 -m apache_beam.examples.inference.decision_models.fraud_review \
  --model=jev --primitive=both --dead-letter-dir="$dead_letter_dir"
find "$dead_letter_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;

A final three-element DirectRunner run returned these values. Beam may print the rows in a different order:

Only the bypass verification message crossed the 0.7 review threshold.

Tests

Install Beam's test dependencies in the same virtual environment:

cd sdks/python
python3 -m pip install -e '.[test]'

The unit tests cover the shared API, routing policy, and mocked Jev responses:

cd sdks/python
python3 -m pytest -q \
  apache_beam/ml/inference/decision_test.py \
  apache_beam/ml/inference/typesafe_inference_test.py \
  apache_beam/examples/inference/decision_models/message_router_test.py \
  apache_beam/examples/inference/decision_models/fraud_review_test.py

With the optional Jev SDK installed and TYPESAFE_API_KEY set, run the live integration test:

cd sdks/python
python3 -m pytest -q -rs -m it_postcommit \
  apache_beam/ml/inference/typesafe_inference_it_test.py

It sends one streaming event through EvaluateDecisions, asking Choice, Boolean/Noul, and Score questions together. It checks the normalized answer types and ranges, and skips when the SDK or API key is missing. The mocked adapter tests run with the SDK installed and require no API key.

Jev benchmark

These measurements were captured on September 22, 2026 with Jev 1.13.0, TypeSafe SDK 0.7.1, and DirectRunner. The final trace used exactly one SDK client, one TCP connection, and one TLS handshake for three events:

Total pipeline time was 986.68 ms, including pipeline construction.

Standalone transport experiment, September 22, 2026

An earlier standalone experiment compared fresh clients with one pooled client:

The trace showed no TCP/TLS spans on reused calls. Reuse pays connection setup once. Warm-call delay is response wait plus network and service time.