Skip to content

Commit 1d94303

Browse files
authored
Merge pull request #8214 from aldbr/v8.0_FEAT_integrate-whole-node-in-slurm
[8.0] feat: integrate `WholeNode` option in `Slurm` plugin
2 parents 80d5ef9 + 6e20412 commit 1d94303

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
@@ -199,3 +199,49 @@ def test_getJobOutputFiles(numberOfNodes, outputContent, expectedContent):
199199

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

0 commit comments

Comments
 (0)