docs(samples): add samples for autologging

PiperOrigin-RevId: 511202564

samples/model-builder/conftest.py

@@ -1091,3 +1091,17 @@ def mock_remove_version_aliases(mock_model_registry):
     ) as mock_remove_version_aliases:
         mock_remove_version_aliases.return_value = None
         yield mock_remove_version_aliases
+
+
+"""
+----------------------------------------------------------------------------
+Autologging Fixtures
+----------------------------------------------------------------------------
+"""
+
+
+@pytest.fixture
+def mock_autolog():
+    with patch.object(aiplatform, "autolog") as mock_autolog_method:
+        mock_autolog_method.return_value = None
+        yield mock_autolog_method

samples/model-builder/experiment_tracking/autologging_with_auto_run_creation_sample.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.
+
+from typing import Union
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_autologging_with_auto_run_creation_sample]
+def autologging_with_auto_run_creation_sample(
+    experiment_name: str,
+    experiment_tensorboard: Union[str, aiplatform.Tensorboard],
+    project: str,
+    location: str,
+):
+    aiplatform.init(
+        experiment_name=experiment_name,
+        project=project,
+        location=location,
+        experiment_tensorboard=experiment_tensorboard,
+    )
+
+    aiplatform.autolog()
+
+    # Your model training code goes here
+
+    aiplatform.autolog(disable=True)
+
+
+#  [END aiplatform_sdk_autologging_with_auto_run_creation_sample]

samples/model-builder/experiment_tracking/autologging_with_auto_run_creation_sample_test.py

@@ -0,0 +1,36 @@
+# 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 experiment_tracking import autologging_with_auto_run_creation_sample
+
+import test_constants as constants
+
+
+def test_autologging_with_auto_run_creation_sample(mock_sdk_init, mock_autolog):
+
+    autologging_with_auto_run_creation_sample.autologging_with_auto_run_creation_sample(
+        experiment_name=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        experiment_tensorboard=constants.TENSORBOARD_NAME,
+    )
+
+    mock_sdk_init.assert_called_with(
+        experiment_name=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        experiment_tensorboard=constants.TENSORBOARD_NAME,
+    )
+
+    assert mock_autolog.call_count == 2

samples/model-builder/experiment_tracking/autologging_with_manual_run_creation_sample.py

@@ -0,0 +1,46 @@
+# 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 Union
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_autologging_with_manual_run_creation_sample]
+def autologging_with_manual_run_creation_sample(
+    experiment_name: str,
+    run_name: str,
+    experiment_tensorboard: Union[str, aiplatform.Tensorboard],
+    project: str,
+    location: str,
+):
+    aiplatform.init(
+        experiment_name=experiment_name,
+        project=project,
+        location=location,
+        experiment_tensorboard=experiment_tensorboard,
+    )
+
+    aiplatform.autolog()
+
+    aiplatform.start_run(run=run_name)
+
+    # Your model training code goes here
+
+    aiplatform.end_run()
+
+    aiplatform.autolog(disable=True)
+
+
+#  [END aiplatform_sdk_autologging_with_manual_run_creation_sample]

samples/model-builder/experiment_tracking/autologging_with_manual_run_creation_sample_test.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 experiment_tracking import autologging_with_manual_run_creation_sample
+
+import test_constants as constants
+
+
+def test_autologging_with_manual_run_creation_sample(
+    mock_sdk_init, mock_start_run, mock_end_run, mock_autolog
+):
+
+    autologging_with_manual_run_creation_sample.autologging_with_manual_run_creation_sample(
+        experiment_name=constants.EXPERIMENT_NAME,
+        run_name=constants.EXPERIMENT_RUN_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        experiment_tensorboard=constants.TENSORBOARD_NAME,
+    )
+
+    mock_sdk_init.assert_called_with(
+        experiment_name=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        experiment_tensorboard=constants.TENSORBOARD_NAME,
+    )
+
+    mock_start_run.assert_called_with(
+        run=constants.EXPERIMENT_RUN_NAME,
+    )
+
+    mock_end_run.assert_called_with()
+
+    assert mock_autolog.call_count == 2