- Add information about GPUs

- Memory limit has to be given in GB now

resources/mlhubspawner/mlhubspawner/mlhubkubernetesspawner.py

@@ -79,7 +79,7 @@ def start(self):
             self.cpu_limit = float(self.user_options.get('cpu_limit'))
 
         if self.user_options.get('mem_limit'):
-            memory = str(self.user_options.get('mem_limit'))
+            memory = str(self.user_options.get('mem_limit')) + "G"
             self.mem_limit = memory.upper().replace("GB", "G").replace("KB", "K").replace("MB", "M").replace("TB", "T")
 
         #if self.user_options.get('is_mount_volume') == 'on':

resources/mlhubspawner/mlhubspawner/mlhubspawner.py

@@ -13,6 +13,7 @@
 from docker.utils import kwargs_from_env
 
 import os
+import subprocess
 import socket
 import ipaddress
 from traitlets import default, Unicode, List
@@ -74,7 +75,8 @@ def __init__(self, *args, **kwargs):
         # Get available resource information
         self.resource_information = {
             "cpu_count": psutil.cpu_count(),
-            "memory_count_in_gb": round(psutil.virtual_memory().total/1024/1024/1024, 1)
+            "memory_count_in_gb": round(psutil.virtual_memory().total/1024/1024/1024, 1),
+            "gpu_count": self.get_gpu_info()
         }
 
     @property
@@ -161,8 +163,8 @@ def start(self) -> (str, int):
             nano_cpus = int(limited_cpus * 1e9)
             extra_host_config['nano_cpus'] = nano_cpus
         if self.user_options.get('mem_limit'):
-            extra_host_config['mem_limit'] = self.user_options.get(
-                'mem_limit')
+            extra_host_config['mem_limit'] = str(self.user_options.get(
+                'mem_limit')) + "gb"
 
         if self.user_options.get('is_mount_volume') == 'on':
             # {username} and {servername} will be automatically replaced by DockerSpawner with the right values as in template_namespace
@@ -316,3 +318,27 @@ def template_namespace(self):
             template["servername"] = "-" + template["servername"]
 
         return template
+
+    def get_gpu_info(self) -> list:
+        count_gpu = 0
+        try:
+            sp = subprocess.Popen(['nvidia-smi', '-q'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            out_str = sp.communicate()
+            out_list = out_str[0].decode("utf-8").split('\n')
+
+            # out_dict = {}
+
+            for item in out_list:
+                try:
+                    key, val = item.split(':')
+                    key, val = key.strip(), val.strip()
+                    if key == 'Product Name':
+                        count_gpu += 1
+                        # gpus.append(val)
+                    #out_dict[key + "_" + str(count_gpu)] = val
+                except:
+                    pass
+        except:
+            pass
+
+        return count_gpu

resources/mlhubspawner/mlhubspawner/spawner_options.py

@@ -8,21 +8,24 @@
 additional_info_style="margin-top: 4px; color: rgb(165,165,165); font-size: 12px;"
 optional_label = "<span style=\"font-size: 12px; font-weight: 400;\">(optional)</span>"
 
-def get_options_form(spawner, additional_cpu_info="", additional_memory_info="") -> str:
+def get_options_form(spawner, additional_cpu_info="", additional_memory_info="", additional_gpu_info="") -> str:
     """Return the spawner options screen"""
 
     # Only show spawner options for named servers (the default server should start with default values)
     if getattr(spawner, "name", "") == "":
         return ''
 
-    description_memory_limit = 'Minimum limit must be 4mb as required by Docker.'
+    description_memory_limit = 'Memory Limit in GB.'
     description_env = 'One name=value pair per line, without quotes'
     description_days_to_live = 'Number of days the container should live'
 
     default_image = getattr(spawner, "image", "mltooling/ml-workspace:latest")
 
     # Show / hide custom image input field when checkbox is clicked
     custom_image_listener = "if(event.target.checked){ $('#image-name').css('display', 'block'); $('.defined-images').css('display', 'none'); }else{ $('#image-name').css('display', 'none'); $('.defined-images').css('display', 'block'); }"
+
+    # Indicate a wrong input value (not a number) by changing the color to red
+    memory_input_listener = "if(isNaN(event.srcElement.value)){ $('#mem-limit').css('color', 'red'); }else{ $('#mem-limit').css('color', 'black'); }"
 
     # Create drop down menu with pre-defined custom images
     image_option_template = """
@@ -55,16 +58,15 @@ def get_options_form(spawner, additional_cpu_info="", additional_memory_info="")
             <div style="{additional_info_style}">{additional_cpu_info}</div>
         </div>
         <div style="{div_style}">
-            <label style="{label_style}" for="mem_limit" title="{description_memory_limit}">Memory Limit {optional_label}</label>
-            <input style="{input_style}" name="mem_limit" title="{description_memory_limit}" placeholder="e.g. 100mb, 8g, ..."></input>
+            <label style="{label_style}" for="mem_limit" title="{description_memory_limit}">Memory Limit in GB {optional_label}</label>
+            <input style="{input_style}" name="mem_limit" id="mem-limit" title="{description_memory_limit}" placeholder="e.g. 1, 2, 15, ..." oninput="{memory_input_listener}"></input>
             <div style="{additional_info_style}">{additional_memory_info}</div>
         </div>
         <div style="{div_style}">
             <label style="{label_style}" for="env" title="{description_env}">Environment Variables {optional_label}</label>
             <textarea style="{input_style}" name="env" title="{description_env}" placeholder="NAME=VALUE"></textarea>
             <div style="{additional_info_style}">{description_env}</div>
         </div>
-        
         <div style="{div_style}">
             <label style="{label_style}" for="days_to_live" title="{description_days_to_live}">Days to live {optional_label}</label>
             <input style="{input_style}" name="days_to_live" title="{description_days_to_live}" placeholder="e.g. 3"></input>
@@ -79,17 +81,20 @@ def get_options_form(spawner, additional_cpu_info="", additional_memory_info="")
         custom_image_listener=custom_image_listener,
         optional_label=optional_label,
         description_memory_limit=description_memory_limit,
+        memory_input_listener=memory_input_listener,
         description_env=description_env,
         description_days_to_live=description_days_to_live,
         additional_cpu_info=additional_cpu_info,
-        additional_memory_info=additional_memory_info
+        additional_memory_info=additional_memory_info,
     )
 
 def get_options_form_docker(spawner):
-    print(spawner.resource_information)
+    description_gpus = 'Empty for no GPU, "all" for all GPUs, or a comma-separated list of indices of the GPUs (e.g 0,2).'
     additional_info = {
         "additional_cpu_info": "Number between 1 and {cpu_count}".format(cpu_count=spawner.resource_information['cpu_count']),
-        "additional_memory_info": "Number between 200mb and {memory_count_in_gb}gb".format(memory_count_in_gb=spawner.resource_information['memory_count_in_gb'])
+        "additional_memory_info": "Number between 1 and {memory_count_in_gb}".format(memory_count_in_gb=spawner.resource_information['memory_count_in_gb']),
+        "additional_gpu_info": "<div>Available GPUs: {gpu_count}</div><div>{description_gpus}</div>".format(gpu_count=spawner.resource_information['gpu_count'], description_gpus=description_gpus)
+
     }
     options_form = get_options_form(spawner, **additional_info)
 
@@ -98,13 +103,16 @@ def get_options_form_docker(spawner):
     # reminding the user of inserting a GPU-leveraging docker image
     show_gpu_info_box = "$('#gpu-info-box').css('display', 'block');"
     hide_gpu_info_box = "$('#gpu-info-box').css('display', 'none');"
-    description_gpus = 'Empty for no GPU-acess. A comma-separted list of numbers describe the indices of the accessible GPUs.'
     gpu_input_listener = "if(event.srcElement.value !== ''){{ {show_gpu_info_box} }}else{{ {hide_gpu_info_box} }}" \
         .format(
             show_gpu_info_box=show_gpu_info_box, 
             hide_gpu_info_box=hide_gpu_info_box
     )
-
+
+    gpu_disabled = ""
+    if spawner.resource_information['gpu_count'] < 1:
+        gpu_disabled = "disabled"
+
     options_form_docker = \
     """
     <div style="{div_style}">
@@ -113,16 +121,20 @@ def get_options_form_docker(spawner):
     </div>
     <div style="{div_style}">
         <label style="{label_style}" for="gpus" title="{description_gpus}">GPUs {optional_label}</label>
-        <input style="{input_style}" name="gpus" title="{description_gpus}" placeholder="e.g. all, 0, 1, 2, ..." oninput="{gpu_input_listener}"></input>
+        <input style="{input_style}" name="gpus" title="{description_gpus}" placeholder="e.g. all, 0, 1, 2, ..." oninput="{gpu_input_listener} {gpu_disabled}"></input>
+        <div style="{additional_info_style}">{additional_gpu_info}</div>
         <div style="background-color: #ffa000; padding: 8px; margin-top: 4px; display: none; {input_style}" id="gpu-info-box">When using GPUs, make sure to use a GPU-supporting Docker image!</div>
     </div>
     """.format(
         div_style=div_style,
         label_style=label_style,
         input_style=input_style,
+        additional_info_style=additional_info_style,
         optional_label=optional_label,
         gpu_input_listener=gpu_input_listener,
-        description_gpus=description_gpus
+        gpu_disabled=gpu_disabled,
+        description_gpus=description_gpus,
+        additional_gpu_info=additional_info['additional_gpu_info']
     )
 
     return options_form + options_form_docker