Merge pull request #36 from truenas/NAS-130121

NAS-130121 / 24.10 / Render GPU feature
truenas · Jul 24, 2024 · 081f8b2 · 081f8b2
2 parents 378e3e7 + eed43d1
commit 081f8b2
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 13 deletions.
diff --git a/catalog_reader/app_utils.py b/catalog_reader/app_utils.py
@@ -31,13 +31,13 @@ def get_app_details_base(retrieve_complete_item_keys: bool = True) -> dict:
 
 def get_default_questions_context() -> dict:
     return {
-        'gpus': {},
         'timezones': {'Asia/Saigon': 'Asia/Saigon', 'Asia/Damascus': 'Asia/Damascus'},
         'ip_choices': {'192.168.0.10': '192.168.0.10', '0.0.0.0': '0.0.0.0'},
         'certificates': [],
         'certificate_authorities': [],
         'system.general.config': {'timezone': 'America/Los_Angeles'},
         'unused_ports': [i for i in range(1025, 65535)],
+        'gpu_choices': {}
     }
 
 

diff --git a/catalog_reader/questions.py b/catalog_reader/questions.py
@@ -35,23 +35,74 @@ def normalize_question(question: dict, version_data: dict, context: dict) -> Non
             data['enum'] = [
                 {'value': i, 'description': f'{i!r} Interface'} for i in context['nic_choices']
             ]
-        elif ref == 'definitions/gpuConfiguration':
+        elif ref == 'definitions/gpu_configuration':
+            # How we will go about this is the following:
+            # If user has only nvidia gpus and we are able to retrieve nvidia gpu's uuid
+            # we will allow user to select any of the available nvidia gpu
+            # If user has any other GPU apart from nvidia vendor, he will have to passthrough all available gpus
+            # If user has nvidia + other gpu's, then we can show a boolean or a string enum
+            # highlighting the nvidia bits
+            gpu_choices = [entry for entry in context['gpu_choices'] if entry['gpu_details']['available_to_host']]
+            show_all_gpus_flag = any(
+                g['vendor'] != 'NVIDIA' or not g['vendor_specific_config'].get('uuid') for g in gpu_choices
+            )
+            show_nvidia_selection = any(
+                g['vendor'] == 'NVIDIA' and g['vendor_specific_config'].get('uuid') for g in gpu_choices
+            )
             data['attrs'] = [
                 {
-                    'variable': gpu,
-                    'label': f'GPU Resource ({gpu})',
-                    'description': 'Please enter the number of GPUs to allocate',
+                    'variable': 'use_all_gpus',
+                    'label': 'Passthrough available (non-NVIDIA) GPUs',
+                    'description': 'Please select this option to passthrough all (non-NVIDIA) GPUs to the app',
                     'schema': {
-                        'type': 'int',
-                        'max': int(quantity),
-                        'enum': [
-                            {'value': i, 'description': f'Allocate {i!r} {gpu} GPU'}
-                            for i in range(int(quantity) + 1)
-                        ],
-                        'default': 0,
+                        'type': 'boolean',
+                        'default': False,
+                        'hidden': not show_all_gpus_flag,
                     }
-                } for gpu, quantity in context['gpus'].items()
+                },
+                {
+                    'variable': 'nvidia_gpu_selection',
+                    'label': 'Select NVIDIA GPU(s)',
+                    'description': 'Please select the NVIDIA GPU(s) to passthrough to the app',
+                    'schema': {
+                        'type': 'dict',
+                        'additional_attrs': True,
+                        'hidden': not show_nvidia_selection,
+                        'attrs': [
+                            {
+                                'variable': gpu['gpu_details']['addr']['pci_slot'],
+                                'label': gpu['description'],
+                                'description': f'NVIDIA gpu {gpu["vendor_specific_config"]["uuid"]}',
+                                'schema': {
+                                    'type': 'dict',
+                                    'attrs': [
+                                        {
+                                            'variable': 'uuid',
+                                            'schema': {
+                                                'type': 'string',
+                                                'default': gpu['vendor_specific_config']['uuid'],
+                                                'hidden': True,
+                                            }
+                                        },
+                                        {
+                                            'variable': 'use_gpu',
+                                            'label': 'Use this GPU',
+                                            'description': 'Use this GPU for the app',
+                                            'schema': {
+                                                'type': 'boolean',
+                                                'default': False,
+                                            }
+                                        }
+                                    ],
+                                }
+                            }
+                            for gpu in (gpu_choices if show_nvidia_selection else [])
+                            if gpu['vendor'] == 'NVIDIA' and gpu['vendor_specific_config'].get('uuid')
+                        ]
+                    },
+                },
             ]
+
         elif ref == 'definitions/timezone':
             data.update({
                 'enum': [{'value': t, 'description': f'{t!r} timezone'} for t in sorted(context['timezones'])],