Merge pull request #942 from linsword13/pkg-prov-cache

Cache package list to reduce hitting the pkgman runner

Author: douglasjacobsen

lib/ramble/ramble/package_manager.py

@@ -220,7 +220,22 @@ def _add_software_to_results(self, workspace, app_inst=None):
             owns the results.
 
         """
-        pass
+        if app_inst.result is None:
+            return
+        prov_cache = workspace.pkg_prov_cache
+        env_name = self.app_inst.expander.expand_var_name(self.keywords.env_name)
+        if env_name in prov_cache[self.name]:
+            # No copy done as this shouldn't be modified once written
+            pkg_list = prov_cache[self.name][env_name]
+        else:
+            pkg_list = self.get_package_list(workspace)
+            prov_cache[self.name][env_name] = pkg_list
+        self.app_inst.result.software[self._spec_prefix] = pkg_list
+
+    def get_package_list(self, workspace):
+        """Method used by add_software_to_results phase to get software provenance info"""
+        del workspace
+        return []
 
 
 class PackageManagerError(RambleError):

lib/ramble/ramble/workspace/workspace.py

@@ -468,6 +468,10 @@ def __init__(self, root, dry_run=False, read_default_template=True):
 
         self.results = self.default_results()
 
+        # A cache structured as {pkg_man: {env_name: pkg_list}}.
+        # It's used to cache package provenance info from different package managers.
+        self.pkg_prov_cache = defaultdict(dict)
+
         self.success_list = ramble.success_criteria.ScopedCriteriaList()
 
         # Key for each application config should be it's filepath

var/ramble/repos/builtin/package_managers/environment-modules/package_manager.py

@@ -111,24 +111,18 @@ def module_list(self):
             f"module list &> {list_file}",
         ]
 
-    def _add_software_to_results(self, workspace, app_inst=None):
+    def get_package_list(self, workspace):
+        del workspace
+        app_inst = self.app_inst
         list_file = app_inst.expander.expand_var(
             f"{{experiment_run_dir}}/{self._list_file}"
         )
 
-        if app_inst.result is None:
-            return
-
         if not os.path.exists(list_file):
-            return
-
-        if self._spec_prefix not in app_inst.result.software:
-            app_inst.result.software[self._spec_prefix] = []
-
-        package_list = app_inst.result.software[self._spec_prefix]
+            return []
 
+        pkg_list = []
         pkg_regex = re.compile(r"\S+$")
-
         with open(list_file) as f:
             packages = re.split(r"[0-9]*\)", f.read())
             for spec in packages:
@@ -138,6 +132,7 @@ def _add_software_to_results(self, workspace, app_inst=None):
                     parts = cleaned.split("/")
                     name = parts[0]
                     version = "/".join(parts[1:]) if len(parts) > 1 else ""
-                    package_list.append(
+                    pkg_list.append(
                         {"name": name, "version": version, "variants": ""}
                     )
+        return pkg_list

var/ramble/repos/builtin/package_managers/pip/package_manager.py

@@ -213,30 +213,20 @@ def _warn_mirror_support(self, workspace, app_inst=None):
             "If a software mirror is required, it needs to be set up outside of Ramble"
         )
 
-    def _add_software_to_results(self, workspace, app_inst=None):
+    def get_package_list(self, workspace):
         """Augment the owning experiment's results with software stack information
 
-        This is a registered phase by the base package manager class, so here
-        we only override its base definition.
-
-        Args:
-            workspace (Workspace): A reference to the workspace that owns the
-                                   current pipeline
-            app_inst (Application): A reference to the application instance for
-                                    the current experiment
+        This is called by the `add_software_to_results` phase registered in the base
+        package manager class.
         """
-
         env_path = self.app_inst.expander.env_path
         self.runner.set_dry_run(workspace.dry_run)
         self.runner.configure_env(env_path)
 
-        if self._spec_prefix not in app_inst.result.software:
-            app_inst.result.software[self._spec_prefix] = []
-
-        package_list = app_inst.result.software[self._spec_prefix]
-
+        pkg_list = []
         for info in self.runner.package_provenance():
-            package_list.append(info)
+            pkg_list.append(info)
+        return pkg_list
 
 
 package_name_regex = re.compile(

var/ramble/repos/builtin/package_managers/spack-lightweight/package_manager.py

@@ -493,32 +493,19 @@ def get_spec_str(self, pkg, all_pkgs, compiler):
             out_str += f" (built with {pkg.compiler})"
         return out_str
 
-    def _add_software_to_results(self, workspace, app_inst=None):
+    def get_package_list(self, workspace):
         """Augment the owning experiment's results with software stack information
 
-        This is a registered phase by the base package manager class, so here
-        we only override its base definition.
-
-        Args:
-            workspace (Workspace): A reference to the workspace that owns the
-                                   current pipeline
-            app_inst (Application): A reference to the application instance for
-                                    the current experiment
+        This is called by the `add_software_to_results` phase registered in the base
+        package manager class.
         """
-
-        # Do not manipulate null results
-        if app_inst.result is None:
-            return
-
-        if self._spec_prefix not in app_inst.result.software:
-            app_inst.result.software[self._spec_prefix] = []
-
-        package_list = app_inst.result.software[self._spec_prefix]
-
+        del workspace
+        pkg_list = []
         self.runner.activate()
         for info in self.runner.package_provenance():
-            package_list.append(info)
+            pkg_list.append(info)
         self.runner.deactivate()
+        return pkg_list
 
 
 spack_namespace = "spack"