@@ -134,13 +134,14 @@ def __init__(self, plugin_args=None):
134134 # GPU found on system
135135 self .n_gpus_visible = MultiProcPlugin .gpu_count ()
136136 # proc per GPU set by user
137- self .n_gpu_procs = plugin_args .get ('n_gpu_procs' , self .n_gpus_visible )
137+ self .n_gpu_procs = self . plugin_args .get ('n_gpu_procs' , self .n_gpus_visible )
138138
139139 # total no. of processes allowed on all gpus
140140 if self .n_gpu_procs > self .n_gpus_visible :
141141 logger .info (
142- 'Total number of GPUs proc requested (%d) exceeds the available number of GPUs (%d) on the system. Using requested GPU slots at your own risk!' % (
143- self .n_gpu_procs , self .n_gpus_visible ))
142+ 'Total number of GPUs proc requested (%d) exceeds the available number of GPUs (%d) on the system. Using requested GPU slots at your own risk!'
143+ % (self .n_gpu_procs , self .n_gpus_visible )
144+ )
144145
145146 # Instantiate different thread pools for non-daemon processes
146147 logger .debug (
@@ -220,9 +221,7 @@ def _prerun_check(self, graph):
220221 if self .raise_insufficient :
221222 raise RuntimeError ("Insufficient resources available for job" )
222223 if np .any (np .array (tasks_gpu_th ) > self .n_gpu_procs ):
223- logger .warning (
224- 'Nodes demand more GPU than allowed (%d).' ,
225- self .n_gpu_procs )
224+ logger .warning ('Nodes demand more GPU than allowed (%d).' , self .n_gpu_procs )
226225 if self .raise_insufficient :
227226 raise RuntimeError ('Insufficient GPU resources available for job' )
228227
@@ -257,7 +256,9 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
257256 )
258257
259258 # Check available resources by summing all threads and memory used
260- free_memory_gb , free_processors , free_gpu_slots = self ._check_resources (self .pending_tasks )
259+ free_memory_gb , free_processors , free_gpu_slots = self ._check_resources (
260+ self .pending_tasks
261+ )
261262
262263 stats = (
263264 len (self .pending_tasks ),
@@ -267,7 +268,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
267268 free_processors ,
268269 self .processors ,
269270 free_gpu_slots ,
270- self .n_gpu_procs
271+ self .n_gpu_procs ,
271272 )
272273 if self ._stats != stats :
273274 tasks_list_msg = ""
@@ -338,8 +339,11 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
338339 is_gpu_node = self .procs [jobid ].is_gpu_node ()
339340
340341 # If node does not fit, skip at this moment
341- if (next_job_th > free_processors or next_job_gb > free_memory_gb
342- or (is_gpu_node and next_job_gpu_th > free_gpu_slots )):
342+ if (
343+ next_job_th > free_processors
344+ or next_job_gb > free_memory_gb
345+ or (is_gpu_node and next_job_gpu_th > free_gpu_slots )
346+ ):
343347 logger .debug (
344348 "Cannot allocate job %d (%0.2fGB, %d threads, %d GPU slots)." ,
345349 jobid ,
@@ -424,6 +428,7 @@ def gpu_count():
424428 n_gpus = 1
425429 try :
426430 import GPUtil
431+
427432 return len (GPUtil .getGPUs ())
428433 except ImportError :
429434 return n_gpus
0 commit comments