nvidia: let nvidia toolkit decide what devices to include

saltyorg · Jan 27, 2025 · bd235a0 · bd235a0
1 parent 764b2f4
commit bd235a0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/resources/tasks/docker/create_docker_container.yml b/resources/tasks/docker/create_docker_container.yml
@@ -37,7 +37,9 @@
     default_host_ip: ""
     device_read_bps: "{{ lookup('vars', _var_prefix + '_docker_device_read_bps', default=omit) }}"
     device_read_iops: "{{ lookup('vars', _var_prefix + '_docker_device_read_iops', default=omit) }}"
-    device_requests: "{{ lookup('vars', _var_prefix + '_docker_device_requests', default=omit) }}"
+    device_requests: "{{ nvidia_docker_device_requests
+                      if use_nvidia
+                      else lookup('vars', _var_prefix + '_docker_device_requests', default=omit) }}"
     device_write_bps: "{{ lookup('vars', _var_prefix + '_docker_device_write_bps', default=omit) }}"
     device_write_iops: "{{ lookup('vars', _var_prefix + '_docker_device_write_iops', default=omit) }}"
     devices: "{{ lookup('vars', _var_prefix + '_docker_devices', default=omit) }}"
@@ -46,9 +48,7 @@
     dns_servers: "{{ lookup('vars', _var_prefix + '_docker_dns_servers', default=omit) }}"
     domainname: "{{ lookup('vars', _var_prefix + '_docker_domainname', default=omit) }}"
     entrypoint: "{{ lookup('vars', _var_prefix + '_docker_entrypoint', default=omit) }}"
-    env: "{{ (nvidia_docker_env if use_nvidia else {})
-             | combine(lookup('vars', _var_prefix + '_docker_envs', default={}))
-             | default(omit) }}"
+    env: "{{ lookup('vars', _var_prefix + '_docker_envs', default=omit) }}"
     env_file: "{{ lookup('vars', _var_prefix + '_docker_env_file', default=omit) }}"
     etc_hosts: "{{ lookup('vars', _var_prefix + '_docker_hosts', default=omit)
                 if not ('container:' in lookup('vars', _var_prefix + '_docker_network_mode', default=docker_networks_name_common))

diff --git a/resources/tasks/docker/set_docker_devices_variable.yml b/resources/tasks/docker/set_docker_devices_variable.yml
@@ -21,5 +21,5 @@
 
     - name: Resources | Tasks | Docker | Set Docker Devices Variable | Set 'docker_devices' variable # noqa var-naming[no-jinja]
       ansible.builtin.set_fact:
-        "{{ docker_devices_var_name }}": "{{ (['/dev/dri:/dev/dri'] if gpu.intel or use_nvidia else [])
+        "{{ docker_devices_var_name }}": "{{ (['/dev/dri:/dev/dri'] if gpu.intel else [])
                                              + lookup('vars', role_name + '_docker_devices_default', default=[]) }}"
diff --git a/roles/nvidia/defaults/main.yml b/roles/nvidia/defaults/main.yml
@@ -86,6 +86,14 @@ nvidia_nvtop_apt_repo_file: "/etc/apt/sources.list.d/flexiondotorg.list"
 # Docker
 ################################
 
+nvidia_docker_device_requests:
+  - driver: nvidia
+    count: -1
+    capabilities:
+      - - gpu
+        - compute
+        - utility
+
 nvidia_docker_env:
   NVIDIA_DRIVER_CAPABILITIES: "compute,video,utility"
   NVIDIA_VISIBLE_DEVICES: "all"