chore: LLM - Removed the model launch stage limitations

Model garden models can have various launch stages: GA, public review, private preview.

Previously, the `ModelGardenModel` base class required a model interface class to be marked as preview to be able to use preview models.
This resulted in maintenance issues and unnecessary proliferations of classes, because each class had to be duplicated up to 3 times to be able to accommodate models in different launch stages. For example, some people with a private preview model had wither change their model launch stage or create a special private preview version of the model interface class just to work around the loading restrictions.

The launch stage restriction has confused some of our users. Some user change `text-bison` to `text-bison-32k` and the code starts throwing error.

With this change, the model interface classes no longer reject preview models. The breaking change guarantees now apply separately to the model interface classes and models. A preview class can change when going to GA. A preview model can change when going to GA.

PiperOrigin-RevId: 606430510

Author: Ark-kun

tests/unit/aiplatform/test_model_garden_models.py

@@ -54,8 +54,6 @@ class TestModelGardenModels:
 
     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):
@@ -84,23 +82,3 @@ 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

@@ -45,21 +45,6 @@
     "codechat-bison-32k": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-chat-model/v3.0.0",
 }
 
-_SDK_PRIVATE_PREVIEW_LAUNCH_STAGE = frozenset(
-    [
-        gca_publisher_model.PublisherModel.LaunchStage.PRIVATE_PREVIEW,
-        gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
-        gca_publisher_model.PublisherModel.LaunchStage.GA,
-    ]
-)
-_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__)
 
 T = TypeVar("T", bound="_ModelGardenModel")
@@ -241,10 +226,6 @@ def _from_pretrained(
             f"{model_name} is of type {model_info.interface_class.__name__} not of type {interface_class.__name__}"
         )
 
-    interface_class._validate_launch_stage(
-        interface_class, model_info.publisher_model_resource
-    )
-
     return model_info.interface_class(
         model_id=model_name,
         endpoint_name=model_info.endpoint_name,
@@ -254,31 +235,6 @@ def _from_pretrained(
 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 has launch stage '{publisher_launch_stage.name}'"
-                f", but the '{type(self).__module__}.{type(self).__name__}' class"
-                f" only supports the following launch stages: {self._LAUNCH_STAGE}."
-                " For preview models please use the classes from the `vertexai.preview.*` namespace."
-            )
-
     # Subclasses override this attribute to specify their instance schema
     _INSTANCE_SCHEMA_URI: Optional[str] = None
 

vertexai/language_models/_language_models.py

@@ -205,8 +205,6 @@ 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,
@@ -1241,8 +1239,6 @@ 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_MAX_OUTPUT_TOKENS = 128
@@ -1984,8 +1980,6 @@ class TextEmbeddingModel(_LanguageModel):
 
     __module__ = "vertexai.language_models"
 
-    _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"
     )
@@ -2121,8 +2115,6 @@ class _PreviewTextEmbeddingModel(
     __name__ = "TextEmbeddingModel"
     __module__ = "vertexai.preview.language_models"
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
-
 
 @dataclasses.dataclass
 class TextEmbeddingStatistics:
@@ -2173,8 +2165,6 @@ class ChatMessage:
 class _ChatModelBase(_LanguageModel):
     """_ChatModelBase is a base class for chat models."""
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
-
     def start_chat(
         self,
         *,
@@ -2251,8 +2241,6 @@ class _PreviewChatModel(ChatModel, _PreviewTunableChatModelMixin):
     __name__ = "ChatModel"
     __module__ = "vertexai.preview.language_models"
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
-
     def start_chat(
         self,
         *,
@@ -2313,7 +2301,6 @@ class CodeChatModel(_ChatModelBase, _TunableChatModelMixin):
     __module__ = "vertexai.language_models"
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/codechat_generation_1.0.0.yaml"
-    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
 
     def start_chat(
         self,
@@ -2351,8 +2338,6 @@ class _PreviewCodeChatModel(CodeChatModel):
     __name__ = "CodeChatModel"
     __module__ = "vertexai.preview.language_models"
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
-
     def start_chat(
         self,
         *,
@@ -3122,7 +3107,6 @@ 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_GA_LAUNCH_STAGE
 
     def _create_prediction_request(
         self,
@@ -3390,8 +3374,6 @@ class _PreviewCodeGenerationModel(CodeGenerationModel, _CountTokensCodeGeneratio
     __name__ = "CodeGenerationModel"
     __module__ = "vertexai.preview.language_models"
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
-
 
 ###### Model tuning
 # Currently, tuning can only work in this location
@@ -3710,5 +3692,3 @@ class _PreviewTextGenerationModel(
     # Do not add docstring so that it's inherited from the base class.
     __name__ = "TextGenerationModel"
     __module__ = "vertexai.preview.language_models"
-
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE

vertexai/preview/vision_models.py

@@ -15,7 +15,6 @@
 """Classes for working with vision models."""
 
 from vertexai.vision_models._vision_models import (
-    _PreviewImageTextModel,
     Image,
     ImageGenerationModel,
     ImageGenerationResponse,
@@ -27,8 +26,6 @@
     MultiModalEmbeddingResponse,
 )
 
-ImageTextModel = _PreviewImageTextModel
-
 __all__ = [
     "Image",
     "ImageGenerationModel",

vertexai/vision_models/_vision_models.py

@@ -599,9 +599,6 @@ class ImageCaptioningModel(
     __module__ = "vertexai.vision_models"
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/vision_reasoning_model_1.0.0.yaml"
-    _LAUNCH_STAGE = (
-        _model_garden_models._SDK_GA_LAUNCH_STAGE  # pylint: disable=protected-access
-    )
 
     def get_captions(
         self,
@@ -667,9 +664,6 @@ class ImageQnAModel(
     __module__ = "vertexai.vision_models"
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/vision_reasoning_model_1.0.0.yaml"
-    _LAUNCH_STAGE = (
-        _model_garden_models._SDK_GA_LAUNCH_STAGE  # pylint: disable=protected-access
-    )
 
     def ask_question(
         self,
@@ -729,10 +723,6 @@ class MultiModalEmbeddingModel(_model_garden_models._ModelGardenModel):
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/vision_embedding_model_1.0.0.yaml"
 
-    _LAUNCH_STAGE = (
-        _model_garden_models._SDK_GA_LAUNCH_STAGE  # pylint: disable=protected-access
-    )
-
     def get_embeddings(
         self,
         image: Optional[Image] = None,
@@ -847,13 +837,3 @@ class ImageTextModel(ImageCaptioningModel, ImageQnAModel):
     # since SDK Model Garden classes should follow the design pattern of exactly 1 SDK class to 1 Model Garden schema URI
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/vision_reasoning_model_1.0.0.yaml"
-    _LAUNCH_STAGE = (
-        _model_garden_models._SDK_GA_LAUNCH_STAGE  # pylint: disable=protected-access
-    )
-
-
-class _PreviewImageTextModel(ImageTextModel):
-
-    __module__ = "vertexai.preview.vision_models"
-
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE