diff --git a/google/cloud/aiplatform/metadata/experiment_resources.py b/google/cloud/aiplatform/metadata/experiment_resources.py index f83f94a537..0d1561eac3 100644 --- a/google/cloud/aiplatform/metadata/experiment_resources.py +++ b/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 ) diff --git a/google/cloud/aiplatform/metadata/experiment_run_resource.py b/google/cloud/aiplatform/metadata/experiment_run_resource.py index 918f37c962..f4d2a85b0d 100644 --- a/google/cloud/aiplatform/metadata/experiment_run_resource.py +++ b/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 {} diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index f1d4997dc3..2b672d2663 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/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. diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index fb865721e8..775a4430a1 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/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",