Finished setting config for gemma sample. it's now working

Author: XrossFox

dataflow/gemma/e2e_test.py

@@ -34,7 +34,7 @@
 
 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,
       add the PYTHONPATH to the "env" in your noxfile_config.py file.
@@ -47,8 +47,11 @@
 
 import pytest
 
-DATAFLOW_MACHINE_TYPE = "g2-standard-4"
-GEMMA_GCS = "gs://perm-dataflow-gemma-example-testdata/gemma_2b"
+DATAFLOW_MACHINE_TYPE = "g2-standard-8"
+# TODO If testing locally, point this to a bucket you control
+# and download the gemma assets
+# GEMMA_GCS = "gs://perm-dataflow-gemma-example-testdata/gemma_2b"
+GEMMA_GCS = "gs://test-bucket-for-gemma/assets_here/gemma_2b"
 NAME = "dataflow/gemma/streaming"
 
 
@@ -111,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",

dataflow/gemma/noxfile_config.py

@@ -19,6 +19,6 @@
     # 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.11)
-    "ignored_versions": ["3.8", "3.9", "3.10"],
+    "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.157.0
+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
+pytest-timeout==2.4.0

dataflow/gemma/requirements.txt

@@ -2,4 +2,4 @@ protobuf==6.33.6
 apache_beam[gcp]==2.74.0
 keras==3.14.1
 keras_nlp==0.29.1
-pyOpenSSL==25.3.0
+pyOpenSSL==25.3.0