Add interactive slurm convenience function. (#101)

spikewrap/configs/backend/hpc.py

@@ -35,3 +35,11 @@ def default_slurm_options() -> Dict:
         "wait": False,
         "env_name": "spikewrap",
     }
+
+
+def default_gpu_partition():
+    return "gpu"
+
+
+def default_cpu_partition():
+    return "cpu"

spikewrap/utils/checks.py

@@ -32,14 +32,14 @@ def check_virtual_machine() -> bool:
     a link to installation instructions.
     """
     if platform.system() == "Linux":
-        has_vm = _system_call_success("singularity version")
+        has_vm = system_call_success("singularity version")
         name = "Singularity"
         link = (
             "https://docs.sylabs.io/guides/main/user-guide/quick_start.html#quick"
             "-installation-steps"
         )
     else:
-        has_vm = _system_call_success("docker -v")
+        has_vm = system_call_success("docker -v")
         name = "Docker"
         if platform.system() == "Windows":
             link = "https://docs.docker.com/desktop/install/windows-install/"
@@ -67,15 +67,15 @@ def docker_desktop_is_running():
     running. "docker ps" requires docker to be running, which can
     be achieved by opening Docker Desktop.
     """
-    return _system_call_success("docker ps")
+    return system_call_success("docker ps")
 
 
 def check_cuda() -> bool:
     """
     Perform a very basic check that NVIDIA drivers are installed. This
     however does not ensure GPU processing will work without error.
     """
-    if _system_call_success("nvidia-smi"):
+    if system_call_success("nvidia-smi"):
         utils.message_user("NVIDIA GPU drivers detected on the system.")
         return True
     else:
@@ -118,7 +118,7 @@ def _check_pip_dependencies() -> None:
         utils.message_user("All python dependencies are installed.")
 
 
-def _system_call_success(command: str) -> bool:
+def system_call_success(command: str) -> bool:
     return (
         subprocess.run(
             command, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL

spikewrap/utils/slurm.py

@@ -5,8 +5,13 @@
 
 import submitit
 
-from ..configs.backend.hpc import default_slurm_options
+from ..configs.backend.hpc import (
+    default_cpu_partition,
+    default_gpu_partition,
+    default_slurm_options,
+)
 from . import utils
+from .checks import system_call_success
 
 
 def run_job(kwargs, command_func: Callable, command_name: str) -> None:
@@ -229,5 +234,31 @@ def send_user_start_message(
 
 
 def is_slurm_installed():
-    slurm_installed = subprocess.run("sinfo -v", shell=True).returncode == 0
+    slurm_installed = system_call_success("sinfo -v")
     return slurm_installed
+
+
+def run_interactive_slurm(partition=None, gpu=True, mem_gb="40GB", cpus=16):
+    """
+    A convenience function to start an interactive SLURM session. Gives quick
+    access for some basic settings - the purpose is not to provide a comprehensive
+    wrapper, of course native `srun` can be called. It is only to provide a
+    quick convenience wrapper to get started quick with some sensible defaults.
+    """
+    if not is_slurm_installed():
+        raise RuntimeError(
+            "Cannot setup interactive SLURM because SLURM is not "
+            "installed on this machine."
+        )
+
+    if partition is None:
+        partition = default_gpu_partition() if gpu else default_cpu_partition()
+
+    gpu_opt = "--gres=gpu:1" if gpu else ""
+
+    default_opts = default_slurm_options()
+    exclude = default_opts["exclude"]
+
+    command = f"srun -p {partition} {gpu_opt} -n {cpus} --mem={mem_gb} --exclude {exclude} --pty bash -i"
+
+    subprocess.run(command, shell=True)