Skip to content

Commit

Permalink
gpu_cmode.lua: add missing initializations in slurm_spank_task_exit()
Browse files Browse the repository at this point in the history
  • Loading branch information
kcgthb committed Mar 9, 2022
1 parent 3eb9901 commit 73805f0
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions gpu_cmode.lua
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,20 @@ function slurm_spank_task_exit(spank)
return SPANK.SUCCESS
end

-- get GPU ids from CUDA_VISIBLE_DEVICES
device_ids = spank:getenv("CUDA_VISIBLE_DEVICES")
if device_ids == nil or device_ids == "" then
SPANK.log_error(myname .. ": CUDA_VISIBLE_DEVICES not set.")
return SPANK.FAILURE
end

-- check for nvidia-smi
nvs_path = exec("which nvidia-smi")
if nvs_path:match("nvidia%-smi$") == nil then
SPANK.log_error(myname .. ": can't find nvidia-smi in PATH.")
return SPANK.FAILURE
end

-- reset compute mode on GPUs
SPANK.log_info(myname .. ": resetting compute mode to default '%s'" ..
" on GPU(s): %s\n", default_cmode, device_ids)
Expand Down

0 comments on commit 73805f0

Please sign in to comment.