Skip to content

Commit

Permalink
add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@jeffnvidia
jeffnvidia committed May 28, 2024
1 parent c8c1adc commit 8415b33
Showing 1 changed file with 81 additions and 16 deletions.
97 changes: 81 additions & 16 deletions tests/test_slurm_system.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,31 @@
@pytest.fixture
def slurm_system():
nodes = [
SlurmNode(name="nodeA001", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="nodeB001", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-115", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-116", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-117", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-118", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-119", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-120", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-121", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node-122", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
]
backup_nodes = [
SlurmNode(name="node01", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node02", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node03", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node04", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node05", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node06", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node07", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
SlurmNode(name="node08", partition="backup", state=SlurmNodeState.UNKNOWN_STATE),
]
system = SlurmSystem(
name="test_system",
install_path="/fake/path",
output_path="/fake/output",
default_partition="main",
partitions={"main": nodes},
partitions={"main": nodes, "backup": backup_nodes},
)
return system

Expand All @@ -29,7 +45,7 @@ def test_parse_squeue_output(slurm_system):


def test_parse_squeue_output_with_node_ranges_and_root_user(slurm_system):
squeue_output = "nodeA[001-008],nodeB[001-008]|root"
squeue_output = "nodeA[001-008]|root\nnodeB[001-008]|root"
user_map = slurm_system.parse_squeue_output(squeue_output)

expected_nodes = [f"nodeA{str(i).zfill(3)}" for i in range(1, 9)] + [f"nodeB{str(i).zfill(3)}" for i in range(1, 9)]
Expand All @@ -39,24 +55,73 @@ def test_parse_squeue_output_with_node_ranges_and_root_user(slurm_system):


def test_parse_sinfo_output(slurm_system):
sinfo_output = (
"PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n"
"main up infinite 1 idle nodeA001\n"
"main up infinite 1 idle nodeB001"
)
node_user_map = {"nodeA001": "root", "nodeB001": "user"}
sinfo_output = """PARTITION AVAIL TIMELIMIT NODES STATE NODELIST
main up 3:00:00 1 inval node-081
main up 3:00:00 5 drain node-[065-066,114,124-125]
main up 3:00:00 2 resv node-[034-035]
main up 3:00:00 88 alloc node-[033,036-064,067-080,082-113,115-123,126-128]
backup up 12:00:00 16 idle node[01-16]
"""
node_user_map = {
"": "user1",
"node-033": "user2",
"node-034": "user3",
"node-056": "user3",
"node-057": "user3",
"node-058": "user3",
"node-049": "user4",
"node-050": "user4",
"node01": "user5",
"node02": "user5",
"node03": "user5",
"node04": "user5",
"node05": "user5",
"node06": "user5",
"node07": "user5",
"node08": "user5",
}
slurm_system.parse_sinfo_output(sinfo_output, node_user_map)
assert slurm_system.partitions["main"][0].state == SlurmNodeState.IDLE
assert slurm_system.partitions["main"][1].state == SlurmNodeState.IDLE
for i in range(len(slurm_system.partitions["main"])):
assert slurm_system.partitions["main"][i].state == SlurmNodeState.ALLOCATED
for i in range(len(slurm_system.partitions["backup"])):
assert slurm_system.partitions["backup"][i].state == SlurmNodeState.IDLE


@patch("cloudai.schema.system.SlurmSystem.get_squeue")
@patch("cloudai.schema.system.SlurmSystem.get_sinfo")
def test_update_node_states_with_mocked_outputs(mock_get_sinfo, mock_get_squeue, slurm_system):
mock_get_squeue.return_value = "nodeA001|root"
mock_get_sinfo.return_value = "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "main up infinite 1 idle nodeA001"
mock_get_squeue.return_value = "node-115|user1"
mock_get_sinfo.return_value = "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "main up infinite 1 idle node-115"

slurm_system.update_node_states()
for node in slurm_system.partitions["main"]:
if node.name == "node-115":
assert node.state == SlurmNodeState.IDLE
assert node.user == "user1"

mock_get_squeue.return_value = "node01|root"
mock_get_sinfo.return_value = (
"PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "backup up infinite 1 allocated node01"
)

slurm_system.update_node_states()
for node in slurm_system.partitions["backup"]:
if node.name == "node01":
assert node.state == SlurmNodeState.ALLOCATED
assert node.user == "root"


def test_parse_node_list(slurm_system):
node_lists = ["node-[048-051]", "node-[055,114]", "", "node-001", "node[1-4]"]

expected_outputs = [
["node-048", "node-049", "node-050", "node-051"],
["node-055", "node-114"],
[],
["node-001"],
["node1", "node2", "node3", "node4"],
]

assert slurm_system.partitions["main"][0].state == SlurmNodeState.IDLE
assert slurm_system.partitions["main"][0].user == "root"
for i, node_list in enumerate(node_lists):
parsed_nodes = slurm_system.parse_node_list(node_list)
assert parsed_nodes == expected_outputs[i]

0 comments on commit 8415b33

Please sign in to comment.