chore: add validation for model launch stage

PiperOrigin-RevId: 538471545

Author: sararob

tests/unit/aiplatform/test_language_models.py

@@ -48,7 +48,7 @@
     pipeline_state as gca_pipeline_state,
     deployed_model_ref_v1,
 )
-from google.cloud.aiplatform_v1.types import (
+from google.cloud.aiplatform.compat.types import (
     publisher_model as gca_publisher_model,
     model as gca_model,
 )
@@ -75,6 +75,7 @@
     "name": "publishers/google/models/text-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.GA,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/text-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/text_generation_1.0.0.yaml",
@@ -87,6 +88,7 @@
     "name": "publishers/google/models/chat-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/chat-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/chat_generation_1.0.0.yaml",
@@ -99,6 +101,7 @@
     "name": "publishers/google/models/codechat-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/codechat-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/codechat_generation_1.0.0.yaml",
@@ -111,6 +114,7 @@
     "name": "publishers/google/models/code-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/code-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml",
@@ -123,6 +127,7 @@
     "name": "publishers/google/models/code-gecko",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/code-gecko@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml",
@@ -135,6 +140,7 @@
     "name": "publishers/google/models/textembedding-gecko",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.GA,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/chat-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/text_embedding_1.0.0.yaml",

tests/unit/aiplatform/test_model_garden_models.py

@@ -28,7 +28,7 @@
     model_garden_service_client_v1,
 )
 
-from google.cloud.aiplatform_v1.types import (
+from google.cloud.aiplatform.compat.types import (
     publisher_model as gca_publisher_model,
 )
 
@@ -38,6 +38,7 @@
     "name": "publishers/google/models/text-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/text-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/text_generation_1.0.0.yaml",
@@ -52,6 +53,9 @@ class TestModelGardenModels:
     """Unit tests for the _ModelGardenModel base class."""
 
     class FakeModelGardenModel(_model_garden_models._ModelGardenModel):
+
+        _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+
         _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/text_generation_1.0.0.yaml"
 
     def setup_method(self):
@@ -80,3 +84,23 @@ def test_init_model_garden_model_with_from_pretrained(self):
                 name="publishers/google/models/text-bison@001",
                 retry=base._DEFAULT_RETRY,
             )
+
+    def test_init_preview_model_raises_with_ga_launch_stage_set(self):
+        """Tests the text generation model."""
+        aiplatform.init(
+            project=test_constants.ProjectConstants._TEST_PROJECT,
+            location=test_constants.ProjectConstants._TEST_LOCATION,
+        )
+        with mock.patch.object(
+            target=model_garden_service_client_v1.ModelGardenServiceClient,
+            attribute="get_publisher_model",
+            return_value=gca_publisher_model.PublisherModel(
+                _TEXT_BISON_PUBLISHER_MODEL_DICT
+            ),
+        ):
+            self.FakeModelGardenModel._LAUNCH_STAGE = (
+                _model_garden_models._SDK_GA_LAUNCH_STAGE
+            )
+
+            with pytest.raises(ValueError):
+                self.FakeModelGardenModel.from_pretrained("text-bison@001")

vertexai/_model_garden/_model_garden_models.py

@@ -15,7 +15,6 @@
 
 """Base class for working with Model Garden models."""
 
-import abc
 import dataclasses
 from typing import Dict, Optional, Type
 
@@ -25,20 +24,32 @@
 from google.cloud.aiplatform import models as aiplatform_models
 from google.cloud.aiplatform import _publisher_models
 
+from google.cloud.aiplatform.compat.types import (
+    publisher_model as gca_publisher_model,
+)
 
 _SUPPORTED_PUBLISHERS = ["google"]
 
 _SHORT_MODEL_ID_TO_TUNING_PIPELINE_MAP = {
     "text-bison": "https://us-kfp.pkg.dev/vertex-ai/large-language-model-pipelines/tune-large-model/sdk-1-25"
 }
 
+_SDK_PUBLIC_PREVIEW_LAUNCH_STAGE = frozenset(
+    [
+        gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
+        gca_publisher_model.PublisherModel.LaunchStage.GA,
+    ]
+)
+_SDK_GA_LAUNCH_STAGE = frozenset([gca_publisher_model.PublisherModel.LaunchStage.GA])
+
 _LOGGER = base.Logger(__name__)
 
 
 @dataclasses.dataclass
 class _ModelInfo:
     endpoint_name: str
     interface_class: Type["_ModelGardenModel"]
+    publisher_model_resource: _publisher_models._PublisherModel
     tuning_pipeline_uri: Optional[str] = None
     tuning_model_id: Optional[str] = None
 
@@ -114,6 +125,7 @@ def _get_model_info(
     return _ModelInfo(
         endpoint_name=endpoint_name,
         interface_class=interface_class,
+        publisher_model_resource=publisher_model_res,
         tuning_pipeline_uri=tuning_pipeline_uri,
         tuning_model_id=tuning_model_id,
     )
@@ -122,6 +134,28 @@ def _get_model_info(
 class _ModelGardenModel:
     """Base class for shared methods and properties across Model Garden models."""
 
+    _LAUNCH_STAGE: gca_publisher_model.PublisherModel.LaunchStage = (
+        _SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+    )
+
+    def _validate_launch_stage(
+        self,
+        publisher_model_resource: gca_publisher_model.PublisherModel,
+    ) -> None:
+        """Validates the model class _LAUNCH_STAGE matches the PublisherModel resource's launch stage.
+
+        Args:
+            publisher_model_resource (gca_publisher_model.PublisherModel
+                The GAPIC PublisherModel resource for this model.
+        """
+
+        publisher_launch_stage = publisher_model_resource.launch_stage
+
+        if publisher_launch_stage not in self._LAUNCH_STAGE:
+            raise ValueError(
+                f"The model you are trying to instantiate does not support the launch stage: {publisher_launch_stage.name}"
+            )
+
     # Subclasses override this attribute to specify their instance schema
     _INSTANCE_SCHEMA_URI: Optional[str] = None
 
@@ -174,6 +208,8 @@ def from_pretrained(cls, model_name: str) -> "_ModelGardenModel":
                 f"{model_name} is of type {model_info.interface_class.__name__} not of type {cls.__name__}"
             )
 
+        cls._validate_launch_stage(cls, model_info.publisher_model_resource)
+
         return model_info.interface_class(
             model_id=model_name,
             endpoint_name=model_info.endpoint_name,

vertexai/language_models/_language_models.py

@@ -126,6 +126,8 @@ def get_tuned_model(cls, tuned_model_name: str) -> "_LanguageModel":
             model_id=base_model_id,
             schema_to_class_map={cls._INSTANCE_SCHEMA_URI: cls},
         )
+        cls._validate_launch_stage(cls, model_info.publisher_model_resource)
+
         model = model_info.interface_class(
             model_id=base_model_id,
             endpoint_name=endpoint_name,
@@ -215,6 +217,8 @@ class TextGenerationModel(_LanguageModel):
         model.predict("What is life?")
     """
 
+    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
+
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/text_generation_1.0.0.yaml"
 
     _DEFAULT_TEMPERATURE = 0.0
@@ -300,7 +304,7 @@ def _batch_predict(
 class _PreviewTextGenerationModel(TextGenerationModel, _TunableModelMixin):
     """Tunable text generation model."""
 
-    pass
+    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
 
 
 class _ChatModel(TextGenerationModel):
@@ -427,6 +431,8 @@ class TextEmbeddingModel(_LanguageModel):
             print(len(vector))
     """
 
+    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
+
     _INSTANCE_SCHEMA_URI = (
         "gs://google-cloud-aiplatform/schema/predict/instance/text_embedding_1.0.0.yaml"
     )
@@ -447,6 +453,12 @@ def get_embeddings(self, texts: List[str]) -> List["TextEmbedding"]:
         ]
 
 
+class _PreviewTextEmbeddingModel(TextEmbeddingModel):
+    """Preview text embedding model."""
+
+    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+
+
 class TextEmbedding:
     """Contains text embedding vector."""
 
@@ -470,6 +482,8 @@ class InputOutputTextPair:
 class _ChatModelBase(_LanguageModel):
     """_ChatModelBase is a base class for chat models."""
 
+    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+
     def start_chat(
         self,
         *,
@@ -777,6 +791,7 @@ class CodeGenerationModel(_LanguageModel):
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml"
 
+    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
     _DEFAULT_TEMPERATURE = 0.0
     _DEFAULT_MAX_OUTPUT_TOKENS = 128
 

vertexai/preview/language_models.py

@@ -15,6 +15,7 @@
 """Classes for working with language models."""
 
 from vertexai.language_models._language_models import (
+    _PreviewTextEmbeddingModel,
     _PreviewTextGenerationModel,
     ChatModel,
     ChatSession,
@@ -23,11 +24,11 @@
     CodeGenerationModel,
     InputOutputTextPair,
     TextEmbedding,
-    TextEmbeddingModel,
     TextGenerationResponse,
 )
 
 TextGenerationModel = _PreviewTextGenerationModel
+TextEmbeddingModel = _PreviewTextEmbeddingModel
 
 __all__ = [
     "ChatModel",