From 2cfbe768166280091179e529643dcf74a49081b3 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Thu, 12 Mar 2026 15:00:41 +0400 Subject: [PATCH 1/3] Cherrypick #37668 --- .../ml/rag/ingestion/milvus_search_it_test.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py index 38b497e8fa71..b6e5083ea728 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py @@ -22,12 +22,6 @@ from typing import cast import pytest -from pymilvus import CollectionSchema -from pymilvus import DataType -from pymilvus import FieldSchema -from pymilvus import MilvusClient -from pymilvus.exceptions import MilvusException -from pymilvus.milvus_client import IndexParams import apache_beam as beam from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig @@ -41,11 +35,21 @@ from apache_beam.ml.rag.utils import unpack_dataclass_with_kwargs from apache_beam.testing.test_pipeline import TestPipeline +# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: + from pymilvus import CollectionSchema + from pymilvus import DataType + from pymilvus import FieldSchema + from pymilvus import MilvusClient + from pymilvus.exceptions import MilvusException + from pymilvus.milvus_client import IndexParams + from apache_beam.ml.rag.ingestion.milvus_search import MilvusVectorWriterConfig from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig -except ImportError as e: - raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}') + PYMILVUS_AVAILABLE = True +except ImportError: + PYMILVUS_AVAILABLE = False +# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports def _construct_index_params(): @@ -158,6 +162,7 @@ def drop_collection(client: MilvusClient, collection_name: str): @pytest.mark.require_docker_in_docker +@unittest.skipIf(not PYMILVUS_AVAILABLE, 'pymilvus is not installed.') @unittest.skipUnless( platform.system() == "Linux", "Test runs only on Linux due to lack of support, as yet, for nested " From 45e8bcaf31336805dde90ae5adc09ec60977f781 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Thu, 12 Mar 2026 17:36:44 +0400 Subject: [PATCH 2/3] Fix lint --- sdks/python/apache_beam/ml/inference/gemini_inference.py | 5 +++-- sdks/python/apache_beam/ml/transforms/base_test.py | 4 ++-- .../apache_beam/ml/transforms/embeddings/huggingface_test.py | 2 +- .../ml/transforms/embeddings/tensorflow_hub_test.py | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference.py b/sdks/python/apache_beam/ml/inference/gemini_inference.py index 8b6d0edfa996..2ba220d0162a 100644 --- a/sdks/python/apache_beam/ml/inference/gemini_inference.py +++ b/sdks/python/apache_beam/ml/inference/gemini_inference.py @@ -22,6 +22,7 @@ from typing import Any from typing import Optional from typing import Union +from typing import cast from google import genai from google.genai import errors @@ -73,7 +74,7 @@ def generate_from_string( call. """ return model.models.generate_content( - model=model_name, contents=batch, **inference_args) + model=model_name, contents=cast(Any, batch), **inference_args) def generate_image_from_strings_and_images( @@ -96,7 +97,7 @@ def generate_image_from_strings_and_images( call. """ return model.models.generate_content( - model=model_name, contents=batch, **inference_args) + model=model_name, contents=cast(Any, batch), **inference_args) class GeminiModelHandler(RemoteModelHandler[Any, PredictionResult, diff --git a/sdks/python/apache_beam/ml/transforms/base_test.py b/sdks/python/apache_beam/ml/transforms/base_test.py index 190381cc2f34..64d0afe955d0 100644 --- a/sdks/python/apache_beam/ml/transforms/base_test.py +++ b/sdks/python/apache_beam/ml/transforms/base_test.py @@ -51,8 +51,8 @@ import PIL from PIL.Image import Image as PIL_Image except ImportError: - PIL = None - PIL_Image = Any + PIL = None # type: ignore[assignment] + PIL_Image = Any # type: ignore[misc, assignment] try: 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 a2358c544781..5e56dc710548 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py @@ -53,7 +53,7 @@ try: from PIL import Image except ImportError: - Image = None + Image = None # type: ignore[assignment] _HF_TOKEN = os.environ.get('HF_INFERENCE_TOKEN') test_query = "This is a test" diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py index 0a4f8c8275c3..b9a14ddb8c53 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py @@ -52,7 +52,7 @@ from apache_beam.ml.transforms.embeddings.tensorflow_hub import TensorflowHubImageEmbeddings except ImportError: TensorflowHubImageEmbeddings = None # type: ignore - Image = None + Image = None # type: ignore[assignment] @unittest.skipIf( From 1ee624f196611805917d5fdaaa2e186d5a5e03dd Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Thu, 12 Mar 2026 19:07:30 +0400 Subject: [PATCH 3/3] Fix lint --- sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py index aa294e98789d..e0da9499faf3 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py @@ -16,11 +16,12 @@ """Tests for apache_beam.ml.rag.embeddings.vertex_ai.""" -import pytest import shutil import tempfile import unittest +import pytest + import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content