job_submit.lua (routes based off gres and cpus per node)

job_submit.lua

--[[
   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