chore(dataflow/gemma): update dependencies and format code

dataflow/gemma/Dockerfile

@@ -15,7 +15,7 @@
 # This uses Ubuntu with Python 3.11
 # You can check the Python version for a given tensorflow
 # container at https://hub.docker.com/r/tensorflow/tensorflow/tags
-ARG SERVING_BUILD_IMAGE=tensorflow/tensorflow:2.16.1-gpu
+ARG SERVING_BUILD_IMAGE=tensorflow/tensorflow:2.20.0-gpu
 
 FROM ${SERVING_BUILD_IMAGE}
 
@@ -29,7 +29,7 @@ RUN pip install --upgrade --no-cache-dir pip \
     && pip install --no-cache-dir -r requirements.txt
 
 # Copy files from official SDK image, including script/dependencies.
-COPY --from=apache/beam_python3.14_sdk:2.73.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.11_sdk:2.74.0 /opt/apache/beam /opt/apache/beam
 
 # Copy the model directory downloaded from Kaggle and the pipeline code.
 COPY gemma_2b gemma_2B

dataflow/gemma/custom_model_gemma.py

@@ -13,29 +13,30 @@
 # limitations under the License.
 
 from collections.abc import Iterable, Sequence
-
-import logging
-
 from typing import Any
 from typing import Optional
 
 import apache_beam as beam
+
 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.base import RunInference
 from apache_beam.options.pipeline_options import PipelineOptions
 
 import keras_nlp
+
 from keras_nlp.src.models.gemma.gemma_causal_lm import GemmaCausalLM
 
+import logging
+
 
 class GemmaModelHandler(ModelHandler[str, PredictionResult, GemmaCausalLM]):
     def __init__(
         self,
         model_name: str = "gemma_2B",
     ):
-        """ Implementation of the ModelHandler interface for Gemma using text as input.
+        """Implementation of the ModelHandler interface for Gemma using text as input.
 
         Example Usage::
 
@@ -48,7 +49,7 @@ def __init__(
         self._env_vars = {}
 
     def share_model_across_processes(self) -> bool:
-        """ Indicates if the model should be loaded once-per-VM rather than
+        """Indicates if the model should be loaded once-per-VM rather than
         once-per-worker-process on a VM. Because Gemma is a large language model,
         this will always return True to avoid OOM errors.
         """
@@ -62,7 +63,7 @@ def run_inference(
         self,
         batch: Sequence[str],
         model: GemmaCausalLM,
-        inference_args: Optional[dict[str, Any]] = None
+        inference_args: Optional[dict[str, Any]] = None,
     ) -> Iterable[PredictionResult]:
         """Runs inferences on a batch of text strings.
 
@@ -85,7 +86,8 @@ def run_inference(
 class FormatOutput(beam.DoFn):
     def process(self, element, *args, **kwargs):
         yield "Input: {input}, Output: {output}".format(
-            input=element.example, output=element.inference)
+            input=element.example, output=element.inference
+        )
 
 
 if __name__ == "__main__":
@@ -119,13 +121,16 @@ def process(self, element, *args, **kwargs):
 
     pipeline = beam.Pipeline(options=beam_options)
     _ = (
-        pipeline | "Read Topic" >>
-        beam.io.ReadFromPubSub(subscription=args.messages_subscription)
+        pipeline
+        | "Read Topic"
+        >> beam.io.ReadFromPubSub(subscription=args.messages_subscription)
         | "Parse" >> beam.Map(lambda x: x.decode("utf-8"))
-        | "RunInference-Gemma" >> RunInference(
+        | "RunInference-Gemma"
+        >> RunInference(
             GemmaModelHandler(args.model_path)
         )  # Send the prompts to the model and get responses.
         | "Format Output" >> beam.ParDo(FormatOutput())  # Format the output.
-        | "Publish Result" >>
-        beam.io.gcp.pubsub.WriteStringsToPubSub(topic=args.responses_topic))
+        | "Publish Result"
+        >> beam.io.gcp.pubsub.WriteStringsToPubSub(topic=args.responses_topic)
+    )
     pipeline.run()

dataflow/gemma/e2e_test.py

@@ -34,19 +34,23 @@
 
 OPTION B: Run tests with nox
     pip install nox
-    nox -s py-3.10
+    nox -s py-3.11
 
-NOTE: For the tests to find the conftest in the testing infrastructure,
+TODO(developer): For the tests to find the conftest in the testing infrastructure,
       add the PYTHONPATH to the "env" in your noxfile_config.py file.
 """
+
 from collections.abc import Callable, Iterator
 
 import conftest  # python-docs-samples/dataflow/conftest.py
 from conftest import Utils
 
 import pytest
 
-DATAFLOW_MACHINE_TYPE = "g2-standard-4"
+DATAFLOW_MACHINE_TYPE = "g2-standard-8"
+# TODO(developer): For local testing, ensure the 'gemma_2b_en' directory is uploaded 
+# to a GCS bucket you manage. Update the constant below to point to 
+# the root path of this uploaded directory (e.g., 'gs://your-bucket-name/path/to/gemma_2b').
 GEMMA_GCS = "gs://perm-dataflow-gemma-example-testdata/gemma_2b"
 NAME = "dataflow/gemma/streaming"
 
@@ -70,8 +74,9 @@ def messages_topic(pubsub_topic: Callable[[str], str]) -> str:
 
 
 @pytest.fixture(scope="session")
-def messages_subscription(pubsub_subscription: Callable[[str, str], str],
-                          messages_topic: str) -> str:
+def messages_subscription(
+    pubsub_subscription: Callable[[str, str], str], messages_topic: str
+) -> str:
     return pubsub_subscription("messages", messages_topic)
 
 
@@ -81,20 +86,21 @@ def responses_topic(pubsub_topic: Callable[[str], str]) -> str:
 
 
 @pytest.fixture(scope="session")
-def responses_subscription(pubsub_subscription: Callable[[str, str], str],
-                           responses_topic: str) -> str:
+def responses_subscription(
+    pubsub_subscription: Callable[[str, str], str], responses_topic: str
+) -> str:
     return pubsub_subscription("responses", responses_topic)
 
 
 @pytest.fixture(scope="session")
 def dataflow_job(
-        project: str,
-        bucket_name: str,
-        location: str,
-        unique_name: str,
-        container_image: str,
-        messages_subscription: str,
-        responses_topic: str,
+    project: str,
+    bucket_name: str,
+    location: str,
+    unique_name: str,
+    container_image: str,
+    messages_subscription: str,
+    responses_topic: str,
 ) -> Iterator[str]:
     # Launch the streaming Dataflow pipeline.
     conftest.run_cmd(
@@ -108,6 +114,7 @@ def dataflow_job(
         f"--temp_location=gs://{bucket_name}/temp",
         f"--region={location}",
         f"--machine_type={DATAFLOW_MACHINE_TYPE}",
+        "--disk_size_gb=100",
         f"--sdk_container_image=gcr.io/{project}/{container_image}",
         "--dataflow_service_options=worker_accelerator=type:nvidia-l4;count:1;install-nvidia-driver:5xx",
         "--requirements_cache=skip",
@@ -127,20 +134,18 @@ def dataflow_job(
 
 @pytest.mark.timeout(3600)
 def test_pipeline_dataflow(
-        project: str,
-        location: str,
-        dataflow_job: str,
-        messages_topic: str,
-        responses_subscription: str,
+    project: str,
+    location: str,
+    dataflow_job: str,
+    messages_topic: str,
+    responses_subscription: str,
 ) -> None:
     print(f"Waiting for the Dataflow workers to start: {dataflow_job}")
     conftest.wait_until(
-        lambda: conftest.dataflow_num_workers(project, location, dataflow_job)
-        > 0,
+        lambda: conftest.dataflow_num_workers(project, location, dataflow_job) > 0,
         "workers are running",
     )
-    num_workers = conftest.dataflow_num_workers(project, location,
-                                                dataflow_job)
+    num_workers = conftest.dataflow_num_workers(project, location, dataflow_job)
     print(f"Dataflow job num_workers: {num_workers}")
 
     messages = ["This is a test for a Python sample."]

dataflow/gemma/noxfile_config.py

@@ -18,9 +18,7 @@
     # You can opt out from the test for specific Python versions.
     # The Python version used is defined by the Dockerfile and the job
     # submission enviornment must match.
-    # Note: Docker-based sample, testing only against version specified in Dockerfile (3.14)
-    "ignored_versions": ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"],
-    "envs": {
-        "PYTHONPATH": ".."
-    },
+    # Note: Docker-based sample, testing only against version specified in Dockerfile (3.11)
+    "ignored_versions": ["3.8", "3.9", "3.10", "3.12", "3.13", "3.14"],
+    "envs": {"PYTHONPATH": ".."},
 }

dataflow/gemma/requirements-test.txt

@@ -1,5 +1,5 @@
-google-cloud-aiplatform==1.49.0
-google-cloud-dataflow-client==0.8.10
-google-cloud-storage==2.16.0
-pytest==9.0.3; python_version >= "3.10"
-pytest-timeout==2.3.1
+google-cloud-aiplatform==1.158.0
+google-cloud-dataflow-client==0.14.0
+google-cloud-storage==3.12.0
+pytest==9.0.3
+pytest-timeout==2.4.0

dataflow/gemma/requirements.txt

@@ -1,4 +1,5 @@
-apache_beam[gcp]==2.54.0
-protobuf==4.25.0
-keras_nlp==0.8.2
-keras==3.0.5
+protobuf==6.33.6
+apache_beam[gcp]==2.74.0
+keras==3.14.1
+keras_nlp==0.29.1
+pyOpenSSL==25.3.0