From 3fe470a8c15f43637a353b91b83ab20f5061e28e Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Mon, 19 Oct 2026 15:46:20 +0000 Subject: [PATCH 1/2] [yaml] - add openai model handler --- .../yaml/tests/runinference_openai.yaml | 60 ++++++++ sdks/python/apache_beam/yaml/yaml_ml.py | 138 ++++++++++++++++++ sdks/python/apache_beam/yaml/yaml_ml_test.py | 71 +++++++++ sdks/python/setup.py | 1 + 4 files changed, 270 insertions(+) create mode 100644 sdks/python/apache_beam/yaml/tests/runinference_openai.yaml diff --git a/sdks/python/apache_beam/yaml/tests/runinference_openai.yaml b/sdks/python/apache_beam/yaml/tests/runinference_openai.yaml new file mode 100644 index 000000000000..a36a42486af0 --- /dev/null +++ b/sdks/python/apache_beam/yaml/tests/runinference_openai.yaml @@ -0,0 +1,60 @@ +# 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. + +pipelines: + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - text: "I love Apache Beam!" + - text: "I hate this error." + - type: RunInference + config: + model_handler: + type: "OpenAIModelHandler" + config: + model_name: "gpt-5.4-mini" + api_key: "fake-api-key" + request_fn: + callable: | + def fake_inference(model_name, batch, client, inference_args): + return [ + { + 'label': 'POSITIVE' if 'love' in text else 'NEGATIVE', + 'model': model_name + } + for text in batch + ] + preprocess: + callable: 'lambda x: x.text' + - type: MapToFields + config: + language: python + fields: + text: text + sentiment: + callable: 'lambda x: x.inference.inference["label"]' + - type: AssertEqual + config: + elements: + - text: "I love Apache Beam!" + sentiment: "POSITIVE" + - text: "I hate this error." + sentiment: "NEGATIVE" + + options: + yaml_experimental_features: ['ML'] diff --git a/sdks/python/apache_beam/yaml/yaml_ml.py b/sdks/python/apache_beam/yaml/yaml_ml.py index a4437b2ebebf..f640f7257744 100644 --- a/sdks/python/apache_beam/yaml/yaml_ml.py +++ b/sdks/python/apache_beam/yaml/yaml_ml.py @@ -374,6 +374,143 @@ def inference_output_type(self): return Any +@ModelHandlerProvider.register_handler_type('OpenAIModelHandler') +class OpenAIModelHandlerProvider(ModelHandlerProvider): + def __init__( + self, + model_name: str, + preprocess: dict[str, str], + request_fn: Optional[dict[str, str]] = None, + api_key: Optional[str] = None, + organization: Optional[str] = None, + project: Optional[str] = None, + base_url: Optional[str] = None, + client_args: Optional[dict[str, Any]] = None, + postprocess: Optional[dict[str, str]] = None, + min_batch_size: Optional[int] = None, + max_batch_size: Optional[int] = None, + max_batch_duration_secs: Optional[int] = None, + max_batch_weight: Optional[int] = None, + element_size_fn: Optional[dict[str, str]] = None, + batch_length_fn: Optional[dict[str, str]] = None, + batch_bucket_boundaries: Optional[list[int]] = None, + **kwargs): + """ + ModelHandler for OpenAI models. + + This Model Handler can be used with RunInference to run inference against + the OpenAI API (or OpenAI-compatible endpoints via `base_url`) using models + such as GPT-5.4, GPT-5.4-mini, or embedding models. + + This Model Handler requires `model_name` and a `preprocess` function to be + specified. Preprocessing and Postprocessing are described in more detail in + the RunInference docs: + https://beam.apache.org/releases/yamldoc/current/#runinference + + For example: :: + + - type: RunInference + config: + model_handler: + type: OpenAIModelHandler + config: + model_name: gpt-5.4-mini + api_key: your-api-key + preprocess: + callable: 'lambda x: x.prompt' + + Args: + model_name: The OpenAI model to send requests to (e.g. 'gpt-5.4', + 'gpt-5.4-mini', 'text-embedding-3-small'). + preprocess: A python callable, defined either inline, or using a file, + that is invoked on the input row before sending to the model to be + loaded by this ModelHandler. This parameter is required by the + `OpenAIModelHandler` ModelHandler. + request_fn: A python callable, defined either inline, or using a file, + used to send requests to OpenAI. Built-in options in + `apache_beam.ml.inference.openai_inference` are + `chat_completion_from_string`, `chat_completion_from_conversation`, and + `embedding_from_string`. Defaults to `chat_completion_from_string`. + api_key: The OpenAI API key. If not provided, the client will look for + the `OPENAI_API_KEY` environment variable. + organization: Optional OpenAI organization ID. If not provided, the + client will look for the `OPENAI_ORG_ID` environment variable. + project: Optional OpenAI project ID. If not provided, the client will + look for the `OPENAI_PROJECT_ID` environment variable. + base_url: Optional base URL for the API requests. This can be used to + target OpenAI-compatible servers (such as vLLM, Ollama, or Azure + OpenAI endpoints). + client_args: Optional dictionary of additional keyword arguments passed + when instantiating the OpenAI client (e.g. `timeout`, + `default_headers`). + postprocess: A python callable, defined either inline, or using a file, + that is invoked on the PredictionResult output by the ModelHandler + before parsing into the output Beam Row. + min_batch_size: Optional minimum batch size to use when batching inputs. + max_batch_size: Optional maximum batch size to use when batching inputs. + max_batch_duration_secs: Optional maximum amount of time to buffer a + batch before emitting; used in streaming contexts. + max_batch_weight: Optional maximum total weight of a batch. + element_size_fn: Optional python callable that returns the size (weight) + of an element. + batch_length_fn: Optional python callable that returns the length of an + element for length-aware batching. + batch_bucket_boundaries: Optional sorted list of positive boundary values + for length-aware batching buckets. + **kwargs: Extra arguments to pass to the model handler. + """ + try: + from apache_beam.ml.inference.openai_inference import OpenAIModelHandler + from apache_beam.ml.inference.openai_inference import chat_completion_from_string + except ImportError: + raise ValueError( + 'Unable to import OpenAIModelHandler. Please ' + 'install openai dependencies: `pip install openai`') + + kwargs = {k: v for k, v in kwargs.items() if not k.startswith('_')} + + request_fn_obj = ( + self.parse_processing_transform(request_fn, 'request_fn') + if request_fn else chat_completion_from_string) + element_size_fn_obj = ( + self.parse_processing_transform(element_size_fn, 'element_size_fn') + if element_size_fn else None) + batch_length_fn_obj = ( + self.parse_processing_transform(batch_length_fn, 'batch_length_fn') + if batch_length_fn else None) + + handler_kwargs: dict[str, Any] = {} + if element_size_fn_obj is not None: + handler_kwargs['element_size_fn'] = element_size_fn_obj + if batch_length_fn_obj is not None: + handler_kwargs['batch_length_fn'] = batch_length_fn_obj + + _handler = OpenAIModelHandler( + model_name=model_name, + request_fn=request_fn_obj, + api_key=api_key, + organization=organization, + project=project, + base_url=base_url, + client_args=client_args, + min_batch_size=min_batch_size, + max_batch_size=max_batch_size, + max_batch_duration_secs=max_batch_duration_secs, + max_batch_weight=max_batch_weight, + batch_bucket_boundaries=batch_bucket_boundaries, + **handler_kwargs, + **kwargs) + + super().__init__(_handler, preprocess, postprocess) + + @staticmethod + def validate(model_handler_spec): + pass + + def inference_output_type(self): + return Any + + @beam.ptransform.ptransform_fn def run_inference( pcoll, @@ -536,6 +673,7 @@ def fn(x: PredictionResult): - [VertexAIModelHandlerJSON](https://beam.apache.org/releases/pydoc/current/apache_beam.yaml.yaml_ml.VertexAIModelHandlerJSONProvider) # pylint: disable=line-too-long - [HuggingFacePipelineModelHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.yaml.yaml_ml.HuggingFacePipelineModelHandlerProvider) # pylint: disable=line-too-long + - [OpenAIModelHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.yaml.yaml_ml.OpenAIModelHandlerProvider) # pylint: disable=line-too-long inference_tag: The tag to use for the returned inference. Default is 'inference'. inference_args: Extra arguments for models whose inference call requires diff --git a/sdks/python/apache_beam/yaml/yaml_ml_test.py b/sdks/python/apache_beam/yaml/yaml_ml_test.py index 0b6be61a5d09..c06ab8e7bce0 100644 --- a/sdks/python/apache_beam/yaml/yaml_ml_test.py +++ b/sdks/python/apache_beam/yaml/yaml_ml_test.py @@ -313,6 +313,77 @@ def test_huggingface_pipeline_model_handler_does_not_use_ml_frameworks(self): self.assertEqual(attempted_imports, []) +class RunInferenceOpenAITest(unittest.TestCase): + def test_openai_model_handler_missing_model_name(self): + ml_opts = beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle', yaml_experimental_features=['ML']) + with self.assertRaisesRegex( + ValueError, + "Unable to instantiate model handler of type OpenAIModelHandler.*" + "model_name"): + with beam.Pipeline(options=ml_opts) as p: + _ = p | YamlTransform( + ''' + type: chain + transforms: + - type: Create + config: + elements: + - {text: "I love Apache Beam!"} + - type: RunInference + config: + model_handler: + type: OpenAIModelHandler + config: + preprocess: + callable: 'lambda row: row.text' + ''') + + def test_openai_model_handler_execution(self): + try: + import openai # pylint: disable=unused-import + except ImportError: + self.skipTest('openai is not installed.') + + ml_opts = beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle', yaml_experimental_features=['ML']) + with beam.Pipeline(options=ml_opts) as p: + result = p | YamlTransform( + ''' + type: chain + transforms: + - type: Create + config: + elements: + - {text: "I love Apache Beam!"} + - {text: "I hate this error."} + - type: RunInference + config: + model_handler: + type: OpenAIModelHandler + config: + model_name: gpt-5.4-mini + api_key: fake-api-key + request_fn: + callable: | + def fake_request(model_name, batch, client, inference_args): + return [ + 'POSITIVE' if 'love' in item else 'NEGATIVE' + for item in batch + ] + preprocess: + callable: 'lambda row: row.text' + postprocess: + callable: 'lambda pred: pred.inference' + ''') + assert_that( + result | beam.Map(lambda row: (row.text, row.inference)), + equal_to([ + ('I love Apache Beam!', 'POSITIVE'), + ('I hate this error.', 'NEGATIVE'), + ])) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 08c5964898de..09c6af08f196 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -204,6 +204,7 @@ def cythonize(*args, **kwargs): 'transformers>=4.28.0,<4.56.0', # Keep tokenizers compatible with this transformers range. 'tokenizers>=0.13.3,<0.22.0', + 'openai>=3.0.0', ] ml_adk_dependency = [ From b60ae5058bba0d6bb5f5c3b6cd15ff96bcbe808a Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Mon, 19 Oct 2026 17:25:52 +0000 Subject: [PATCH 2/2] fix import error --- sdks/python/apache_beam/yaml/yaml_ml_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/yaml/yaml_ml_test.py b/sdks/python/apache_beam/yaml/yaml_ml_test.py index c06ab8e7bce0..79cdcaf64e6b 100644 --- a/sdks/python/apache_beam/yaml/yaml_ml_test.py +++ b/sdks/python/apache_beam/yaml/yaml_ml_test.py @@ -341,7 +341,7 @@ def test_openai_model_handler_missing_model_name(self): def test_openai_model_handler_execution(self): try: - import openai # pylint: disable=unused-import + from apache_beam.ml.inference.openai_inference import OpenAIModelHandler # pylint: disable=unused-import except ImportError: self.skipTest('openai is not installed.')