Merge pull request #705 from douglasjacobsen/modifier-inventory

Allow modifiers to add attributes into experiment inventories

lib/ramble/ramble/application.py

@@ -162,6 +162,7 @@ def __init__(self, file_path):
             "software": [],
             "templates": [],
             "package_manager": [],
+            "modifier_artifacts": [],
         }
         self.experiment_hash = None
 
@@ -1492,6 +1493,16 @@ def _write_inventory(self, workspace, app_inst=None):
         experiment_run_dir = self.expander.experiment_run_dir
         inventory_file = os.path.join(experiment_run_dir, self._inventory_file_name)
 
+        # Populate modifier artifacts portion of inventory
+        # This happens here to allow modifiers to hash files
+        # that are downloaded within phases earlier than this.
+        for mod_inst in self._modifier_instances:
+            inventory = mod_inst.artifact_inventory(workspace, app_inst)
+            if inventory:
+                self.hash_inventory["modifier_artifacts"].append(
+                    {"name": mod_inst.name, "mode": mod_inst._usage_mode, "artifacts": inventory}
+                )
+
         with lk.WriteTransaction(self.experiment_lock()):
             with open(inventory_file, "w+") as f:
                 spack.util.spack_json.dump(self.hash_inventory, f)

lib/ramble/ramble/modifier.py

@@ -244,6 +244,20 @@ def run_phase_hook(self, workspace, pipeline, hook_name):
 
                 phase_func(workspace)
 
+    def artifact_inventory(self, workspace, app_inst=None):
+        """Return an inventory of modifier artifacts
+
+        Artifact inventories are up to the individual modifier to define the
+        format of.
+
+        This will then show up in an experiment inventory.
+
+        Returns:
+            (Any) Artifact inventory for this modifier
+        """
+
+        return None
+
     def _prepare_analysis(self, workspace):
         """Hook to perform analysis that a modifier defines.
 

var/ramble/repos/builtin/modifiers/pyxis-enroot/modifier.py

@@ -9,6 +9,7 @@
 import os
 
 from ramble.modkit import *
+from ramble.util.hashing import hash_file, hash_string
 
 import llnl.util.filesystem as fs
 
@@ -270,6 +271,41 @@ def _extract_from_sqsh(self, workspace, app_inst=None):
                     unsquash_args + [expanded_path],
                 )
 
+    def artifact_inventory(self, workspace, app_inst=None):
+        """Return hash of container uri and sqsh file if they exist
+
+        Args:
+            workspace (Workspace): Reference to workspace
+            app_inst (ApplicationBase): Reference to application instance
+
+        Returns:
+            (dict): Artifact inventory for container attributes
+        """
+        container_name = self.expander.expand_var_name("container_name")
+        container_path = self.expander.expand_var_name("container_path")
+        container_uri = self.expander.expand_var_name("container_uri")
+        inventory = []
+
+        if self._usage_mode == "disabled":
+            return inventory
+
+        inventory.append(
+            {
+                "container_uri": container_uri,
+                "digest": hash_string(container_uri),
+            }
+        )
+
+        if os.path.isfile(container_path):
+            inventory.append(
+                {
+                    "container_name": container_name,
+                    "digest": hash_file(container_path),
+                }
+            )
+
+        return inventory
+
     # TODO: Decide on backing up sqsh files.
     #       The following code works. But there's not a nice way to auto-extract the sqsh file out of the mirror
     #       This is because the import functionality uses `enroot` directly, which bypasses