From 193bfc350196b4c8e36345a44b1485dcc1519be3 Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Mon, 12 Oct 2026 14:33:50 -0700 Subject: [PATCH] Modifies Hugging Face model handlers to not use ML packages during job submission. --- .../ml/inference/huggingface_inference.py | 126 +++++++++----- .../inference/huggingface_inference_test.py | 157 +++++++++++++++++- sdks/python/apache_beam/yaml/yaml_ml_test.py | 39 +++++ 3 files changed, 276 insertions(+), 46 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py b/sdks/python/apache_beam/ml/inference/huggingface_inference.py index a9893ea9290c..93ffc3fc1f52 100644 --- a/sdks/python/apache_beam/ml/inference/huggingface_inference.py +++ b/sdks/python/apache_beam/ml/inference/huggingface_inference.py @@ -17,6 +17,9 @@ # pytype: skip-file +# Support forward references +from __future__ import annotations + import logging import sys from collections import defaultdict @@ -24,21 +27,27 @@ from collections.abc import Iterable from collections.abc import Sequence from enum import Enum +from typing import TYPE_CHECKING from typing import Any from typing import Optional from typing import Union -import tensorflow as tf -import torch -from transformers import AutoModel -from transformers import Pipeline -from transformers import TFAutoModel -from transformers import pipeline - from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult -from apache_beam.ml.inference.pytorch_inference import _convert_to_device + +# tensorflow, torch and transformers are intentionally not imported at module +# import time. They are only needed for loading models and running inference +# when the pipeline is executed (for example, on the workers), and should not +# be required for constructing and submitting pipelines. +# Following imports are kept for lint but not actually imported by Python +# during execution. +if TYPE_CHECKING: + import tensorflow as tf + import torch + from transformers import AutoModel + from transformers import Pipeline + from transformers import TFAutoModel _LOGGER = logging.getLogger(__name__) @@ -50,8 +59,10 @@ TensorInferenceFn = Callable[ [ - Sequence[Union[torch.Tensor, tf.Tensor]], - Union[AutoModel, TFAutoModel], + # Type aliases are evaluated at import time, so framework types are + # quoted (as forward references) to avoid importing them. + Sequence[Union['torch.Tensor', 'tf.Tensor']], + Union['AutoModel', 'TFAutoModel'], str, Optional[dict[str, Any]], Optional[str], @@ -60,8 +71,8 @@ ] KeyedTensorInferenceFn = Callable[[ - Sequence[dict[str, Union[torch.Tensor, tf.Tensor]]], - Union[AutoModel, TFAutoModel], + Sequence[dict[str, Union['torch.Tensor', 'tf.Tensor']]], + Union['AutoModel', 'TFAutoModel'], str, Optional[dict[str, Any]], Optional[str] @@ -69,7 +80,7 @@ Iterable[PredictionResult]] PipelineInferenceFn = Callable[ - [Sequence[str], Pipeline, Optional[dict[str, Any]]], + [Sequence[str], 'Pipeline', Optional[dict[str, Any]]], Iterable[PredictionResult]] @@ -135,6 +146,7 @@ def no_gpu_available_warning(): def is_gpu_available_torch(): + import torch if torch.cuda.is_available(): return True else: @@ -143,12 +155,14 @@ def is_gpu_available_torch(): def get_device_torch(device): + import torch if device == "GPU" and is_gpu_available_torch(): return torch.device("cuda") return torch.device("cpu") def is_gpu_available_tensorflow(device): + import tensorflow as tf gpu_devices = tf.config.list_physical_devices(device) if len(gpu_devices) == 0: no_gpu_available_warning() @@ -156,6 +170,27 @@ def is_gpu_available_tensorflow(device): return True +def _is_tensorflow_tensor(value) -> bool: + # A TensorFlow tensor can only exist if tensorflow was already imported. + # Checking sys.modules, instead of importing tensorflow, avoids requiring + # tensorflow to be installed (and the cost of importing it) for PyTorch. + tensor_type = getattr(sys.modules.get('tensorflow'), 'Tensor', None) + return tensor_type is not None and isinstance(value, tensor_type) + + +def _import_transformers_pipeline() -> Callable[..., Pipeline]: + """Returns transformers.pipeline, importing it lazily.""" + try: + from transformers import pipeline + except ImportError as e: + raise ImportError( + 'transformers is required to load Hugging Face pipelines, e.g. for ' + 'HuggingFacePipelineModelHandler. Please install it using ' + '`pip install transformers` in the environment where the pipeline is ' + 'executed (for example, on the workers).') from e + return pipeline + + def _validate_constructor_args_hf_pipeline(task, model): if not task and not model: raise RuntimeError( @@ -170,6 +205,9 @@ def _run_inference_torch_keyed_tensor( device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: + import torch + + from apache_beam.ml.inference.pytorch_inference import _convert_to_device device = get_device_torch(device) key_to_tensor_list = defaultdict(list) # torch.no_grad() mitigates GPU memory issues @@ -193,6 +231,7 @@ def _run_inference_tensorflow_keyed_tensor( device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: + import tensorflow as tf if device == "GPU": is_gpu_available_tensorflow(device) key_to_tensor_list = defaultdict(list) @@ -207,12 +246,13 @@ def _run_inference_tensorflow_keyed_tensor( return utils._convert_to_result(batch, predictions, model_id) -class HuggingFaceModelHandlerKeyedTensor(ModelHandler[dict[str, - Union[tf.Tensor, - torch.Tensor]], - PredictionResult, - Union[AutoModel, - TFAutoModel]]): +# Unlike annotations, generic base class arguments are evaluated when the class +# is defined, so framework types are quoted (as forward references) to avoid +# importing them. Type checkers resolve them via the TYPE_CHECKING imports. +class HuggingFaceModelHandlerKeyedTensor( + ModelHandler[dict[str, Union['tf.Tensor', 'torch.Tensor']], + PredictionResult, + Union['AutoModel', 'TFAutoModel']]): def __init__( self, model_uri: str, @@ -305,6 +345,7 @@ def load_model(self): self._model_uri, **self._model_config_args) if self._framework == 'pt': if self._device == "GPU" and is_gpu_available_torch(): + import torch model.to(torch.device("cuda")) if callable(getattr(model, 'requires_grad_', None)): model.requires_grad_(False) @@ -377,6 +418,9 @@ def _default_inference_fn_torch( device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: + import torch + + from apache_beam.ml.inference.pytorch_inference import _convert_to_device device = get_device_torch(device) # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 @@ -393,6 +437,7 @@ def _default_inference_fn_tensorflow( device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: + import tensorflow as tf if device == "GPU": is_gpu_available_tensorflow(device) batched_tensors = tf.stack(batch, axis=0) @@ -400,10 +445,11 @@ def _default_inference_fn_tensorflow( return utils._convert_to_result(batch, predictions, model_id) -class HuggingFaceModelHandlerTensor(ModelHandler[Union[tf.Tensor, torch.Tensor], +class HuggingFaceModelHandlerTensor(ModelHandler[Union['tf.Tensor', + 'torch.Tensor'], PredictionResult, - Union[AutoModel, - TFAutoModel]]): + Union['AutoModel', + 'TFAutoModel']]): def __init__( self, model_uri: str, @@ -525,13 +571,14 @@ def run_inference( """ inference_args = {} if not inference_args else inference_args if not self._framework: - if isinstance(batch[0], tf.Tensor): + if _is_tensorflow_tensor(batch[0]): self._framework = "tf" else: self._framework = "pt" if (self._framework == "pt" and self._device == "GPU" and is_gpu_available_torch()): + import torch model.to(torch.device("cuda")) if self._inference_fn: @@ -576,7 +623,7 @@ def _default_pipeline_inference_fn( class HuggingFacePipelineModelHandler(ModelHandler[str, PredictionResult, - Pipeline]): + 'Pipeline']): def __init__( self, task: Union[str, PipelineTask] = "", @@ -619,9 +666,11 @@ def __init__( device (str): the device (`"CPU"` or `"GPU"`) on which you wish to run the pipeline. Defaults to GPU. If GPU is not available then it falls - back to CPU. You can also use advanced option like `device_map` with - key-value pair as you would do in the usual Hugging Face pipeline using - `load_pipeline_args`. Ex: load_pipeline_args={'device_map':auto}). + back to CPU. Whether a GPU is available is determined when the model + is loaded, i.e. on the workers that execute the pipeline. You can also + use advanced option like `device_map` with key-value pair as you would + do in the usual Hugging Face pipeline using `load_pipeline_args`. + Ex: load_pipeline_args={'device_map':auto}). inference_fn: the inference function to use during RunInference. Default is _default_pipeline_inference_fn. load_pipeline_args (dict[str, Any]): keyword arguments to provide load @@ -647,7 +696,9 @@ def __init__( before loading the model. **Supported Versions:** HuggingFacePipelineModelHandler supports - transformers>=4.18.0. + transformers>=4.18.0. transformers and torch are only needed in the + environment where the pipeline is executed (for example, on the workers), + not for constructing and submitting the pipeline. """ super().__init__( min_batch_size=min_batch_size, @@ -681,14 +732,9 @@ def _deduplicate_device_value(self, device: Optional[str]): if 'device' not in self._load_pipeline_args: if current_device == 'CPU': self._load_pipeline_args['device'] = 'cpu' - else: - if is_gpu_available_torch(): - self._load_pipeline_args['device'] = 'cuda:0' - else: - _LOGGER.warning( - "HuggingFaceModelHandler specified a 'GPU' device, " - "but GPUs are not available. Switching to CPU.") - self._load_pipeline_args['device'] = 'cpu' + # Otherwise, a GPU is used if one is available. This is determined in + # load_model(), i.e. on the workers that execute the pipeline, instead of + # on the machine that constructs and submits the pipeline. else: if current_device: raise ValueError( @@ -697,8 +743,12 @@ def _deduplicate_device_value(self, device: Optional[str]): def load_model(self): """Loads and initializes the pipeline for processing.""" - return pipeline( - task=self._task, model=self._model, **self._load_pipeline_args) + pipeline = _import_transformers_pipeline() + load_pipeline_args = dict(self._load_pipeline_args) + if 'device' not in load_pipeline_args: + load_pipeline_args['device'] = ( + 'cuda:0' if is_gpu_available_torch() else 'cpu') + return pipeline(task=self._task, model=self._model, **load_pipeline_args) def run_inference( self, diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference_test.py b/sdks/python/apache_beam/ml/inference/huggingface_inference_test.py index 5cd55f1b1f5f..f78b72aec51e 100644 --- a/sdks/python/apache_beam/ml/inference/huggingface_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/huggingface_inference_test.py @@ -17,8 +17,12 @@ # pytype: skip-file +import contextlib +import importlib.util import shutil +import sys import tempfile +import types import unittest from collections.abc import Iterable from collections.abc import Sequence @@ -27,11 +31,16 @@ from typing import Union import pytest +from parameterized import parameterized +import apache_beam as beam +from apache_beam.ml.inference import huggingface_inference from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult -from apache_beam.ml.inference.tensorflow_inference_test import FakeTFTensorModel -from apache_beam.ml.inference.tensorflow_inference_test import _compare_tensor_prediction_result +from apache_beam.ml.inference.base import RunInference +from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerTensor +from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler +from apache_beam.testing import test_utils # pylint: disable=ungrouped-imports try: @@ -40,14 +49,18 @@ from transformers import AutoModel from transformers import TFAutoModel - from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerTensor -except ImportError: - raise unittest.SkipTest('Transformers dependencies are not installed.') + from apache_beam.ml.inference.tensorflow_inference_test import FakeTFTensorModel + from apache_beam.ml.inference.tensorflow_inference_test import _compare_tensor_prediction_result +except (ImportError, unittest.SkipTest): + # tensorflow_inference_test raises SkipTest if its dependencies are missing. + tf = None + +_ML_FRAMEWORKS = ('tensorflow', 'torch', 'transformers') def fake_inference_fn_tensor( - batch: Sequence[Union[tf.Tensor, torch.Tensor]], - model: Union[AutoModel, TFAutoModel], + batch: Sequence[Union['tf.Tensor', 'torch.Tensor']], + model: Union['AutoModel', 'TFAutoModel'], device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: @@ -56,11 +69,12 @@ def fake_inference_fn_tensor( class FakeTorchModel: - def predict(self, input: torch.Tensor): + def predict(self, input: 'torch.Tensor'): return input @pytest.mark.uses_transformers +@unittest.skipIf(tf is None, 'Transformers dependencies are not installed.') class HuggingFaceInferenceTest(unittest.TestCase): def setUp(self) -> None: self.tmpdir = tempfile.mkdtemp() @@ -152,5 +166,132 @@ def test_convert_to_result_batch_processing(self): self.assertEqual(results[1].inference, {"translation_text": "output2"}) +class _FakePipeline: + """A stand-in for a transformers Pipeline, created by a fake pipeline().""" + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def _fake_transformers_module(): + module = types.ModuleType('transformers') + module.pipeline = _FakePipeline + return module + + +def _fake_torch_module(cuda_available): + module = types.ModuleType('torch') + module.cuda = types.SimpleNamespace(is_available=lambda: cuda_available) + return module + + +@contextlib.contextmanager +def _fake_modules(**fake_modules): + """Replaces the given modules with fake ones.""" + # Also hides the real modules, if already imported, and restores them + # afterwards. + with test_utils.block_imports(*fake_modules): + sys.modules.update(fake_modules) + try: + yield + finally: + for name in fake_modules: + del sys.modules[name] + + +class HuggingFaceJobSubmissionTest(unittest.TestCase): + """Tests that ML frameworks are only used when executing pipelines. + + tensorflow, torch and transformers should only be needed for loading models + and running inference when the pipeline is executed (for example, on the + workers), but not for constructing and submitting the pipeline. These tests + don't require any of them to be installed. + """ + def test_importing_module_does_not_import_ml_frameworks(self): + # Executes the module in a new module object, i.e. without replacing the + # already imported module. + spec = importlib.util.find_spec(huggingface_inference.__name__) + with test_utils.block_imports(*_ML_FRAMEWORKS) as attempted_imports: + spec.loader.exec_module(importlib.util.module_from_spec(spec)) + self.assertEqual(attempted_imports, []) + + @parameterized.expand([('default', None), ('gpu', 'GPU'), ('cpu', 'CPU')]) + def test_pipeline_job_submission_does_not_use_ml_frameworks( + self, unused_name, device): + with test_utils.block_imports(*_ML_FRAMEWORKS) as attempted_imports: + model_handler = HuggingFacePipelineModelHandler( + task='text-classification', + model='distilbert-base-uncased-finetuned-sst-2-english', + device=device) + pipeline = beam.Pipeline() + _ = ( + pipeline + | beam.Create(['I love Apache Beam!']) + | RunInference(model_handler)) + # Serializes the pipeline, as done during job submission, and + # deserializes it again. Deserializing fails if the pipeline references + # (i.e. requires) a module that can't be imported. + beam.Pipeline.from_runner_api( + pipeline.to_runner_api(), pipeline.runner, pipeline.options) + self.assertEqual(attempted_imports, []) + + @parameterized.expand([ + ('default_with_gpu', None, True, 'cuda:0'), + ('default_without_gpu', None, False, 'cpu'), + ('gpu_with_gpu', 'GPU', True, 'cuda:0'), + ('gpu_without_gpu', 'GPU', False, 'cpu'), + ('cpu_with_gpu', 'CPU', True, 'cpu'), + ]) + def test_pipeline_device_is_determined_when_loading_model( + self, unused_name, device, cuda_available, expected_device): + # Whether a GPU is available is determined on the workers that load the + # model, not on the machine that constructs the model handler (and submits + # the pipeline). So the GPU availability is inverted for the latter. + with _fake_modules(torch=_fake_torch_module(not cuda_available)): + model_handler = HuggingFacePipelineModelHandler( + task='text-classification', device=device) + with _fake_modules(transformers=_fake_transformers_module(), + torch=_fake_torch_module(cuda_available)): + pipeline = model_handler.load_model() + self.assertEqual( + pipeline.kwargs, { + 'task': 'text-classification', + 'model': '', + 'device': expected_device + }) + + def test_pipeline_load_pipeline_args_device_is_used(self): + model_handler = HuggingFacePipelineModelHandler( + task='text-classification', load_pipeline_args={'device': 'cuda:1'}) + with _fake_modules(transformers=_fake_transformers_module()): + with test_utils.block_imports('torch') as attempted_imports: + pipeline = model_handler.load_model() + self.assertEqual(pipeline.kwargs['device'], 'cuda:1') + self.assertEqual(attempted_imports, []) + + def test_pipeline_invalid_device(self): + with self.assertRaisesRegex(ValueError, 'Invalid device value'): + HuggingFacePipelineModelHandler(task='text-classification', device='TPU') + + def test_pipeline_device_specified_twice(self): + with self.assertRaisesRegex(ValueError, + 'specified in `load_pipeline_args`'): + HuggingFacePipelineModelHandler( + task='text-classification', + device='GPU', + load_pipeline_args={'device': 'cuda:0'}) + + def test_pipeline_loading_model_requires_transformers(self): + model_handler = HuggingFacePipelineModelHandler( + task='text-classification', device='CPU') + with test_utils.block_imports('transformers'): + with self.assertRaisesRegex(ImportError, 'transformers is required'): + model_handler.load_model() + + def test_framework_detection_does_not_import_tensorflow(self): + with test_utils.block_imports('tensorflow') as attempted_imports: + self.assertFalse(huggingface_inference._is_tensorflow_tensor([1])) + self.assertEqual(attempted_imports, []) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/yaml/yaml_ml_test.py b/sdks/python/apache_beam/yaml/yaml_ml_test.py index c4fc179b9a3c..0b6be61a5d09 100644 --- a/sdks/python/apache_beam/yaml/yaml_ml_test.py +++ b/sdks/python/apache_beam/yaml/yaml_ml_test.py @@ -274,6 +274,45 @@ def test_sentence_transformer_embedding_does_not_use_sentence_transformers( self.assertEqual(attempted_imports, []) +class RunInferenceJobSubmissionTest(unittest.TestCase): + def test_huggingface_pipeline_model_handler_does_not_use_ml_frameworks(self): + # tensorflow, torch and transformers should only be needed when executing + # the pipeline (for example, on the workers), not for constructing and + # submitting it. + # Imports the YAML ML module before blocking imports, since importing it + # loads all available MLTransform transforms, and some of these (e.g. the + # TFT based ones) do require tensorflow. + from apache_beam.yaml import yaml_ml # pylint: disable=unused-import + ml_opts = beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle', yaml_experimental_features=['ML']) + with test_utils.block_imports('tensorflow', 'torch', + 'transformers') as attempted_imports: + pipeline = beam.Pipeline(options=ml_opts) + _ = pipeline | YamlTransform( + ''' + type: chain + transforms: + - type: Create + config: + elements: + - {text: "I love Apache Beam!"} + - type: RunInference + config: + model_handler: + type: HuggingFacePipelineModelHandler + config: + task: text-classification + preprocess: + callable: 'lambda row: row.text' + ''') + # Serializes the pipeline, as done during job submission, and + # deserializes it again. Deserializing fails if the pipeline references + # (i.e. requires) a module that can't be imported. + beam.Pipeline.from_runner_api( + pipeline.to_runner_api(), pipeline.runner, pipeline.options) + self.assertEqual(attempted_imports, []) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main()