diff --git a/google/cloud/aiplatform/metadata/constants.py b/google/cloud/aiplatform/metadata/constants.py index 2dd0c4db86..01c4b25b1a 100644 --- a/google/cloud/aiplatform/metadata/constants.py +++ b/google/cloud/aiplatform/metadata/constants.py @@ -22,6 +22,9 @@ SYSTEM_PIPELINE = "system.Pipeline" SYSTEM_PIPELINE_RUN = "system.PipelineRun" SYSTEM_METRICS = "system.Metrics" +GOOGLE_CLASSIFICATION_METRICS = "google.ClassificationMetrics" +GOOGLE_REGRESSION_METRICS = "google.RegressionMetrics" +GOOGLE_FORECASTING_METRICS = "google.ForecastingMetrics" _EXPERIMENTS_V2_TENSORBOARD_RUN = "google.VertexTensorboardRun" diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index 2a3914efea..afc5dd86ad 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -639,10 +639,14 @@ def _query_experiment_row( credentials=node.credentials, filter=metadata_utils._make_filter_string( in_context=[node.resource_name], - schema_title=metadata_constants.SYSTEM_METRICS, + schema_title=[ + metadata_constants.SYSTEM_METRICS, + metadata_constants.GOOGLE_CLASSIFICATION_METRICS, + metadata_constants.GOOGLE_REGRESSION_METRICS, + metadata_constants.GOOGLE_FORECASTING_METRICS, + ], ), ) - row = experiment_resources._ExperimentRow( experiment_run_type=node.schema_title, name=node.display_name ) diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index 514bdb84bd..ffade4976d 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -1352,7 +1352,12 @@ def test_get_experiment_df( expected_filter = metadata_utils._make_filter_string( in_context=[_TEST_PIPELINE_CONTEXT.name], - schema_title=constants.SYSTEM_METRICS, + schema_title=[ + constants.SYSTEM_METRICS, + constants.GOOGLE_CLASSIFICATION_METRICS, + constants.GOOGLE_REGRESSION_METRICS, + constants.GOOGLE_FORECASTING_METRICS, + ], ) list_artifact_mock_for_experiment_dataframe.assert_has_calls(