Automatically remove G from memory commands in GPU queues

gridtk/script/jman.py

@@ -24,7 +24,26 @@ from ..tools import make_shell, logger
 from .. import local, sge
 from ..models import Status
 
-QUEUES = ['all.q', 'q1d', 'q1w', 'q1m', 'q1dm', 'q1wm', 'gpu', 'lgpu', 'sgpu', 'gpum']
+GPU_QUEUES = ['gpu', 'lgpu', 'sgpu', 'gpum']
+QUEUES = ['all.q', 'q1d', 'q1w', 'q1m', 'q1dm', 'q1wm'] + GPU_QUEUES
+
+
+def appropriate_for_gpu(args, kwargs):
+  # don't set these for GPU processing or the maximum virtual memory will be
+  # set on ulimit
+  kwargs.pop('memfree', None)
+  kwargs.pop('hvmem', None)
+
+  if args.memory is None:
+    return
+
+  # if this is a GPU queue and args.memory is provided, we set gpumem flag
+  # remove 'G' last character from the args.memory string
+  if args.memory.isdigit():
+    kwargs['gpumem'] = args.memory  # assign directly
+  elif args.memory.endswith('G'):
+    kwargs['gpumem'] = args.memory[:-1]  # remove G at the end
+
 
 def setup(args):
   """Returns the JobManager and sets up the basic infrastructure"""
@@ -128,14 +147,8 @@ def submit(args):
   if args.log_dir is not None:       kwargs['log_dir'] = args.log_dir
   if args.dependencies is not None:  kwargs['dependencies'] = args.dependencies
   if args.qname != 'all.q':          kwargs['hvmem'] = args.memory
-  # if this is a GPU queue and args.memory is provided, we set gpumem flag
-  # remove 'G' last character from the args.memory string
-  if args.qname in ('gpu', 'lgpu', 'sgpu', 'gpum') and args.memory is not None:
-    kwargs['gpumem'] = args.memory
-    # don't set these for GPU processing or the maximum virtual memroy will be
-    # set on ulimit
-    kwargs.pop('memfree', None)
-    kwargs.pop('hvmem', None)
+  if args.qname in GPU_QUEUES:
+    appropriate_for_gpu(args, kwargs)
   if args.parallel is not None:
     kwargs['pe_opt'] = "pe_mth %d" % args.parallel
     if args.memory is not None:
@@ -168,14 +181,8 @@ def resubmit(args):
     kwargs['memfree'] = args.memory
     if args.qname not in (None, 'all.q'):
       kwargs['hvmem'] = args.memory
-    # if this is a GPU queue and args.memory is provided, we set gpumem flag
-    # remove 'G' last character from the args.memory string
-    if args.qname in ('gpu', 'lgpu', 'sgpu', 'gpum') and args.memory is not None:
-      kwargs['gpumem'] = args.memory
-      # don't set these for GPU processing or the maximum virtual memroy will be
-      # set on ulimit
-      kwargs.pop('memfree', None)
-      kwargs.pop('hvmem', None)
+    if args.qname in GPU_QUEUES:
+      appropriate_for_gpu(args, kwargs)
   if args.parallel is not None:
     kwargs['pe_opt'] = "pe_mth %d" % args.parallel
     kwargs['memfree'] = get_memfree(args.memory, args.parallel)