From cc38f438a156443f7e077b7520aed559a48bfde8 Mon Sep 17 00:00:00 2001
From: Michael Dawson <mdawson@devrus.com>
Date: Wed, 11 Dec 2024 05:30:28 -0500
Subject: [PATCH] feat: force use of libPod if nvidia device requested (#10251)

Requesting an nvidia device through the container toolkit
is only supported through the libPod API not the compatibility
API. Recognize the request and for libPod in a similar manner
to what's done when the container is associated with a pod.

Signed-off-by: Michael Dawson <mdawson@devrus.com>
---
 .../src/plugin/container-registry.spec.ts     | 98 ++++++++++++++++++-
 .../main/src/plugin/container-registry.ts     | 13 ++-
 2 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/packages/main/src/plugin/container-registry.spec.ts b/packages/main/src/plugin/container-registry.spec.ts
index 2ef851653..4dae3d9c2 100644
--- a/packages/main/src/plugin/container-registry.spec.ts
+++ b/packages/main/src/plugin/container-registry.spec.ts
@@ -4098,7 +4098,7 @@ describe('createContainerLibPod', () => {
       HostConfig: {
         Devices: [
           {
-            PathOnHost: 'nvidia.com/gpu=all',
+            PathOnHost: 'device1',
             PathInContainer: '',
             CgroupPermissions: '',
           },
@@ -4150,7 +4150,7 @@ describe('createContainerLibPod', () => {
     const expectedOptions: PodmanContainerCreateOptions = {
       name: options.name,
       command: options.Cmd,
-      devices: [{ path: 'nvidia.com/gpu=all' }],
+      devices: [{ path: 'device1' }],
       entrypoint: [options.Entrypoint as string],
       env: {
         key: 'value',
@@ -4228,6 +4228,100 @@ describe('createContainerLibPod', () => {
     await containerRegistry.createContainer('podman1', options);
     expect(createPodmanContainerMock).toBeCalledWith(expectedOptions);
   });
+
+  test('check that use of libPod is forced by request for nvidia device', async () => {
+    const dockerAPI = new Dockerode({ protocol: 'http', host: 'localhost' });
+
+    const libpod = new LibpodDockerode();
+    libpod.enhancePrototypeWithLibPod();
+    containerRegistry.addInternalProvider('podman1', {
+      name: 'podman',
+      id: 'podman1',
+      api: dockerAPI,
+      libpodApi: dockerAPI,
+      connection: {
+        type: 'podman',
+      },
+    } as unknown as InternalContainerProvider);
+
+    const createPodmanContainerMock = vi
+      .spyOn(dockerAPI as unknown as LibPod, 'createPodmanContainer')
+      .mockImplementation(_options =>
+        Promise.resolve({
+          Id: 'id',
+          Warnings: [],
+        }),
+      );
+    vi.spyOn(dockerAPI as unknown as Dockerode, 'getContainer').mockImplementation((_id: string) => {
+      return {
+        start: () => {},
+      } as unknown as Dockerode.Container;
+    });
+    // use minimum set as the full of options is validated in the previous test
+    const options: ContainerCreateOptions = {
+      Image: 'image',
+      name: 'name',
+      HostConfig: {
+        Devices: [
+          {
+            PathOnHost: 'nvidia.com/gpu=all',
+            PathInContainer: '',
+            CgroupPermissions: '',
+          },
+        ],
+        NetworkMode: 'mode',
+        AutoRemove: true,
+      },
+      Cmd: ['cmd'],
+      Entrypoint: 'entrypoint',
+      User: 'user',
+    };
+    const expectedOptions: PodmanContainerCreateOptions = {
+      image: options.Image,
+      name: options.name,
+      devices: [{ path: 'nvidia.com/gpu=all' }],
+      netns: {
+        nsmode: 'mode',
+      },
+      command: options.Cmd,
+      entrypoint: [options.Entrypoint as string],
+      user: options.User,
+      cap_add: undefined,
+      cap_drop: undefined,
+      dns_server: undefined,
+      env: undefined,
+      healthconfig: undefined,
+      hostadd: undefined,
+      hostname: undefined,
+      labels: undefined,
+      mounts: undefined,
+      pod: undefined,
+      portmappings: undefined,
+      privileged: undefined,
+      read_only_filesystem: undefined,
+      remove: true,
+      restart_policy: undefined,
+      restart_tries: undefined,
+      seccomp_policy: undefined,
+      seccomp_profile_path: undefined,
+      selinux_opts: [],
+      stop_timeout: undefined,
+      userns: undefined,
+      work_dir: undefined,
+    };
+    vi.spyOn(containerRegistry, 'attachToContainer').mockImplementation(
+      (
+        _engine: InternalContainerProvider,
+        _container: Dockerode.Container,
+        _hasTty?: boolean,
+        _openStdin?: boolean,
+      ) => {
+        return Promise.resolve();
+      },
+    );
+    await containerRegistry.createContainer('podman1', options);
+    expect(createPodmanContainerMock).toBeCalledWith(expectedOptions);
+  });
 });
 
 describe('getContainerCreateMountOptionFromBind', () => {
diff --git a/packages/main/src/plugin/container-registry.ts b/packages/main/src/plugin/container-registry.ts
index d12b58055..38e1c7e8f 100644
--- a/packages/main/src/plugin/container-registry.ts
+++ b/packages/main/src/plugin/container-registry.ts
@@ -1941,7 +1941,18 @@ export class ContainerProviderRegistry {
     let telemetryOptions = {};
     try {
       let container: Dockerode.Container;
-      if (options.pod) {
+      let forceLibPod = false;
+
+      // the device option requesting an nvidia gpu on linux only works
+      // if the LibPod API is used. Check if such a device is requested
+      // and if so force the use of LibPod
+      for (const device of options.HostConfig?.Devices ?? []) {
+        if (device.PathOnHost === 'nvidia.com/gpu=all') {
+          forceLibPod = true;
+          break;
+        }
+      }
+      if (options.pod ?? forceLibPod) {
         container = await this.createContainerLibPod(engineId, options);
       } else {
         container = await this.createContainerDockerode(engineId, options);