feat: GenAI - Add support for self-hosted OSS models in Batch Prediction.

PiperOrigin-RevId: 752487510

Author: vertex-sdk-bot

tests/unit/vertexai/test_batch_prediction.py

@@ -56,6 +56,9 @@
 _TEST_CLAUDE_MODEL_RESOURCE_NAME = (
     f"publishers/anthropic/models/{_TEST_CLAUDE_MODEL_NAME}"
 )
+_TEST_SELF_HOSTED_GEMMA_MODEL_RESOURCE_NAME = (
+    "publishers/google/models/gemma@gemma-2b-it"
+)
 
 _TEST_GCS_INPUT_URI = "gs://test-bucket/test-input.jsonl"
 _TEST_GCS_INPUT_URI_2 = "gs://test-bucket/test-input-2.jsonl"
@@ -589,6 +592,39 @@ def test_submit_batch_prediction_job_with_tuned_model(
             retry=aiplatform_base._DEFAULT_RETRY,
         )
 
+    def test_submit_batch_prediction_job_with_self_hosted_gemma_model(
+        self,
+        create_batch_prediction_job_mock,
+    ):
+        job = batch_prediction.BatchPredictionJob.submit(
+            source_model=_TEST_SELF_HOSTED_GEMMA_MODEL_RESOURCE_NAME,
+            input_dataset=_TEST_BQ_INPUT_URI,
+        )
+
+        assert job.gca_resource == _TEST_GAPIC_BATCH_PREDICTION_JOB
+
+        expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob(
+            display_name=_TEST_DISPLAY_NAME,
+            model=_TEST_SELF_HOSTED_GEMMA_MODEL_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",
+            ),
+        )
+        create_batch_prediction_job_mock.assert_called_once_with(
+            parent=_TEST_PARENT,
+            batch_prediction_job=expected_gapic_batch_prediction_job,
+            timeout=None,
+        )
+
     def test_submit_batch_prediction_job_with_invalid_source_model(self):
         with pytest.raises(
             ValueError,

vertexai/batch_prediction/_batch_prediction.py

@@ -116,6 +116,11 @@ def submit(
         *,
         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,
     ) -> "BatchPredictionJob":
         """Submits a batch prediction job for a GenAI model.
 
@@ -142,6 +147,16 @@ def submit(
                 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 (str):
+                The type of machine for running batch prediction job.
+            accelerator_type (str):
+                The type of accelerator for running batch prediction job.
+            accelerator_count (int):
+                The number of accelerators for running batch prediction job.
+            starting_replica_count (int):
+                The starting number of replica for running batch prediction job.
+            max_replica_count (int):
+                The maximum number of replica for running batch prediction job.
 
         Returns:
             Instantiated BatchPredictionJob.
@@ -219,6 +234,11 @@ def submit(
                 bigquery_source=bigquery_source,
                 gcs_destination_prefix=gcs_destination_prefix,
                 bigquery_destination_prefix=bigquery_destination_prefix,
+                machine_type=machine_type,
+                accelerator_type=accelerator_type,
+                accelerator_count=accelerator_count,
+                starting_replica_count=starting_replica_count,
+                max_replica_count=max_replica_count,
             )
             job = cls._empty_constructor()
             job._gca_resource = aiplatform_job._gca_resource
@@ -281,27 +301,29 @@ def _reconcile_model_name(cls, model_name: str) -> str:
         if "/" not in model_name:
             # model name (e.g., gemini-1.0-pro)
             if model_name.startswith("gemini"):
-                model_name = "publishers/google/models/" + model_name
+                return "publishers/google/models/" + model_name
             else:
                 raise ValueError(
                     "Abbreviated model names are only supported for Gemini models. "
                     "Please provide the full publisher model name."
                 )
         elif model_name.startswith("models/"):
             # publisher model name (e.g., models/gemini-1.0-pro)
-            model_name = "publishers/google/" + model_name
+            return "publishers/google/" + model_name
         elif (
-            # publisher model full name
-            not model_name.startswith("publishers/google/models/")
-            and not model_name.startswith("publishers/meta/models/")
-            and not model_name.startswith("publishers/anthropic/models/")
-            # tuned model full resource name
-            and not re.search(_GEMINI_TUNED_MODEL_PATTERN, model_name)
+            re.match(
+                r"^publishers/(?P<publisher>[^/]+)/models/(?P<model>[^@]+)@(?P<version>[^@]+)$",
+                model_name,
+            )
+            or model_name.startswith("publishers/google/models/")
+            or model_name.startswith("publishers/meta/models/")
+            or model_name.startswith("publishers/anthropic/models/")
+            or re.search(_GEMINI_TUNED_MODEL_PATTERN, model_name)
         ):
+            return model_name
+        else:
             raise ValueError(f"Invalid format for model name: {model_name}.")
 
-        return model_name
-
     @classmethod
     def _is_genai_model(cls, model_name: str) -> bool:
         """Validates if a given model_name represents a GenAI model."""
@@ -326,6 +348,13 @@ def _is_genai_model(cls, model_name: str) -> bool:
             # Model is a claude model.
             return True
 
+        if re.match(
+            r"^publishers/(?P<publisher>[^/]+)/models/(?P<model>[^@]+)@(?P<version>[^@]+)$",
+            model_name,
+        ):
+            # Model is a self-hosted model.
+            return True
+
         return False
 
     @classmethod