diff --git a/CHANGES.md b/CHANGES.md index fda4c9edc218..014fb962d0d6 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ * (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)). * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)). +* (Python) Added typed decision questions and the `EvaluateDecisions` transform, with an optional Jev model adapter ([#40220](https://github.com/apache/beam/pull/40220)). * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)). * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md index c43ad4583479..f14b531c21a2 100644 --- a/sdks/python/apache_beam/examples/inference/README.md +++ b/sdks/python/apache_beam/examples/inference/README.md @@ -17,7 +17,13 @@ under the License. --> -# Example RunInference API pipelines +# Example inference pipelines + +For typed classification, Boolean judgments, and scoring, see the +[decision model examples](decision_models/README.md). They use +`EvaluateDecisions` with a swappable model adapter, including Jev and local rules. + +## RunInference This module contains example pipelines that use the Beam RunInference API. diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md new file mode 100644 index 000000000000..637b8159c555 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -0,0 +1,304 @@ + + +# 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`. + +| Question | Normalized answer | +| --- | --- | +| `ChoiceQuestion(instructions, criteria)` | `ChoiceAnswer.choice`, with optional label probabilities and confidence | +| `BooleanQuestion(instructions)` | `BooleanAnswer.probability` in `[0, 1]` | +| `ScoreQuestion(instructions, criteria)` | `ScoreAnswer.score` and `ScoreAnswer.probabilities`, keyed by zero-based integer level, with optional confidence | + +`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 + +```mermaid +flowchart LR + input["PCollection[state]"] --> evaluate["EvaluateDecisions(model, questions)
PTransform"] + model["DecisionModel adapter"] --> evaluate + evaluate --> result["PCollection[DecisionResult]"] + result --> policy["route_message
policy and sink selection"] + policy --> sink["JSONL files
or BigQuery tables"] +``` + +![Rendered Beam message router graph](message_router.svg) + +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`: + +```python +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: + +```sh +cd sdks/python +python3 -m pip install -e . +``` + +Install the optional Jev SDK used by the adapter and examples: + +```sh +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: + +```sh +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: + +```sh +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: + +```sh +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: + +```sh +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: + +```sh +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: + +```sh +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: + +```sh +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: + +| message | Noul probability | Score | +| --- | ---: | ---: | +| billing address | 0.42 | 0.80 | +| duplicate charge | 0.32 | 0.85 | +| bypass verification | 0.97 | 3.00 | + +Only the bypass verification message crossed the `0.7` review threshold. + +## Tests + +Install Beam's test dependencies in the same virtual environment: + +```sh +cd sdks/python +python3 -m pip install -e '.[test]' +``` + +The unit tests cover the shared API, routing policy, and mocked Jev responses: + +```sh +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: + +```sh +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: + +| event | request latency | +| --- | ---: | +| billing | 408.27 ms | +| technical | 217.50 ms | +| sales | 190.43 ms | + +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: + +| client setup | observations | +| --- | --- | +| fresh client, `n=3` | 602.44 ms, 433.56 ms, 537.60 ms | +| one pooled client | first 416.77 ms; next five 170.98 ms, 211.24 ms, 222.02 ms, 240.40 ms, 226.82 ms | + +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. diff --git a/sdks/python/apache_beam/examples/inference/decision_models/__init__.py b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py new file mode 100644 index 000000000000..fc00cf8225fd --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py @@ -0,0 +1,18 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Decision model routing examples.""" diff --git a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py new file mode 100644 index 000000000000..5afcb6f34f2d --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py @@ -0,0 +1,126 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Small Beam example: flag messages for fraud review with Noul or Score.""" + +import argparse +import json +import os + +import apache_beam as beam +from apache_beam.examples.inference.decision_models.local_model import ( + LocalDecisionModel) +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import EvaluateDecisions +from apache_beam.ml.inference.decision import ScoreQuestion +from apache_beam.ml.inference.typesafe_inference import JevDecisionModel +from apache_beam.options.pipeline_options import PipelineOptions + +SAMPLE_MESSAGES = [ + 'Please update my billing address before the next renewal.', + 'I was charged twice; please look into the duplicate payment.', + 'Transfer the refund to this new account immediately and ' + 'skip verification.', +] + +FRAUD_NOUL = BooleanQuestion( + instructions='Does this message contain signals that warrant fraud review?') +FRAUD_SCORE = ScoreQuestion( + instructions='How strongly does this message warrant fraud review?', + criteria=( + 'No visible fraud cue', + 'A weak cue that merits routine checking', + 'A concrete suspicious instruction', + 'An explicit attempt to bypass verification')) + + +def review_fraud_cues(result): + response = result.response + row = { + 'message': result.state['message'], + 'model': response.model, + 'provider': response.provider, + 'latency_ms': round(result.latency_ms, 2), + } + if 'fraud_cue' in response.answers: + probability = response.answers['fraud_cue'].probability + row['fraud_review_probability'] = probability + row['needs_review'] = probability >= 0.7 + if 'risk_level' in response.answers: + score = response.answers['risk_level'] + row['risk_score'] = score.score + row['risk_confidence'] = score.confidence + row['risk_probabilities'] = score.probabilities + return row + + +def dead_letter_record(dead_letter): + """Keep the input and Beam's exception representation for a failed row.""" + original_input, exception_info = dead_letter + _, error, _ = exception_info + return {'input': original_input, 'error': error} + + +def run(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model', choices=('jev', 'local'), default='jev') + parser.add_argument( + '--primitive', choices=('noul', 'score', 'both'), default='noul') + parser.add_argument( + '--dead-letter-dir', + help=( + 'Optional local JSONL output directory for per-element ' + 'ValueError failures during decision evaluation.')) + args, pipeline_args = parser.parse_known_args(argv) + if args.model == 'jev' and not os.environ.get('TYPESAFE_API_KEY'): + parser.error('TYPESAFE_API_KEY is required for --model=jev') + model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel() + questions = {} + if args.primitive in ('noul', 'both'): + questions['fraud_cue'] = FRAUD_NOUL + if args.primitive in ('score', 'both'): + questions['risk_level'] = FRAUD_SCORE + with beam.Pipeline(options=PipelineOptions(pipeline_args)) as pipeline: + states = pipeline | 'Sample payment messages' >> beam.Create([{ + 'message': message + } for message in SAMPLE_MESSAGES]) + evaluation = EvaluateDecisions(model, questions) + if args.dead_letter_dir: + good, bad = ( + states + | 'Evaluate fraud cues' >> evaluation.with_exception_handling( + exc_class=ValueError)) + ( + bad + | 'Format malformed responses' >> beam.Map(dead_letter_record) + | 'Encode malformed responses' >> + beam.Map(lambda row: json.dumps(row, sort_keys=True)) + | 'Write malformed responses' >> beam.io.WriteToText( + os.path.join(args.dead_letter_dir, 'malformed'), + file_name_suffix='.jsonl')) + decisions = good + else: + decisions = states | 'Evaluate fraud cues' >> evaluation + ( + decisions + | 'Apply review threshold' >> beam.Map(review_fraud_cues) + | 'Show review decisions' >> + beam.Map(lambda row: print(json.dumps(row, sort_keys=True)))) + + +if __name__ == '__main__': + run() diff --git a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py new file mode 100644 index 000000000000..ec6f8160ea6e --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py @@ -0,0 +1,86 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for the fraud example's opt-in malformed-response dead letter.""" + +import glob +import json +import os +import tempfile +import unittest + +import apache_beam as beam +from apache_beam.examples.inference.decision_models.fraud_review import ( + dead_letter_record) +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import DecisionModel +from apache_beam.ml.inference.decision import DecisionResponse +from apache_beam.ml.inference.decision import EvaluateDecisions +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +_QUESTION = {'fraud_cue': BooleanQuestion('Does this need fraud review?')} +_VALID_INPUT = {'message': 'Looks like an ordinary payment question.'} +_MALFORMED_INPUT = {'message': 'Provider returned an invalid probability.'} + + +class _MalformedResponseModel(DecisionModel): + def evaluate(self, state, questions): + if state == _MALFORMED_INPUT: + raise ValueError('Invalid response probability: 1.2') + return DecisionResponse( + answers={'fraud_cue': BooleanAnswer(0.2)}, + model='fixed', + provider='test') + + +class FraudReviewDeadLetterTest(unittest.TestCase): + def test_malformed_response_is_routed_with_original_input(self): + with tempfile.TemporaryDirectory() as output_dir: + with TestPipeline() as pipeline: + good, bad = ( + pipeline + | beam.Create([_VALID_INPUT, _MALFORMED_INPUT]) + | EvaluateDecisions(_MalformedResponseModel(), _QUESTION). + with_exception_handling(exc_class=ValueError)) + + assert_that( + good | beam.Map(lambda result: result.state), + equal_to([_VALID_INPUT])) + ( + bad + | beam.Map(dead_letter_record) + | beam.Map(lambda row: json.dumps(row, sort_keys=True)) + | beam.io.WriteToText( + os.path.join(output_dir, 'malformed'), + file_name_suffix='.jsonl', + num_shards=1)) + + output_files = glob.glob(os.path.join(output_dir, 'malformed-*.jsonl')) + self.assertEqual(1, len(output_files)) + with open(output_files[0], encoding='utf-8') as output_file: + records = [json.loads(line) for line in output_file] + self.assertEqual(1, len(records)) + self.assertEqual(_MALFORMED_INPUT, records[0]['input']) + self.assertTrue(records[0]['error'].startswith('ValueError(')) + self.assertIn('1.2', records[0]['error']) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/examples/inference/decision_models/local_model.py b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py new file mode 100644 index 000000000000..55c30216f2c2 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py @@ -0,0 +1,78 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Deterministic rules for the decision model examples.""" + +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import DecisionModel +from apache_beam.ml.inference.decision import DecisionResponse +from apache_beam.ml.inference.decision import ScoreAnswer +from apache_beam.ml.inference.decision import ScoreQuestion + +_ROUTE_TERMS = { + 'billing': ('invoice', 'charge', 'refund', 'payment', 'subscription'), + 'technical': ('api', '503', 'error', 'integration', 'outage'), + 'sales': ('quote', 'seat', 'pricing', 'price', 'plan'), +} + + +class LocalDecisionModel(DecisionModel): + """Transparent rules for running these examples without network access.""" + def evaluate(self, state, questions): + message = state['message'].lower() + fraud_cue = ( + 'skip verification' in message or 'bypass verification' in message or + 'unauthorized' in message or 'phishing' in message) + answers = {} + for name, question in questions.items(): + if isinstance(question, ChoiceQuestion): + matches = { + label: sum(term in message for term in _ROUTE_TERMS.get(label, ())) + for label in question.criteria + } + if not matches: + raise ValueError('Choice needs at least one option') + best = max(matches.values()) + winners = [label for label, count in matches.items() if count == best] + probabilities = { + label: 1 / len(winners) if label in winners else 0.0 + for label in matches + } + answers[name] = ChoiceAnswer( + choice=winners[0], + probabilities=probabilities, + confidence=0.0 if best == 0 else 1 / len(winners)) + elif isinstance(question, BooleanQuestion): + answers[name] = BooleanAnswer(probability=float(fraud_cue)) + elif isinstance(question, ScoreQuestion): + if not question.criteria: + raise ValueError('Score needs at least one level') + score = min(3, len(question.criteria) - 1) if fraud_cue else 0 + answers[name] = ScoreAnswer( + score=float(score), + probabilities={ + level: float(level == score) + for level in range(len(question.criteria)) + }, + confidence=1.0) + else: + raise TypeError('Unsupported decision question: %r' % type(question)) + return DecisionResponse( + answers=answers, model='local-rules', provider='local') diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router.py b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py new file mode 100644 index 000000000000..8413d1e0b0ef --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py @@ -0,0 +1,199 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Route a short stream of messages to local files or BigQuery with Choice.""" + +import argparse +import json +import os + +import apache_beam as beam +from apache_beam.examples.inference.decision_models.local_model import LocalDecisionModel +from apache_beam.io import fileio +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import EvaluateDecisions +from apache_beam.ml.inference.typesafe_inference import JevDecisionModel +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.runners import render +from apache_beam.testing.test_stream import TestStream +from apache_beam.transforms.window import FixedWindows +from apache_beam.transforms.window import TimestampedValue + +DESTINATIONS = { + 'billing': 'Invoices, charges, payments, and refunds', + 'technical': 'Bugs, outages, and integration problems', + 'sales': 'Pricing, plans, and buying questions', +} + +QUESTION = ChoiceQuestion( + instructions='Which team should handle this customer message?', + criteria=DESTINATIONS) + +SAMPLE_MESSAGES = [ + { + 'event_id': 'msg-001', + 'message': 'My invoice has two charges for the same subscription. ' + 'Please refund one.' + }, + { + 'event_id': 'msg-002', + 'message': 'The API returns 503 for every request after I rotated ' + 'my key.' + }, + { + 'event_id': 'msg-003', + 'message': 'Can I get a quote for 50 seats next month?' + }, +] + +ROW_SCHEMA = ( + 'event_id:STRING,message:STRING,decision:STRING,destination:STRING,' + 'confidence:FLOAT,probabilities_json:STRING,model:STRING,' + 'provider:STRING,request_id:STRING,latency_ms:FLOAT') + + +def route_message(result, min_confidence): + event = result.state + response = result.response + answer = response.answers['destination'] + if answer.choice not in DESTINATIONS: + raise ValueError('Decision model returned an unknown destination') + destination = ( + answer.choice if answer.confidence is not None and + answer.confidence >= min_confidence else 'review') + return { + 'event_id': event['event_id'], + 'message': event['message'], + 'decision': answer.choice, + 'destination': destination, + 'confidence': answer.confidence, + 'probabilities_json': json.dumps(answer.probabilities, sort_keys=True), + 'model': response.model, + 'provider': response.provider, + 'request_id': response.request_id, + 'latency_ms': round(result.latency_ms, 2), + } + + +class JsonSink(fileio.FileSink): + def open(self, fh): + self.fh = fh + + def write(self, row): + self.fh.write(json.dumps(row, sort_keys=True).encode('utf-8') + b'\n') + + def flush(self): + self.fh.flush() + + +def decode_event(payload): + event = json.loads(payload.decode('utf-8')) + if not isinstance(event, dict) or not isinstance(event.get('message'), str): + raise ValueError('Pub/Sub event must be JSON with a message string') + if not isinstance(event.get('event_id'), str): + raise ValueError('Pub/Sub event must have an event_id string') + return event + + +def table_for(row, dataset): + return '%s.messages_%s' % (dataset, row['destination']) + + +def build_pipeline(pipeline, args): + if args.pubsub_subscription: + events = ( + pipeline + | 'Read Pub/Sub' >> + beam.io.ReadFromPubSub(subscription=args.pubsub_subscription) + | 'Decode events' >> beam.Map(decode_event)) + else: + stream = TestStream().add_elements([ + TimestampedValue(message, index) + for index, message in enumerate(SAMPLE_MESSAGES) + ]).advance_watermark_to_infinity() + events = pipeline | 'Sample stream' >> stream + + model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel() + rows = ( + events + | 'Ask Choice' >> EvaluateDecisions(model, {'destination': QUESTION}) + | 'Select destination' >> beam.Map(route_message, args.min_confidence) + | 'Window routed messages' >> beam.WindowInto(FixedWindows(60))) + + if args.bq_dataset: + rows | 'Write routed BigQuery tables' >> beam.io.WriteToBigQuery( + table=lambda row: table_for(row, args.bq_dataset), + schema=ROW_SCHEMA, + method=beam.io.WriteToBigQuery.Method.STREAMING_INSERTS, + write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, + create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED) + else: + rows | 'Write routed local files' >> fileio.WriteToFiles( + path=args.output_dir, + destination=lambda row: row['destination'], + file_naming=fileio.destination_prefix_naming('.jsonl'), + sink=JsonSink) + return rows + + +def run(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model', choices=('jev', 'local'), default='jev') + parser.add_argument('--min-confidence', type=float, default=0.5) + parser.add_argument('--pubsub-subscription') + parser.add_argument('--bq-dataset', help='BigQuery project:dataset') + parser.add_argument('--output-dir', help='Local output directory') + parser.add_argument('--graph', help='Render the Beam graph to .dot or .svg') + parser.add_argument( + '--graph-only', + action='store_true', + help='Render without running the job') + args, pipeline_args = parser.parse_known_args(argv) + if not 0 <= args.min_confidence <= 1: + parser.error('--min-confidence must be between 0 and 1') + if not args.bq_dataset and not args.output_dir: + parser.error('Set --output-dir or --bq-dataset') + if args.bq_dataset and args.output_dir: + parser.error('Use only one sink') + if (args.model == 'jev' and not args.graph_only and + not os.environ.get('TYPESAFE_API_KEY')): + parser.error('TYPESAFE_API_KEY is required for --model=jev') + if args.graph_only and not args.graph: + parser.error('--graph-only requires --graph') + + options = PipelineOptions(pipeline_args) + options.view_as(StandardOptions).streaming = True + pipeline = beam.Pipeline(options=options) + build_pipeline(pipeline, args) + if args.graph: + render_options = render.RenderOptions([ + '--render_output=' + args.graph, + '--render_leaf_composite_nodes=^Sample stream$', + '--render_leaf_composite_nodes=^Ask Choice$', + '--render_leaf_composite_nodes=^Write routed .*$', + ]) + render.RenderRunner().run_pipeline(pipeline, render_options) + if args.graph_only: + return None + result = pipeline.run() + result.wait_until_finish() + return result + + +if __name__ == '__main__': + run() diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg new file mode 100644 index 000000000000..e37964497cf8 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg @@ -0,0 +1,98 @@ + + + + + + + +G + + + +ref_AppliedPTransform_Sample-stream_2 + + +Sample stream + + + + + +ref_AppliedPTransform_Ask-Choice_3 + + +Ask Choice + + + + + +ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Ask-Choice_3 + + + + + +ref_AppliedPTransform_Select-destination_5 + + +Select destination + + + + + +ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_5 + + + + + +ref_AppliedPTransform_Window-routed-messages_6 + + +Window routed messages + + + + + +ref_AppliedPTransform_Select-destination_5->ref_AppliedPTransform_Window-routed-messages_6 + + + + + +ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + + +Write routed BigQuery tables + + + + + +ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + + + + + diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py new file mode 100644 index 000000000000..c5040599b081 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py @@ -0,0 +1,75 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for the policy boundary after a Choice decision.""" + +import unittest + +from apache_beam.examples.inference.decision_models.message_router import route_message +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import DecisionResponse +from apache_beam.ml.inference.decision import DecisionResult + + +def result_for(answer): + return DecisionResult( + state={ + 'event_id': '1', 'message': 'I have an account question' + }, + response=DecisionResponse( + answers={'destination': answer}, model='fixed', provider='test'), + latency_ms=1.0) + + +class RouteMessageTest(unittest.TestCase): + def test_uncertain_choice_goes_to_review(self): + result = result_for( + ChoiceAnswer( + choice='billing', + probabilities={ + 'billing': 0.6, 'sales': 0.4 + }, + confidence=0.4)) + + row = route_message(result, min_confidence=0.65) + + self.assertEqual('billing', row['decision']) + self.assertEqual('review', row['destination']) + + def test_confident_choice_selects_destination(self): + result = result_for(ChoiceAnswer(choice='sales', confidence=0.65)) + + row = route_message(result, min_confidence=0.65) + + self.assertEqual('sales', row['destination']) + + def test_unknown_choice_cannot_select_a_table(self): + result = result_for(ChoiceAnswer(choice='other', confidence=1.0)) + + with self.assertRaisesRegex(ValueError, 'unknown destination'): + route_message(result, min_confidence=0.5) + + def test_missing_confidence_goes_to_review(self): + result = result_for(ChoiceAnswer(choice='sales')) + + row = route_message(result, min_confidence=0.5) + + self.assertEqual('review', row['destination']) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt new file mode 100644 index 000000000000..6b99d3ce92cd --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt @@ -0,0 +1,19 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Optional client for apache_beam.ml.inference.typesafe_inference. +-r ../../../ml/inference/typesafe_tests_requirements.txt diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py index 5a47d9289012..8898cc8555f9 100644 --- a/sdks/python/apache_beam/examples/snippets/snippets_test.py +++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py @@ -50,6 +50,7 @@ from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that @@ -1509,8 +1510,11 @@ def check_enriched_elements(actual): class ValueProviderInfoTest(unittest.TestCase): """Tests for accessing value provider info after run.""" def test_accessing_valueprovider_info_after_run(self): - with self.assertLogs(level='INFO') as log_capture: - snippets.accessing_valueprovider_info_after_run() + # Prism's loopback worker runs in this process and does not initialize + # RuntimeValueProvider through sdk_worker_main. + with mock.patch.object(RuntimeValueProvider, 'runtime_options', {}): + with self.assertLogs(level='INFO') as log_capture: + snippets.accessing_valueprovider_info_after_run() expected_log_message = "The string value is" self.assertTrue( any(expected_log_message in log for log in log_capture.output), diff --git a/sdks/python/apache_beam/ml/inference/decision.py b/sdks/python/apache_beam/ml/inference/decision.py new file mode 100644 index 000000000000..fdfdc5801dc4 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/decision.py @@ -0,0 +1,246 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Typed questions and a Beam transform for decision models. + +A decision model evaluates named questions against each input state. Choice +selects a label, Boolean estimates the probability of a true statement, and +Score evaluates an ordered rubric. Adapters translate these types to a local +model or a service's API. + +Use the same questions and transform when changing providers:: + + from apache_beam.ml.inference.decision import ChoiceQuestion + from apache_beam.ml.inference.decision import EvaluateDecisions + + questions = { + 'team': ChoiceQuestion( + 'Which team should handle this message?', + {'billing': 'Invoices and refunds', 'support': 'Technical problems'}) + } + decisions = events | EvaluateDecisions(model, questions) + +Each output is a :class:`DecisionResult` containing the original state, +normalized answers, provider metadata, and request latency. Applications apply +thresholds and route these results with ordinary Beam transforms. +""" + +import abc +import time +from collections.abc import Mapping +from collections.abc import Sequence +from dataclasses import dataclass +from dataclasses import field +from typing import Any +from typing import Union + +import apache_beam as beam +from apache_beam.metrics import Metrics + +__all__ = [ + 'ChoiceQuestion', + 'BooleanQuestion', + 'ScoreQuestion', + 'Question', + 'ChoiceAnswer', + 'BooleanAnswer', + 'ScoreAnswer', + 'Answer', + 'DecisionResponse', + 'DecisionResult', + 'DecisionModel', + 'EvaluateDecisions', +] + + +@dataclass(frozen=True) +class ChoiceQuestion: + """Select one label from ``criteria``, a mapping of labels to descriptions.""" + instructions: str + criteria: Mapping[str, str] + + +@dataclass(frozen=True) +class BooleanQuestion: + """Estimate the probability that the statement in ``instructions`` is true.""" + instructions: str + + +@dataclass(frozen=True) +class ScoreQuestion: + """Score an ordered rubric whose ``criteria`` levels are indexed from zero. + + Providers may interpolate between levels, returning a fractional score. + """ + instructions: str + criteria: Sequence[str] + + +Question = Union[ChoiceQuestion, BooleanQuestion, ScoreQuestion] + + +@dataclass(frozen=True) +class ChoiceAnswer: + """A selected label, with optional label probabilities and confidence. + + Probabilities and confidence use [0, 1]. Confidence is a provider's estimate + of answer certainty; ``None`` means the provider did not supply it. + """ + choice: str + probabilities: Mapping[str, float] = field(default_factory=dict) + confidence: float | None = None + + +@dataclass(frozen=True) +class BooleanAnswer: + """Probability in [0, 1] that the statement is true. + + Applications choose their own threshold for acting on this probability. + """ + probability: float + + +@dataclass(frozen=True) +class ScoreAnswer: + """Rubric score, optional probabilities by level index, and confidence. + + The score uses the question's zero-based scale. Probabilities and confidence + use [0, 1]; ``None`` confidence means the provider did not supply it. + """ + score: float + probabilities: Mapping[int, float] = field(default_factory=dict) + confidence: float | None = None + + +Answer = Union[ChoiceAnswer, BooleanAnswer, ScoreAnswer] + + +@dataclass(frozen=True) +class DecisionResponse: + """One answer per question name, plus the model and provider identifiers.""" + answers: Mapping[str, Answer] + model: str + provider: str + request_id: str | None = None + + +@dataclass(frozen=True) +class DecisionResult: + """An input state, its decision response, and evaluation time in milliseconds.""" + state: Any + response: DecisionResponse + latency_ms: float + + +class DecisionModel(abc.ABC): + """Adapter contract for hosted and local decision models. + + Implement :meth:`evaluate` to return the shared answer types. Override the + context manager methods to open and close clients or load model weights. + Constructors hold serializable configuration; Beam enters the context on + the worker and reuses it across elements and bundles on that DoFn instance. + + An adapter owns provider timeouts, retries, and conversion to the shared + types. Input states must be serializable by Beam and accepted by the adapter. + """ + def __enter__(self) -> 'DecisionModel': + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + pass + + @abc.abstractmethod + def evaluate( + self, state: Any, questions: Mapping[str, Question]) -> DecisionResponse: + """Evaluate each named question against the same state. + + Return exactly one answer of the corresponding type for each question. + Choice labels must belong to the question's criteria; score levels use + their zero-based indices. Raise an exception if evaluation fails. + """ + raise NotImplementedError + + +class EvaluateDecisions(beam.PTransform): + """Evaluate typed questions for every input element with a reusable model. + + Args: + model: A serializable :class:`DecisionModel` adapter. Its context is entered + once per DoFn instance and closed during teardown, which is best effort. + questions: A mapping of question names to Choice, Boolean, or Score + questions. All questions are evaluated against each input state. + + Returns: + A PCollection of :class:`DecisionResult` objects, preserving each element's + timestamp and window. Calls run synchronously on the processing thread. + The ``EvaluateDecisions`` metrics namespace records ``requests``, + ``failures``, and the ``request_latency_ms`` distribution for successful + evaluations. Exceptions propagate to the runner. + """ + def __init__(self, model: DecisionModel, questions: Mapping[str, Question]): + super().__init__() + self._model = model + self._questions = dict(questions) + self._exception_handling = None + + def with_exception_handling(self, **kwargs): + """Return successful decisions and failed states as tagged PCollections. + + Uses :meth:`apache_beam.ParDo.with_exception_handling` with the supplied + options. The default tags are ``good`` and ``bad``; bad records contain + the original state and exception details. Only processing failures are + caught. Model setup and teardown failures still propagate to the runner. + """ + self._exception_handling = kwargs + return self + + def expand(self, states): + evaluate = beam.ParDo(_EvaluateDecisionsDoFn(self._model, self._questions)) + if self._exception_handling is not None: + evaluate = evaluate.with_exception_handling(**self._exception_handling) + return states | evaluate + + +class _EvaluateDecisionsDoFn(beam.DoFn): + def __init__(self, model, questions): + self._model = model + self._questions = questions + self._active_model = None + self._requests = Metrics.counter('EvaluateDecisions', 'requests') + self._failures = Metrics.counter('EvaluateDecisions', 'failures') + self._latency = Metrics.distribution( + 'EvaluateDecisions', 'request_latency_ms') + + def setup(self): + self._active_model = self._model.__enter__() + + def process(self, state): + self._requests.inc() + started = time.perf_counter() + try: + response = self._active_model.evaluate(state, self._questions) + except Exception: + self._failures.inc() + raise + elapsed_ms = (time.perf_counter() - started) * 1000 + self._latency.update(int(elapsed_ms)) + yield DecisionResult(state, response, elapsed_ms) + + def teardown(self): + if self._active_model is not None: + self._model.__exit__(None, None, None) + self._active_model = None diff --git a/sdks/python/apache_beam/ml/inference/decision_test.py b/sdks/python/apache_beam/ml/inference/decision_test.py new file mode 100644 index 000000000000..f596acd6e3ec --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/decision_test.py @@ -0,0 +1,254 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import DecisionModel +from apache_beam.ml.inference.decision import DecisionResponse +from apache_beam.ml.inference.decision import EvaluateDecisions +from apache_beam.ml.inference.decision import ScoreAnswer +from apache_beam.ml.inference.decision import ScoreQuestion +from apache_beam.ml.inference.decision import _EvaluateDecisionsDoFn +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.test_stream import TestStream +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.transforms.window import FixedWindows + +_QUESTIONS = { + 'team': ChoiceQuestion( + 'Which team should handle this message?', { + 'billing': 'Invoices and refunds', + 'support': 'Technical problems', + }), + 'eligible': BooleanQuestion('Is this message eligible for automation?'), + 'severity': ScoreQuestion( + 'How severe is this message?', ('low', 'medium', 'high')), +} + + +def _response(model, provider, state, eligible, severity): + return DecisionResponse( + answers={ + 'team': ChoiceAnswer( + 'support' if eligible > .5 else 'billing', + probabilities={ + 'support': eligible, + 'billing': 1 - eligible, + }), + 'eligible': BooleanAnswer(eligible), + 'severity': ScoreAnswer(float(severity)), + }, + model=model, + provider=provider, + request_id=state['id']) + + +class _KeywordModel(DecisionModel): + def evaluate(self, state, questions): + if set(questions) != set(_QUESTIONS): + raise AssertionError('the model received a different question set') + eligible = .9 if 'bug' in state['text'] else .2 + return _response( + 'keyword-v1', 'local-keyword', state, eligible, state['priority']) + + +class _PriorityModel(DecisionModel): + def evaluate(self, state, questions): + if set(questions) != set(_QUESTIONS): + raise AssertionError('the model received a different question set') + eligible = .8 if state['priority'] >= 2 else .1 + return _response( + 'priority-v1', 'local-priority', state, eligible, state['priority']) + + +class _RejectingModel(_KeywordModel): + def evaluate(self, state, questions): + if state['id'] == 'bad': + raise ValueError('Invalid answer probability') + return super().evaluate(state, questions) + + +class _SummarizeDecision(beam.DoFn): + def process( + self, + result, + timestamp=beam.DoFn.TimestampParam, + window=beam.DoFn.WindowParam): + answers = result.response.answers + yield ( + result.state, + result.response, + answers['team'].choice, + answers['eligible'].probability >= .5, + answers['severity'].score, + timestamp.micros, (window.start.micros, window.end.micros)) + + +class _LifecycleModel(DecisionModel): + def __init__(self): + self.enter_count = 0 + self.exit_calls = [] + self.active = mock.Mock() + + def __enter__(self): + self.enter_count += 1 + return self.active + + def __exit__(self, exc_type, exc_value, traceback): + self.exit_calls.append((exc_type, exc_value, traceback)) + + def evaluate(self, state, questions): + raise AssertionError('the context returned model must evaluate') + + +def _summarize_failure( + failure, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam): + return ( + failure[0], + failure[1][1], + timestamp.micros, (window.start.micros, window.end.micros)) + + +class DecisionTest(unittest.TestCase): + def test_dead_letter_output_preserves_failed_state_and_window(self): + good = {'id': 'good', 'text': 'bug in invoice', 'priority': 2} + bad = {'id': 'bad', 'text': 'refund request', 'priority': 0} + with TestPipeline() as pipeline: + states = ( + pipeline + | beam.Create([ + beam.window.TimestampedValue(good, 1), + beam.window.TimestampedValue(bad, 11), + ]) + | beam.WindowInto(FixedWindows(10))) + decisions, failures = states | EvaluateDecisions( + _RejectingModel(), _QUESTIONS).with_exception_handling( + exc_class=ValueError) + assert_that( + decisions | beam.Map(lambda result: result.state), + equal_to([good]), + label='successful decisions') + assert_that( + failures | beam.Map(_summarize_failure), + equal_to([( + bad, + "ValueError('Invalid answer probability')", + 11_000_000, (10_000_000, 20_000_000))]), + label='failed decisions') + + def test_models_share_questions_downstream_and_preserve_window_metadata(self): + first = {'id': 'first', 'text': 'bug in invoice', 'priority': 2} + second = {'id': 'second', 'text': 'refund request', 'priority': 0} + stream = ( + TestStream().add_elements([ + beam.window.TimestampedValue(first, 1), + beam.window.TimestampedValue(second, 11), + ]).advance_watermark_to_infinity()) + + expected = { + 'keyword': [ + ( + first, + _response('keyword-v1', 'local-keyword', first, .9, 2), + 'support', + True, + 2.0, + 1_000_000, (0, 10_000_000)), + ( + second, + _response('keyword-v1', 'local-keyword', second, .2, 0), + 'billing', + False, + 0.0, + 11_000_000, (10_000_000, 20_000_000)), + ], + 'priority': [ + ( + first, + _response('priority-v1', 'local-priority', first, .8, 2), + 'support', + True, + 2.0, + 1_000_000, (0, 10_000_000)), + ( + second, + _response('priority-v1', 'local-priority', second, .1, 0), + 'billing', + False, + 0.0, + 11_000_000, (10_000_000, 20_000_000)), + ], + } + + with TestPipeline() as pipeline: + states = pipeline | stream | beam.WindowInto(FixedWindows(10)) + for label, model in ( + ('keyword', _KeywordModel()), + ('priority', _PriorityModel()), + ): + decisions = states | f'{label} decisions' >> EvaluateDecisions( + model, _QUESTIONS) + summarized = decisions | f'{label} downstream' >> beam.ParDo( + _SummarizeDecision()) + assert_that( + summarized, equal_to(expected[label]), label=f'check {label}') + + def test_context_model_is_reused_across_bundles_and_closed(self): + model = _LifecycleModel() + model.active.evaluate.side_effect = [ + _response('test', 'mock', {'id': 'one'}, .9, 1), + _response('test', 'mock', {'id': 'two'}, .1, 0), + ] + dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS) + + dofn.setup() + dofn.start_bundle() + first = list(dofn.process({'id': 'one'})) + dofn.finish_bundle() + dofn.start_bundle() + second = list(dofn.process({'id': 'two'})) + dofn.finish_bundle() + dofn.teardown() + + self.assertEqual(1, model.enter_count) + self.assertEqual(2, model.active.evaluate.call_count) + self.assertEqual(1, len(model.exit_calls)) + self.assertEqual((None, None, None), model.exit_calls[0]) + self.assertEqual('one', first[0].state['id']) + self.assertEqual('two', second[0].state['id']) + + def test_model_exception_propagates_from_process(self): + model = _LifecycleModel() + model.active.evaluate.side_effect = RuntimeError('provider failed') + dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS) + + dofn.setup() + with self.assertRaisesRegex(RuntimeError, 'provider failed'): + list(dofn.process({'id': 'failed'})) + dofn.teardown() + self.assertEqual(1, len(model.exit_calls)) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference.py b/sdks/python/apache_beam/ml/inference/typesafe_inference.py new file mode 100644 index 000000000000..130a0cafb5e9 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference.py @@ -0,0 +1,164 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Optional TypeSafe adapter for Beam decision models.""" + +from collections.abc import Mapping + +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import DecisionModel +from apache_beam.ml.inference.decision import DecisionResponse +from apache_beam.ml.inference.decision import ScoreAnswer +from apache_beam.ml.inference.decision import ScoreQuestion + +__all__ = ['JevDecisionModel'] + + +def _number(name, field, value, maximum=1): + if (isinstance(value, bool) or not isinstance(value, (int, float)) or + not 0 <= value <= maximum): + raise ValueError( + f'Invalid {field} for question {name!r}: ' + f'expected a finite number in [0, {maximum}]') + return value + + +def _probabilities(name, probabilities, keys, key_type): + if not isinstance(probabilities, Mapping): + raise ValueError( + f'Invalid probabilities for question {name!r}: ' + 'expected a mapping') + for key, probability in probabilities.items(): + if (not isinstance(key, key_type) or isinstance(key, bool) or + key not in keys): + raise ValueError( + f'Invalid probability key for question {name!r}: {key!r}') + _number(name, 'probability', probability) + return dict(probabilities) + + +def _confidence(name, confidence): + return None if confidence is None else _number(name, 'confidence', confidence) + + +def _answer_type(name, answer, expected_type): + if not isinstance(answer, expected_type): + raise ValueError( + f'Invalid answer type for question {name!r}: ' + f'expected {expected_type.__name__}') + + +class JevDecisionModel(DecisionModel): + """Evaluate Beam decision questions with Jev through the TypeSafe SDK. + + Install the optional SDK from ``typesafe_tests_requirements.txt`` and set + ``TYPESAFE_API_KEY`` on the workers. The client is created when Beam + enters the adapter context; its HTTP pool is reused until teardown. + Choice and Score map directly to the SDK primitives. Boolean maps to Noul, + whose ``noul`` field becomes :attr:`BooleanAnswer.probability`. + + The SDK retries 429 and server errors with exponential backoff, jitter, and + ``Retry-After`` support. Calls and retries block the processing thread; they + do not coordinate rate limits across workers. Invalid answer keys, types, + labels, numeric ranges, or probability keys raise ``ValueError``. + + Args: + model: TypeSafe model identifier, defaulting to ``jev-latest``. + timeout: Timeout in seconds for HTTP operations. The SDK handles retries. + """ + def __init__(self, model='jev-latest', timeout=10.0): + self.model = model + self.timeout = timeout + self.client = None + + def __enter__(self): + try: + from typesafe_sdk import TypeSafeClient + except ImportError as error: + raise ImportError( + 'Install typesafe-sdk to use JevDecisionModel') from error + self.client = TypeSafeClient(timeout=self.timeout) + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.client.close() + self.client = None + + def evaluate(self, state, questions): + if self.client is None: + raise RuntimeError('Enter JevDecisionModel before evaluating questions') + import typesafe_sdk + + sdk_questions = {} + for name, question in questions.items(): + if isinstance(question, ChoiceQuestion): + sdk_questions[name] = typesafe_sdk.Choice( + instructions=question.instructions, criteria=question.criteria) + elif isinstance(question, BooleanQuestion): + sdk_questions[name] = typesafe_sdk.Noul( + instructions=question.instructions) + elif isinstance(question, ScoreQuestion): + sdk_questions[name] = typesafe_sdk.Score( + instructions=question.instructions, criteria=question.criteria) + else: + raise TypeError('Unsupported decision question: %r' % type(question)) + + try: + response = self.client.system_one( + state=state, questions=sdk_questions, model=self.model) + except typesafe_sdk.TypeSafeAPIResponseValidationError as error: + raise ValueError('Invalid Jev response') from error + if (not isinstance(response.answers, Mapping) or + response.answers.keys() != questions.keys()): + raise ValueError( + 'Jev answer question names must match the requested ' + 'question names') + answers = {} + for name, question in questions.items(): + answer = response.answers[name] + if isinstance(question, ChoiceQuestion): + _answer_type(name, answer, typesafe_sdk.ChoiceAnswer) + if (not isinstance(answer.choice, str) or + answer.choice not in question.criteria): + raise ValueError( + f'Invalid choice for question {name!r}: ' + f'{answer.choice!r}') + answers[name] = ChoiceAnswer( + choice=answer.choice, + probabilities=_probabilities( + name, answer.probabilities, question.criteria, str), + confidence=_confidence(name, answer.confidence)) + elif isinstance(question, BooleanQuestion): + _answer_type(name, answer, typesafe_sdk.NoulAnswer) + answers[name] = BooleanAnswer( + probability=_number(name, 'probability', answer.noul)) + else: + _answer_type(name, answer, typesafe_sdk.ScoreAnswer) + answers[name] = ScoreAnswer( + score=_number( + name, 'score', answer.score, len(question.criteria) - 1), + probabilities=_probabilities( + name, answer.probabilities, range(len(question.criteria)), int), + confidence=_confidence(name, answer.confidence)) + return DecisionResponse( + answers=answers, + model=response.model, + provider='typesafe', + request_id=response.request_id) diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py new file mode 100644 index 000000000000..acd2234f2ddf --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py @@ -0,0 +1,93 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Live Jev integration test, enabled with TYPESAFE_API_KEY.""" + +import os +import unittest + +import pytest + +try: + import typesafe_sdk +except ImportError: + typesafe_sdk = None + +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import EvaluateDecisions +from apache_beam.ml.inference.decision import ScoreAnswer +from apache_beam.ml.inference.decision import ScoreQuestion +from apache_beam.ml.inference.typesafe_inference import JevDecisionModel +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.test_stream import TestStream +from apache_beam.testing.util import assert_that + + +@pytest.mark.it_postcommit +@unittest.skipIf(typesafe_sdk is None, 'typesafe-sdk is not installed') +@unittest.skipUnless( + os.environ.get('TYPESAFE_API_KEY'), + 'TYPESAFE_API_KEY environment variable is not set') +class JevInferenceIT(unittest.TestCase): + def test_typed_decisions_in_streaming_pipeline(self): + state = { + 'message': 'My invoice was charged twice. Please refund the duplicate.' + } + questions = { + 'team': ChoiceQuestion( + 'Which team should handle this message?', + { + 'billing': 'Invoices and refunds', + 'technical': 'Bugs and outages', + }), + 'refund': BooleanQuestion('Does this message ask for a refund?'), + 'urgency': ScoreQuestion( + 'How urgent is this message?', ('Routine', 'Urgent', 'Emergency')), + } + + def check_results(results): + assert len(results) == 1, results + result = results[0] + assert result.state == state + assert result.response.provider == 'typesafe' + assert result.response.model + assert set(result.response.answers) == set(questions) + choice = result.response.answers['team'] + boolean = result.response.answers['refund'] + score = result.response.answers['urgency'] + assert isinstance(choice, ChoiceAnswer) + assert choice.choice in questions['team'].criteria + assert isinstance(boolean, BooleanAnswer) + assert 0 <= boolean.probability <= 1 + assert isinstance(score, ScoreAnswer) + assert 0 <= score.score <= 2 + assert set(score.probabilities) <= {0, 1, 2} + + stream = TestStream().add_elements([state]).advance_watermark_to_infinity() + with TestPipeline() as pipeline: + results = ( + pipeline + | stream + | EvaluateDecisions(JevDecisionModel(), questions)) + assert_that(results, check_results) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py new file mode 100644 index 000000000000..5776fc9133e8 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py @@ -0,0 +1,256 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for TypeSafe's translation to the shared decision types.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +from parameterized import parameterized + +try: + import typesafe_sdk +except ImportError: + raise unittest.SkipTest('typesafe-sdk is not installed') + +from apache_beam.ml.inference.decision import BooleanAnswer +from apache_beam.ml.inference.decision import BooleanQuestion +from apache_beam.ml.inference.decision import ChoiceAnswer +from apache_beam.ml.inference.decision import ChoiceQuestion +from apache_beam.ml.inference.decision import ScoreAnswer +from apache_beam.ml.inference.decision import ScoreQuestion +from apache_beam.ml.inference.typesafe_inference import JevDecisionModel + + +def _choice_answer(**fields): + return typesafe_sdk.ChoiceAnswer.model_construct( + **{ + 'choice': 'billing', + 'probabilities': { + 'billing': 1.0 + }, + 'confidence': 0.9, + **fields, + }) + + +def _score_answer(**fields): + return typesafe_sdk.ScoreAnswer.model_construct( + **{ + 'score': 0.7, + 'probabilities': { + 0: 0.3, 1: 0.7 + }, + 'confidence': 0.6, + **fields, + }) + + +class JevDecisionModelTest(unittest.TestCase): + def test_normalizes_all_primitives_and_reuses_client(self): + client = mock.Mock() + client.system_one.return_value = SimpleNamespace( + answers={ + 'team': _choice_answer(), + 'review': typesafe_sdk.NoulAnswer.model_construct(noul=0.8), + 'urgency': _score_answer(), + }, + model='jev-test', + request_id='request-1') + questions = { + 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}), + 'review': BooleanQuestion('Does this need review?'), + 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']), + } + state = {'message': 'Review my invoice'} + with mock.patch.object(typesafe_sdk, 'TypeSafeClient', + return_value=client) as factory: + adapter = JevDecisionModel(model='jev-test', timeout=2.0) + factory.assert_not_called() + with adapter: + response = adapter.evaluate(state, questions) + adapter.evaluate(state, questions) + factory.assert_called_once_with(timeout=2.0) + client.close.assert_called_once_with() + self.assertEqual(2, client.system_one.call_count) + request = client.system_one.call_args.kwargs + self.assertEqual(state, request['state']) + self.assertEqual('jev-test', request['model']) + sent = request['questions'] + self.assertIsInstance(sent['team'], typesafe_sdk.Choice) + self.assertEqual(questions['team'].criteria, sent['team'].criteria) + self.assertIsInstance(sent['review'], typesafe_sdk.Noul) + self.assertEqual( + questions['review'].instructions, sent['review'].instructions) + self.assertIsInstance(sent['urgency'], typesafe_sdk.Score) + self.assertEqual(questions['urgency'].criteria, sent['urgency'].criteria) + self.assertEqual( + ChoiceAnswer('billing', {'billing': 1.0}, 0.9), + response.answers['team']) + self.assertEqual(BooleanAnswer(0.8), response.answers['review']) + self.assertEqual( + ScoreAnswer(0.7, { + 0: 0.3, 1: 0.7 + }, 0.6), response.answers['urgency']) + self.assertEqual('jev-test', response.model) + self.assertEqual('typesafe', response.provider) + self.assertEqual('request-1', response.request_id) + + @parameterized.expand([ + ('choice_label', 'team', _choice_answer(choice='sales')), + ('choice_type', 'team', typesafe_sdk.NoulAnswer.model_construct(noul=.5)), + ( + 'choice_probability_key', + 'team', + _choice_answer(probabilities={'sales': 1.0})), + ('choice_probabilities_type', 'team', _choice_answer(probabilities=[])), + ( + 'choice_probability_range', + 'team', + _choice_answer(probabilities={'billing': 1.2})), + ( + 'choice_probability_nan', + 'team', + _choice_answer(probabilities={'billing': float('nan')})), + ('choice_confidence', 'team', _choice_answer(confidence=1.3)), + ('choice_confidence_type', 'team', _choice_answer(confidence=True)), + ( + 'boolean_range', + 'review', + typesafe_sdk.NoulAnswer.model_construct(noul=1.4)), + ( + 'boolean_nan', + 'review', + typesafe_sdk.NoulAnswer.model_construct(noul=float('nan'))), + ( + 'boolean_type', + 'review', + typesafe_sdk.NoulAnswer.model_construct(noul='0.5')), + ('boolean_answer_type', 'review', _choice_answer()), + ('score_range', 'urgency', _score_answer(score=3.5)), + ('score_negative', 'urgency', _score_answer(score=-.1)), + ('score_infinite', 'urgency', _score_answer(score=float('inf'))), + ('score_boolean', 'urgency', _score_answer(score=True)), + ('score_answer_type', 'urgency', _choice_answer()), + ( + 'score_probability_key', + 'urgency', + _score_answer(probabilities={999: 1.0})), + ( + 'score_probability_key_type', + 'urgency', + _score_answer(probabilities={'1': 1.0})), + ( + 'score_probability_fractional_key', + 'urgency', + _score_answer(probabilities={.5: 1.0})), + ( + 'score_probability_negative', + 'urgency', + _score_answer(probabilities={0: -.1})), + ('score_confidence', 'urgency', _score_answer(confidence=-.1)), + ]) + def test_rejects_malformed_answer(self, unused_name, name, answer): + questions = { + 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}), + 'review': BooleanQuestion('Does this need review?'), + 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']), + } + client = mock.Mock() + client.system_one.return_value = SimpleNamespace( + answers={name: answer}, model='jev-test', request_id='request-1') + with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client): + with JevDecisionModel() as adapter: + with self.assertRaisesRegex(ValueError, name): + adapter.evaluate('Hello', {name: questions[name]}) + + @parameterized.expand([ + ('missing', {}), + ( + 'extra', + { + 'review': typesafe_sdk.NoulAnswer.model_construct(noul=.5), + 'unexpected': typesafe_sdk.NoulAnswer.model_construct(noul=.5), + }), + ('wrong_container', []), + ]) + def test_rejects_answer_keys(self, unused_name, answers): + client = mock.Mock() + client.system_one.return_value = SimpleNamespace( + answers=answers, model='jev-test', request_id='request-1') + with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client): + with JevDecisionModel() as adapter: + with self.assertRaisesRegex(ValueError, 'question names'): + adapter.evaluate( + 'Hello', {'review': BooleanQuestion('Needs review?')}) + + def test_accepts_boundaries_and_optional_probabilities(self): + questions = { + 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}), + 'review': BooleanQuestion('Does this need review?'), + 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']), + } + client = mock.Mock() + for probability in (0, 1): + with self.subTest(probability=probability): + client.system_one.return_value = SimpleNamespace( + answers={ + 'team': _choice_answer(probabilities={}, confidence=None), + 'review': typesafe_sdk.NoulAnswer.model_construct( + noul=probability), + 'urgency': _score_answer( + score=probability, probabilities={}, confidence=None), + }, + model='jev-test', + request_id='request-1') + with mock.patch.object(typesafe_sdk, + 'TypeSafeClient', + return_value=client): + with JevDecisionModel() as adapter: + response = adapter.evaluate('Hello', questions) + self.assertEqual(probability, response.answers['review'].probability) + self.assertEqual(probability, response.answers['urgency'].score) + self.assertIsNone(response.answers['team'].confidence) + + def test_provider_error_propagates_and_client_closes(self): + client = mock.Mock() + client.system_one.side_effect = TimeoutError('request timed out') + with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client): + with self.assertRaisesRegex(TimeoutError, 'request timed out'): + with JevDecisionModel() as adapter: + adapter.evaluate( + 'Hello', {'review': BooleanQuestion('Needs review?')}) + client.close.assert_called_once_with() + + def test_sdk_response_validation_error_is_a_value_error(self): + error = typesafe_sdk.TypeSafeAPIResponseValidationError( + status=200, body={}, headers={}, field_path='answers.review.noul') + client = mock.Mock() + client.system_one.side_effect = error + with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client): + with JevDecisionModel() as adapter: + with self.assertRaisesRegex(ValueError, + 'Invalid Jev response') as raised: + adapter.evaluate( + 'Hello', {'review': BooleanQuestion('Needs review?')}) + self.assertIs(error, raised.exception.__cause__) + client.close.assert_called_once_with() + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt new file mode 100644 index 000000000000..982081593f5e --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt @@ -0,0 +1,18 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +typesafe-sdk>=0.7.1,<0.8 diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 24e9440361ca..d18375d6129d 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -613,6 +613,7 @@ def get_portability_package_data(): ] + ml_base + qdrant_dependency, 'p312_ml_test': [ 'datatable', + 'typesafe-sdk>=0.7.1,<0.8', ] + ml_base + qdrant_dependency, # maintainer: milvus tests only run with this extension. Make sure it # is covered by docker-in-docker test when changing py version