Merge pull request #1124 from linsword13/inmem-fom

Extend to support in memory FOMs

Author: douglasjacobsen

lib/ramble/ramble/language/shared_language.py

@@ -98,13 +98,14 @@ def _execute_figure_of_merit_context(obj):
 @shared_directive("figures_of_merit")
 def figure_of_merit(
     name,
-    fom_regex,
-    group_name,
+    fom_regex=None,
+    group_name=None,
     log_file="{log_file}",
     units="",
     contexts=None,
     fom_type: FomType = FomType.UNDEFINED,
     when=None,
+    fom_map_key=None,
     **kwargs,
 ):
     """Adds a figure of merit to track for this object
@@ -122,9 +123,16 @@ def figure_of_merit(
                                    should exist in.
       fom_type (ramble.util.foms.FomType): The type of figure of merit
       when (list | None): List of when conditions to apply to directive
+      fom_map_key: If supplied, this is treated as an in-memory (as opposed to file-based)
+                   figure of merit, and its value is extracted using this key
     """
 
     def _execute_figure_of_merit(obj):
+        if fom_map_key is None:
+            if fom_regex is None or group_name is None:
+                raise ramble.language.language_base.DirectiveError(
+                    "`fom_regex` and `group_name` are required for defining file-based FOM"
+                )
         when_list = ramble.language.language_helpers.build_when_list(
             when, obj, name, "figure_of_merit"
         )
@@ -146,6 +154,7 @@ def _execute_figure_of_merit(obj):
             "fom_type": fom_type,
             "when": when_list,
             "origin_type": obj.origin_type if hasattr(obj, "origin_type") else "",
+            "fom_map_key": fom_map_key,
         }
 
     return _execute_figure_of_merit

lib/ramble/ramble/pipeline.py

@@ -634,7 +634,7 @@ def print_archive_files(app_inst, pattern_title, patterns):
             logger.all_msg(f"Experiment: {exp}")
             logger.all_msg(f"    Experiment log file: {log_file}")
 
-            analysis_logs, _ = app_inst._analysis_dicts(self.workspace.success_list)
+            analysis_logs, _, _ = app_inst._analysis_dicts(self.workspace.success_list)
 
             logger.all_msg("    Auxiliary experiment logs:")
             for log in analysis_logs:

lib/ramble/ramble/test/application_language.py

@@ -10,6 +10,7 @@
 import deprecation
 import pytest
 
+from ramble import language
 from ramble.appkit import *  # noqa
 
 app_types = [
@@ -280,6 +281,25 @@ def test_figure_of_merit_directive(app_class):
             assert app_inst.figures_of_merit[_FS][_FS][fom_name][conf_name] == conf_val
 
 
+def test_figure_of_merit_directive_required_args():
+    app_inst = ExecutableApplication("/not/a/path")  # noqa: F405
+    with pytest.raises(
+        language.language_base.DirectiveError, match="required for defining file-based FOM"
+    ):
+        app_inst.figure_of_merit(
+            "test_fom",
+            units="s",
+        )
+    app_inst.figure_of_merit(
+        "test_inmem_fom",
+        units="s",
+        fom_map_key="test_fom_map_key",
+    )
+    foms = list(list(app_inst.figures_of_merit.values())[0].values())[0]
+    assert len(foms) == 1
+    assert foms["test_inmem_fom"]["fom_map_key"] == "test_fom_map_key"
+
+
 @pytest.mark.parametrize("app_class", app_types)
 def test_input_file_directive(app_class):
     test_defs = {}

lib/ramble/ramble/test/end_to_end/expanded_fom_dry_run.py

@@ -29,6 +29,7 @@ def test_expanded_foms_dry_run(
     mpi_command: 'mpirun -n {n_ranks} -ppn {processes_per_node}'
     batch_submit: 'batch_submit {execute_experiment}'
     n_threads: '1'
+    my_var: 'testvar'
   applications:
     expanded_foms:
       workloads:
@@ -78,5 +79,7 @@ def test_expanded_foms_dry_run(
         text_results_files = glob.glob(os.path.join(ws1.root, "results*.txt"))
         with open(text_results_files[0]) as f:
             data = f.read()
+            assert "Status = SUCCESS" in data
             for expected in expected_expansions:
                 assert f"test_fom {expected} = 567.8 {unit}" in data
+            assert "test_inmem_testvar = inmem_val" in data

var/ramble/repos/builtin.mock/applications/expanded_foms/application.py

@@ -37,4 +37,20 @@ class ExpandedFoms(ExecutableApplication):
         units="{unit}",
     )
 
+    figure_of_merit(
+        "test_inmem_{my_var}",
+        units="",
+        fom_map_key="test_inmem",
+    )
+
     success_criteria("Run", mode="string", match=r"Collect", file="{log_file}")
+
+    success_criteria(
+        "Inmem FOM matched",
+        mode="fom_comparison",
+        fom_name="test_inmem_{my_var}",
+        formula="{value} == 'inmem_val'",
+    )
+
+    def _prepare_analysis(self, workspace, app_inst):
+        self.add_inmem_fom_value("test_inmem", "inmem_val")

var/ramble/repos/builtin/base_classes/application-base/base_class.py

@@ -196,6 +196,8 @@ def __init__(self, file_path):
         self._input_lock = None
         self._software_lock = None
         self._experiment_graph = None
+        # A dict storing fom values, currently it only stores inmem FOMs
+        self._fom_map = {}
 
         # Ensure we always have the application name, and this is never empty
         self.license_names = self.license_names + [self.name]
@@ -1427,12 +1429,8 @@ def _define_commands(self, exec_graph, success_list=None):
                         )
                         logs.append(expanded_log)
 
-            analysis_logs, _ = self._analysis_dicts(success_list)
-
-            for log in analysis_logs:
-                logs.append(log)
-
-            logs = list(dict.fromkeys(logs))
+            analysis_logs, _, _ = self._analysis_dicts(success_list)
+            logs = list(set(logs) | analysis_logs.keys())
 
             for log in logs:
                 self._command_list.append('rm -f "%s"' % log)
@@ -2156,7 +2154,7 @@ def _archive_experiments(self, workspace, app_inst=None):
 
             # Copy all figure of merit files
             criteria_list = self.success_list
-            analysis_files, _ = self._analysis_dicts(criteria_list)
+            analysis_files, _, _ = self._analysis_dicts(criteria_list)
             for file in analysis_files.keys():
                 if os.path.exists(file):
                     shutil.copy(file, archive_experiment_dir)
@@ -2201,6 +2199,30 @@ def _prepare_analysis(self, workspace, app_inst=None):
         """
         pass
 
+    def _extract_inmem_foms(self, inmem_fom_defs, fom_values):
+        """Extract in-memory FOMs"""
+        for context, foms in inmem_fom_defs.items():
+            if context not in fom_values:
+                fom_values[context] = {}
+            foms = inmem_fom_defs[context]["foms"]
+            for fom in foms:
+                fom_conf = inmem_fom_defs[context]["foms"][fom]
+                # Currently inmem FOM does not have semantics for expanded vars,
+                # so use the already expanded name and unit
+                fom_name = fom_conf["fom_name_expanded"]
+                # TODO: this can be extended to support derived FOMs,
+                # since the `fom_values` contains resolved file-based FOMs
+                fom_map_key = fom_conf["fom_map_key"]
+                fom_value = self._fom_map.get(fom_map_key)
+                expanded_fom_value = self.expander.expand_var(fom_value)
+                fom_values[context][fom_name] = {
+                    "value": expanded_fom_value,
+                    "units": fom_conf["units_expanded"],
+                    "origin": fom_conf["origin"],
+                    "origin_type": fom_conf["origin_type"],
+                    "fom_type": fom_conf["fom_type"],
+                }
+
     register_phase(
         "analyze_experiments",
         pipeline="analyze",
@@ -2247,17 +2269,16 @@ def format_context(context_match, context_format):
             context_string = context_format.format(**context_val)
             return context_string
 
-        fom_values = {}
-
         criteria_list = self.success_list
         if not criteria_list:
             criteria_list = ramble.success_criteria.ScopedCriteriaList()
         criteria_list.reset()
 
-        files, definitions = self._analysis_dicts(criteria_list)
+        files, f_defs, inmem_defs = self._analysis_dicts(criteria_list)
 
         exp_lock = self.experiment_lock()
 
+        fom_values = {}
         # Iterate over files. We already know they exist
         with lk.ReadTransaction(exp_lock):
             for file, file_conf in files.items():
@@ -2292,9 +2313,7 @@ def format_context(context_match, context_format):
                         # Iterate over contexts and add matched contexts to active_contexts
                         for context, foms in file_conf["contexts"].items():
                             if not context == _NULL_CONTEXT:
-                                context_conf = definitions[context][
-                                    "definition"
-                                ]
+                                context_conf = f_defs[context]["definition"]
                                 context_match = context_conf["regex"].match(
                                     line
                                 )
@@ -2314,7 +2333,7 @@ def format_context(context_match, context_format):
                                         fom_values[context_name] = {}
 
                             for fom in foms:
-                                fom_conf = definitions[context]["foms"][fom]
+                                fom_conf = f_defs[context]["foms"][fom]
                                 fom_match = fom_conf["regex"].match(line)
 
                                 if fom_match:
@@ -2389,6 +2408,7 @@ def format_context(context_match, context_format):
                                                     "fom_type"
                                                 ],
                                             }
+        self._extract_inmem_foms(inmem_defs, fom_values)
 
         # Test all non-file based success criteria
         for criteria_obj in criteria_list.all_criteria():
@@ -2397,7 +2417,7 @@ def format_context(context_match, context_format):
                     criteria_obj.mark_found()
 
         # If an app has no FOMs defined, don't fail it for that
-        success = (not definitions) or False
+        success = (not f_defs and not inmem_defs) or False
         for fom in fom_values.values():
             for value in fom.values():
                 if (
@@ -2733,12 +2753,13 @@ def _analysis_dicts(self, criteria_list):
 
         Returns:
             files (dict): All files that need to be processed
-            contexts (dict): Any contexts that have been defined
-            foms (dict): All figures of merit that need to be extracted
+            file_fom_defs (dict): Definitions of all file-backed FOMs to be extracted
+            inmem_fom_defs (dict): Definitions of all in-memory FOMs to be extracted
         """
 
         files = {}
-        definitions = {}
+        file_fom_defs = {}
+        inmem_fom_defs = {}
 
         # Add the application defined criteria
         criteria_list.flush_scope("application_definition")
@@ -2861,32 +2882,38 @@ def _analysis_dicts(self, criteria_list):
                                 "definitions and 'when' conditions."
                             )
 
-                        # Copy context definition for contexts used by a FOM
-                        if context not in definitions:
-                            definitions[context] = {
-                                "definition": {},
-                                "foms": {},
-                            }
-                            if context != _NULL_CONTEXT:
-                                regex_str = self.expander.expand_var(
-                                    all_contexts[context]["regex"]
-                                )
-                                definitions[context]["definition"] = {
-                                    "regex": re.compile(regex_str),
-                                    "format": all_contexts[context][
-                                        "output_format"
-                                    ],
+                        def _preset_context_dict(dest_def_dict, context):
+                            # Copy context definition for contexts used by a FOM
+                            if context not in dest_def_dict:
+                                dest_def_dict[context] = {
+                                    "definition": {},
+                                    "foms": {},
                                 }
+                                if context != _NULL_CONTEXT:
+                                    regex_str = self.expander.expand_var(
+                                        all_contexts[context]["regex"]
+                                    )
+                                    dest_def_dict[context]["definition"] = {
+                                        "regex": re.compile(regex_str),
+                                        "format": all_contexts[context][
+                                            "output_format"
+                                        ],
+                                    }
 
                         for fom, source_def in source_foms.items():
-                            if fom in definitions[context]["foms"]:
+                            is_inmem = source_def["fom_map_key"] is not None
+                            dest_def_dict = (
+                                inmem_fom_defs if is_inmem else file_fom_defs
+                            )
+                            _preset_context_dict(dest_def_dict, context)
+                            if fom in dest_def_dict[context]["foms"]:
                                 logger.warn(
                                     f"FOM {fom} already defined in context {context} by "
-                                    f"{definitions[context]['foms'][fom]['origin']}. "
+                                    f"{dest_def_dict[context]['foms'][fom]['origin']}. "
                                     f"Overwriting with new definition from {source.name}"
                                 )
                             else:
-                                definitions[context]["foms"][fom] = {}
+                                dest_def_dict[context]["foms"][fom] = {}
 
                             def _expand_var(var):
                                 return self.expander.expand_var(
@@ -2905,12 +2932,21 @@ def _try_expand_var_or_none(var: str, expander):
                                 "origin": source.name,
                                 "origin_type": source.origin_type,
                                 "contexts": set(source_def["contexts"]),
-                                "group": _expand_var(source_def["group_name"]),
+                                "group": (
+                                    ""
+                                    if is_inmem
+                                    else _expand_var(source_def["group_name"])
+                                ),
                                 "units": _expand_var(source_def["units"]),
-                                "regex": re.compile(
-                                    _expand_var(source_def["regex"])
+                                "regex": (
+                                    ""
+                                    if is_inmem
+                                    else re.compile(
+                                        _expand_var(source_def["regex"])
+                                    )
                                 ),
                                 "fom_type": source_def["fom_type"].to_dict(),
+                                "fom_map_key": source_def["fom_map_key"],
                                 # If expansion works (i.e., it doesn't rely on the matched fom
                                 # groups), then cache it here to avoid repeated expansion later.
                                 "units_expanded": _try_expand_var_or_none(
@@ -2921,8 +2957,10 @@ def _try_expand_var_or_none(var: str, expander):
                                 ),
                             }
 
-                            definitions[context]["foms"][fom] = fom_def
+                            dest_def_dict[context]["foms"][fom] = fom_def
 
+                            if is_inmem:
+                                continue
                             log_path = _expand_var(source_def["log_file"])
                             # Ensure log path is absolute. If not, prepend the experiment run dir
                             if (
@@ -2944,7 +2982,11 @@ def _try_expand_var_or_none(var: str, expander):
                             logger.debug("Log = %s" % log_path)
                             logger.debug("Conf = %s" % fom_def)
 
-        return files, definitions
+        return files, file_fom_defs, inmem_fom_defs
+
+    def add_inmem_fom_value(self, fom_map_key, value):
+        """Add value to an in-memory FOM"""
+        self._fom_map[fom_map_key] = value
 
     def read_status(self):
         """Read status from an experiment's status file, if possible.