feat: add eval metrics types to get_experiment_df (#1648)

Co-authored-by: sina chavoshi <sina.chavoshi@gmail.com>
googleapis · Sep 9, 2022 · 944b03f · 944b03f
1 parent 83da8c5
commit 944b03f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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(