diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 6605610..dec3cb9 100644 --- a/src/glicid_spawner/slurm.py +++ b/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: { diff --git a/tests/test_slurm.py b/tests/test_slurm.py index ebfff79..2039602 100644 --- a/tests/test_slurm.py +++ b/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."""