WIP-Adding Prov env info

pydra/engine/audit.py

@@ -6,6 +6,8 @@
 from ..utils.messenger import send_message, make_message, gen_uuid, now, AuditFlag
 from .helpers import ensure_list, gather_runtime_info, hash_file
 from .specs import attr_fields, File, Directory
+import subprocess as sp
+import platform
 
 
 class Audit:
@@ -171,7 +173,6 @@ def audit_check(self, flag):
         return self.audit_flags & flag
 
     def audit_task(self, task):
-        import subprocess as sp
 
         label = task.name
 
@@ -221,3 +222,25 @@ def audit_task(self, task):
         }
 
         self.audit_message(start_message, AuditFlag.PROV)
+
+        env_id = f"uid:{gen_uuid()}"
+        os_plat = str(platform.platform())
+        env_vars = str(os.environ)
+        # check if using conda environment
+        if "CONDA_PREFIX" in os.environ:
+            conda_env_path = str(os.environ["CONDA_PREFIX"])
+            conda_env_name = conda_env_path.split("/")[-1]
+
+        else:
+            conda_env_name = str(None)
+
+        env_message = {
+            "@id": env_id,
+            "Label": f"Conda environment: {conda_env_name}",
+            "EnvVars": env_vars,
+            "OperatingSystem": os_plat,
+            "Dependencies": "test",
+        }
+
+        # Fetch env info
+        self.audit_message(env_message, AuditFlag.PROV)

pydra/engine/tests/test_task.py

@@ -1002,12 +1002,20 @@ def testfunc(a: int, b: float = 0.1) -> ty.NamedTuple("Output", [("out", float)]
         return a + b
 
     from glob import glob
+    import platform
 
     funky = testfunc(a=2, audit_flags=AuditFlag.PROV, messengers=FileMessenger())
     funky.cache_dir = tmpdir
     funky()
     message_path = tmpdir / funky.checksum / "messages"
-
+    op_sys = platform.platform()
+    env_vars = str(os.environ)
+    if "CONDA_PREFIX" in os.environ:
+        conda_env_path = str(os.environ["CONDA_PREFIX"])
+        conda_env_name = conda_env_path.split("/")[-1]
+
+    else:
+        conda_env_name = str(None)
     for file in glob(str(message_path) + "/*.jsonld"):
         with open(file, "r") as f:
             data = json.load(f)
@@ -1020,8 +1028,12 @@ def testfunc(a: int, b: float = 0.1) -> ty.NamedTuple("Output", [("out", float)]
                     assert None == data["Label"]
             if "AssociatedWith" in data:
                 assert None == data["AssociatedWith"]
-
-    # assert any(json_content)
+            if "OperatingSystem" in data:
+                assert op_sys == data["OperatingSystem"]
+            if "EnvVars" in data:
+                assert env_vars in data["EnvVars"]
+                if "Label" in data:
+                    assert conda_env_name in data["Label"]
 
 
 def test_audit_shellcommandtask(tmpdir):
@@ -1079,7 +1091,6 @@ def test_audit_shellcommandtask_file(tmpdir):
     shutil.copy("test.txt", tmpdir)
     shutil.copy("test2.txt", tmpdir)
 
-
     cmd = "cat"
     file_in = tmpdir / "test.txt"
     file_in_2 = tmpdir / "test2.txt"
@@ -1100,7 +1111,7 @@ def test_audit_shellcommandtask_file(tmpdir):
                     },
                 ),
             ),
-                    (
+            (
                 "in_file_2",
                 attr.ib(
                     type=File,
@@ -1139,6 +1150,7 @@ def test_audit_shellcommandtask_file(tmpdir):
                         assert data["AtLocation"] == str(file_in_2)
                         assert data["digest"] == test_file_hash_2
 
+
 def test_audit_shellcommandtask_version(tmpdir):
     import subprocess as sp
 
@@ -1268,7 +1280,7 @@ def testfunc(a: int, b: float = 0.1) -> ty.NamedTuple("Output", [("out", float)]
     from glob import glob
 
     assert len(glob(str(tmpdir / funky.checksum / "proc*.log"))) == 1
-    assert len(glob(str(message_path / "*.jsonld"))) == 7
+    assert len(glob(str(message_path / "*.jsonld"))) == 8
 
     # commented out to speed up testing
     collect_messages(tmpdir / funky.checksum, message_path, ld_op="compact")