docs(samples): add sample for experiment run state update.

PiperOrigin-RevId: 521002287

samples/model-builder/conftest.py

@@ -1034,6 +1034,13 @@ def mock_register_model(mock_experiment_model, mock_model):
         yield mock_register_model
 
 
+@pytest.fixture
+def mock_update_run_state(mock_experiment_run):
+    with patch.object(mock_experiment_run, "update_state") as mock_update_run_state:
+        mock_update_run_state.return_value = None
+        yield mock_update_run_state
+
+
 """
 ----------------------------------------------------------------------------
 Model Versioning Fixtures

samples/model-builder/experiment_tracking/update_experiment_run_state_sample.py

@@ -0,0 +1,38 @@
+# 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_update_experiment_run_state_sample]
+def update_experiment_run_state_sample(
+    run_name: str,
+    experiment: Union[str, aiplatform.Experiment],
+    project: str,
+    location: str,
+    state: aiplatform.gapic.Execution.State,
+) -> None:
+    experiment_run = aiplatform.ExperimentRun(
+        run_name=run_name,
+        experiment=experiment,
+        project=project,
+        location=location,
+    )
+
+    experiment_run.update_state(state)
+
+
+#  [END aiplatform_sdk_update_experiment_run_state_sample]

samples/model-builder/experiment_tracking/update_experiment_run_state_sample_test.py

@@ -0,0 +1,29 @@
+# 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
+
+from experiment_tracking import update_experiment_run_state_sample
+import test_constants as constants
+
+
+@pytest.mark.usefixtures("mock_get_run")
+def test_update_experiment_run_state_sample(mock_update_run_state):
+
+    update_experiment_run_state_sample.update_experiment_run_state_sample(
+        run_name=constants.EXPERIMENT_RUN_NAME,
+        experiment=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        state=constants.EXPERIMENT_RUN_STATE,
+    )
+
+    mock_update_run_state.assert_called_once_with(constants.EXPERIMENT_RUN_STATE)

samples/model-builder/test_constants.py

@@ -279,6 +279,7 @@
 METADATA = {}
 
 EXPERIMENT_RUN_NAME = "my-run"
+EXPERIMENT_RUN_STATE = aiplatform.gapic.Execution.State.RUNNING
 
 METRICS = {"accuracy": 0.1}
 PARAMS = {"learning_rate": 0.1}