prince-lua/princeReservation.lua at master · wangsl/prince-lua · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
#!/bin/env lua

local princeReservation = { }

local princeUtils = require "princeUtils"
local princeUsers = require "princeUsers"

local slurm_log = princeUtils.slurm_log
local user_log = princeUtils.user_log

-- local bigIntNumber = 10240*slurm.NO_VAL
local bigIntNumber = princeUtils.bigIntNumber

local function memory_is_specified(mem)
   if mem == nil or mem > bigIntNumber then
      return false
   else
      return true
   end
end

local job_desc = nil

local function check_reservation_morari_is_OK()
   local res_name = "morari"
   slurm_log("Reservation: %s", res_name)

   if job_desc.reservation ~= "morari" then return false end

   if job_desc.min_nodes ~= uint32_NO_VAL then
      if job_desc.min_nodes ~= 1 then
	 user_log("Reservation morari: please specify: --nodes=1")
	 return false
      end
   end

   if job_desc.time_limit > 30 then
      user_log("Reservation morari: time limit 30 mins")
      return false
   end

   if job_desc.shared ~= 0 then
      user_log("Reservation morari: please specify: --exclusive")
      return false
   end

   if job_desc.gres == nil then
      -- CPU only jobs

      if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory > 62*1024 then
	 user_log("Reservation morari: maximum memory for CPU only job is 62GB")
	 return false
      end

      if job_desc.cpus_per_task ~= uint16_NO_VAL then
	 if job_desc.cpus_per_task ~= 20 then
	    user_log("Reservation morari: --cpus-per-task=20")
	    return false
	 end
      end

      if job_desc.partition ~= nil then
	 if job_desc.partition ~= "c26" then
	    user_log("Reservation morari: --partition=c26")
	    return false
	 end
      else
	 job_desc.partition = "c26"
      end
   else
      -- GPU jobs

      if job_desc.gres ~= "gpu:k80:4" then
	 user_log("Reservation morari: --gres=gpu:k80:4")
	 return false
      end

      if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory > 250*1024 then
	 user_log("Reservation morari: maximum memory for GPU only job is 250GB")
	 return false
      end

      if job_desc.cpus_per_task ~= uint16_NO_VAL then
	 if job_desc.cpus_per_task ~= 28 then
	    user_log("Reservation morari: --cpus-per-task=28")
	    return false
	 end
      end

      if job_desc.partition ~= nil then
	 if job_desc.partition ~= "k80_4" then
	    user_log("Reservation morari: --partition=k80_4")
	    return false
	 end
      else
	 job_desc.partition = "k80_4"
      end
   end

   return true
end

local function check_reservation_jupyter_cpu_is_OK()
   local res_name = "jupyter_cpu"
   slurm_log("Reservation: %s", res_name)

   if job_desc.reservation ~= res_name then return false end

   if job_desc.min_nodes ~= princeUtils.uint32_NO_VAL and job_desc.min_nodes ~= 1 then
      user_log("Reservation jupyter_cpu: please specify: --nodes=1")
      return false
   end

   if job_desc.time_limit ~= princeUtils.uint32_NO_VAL and job_desc.time_limit > 240 then
      user_log("Reservation jupyter_cpu: time limit 4 hours")
      return false
   end

   if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory > 4*1024 then
      user_log("Reservation jupyter_cpu: maximum memory for CPU only job is 4GB")
      return false
   end

   if job_desc.cpus_per_task ~= princeUtils.uint16_NO_VAL and job_desc.cpus_per_task > 4 then
      user_log("Reservation jupyter_cpu: --cpus-per-task=4")
      return false
   end

   if job_desc.partition ~= nil and job_desc.partition ~= "jupyterhub_cpu" then
      user_log("Reservation jupyter_cpu: --partition=jupyterhub_cpu")
      return false
   end

   return true
end

local function check_reservation_chung_is_OK()
   local res_name = "chung"
   slurm_log("Reservation: %s", res_name)

   if job_desc.reservation ~= res_name then return false end

   if job_desc.min_nodes ~= uint32_NO_VAL then
      if job_desc.min_nodes ~= 1 then
	 user_log("Reservation chung: please specify: --nodes=1")
	 return false
      end
   end

   if job_desc.time_limit > 240 then
      user_log("Reservation chung: time limit 240 mins")
      return false
   end

   --[[
   if job_desc.shared ~= 0 then
      user_log("Reservation: please specify: --exclusive")
      return false
   end
   --]]

   if job_desc.gres == nil then return false end

   if job_desc.gres ~= "gpu:k80:4" and job_desc.gres ~= "gpu:p40:4"   then
      user_log("Reservation chung: --gres=gpu:k80:4 or --gres=gpu:p40:4")
      return false
   end

   if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory > 250*1024 then
      user_log("Reservation chung: maximum memory for GPU only job is 250GB")
      return false
   end

   if job_desc.cpus_per_task ~= uint16_NO_VAL then
      if job_desc.cpus_per_task ~= 28 then
	 user_log("Reservation chung: --cpus-per-task=28")
	 return false
      end
   end

   return true
end

local function check_reservation_zhang_is_OK()
   local res_name = "zhang"
   slurm_log("Reservation: %s", res_name)

   if job_desc.reservation ~= res_name then return false end

   if job_desc.min_nodes ~= uint32_NO_VAL then
      if job_desc.min_nodes ~= 1 then
	 user_log("Reservation zhang: please specify: --nodes=1")
	 return false
      end
   end

   --[[
   if job_desc.time_limit > 240 then
      user_log("Reservation zhang: time limit 240 mins")
      return false
   end
   --]]
   --[[
   if job_desc.shared ~= 0 then
      user_log("Reservation: please specify: --exclusive")
      return false
   end
   --]]

   if job_desc.gres == nil then return false end

   if job_desc.gres ~= "gpu:k80:8" then
      user_log("Reservation zhang: --gres=gpu:k80:8")
      return false
   end

   if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory ~= 120*1024 then
      user_log("Reservation zhang: --mem=120GB")
      return false
   end

   if job_desc.cpus_per_task ~= uint16_NO_VAL then
      if job_desc.cpus_per_task ~= 10 then
	 user_log("Reservation zhang: --cpus-per-task=10")
	 return false
      end
   end

   if job_desc.partition ~= nil and job_desc.partition ~= "k80_8" then
      user_log("Reservation zhang: --partition=k80_8")
      return false
   end

   return true
end

local function check_reservation_cds_courses_is_OK()
   local res_name = "cds-courses"
   slurm_log("Reservation: %s", res_name)

   if job_desc.reservation ~= res_name then return false end

   if job_desc.min_nodes ~= uint32_NO_VAL then
      if job_desc.min_nodes ~= 1 then
	 user_log("Reservation cds-courses: please specify: --nodes=1")
	 return false
      end
   end

   if job_desc.time_limit > 240 then
      user_log("Reservation cds-courses: time limit 240 mins")
      return false
   end

   --[[
   if job_desc.shared ~= 0 then
      user_log("Reservation: please specify: --exclusive")
      return false
   end
   --]]

   if job_desc.gres == nil then return false end

   if job_desc.gres ~= "gpu:k80:4" and job_desc.gres ~= "gpu:p40:4"   then
      user_log("Reservation cds-courses: --gres=gpu:k80:4 or --gres=gpu:p40:4")
      return false
   end

   if memory_is_specified(job_desc.pn_min_memory) and job_desc.pn_min_memory > 250*1024 then
      user_log("Reservation cds-courses: maximum memory for GPU only job is 250GB")
      return false
   end

   if job_desc.cpus_per_task ~= uint16_NO_VAL then
      if job_desc.cpus_per_task ~= 28 then
	 user_log("Reservation cds-courses: --cpus-per-task=28")
	 return false
      end
   end

   return true
end

local function check_reservation_is_OK(job_desc_)
   job_desc = job_desc_

   -- if job_desc.reservation == "morari" then return check_reservation_morari_is_OK() end

   if job_desc.reservation == "jupyter_cpu" then return check_reservation_jupyter_cpu_is_OK() end

   if job_desc.reservation == "chung" then return check_reservation_chung_is_OK() end

   if job_desc.reservation == "zhang" then return check_reservation_zhang_is_OK() end

   if job_desc.reservation == "cds-courses" then return check_reservation_cds_courses_is_OK() end

   return true
end

-- functions

princeReservation.check_reservation_is_OK = check_reservation_is_OK

slurm_log("To load princeReservation.lua")

return princeReservation