This is an automated email from the ASF dual-hosted git repository.
chamikaramj pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 98fe1763084 Modifies Hugging Face model handlers to not use ML
packages... (#40541)
98fe1763084 is described below
commit 98fe1763084e1985a86f91a168fb42f547b7ea22
Author: Chamikara Jayalath <[email protected]>
AuthorDate: Tue Oct 13 08:50:42 2026 -0700
Modifies Hugging Face model handlers to not use ML packages... (#40541)
---
.../ml/inference/huggingface_inference.py | 126 ++++++++++++-----
.../ml/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 a9893ea9290..93ffc3fc1f5 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 Callable
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 @@ __all__ = [
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 @@ TensorInferenceFn = Callable[
]
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 @@ KeyedTensorInferenceFn = Callable[[
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 @@ class
HuggingFaceModelHandlerKeyedTensor(ModelHandler[dict[str,
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 @@ class
HuggingFaceModelHandlerTensor(ModelHandler[Union[tf.Tensor, torch.Tensor],
"""
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 @@ class HuggingFacePipelineModelHandler(ModelHandler[str,
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 @@ class HuggingFacePipelineModelHandler(ModelHandler[str,
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 @@ class HuggingFacePipelineModelHandler(ModelHandler[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 @@ class HuggingFacePipelineModelHandler(ModelHandler[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 5cd55f1b1f5..f78b72aec51 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 Optional
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 @@ try:
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
[email protected](tf is None, 'Transformers dependencies are not installed.')
class HuggingFaceInferenceTest(unittest.TestCase):
def setUp(self) -> None:
self.tmpdir = tempfile.mkdtemp()
@@ -152,5 +166,132 @@ class HuggingFaceInferenceTest(unittest.TestCase):
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
+
+
[email protected]
+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 c4fc179b9a3..0b6be61a5d0 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 @@ class MLTransformJobSubmissionTest(unittest.TestCase):
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()
