feat: Vertex AI Model Garden open model export SDK Public Preview

PiperOrigin-RevId: 742429430

Author: vertex-sdk-bot

tests/unit/vertexai/model_garden/test_model_garden.py

@@ -48,7 +48,7 @@
     "publishers/hf-meta-llama/models/llama-3.3-70b-instruct@001"
 )
 _TEST_HUGGING_FACE_ACCESS_TOKEN = "test-access-token"
-
+_TEST_GCS_URI = "gs://some-bucket/some-model"
 _TEST_ENDPOINT_NAME = "projects/test-project/locations/us-central1/endpoints/1234567890"
 _TEST_MODEL_NAME = "projects/test-project/locations/us-central1/models/9876543210"
 _TEST_MODEL_CONTAINER_SPEC = types.ModelContainerSpec(
@@ -85,6 +85,22 @@ def google_auth_mock():
         yield google_auth_mock
 
 
+@pytest.fixture
+def export_publisher_model_mock():
+    """Mocks the export_publisher_model method."""
+    with mock.patch.object(
+        model_garden_service.ModelGardenServiceClient,
+        "export_publisher_model",
+    ) as export_publisher_model:
+        mock_export_lro = mock.Mock(ga_operation.Operation)
+        mock_export_lro.result.return_value = types.ExportPublisherModelResponse(
+            publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+            destination_uri=_TEST_GCS_URI,
+        )
+        export_publisher_model.return_value = mock_export_lro
+        yield export_publisher_model
+
+
 @pytest.fixture
 def deploy_mock():
     """Mocks the deploy method."""
@@ -338,6 +354,7 @@ def list_publisher_models_mock():
     "deploy_mock",
     "get_publisher_model_mock",
     "list_publisher_models_mock",
+    "export_publisher_model_mock",
 )
 class TestModelGarden:
     """Test cases for ModelGarden class."""
@@ -350,6 +367,54 @@ def setup_method(self):
     def teardown_method(self):
         aiplatform.initializer.global_pool.shutdown(wait=True)
 
+    def test_export_full_resource_name_success(self, export_publisher_model_mock):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
+        model.export(_TEST_GCS_URI)
+        export_publisher_model_mock.assert_called_once_with(
+            types.ExportPublisherModelRequest(
+                parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
+                name=_TEST_MODEL_FULL_RESOURCE_NAME,
+                destination=types.GcsDestination(output_uri_prefix=_TEST_GCS_URI),
+            ),
+            metadata=[("x-goog-user-project", "test-project")],
+        )
+
+    def test_export_simplified_resource_name_success(self, export_publisher_model_mock):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        model = model_garden.OpenModel(model_name=_TEST_MODEL_SIMPLIFIED_RESOURCE_NAME)
+        model.export(_TEST_GCS_URI)
+        export_publisher_model_mock.assert_called_once_with(
+            types.ExportPublisherModelRequest(
+                parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
+                name=_TEST_MODEL_FULL_RESOURCE_NAME,
+                destination=types.GcsDestination(output_uri_prefix=_TEST_GCS_URI),
+            ),
+            metadata=[("x-goog-user-project", "test-project")],
+        )
+
+    def test_export_hugging_face_id_success(self, export_publisher_model_mock):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        model = model_garden.OpenModel(model_name=_TEST_MODEL_HUGGING_FACE_ID)
+        model.export(_TEST_GCS_URI)
+        export_publisher_model_mock.assert_called_once_with(
+            types.ExportPublisherModelRequest(
+                parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
+                name=_TEST_HUGGING_FACE_MODEL_FULL_RESOURCE_NAME,
+                destination=types.GcsDestination(output_uri_prefix=_TEST_GCS_URI),
+            ),
+            metadata=[("x-goog-user-project", "test-project")],
+        )
+
     def test_deploy_full_resource_name_success(self, deploy_mock):
         aiplatform.init(
             project=_TEST_PROJECT,

vertexai/model_garden/_model_garden.py

@@ -36,6 +36,7 @@
 _LOGGER = base.Logger(__name__)
 _DEFAULT_VERSION = compat.V1BETA1
 _DEFAULT_TIMEOUT = 2 * 60 * 60  # 2 hours, same as UI one-click deployment.
+_DEFAULT_EXPORT_TIMEOUT = 1 * 60 * 60  # 1 hour.
 _HF_WILDCARD_FILTER = "is_hf_wildcard(true)"
 _NATIVE_MODEL_FILTER = "is_hf_wildcard(false)"
 _VERIFIED_DEPLOYMENT_FILTER = (
@@ -308,6 +309,53 @@ def _us_central1_model_garden_client(
             location_override="us-central1",
         )
 
+    def export(
+        self,
+        target_gcs_path: str = "",
+        export_request_timeout: Optional[float] = None,
+    ) -> str:
+        """Exports an Open Model to a google cloud storage bucket.
+
+        Args:
+            target_gcs_path: target gcs path.
+            export_request_timeout: The timeout for the deploy request. Default is 2
+              hours.
+
+        Returns:
+            str: the target gcs bucket where the model weights are downloaded to
+
+
+        Raises:
+            ValueError: If ``target_gcs_path`` is not specified
+        """
+        if not target_gcs_path:
+            raise ValueError("target_gcs_path is required.")
+
+        request = types.ExportPublisherModelRequest(
+            parent=f"projects/{self._project}/locations/{self._location}",
+            name=self._publisher_model_name,
+            destination=types.GcsDestination(output_uri_prefix=target_gcs_path),
+        )
+        request_headers = [
+            ("x-goog-user-project", "{}".format(initializer.global_config.project)),
+        ]
+
+        _LOGGER.info(f"Exporting model weights: {self._model_name}")
+
+        operation_future = self._model_garden_client.export_publisher_model(
+            request, metadata=request_headers
+        )
+        _LOGGER.info(f"LRO: {operation_future.operation.name}")
+
+        _LOGGER.info(f"Start time: {datetime.datetime.now()}")
+        export_publisher_model_response = operation_future.result(
+            timeout=export_request_timeout or _DEFAULT_EXPORT_TIMEOUT
+        )
+        _LOGGER.info(f"End time: {datetime.datetime.now()}")
+        _LOGGER.info(f"Response: {export_publisher_model_response}")
+
+        return target_gcs_path
+
     def deploy(
         self,
         accept_eula: bool = False,