docs: new samples for model serialization

PiperOrigin-RevId: 504654663

Author: jaycee-li

samples/model-builder/conftest.py

@@ -896,9 +896,16 @@ def mock_log_classification_metrics():
 
 @pytest.fixture
 def mock_log_model():
-    with patch.object(aiplatform, "log_model") as mock_log_metrics:
-        mock_log_metrics.return_value = None
-        yield mock_log_metrics
+    with patch.object(aiplatform, "log_model") as mock_log_model:
+        mock_log_model.return_value = None
+        yield mock_log_model
+
+
+@pytest.fixture
+def mock_save_model():
+    with patch.object(aiplatform, "save_model") as mock_save_model:
+        mock_save_model.return_value = None
+        yield mock_save_model
 
 
 @pytest.fixture

samples/model-builder/experiment_tracking/get_experiment_model_sample.py

@@ -0,0 +1,30 @@
+# 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
+#
+#     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_experiment_model_sample]
+def get_experiment_model_sample(
+    project: str,
+    location: str,
+    artifact_id: str,
+) -> "ExperimentModel":  # noqa: F821
+    aiplatform.init(project=project, location=location)
+    experiment_model = aiplatform.get_experiment_model(artifact_id=artifact_id)
+
+    return experiment_model
+
+
+#  [END aiplatform_sdk_get_experiment_model_sample]

samples/model-builder/experiment_tracking/get_experiment_model_sample_test.py

@@ -0,0 +1,33 @@
+# 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
+#
+#     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_experiment_model_sample
+
+import pytest
+
+import test_constants as constants
+
+
+@pytest.mark.usefixtures("mock_sdk_init")
+def test_get_experiment_model_sample(mock_get_experiment_model):
+
+    get_experiment_model_sample.get_experiment_model_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        artifact_id=constants.EXPERIMENT_MODEL_ID,
+    )
+
+    mock_get_experiment_model.assert_called_once_with(
+        artifact_id=constants.EXPERIMENT_MODEL_ID,
+    )

samples/model-builder/experiment_tracking/save_model_sample.py

@@ -0,0 +1,45 @@
+# 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
+#
+#     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 Optional, Union
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_save_model_sample]
+def save_model_sample(
+    project: str,
+    location: str,
+    model: Union[
+        "sklearn.base.BaseEstimator", "xgb.Booster", "tf.Module"  # noqa: F821
+    ],
+    artifact_id: Optional[str] = None,
+    uri: Optional[str] = None,
+    input_example: Optional[
+        Union[list, dict, "pd.DataFrame", "np.ndarray"]  # noqa: F821
+    ] = None,
+    display_name: Optional[str] = None,
+) -> None:
+    aiplatform.init(project=project, location=location)
+
+    aiplatform.save_model(
+        model=model,
+        artifact_id=artifact_id,
+        uri=uri,
+        input_example=input_example,
+        display_name=display_name,
+    )
+
+
+#  [END aiplatform_sdk_save_model_sample]

samples/model-builder/experiment_tracking/save_model_sample_test.py

@@ -0,0 +1,41 @@
+# 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
+#
+#     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 pytest
+
+import save_model_sample
+
+import test_constants as constants
+
+
+@pytest.mark.usefixtures("mock_sdk_init")
+def test_save_model_sample(mock_save_model):
+
+    save_model_sample.save_model_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        model=constants.ML_MODEL,
+        artifact_id=constants.EXPERIMENT_MODEL_ID,
+        uri=constants.MODEL_ARTIFACT_URI,
+        input_example=constants.EXPERIMENT_MODEL_INPUT_EXAMPLE,
+        display_name=constants.DISPLAY_NAME,
+    )
+
+    mock_save_model.assert_called_once_with(
+        model=constants.ML_MODEL,
+        artifact_id=constants.EXPERIMENT_MODEL_ID,
+        uri=constants.MODEL_ARTIFACT_URI,
+        input_example=constants.EXPERIMENT_MODEL_INPUT_EXAMPLE,
+        display_name=constants.DISPLAY_NAME,
+    )