feat: Added deploy_model_with_dedicated_resources_sample, deploy_mode…

…l_with_automatic_resources_sample, upload_model and get_model samples (#337) * Added deploy_model and upload_model samples * Fixed tests * Fixed tests * Added get_model sample and fixed errors * Added get_model sample and fixed errors * Added get_model sample and fixed errors * Added deploy_model_custom_trained_model_sample.py * Small tweaks * Fixed arguments * Removed explain args from custom training sample * Removed comment about AutoML * Renamed samples * Added note about explanations * Fixed test * Fixed tests * Ran linter
googleapis · Apr 30, 2021 · ef4f6f8 · ef4f6f8
1 parent 4241738
commit ef4f6f8
Show file tree

Hide file tree

Showing 11 changed files with 505 additions and 3 deletions.
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -220,7 +220,7 @@ def mock_model():
 
 
 @pytest.fixture
-def mock_get_model(mock_model):
+def mock_init_model(mock_model):
     with patch.object(aiplatform, "Model") as mock:
         mock.return_value = mock_model
         yield mock
@@ -232,6 +232,19 @@ def mock_batch_predict_model(mock_model):
         yield mock
 
 
+@pytest.fixture
+def mock_upload_model():
+    with patch.object(aiplatform.models.Model, "upload") as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_deploy_model(mock_model, mock_endpoint):
+    with patch.object(mock_model, "deploy") as mock:
+        mock.return_value = mock_endpoint
+        yield mock
+
+
 """
 ----------------------------------------------------------------------------
 Job Fixtures

diff --git a/samples/model-builder/create_batch_prediction_job_sample_test.py b/samples/model-builder/create_batch_prediction_job_sample_test.py
@@ -18,7 +18,7 @@
 
 
 def test_create_batch_prediction_job_sample(
-    mock_sdk_init, mock_get_model, mock_batch_predict_model
+    mock_sdk_init, mock_model, mock_init_model, mock_batch_predict_model
 ):
 
     create_batch_prediction_job_sample.create_batch_prediction_job_sample(
@@ -33,7 +33,7 @@ def test_create_batch_prediction_job_sample(
     mock_sdk_init.assert_called_once_with(
         project=constants.PROJECT, location=constants.LOCATION
     )
-    mock_get_model.assert_called_once_with(constants.MODEL_NAME)
+    mock_init_model.assert_called_once_with(constants.MODEL_NAME)
     mock_batch_predict_model.assert_called_once_with(
         job_display_name=constants.DISPLAY_NAME,
         gcs_source=constants.GCS_SOURCES,

diff --git a/samples/model-builder/deploy_model_with_automatic_resources_sample.py b/samples/model-builder/deploy_model_with_automatic_resources_sample.py
@@ -0,0 +1,57 @@
+# Copyright 2021 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
+#
+#     https://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.
+
+from typing import Dict, Optional, Sequence, Tuple
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_deploy_model_with_automatic_resources_sample]
+def deploy_model_with_automatic_resources_sample(
+    project,
+    location,
+    model_name: str,
+    endpoint: Optional[aiplatform.Endpoint] = None,
+    deployed_model_display_name: Optional[str] = None,
+    traffic_percentage: Optional[int] = 0,
+    traffic_split: Optional[Dict[str, int]] = None,
+    min_replica_count: int = 1,
+    max_replica_count: int = 1,
+    metadata: Optional[Sequence[Tuple[str, str]]] = (),
+    sync: bool = True,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    model = aiplatform.Model(model_name=model_name)
+
+    model.deploy(
+        endpoint=endpoint,
+        deployed_model_display_name=deployed_model_display_name,
+        traffic_percentage=traffic_percentage,
+        traffic_split=traffic_split,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+        metadata=metadata,
+        sync=sync,
+    )
+
+    model.wait()
+
+    print(model.display_name)
+    print(model.resource_name)
+    return model
+
+
+#  [END aiplatform_sdk_deploy_model_with_automatic_resources_sample]
diff --git a/samples/model-builder/deploy_model_with_automatic_resources_test.py b/samples/model-builder/deploy_model_with_automatic_resources_test.py
@@ -0,0 +1,52 @@
+# Copyright 2021 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
+#
+#     https://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 deploy_model_with_automatic_resources_sample
+import test_constants as constants
+
+
+def test_deploy_model_with_automatic_resources_sample(
+    mock_sdk_init, mock_model, mock_init_model, mock_deploy_model,
+):
+
+    deploy_model_with_automatic_resources_sample.deploy_model_with_automatic_resources_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        model_name=constants.MODEL_NAME,
+        endpoint=constants.ENDPOINT_NAME,
+        deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
+        traffic_percentage=constants.TRAFFIC_PERCENTAGE,
+        traffic_split=constants.TRAFFIC_SPLIT,
+        min_replica_count=constants.MIN_REPLICA_COUNT,
+        max_replica_count=constants.MAX_REPLICA_COUNT,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_init_model.assert_called_once_with(model_name=constants.MODEL_NAME)
+
+    mock_deploy_model.assert_called_once_with(
+        endpoint=constants.ENDPOINT_NAME,
+        deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
+        traffic_percentage=constants.TRAFFIC_PERCENTAGE,
+        traffic_split=constants.TRAFFIC_SPLIT,
+        min_replica_count=constants.MIN_REPLICA_COUNT,
+        max_replica_count=constants.MAX_REPLICA_COUNT,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
+        sync=True,
+    )
diff --git a/samples/model-builder/deploy_model_with_dedicated_resources_sample.py b/samples/model-builder/deploy_model_with_dedicated_resources_sample.py
@@ -0,0 +1,70 @@
+# Copyright 2021 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
+#
+#     https://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.
+
+from typing import Dict, Optional, Sequence, Tuple
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import explain
+
+
+#  [START aiplatform_sdk_deploy_model_with_dedicated_resources_sample]
+def deploy_model_with_dedicated_resources_sample(
+    project,
+    location,
+    model_name: str,
+    machine_type: str,
+    endpoint: Optional[aiplatform.Endpoint] = None,
+    deployed_model_display_name: Optional[str] = None,
+    traffic_percentage: Optional[int] = 0,
+    traffic_split: Optional[Dict[str, int]] = None,
+    min_replica_count: int = 1,
+    max_replica_count: int = 1,
+    accelerator_type: Optional[str] = None,
+    accelerator_count: Optional[int] = None,
+    explanation_metadata: Optional[explain.ExplanationMetadata] = None,
+    explanation_parameters: Optional[explain.ExplanationParameters] = None,
+    metadata: Optional[Sequence[Tuple[str, str]]] = (),
+    sync: bool = True,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    model = aiplatform.Model(model_name=model_name)
+
+    # The explanation_metadata and explanation_parameters should only be
+    # provided for a custom trained model and not an AutoML model.
+    model.deploy(
+        endpoint=endpoint,
+        deployed_model_display_name=deployed_model_display_name,
+        traffic_percentage=traffic_percentage,
+        traffic_split=traffic_split,
+        machine_type=machine_type,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+        accelerator_type=accelerator_type,
+        accelerator_count=accelerator_count,
+        explanation_metadata=explanation_metadata,
+        explanation_parameters=explanation_parameters,
+        metadata=metadata,
+        sync=sync,
+    )
+
+    model.wait()
+
+    print(model.display_name)
+    print(model.resource_name)
+    return model
+
+
+#  [END aiplatform_sdk_deploy_model_with_dedicated_resources_sample]
diff --git a/samples/model-builder/deploy_model_with_dedicated_resources_test.py b/samples/model-builder/deploy_model_with_dedicated_resources_test.py
@@ -0,0 +1,62 @@
+# Copyright 2021 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
+#
+#     https://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 deploy_model_with_dedicated_resources_sample
+import test_constants as constants
+
+
+def test_deploy_model_with_dedicated_resources_sample(
+    mock_sdk_init, mock_model, mock_init_model, mock_deploy_model
+):
+
+    deploy_model_with_dedicated_resources_sample.deploy_model_with_dedicated_resources_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        machine_type=constants.MACHINE_TYPE,
+        model_name=constants.MODEL_NAME,
+        endpoint=constants.ENDPOINT_NAME,
+        deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
+        traffic_percentage=constants.TRAFFIC_PERCENTAGE,
+        traffic_split=constants.TRAFFIC_SPLIT,
+        min_replica_count=constants.MIN_REPLICA_COUNT,
+        max_replica_count=constants.MAX_REPLICA_COUNT,
+        accelerator_type=constants.ACCELERATOR_TYPE,
+        accelerator_count=constants.ACCELERATOR_COUNT,
+        explanation_metadata=constants.EXPLANATION_METADATA,
+        explanation_parameters=constants.EXPLANATION_PARAMETERS,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_init_model.assert_called_once_with(model_name=constants.MODEL_NAME)
+
+    mock_deploy_model.assert_called_once_with(
+        endpoint=constants.ENDPOINT_NAME,
+        deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
+        traffic_percentage=constants.TRAFFIC_PERCENTAGE,
+        traffic_split=constants.TRAFFIC_SPLIT,
+        machine_type=constants.MACHINE_TYPE,
+        min_replica_count=constants.MIN_REPLICA_COUNT,
+        max_replica_count=constants.MAX_REPLICA_COUNT,
+        accelerator_type=constants.ACCELERATOR_TYPE,
+        accelerator_count=constants.ACCELERATOR_COUNT,
+        explanation_metadata=constants.EXPLANATION_METADATA,
+        explanation_parameters=constants.EXPLANATION_PARAMETERS,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
+        sync=True,
+    )
diff --git a/samples/model-builder/get_model_sample.py b/samples/model-builder/get_model_sample.py
@@ -0,0 +1,31 @@
+# Copyright 2021 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
+#
+#     https://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.
+
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_get_model_sample]
+def get_model_sample(project: str, location: str, model_name: str):
+
+    aiplatform.init(project=project, location=location)
+
+    model = aiplatform.Model(model_name=model_name)
+
+    print(model.display_name)
+    print(model.resource_name)
+    return model
+
+
+#  [END aiplatform_sdk_get_model_sample]
diff --git a/samples/model-builder/get_model_test.py b/samples/model-builder/get_model_test.py
@@ -0,0 +1,32 @@
+# Copyright 2021 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
+#
+#     https://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 get_model_sample
+import test_constants as constants
+
+
+def test_get_model_sample(mock_sdk_init, mock_init_model):
+
+    get_model_sample.get_model_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        model_name=constants.MODEL_NAME,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_init_model.assert_called_once_with(model_name=constants.MODEL_NAME)