feat: add ExperimentRun.get_logged_custom_jobs method

PiperOrigin-RevId: 523824096

Author: jaycee-li

google/cloud/aiplatform/constants/base.py

@@ -114,11 +114,9 @@
 
 # Used in CustomJob.from_local_script for experiments integration in training
 AIPLATFORM_DEPENDENCY_PATH = (
-    "google-cloud-aiplatform[metadata,tensorboard]"
-    + f"=={aiplatform_version.__version__}"
+    f"google-cloud-aiplatform=={aiplatform_version.__version__}"
 )
 
 AIPLATFORM_AUTOLOG_DEPENDENCY_PATH = (
-    "google-cloud-aiplatform[metadata,tensorboard,autologging]"
-    + f"=={aiplatform_version.__version__}"
+    f"google-cloud-aiplatform[autologging]=={aiplatform_version.__version__}"
 )

google/cloud/aiplatform/jobs.py

@@ -1488,10 +1488,12 @@ def from_local_script(
             ).pool_specs
         )
 
+        # if users enable autolog, automatically install SDK in their container image
+        # otherwise users need to manually install SDK
         if enable_autolog:
             experiment_requirements = [constants.AIPLATFORM_AUTOLOG_DEPENDENCY_PATH]
         else:
-            experiment_requirements = [constants.AIPLATFORM_DEPENDENCY_PATH]
+            experiment_requirements = []
 
         if requirements:
             requirements.extend(experiment_requirements)

google/cloud/aiplatform/metadata/experiment_run_resource.py

@@ -25,6 +25,7 @@
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import pipeline_jobs
+from google.cloud.aiplatform import jobs
 from google.cloud.aiplatform.compat.types import artifact as gca_artifact
 from google.cloud.aiplatform.compat.types import execution as gca_execution
 from google.cloud.aiplatform.compat.types import (
@@ -1304,6 +1305,24 @@ def get_logged_pipeline_jobs(self) -> List[pipeline_jobs.PipelineJob]:
             for c in pipeline_job_contexts
         ]
 
+    @_v1_not_supported
+    def get_logged_custom_jobs(self) -> List[jobs.CustomJob]:
+        """Get all CustomJobs associated to this experiment run.
+
+        Returns:
+            List of CustomJobs associated this run.
+        """
+
+        custom_jobs = self._metadata_node.metadata.get(constants._CUSTOM_JOB_KEY)
+
+        return [
+            jobs.CustomJob.get(
+                resource_name=custom_job.get(constants._CUSTOM_JOB_RESOURCE_NAME),
+                credentials=self.credentials,
+            )
+            for custom_job in custom_jobs
+        ]
+
     def __enter__(self):
         return self
 

tests/system/aiplatform/test_custom_job.py

@@ -40,15 +40,15 @@
 @mock.patch.object(
     constants,
     "AIPLATFORM_DEPENDENCY_PATH",
-    "google-cloud-aiplatform[metadata,tensorboard] @ git+https://github.com/googleapis/"
+    "google-cloud-aiplatform @ git+https://github.com/googleapis/"
     f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}#egg=google-cloud-aiplatform"
     if os.environ.get("KOKORO_GIT_COMMIT")
     else constants.AIPLATFORM_DEPENDENCY_PATH,
 )
 @mock.patch.object(
     constants,
     "AIPLATFORM_AUTOLOG_DEPENDENCY_PATH",
-    "google-cloud-aiplatform[metadata,tensorboard,autologging] @ git+https://github.com/googleapis/"
+    "google-cloud-aiplatform[autologging] @ git+https://github.com/googleapis/"
     f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}#egg=google-cloud-aiplatform"
     if os.environ.get("KOKORO_GIT_COMMIT")
     else constants.AIPLATFORM_AUTOLOG_DEPENDENCY_PATH,
@@ -88,7 +88,7 @@ def test_from_local_script_prebuilt_container(self, shared_state):
             display_name=display_name,
             script_path=_LOCAL_TRAINING_SCRIPT_PATH,
             container_uri=_PREBUILT_CONTAINER_IMAGE,
-            requirements=["scikit-learn"],
+            requirements=["scikit-learn", "pandas"],
         )
         custom_job.run()
 
@@ -110,7 +110,7 @@ def test_from_local_script_custom_container(self, shared_state):
             display_name=display_name,
             script_path=_LOCAL_TRAINING_SCRIPT_PATH,
             container_uri=_CUSTOM_CONTAINER_IMAGE,
-            requirements=["scikit-learn"],
+            requirements=["scikit-learn", "pandas"],
         )
         custom_job.run()
 
@@ -139,7 +139,7 @@ def test_from_local_script_enable_autolog_prebuilt_container(self, shared_state)
             display_name=display_name,
             script_path=_LOCAL_TRAINING_SCRIPT_PATH,
             container_uri=_PREBUILT_CONTAINER_IMAGE,
-            requirements=["scikit-learn"],
+            requirements=["scikit-learn", "pandas"],
             enable_autolog=True,
         )
 
@@ -169,7 +169,7 @@ def test_from_local_script_enable_autolog_custom_container(self, shared_state):
             display_name=display_name,
             script_path=_LOCAL_TRAINING_SCRIPT_PATH,
             container_uri=_CUSTOM_CONTAINER_IMAGE,
-            requirements=["scikit-learn"],
+            requirements=["scikit-learn", "pandas"],
             enable_autolog=True,
         )
 

tests/unit/aiplatform/test_metadata.py

@@ -42,6 +42,7 @@
     Artifact as GapicArtifact,
     Context as GapicContext,
     Execution as GapicExecution,
+    JobServiceClient,
     MetadataServiceClient,
     AddExecutionEventsResponse,
     MetadataStore as GapicMetadataStore,
@@ -686,6 +687,21 @@ def test_get_pipeline_df_wrong_schema(self):
 _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT = copy.deepcopy(_EXPERIMENT_RUN_MOCK)
 _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT.parent_contexts = [_TEST_CONTEXT_NAME]
 
+_TEST_CUSTOM_JOB_NAME = (
+    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/customJobs/12345"
+)
+_TEST_CUSTOM_JOB_CONSOLE_URI = "test-custom-job-console-uri"
+
+_EXPERIMENT_RUN_MOCK_WITH_CUSTOM_JOBS = copy.deepcopy(
+    _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT
+)
+_EXPERIMENT_RUN_MOCK_WITH_CUSTOM_JOBS.metadata[constants._CUSTOM_JOB_KEY] = [
+    {
+        constants._CUSTOM_JOB_RESOURCE_NAME: _TEST_CUSTOM_JOB_NAME,
+        constants._CUSTOM_JOB_CONSOLE_URI: _TEST_CUSTOM_JOB_CONSOLE_URI,
+    },
+]
+
 
 @pytest.fixture
 def get_experiment_mock():
@@ -724,6 +740,17 @@ def get_experiment_run_mock():
         yield get_context_mock
 
 
+@pytest.fixture
+def get_experiment_run_with_custom_jobs_mock():
+    with patch.object(MetadataServiceClient, "get_context") as get_context_mock:
+        get_context_mock.side_effect = [
+            _EXPERIMENT_MOCK,
+            _EXPERIMENT_RUN_MOCK_WITH_CUSTOM_JOBS,
+        ]
+
+        yield get_context_mock
+
+
 @pytest.fixture
 def get_experiment_run_not_found_mock():
     with patch.object(MetadataServiceClient, "get_context") as get_context_mock:
@@ -831,6 +858,12 @@ def add_context_children_mock():
         yield add_context_children_mock
 
 
+@pytest.fixture
+def get_custom_job_mock():
+    with patch.object(JobServiceClient, "get_custom_job") as get_custom_job_mock:
+        yield get_custom_job_mock
+
+
 _EXPERIMENT_RUN_MOCK_POPULATED_1 = copy.deepcopy(
     _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT
 )
@@ -1869,3 +1902,19 @@ def test_get_experiment_df_wrong_schema(self):
         aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
         with pytest.raises(ValueError):
             aiplatform.get_experiment_df(_TEST_EXPERIMENT)
+
+    @pytest.mark.usefixtures(
+        "get_experiment_run_with_custom_jobs_mock",
+        "get_metadata_store_mock",
+        "get_tensorboard_run_artifact_not_found_mock",
+    )
+    def test_experiment_run_get_logged_custom_jobs(self, get_custom_job_mock):
+        aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+        run = aiplatform.ExperimentRun(_TEST_RUN, experiment=_TEST_EXPERIMENT)
+        jobs = run.get_logged_custom_jobs()
+
+        assert len(jobs) == 1
+        get_custom_job_mock.assert_called_once_with(
+            name=_TEST_CUSTOM_JOB_NAME,
+            retry=base._DEFAULT_RETRY,
+        )