feature: reshape Parents into experiment analytics dataframe

aws/sagemaker-experiments#99

Author: danabens

src/sagemaker/analytics.py

@@ -536,6 +536,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:
@@ -563,6 +580,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/unit/test_experiments_analytics.py

@@ -48,10 +48,16 @@ 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 +94,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 +149,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 +199,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 +255,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"]]),
             ]
         )
     )