Remove Vertex LLM Inference Integration Test (#36172)

jrmccluskey · gemini-code-assist[bot] · web-flow · commit 11f75687d680 · 2025-09-22T10:40:15.000-04:00
* Move Vertex LLM Inference IT Test to new Endpoint

* trigger python postcommit

* Remove LLM IT test, add note to example

* remove unused var

* Remove unused import

* Update sdks/python/apache_beam/examples/inference/vertex_ai_llm_text_classification.py

Co-authored-by: gemini-code-assist[bot] &lt;176961590+gemini-code-assist[bot]@users.noreply.github.com&gt;

* linting

* Soften disclaimer language a bit, add call-out to README

---------

Co-authored-by: gemini-code-assist[bot] &lt;176961590+gemini-code-assist[bot]@users.noreply.github.com&gt;
diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json
@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run.",
-  "modification": 29
+  "modification": 30
 }
 
diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md
@@ -856,6 +856,12 @@ Each line represents a prediction of the flower type along with the confidence i
 
 ## Text classifcation with a Vertex AI LLM
 
+**NOTE**
+Google has deprecated PaLM LLMs like text-bison and no longer supports querying them on Vertex AI endpoints. Separately, the use of the Vertex AI Predict API is
+not supported for Gemini models in favor of use of the google-genai API. As a result, this example no longer works as-written. To perform inference with
+Gemini models deployed on Google infrastructure, please see the `GeminiModelHandler` (in `apache_beam.ml.inference.gemini_inference`) and the
+[`gemini_text_classification.py` example](./gemini_text_classification.py). For custom LLMs, you may still follow this design pattern.
+
 [`vertex_ai_llm_text_classification.py`](./vertex_ai_llm_text_classification.py) contains an implementation for a RunInference pipeline that performs image classification using a model hosted on Vertex AI (based on https://cloud.google.com/vertex-ai/docs/tutorials/image-recognition-custom).
 
 The pipeline reads image urls, performs basic preprocessing to convert them into a List of floats, passes the masked sentence to the Vertex AI implementation of RunInference, and then writes the predictions to a text file.
diff --git a/sdks/python/apache_beam/examples/inference/vertex_ai_llm_text_classification.py b/sdks/python/apache_beam/examples/inference/vertex_ai_llm_text_classification.py
@@ -21,6 +21,16 @@
 model can be generated by fine tuning the text-bison model or another similar
 model (see
 https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models#supervised-fine-tuning)
+
+**NOTE**
+Google has deprecated PaLM LLMs and no longer supports querying them on
+Vertex AI endpoints. Separately, the use of the Vertex AI Predict API is not
+supported for Gemini models in favor of use of the google-genai API. As a
+result, this example no longer works as-written. To perform inference with
+Gemini models deployed on Google infrastructure, please see the
+`GeminiModelHandler` (in `apache_beam.ml.inference.gemini_inference`) and the
+`gemini_text_classification.py` example. For custom LLMs, you may still follow
+this design pattern.
 """
 
 import argparse
diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py
@@ -29,15 +29,13 @@
 # pylint: disable=ungrouped-imports
 try:
   from apache_beam.examples.inference import vertex_ai_image_classification
-  from apache_beam.examples.inference import vertex_ai_llm_text_classification
 except ImportError as e:
   raise unittest.SkipTest(
       "Vertex AI model handler dependencies are not installed")
 
 _INPUT = "gs://apache-beam-ml/testing/inputs/vertex_images/*/*.jpg"
 _OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/vertex_images"
 _FLOWER_ENDPOINT_ID = "5384055553544683520"
-_LLM_ENDPOINT_ID = "9157860935048626176"
 _ENDPOINT_PROJECT = "apache-beam-testing"
 _ENDPOINT_REGION = "us-central1"
 _ENDPOINT_NETWORK = "projects/844138762903/global/networks/beam-test-vpc"
@@ -65,21 +63,6 @@ def test_vertex_ai_run_flower_image_classification(self):
         test_pipeline.get_full_options_as_args(**extra_opts))
     self.assertEqual(FileSystems().exists(output_file), True)
 
-  @pytest.mark.vertex_ai_postcommit
-  def test_vertex_ai_run_llm_text_classification(self):
-    output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt'])
-
-    test_pipeline = TestPipeline(is_integration_test=True)
-    extra_opts = {
-        'output': output_file,
-        'endpoint_id': _LLM_ENDPOINT_ID,
-        'endpoint_project': _ENDPOINT_PROJECT,
-        'endpoint_region': _ENDPOINT_REGION
-    }
-    vertex_ai_llm_text_classification.run(
-        test_pipeline.get_full_options_as_args(**extra_opts))
-    self.assertEqual(FileSystems().exists(output_file), True)
-
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.DEBUG)

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`{`
`2`	`2`	`"comment": "Modify this file in a trivial way to cause this test suite to run.",`
`3`		`- "modification": 29`
	`3`	`+ "modification": 30`
`4`	`4`	`}`
`5`	`5`