feat: Release Vertex SDK Preview

PiperOrigin-RevId: 559477466

.kokoro/continuous/system.cfg

@@ -2,7 +2,7 @@
 
 env_vars: {
     key: "NOX_SESSION"
-    value: "system-3.8"
+    value: "system-3.11"
 }
 
 # Run system tests in parallel, splitting up by file

.kokoro/presubmit/system.cfg

@@ -3,7 +3,7 @@
 # Run system tests when test files are modified
 env_vars: {
     key: "NOX_SESSION"
-    value: "system-3.8"
+    value: "system-3.11"
 }
 
 # Run system tests in parallel, splitting up by file

noxfile.py

@@ -48,7 +48,7 @@
 ]
 UNIT_TEST_EXTRAS_BY_PYTHON = {}
 
-SYSTEM_TEST_PYTHON_VERSIONS = ["3.8"]
+SYSTEM_TEST_PYTHON_VERSIONS = ["3.11"]
 SYSTEM_TEST_STANDARD_DEPENDENCIES = [
     "mock",
     "pytest",

setup.py

@@ -86,6 +86,12 @@
 
 autologging_extra_require = ["mlflow>=1.27.0,<=2.1.1"]
 
+preview_extra_require = [
+    "cloudpickle < 3.0",
+    "google-cloud-logging < 4.0",
+    "importlib-metadata < 7.0; python_version<'3.8'",
+]
+
 full_extra_require = list(
     set(
         tensorboard_extra_require
@@ -100,19 +106,24 @@
         + prediction_extra_require
         + private_endpoints_extra_require
         + autologging_extra_require
+        + preview_extra_require
     )
 )
 testing_extra_require = (
     full_extra_require
     + profiler_extra_require
     + [
         "grpcio-testing",
-        "pytest-asyncio",
-        "pytest-xdist",
         "ipython",
         "kfp",
-        "xgboost",
+        "pyfakefs",
+        "pytest-asyncio",
+        "pytest-xdist",
         "scikit-learn",
+        "tensorflow >=2.3.0, < 2.13.0",
+        "torch >= 2.0.0; python_version>='3.8'",
+        "torch; python_version<'3.8'",
+        "xgboost",
     ]
 )
 
@@ -160,6 +171,7 @@
         "datasets": datasets_extra_require,
         "private_endpoints": private_endpoints_extra_require,
         "autologging": autologging_extra_require,
+        "preview": preview_extra_require,
     },
     python_requires=">=3.7",
     classifiers=[

testing/constraints-3.10.txt

@@ -1,6 +1,10 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
-proto-plus
-protobuf
+google-api-core==1.32.0
+proto-plus==1.22.0
+protobuf==3.19.5
+mock==4.0.2
+google-cloud-storage==2.0.0
+packaging==20.0 # Increased for compatibility with MLFlow
+grpcio-testing==1.34.0

testing/constraints-3.7.txt

@@ -9,6 +9,6 @@ google-api-core==1.32.0
 proto-plus==1.22.0
 protobuf==3.19.5
 mock==4.0.2
-google-cloud-storage==1.32.0
+google-cloud-storage==2.0.0
 packaging==20.0 # Increased for compatibility with MLFlow
 grpcio-testing==1.34.0

testing/constraints-3.8.txt

@@ -1,6 +1,10 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
-proto-plus
-protobuf
+google-api-core==1.32.0
+proto-plus==1.22.0
+protobuf==3.19.5
+mock==4.0.2
+google-cloud-storage==2.0.0
+packaging==20.0 # Increased for compatibility with MLFlow
+grpcio-testing==1.34.0

testing/constraints-3.9.txt

@@ -1,6 +1,10 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
-proto-plus
-protobuf
+google-api-core==1.32.0
+proto-plus==1.22.0
+protobuf==3.19.5
+mock==4.0.2
+google-cloud-storage==2.0.0
+packaging==20.0 # Increased for compatibility with MLFlow
+grpcio-testing==1.34.0

tests/system/aiplatform/e2e_base.py

@@ -26,6 +26,7 @@
 
 from google.api_core import exceptions
 from google.cloud import aiplatform
+import vertexai
 from google.cloud import bigquery
 from google.cloud import resourcemanager
 from google.cloud import storage
@@ -62,6 +63,7 @@ def _make_display_name(cls, key: str) -> str:
     def setup_method(self):
         importlib.reload(initializer)
         importlib.reload(aiplatform)
+        importlib.reload(vertexai)
 
     @pytest.fixture(scope="class")
     def shared_state(self) -> Generator[Dict[str, Any], None, None]:

tests/system/vertexai/test_pytorch.py

@@ -0,0 +1,158 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from unittest import mock
+
+import vertexai
+from tests.system.aiplatform import e2e_base
+from vertexai.preview._workflow.executor import training
+import pytest
+from sklearn.datasets import load_iris
+import torch
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+
+
+@mock.patch.object(
+    training,
+    "VERTEX_AI_DEPENDENCY_PATH",
+    "google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/"
+    f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
+    if os.environ.get("KOKORO_GIT_COMMIT")
+    else "google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/python-aiplatform.git@copybara_557913723",
+)
+@mock.patch.object(
+    training,
+    "VERTEX_AI_DEPENDENCY_PATH_AUTOLOGGING",
+    "google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/"
+    f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
+    if os.environ.get("KOKORO_GIT_COMMIT")
+    else "google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/python-aiplatform.git@copybara_557913723",
+)
+@pytest.mark.usefixtures(
+    "prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources"
+)
+class TestRemoteExecutionPytorch(e2e_base.TestEndToEnd):
+
+    _temp_prefix = "temp-vertexai-remote-execution"
+
+    def test_remote_execution_pytorch(self, shared_state):
+        # Define the pytorch custom model
+        class TorchLogisticRegression(vertexai.preview.VertexModel, torch.nn.Module):
+            def __init__(self, input_size: int, output_size: int):
+                torch.nn.Module.__init__(self)
+                vertexai.preview.VertexModel.__init__(self)
+                self.linear = torch.nn.Linear(input_size, output_size)
+                self.softmax = torch.nn.Softmax(dim=1)
+
+            def forward(self, x):
+                return self.softmax(self.linear(x))
+
+            @vertexai.preview.developer.mark.train()
+            def train(self, dataloader, num_epochs, lr):
+                criterion = torch.nn.CrossEntropyLoss()
+                optimizer = torch.optim.SGD(self.parameters(), lr=lr)
+
+                for t in range(num_epochs):
+                    for idx, batch in enumerate(dataloader):
+                        # move data to the same device as model
+                        device = next(self.parameters()).device
+                        x, y = batch[0].to(device), batch[1].to(device)
+
+                        optimizer.zero_grad()
+                        pred = self(x)
+                        loss = criterion(pred, y)
+                        loss.backward()
+                        optimizer.step()
+
+            @vertexai.preview.developer.mark.predict()
+            def predict(self, X):
+                X = torch.tensor(X).to(torch.float32)
+                with torch.no_grad():
+                    pred = torch.argmax(self(X), dim=1)
+                return pred
+
+        # Initialize vertexai
+        vertexai.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+            staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
+        )
+
+        # Prepare dataset
+        dataset = load_iris()
+
+        X, X_retrain, y, y_retrain = train_test_split(
+            dataset.data, dataset.target, test_size=0.60, random_state=42
+        )
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=0.20, random_state=42
+        )
+
+        transformer = StandardScaler()
+        X_train = transformer.fit_transform(X_train)
+        X_test = transformer.transform(X_test)
+        X_retrain = transformer.transform(X_retrain)
+
+        train_loader = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(
+                torch.tensor(X_train).to(torch.float32),
+                torch.tensor(y_train),
+            ),
+            batch_size=10,
+            shuffle=True,
+        )
+
+        retrain_loader = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(
+                torch.tensor(X_retrain).to(torch.float32),
+                torch.tensor(y_retrain),
+            ),
+            batch_size=10,
+            shuffle=True,
+        )
+
+        # Remote CPU training on Torch custom model
+        vertexai.preview.init(remote=True)
+
+        model = TorchLogisticRegression(4, 3)
+        model.train.vertex.remote_config.display_name = self._make_display_name(
+            "pytorch-cpu-training"
+        )
+        model.train(train_loader, num_epochs=100, lr=0.05)
+
+        # Remote prediction on Torch custom model
+        model.predict.vertex.remote_config.display_name = self._make_display_name(
+            "pytorch-prediction"
+        )
+        model.predict(X_test)
+
+        # Register trained model
+        registered_model = vertexai.preview.register(model)
+        shared_state["resources"] = [registered_model]
+
+        # Load the registered model
+        pulled_model = vertexai.preview.from_pretrained(
+            model_name=registered_model.resource_name
+        )
+
+        # Uptrain the pretrained model on CPU
+        pulled_model.train.vertex.remote_config.display_name = self._make_display_name(
+            "pytorch-cpu-uptraining"
+        )
+        pulled_model.train(retrain_loader, num_epochs=100, lr=0.05)