[GitHub] [madlib] reductionista commented on a change in pull request #379: DL: Replace use_gpu flag with gpus_per_host integer

[GitBox <gi...@apache.org>]

reductionista commented on a change in pull request #379:  DL: Replace use_gpu flag with gpus_per_host integer 
URL: https://github.com/apache/madlib/pull/379#discussion_r279906423
 
 

 ##########
 File path: src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
 ##########
 @@ -35,27 +35,45 @@ import madlib_keras_serializer
 from utilities.utilities import _assert
 from utilities.utilities import is_platform_pg
 
+CUDA_VISIBLE_DEVICES_KEY = 'CUDA_VISIBLE_DEVICES'
 #######################################################################
 ########### Keras specific functions #####
 #######################################################################
-def get_device_name_and_set_cuda_env(use_gpu, seg):
-    gpus_per_host = 4
-    if use_gpu:
+
+def set_cuda_env(value):
+    """
+    :param value: -1 to disable gpu
+    :return:
+    """
+    os.environ[CUDA_VISIBLE_DEVICES_KEY] = value
+
+def reset_cuda_env(value):
+    """
+    This function will reset the cuda env variable. This should only be called
+    if set_cuda_env was called previously.
+    :param value:
+    """
+    if value:
+        os.environ[CUDA_VISIBLE_DEVICES_KEY] = value
+    else:
+        del os.environ[CUDA_VISIBLE_DEVICES_KEY]
+
 
 Review comment:
   If we're going to have both set and reset functionality in one function, why have a second function just for duplicating half of the behavior?
   
   I'd suggest deleting the current `set_cuda_env` function, renaming `reset_cuda_env` to `set_cuda_env`, and making `value` an optional parameter.  Then you can use `set_cuda_env(value)` to set it to something and `set_cuda_env()` to clear it.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services