diff --git a/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py b/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py index e17a5d66fe91..8ad34ca690b6 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py @@ -32,11 +32,6 @@ from apache_beam.ml.transforms.base import _TextEmbeddingHandler from apache_beam.ml.transforms.embeddings.huggingface import _SentenceTransformerModelHandler -try: - from sentence_transformers import SentenceTransformer -except ImportError: - SentenceTransformer = None - try: from PIL import Image as PILImage except ImportError: @@ -66,20 +61,15 @@ def __init__( ``model.encode()`` (e.g. ``normalize_embeddings``). """ - if not SentenceTransformer: - raise ImportError( - "sentence-transformers is required to use " - "HuggingfaceTextEmbeddings." - "Please install it with using `pip install sentence-transformers`.") super().__init__(type_adapter=create_text_adapter(), **kwargs) self.model_name = model_name self.max_seq_length = max_seq_length - self.model_class = SentenceTransformer def get_model_handler(self): """Returns model handler configured with RAG adapter.""" + # model_class is not set, so that sentence-transformers is only imported + # when the model is loaded during pipeline execution. return _SentenceTransformerModelHandler( - model_class=self.model_class, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, @@ -165,11 +155,6 @@ def __init__( ``model.encode()`` (e.g. ``normalize_embeddings``). """ - if not SentenceTransformer: - raise ImportError( - "sentence-transformers is required to use " - "HuggingfaceImageEmbeddings. " - "Please install it with `pip install sentence-transformers`.") if not PILImage: raise ImportError( "Pillow is required to use HuggingfaceImageEmbeddings. " @@ -177,12 +162,12 @@ def __init__( super().__init__(type_adapter=_create_hf_image_adapter(), **kwargs) self.model_name = model_name self.max_seq_length = max_seq_length - self.model_class = SentenceTransformer def get_model_handler(self): """Returns model handler configured with RAG adapter.""" + # model_class is not set, so that sentence-transformers is only imported + # when the model is loaded during pipeline execution. return _SentenceTransformerModelHandler( - model_class=self.model_class, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, diff --git a/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py b/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py index 65c21771524d..de93df956270 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py @@ -16,6 +16,7 @@ """Tests for apache_beam.ml.rag.embeddings.huggingface.""" +import importlib.util import io import os import shutil @@ -33,6 +34,7 @@ from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.transforms.base import MLTransform +from apache_beam.testing import test_utils from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @@ -234,6 +236,60 @@ def test_image_embedding_pipeline_from_path(self): shutil.rmtree(artifact_location) +class HuggingfaceEmbeddingsJobSubmissionTest(unittest.TestCase): + """Tests that sentence-transformers is only used when executing pipelines. + + sentence-transformers should only be needed for loading the model when the + pipeline is executed (for example, on the workers), but not for constructing + and submitting the pipeline. These tests don't require sentence-transformers + to be installed. + """ + def setUp(self): + self.artifact_location = tempfile.mkdtemp(prefix='hf_job_submission_') + + def tearDown(self): + shutil.rmtree(self.artifact_location) + + def _construct_and_serialize_pipeline(self, items, embedder): + pipeline = beam.Pipeline() + _ = ( + pipeline + | beam.Create(items) + | MLTransform(write_artifact_location=self.artifact_location). + with_transform(embedder)) + # 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) + + def test_importing_module_does_not_import_sentence_transformers(self): + # Executes the module in a new module object, i.e. without replacing the + # already imported module. + spec = importlib.util.find_spec(HuggingfaceTextEmbeddings.__module__) + with test_utils.block_imports('sentence_transformers') as attempted_imports: + spec.loader.exec_module(importlib.util.module_from_spec(spec)) + self.assertEqual(attempted_imports, []) + + def test_text_embeddings_job_submission_does_not_use_sentence_transformers( + self): + items = [Chunk(content=Content(text="This is a test sentence."), id="1")] + with test_utils.block_imports('sentence_transformers') as attempted_imports: + embedder = HuggingfaceTextEmbeddings( + model_name="sentence-transformers/all-MiniLM-L6-v2") + self._construct_and_serialize_pipeline(items, embedder) + self.assertEqual(attempted_imports, []) + + @unittest.skipIf(not PIL_AVAILABLE, "PIL not available") + def test_image_embeddings_job_submission_does_not_use_sentence_transformers( + self): + items = [EmbeddableItem.from_image(_create_png_bytes(), id="img1")] + with test_utils.block_imports('sentence_transformers') as attempted_imports: + embedder = HuggingfaceImageEmbeddings(model_name="clip-ViT-B-32") + self._construct_and_serialize_pipeline(items, embedder) + self.assertEqual(attempted_imports, []) + + def _create_png_bytes(): """Create a small valid RGB PNG image.""" from PIL import Image diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py index a6905c375451..c3eeccf2a0f0 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py @@ -21,6 +21,7 @@ from collections.abc import Callable from collections.abc import Mapping from collections.abc import Sequence +from typing import TYPE_CHECKING from typing import Any from typing import Optional @@ -33,14 +34,29 @@ from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler -try: +# sentence-transformers is intentionally not imported at module import time. +# It is only needed for loading the model when the pipeline is executed (for +# example, on the workers) and should not be required for constructing and +# submitting pipelines. +if TYPE_CHECKING: from sentence_transformers import SentenceTransformer -except ImportError: - SentenceTransformer = None _LOGGER = logging.getLogger(__name__) +def _import_sentence_transformer() -> Callable: + """Returns sentence_transformers.SentenceTransformer, importing it lazily.""" + try: + from sentence_transformers import SentenceTransformer + except ImportError as e: + raise ImportError( + "sentence-transformers is required to load SentenceTransformer " + "models, e.g. for SentenceTransformerEmbeddings. Please install it " + "using `pip install sentence-transformers` in the environment where " + "the pipeline is executed (for example, on the workers).") from e + return SentenceTransformer + + # TODO: https://github.com/apache/beam/issues/29621 # Use HuggingFaceModelHandlerTensor once the import issue is fixed. # Right now, the hugging face model handler import torch and tensorflow @@ -52,7 +68,7 @@ class _SentenceTransformerModelHandler(ModelHandler): def __init__( self, model_name: str, - model_class: Callable, + model_class: Optional[Callable] = None, load_model_args: Optional[dict] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, @@ -61,6 +77,9 @@ def __init__( **kwargs): self._max_seq_length = max_seq_length self.model_name = model_name + # If not set, sentence_transformers.SentenceTransformer is used. It is + # imported in load_model(), i.e. only when the pipeline is executed, so that + # sentence-transformers is not required for job submission. self._model_class = model_class self._load_model_args = load_model_args self._min_batch_size = min_batch_size @@ -68,23 +87,25 @@ def __init__( self._large_model = large_model self._kwargs = kwargs - if not SentenceTransformer: - raise ImportError( - "sentence-transformers is required to use " - "SentenceTransformerEmbeddings." - "Please install it with using `pip install sentence-transformers`.") - def run_inference( self, batch: Sequence[str], - model: SentenceTransformer, + model: 'SentenceTransformer', inference_args: Optional[dict[str, Any]] = None, ): inference_args = inference_args or {} return model.encode(batch, **inference_args) def load_model(self): - model = self._model_class(self.model_name, **self._load_model_args) + model_class = self._model_class + # Checking callable() instead of None also covers artifacts written by + # Beam 2.77.0 and earlier, which stored a reference to the + # SentenceTransformer class. jsonpickle restores such a reference as a + # non-callable placeholder if sentence-transformers was not installed when + # the artifacts were read (for example, during job submission). + if not callable(model_class): + model_class = _import_sentence_transformer() + model = model_class(self.model_name, **(self._load_model_args or {})) if self._max_seq_length: model.max_seq_length = self._max_seq_length return model @@ -139,8 +160,9 @@ def __init__( self.image_model = image_model def get_model_handler(self): + # model_class is not set, so that sentence-transformers is only imported + # when the model is loaded during pipeline execution. return _SentenceTransformerModelHandler( - model_class=SentenceTransformer, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py index 5e56dc710548..12990cb4c2c8 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py @@ -14,12 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib +import importlib.util import os import shutil +import sys import tempfile +import types import unittest import uuid +import jsonpickle import numpy as np import pytest from parameterized import parameterized @@ -28,6 +33,8 @@ from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.embeddings import huggingface +from apache_beam.testing import test_utils from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @@ -330,6 +337,141 @@ def test_sentence_transformer_images_with_str_data_types(self): embedding_config)) +class _FakeSentenceTransformer: + """A stand-in for sentence_transformers.SentenceTransformer.""" + def __init__(self, model_name, **kwargs): + self.model_name = model_name + self.kwargs = kwargs + self.max_seq_length = None + + +@contextlib.contextmanager +def _fake_sentence_transformers_module(): + """Replaces the sentence_transformers module with a fake one.""" + fake_module = types.ModuleType('sentence_transformers') + fake_module.SentenceTransformer = _FakeSentenceTransformer + # Also hides the real sentence_transformers modules, if already imported, and + # restores them afterwards. + with test_utils.block_imports('sentence_transformers'): + sys.modules['sentence_transformers'] = fake_module + try: + yield + finally: + del sys.modules['sentence_transformers'] + + +def _serialize_and_deserialize(pipeline): + """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.""" + return beam.Pipeline.from_runner_api( + pipeline.to_runner_api(), pipeline.runner, pipeline.options) + + +class SentenceTransformerJobSubmissionTest(unittest.TestCase): + """Tests that sentence-transformers is only used when executing pipelines. + + sentence-transformers (and its heavy dependencies, like torch) should only be + needed for loading the model when the pipeline is executed (for example, on + the workers), but not for constructing and submitting the pipeline. These + tests don't require sentence-transformers to be installed. + """ + def setUp(self): + self.artifact_location = tempfile.mkdtemp(prefix='sentence_transformers_') + + def tearDown(self): + shutil.rmtree(self.artifact_location) + + def _apply_ml_transform(self, pipeline, ml_transform): + _ = ( + pipeline + | beam.Create([{ + test_query_column: test_query + }]) + | ml_transform) + + def test_importing_module_does_not_import_sentence_transformers(self): + # Executes the module in a new module object, i.e. without replacing the + # already imported module. + spec = importlib.util.find_spec(huggingface.__name__) + with test_utils.block_imports('sentence_transformers') as attempted_imports: + spec.loader.exec_module(importlib.util.module_from_spec(spec)) + self.assertEqual(attempted_imports, []) + + @parameterized.expand([('text_model', False), ('image_model', True)]) + def test_job_submission_does_not_use_sentence_transformers( + self, unused_name, image_model): + with test_utils.block_imports('sentence_transformers') as attempted_imports: + embedding_config = huggingface.SentenceTransformerEmbeddings( + model_name=DEFAULT_MODEL_NAME, + columns=[test_query_column], + image_model=image_model) + pipeline = beam.Pipeline() + self._apply_ml_transform( + pipeline, + MLTransform( + write_artifact_location=self.artifact_location).with_transform( + embedding_config)) + _serialize_and_deserialize(pipeline) + self.assertEqual(attempted_imports, []) + + def test_reading_artifacts_does_not_use_sentence_transformers(self): + with test_utils.block_imports('sentence_transformers') as attempted_imports: + embedding_config = huggingface.SentenceTransformerEmbeddings( + model_name=DEFAULT_MODEL_NAME, columns=[test_query_column]) + # Artifacts are written when constructing the pipeline. + self._apply_ml_transform( + beam.Pipeline(), + MLTransform( + write_artifact_location=self.artifact_location).with_transform( + embedding_config)) + pipeline = beam.Pipeline() + self._apply_ml_transform( + pipeline, MLTransform(read_artifact_location=self.artifact_location)) + _serialize_and_deserialize(pipeline) + self.assertEqual(attempted_imports, []) + + def test_loading_model_requires_sentence_transformers(self): + model_handler = huggingface.SentenceTransformerEmbeddings( + model_name=DEFAULT_MODEL_NAME, + columns=[test_query_column]).get_model_handler() + with test_utils.block_imports('sentence_transformers'): + with self.assertRaisesRegex(ImportError, + 'sentence-transformers is required'): + model_handler.load_model() + + def test_loading_model_uses_sentence_transformers(self): + model_handler = huggingface.SentenceTransformerEmbeddings( + model_name=DEFAULT_MODEL_NAME, + columns=[test_query_column], + max_seq_length=128, + load_model_args={ + 'device': 'cpu' + }).get_model_handler() + with _fake_sentence_transformers_module(): + model = model_handler.load_model() + self.assertIsInstance(model, _FakeSentenceTransformer) + self.assertEqual(model.model_name, DEFAULT_MODEL_NAME) + self.assertEqual(model.kwargs, {'device': 'cpu'}) + self.assertEqual(model.max_seq_length, 128) + + def test_loading_model_from_artifacts_written_by_older_versions(self): + model_handler = huggingface.SentenceTransformerEmbeddings( + model_name=DEFAULT_MODEL_NAME, + columns=[test_query_column]).get_model_handler() + # Beam 2.77.0 and earlier stored a reference to the SentenceTransformer + # class in the artifacts. This is how jsonpickle restores the reference if + # sentence-transformers is not installed when reading the artifacts. + with test_utils.block_imports('sentence_transformers'): + model_handler._model_class = jsonpickle.decode( + '{"py/type": ' + '"sentence_transformers.SentenceTransformer.SentenceTransformer"}') + with _fake_sentence_transformers_module(): + model = model_handler.load_model() + self.assertIsInstance(model, _FakeSentenceTransformer) + self.assertEqual(model.model_name, DEFAULT_MODEL_NAME) + + @unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') class HuggingfaceInferenceAPITest(unittest.TestCase): def setUp(self): diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py index e31c53765b81..425b4c6dfac4 100644 --- a/sdks/python/apache_beam/testing/test_utils.py +++ b/sdks/python/apache_beam/testing/test_utils.py @@ -22,11 +22,14 @@ # pytype: skip-file +import contextlib import hashlib import importlib +import importlib.abc import os import shutil import struct +import sys import tempfile import numpy as np @@ -125,6 +128,64 @@ def remove_patches(): testcase.addCleanup(remove_patches) +class _ImportBlocker(importlib.abc.MetaPathFinder): + """A meta path finder that blocks, and records, imports of given packages.""" + def __init__(self, packages): + self._packages = tuple(packages) + self.attempted_imports = [] + + def blocks(self, module_name): + return any( + module_name == package or module_name.startswith(package + '.') + for package in self._packages) + + def find_spec(self, fullname, path, target=None): + if self.blocks(fullname): + self.attempted_imports.append(fullname) + raise ModuleNotFoundError( + 'Import of %s is blocked for testing.' % fullname, name=fullname) + return None + + +@contextlib.contextmanager +def block_imports(*packages): + """Context manager that makes the given packages unimportable. + + Within the context, importing any of the given packages, or their + submodules, raises a ModuleNotFoundError as if they were not installed. This + also applies to packages that were already imported. Such packages are + restored when exiting the context. + + This is useful, for example, for verifying that constructing and submitting + a pipeline neither requires, nor tries to import, a dependency that is only + needed when executing the pipeline on workers:: + + with block_imports('sentence_transformers') as attempted_imports: + ... # Construct and serialize the pipeline. + assert not attempted_imports + + Args: + *packages: names of the top-level packages, or modules, to block. + + Yields: + A list with the names of all modules whose import was attempted, and + blocked, within the context. + """ + blocker = _ImportBlocker(packages) + hidden_modules = { + name: module + for name, module in list(sys.modules.items()) if blocker.blocks(name) + } + for name in hidden_modules: + del sys.modules[name] + sys.meta_path.insert(0, blocker) + try: + yield blocker.attempted_imports + finally: + sys.meta_path.remove(blocker) + sys.modules.update(hidden_modules) + + @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_beam_io_error_filter) def delete_files(file_paths): diff --git a/sdks/python/apache_beam/testing/test_utils_test.py b/sdks/python/apache_beam/testing/test_utils_test.py index 1a89ada2a709..9c379cf35aab 100644 --- a/sdks/python/apache_beam/testing/test_utils_test.py +++ b/sdks/python/apache_beam/testing/test_utils_test.py @@ -19,8 +19,10 @@ # pytype: skip-file +import importlib import logging import os +import sys import tempfile import unittest @@ -95,6 +97,35 @@ def test_cleanup_topics(self): pub_client.delete_topic.assert_called_with(topic=topic.name) +class BlockImportsTest(unittest.TestCase): + def test_block_imports_blocks_and_records_imports(self): + with utils.block_imports('colorsys') as attempted_imports: + with self.assertRaises(ImportError): + importlib.import_module('colorsys') + self.assertEqual(attempted_imports, ['colorsys']) + # The module can be imported again after exiting the context. + importlib.import_module('colorsys') + + def test_block_imports_hides_and_restores_imported_modules(self): + wsgiref_util = importlib.import_module('wsgiref.util') + wsgiref = sys.modules['wsgiref'] + with utils.block_imports('wsgiref') as attempted_imports: + self.assertNotIn('wsgiref', sys.modules) + self.assertNotIn('wsgiref.util', sys.modules) + with self.assertRaises(ImportError): + importlib.import_module('wsgiref.util') + # Importing a submodule first imports its parent package. + self.assertEqual(attempted_imports, ['wsgiref']) + self.assertIs(sys.modules['wsgiref'], wsgiref) + self.assertIs(sys.modules['wsgiref.util'], wsgiref_util) + + def test_block_imports_only_blocks_given_packages(self): + with utils.block_imports('wsgi', 'colorsys.foo') as attempted_imports: + importlib.import_module('wsgiref') + importlib.import_module('colorsys') + self.assertEqual(attempted_imports, []) + + class LCGeneratorTest(unittest.TestCase): Generator = utils.LCGenerator diff --git a/sdks/python/apache_beam/yaml/yaml_ml_test.py b/sdks/python/apache_beam/yaml/yaml_ml_test.py index d8b1bdbae1b2..c4fc179b9a3c 100644 --- a/sdks/python/apache_beam/yaml/yaml_ml_test.py +++ b/sdks/python/apache_beam/yaml/yaml_ml_test.py @@ -20,6 +20,7 @@ import unittest import apache_beam as beam +from apache_beam.testing import test_utils from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.yaml_transform import YamlTransform @@ -28,7 +29,7 @@ # pylint: disable=wrong-import-order, wrong-import-position, unused-import from apache_beam.ml.transforms import tft except ImportError: - raise unittest.SkipTest('tensorflow_transform is not installed.') + tft = None TRAIN_DATA = [ beam.Row(num=0, text='And God said, Let there be light,'), @@ -41,6 +42,7 @@ ] +@unittest.skipIf(tft is None, 'tensorflow_transform is not installed.') class MLTransformTest(unittest.TestCase): def test_ml_transform(self): ml_opts = beam.options.pipeline_options.PipelineOptions( @@ -236,6 +238,42 @@ def check_row(row): assert_that(result | beam.Map(check_row), equal_to([1, 2, 3])) +class MLTransformJobSubmissionTest(unittest.TestCase): + def test_sentence_transformer_embedding_does_not_use_sentence_transformers( + self): + # sentence-transformers should only be needed when executing the pipeline + # (for example, on the workers), not for constructing and submitting it. + ml_opts = beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle', yaml_experimental_features=['ML']) + with tempfile.TemporaryDirectory() as tempdir: + with test_utils.block_imports( + 'sentence_transformers') as attempted_imports: + pipeline = beam.Pipeline(options=ml_opts) + _ = pipeline | YamlTransform( + f''' + type: chain + transforms: + - type: Create + config: + elements: + - {{id: 1, log_message: "Error in module A"}} + - type: MLTransform + config: + write_artifact_location: {tempdir} + transforms: + - type: SentenceTransformerEmbeddings + config: + model_name: all-MiniLM-L6-v2 + columns: [log_message] + ''') + # 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()