layout changes to whatenv, python3 in bash composer for module

DiamondLightSource · Jun 20, 2023 · c2891a2 · c2891a2
1 parent b3d0e4c
commit c2891a2
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 21 deletions.
diff --git a/src/dls_utilpack/bash_composer.py b/src/dls_utilpack/bash_composer.py
@@ -75,6 +75,7 @@ def __init__(self, message):
     def compose(self) -> List[str]:
         lines = []
 
+        lines.append("")
         lines.append(f"eval echo '{self.__message}'")
 
         return lines
@@ -119,7 +120,7 @@ def compose(self) -> List[str]:
 
         # Put a few things in the log.
         composer.add(Command("module list 2>&1"))
-        composer.add(Command("python --version"))
+        composer.add(Command("python3 --version"))
         composer.add(Print("------ end of modules_load_bash_lines ------"))
 
         return composer.compose_lines()

diff --git a/src/dls_utilpack/whatenv.py b/src/dls_utilpack/whatenv.py
@@ -64,30 +64,42 @@ def compose_as_yaml(self) -> str:
         self.__compose_scalar("HOSTNAME", output_dict)
         self.__compose_scalar("USER", output_dict)
         output_dict["os.getcwd"] = os.getcwd()
-        self.__compose_scalar("SGE_CELL", output_dict)
-        self.__compose_scalar("SGE_EXECD_PORT", output_dict)
-        self.__compose_scalar("SGE_QMASTER_PORT", output_dict)
-
-        self.__compose_scalar("SGE_HGR_gpu", output_dict)
-        self.__compose_scalar("SGE_HGR_m_mem_free", output_dict)
-        self.__compose_scalar("SGE_HGR_TASK_gpu", output_dict)
-        self.__compose_scalar("SGE_HGR_TASK_m_mem_free", output_dict)
-        self.__compose_scalar("PE", output_dict)
-        self.__compose_scalar("QUEUE", output_dict)
-
-        self.__compose_scalar("JOB_ID", output_dict)
-        self.__compose_scalar("JOB_NAME", output_dict)
 
         self.__compose_scalar("VIRTUAL_ENV", output_dict)
 
-        self.__compose_scalar("CONDA_PREFIX", output_dict)
+        sge_dict = {}
+        self.__compose_scalar("SGE_CELL", sge_dict)
+        self.__compose_scalar("SGE_EXECD_PORT", sge_dict)
+        self.__compose_scalar("SGE_QMASTER_PORT", sge_dict)
+        self.__compose_scalar("SGE_HGR_gpu", sge_dict)
+        self.__compose_scalar("SGE_HGR_m_mem_free", sge_dict)
+        self.__compose_scalar("SGE_HGR_TASK_gpu", sge_dict)
+        self.__compose_scalar("SGE_HGR_TASK_m_mem_free", sge_dict)
+        self.__compose_scalar("PE", sge_dict)
+        self.__compose_scalar("QUEUE", sge_dict)
+        self.__compose_scalar("JOB_ID", sge_dict)
+        self.__compose_scalar("JOB_NAME", sge_dict)
+        output_dict["sge"] = sge_dict
+
+        slurm_dict = {}
+        for k in os.environ.keys():
+            if k.startswith("SLURM"):
+                self.__compose_scalar(k, slurm_dict)
+        output_dict["slurm"] = slurm_dict
+
+        conda_dict = {}
+        self.__compose_scalar("CONDA_PREFIX", conda_dict)
         conda_shlvl = int(os.environ.get("CONDA_SHLVL", 0))
         for i in range(conda_shlvl - 1, 0, -1):
-            self.__compose_scalar(f"CONDA_PREFIX_{i}", output_dict)
-
-        self.__compose_paths("MODULEPATH", output_dict)
-        self.__compose_paths("LOADEDMODULES", output_dict)
-        self.__compose_paths("PATH", output_dict)
+            self.__compose_scalar(f"CONDA_PREFIX_{i}", conda_dict)
+        output_dict["conda"] = conda_dict
+
+        modules_dict = {}
+        self.__compose_paths("MODULEPATH", modules_dict)
+        self.__compose_paths("MODULESHOME", modules_dict)
+        self.__compose_paths("LOADEDMODULES", modules_dict)
+        self.__compose_paths("PATH", modules_dict)
+        output_dict["modules"] = modules_dict
 
         uname = platform.uname()
         platform_dict = {}

diff --git a/tests/test_server_context.py b/tests/test_server_context.py
@@ -35,7 +35,7 @@ async def aenter(self):
         """ """
         self.was_entered += 1
 
-    async def aexit(self, type, value, traceback):
+    async def aexit(self, type=None, value=None, traceback=None):
         """ """
         self.was_exited += 1