Update docker_wrapper.py

Added gpu flag and update trigger to include docker create calls

scripts/docker_wrapper.py

@@ -49,7 +49,7 @@ def main():
     cmd = [DOCKER_BIN] + sys.argv[1:]
 
     # Are we handling 'docker run'?
-    if args.exe == 'create':
+    if args.exe in ['run', 'create']:
         # Yes, parse the options to 'docker run'
         parser = ArgumentParser(add_help=False)
         parser.add_argument('exe')
@@ -70,7 +70,11 @@ def main():
         parser.add_argument('--volume', action='append')
         parser.add_argument('--workdir', action='append')
         parser.add_argument('--user', action='append')
+        parser.add_argument('--cpu-shares', dest='cpu-shares', action='append')
+        parser.add_argument('--memory', action='append')
+        parser.add_argument('--cap-drop', dest='cap-drop', action='append')
         parser.add_argument('--group-add', dest='group-add', action='append')
+        parser.add_argument('--gpus', action='append')
 
         args, docker_opts = parser.parse_known_args()
 
@@ -83,6 +87,9 @@ def main():
                 if arg=='memory':
                     docker_opts += ['--memory-swap', val]
 
+        # Enable GPUS if CUDA_VISIBLE_DEVICES is set to a GPU but --gpus flag is not present
+        if [x for x in docker_opts if (('CUDA_VISIBLE_DEVICES=' in x) and not ('CUDA_VISIBLE_DEVICES=10000' in x))] and ('--gpus' not in docker_opts):
+            docker_opts += ['--gpus', 'all']
 
         # Did we successfully parse the docker run command?
         if args.image and args.command: