Repository: snek5000/snek5000 (language: python)

Issue:
Include schedulers in cluster utilities
It would be nice to improve a bit the cluster utilities to at least have OAR scheduler.


Discussion / hints:
The only cluster utility we have at the moment is `nproc_available`. I tried to read up the OAR docs, but I couldn't find an equivalent for `SLURM_NPROCS`. Can you start an interactive job of 1 or 2 nodes and check by running `/usr/bin/env`? Once we find that out, you can add it here:

https://github.com/exabl/snek5000/blob/a0c9a056fa52b5e0b5fe3ff89ecb4b7e931bc850/src/snek5000/clusters.py#L12
I noticed that if I use SLURM with 2 nodes, the env variable `SLURM_NPROCS` is missing, so the current implementation is not enough. Some alternatives :

```sh
SLURM_NNODES=2
SLURM_CPUS_ON_NODE=32
SLURM_TASKS_PER_NODE=32(x2)
SLURM_JOB_CPUS_PER_NODE=32(x2)
SLURM_JOB_NUM_NODES=2
```
I use OAR like this:

>
>''' #!/bin/bash
>#OAR -n ....name...
>#OAR -l /nodes=1/core=20,walltime=24:00:00
>
>source /etc/profile
>
>export NEK_SOURCE_ROOT=".../Nek5000"
>export PATH="$PATH:$NEK_SOURCE_ROOT/bin"
>
>echo cbox        >  SESSION.NAME
>echo `pwd`'/' >>  SESSION.NAME
>mpiexec -np 20 ./nek5000 '''
>

I can find something about SLURM.

What I mean is detecting the `20` that you set for the job. We can automate it via `os.environ` / `os.getenv`, so that we can do something like:

    python run_simul.py

which runs a snakemake rule, instead of 

    mpiexec -np 20 ./nek5000

Then the question is which environment variable? You don't have to fix the SLURM thing. I can do that.
There no environment variable like SLURM_NPROCS in OAR. However, something like this should work: `NBCPUS=$(cat ${OAR_NODEFILE} | wc -l)`.

Or something similar in Python!

Fix patch:
diff --git a/src/snek5000/clusters.py b/src/snek5000/clusters.py
index 33a67c48..432be77b 100644
--- a/src/snek5000/clusters.py
+++ b/src/snek5000/clusters.py
@@ -4,18 +4,52 @@
 from .log import logger
 
 
+def _int_getenv(key, default=0):
+    return int(os.getenv(key, default))
+
+
+def _slurm():
+    """Check different environment variables available within a SLURM job.
+
+    .. seealso::
+
+        https://slurm.schedmd.com/srun.html#SECTION_OUTPUT-ENVIRONMENT-VARIABLES
+
+    """
+    return _int_getenv("SLURM_NPROCS") or (
+        _int_getenv("SLURM_NNODES") * _int_getenv("SLURM_CPUS_ON_NODE")
+    )
+
+
+def _oar():
+    """Equivalent to  the shell command::
+
+    NPROC=$(cat ${OAR_NODEFILE} | wc -l).
+
+    """
+    nodefile = os.getenv("OAR_NODEFILE")
+    if nodefile and os.path.exists(nodefile):
+        with open(nodefile) as file:
+            return len(file.readlines())
+    else:
+        return 0
+
+
 def nproc_available():
-    """Detect maximum number of processors available"""
-
-    # Inside job schedulers, number of processors assigened will vary
-    # TODO: extend this loop for other schedulers OAR, PBS etc.
-    for var in ("SLURM_NPROCS",):
-        try:
-            nproc = os.environ[var]
-        except KeyError:
-            continue
-        else:
-            break
+    """Detect maximum number of processors available
+
+    Returns
+    -------
+    int
+
+    """
+    # Inside job schedulers, number of processors assigned will vary
+    # TODO: extend this function for other schedulers PBS etc.
+    cluster_nproc = _slurm() or _oar()
+
+    if cluster_nproc:
+        var = "Cluster environment variables"
+        nproc = cluster_nproc
     else:
         var = "os.cpu_count()"
         nproc = os.cpu_count()


Test patch:
diff --git a/tests/test_clusters.py b/tests/test_clusters.py
new file mode 100644
index 00000000..9c539f31
--- /dev/null
+++ b/tests/test_clusters.py
@@ -0,0 +1,7 @@
+from snek5000.clusters import nproc_available
+
+
+def test_nproc_available():
+    nproc = nproc_available()
+    assert nproc > 0
+    assert isinstance(nproc, int)