From e6694bc9bc9b8e416095f8e1311f3d896c4c82c3 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Thu, 2 Oct 2025 23:09:26 -0400 Subject: [PATCH 1/8] Intial commit --- .../ml/inference/TritonModelHandler.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 sdks/python/apache_beam/ml/inference/TritonModelHandler.py diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py new file mode 100644 index 000000000000..2dcced66fbd5 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -0,0 +1,54 @@ +from typing import Sequence, Dict, Any, Iterable, Optional +import logging +import json +import tritonserver +from tritonserver import Model +from apache_beam.ml.inference.base import ModelHandler, PredictionResult + + +class TritonModelHandler(ModelHandler[str, + PredictionResult, + Model]): + + def __init__( + self, + model_repository: str = "/workspace/models", + model_name: str = "doc_tagging" + ): + """ + + """ + + self._model_repository = model_repository + self._model_name = model_name + + def load_model(self) -> Model: + """Loads and initializes a model for processing.""" + server = tritonserver.Server(model_repository=self._model_repository) + server.start() + model = server.model(self._model_name) + return model + + def run_inference( + self, + batch: Sequence[str], + model: Model, + inference_args: Optional[Dict[str, Any]] = None + ) -> Iterable[PredictionResult]: + """Runs inferences on a batch of text strings. + + Args: + batch: A sequence of examples as text strings. + model: Model returned by Tritonserver. + inference_args: Any additional arguments for an inference. + + Returns: + An Iterable of type PredictionResult. + """ + # Loop each text string, and use a tuple to store the inference results. + logging.info(f"{type(batch)}") + responses = model.infer(inputs={"string_input":batch}) + for response in responses: + predictions = [json.loads(answer) for answer in response.outputs["string_output"].to_string_array().tolist()] + return [PredictionResult(x, y) for x, y in zip(batch, predictions)] + From 9cf745646da3da6f8bedc3819e3fb631ca4dc020 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 4 Oct 2025 21:34:29 -0400 Subject: [PATCH 2/8] Improve TritonModelHandler with proper lifecycle management and tests - Add TritonModelWrapper for server cleanup on deletion - Make tensor names configurable via constructor and inference_args - Add comprehensive error handling with descriptive messages - Complete all docstrings with usage examples - Add custom output parsing function support - Create full unit test suite with mocks - Apply Apache license header and yapf formatting --- .../ml/inference/TritonModelHandler.py | 237 +++++++++++++---- .../ml/inference/triton_inference_test.py | 242 ++++++++++++++++++ 2 files changed, 432 insertions(+), 47 deletions(-) create mode 100644 sdks/python/apache_beam/ml/inference/triton_inference_test.py diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py index 2dcced66fbd5..e670a5d212fc 100644 --- a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -1,54 +1,197 @@ +# +# 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. +# + +"""Apache Beam ModelHandler implementation for Triton Inference Server.""" + from typing import Sequence, Dict, Any, Iterable, Optional import logging import json -import tritonserver -from tritonserver import Model +import atexit + from apache_beam.ml.inference.base import ModelHandler, PredictionResult +try: + import tritonserver + from tritonserver import Model, Server +except ImportError: + tritonserver = None # type: ignore + +LOGGER = logging.getLogger(__name__) + + +class TritonModelWrapper: + """Wrapper to manage Triton Server lifecycle with the model.""" + def __init__(self, server: 'Server', model: 'Model'): + self.server = server + self.model = model + + def __del__(self): + """Cleanup server when model is garbage collected.""" + try: + if self.server: + self.server.stop() + except Exception as e: + LOGGER.warning(f"Error stopping Triton server: {e}") + + +class TritonModelHandler(ModelHandler[Any, PredictionResult, + TritonModelWrapper]): + """Beam ModelHandler for Triton Inference Server. + + This handler supports loading models from a Triton model repository and + running inference using the Triton Python API. + + Example usage:: + + pcoll | RunInference( + TritonModelHandler( + model_repository="/workspace/models", + model_name="my_model", + input_tensor_name="input", + output_tensor_name="output" + ) + ) + + Args: + model_repository: Path to the Triton model repository directory. + model_name: Name of the model to load from the repository. + input_tensor_name: Name of the input tensor (default: "INPUT"). + output_tensor_name: Name of the output tensor (default: "OUTPUT"). + parse_output_fn: Optional custom function to parse model outputs. + Should take (outputs_dict, output_tensor_name) and return parsed result. + """ + def __init__( + self, + model_repository: str, + model_name: str, + input_tensor_name: str = "INPUT", + output_tensor_name: str = "OUTPUT", + parse_output_fn: Optional[callable] = None, + ): + if tritonserver is None: + raise ImportError( + "tritonserver is not installed. " + "Install it with: pip install tritonserver") + + self._model_repository = model_repository + self._model_name = model_name + self._input_tensor_name = input_tensor_name + self._output_tensor_name = output_tensor_name + self._parse_output_fn = parse_output_fn + + def load_model(self) -> TritonModelWrapper: + """Loads and initializes a Triton model for processing. + + Returns: + TritonModelWrapper containing the server and model instances. + + Raises: + RuntimeError: If server fails to start or model fails to load. + """ + try: + server = tritonserver.Server(model_repository=self._model_repository) + server.start() + except Exception as e: + raise RuntimeError( + f"Failed to start Triton server with repository " + f"'{self._model_repository}': {e}") from e + + try: + model = server.model(self._model_name) + if model is None: + raise RuntimeError( + f"Model '{self._model_name}' not found in repository") + except Exception as e: + server.stop() + raise RuntimeError( + f"Failed to load model '{self._model_name}': {e}") from e + + return TritonModelWrapper(server, model) + + def run_inference( + self, + batch: Sequence[Any], + model: TritonModelWrapper, + inference_args: Optional[Dict[str, Any]] = None + ) -> Iterable[PredictionResult]: + """Runs inferences on a batch of inputs. + + Args: + batch: A sequence of examples (can be strings, arrays, etc.). + model: TritonModelWrapper returned by load_model(). + inference_args: Optional dict with 'input_tensor_name' and/or + 'output_tensor_name' to override defaults for this batch. + + Returns: + An Iterable of PredictionResult objects. + + Raises: + RuntimeError: If inference fails. + """ + # Allow per-batch tensor name overrides + input_name = self._input_tensor_name + output_name = self._output_tensor_name + if inference_args: + input_name = inference_args.get('input_tensor_name', input_name) + output_name = inference_args.get('output_tensor_name', output_name) + + try: + responses = model.model.infer(inputs={input_name: batch}) + except Exception as e: + raise RuntimeError( + f"Triton inference failed for model '{self._model_name}': {e}") from e + + # Parse outputs + predictions = [] + try: + for response in responses: + if output_name not in response.outputs: + raise RuntimeError( + f"Output tensor '{output_name}' not found in response. " + f"Available outputs: {list(response.outputs.keys())}") + + output_tensor = response.outputs[output_name] + + # Use custom parser if provided + if self._parse_output_fn: + parsed = self._parse_output_fn(response.outputs, output_name) + else: + # Default parsing: try string array, fallback to raw + try: + parsed = [ + json.loads(val) + for val in output_tensor.to_string_array().tolist() + ] + except Exception: + # If JSON parsing fails, return raw output + parsed = output_tensor.to_bytes_array().tolist() + + predictions.extend(parsed if isinstance(parsed, list) else [parsed]) + + except Exception as e: + raise RuntimeError(f"Failed to parse model outputs: {e}") from e + + if len(predictions) != len(batch): + LOGGER.warning( + f"Prediction count ({len(predictions)}) doesn't match " + f"batch size ({len(batch)}). Truncating or padding.") -class TritonModelHandler(ModelHandler[str, - PredictionResult, - Model]): - - def __init__( - self, - model_repository: str = "/workspace/models", - model_name: str = "doc_tagging" - ): - """ - - """ - - self._model_repository = model_repository - self._model_name = model_name - - def load_model(self) -> Model: - """Loads and initializes a model for processing.""" - server = tritonserver.Server(model_repository=self._model_repository) - server.start() - model = server.model(self._model_name) - return model - - def run_inference( - self, - batch: Sequence[str], - model: Model, - inference_args: Optional[Dict[str, Any]] = None - ) -> Iterable[PredictionResult]: - """Runs inferences on a batch of text strings. - - Args: - batch: A sequence of examples as text strings. - model: Model returned by Tritonserver. - inference_args: Any additional arguments for an inference. - - Returns: - An Iterable of type PredictionResult. - """ - # Loop each text string, and use a tuple to store the inference results. - logging.info(f"{type(batch)}") - responses = model.infer(inputs={"string_input":batch}) - for response in responses: - predictions = [json.loads(answer) for answer in response.outputs["string_output"].to_string_array().tolist()] - return [PredictionResult(x, y) for x, y in zip(batch, predictions)] + return [PredictionResult(x, y) for x, y in zip(batch, predictions)] + def get_metrics_namespace(self) -> str: + """Returns namespace for metrics.""" + return "BeamML_Triton" diff --git a/sdks/python/apache_beam/ml/inference/triton_inference_test.py b/sdks/python/apache_beam/ml/inference/triton_inference_test.py new file mode 100644 index 000000000000..d23efe503864 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/triton_inference_test.py @@ -0,0 +1,242 @@ +# +# 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. +# + +# pytype: skip-file + +import unittest +from unittest.mock import Mock, MagicMock, patch + +import apache_beam as beam +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +# Protect against environments where tritonserver library is not available. +# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports +try: + from apache_beam.ml.inference.base import PredictionResult, RunInference + from apache_beam.ml.inference.TritonModelHandler import TritonModelHandler + from apache_beam.ml.inference.TritonModelHandler import TritonModelWrapper +except ImportError: + raise unittest.SkipTest('Triton dependencies are not installed') + + +class TritonModelHandlerTest(unittest.TestCase): + def test_handler_initialization(self): + """Test that handler initializes correctly with required parameters.""" + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + self.assertEqual(handler._model_repository, "/workspace/models") + self.assertEqual(handler._model_name, "test_model") + self.assertEqual(handler._input_tensor_name, "INPUT") + self.assertEqual(handler._output_tensor_name, "OUTPUT") + + def test_handler_custom_tensor_names(self): + """Test handler with custom tensor names.""" + handler = TritonModelHandler( + model_repository="/workspace/models", + model_name="test_model", + input_tensor_name="custom_input", + output_tensor_name="custom_output") + self.assertEqual(handler._input_tensor_name, "custom_input") + self.assertEqual(handler._output_tensor_name, "custom_output") + + def test_handler_missing_tritonserver(self): + """Test that handler raises ImportError if tritonserver is not available.""" + with patch('apache_beam.ml.inference.TritonModelHandler.tritonserver', + None): + with self.assertRaises(ImportError): + TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_load_model_success(self, mock_tritonserver): + """Test successful model loading.""" + mock_server = Mock() + mock_model = Mock() + mock_tritonserver.Server.return_value = mock_server + mock_server.model.return_value = mock_model + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + wrapper = handler.load_model() + + self.assertIsInstance(wrapper, TritonModelWrapper) + self.assertEqual(wrapper.server, mock_server) + self.assertEqual(wrapper.model, mock_model) + mock_server.start.assert_called_once() + mock_server.model.assert_called_once_with("test_model") + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_load_model_server_start_fails(self, mock_tritonserver): + """Test model loading when server fails to start.""" + mock_tritonserver.Server.side_effect = Exception("Server start failed") + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + with self.assertRaises(RuntimeError) as context: + handler.load_model() + self.assertIn("Failed to start Triton server", str(context.exception)) + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_load_model_model_not_found(self, mock_tritonserver): + """Test model loading when model is not found.""" + mock_server = Mock() + mock_tritonserver.Server.return_value = mock_server + mock_server.model.return_value = None + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + with self.assertRaises(RuntimeError) as context: + handler.load_model() + self.assertIn("Model 'test_model' not found", str(context.exception)) + mock_server.stop.assert_called_once() + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_run_inference_success(self, mock_tritonserver): + """Test successful inference.""" + # Setup mocks + mock_model = Mock() + mock_server = Mock() + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + + mock_response = Mock() + mock_output = Mock() + mock_output.to_string_array.return_value.tolist.return_value = [ + '{"result": 1}', '{"result": 2}' + ] + mock_response.outputs = {"OUTPUT": mock_output} + mock_model.infer.return_value = [mock_response] + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + batch = ["input1", "input2"] + results = list(handler.run_inference(batch, wrapper)) + + self.assertEqual(len(results), 2) + self.assertIsInstance(results[0], PredictionResult) + mock_model.infer.assert_called_once_with(inputs={"INPUT": batch}) + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_run_inference_with_custom_tensor_names(self, mock_tritonserver): + """Test inference with custom tensor names via inference_args.""" + mock_model = Mock() + mock_server = Mock() + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + + mock_response = Mock() + mock_output = Mock() + mock_output.to_string_array.return_value.tolist.return_value = [ + '{"result": 1}' + ] + mock_response.outputs = {"custom_out": mock_output} + mock_model.infer.return_value = [mock_response] + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + batch = ["input1"] + inference_args = { + 'input_tensor_name': 'custom_in', 'output_tensor_name': 'custom_out' + } + results = list(handler.run_inference(batch, wrapper, inference_args)) + + self.assertEqual(len(results), 1) + mock_model.infer.assert_called_once_with(inputs={"custom_in": batch}) + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_run_inference_fails(self, mock_tritonserver): + """Test inference failure handling.""" + mock_model = Mock() + mock_server = Mock() + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + mock_model.infer.side_effect = Exception("Inference error") + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + batch = ["input1"] + with self.assertRaises(RuntimeError) as context: + list(handler.run_inference(batch, wrapper)) + self.assertIn("Triton inference failed", str(context.exception)) + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_run_inference_output_not_found(self, mock_tritonserver): + """Test error when expected output tensor is not in response.""" + mock_model = Mock() + mock_server = Mock() + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + + mock_response = Mock() + mock_response.outputs = {} # Empty outputs + mock_model.infer.return_value = [mock_response] + + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + + batch = ["input1"] + with self.assertRaises(RuntimeError) as context: + list(handler.run_inference(batch, wrapper)) + self.assertIn("Output tensor 'OUTPUT' not found", str(context.exception)) + + @patch('apache_beam.ml.inference.TritonModelHandler.tritonserver') + def test_custom_parse_function(self, mock_tritonserver): + """Test using a custom output parsing function.""" + def custom_parser(outputs, output_name): + return ["custom_parsed_result"] + + mock_model = Mock() + mock_server = Mock() + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + + mock_response = Mock() + mock_response.outputs = {"OUTPUT": Mock()} + mock_model.infer.return_value = [mock_response] + + handler = TritonModelHandler( + model_repository="/workspace/models", + model_name="test_model", + parse_output_fn=custom_parser) + + batch = ["input1"] + results = list(handler.run_inference(batch, wrapper)) + + self.assertEqual(len(results), 1) + self.assertEqual(results[0].inference, "custom_parsed_result") + + def test_get_metrics_namespace(self): + """Test metrics namespace.""" + handler = TritonModelHandler( + model_repository="/workspace/models", model_name="test_model") + self.assertEqual(handler.get_metrics_namespace(), "BeamML_Triton") + + def test_wrapper_cleanup(self): + """Test that TritonModelWrapper cleans up server on deletion.""" + mock_server = Mock() + mock_model = Mock() + + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + wrapper.__del__() + + mock_server.stop.assert_called_once() + + +if __name__ == '__main__': + unittest.main() From 8d991ce6de1baa4616f43cbcfcb640c315333c22 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 11 Oct 2025 19:29:52 -0400 Subject: [PATCH 3/8] Address review feedback: Fix linting and improve resource cleanup - Remove unused atexit import - Fix logging f-string usage (use lazy % formatting) - Add explicit cleanup() method to TritonModelWrapper for reliable resource cleanup - Improve __del__ to prevent double cleanup - Remove unused imports from test file (MagicMock, beam, TestPipeline, assert_that, equal_to, RunInference) - Add tests for explicit cleanup and idempotent behavior All files now pass pylint 10.00/10 and yapf formatting checks. --- .../ml/inference/TritonModelHandler.py | 36 ++++++++++++++--- .../ml/inference/triton_inference_test.py | 39 +++++++++++++++---- 2 files changed, 62 insertions(+), 13 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py index e670a5d212fc..ad4d264b3b16 100644 --- a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -20,7 +20,6 @@ from typing import Sequence, Dict, Any, Iterable, Optional import logging import json -import atexit from apache_beam.ml.inference.base import ModelHandler, PredictionResult @@ -38,14 +37,37 @@ class TritonModelWrapper: def __init__(self, server: 'Server', model: 'Model'): self.server = server self.model = model + self._cleaned_up = False + + def cleanup(self): + """Explicitly cleanup server resources. + + This method should be called when the model is no longer needed. + It's safe to call multiple times. + """ + if self._cleaned_up: + return - def __del__(self): - """Cleanup server when model is garbage collected.""" try: if self.server: self.server.stop() + self._cleaned_up = True except Exception as e: - LOGGER.warning(f"Error stopping Triton server: {e}") + LOGGER.warning("Error stopping Triton server: %s", e) + raise + + def __del__(self): + """Cleanup server when model is garbage collected. + + Note: __del__ is not guaranteed to be called. Prefer using cleanup() + explicitly when possible. + """ + if not self._cleaned_up: + try: + if self.server: + self.server.stop() + except Exception as e: + LOGGER.warning("Error stopping Triton server in __del__: %s", e) class TritonModelHandler(ModelHandler[Any, PredictionResult, @@ -187,8 +209,10 @@ def run_inference( if len(predictions) != len(batch): LOGGER.warning( - f"Prediction count ({len(predictions)}) doesn't match " - f"batch size ({len(batch)}). Truncating or padding.") + "Prediction count (%d) doesn't match " + "batch size (%d). Truncating or padding.", + len(predictions), + len(batch)) return [PredictionResult(x, y) for x, y in zip(batch, predictions)] diff --git a/sdks/python/apache_beam/ml/inference/triton_inference_test.py b/sdks/python/apache_beam/ml/inference/triton_inference_test.py index d23efe503864..bab96f3a70e6 100644 --- a/sdks/python/apache_beam/ml/inference/triton_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/triton_inference_test.py @@ -18,17 +18,12 @@ # pytype: skip-file import unittest -from unittest.mock import Mock, MagicMock, patch - -import apache_beam as beam -from apache_beam.testing.test_pipeline import TestPipeline -from apache_beam.testing.util import assert_that -from apache_beam.testing.util import equal_to +from unittest.mock import Mock, patch # Protect against environments where tritonserver library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: - from apache_beam.ml.inference.base import PredictionResult, RunInference + from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.TritonModelHandler import TritonModelHandler from apache_beam.ml.inference.TritonModelHandler import TritonModelWrapper except ImportError: @@ -237,6 +232,36 @@ def test_wrapper_cleanup(self): mock_server.stop.assert_called_once() + def test_wrapper_explicit_cleanup(self): + """Test explicit cleanup method.""" + mock_server = Mock() + mock_model = Mock() + + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + wrapper.cleanup() + + mock_server.stop.assert_called_once() + self.assertTrue(wrapper._cleaned_up) + + # Calling cleanup again should not call stop again + wrapper.cleanup() + mock_server.stop.assert_called_once() + + def test_wrapper_cleanup_idempotent(self): + """Test that cleanup after __del__ doesn't cause double cleanup.""" + mock_server = Mock() + mock_model = Mock() + + wrapper = TritonModelWrapper(server=mock_server, model=mock_model) + + # Call cleanup explicitly first + wrapper.cleanup() + mock_server.stop.assert_called_once() + + # __del__ should not call stop again + wrapper.__del__() + mock_server.stop.assert_called_once() + if __name__ == '__main__': unittest.main() From cf06fa605bd61ddb5566b84d399f89c58c25bebb Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 11 Oct 2025 19:37:12 -0400 Subject: [PATCH 4/8] Add Triton dependencies to test configuration - Add tritonserver to ml_test and p312_ml_test extras in setup.py - Create triton_tests_requirements.txt for test-specific dependencies This ensures CI environments have tritonserver installed when running tests. --- .../ml/inference/triton_tests_requirements.txt | 18 ++++++++++++++++++ sdks/python/setup.py | 2 ++ 2 files changed, 20 insertions(+) create mode 100644 sdks/python/apache_beam/ml/inference/triton_tests_requirements.txt diff --git a/sdks/python/apache_beam/ml/inference/triton_tests_requirements.txt b/sdks/python/apache_beam/ml/inference/triton_tests_requirements.txt new file mode 100644 index 000000000000..614b2fac4956 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/triton_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. +# + +tritonserver>=2.41.0 diff --git a/sdks/python/setup.py b/sdks/python/setup.py index c23d69225d52..b4e5050fed08 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -537,6 +537,7 @@ def get_portability_package_data(): 'tf2onnx', 'torch', 'transformers', + 'tritonserver', # Comment out xgboost as it is breaking presubmit python ml # tests due to tag check introduced since pip 24.2 # https://github.com/apache/beam/issues/31285 @@ -556,6 +557,7 @@ def get_portability_package_data(): 'tf2onnx', 'torch', 'transformers', + 'tritonserver', ], 'aws': ['boto3>=1.9,<2'], 'azure': [ From 267a111cb148f15bb0f85b00b58d5a7ff7218d1f Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 11 Oct 2025 23:12:14 -0400 Subject: [PATCH 5/8] Fix import ordering and remove unused type: ignore - Reorder imports: stdlib (json, logging) before typing before apache_beam - Separate each typing import on its own line - Separate base imports on individual lines - Remove unused type: ignore comment on line 30 - Apply isort formatting for consistency --- sdks/python/apache_beam/ml/inference/TritonModelHandler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py index ad4d264b3b16..ca26fd2b9ad7 100644 --- a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -17,9 +17,9 @@ """Apache Beam ModelHandler implementation for Triton Inference Server.""" -from typing import Sequence, Dict, Any, Iterable, Optional -import logging import json +import logging +from typing import Any, Dict, Iterable, Optional, Sequence from apache_beam.ml.inference.base import ModelHandler, PredictionResult @@ -27,7 +27,7 @@ import tritonserver from tritonserver import Model, Server except ImportError: - tritonserver = None # type: ignore + tritonserver = None LOGGER = logging.getLogger(__name__) From 0eacbd1e7f3bf00a3d1eeca6225a211170d8da17 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 11 Oct 2025 23:13:43 -0400 Subject: [PATCH 6/8] Address Gemini code review feedback - Set _cleaned_up flag in __del__ to ensure idempotent cleanup - Use specific exception types (JSONDecodeError, TypeError, AttributeError) instead of broad Exception catch in JSON parsing - Improves resource management reliability and error handling precision --- sdks/python/apache_beam/ml/inference/TritonModelHandler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py index ca26fd2b9ad7..5c599f9c6665 100644 --- a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -66,6 +66,7 @@ def __del__(self): try: if self.server: self.server.stop() + self._cleaned_up = True except Exception as e: LOGGER.warning("Error stopping Triton server in __del__: %s", e) @@ -198,7 +199,7 @@ def run_inference( json.loads(val) for val in output_tensor.to_string_array().tolist() ] - except Exception: + except (json.JSONDecodeError, TypeError, AttributeError): # If JSON parsing fails, return raw output parsed = output_tensor.to_bytes_array().tolist() From aab0d2f374103a98cca9536f1811c50baa82e088 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sat, 11 Oct 2025 23:27:29 -0400 Subject: [PATCH 7/8] Add Triton Inference Server entry to CHANGES.md --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 07f3b4f5accc..59794ca8f777 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -78,6 +78,7 @@ * Python examples added for CloudSQL enrichment handler on [Beam website](https://beam.apache.org/documentation/transforms/python/elementwise/enrichment-cloudsql/) (Python) ([#35473](https://github.com/apache/beam/issues/36095)). * Support for batch mode execution in WriteToPubSub transform added (Python) ([#35990](https://github.com/apache/beam/issues/35990)). * Added official support for Python 3.13 ([#34869](https://github.com/apache/beam/issues/34869)). +* Added Triton Inference Server ModelHandler for ML inference (Python) ([#36369](https://github.com/apache/beam/issues/36369)). ## Breaking Changes From 65425b5d9a80db3d6e2b702000fcc43f9f0e1aa7 Mon Sep 17 00:00:00 2001 From: Sai Shashank Mudliar Date: Sun, 12 Oct 2025 00:08:46 -0400 Subject: [PATCH 8/8] Fix Callable type hint and add CHANGES.md entry - Change Optional[callable] to Optional[Callable] in TritonModelHandler - Add Triton Inference Server entry to CHANGES.md --- sdks/python/apache_beam/ml/inference/TritonModelHandler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py index 5c599f9c6665..fc83d0aa8202 100644 --- a/sdks/python/apache_beam/ml/inference/TritonModelHandler.py +++ b/sdks/python/apache_beam/ml/inference/TritonModelHandler.py @@ -19,7 +19,7 @@ import json import logging -from typing import Any, Dict, Iterable, Optional, Sequence +from typing import Any, Callable, Dict, Iterable, Optional, Sequence from apache_beam.ml.inference.base import ModelHandler, PredictionResult @@ -103,7 +103,7 @@ def __init__( model_name: str, input_tensor_name: str = "INPUT", output_tensor_name: str = "OUTPUT", - parse_output_fn: Optional[callable] = None, + parse_output_fn: Optional[Callable] = None, ): if tritonserver is None: raise ImportError(