chore: support full resource name in environment experiment run

jaycee-li · copybara-github · commit f041f1af0bb0 · 2024-05-16T00:54:47.000-07:00
PiperOrigin-RevId: 634250315
diff --git a/google/cloud/aiplatform/metadata/metadata.py b/google/cloud/aiplatform/metadata/metadata.py
@@ -276,12 +276,21 @@ def experiment_run(self) -> Optional[experiment_run_resource.ExperimentRun]:
         """Returns the currently set experiment run or experiment run set via env variable AIP_EXPERIMENT_RUN_NAME."""
         if self._experiment_run:
             return self._experiment_run
-        if os.getenv(constants.ENV_EXPERIMENT_RUN_KEY):
-            self._experiment_run = experiment_run_resource.ExperimentRun.get(
-                os.getenv(constants.ENV_EXPERIMENT_RUN_KEY),
+
+        env_experiment_run = os.getenv(constants.ENV_EXPERIMENT_RUN_KEY)
+        if env_experiment_run and self.experiment:
+            # The run could be run name or full resource name,
+            # so we remove the experiment resource prefix if necessary.
+            env_experiment_run = env_experiment_run.replace(
+                f"{self.experiment.resource_name}-",
+                "",
+            )
+            self._experiment_run = experiment_run_resource.ExperimentRun(
+                env_experiment_run,
                 experiment=self.experiment,
             )
             return self._experiment_run
+
         return None
 
     def set_experiment(
diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py
@@ -1487,7 +1487,7 @@ def test_start_run_from_env_experiment(
         "get_tensorboard_run_artifact_not_found_mock",
         "get_or_create_default_tb_none_mock",
     )
-    def test_init_experiment_run_from_env(self):
+    def test_init_experiment_run_from_env_run_name(self):
         os.environ["AIP_EXPERIMENT_RUN_NAME"] = _TEST_RUN
 
         aiplatform.init(
@@ -1501,6 +1501,26 @@ def test_init_experiment_run_from_env(self):
 
         del os.environ["AIP_EXPERIMENT_RUN_NAME"]
 
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_mock",
+        "get_tensorboard_run_artifact_not_found_mock",
+        "get_or_create_default_tb_none_mock",
+    )
+    def test_init_experiment_run_from_env_run_resource_name(self):
+        os.environ["AIP_EXPERIMENT_RUN_NAME"] = _TEST_EXPERIMENT_RUN_CONTEXT_NAME
+
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+        )
+
+        run = metadata._experiment_tracker.experiment_run
+        assert run.name == _TEST_RUN
+
+        del os.environ["AIP_EXPERIMENT_RUN_NAME"]
+
     def test_get_experiment(self, get_experiment_mock):
         aiplatform.init(
             project=_TEST_PROJECT,
