Skip to content

Commit 0f1e2d0

Browse files
aldbrweb-flow
authored andcommitted
sweep: DIRACGrid#8214 feat: integrate WholeNode option in Slurm plugin
1 parent 3bb6ac8 commit 0f1e2d0

File tree

2 files changed

+51
-1
lines changed

2 files changed

+51
-1
lines changed

src/DIRAC/Resources/Computing/BatchSystems/SLURM.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ def submitJob(self, **kwargs):
4040
executable = kwargs["Executable"]
4141
account = kwargs.get("Account", "")
4242
numberOfProcessors = kwargs.get("NumberOfProcessors", 1)
43+
wholeNode = kwargs.get("WholeNode", False)
4344
# numberOfNodes is treated as a string as it can contain values such as "2-4"
4445
# where 2 would represent the minimum number of nodes to allocate, and 4 the maximum
4546
numberOfNodes = kwargs.get("NumberOfNodes", "1")
@@ -72,7 +73,10 @@ def submitJob(self, **kwargs):
7273
# One pilot (task) per node, allocating a certain number of processors
7374
cmd += "--ntasks-per-node=1 "
7475
cmd += "--nodes=%s " % numberOfNodes
75-
cmd += "--cpus-per-task=%d " % numberOfProcessors
76+
if wholeNode:
77+
cmd += "--exclusive "
78+
else:
79+
cmd += "--cpus-per-task=%d " % numberOfProcessors
7680
if numberOfGPUs:
7781
cmd += "--gpus-per-task=%d " % int(numberOfGPUs)
7882
# Additional options

src/DIRAC/Resources/Computing/BatchSystems/test/Test_SLURM.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,3 +198,49 @@ def test_getJobOutputFiles(numberOfNodes, outputContent, expectedContent):
198198

199199
os.remove(outputFile)
200200
os.remove(errorFile)
201+
202+
203+
def test_submitJob_cmd_generation(mocker):
204+
"""Test submitJob() command string generation for various kwargs"""
205+
slurm = SLURM()
206+
# Mock subprocess.Popen to capture the command
207+
popen_mock = mocker.patch("subprocess.Popen")
208+
process_mock = popen_mock.return_value
209+
process_mock.communicate.return_value = ("Submitted batch job 1234\n", "")
210+
process_mock.returncode = 0
211+
212+
# Minimal kwargs
213+
kwargs = {
214+
"Executable": "/bin/echo",
215+
"OutputDir": "/tmp",
216+
"ErrorDir": "/tmp",
217+
"Queue": "testq",
218+
"SubmitOptions": "",
219+
"JobStamps": ["stamp1"],
220+
"NJobs": 1,
221+
}
222+
# Test default (WholeNode False)
223+
slurm.submitJob(**kwargs)
224+
cmd = popen_mock.call_args[0][0]
225+
assert "--cpus-per-task=1" in cmd
226+
assert "--exclusive" not in cmd
227+
228+
# Test WholeNode True disables --cpus-per-task and adds --exclusive
229+
kwargs["WholeNode"] = True
230+
slurm.submitJob(**kwargs)
231+
cmd = popen_mock.call_args[0][0]
232+
assert "--exclusive" in cmd
233+
assert "--cpus-per-task" not in cmd
234+
235+
# Test NumberOfProcessors
236+
kwargs["WholeNode"] = False
237+
kwargs["NumberOfProcessors"] = 8
238+
slurm.submitJob(**kwargs)
239+
cmd = popen_mock.call_args[0][0]
240+
assert "--cpus-per-task=8" in cmd
241+
242+
# Test NumberOfGPUs
243+
kwargs["NumberOfGPUs"] = 2
244+
slurm.submitJob(**kwargs)
245+
cmd = popen_mock.call_args[0][0]
246+
assert "--gpus-per-task=2" in cmd

0 commit comments

Comments
 (0)