Merge branch 'master' into master

percyfal · Sep 1, 2022 · 625e956 · 625e956
2 parents f3899a2 + f773e01
commit 625e956
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/{{cookiecutter.profile_name}}/CookieCutter.py b/{{cookiecutter.profile_name}}/CookieCutter.py
@@ -9,11 +9,19 @@
     settings = json.load(fh)
 
 
+def from_entry_or_env(values, key):
+    """Return value from ``values`` and override with environment variables."""
+    if key in os.environ:
+        return os.environ[key]
+    else:
+        return values[key]
+
+
 class CookieCutter:
 
-    SBATCH_DEFAULTS = settings['SBATCH_DEFAULTS']
-    CLUSTER_NAME = settings['CLUSTER_NAME']
-    CLUSTER_CONFIG = settings['CLUSTER_CONFIG']
+    SBATCH_DEFAULTS = from_entry_or_env(settings, "SBATCH_DEFAULTS")
+    CLUSTER_NAME = from_entry_or_env(settings, "CLUSTER_NAME")
+    CLUSTER_CONFIG = from_entry_or_env(settings, "CLUSTER_CONFIG")
 
     @staticmethod
     def get_cluster_option() -> str:

diff --git a/{{cookiecutter.profile_name}}/slurm-sidecar.py b/{{cookiecutter.profile_name}}/slurm-sidecar.py
@@ -39,8 +39,10 @@
 from CookieCutter import CookieCutter
 
 
-#: Enables debug messages for slurm sidecard.
+#: Enables debug messages for slurm sidecar.
 DEBUG = bool(int(os.environ.get("SNAKEMAKE_SLURM_DEBUG", "0")))
+#: Enables HTTP request logging in sidecar.
+LOG_REQUESTS = bool(int(os.environ.get("SNAKEMAKE_SLURM_LOG_REQUESTS", "0")))
 #: Command to call when calling squeue
 SQUEUE_CMD = os.environ.get("SNAKEMAKE_SLURM_SQUEUE_CMD", "squeue")
 #: Number of seconds to wait between ``squeue`` calls.
@@ -247,6 +249,10 @@ def do_POST(self):
         self.end_headers()
         logger.debug("--- END POST")
 
+    def log_request(self, *args, **kwargs):
+        if LOG_REQUESTS:
+            super().log_request(*args, **kwargs)
+
 
 class JobStateHttpServer(http.server.HTTPServer):
     """The HTTP server class"""

diff --git a/{{cookiecutter.profile_name}}/slurm_utils.py b/{{cookiecutter.profile_name}}/slurm_utils.py
@@ -10,6 +10,8 @@
 from time import time as unix_time
 from typing import Union
 from uuid import uuid4
+import shlex
+from io import StringIO
 
 from CookieCutter import CookieCutter
 from snakemake import io
@@ -46,7 +48,7 @@ def parse_jobscript():
 
 def parse_sbatch_defaults(parsed):
     """Unpack SBATCH_DEFAULTS."""
-    d = parsed.split() if type(parsed) == str else parsed
+    d = shlex.split(parsed) if type(parsed) == str else parsed
     args = {}
     for keyval in [a.split("=") for a in d]:
         k = keyval[0].strip().strip("-")