fix: change to N gpus

worker/utils/gpuinfo.py

@@ -14,7 +14,7 @@ def __init__(self):
         # Look out for device env var hack
         self.forced_gpu = os.getenv("CUDA_VISIBLE_DEVICES", None) is not None
         self.device = int(os.getenv("CUDA_VISIBLE_DEVICES", 0))
-        self.ui_show_one_gpu = os.getenv("AIWORKER_UI_SHOW_ONE_GPU", None) is not None
+        self.ui_show_n_gpus = int(os.getenv("AIWORKER_UI_SHOW_N_GPUS", 1))
 
     # Return a value from the given dictionary supporting dot notation
     def get(self, data, key, default=""):
@@ -36,9 +36,12 @@ def get(self, data, key, default=""):
 
     def get_num_gpus(self):
         """How many GPUs in this system?"""
-        if self.forced_gpu or self.ui_show_one_gpu:
+        if self.forced_gpu:
             return 1
 
+        if self.ui_show_n_gpus:
+            return self.ui_show_n_gpus
+
         with contextlib.suppress(Exception):
             nvsmi = nvidia_smi.getInstance()
             data = nvsmi.DeviceQuery()