feat: Support referenced models in SDK.

PiperOrigin-RevId: 623654808
googleapis · Apr 11, 2024 · c9b6b8b · c9b6b8b
1 parent 2d43438
commit c9b6b8b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -3251,7 +3251,9 @@ def upload(
             container_spec = local_model.get_serving_container_spec()
             appended_user_agent = [prediction_constants.CUSTOM_PREDICTION_ROUTINES]
         else:
-            if not serving_container_image_uri:
+            # Referenced models do not have container_image and artifact_uri
+            # Skip the container_image if this is a referenced model
+            if not serving_container_image_uri and artifact_uri:
                 raise ValueError(
                     "The parameter `serving_container_image_uri` is required "
                     "if no `local_model` is provided."

diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py
@@ -1293,6 +1293,7 @@ def test_upload_without_serving_container_image_uri_throw_error(
                 display_name=_TEST_MODEL_NAME,
                 serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE,
                 serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE,
+                artifact_uri=_TEST_ARTIFACT_URI,
             )
 
         assert str(exception.value) == expected_message