Add support for enabling GPU access

Signed-off-by: YISH <[email protected]>

.gitignore

@@ -105,3 +105,6 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+
+.vscode

examples/nvidia-smi/docker-compose.yaml

@@ -0,0 +1,11 @@
+services:
+ test:
+ image: nvidia/cuda:12.3.1-base-ubuntu20.04
+ command: nvidia-smi
+ deploy:
+ resources:
+ reservations:
+ devices:
+ - driver: nvidia
+ count: 1
+ capabilities: [gpu]

podman_compose.py

@@ -635,6 +635,62 @@ def get_secret_args(compose, cnt, secret, podman_is_building=False):
 
 
 def container_to_res_args(cnt, podman_args):
+ container_to_cpu_res_args(cnt, podman_args)
+ container_to_gpu_res_args(cnt, podman_args)
+
+
+def container_to_gpu_res_args(cnt, podman_args):
+ # https://docs.docker.com/compose/gpu-support/
+ # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/cdi-support.html
+
+ deploy = cnt.get("deploy", None) or {}
+ res = deploy.get("resources", None) or {}
+ reservations = res.get("reservations", None) or {}
+ devices = reservations.get("devices", [])
+ gpu_on = False
+ for device in devices:
+ driver = device.get("driver", None)
+ if driver is None:
+ continue
+
+ capabilities = device.get("capabilities", None)
+ if capabilities is None:
+ continue
+
+ if driver != "nvidia" or "gpu" not in capabilities:
+ continue
+
+ count = device.get("count", "all")
+ device_ids = device.get("device_ids", "all")
+ if device_ids != "all" and len(device_ids) > 0:
+ for device_id in device_ids:
+ podman_args.extend((
+ "--device",
+ f"nvidia.com/gpu={device_id}",
+ ))
+ gpu_on = True
+ continue
+
+ if count != "all":
+ for device_id in range(count):
+ podman_args.extend((
+ "--device",
+ f"nvidia.com/gpu={device_id}",
+ ))
+ gpu_on = True
+ continue
+
+ podman_args.extend((
+ "--device",
+ "nvidia.com/gpu=all",
+ ))
+ gpu_on = True
+
+ if gpu_on:
+ podman_args.append("--security-opt=label=disable")
+
+
+def container_to_cpu_res_args(cnt, podman_args):
  # v2: https://docs.docker.com/compose/compose-file/compose-file-v2/#cpu-and-other-resources
  # cpus, cpu_shares, mem_limit, mem_reservation
  cpus_limit_v2 = try_float(cnt.get("cpus", None), None)

pytests/test_container_to_args.py

@@ -325,3 +325,106 @@ async def test_env_file_obj_optional(self):
  "busybox",
  ],
  )
+
+ async def test_gpu(self):
+ c = create_compose_mock()
+
+ cnt = get_minimal_container()
+ cnt["command"] = ["nvidia-smi"]
+ cnt["deploy"] = {"resources": {"reservations": {"devices": [{}]}}}
+
+ # count: all
+ cnt["deploy"]["resources"]["reservations"]["devices"][0] = {
+ "driver": "nvidia",
+ "count": "all",
+ "capabilities": ["gpu"],
+ }
+
+ args = await container_to_args(c, cnt)
+ self.assertEqual(
+ args,
+ [
+ "--name=project_name_service_name1",
+ "-d",
+ "--network=bridge",
+ "--network-alias=service_name",
+ "--device",
+ "nvidia.com/gpu=all",
+ "--security-opt=label=disable",
+ "busybox",
+ "nvidia-smi",
+ ],
+ )
+
+ # count: 2
+ cnt["deploy"]["resources"]["reservations"]["devices"][0] = {
+ "driver": "nvidia",
+ "count": 2,
+ "capabilities": ["gpu"],
+ }
+
+ args = await container_to_args(c, cnt)
+ self.assertEqual(
+ args,
+ [
+ "--name=project_name_service_name1",
+ "-d",
+ "--network=bridge",
+ "--network-alias=service_name",
+ "--device",
+ "nvidia.com/gpu=0",
+ "--device",
+ "nvidia.com/gpu=1",
+ "--security-opt=label=disable",
+ "busybox",
+ "nvidia-smi",
+ ],
+ )
+
+ # device_ids: all
+ cnt["deploy"]["resources"]["reservations"]["devices"][0] = {
+ "driver": "nvidia",
+ "device_ids": "all",
+ "capabilities": ["gpu"],
+ }
+
+ args = await container_to_args(c, cnt)
+ self.assertEqual(
+ args,
+ [
+ "--name=project_name_service_name1",
+ "-d",
+ "--network=bridge",
+ "--network-alias=service_name",
+ "--device",
+ "nvidia.com/gpu=all",
+ "--security-opt=label=disable",
+ "busybox",
+ "nvidia-smi",
+ ],
+ )
+
+ # device_ids: 1,3
+ cnt["deploy"]["resources"]["reservations"]["devices"][0] = {
+ "driver": "nvidia",
+ "device_ids": [1, 3],
+ "capabilities": ["gpu"],
+ }
+
+ args = await container_to_args(c, cnt)
+ self.assertEqual(
+ args,
+ [
+ "--name=project_name_service_name1",
+ "-d",
+ "--network=bridge",
+ "--network-alias=service_name",
+ "--device",
+ "nvidia.com/gpu=1",
+ "--device",
+ "nvidia.com/gpu=3",
+ "--security-opt=label=disable",
+ "busybox",
+ "nvidia-smi",
+ ],
+ )