feat: batch predict with GCS

Author: Valeriy-Burlaka

genai/batch_predict/batch_prediction_with_gcs.py

@@ -0,0 +1,66 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.genai.types import BatchJob
+
+
+def create_job(output_uri: str) -> BatchJob:
+    # [START googlegenaisdk_batch_prediction_with_gcs]
+    import time
+
+    from google import genai
+
+    client = genai.Client()
+    # TODO(developer): Update and un-comment below line
+    # output_uri = "gs://your-bucket/your-prefix/..."
+
+    job = client.batches.create(
+        model="gemini-2.0-flash-001",
+        src="gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl",
+        config={
+            "dest": output_uri
+        }
+    )
+    print(f"Job name: {job.name}")
+    print(f"Job state: {job.state}")
+    # Example response:
+    # Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
+    # Job state: JOB_STATE_PENDING
+
+    # See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
+    completed_states = [
+        "JOB_STATE_SUCCEEDED",
+        "JOB_STATE_FAILED",
+        "JOB_STATE_CANCELLED",
+        "JOB_STATE_PAUSED",
+    ]
+    while job.state not in completed_states:
+        time.sleep(30)
+        job = client.batches.get(name=job.name)
+        print(f"Job state: {job.state}")
+    # Example response:
+    # Job state: JOB_STATE_PENDING
+    # Job state: JOB_STATE_RUNNING
+    # Job state: JOB_STATE_RUNNING
+    # ...
+    # Job state: JOB_STATE_SUCCEEDED
+
+    # [END googlegenaisdk_batch_prediction_with_gcs]
+    return job
+
+
+if __name__ == "__main__":
+    create_job(
+        output_uri="gs://your-bucket/your-prefix/..."
+    )

genai/batch_predict/requirements-test.txt

@@ -0,0 +1,3 @@
+google-api-core==2.24.0
+google-cloud-storage==2.19.0
+pytest==8.2.0

genai/batch_predict/requirements.txt

@@ -0,0 +1 @@
+google-genai==0.7.0

genai/batch_predict/test_batch_predict.py

@@ -0,0 +1,51 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datetime import datetime as dt, UTC
+import os
+
+from google.cloud import storage
+from google.genai.types import JobState
+import pytest
+
+import batch_prediction_with_gcs
+
+
+os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
+os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
+# The project name is included in the CICD pipeline
+# os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
+GCS_OUTPUT_BUCKET = "python-docs-samples-tests"
+GCS_OUTPUT_BUCKET = "gemini-batch-prediction-results"
+
+
+@pytest.fixture(scope="session")
+def gcs_output_uri():
+    prefix = f"text_output/{dt.now(UTC)}"
+
+    yield f"gs://{GCS_OUTPUT_BUCKET}/{prefix}"
+
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(GCS_OUTPUT_BUCKET)
+    blobs = bucket.list_blobs(prefix=prefix)
+    for blob in blobs:
+        blob.delete()
+
+
+def test_batch_prediction_with_gcs(gcs_output_uri) -> None:
+    job = batch_prediction_with_gcs.create_job(output_uri=gcs_output_uri)
+    assert job
+    assert job.state == "JOB_STATE_SUCCEEDED"
+    assert job.dest.gcs_uri == gcs_output_uri
+    assert job.state == JobState.JOB_STATE_SUCCEEDED