135 lines
3.6 KiB
Python
135 lines
3.6 KiB
Python
"""Test SLURM module."""
|
|
|
|
from pathlib import Path
|
|
|
|
from glicid_spawner.slurm import (
|
|
SlurmCpu,
|
|
SlurmGpu,
|
|
SlurmNode,
|
|
_sinfo_reader,
|
|
_sinfo_run,
|
|
sinfo,
|
|
subprocess,
|
|
)
|
|
|
|
DATA = Path(__file__).parent / 'data'
|
|
|
|
|
|
def test_slurm_dataclasses():
|
|
"""Test SLURM dataclasses formatter."""
|
|
cpu = SlurmCpu(1, '2', 4.0)
|
|
|
|
assert cpu.allocated == 1
|
|
assert cpu.idle == 2
|
|
assert cpu.total == 4
|
|
|
|
assert isinstance(cpu.allocated, int)
|
|
assert isinstance(cpu.idle, int)
|
|
assert isinstance(cpu.total, int)
|
|
|
|
gpu = SlurmGpu('fOo', '1')
|
|
|
|
assert gpu # __bool__
|
|
assert gpu.name == 'Foo'
|
|
assert gpu.nb == 1
|
|
|
|
assert isinstance(gpu.name, str)
|
|
assert isinstance(gpu.nb, int)
|
|
|
|
# Default values
|
|
gpu = SlurmGpu()
|
|
|
|
assert not gpu # __bool__
|
|
assert gpu.name == 'None'
|
|
assert gpu.nb == 0
|
|
|
|
|
|
def test_slurm_sinfo_run(monkeypatch):
|
|
"""Test SLURM SINFO run command."""
|
|
monkeypatch.setattr(subprocess, 'check_output', lambda cmd: ' '.join(cmd).encode())
|
|
|
|
assert _sinfo_run() == (
|
|
'sinfo '
|
|
'--federation '
|
|
'--noheader '
|
|
'--responding '
|
|
'--Format=Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
|
|
)
|
|
|
|
|
|
def test_slurm_sinfo_reader():
|
|
"""Test SLURM SINFO reader."""
|
|
nodes = _sinfo_reader((DATA / 'sinfo.txt').read_text())
|
|
|
|
for node in nodes:
|
|
assert isinstance(node, SlurmNode)
|
|
|
|
node = nodes[0]
|
|
|
|
assert node.cluster is None
|
|
assert node.partition == 'Devel'
|
|
assert node.hostname == 'nazare001'
|
|
assert node.state == 'idle'
|
|
assert node.cpu.allocated == 0
|
|
assert node.cpu.idle == node.cpu.total == 20
|
|
assert node.mem == 128
|
|
assert not node.gpu
|
|
|
|
assert [node.cluster for node in nodes] == 7 * [None] + 6 * ['nautilus'] + 4 * ['waves']
|
|
|
|
assert len([node for node in nodes if node.state in ('idle', 'mixed')]) == 7
|
|
|
|
for node in nodes:
|
|
if node.state == 'idle':
|
|
assert node.cpu.allocated == 0
|
|
assert node.cpu.idle > 0
|
|
elif node.state == 'mixed':
|
|
assert node.cpu.allocated > 0
|
|
assert node.cpu.idle > 0
|
|
elif node.state == 'allocated':
|
|
assert node.cpu.allocated > 0
|
|
assert node.cpu.idle == 0
|
|
|
|
assert sum(node.mem for node in nodes) == 4_672
|
|
|
|
assert [node.gpu.name for node in nodes if node.gpu] == [
|
|
'T4',
|
|
'A40',
|
|
'P100',
|
|
'K80',
|
|
'P100',
|
|
] + 4 * ['A100']
|
|
|
|
assert [node.gpu.nb for node in nodes if node.gpu] == [2, 2, 2, 4, 1, 1, 2, 4, 2]
|
|
|
|
|
|
def test_slurm_sinfo_resources(monkeypatch):
|
|
"""Test SLURM SINFO resources."""
|
|
monkeypatch.setattr(subprocess, 'check_output', lambda _: (DATA / 'sinfo.txt').read_bytes())
|
|
|
|
clusters = sinfo()
|
|
|
|
assert isinstance(clusters, dict)
|
|
assert len(clusters) == 3
|
|
assert list(clusters) == [None, 'nautilus', 'waves']
|
|
|
|
assert [len(partitions) for partitions in clusters.values()] == [2, 2, 2]
|
|
|
|
nautilus = clusters['nautilus']
|
|
|
|
assert isinstance(nautilus, dict)
|
|
assert len(nautilus) == 2
|
|
assert list(nautilus) == ['gpu', 'all']
|
|
|
|
gpus = nautilus['gpu']
|
|
|
|
assert len(gpus) == 2
|
|
assert [partition.hostname for partition in gpus] == ['gnode1', 'gnode2']
|
|
assert [partition.cpu.allocated for partition in gpus] == [4, 0]
|
|
assert [partition.cpu.idle for partition in gpus] == [92, 96]
|
|
assert [partition.mem for partition in gpus] == [768, 256]
|
|
assert [partition.gpu.name for partition in gpus] == ['A100', 'A100']
|
|
assert [partition.gpu.nb for partition in gpus] == [1, 2]
|
|
|
|
# Get only `idle` nodes
|
|
assert [len(partitions) for partitions in sinfo(with_states=('idle')).values()] == [1, 2, 1]
|