GpuResourceHandlerImpl implements ResourceHandler { @Override public synchronized List<PrivilegedOperation> postComplete( ContainerId containerId) throws ResourceHandlerException { gpuAllocator.unassignGpus(containerId); cGroupsHandler.deleteCGroup(CGroupsHandler.CGroupController.DEVICES, containerId.toString()); 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 testAllocationWhenDockerContainerEnabled() throws Exception { initializeGpus(); startContainerWithGpuRequestsDocker(1, 3); verifyDeniedDevices(getContainerId(1), Collections.emptyList()); boolean failedToAllocate = false; try { startContainerWithGpuRequestsDocker(2, 2); } catch (ResourceHandlerException e) { failedToAllocate = true; } assertTrue("Container allocation is expected to fail!", failedToAllocate); startContainerWithGpuRequestsDocker(3, 1); verifyDeniedDevices(getContainerId(3), Collections.emptyList()); startContainerWithGpuRequestsDocker(4, 0); verifyDeniedDevices(getContainerId(4), Collections.emptyList()); gpuResourceHandler.postComplete(getContainerId(1)); verifyCgroupsDeletedForContainer(1); verifyNumberOfAvailableGpus(3, gpuResourceHandler); gpuResourceHandler.postComplete(getContainerId(3)); verifyCgroupsDeletedForContainer(3); verifyNumberOfAvailableGpus(4, gpuResourceHandler); }
@Test public void testAllocation() throws Exception { initializeGpus(); startContainerWithGpuRequests(1, 3); verifyDeniedDevices(getContainerId(1), Collections.singletonList(new GpuDevice(3, 4))); boolean failedToAllocate = false; try { startContainerWithGpuRequests(2, 2); } catch (ResourceHandlerException e) { failedToAllocate = true; } assertTrue("Container allocation is expected to fail!", failedToAllocate); startContainerWithGpuRequests(3, 1); verifyDeniedDevices(getContainerId(3), Arrays.asList(new GpuDevice(0, 0), new GpuDevice(1, 1), new GpuDevice(2, 3))); startContainerWithGpuRequests(4, 0); verifyDeniedDevices(getContainerId(4), Arrays.asList(new GpuDevice(0, 0), new GpuDevice(1, 1), new GpuDevice(2, 3), new GpuDevice(3, 4))); gpuResourceHandler.postComplete(getContainerId(1)); verifyCgroupsDeletedForContainer(1); verifyNumberOfAvailableGpus(3, gpuResourceHandler); gpuResourceHandler.postComplete(getContainerId(3)); verifyCgroupsDeletedForContainer(3); verifyNumberOfAvailableGpus(4, gpuResourceHandler); }