feature: reshape Parents into experiment analytics dataframe (#1884)

aws/sagemaker-experiments#99

Author: danabens

src/sagemaker/analytics.py

@@ -547,6 +547,23 @@ def _reshape_artifacts(self, artifacts, _artifact_names):
             out["{} - {}".format(name, "Value")] = value.get("Value")
         return out
 
+    def _reshape_parents(self, parents):
+        """Reshape trial component parents to a pandas column
+        Args:
+            parents: trial component parents (trials and experiments)
+        Returns:
+            dict: Key: artifacts name, Value: artifacts value
+        """
+        out = OrderedDict()
+        trials = []
+        experiments = []
+        for parent in parents:
+            trials.append(parent["TrialName"])
+            experiments.append(parent["ExperimentName"])
+        out["Trials"] = trials
+        out["Experiments"] = experiments
+        return out
+
     def _reshape(self, trial_component):
         """Reshape trial component data to pandas columns
         Args:
@@ -574,6 +591,7 @@ def _reshape(self, trial_component):
                 trial_component.get("OutputArtifacts", []), self._output_artifact_names
             )
         )
+        out.update(self._reshape_parents(trial_component.get("Parents", [])))
         return out
 
     def _fetch_dataframe(self):

tests/integ/test_experiments_analytics.py

@@ -99,6 +99,8 @@ def test_experiment_analytics_artifacts(sagemaker_session):
             "inputArtifacts1 - Value",
             "outputArtifacts1 - MediaType",
             "outputArtifacts1 - Value",
+            "Trials",
+            "Experiments",
         ]
 
 
@@ -109,7 +111,13 @@ def test_experiment_analytics(sagemaker_session):
             experiment_name=experiment_name, sagemaker_session=sagemaker_session
         )
 
-        assert list(analytics.dataframe().columns) == ["TrialComponentName", "DisplayName", "hp1"]
+        assert list(analytics.dataframe().columns) == [
+            "TrialComponentName",
+            "DisplayName",
+            "hp1",
+            "Trials",
+            "Experiments",
+        ]
 
 
 def test_experiment_analytics_pagination(sagemaker_session):
@@ -118,7 +126,13 @@ def test_experiment_analytics_pagination(sagemaker_session):
             experiment_name=experiment_name, sagemaker_session=sagemaker_session
         )
 
-        assert list(analytics.dataframe().columns) == ["TrialComponentName", "DisplayName", "hp1"]
+        assert list(analytics.dataframe().columns) == [
+            "TrialComponentName",
+            "DisplayName",
+            "hp1",
+            "Trials",
+            "Experiments",
+        ]
         assert (
             len(analytics.dataframe()) > 10
         )  # TODO [owen-t] Replace with == 20 and put test in retry block
@@ -137,7 +151,13 @@ def test_experiment_analytics_search_by_nested_filter(sagemaker_session):
             sagemaker_session=sagemaker_session, search_expression=search_exp
         )
 
-        assert list(analytics.dataframe().columns) == ["TrialComponentName", "DisplayName", "hp1"]
+        assert list(analytics.dataframe().columns) == [
+            "TrialComponentName",
+            "DisplayName",
+            "hp1",
+            "Trials",
+            "Experiments",
+        ]
         assert (
             len(analytics.dataframe()) > 5
         )  # TODO [owen-t] Replace with == 10 and put test in retry block
@@ -159,7 +179,13 @@ def test_experiment_analytics_search_by_nested_filter_sort_ascending(sagemaker_s
             sort_order="Ascending",
         )
 
-        assert list(analytics.dataframe().columns) == ["TrialComponentName", "DisplayName", "hp1"]
+        assert list(analytics.dataframe().columns) == [
+            "TrialComponentName",
+            "DisplayName",
+            "hp1",
+            "Trials",
+            "Experiments",
+        ]
         assert (
             len(analytics.dataframe()) > 5
         )  # TODO [owen-t] Replace with == 10 and put test in retry block
@@ -183,7 +209,13 @@ def test_experiment_analytics_search_by_nested_filter_sort_descending(sagemaker_
             sort_by="Parameters.hp1",
         )
 
-        assert list(analytics.dataframe().columns) == ["TrialComponentName", "DisplayName", "hp1"]
+        assert list(analytics.dataframe().columns) == [
+            "TrialComponentName",
+            "DisplayName",
+            "hp1",
+            "Trials",
+            "Experiments",
+        ]
         assert (
             len(analytics.dataframe()) > 5
         )  # TODO [owen-t] Replace with == 10 and put test in retry block

tests/unit/test_experiments_analytics.py

@@ -48,10 +48,11 @@ def trial_component(trial_component_name):
             "outputArtifacts1": {"MediaType": "text/csv", "Value": "s3:/sky/far1"},
             "outputArtifacts2": {"MediaType": "text/csv", "Value": "s3:/sky/far2"},
         },
+        "Parents": [{"TrialName": "trial1", "ExperimentName": "experiment1"}],
     }
 
 
-def test_trial_analytics_dataframe_all_metrics_hyperparams(mock_session):
+def test_trial_analytics_dataframe_all(mock_session):
     mock_session.sagemaker_client.search.return_value = {
         "Results": [
             {"TrialComponent": trial_component("trial-1")},
@@ -88,6 +89,8 @@ def test_trial_analytics_dataframe_all_metrics_hyperparams(mock_session):
                 ("outputArtifacts1 - Value", ["s3:/sky/far1", "s3:/sky/far1"]),
                 ("outputArtifacts2 - MediaType", ["text/csv", "text/csv"]),
                 ("outputArtifacts2 - Value", ["s3:/sky/far2", "s3:/sky/far2"]),
+                ("Trials", [["trial1"], ["trial1"]]),
+                ("Experiments", [["experiment1"], ["experiment1"]]),
             ]
         )
     )
@@ -141,6 +144,8 @@ def test_trial_analytics_dataframe_selected_hyperparams(mock_session):
                 ("outputArtifacts1 - Value", ["s3:/sky/far1", "s3:/sky/far1"]),
                 ("outputArtifacts2 - MediaType", ["text/csv", "text/csv"]),
                 ("outputArtifacts2 - Value", ["s3:/sky/far2", "s3:/sky/far2"]),
+                ("Trials", [["trial1"], ["trial1"]]),
+                ("Experiments", [["experiment1"], ["experiment1"]]),
             ]
         )
     )
@@ -189,6 +194,8 @@ def test_trial_analytics_dataframe_selected_metrics(mock_session):
                 ("outputArtifacts1 - Value", ["s3:/sky/far1", "s3:/sky/far1"]),
                 ("outputArtifacts2 - MediaType", ["text/csv", "text/csv"]),
                 ("outputArtifacts2 - Value", ["s3:/sky/far2", "s3:/sky/far2"]),
+                ("Trials", [["trial1"], ["trial1"]]),
+                ("Experiments", [["experiment1"], ["experiment1"]]),
             ]
         )
     )
@@ -243,6 +250,8 @@ def test_trial_analytics_dataframe_search_pagination(mock_session):
                 ("outputArtifacts1 - Value", ["s3:/sky/far1", "s3:/sky/far1"]),
                 ("outputArtifacts2 - MediaType", ["text/csv", "text/csv"]),
                 ("outputArtifacts2 - Value", ["s3:/sky/far2", "s3:/sky/far2"]),
+                ("Trials", [["trial1"], ["trial1"]]),
+                ("Experiments", [["experiment1"], ["experiment1"]]),
             ]
         )
     )