feat: GenAI - Support batch prediction in Model Garden OpenModel.

vertex-sdk-bot · copybara-github · commit 8bf30b74828c · 2025-04-29T18:06:08.000-07:00
PiperOrigin-RevId: 752953157
diff --git a/tests/unit/vertexai/model_garden/test_model_garden.py b/tests/unit/vertexai/model_garden/test_model_garden.py
@@ -20,13 +20,25 @@
 from google.api_core import operation as ga_operation
 from google.auth import credentials as auth_credentials
 from google.cloud import aiplatform
+from google.cloud.aiplatform.compat.services import job_service_client
+from google.cloud.aiplatform.compat.types import (
+    batch_prediction_job as gca_batch_prediction_job_compat,
+)
+from google.cloud.aiplatform.compat.types import io as gca_io_compat
+from google.cloud.aiplatform.compat.types import (
+    job_state as gca_job_state_compat,
+)
+from google.cloud.aiplatform_v1.types import machine_resources
+from google.cloud.aiplatform_v1.types import manual_batch_tuning_parameters
 from google.cloud.aiplatform_v1beta1 import types
 from google.cloud.aiplatform_v1beta1.services import model_garden_service
+from vertexai import batch_prediction
 from vertexai.preview import model_garden
 import pytest
 
 from google.protobuf import duration_pb2
 
+
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"
 
@@ -73,6 +85,24 @@
         timeout_seconds=10,
     ),
 )
+_TEST_BATCH_PREDICTION_JOB_ID = "123456789"
+_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}"
+_TEST_BATCH_PREDICTION_JOB_NAME = (
+    f"{_TEST_PARENT}/batchPredictionJobs/{_TEST_BATCH_PREDICTION_JOB_ID}"
+)
+_TEST_BATCH_PREDICTION_MODEL_FULL_RESOURCE_NAME = (
+    "publishers/google/models/gemma@gemma-2b-it"
+)
+_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME = "test-batch-prediction-job"
+_TEST_JOB_STATE_RUNNING = gca_job_state_compat.JobState(3)
+_TEST_GAPIC_BATCH_PREDICTION_JOB = gca_batch_prediction_job_compat.BatchPredictionJob(
+    name=_TEST_BATCH_PREDICTION_JOB_NAME,
+    display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME,
+    model=_TEST_BATCH_PREDICTION_MODEL_FULL_RESOURCE_NAME,
+    state=_TEST_JOB_STATE_RUNNING,
+)
+_TEST_BQ_INPUT_URI = "bq://test-project.test-dataset.test-input"
+_TEST_BQ_OUTPUT_PREFIX = "bq://test-project.test-dataset.test-output"
 
 
 @pytest.fixture(scope="module")
@@ -117,6 +147,25 @@ def deploy_mock():
         yield deploy
 
 
+@pytest.fixture
+def batch_prediction_mock():
+    """Mocks the create_batch_prediction_job method."""
+    with mock.patch.object(
+        job_service_client.JobServiceClient, "create_batch_prediction_job"
+    ) as create_batch_prediction_job_mock:
+        create_batch_prediction_job_mock.return_value = _TEST_GAPIC_BATCH_PREDICTION_JOB
+        yield create_batch_prediction_job_mock
+
+
+@pytest.fixture
+def complete_bq_uri_mock():
+    with mock.patch.object(
+        batch_prediction.BatchPredictionJob, "_complete_bq_uri"
+    ) as complete_bq_uri_mock:
+        complete_bq_uri_mock.return_value = _TEST_BQ_OUTPUT_PREFIX
+        yield complete_bq_uri_mock
+
+
 @pytest.fixture
 def get_publisher_model_mock():
     with mock.patch.object(
@@ -355,6 +404,8 @@ def list_publisher_models_mock():
     "get_publisher_model_mock",
     "list_publisher_models_mock",
     "export_publisher_model_mock",
+    "batch_prediction_mock",
+    "complete_bq_uri_mock",
 )
 class TestModelGarden:
     """Test cases for ModelGarden class."""
@@ -897,3 +948,54 @@ def test_list_deployable_models(self, list_publisher_models_mock):
             "google/gemma-2-2b",
             "google/gemma-2-2b",
         ]
+
+    def test_batch_prediction_success(self, batch_prediction_mock):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        model = model_garden.OpenModel(
+            model_name=_TEST_BATCH_PREDICTION_MODEL_FULL_RESOURCE_NAME
+        )
+        job = model.batch_predict(
+            input_dataset=_TEST_BQ_INPUT_URI,
+            job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME,
+            machine_type="g2-standard-12",
+            accelerator_type="NVIDIA_L4",
+            accelerator_count=1,
+            starting_replica_count=1,
+        )
+
+        assert job.gca_resource == _TEST_GAPIC_BATCH_PREDICTION_JOB
+
+        expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob(
+            display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME,
+            model=_TEST_BATCH_PREDICTION_MODEL_FULL_RESOURCE_NAME,
+            input_config=gca_batch_prediction_job_compat.BatchPredictionJob.InputConfig(
+                instances_format="bigquery",
+                bigquery_source=gca_io_compat.BigQuerySource(
+                    input_uri=_TEST_BQ_INPUT_URI
+                ),
+            ),
+            output_config=gca_batch_prediction_job_compat.BatchPredictionJob.OutputConfig(
+                bigquery_destination=gca_io_compat.BigQueryDestination(
+                    output_uri=_TEST_BQ_OUTPUT_PREFIX
+                ),
+                predictions_format="bigquery",
+            ),
+            dedicated_resources=machine_resources.BatchDedicatedResources(
+                machine_spec=machine_resources.MachineSpec(
+                    machine_type="g2-standard-12",
+                    accelerator_type="NVIDIA_L4",
+                    accelerator_count=1,
+                ),
+                starting_replica_count=1,
+            ),
+            manual_batch_tuning_parameters=manual_batch_tuning_parameters.ManualBatchTuningParameters(),
+        )
+
+        batch_prediction_mock.assert_called_once_with(
+            parent=_TEST_PARENT,
+            batch_prediction_job=expected_gapic_batch_prediction_job,
+            timeout=None,
+        )
diff --git a/vertexai/model_garden/_model_garden.py b/vertexai/model_garden/_model_garden.py
@@ -19,7 +19,7 @@
 import datetime
 import functools
 import re
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, List, Optional, Sequence, Union
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
@@ -29,6 +29,8 @@
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform_v1beta1 import types
 from google.cloud.aiplatform_v1beta1.services import model_garden_service
+from vertexai import batch_prediction
+
 
 from google.protobuf import duration_pb2
 
@@ -656,3 +658,60 @@ def list_deploy_options(
                 " to find out which ones currently support deployment."
             )
         return multi_deploy
+
+    def batch_predict(
+        self,
+        input_dataset: Union[str, List[str]],
+        *,
+        output_uri_prefix: Optional[str] = None,
+        job_display_name: Optional[str] = None,
+        machine_type: Optional[str] = None,
+        accelerator_type: Optional[str] = None,
+        accelerator_count: Optional[int] = None,
+        starting_replica_count: Optional[int] = None,
+        max_replica_count: Optional[int] = None,
+    ) -> batch_prediction.BatchPredictionJob:
+        """Perform batch prediction on the model.
+
+        Args:
+            input_dataset (Union[str, List[str]]):
+                GCS URI(-s) or BigQuery URI to your input data to run batch
+                prediction on. Example: "gs://path/to/input/data.jsonl" or
+                "bq://projectId.bqDatasetId.bqTableId"
+            output_uri_prefix (Optional[str]):
+                GCS or BigQuery URI prefix for the output predictions. Example:
+                "gs://path/to/output/data" or "bq://projectId.bqDatasetId"
+                If not specified, f"{STAGING_BUCKET}/gen-ai-batch-prediction" will
+                be used for GCS source and
+                f"bq://projectId.gen_ai_batch_prediction.predictions_{TIMESTAMP}"
+                will be used for BigQuery source.
+            job_display_name (Optional[str]):
+                The user-defined name of the BatchPredictionJob.
+                The name can be up to 128 characters long and can be consist
+                of any UTF-8 characters.
+            machine_type (Optional[str]):
+                The machine type for the batch prediction job.
+            accelerator_type (Optional[str]):
+                The accelerator type for the batch prediction job.
+            accelerator_count (Optional[int]):
+                The accelerator count for the batch prediction job.
+            starting_replica_count (Optional[int]):
+                The starting replica count for the batch prediction job.
+            max_replica_count (Optional[int]):
+                The maximum replica count for the batch prediction job.
+
+        Returns:
+            batch_prediction.BatchPredictionJob:
+                The batch prediction job.
+        """
+        return batch_prediction.BatchPredictionJob.submit(
+            source_model=self._publisher_model_name,
+            input_dataset=input_dataset,
+            output_uri_prefix=output_uri_prefix,
+            job_display_name=job_display_name,
+            machine_type=machine_type,
+            accelerator_type=accelerator_type,
+            accelerator_count=accelerator_count,
+            starting_replica_count=starting_replica_count,
+            max_replica_count=max_replica_count,
+        )
