chore: Add PyTorch first-party images to prediction first-party uris and allow PyTorch artifacts.

PiperOrigin-RevId: 496359092

Author: vertex-sdk-bot

google/cloud/aiplatform/constants/prediction.py

@@ -25,7 +25,7 @@
 CONTAINER_URI_REGEX = (
     r"^(us|europe|asia)-docker.pkg.dev/"
     r"vertex-ai/prediction/"
-    r"(tf|sklearn|xgboost).+$"
+    r"(tf|sklearn|xgboost|pytorch).+$"
 )
 
 SKLEARN = "sklearn"
@@ -147,8 +147,20 @@
     "asia-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest",
 ]
 
+PYTORCH_CONTAINER_URIS = [
+    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest",
+    "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest",
+    "asia-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest",
+    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-11:latest",
+    "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-11:latest",
+    "asia-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-11:latest",
+]
+
 SERVING_CONTAINER_URIS = (
-    SKLEARN_CONTAINER_URIS + TF_CONTAINER_URIS + XGBOOST_CONTAINER_URIS
+    SKLEARN_CONTAINER_URIS
+    + TF_CONTAINER_URIS
+    + XGBOOST_CONTAINER_URIS
+    + PYTORCH_CONTAINER_URIS
 )
 
 # Map of all first-party prediction containers

google/cloud/aiplatform/models.py

@@ -88,6 +88,7 @@
     "model.pkl",
     "model.joblib",
     "model.bst",
+    "model.mar",
     "saved_model.pb",
     "saved_model.pbtxt",
 ]

tests/unit/aiplatform/test_helpers.py

@@ -109,8 +109,11 @@ def test_correct_prediction_uri_args_with_init_location(self):
         "args, expected_error_msg",
         [
             (
-                ("pytorch", "1.10", None, None),
-                "No containers found for framework `pytorch`. Supported frameworks are",
+                ("pytorch", "1.0", None, None),
+                (
+                    "No serving container for `pytorch` version `1.0` with accelerator "
+                    "`cpu` found. Supported versions include"
+                ),
             ),
             (
                 ("tensorflow", "9.15", None, None),
@@ -177,6 +180,10 @@ def test_invalid_prediction_uri_args(self, args, expected_error_msg):
             ),
             (
                 "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-0:latest",
+                True,
+            ),
+            (
+                "europe-docker.pkg.dev/vertex-ai/prediction/onnx-cpu.1-0:latest",
                 False,
             ),
             (