Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
job_submit.lua (routes based off gres and cpus per node)
--[[
For use, this script should be copied into a file name "job_submit.lua"
in the same directory as the SLURM configuration file, slurm.conf.
--]]
function slurm_job_submit(job_desc, part_list, submit_uid)
test_enabled = (submit_uid == __ENTER_UID_HERE__) -- user is 'renfro'
-- test_enabled = false
if (test_enabled) then
-- As the default partition is set later by SLURM we need to set it
-- here using the same logic
if (job_desc.partition == nil) then
local default_partition = "batch"
job_desc.partition = default_partition
slurm.log_info(
"slurm_job_submit: No partition specified, moved to batch.")
end
-- If we reserved a GPU,
if (job_desc.gres ~= nil) then
-- we'll route to the appropriate GPU partition:
-- batch -> gpu
-- interactive -> gpu-interactive
-- debug -> gpu-debug
local partition = ""
if (job_desc.partition == 'batch') then
partition = "gpu"
else
partition = "gpu-"..job_desc.partition
end
slurm.log_info("slurm_job_submit: for user %u, setting partition: %s", submit_uid, partition)
job_desc.partition = partition
else
-- Make default reservation values explicit for easier decision-making
if (job_desc.cpus_per_task == 65534) then
job_desc.cpus_per_task = 1
slurm.log_info("slurm_job_submit: setting cpus_per_task = 1.")
end
local ntasks_per_node_specified = true
local ntasks_specified = true
if ((job_desc.ntasks_per_node == 65534) and
(job_desc.ntasks == 65534)) then
job_desc.ntasks_per_node = 1
job_desc.ntasks = 1
ntasks_per_node_specified = false
ntasks_specified = false
slurm.log_info(
"slurm_job_submit: setting ntasks, ntasks_per_node = 1.")
elseif ((job_desc.ntasks_per_node == 65534) and
(job_desc.ntasks ~= 65534)) then
ntasks_per_node_specified = false
slurm.log_info(
"slurm_job_submit: ntasks specified, ntasks_per_node not specified.")
elseif ((job_desc.ntasks_per_node ~= 65534) and
(job_desc.ntasks == 65534)) then
ntasks_specified = false
slurm.log_info(
"slurm_job_submit: ntasks_per_node specified, ntasks not specified.")
end
--[[
How many CPUs are requested per node?
If ntasks_per_node was explicitly specified, could be
cpus_per_tasks*ntasks_per_node
If ntasks was explicitly specified, could be
cpus_per_task <= cpus_per_node <= cpus_per_task*ntasks
--]]
if (ntasks_per_node_specified) then
if ((job_desc.cpus_per_task)*(job_desc.ntasks_per_node) <= 12) then
slurm.log_info(
string.format("slurm_job_submit: candidate for anywhere queue, (cpus/task)*(task/node)=(%d)*(%d)=%d.",
job_desc.cpus_per_task, job_desc.ntasks_per_node,
(job_desc.cpus_per_task)*(job_desc.ntasks_per_node)))
-- we'll route to the appropriate anywhere partition:
-- batch -> any
-- interactive -> any-interactive
-- debug -> any-debug
local partition = ""
if (job_desc.partition == 'batch') then
partition = "any"
else
partition = "any-"..job_desc.partition
end
slurm.log_info("slurm_job_submit: for user %u, setting partition: %s", submit_uid, partition)
job_desc.partition = partition
end
elseif (ntasks_specified and (job_desc.cpus_per_task <= 12)) then
slurm.log_info(
string.format("slurm_job_submit: candidate for anywhere queue, (cpus/task)=%d.",
job_desc.cpus_per_task))
-- we'll route to the appropriate anywhere partition:
-- batch -> any
-- interactive -> any-interactive
-- debug -> any-debug
local partition = ""
if (job_desc.partition == 'batch') then
partition = "any"
else
partition = "any-"..job_desc.partition
end
slurm.log_info("slurm_job_submit: for user %u, setting partition: %s", submit_uid, partition)
job_desc.partition = partition
end -- if job is small enough for anywhere queue
end -- if we reserved a GPU
else -- use default logic for production
-- Nothing to do here
local x = 1
end -- detect if testing or production
return slurm.SUCCESS
end
function slurm_job_modify(job_desc, job_rec, part_list, modify_uid)
return slurm.SUCCESS
end
slurm.log_info("initialized")
return slurm.SUCCESS
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment