NvidiaDockerV2CommandPlugin implements DockerCommandPlugin { @Override public synchronized void updateDockerRunCommand( DockerRunCommand dockerRunCommand, Container container) throws ContainerExecutionException { if (!requestsGpu(container)) { return; } Set<GpuDevice> assignedResources = getAssignedGpus(container); if (assignedResources == null || assignedResources.isEmpty()) { return; } Map<String, String> environment = new HashMap<>(); String gpuIndexList = ""; for (GpuDevice gpuDevice : assignedResources) { gpuIndexList = gpuIndexList + gpuDevice.getIndex() + ","; LOG.info("nvidia docker2 assigned gpu index: " + gpuDevice.getIndex()); } dockerRunCommand.addRuntime(nvidiaRuntime); environment.put(nvidiaVisibleDevices, gpuIndexList.substring(0, gpuIndexList.length() - 1)); dockerRunCommand.addEnv(environment); } NvidiaDockerV2CommandPlugin(); @Override synchronized void updateDockerRunCommand(
      DockerRunCommand dockerRunCommand, Container container); @Override DockerVolumeCommand getCreateDockerVolumeCommand(Container container); @Override DockerVolumeCommand getCleanupDockerVolumesCommand(Container container);  }
@Test public void testPlugin() throws Exception { DockerRunCommand runCommand = new DockerRunCommand("container_1", "user", "fakeimage"); Map<String, List<String>> originalCommandline = copyCommandLine( runCommand.getDockerCommandWithArguments()); MyNvidiaDockerV2CommandPlugin commandPlugin = new MyNvidiaDockerV2CommandPlugin(); Container nmContainer = mock(Container.class); commandPlugin.updateDockerRunCommand(runCommand, nmContainer); Assert.assertTrue(commandlinesEquals(originalCommandline, runCommand.getDockerCommandWithArguments())); ResourceMappings resourceMappings = new ResourceMappings(); when(nmContainer.getResourceMappings()).thenReturn(resourceMappings); commandPlugin.updateDockerRunCommand(runCommand, nmContainer); Assert.assertTrue(commandlinesEquals(originalCommandline, runCommand.getDockerCommandWithArguments())); ResourceMappings.AssignedResources assigned = new ResourceMappings.AssignedResources(); assigned.updateAssignedResources( ImmutableList.of(new GpuDevice(0, 0), new GpuDevice(1, 1))); resourceMappings.addAssignedResources(ResourceInformation.GPU_URI, assigned); commandPlugin.setRequestsGpu(true); commandPlugin.updateDockerRunCommand(runCommand, nmContainer); Map<String, List<String>> newCommandLine = runCommand.getDockerCommandWithArguments(); Assert.assertFalse(commandlinesEquals(originalCommandline, newCommandLine)); Assert.assertTrue( runCommand.getEnv().get("NVIDIA_VISIBLE_DEVICES").equals("0,1")); Assert.assertTrue(newCommandLine.containsKey("runtime")); }