GpuResourceAllocator { public GpuAllocation assignGpus(Container container) throws ResourceHandlerException { GpuAllocation allocation = internalAssignGpus(container); int timeWaiting = 0; while (allocation == null) { if (timeWaiting >= waitPeriodForResource) { break; } try { LOG.info("Container : " + container.getContainerId() + " is waiting for free GPU devices."); Thread.sleep(WAIT_MS_PER_LOOP); timeWaiting += WAIT_MS_PER_LOOP; allocation = internalAssignGpus(container); } catch (InterruptedException e) { Thread.currentThread().interrupt(); LOG.warn("Interrupted while waiting for available GPU"); break; } } if(allocation == null) { String message = "Could not get valid GPU device for container '" + container.getContainerId() + "' as some other containers might not releasing GPUs."; LOG.warn(message); throw new ResourceHandlerException(message); } return allocation; } GpuResourceAllocator(Context ctx); @VisibleForTesting  GpuResourceAllocator(Context ctx, int waitPeriodForResource); synchronized void addGpu(GpuDevice gpuDevice); @VisibleForTesting synchronized int getAvailableGpus(); synchronized void recoverAssignedGpus(ContainerId containerId); static int getRequestedGpus(Resource requestedResource); GpuAllocation assignGpus(Container container); synchronized void unassignGpus(ContainerId containerId); @VisibleForTesting synchronized Map<GpuDevice, ContainerId> getDeviceAllocationMapping(); synchronized List<GpuDevice> getAllowedGpus(); synchronized List<AssignedGpuDevice> getAssignedGpus(); @Override String toString();  }
@Test public void testRequestZeroGpu() throws ResourceHandlerException { addGpus(new GpuDevice(1, 1)); Container container = createMockContainer(0, 5L); GpuAllocation allocation = testSubject.assignGpus(container); assertNoAllocation(allocation); }
@Test public void testRequestMoreThanAvailableGpu() throws ResourceHandlerException { addGpus(new GpuDevice(1, 1)); Container container = createMockContainer(2, 5L); exception.expect(ResourceHandlerException.class); exception.expectMessage("Failed to find enough GPUs"); testSubject.assignGpus(container); }
@Test public void testGpuGetsUnassignedWhenStateStoreThrowsException() throws ResourceHandlerException, IOException { doThrow(new IOException("Failed to save container mappings " + "to NM state store!")) .when(nmStateStore).storeAssignedResources(any(Container.class), anyString(), anyList()); createAndAddGpus(1); exception.expect(ResourceHandlerException.class); exception.expectMessage("Failed to save container mappings " + "to NM state store"); Container container = createMockContainer(1, 5L); testSubject.assignGpus(container); }