add trace tagging support

Author: ptooley

pypop/cli.py

@@ -133,6 +133,10 @@ def _preprocess_traces_parse_args():
     parser.add_argument(
         "--dimemas-path", type=str, metavar="PATH", help="Path to Dimemas executable"
     )
+
+    parser.add_argument(
+        "--tag", type=str, metavar="TAG", help="Tag to apply to trace(s)"
+    )
     parser.add_argument(
         "--outfile-path",
         type=str,
@@ -274,6 +278,7 @@ def preprocess_traces():
         force_recalculation=config.force_recalculation,
         chop_to_roi=config.chop_to_roi,
         outpath=config.outfile_path,
+        tag=config.tag,
     )
 
 

pypop/metrics/metricset.py

@@ -59,6 +59,7 @@ class MetricSet:
         "Threads per Process": "",
         "Total Threads": "",
         "Hybrid Layout": "",
+        "Tag": "",
     }
 
     def __init__(self, stats_data, ref_key=None, sort_keys=True):
@@ -102,9 +103,10 @@ def _choose_ref_key(stats_dict):
 
         return min(
             stats_dict.items(),
-            key=lambda x: "{:05}_{:05}".format(
+            key=lambda x: "{:05}_{:05}_{}".format(
                 sum(x[1].metadata.threads_per_process),
                 max(x[1].metadata.threads_per_process),
+                x[1].metadata.tag,
             ),
         )[0]
 
@@ -165,6 +167,7 @@ def _create_subdataframe(self, metadata, idxkey):
                 ],
                 index=[idxkey],
             ),
+            "Tag": pandas.Series(data=[metadata.tag], index=[idxkey]),
         }
 
         for metric in self._metric_list:
@@ -182,6 +185,7 @@ def plot_table(
         **kwargs
     ):
         from pypop.notebook_interface.plotting import MetricTable
+
         return MetricTable(
             self,
             columns_key=columns_key,
@@ -201,6 +205,7 @@ def plot_scaling(
         **kwargs
     ):
         from pypop.notebook_interface.plotting import ScalingPlot
+
         return ScalingPlot(
             self,
             scaling_variable=scaling_variable,

pypop/trace/trace.py

@@ -42,17 +42,19 @@ def __new__(cls, filename, **kwargs):
         return super().__new__(cls)
 
     @staticmethod
-    def load(filename, force_recalculation=False, **kwargs):
+    def load(filename, force_recalculation=False, tag=None, **kwargs):
 
         if force_recalculation is False:
             # First try to open file as if it *is* a summary file
             try:
-                return Trace._load_from_summary_file(filename, **kwargs)
+                summaryloaded = Trace._load_from_summary_file(
+                    filename, tag=tag, **kwargs
+                )
             except WrongLoaderError:
                 # Otherwise, see if it has a summaryfile present
                 try:
-                    return Trace._load_from_summary_file(
-                        Trace.get_summary_filename(filename), **kwargs
+                    summaryloaded = Trace._load_from_summary_file(
+                        Trace.get_summary_filename(filename), tag=tag, **kwargs
                     )
                 # Fallback - assume valid trace with no summary
                 except (WrongLoaderError, FileNotFoundError):
@@ -62,21 +64,21 @@ def load(filename, force_recalculation=False, **kwargs):
         for loader in Trace._loader_registry.values():
             try:
                 return loader(
-                    filename, force_recalculation=force_recalculation, **kwargs
+                    filename, force_recalculation=force_recalculation, tag=tag, **kwargs
                 )
             except WrongLoaderError:
                 pass
 
         raise ValueError("Invalid or unsupported trace {}".format(filename))
 
     @staticmethod
-    def _load_from_summary_file(filename, **kwargs):
+    def _load_from_summary_file(filename, tag=None, **kwargs):
 
         # This will raise a WrongLoader error already so no need to trap and reraise
         metadata, _ = Trace._parse_summary_file(filename)
 
         try:
-            return Trace._loader_registry[metadata.trace_subclass_name](
+            trace = Trace._loader_registry[metadata.trace_subclass_name](
                 filename, **kwargs
             )
         except KeyError:
@@ -86,7 +88,12 @@ def _load_from_summary_file(filename, **kwargs):
                 "".format(metadata.trace_subclass_name)
             )
 
-    def __init__(self, filename, force_recalculation=False, **kwargs):
+        if tag is not None:
+            warn("Loading from summary, tag argument ignored")
+
+        return trace
+
+    def __init__(self, filename, force_recalculation=False, tag=None, **kwargs):
         # Save keyword args for later
         self._kwargs = kwargs
 
@@ -115,6 +122,8 @@ def __init__(self, filename, force_recalculation=False, **kwargs):
                     pass
 
         self._metadata = TraceMetadata(self)
+        if tag:
+            self._metadata.tag = tag
         self._statistics = None
 
         # Loading from a tracefile so

pypop/trace/tracemetadata.py

@@ -35,8 +35,11 @@ class TraceMetadata:
         # "layout": str,
         "tracefile_name": str,
         "fingerprint": str,
+        "tag": str,
     }
 
+    _optional = {"tag"}
+
     _needs_packing = ["layout", "threads_per_process"]
 
     def __init__(self, parent):
@@ -58,7 +61,9 @@ def unpack_dataframe(dataframe):
             try:
                 packed_data = dataframe[var]
             except KeyError as err:
-                warn("Missing data value: {}".format(err))
+                if var not in metadata._optional:
+                    warn("Missing metadata entry: {}".format(err))
+                setattr(metadata, var, None)
                 continue
             if var in metadata._needs_packing:
                 unpack = getattr(metadata, "unpack_{}".format(var))
@@ -73,6 +78,8 @@ def pack_dataframe(self):
         data = {}
 
         for var, vartype in self._datavars.items():
+            if getattr(self, var) is None:
+                continue
             if var in self._needs_packing:
                 pack = getattr(self, "pack_{}".format(var))
                 packed = pack(getattr(self, var), vartype)

pypop/traceset.py

@@ -66,6 +66,7 @@ def __init__(
         no_progress=False,
         outpath=None,
         ignore_cache=None,
+        tag=None,
     ):
         # Setup data structures
         self.traces = set()
@@ -79,7 +80,7 @@ def __init__(
 
         # Add traces
         self.add_traces(
-            path_list, force_recalculation, chop_to_roi, no_progress, outpath
+            path_list, force_recalculation, chop_to_roi, no_progress, outpath, tag
         )
 
     def add_traces(
@@ -89,6 +90,7 @@ def add_traces(
         chop_to_roi=False,
         no_progress=False,
         outpath=None,
+        tag=None,
     ):
         """Collect statistics from provided trace files, currently only Extrae .prv files
         are supported.
@@ -137,6 +139,7 @@ def add_traces(
                         force_recalculation=force_recalculation,
                         chop_to_roi=chop_to_roi,
                         outpath=outpath,
+                        tag=tag,
                     )
                 )