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

vertex-sdk-bot · copybara-github · commit cd167b82c024 · 2022-12-19T03:41:11.000-08:00
PiperOrigin-RevId: 496359092
diff --git a/google/cloud/aiplatform/constants/prediction.py b/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
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -88,6 +88,7 @@
     "model.pkl",
     "model.joblib",
     "model.bst",
+    "model.mar",
     "saved_model.pb",
     "saved_model.pbtxt",
 ]
diff --git a/tests/unit/aiplatform/test_helpers.py b/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,
             ),
             (

Original file line number	Diff line number	Diff line change
`@@ -88,6 +88,7 @@`
`88`	`88`	`"model.pkl",`
`89`	`89`	`"model.joblib",`
`90`	`90`	`"model.bst",`
	`91`	`+ "model.mar",`
`91`	`92`	`"saved_model.pb",`
`92`	`93`	`"saved_model.pbtxt",`
`93`	`94`	`]`
Original file line number	Diff line number	Diff line change
`@@ -109,8 +109,11 @@ def test_correct_prediction_uri_args_with_init_location(self):`
`109`	`109`	`"args, expected_error_msg",`
`110`	`110`	`[`
`111`	`111`	`(`
`112`		`- ("pytorch", "1.10", None, None),`
`113`		- "No containers found for framework `pytorch`. Supported frameworks are",
	`112`	`+ ("pytorch", "1.0", None, None),`
	`113`	`+ (`
	`114`	+ "No serving container for `pytorch` version `1.0` with accelerator "
	`115`	+ "`cpu` found. Supported versions include"
	`116`	`+ ),`
`114`	`117`	`),`
`115`	`118`	`(`
`116`	`119`	`("tensorflow", "9.15", None, None),`
`@@ -177,6 +180,10 @@ def test_invalid_prediction_uri_args(self, args, expected_error_msg):`
`177`	`180`	`),`
`178`	`181`	`(`
`179`	`182`	`"europe-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-0:latest",`
	`183`	`+ True,`
	`184`	`+ ),`
	`185`	`+ (`
	`186`	`+ "europe-docker.pkg.dev/vertex-ai/prediction/onnx-cpu.1-0:latest",`
`180`	`187`	`False,`
`181`	`188`	`),`
`182`	`189`	`(`