Add username to sinfo command

src/glicid_spawner/slurm.py

@@ -59,13 +59,16 @@ class SlurmNode:
         self.gpu = SlurmGpu(*re.findall(r'gpu:(\w+):(\d+)', gres)[0] if 'gpu:' in gres else [])
 
 
-def _sinfo_run() -> str:
+def _sinfo_run(username: str = None) -> str:
     """SLURM SINFO run command."""
-    flags = '--federation --noheader --responding --cluster=all'
+    flags = '--federation --noheader --responding'
     fmt = 'Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
-    cmd = shlex.split(f'sinfo {flags} --Format={fmt}')
+    cmd = f'sinfo {flags} --Format={fmt}'
 
-    return subprocess.check_output(cmd).decode('utf-8')
+    if username:
+        cmd = f'su - {username} -c "{cmd}"'
+
+    return subprocess.check_output(shlex.split(cmd, posix=False)).decode('utf-8')
 
 
 def _sinfo_reader(result) -> list:
@@ -73,13 +76,13 @@ def _sinfo_reader(result) -> list:
     return [SlurmNode(*re.findall('.{20}', node)) for node in result.splitlines()]
 
 
-def sinfo(with_states=('idle', 'mixed')) -> dict:
+def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict:
     """SLURM SINFO resources available with a given state(s).
 
     Grouped by cluster and partition names.
 
     """
-    resources = _sinfo_reader(_sinfo_run())
+    resources = _sinfo_reader(_sinfo_run(username=username))
 
     return {
         cluster: {

tests/test_slurm.py

@@ -53,10 +53,19 @@ def test_slurm_sinfo_run(monkeypatch):
         '--federation '
         '--noheader '
         '--responding '
-        '--cluster=all '
         '--Format=Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
     )
 
+    assert _sinfo_run(username='john-doe') == (
+        'su - john-doe -c "'
+        'sinfo '
+        '--federation '
+        '--noheader '
+        '--responding '
+        '--Format=Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
+        '"'
+    )
+
 
 def test_slurm_sinfo_reader():
     """Test SLURM SINFO reader."""