fix: Update get_experiment_df to pass Experiment and allow empty metrics.

PiperOrigin-RevId: 634867806

Author: vertex-sdk-bot

google/cloud/aiplatform/metadata/experiment_resources.py

@@ -482,7 +482,8 @@ def get_data_frame(
                             metadata_context.schema_title
                         ]._query_experiment_row,
                         metadata_context,
-                        include_time_series,
+                        experiment=self,
+                        include_time_series=include_time_series,
                     )
                     for metadata_context in contexts
                 ]
@@ -494,7 +495,8 @@ def get_data_frame(
                             metadata_execution.schema_title
                         ]._query_experiment_row,
                         metadata_execution,
-                        include_time_series,
+                        experiment=self,
+                        include_time_series=include_time_series,
                     )
                     for metadata_execution in executions
                 )

google/cloud/aiplatform/metadata/experiment_run_resource.py

@@ -454,6 +454,8 @@ def _initialize_experiment_run(
             self._metadata_metric_artifact = self._v1_get_metric_artifact()
         if not self._is_legacy_experiment_run() and lookup_tensorboard_run:
             self._backing_tensorboard_run = self._lookup_tensorboard_run_artifact()
+            if not self._backing_tensorboard_run:
+                self._assign_to_experiment_backing_tensorboard()
 
     @classmethod
     def list(
@@ -553,13 +555,16 @@ def _create_v1_experiment_run(
     def _query_experiment_row(
         cls,
         node: Union[context.Context, execution.Execution],
-        include_time_series: Optional[bool] = True,
+        experiment: Optional[experiment_resources.Experiment] = None,
+        include_time_series: bool = True,
     ) -> experiment_resources._ExperimentRow:
         """Retrieves the runs metric and parameters into an experiment run row.
 
         Args:
             node (Union[context._Context, execution.Execution]):
                 Required. Metadata node instance that represents this run.
+            experiment:
+                Optional. Experiment associated with this run.
             include_time_series (bool):
                 Optional. Whether or not to include time series metrics in df.
                 Default is True.
@@ -568,7 +573,7 @@ def _query_experiment_row(
         """
         this_experiment_run = cls.__new__(cls)
         this_experiment_run._initialize_experiment_run(
-            node, lookup_tensorboard_run=include_time_series
+            node, experiment=experiment, lookup_tensorboard_run=include_time_series
         )
 
         row = experiment_resources._ExperimentRow(
@@ -620,8 +625,11 @@ def _get_latest_time_series_metric_columns(self) -> Dict[str, Union[float, int]]
             return {
                 display_name: data.values[-1].scalar.value
                 for display_name, data in time_series_metrics.items()
-                if data.value_type
-                == gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR
+                if (
+                    data.values
+                    and data.value_type
+                    == gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR
+                )
             }
         return {}
 

google/cloud/aiplatform/pipeline_jobs.py

@@ -887,17 +887,24 @@ def _get_context(self) -> context.Context:
 
     @classmethod
     def _query_experiment_row(
-        cls, node: context.Context, include_time_series: Optional[bool] = True
+        cls,
+        node: context.Context,
+        experiment: Optional[experiment_resources.Experiment] = None,
+        include_time_series: bool = True,
     ) -> experiment_resources._ExperimentRow:
         """Queries the PipelineJob metadata as an experiment run parameter and metric row.
 
-        Parameters are retrieved from the system.Run Execution.metadata of the PipelineJob.
+        Parameters are retrieved from the system.Run Execution.metadata of the
+        PipelineJob.
 
-        Metrics are retrieved from the system.Metric Artifacts.metadata produced by this PipelineJob.
+        Metrics are retrieved from the system.Metric Artifacts.metadata produced
+        by this PipelineJob.
 
         Args:
             node (context._Context):
                 Required. System.PipelineRun context that represents a PipelineJob Run.
+            experiment:
+                Optional. Experiment associated with this run.
             include_time_series (bool):
                 Optional. Whether or not to include time series metrics in df.
                 Default is True.

tests/unit/aiplatform/test_metadata.py

@@ -18,7 +18,7 @@
 import os
 import copy
 from importlib import reload
-from unittest import mock
+from unittest import TestCase, mock
 from unittest.mock import patch, call
 
 import numpy as np
@@ -2048,6 +2048,32 @@ def test_log_pipeline_job(
             ],
         )
 
+    @pytest.mark.usefixtures(
+        "get_experiment_mock",
+    )
+    def test_get_experiment_df_passes_experiment_variable(
+        self,
+        list_context_mock_for_experiment_dataframe_mock,
+        list_artifact_mock_for_experiment_dataframe,
+        list_executions_mock_for_experiment_dataframe,
+        get_tensorboard_run_artifact_mock,
+        get_tensorboard_run_mock,
+    ):
+        aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+        with patch.object(
+            experiment_run_resource.ExperimentRun, "_query_experiment_row"
+        ) as query_experiment_row_mock:
+            row = experiment_resources._ExperimentRow(
+                experiment_run_type=constants.SYSTEM_EXPERIMENT_RUN,
+                name=_TEST_EXPERIMENT,
+            )
+            query_experiment_row_mock.return_value = row
+
+            aiplatform.get_experiment_df(_TEST_EXPERIMENT)
+            _, kwargs = query_experiment_row_mock.call_args_list[0]
+            TestCase.assertTrue(self, kwargs["experiment"].name == _TEST_EXPERIMENT)
+
     @pytest.mark.usefixtures(
         "get_experiment_mock",
         "list_tensorboard_time_series_mock",