GpuResourceHandlerImpl implements ResourceHandler { @Override public synchronized List<PrivilegedOperation> preStart(Container container) throws ResourceHandlerException { String containerIdStr = container.getContainerId().toString(); GpuResourceAllocator.GpuAllocation allocation = gpuAllocator.assignGpus( container); cGroupsHandler.createCGroup(CGroupsHandler.CGroupController.DEVICES, containerIdStr); if (!OCIContainerRuntime.isOCICompliantContainerRequested( nmContext.getConf(), container.getLaunchContext().getEnvironment())) { try { PrivilegedOperation privilegedOperation = new PrivilegedOperation( PrivilegedOperation.OperationType.GPU, Arrays.asList(CONTAINER_ID_CLI_OPTION, containerIdStr)); if (!allocation.getDeniedGPUs().isEmpty()) { List<Integer> minorNumbers = new ArrayList<>(); for (GpuDevice deniedGpu : allocation.getDeniedGPUs()) { minorNumbers.add(deniedGpu.getMinorNumber()); } privilegedOperation.appendArgs(Arrays.asList(EXCLUDED_GPUS_CLI_OPTION, StringUtils.join(",", minorNumbers))); } privilegedOperationExecutor.executePrivilegedOperation( privilegedOperation, true); } catch (PrivilegedOperationException e) { cGroupsHandler.deleteCGroup(CGroupsHandler.CGroupController.DEVICES, containerIdStr); LOG.warn("Could not update cgroup for container", e); throw new ResourceHandlerException(e); } List<PrivilegedOperation> ret = new ArrayList<>(); ret.add(new PrivilegedOperation( PrivilegedOperation.OperationType.ADD_PID_TO_CGROUP, PrivilegedOperation.CGROUP_ARG_PREFIX + cGroupsHandler .getPathForCGroupTasks(CGroupsHandler.CGroupController.DEVICES, containerIdStr))); return ret; } return null; } GpuResourceHandlerImpl(Context nmContext,
      CGroupsHandler cGroupsHandler,
      PrivilegedOperationExecutor privilegedOperationExecutor,
      GpuDiscoverer gpuDiscoverer); @Override List<PrivilegedOperation> bootstrap(Configuration configuration); @Override synchronized List<PrivilegedOperation> preStart(Container container); GpuResourceAllocator getGpuAllocator(); @Override List<PrivilegedOperation> reacquireContainer(ContainerId containerId); @Override List<PrivilegedOperation> updateContainer(Container container); @Override synchronized List<PrivilegedOperation> postComplete(
      ContainerId containerId); @Override List<PrivilegedOperation> teardown(); @Override String toString(); static final String EXCLUDED_GPUS_CLI_OPTION; static final String CONTAINER_ID_CLI_OPTION; }
@Test public void testAssignedGpuWillBeCleanedUpWhenStoreOpFails() throws Exception { initializeGpus(); doThrow(new IOException("Exception ...")).when(mockNMStateStore) .storeAssignedResources( any(Container.class), anyString(), anyList()); boolean exception = false; try { gpuResourceHandler.preStart(mockContainerWithGpuRequest(1, createResourceRequest(3))); } catch (ResourceHandlerException e) { exception = true; } assertTrue("preStart should throw exception", exception); verifyNumberOfAvailableGpus(4, gpuResourceHandler); }
@Test public void testAllocationStored() throws Exception { initializeGpus(); Container container = mockContainerWithGpuRequest(1, createResourceRequest(3)); gpuResourceHandler.preStart(container); verify(mockNMStateStore).storeAssignedResources(container, ResourceInformation.GPU_URI, Arrays .asList(new GpuDevice(0, 0), new GpuDevice(1, 1), new GpuDevice(2, 3))); verifyDeniedDevices(getContainerId(1), Collections.singletonList(new GpuDevice(3, 4))); container = mockContainerWithGpuRequest(2, createResourceRequest(0)); gpuResourceHandler.preStart(container); verifyDeniedDevices(getContainerId(2), Arrays .asList(new GpuDevice(0, 0), new GpuDevice(1, 1), new GpuDevice(2, 3), new GpuDevice(3, 4))); assertEquals("Number of GPU device allocations is not the expected!", 0, container.getResourceMappings() .getAssignedResources(ResourceInformation.GPU_URI).size()); verify(mockNMStateStore, never()).storeAssignedResources( eq(container), eq(ResourceInformation.GPU_URI), anyList()); }