GpuResourceHandlerImpl implements ResourceHandler { @Override public List<PrivilegedOperation> bootstrap(Configuration configuration) throws ResourceHandlerException { List<GpuDevice> usableGpus; try { usableGpus = gpuDiscoverer.getGpusUsableByYarn(); if (usableGpus == null || usableGpus.isEmpty()) { String message = "GPU is enabled on the NodeManager, but couldn't find " + "any usable GPU devices, please double check configuration!"; LOG.error(message); throwIfNecessary(new ResourceHandlerException(message), configuration); } } catch (YarnException e) { LOG.error("Exception when trying to get usable GPU device", e); throw new ResourceHandlerException(e); } for (GpuDevice gpu : usableGpus) { gpuAllocator.addGpu(gpu); } this.cGroupsHandler.initializeCGroupController( CGroupsHandler.CGroupController.DEVICES); 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 testBootstrapWithMockGpuDiscoverer() throws Exception { GpuDiscoverer mockDiscoverer = mock(GpuDiscoverer.class); Configuration conf = new YarnConfiguration(); mockDiscoverer.initialize(conf, nvidiaBinaryHelper); expected.expect(ResourceHandlerException.class); gpuResourceHandler.bootstrap(conf); }