FpgaResourceHandlerImpl implements ResourceHandler { @Override public List<PrivilegedOperation> bootstrap(Configuration configuration) throws ResourceHandlerException { if (!vendorPlugin.initPlugin(configuration)) { throw new ResourceHandlerException("FPGA plugin initialization failed"); } LOG.info("FPGA Plugin bootstrap success."); List<FpgaDevice> fpgaDeviceList = fpgaDiscoverer.discover(); allocator.addFpgaDevices(vendorPlugin.getFpgaType(), fpgaDeviceList); this.cGroupsHandler.initializeCGroupController( CGroupsHandler.CGroupController.DEVICES); return null; } @VisibleForTesting  FpgaResourceHandlerImpl(Context nmContext,
      CGroupsHandler cGroupsHandler,
      PrivilegedOperationExecutor privilegedOperationExecutor,
      AbstractFpgaVendorPlugin plugin,
      FpgaDiscoverer fpgaDiscoverer); String getRequestedIPID(Container container); @Override List<PrivilegedOperation> bootstrap(Configuration configuration); @Override List<PrivilegedOperation> preStart(Container container); @Override List<PrivilegedOperation> reacquireContainer(ContainerId containerId); @Override List<PrivilegedOperation> updateContainer(Container container); @Override List<PrivilegedOperation> postComplete(ContainerId containerId); @Override List<PrivilegedOperation> teardown(); @Override String toString(); static final String EXCLUDED_FPGAS_CLI_OPTION; static final String CONTAINER_ID_CLI_OPTION; }
@Test public void testBootstrap() throws ResourceHandlerException { String allowed = "auto"; configuration.set(YarnConfiguration.NM_FPGA_ALLOWED_DEVICES, allowed); fpgaResourceHandler.bootstrap(configuration); verify(mockVendorPlugin, times(2)).initPlugin(configuration); verify(mockCGroupsHandler, times(1)).initializeCGroupController( CGroupsHandler.CGroupController.DEVICES); Assert.assertEquals(5, fpgaResourceHandler.getFpgaAllocator() .getAvailableFpgaCount()); Assert.assertEquals(5, fpgaResourceHandler.getFpgaAllocator() .getAllowedFpga().size()); fpgaResourceHandler = new FpgaResourceHandlerImpl(mockContext, mockCGroupsHandler, mockPrivilegedExecutor, mockVendorPlugin, fpgaDiscoverer); allowed = "0,1,2"; configuration.set(YarnConfiguration.NM_FPGA_ALLOWED_DEVICES, allowed); fpgaResourceHandler.bootstrap(configuration); Assert.assertEquals(3, fpgaResourceHandler.getFpgaAllocator().getAllowedFpga().size()); List<FpgaDevice> allowedDevices = fpgaResourceHandler.getFpgaAllocator().getAllowedFpga(); for (String s : allowed.split(",")) { boolean check = false; for (FpgaDevice device : allowedDevices) { if (String.valueOf(device.getMinor()).equals(s)) { check = true; } } Assert.assertTrue("Minor:" + s +" found", check); } Assert.assertEquals(3, fpgaResourceHandler.getFpgaAllocator().getAvailableFpgaCount()); fpgaResourceHandler = new FpgaResourceHandlerImpl(mockContext, mockCGroupsHandler, mockPrivilegedExecutor, mockVendorPlugin, fpgaDiscoverer); allowed = "0,1,7"; configuration.set(YarnConfiguration.NM_FPGA_ALLOWED_DEVICES, allowed); fpgaResourceHandler.bootstrap(configuration); Assert.assertEquals(2, fpgaResourceHandler.getFpgaAllocator().getAvailableFpgaCount()); Assert.assertEquals(2, fpgaResourceHandler.getFpgaAllocator().getAllowedFpga().size()); }
@Test public void testBootStrapWithEmptyUserConfiguration() throws ResourceHandlerException { String allowed = ""; boolean invalidConfiguration = false; configuration.set(YarnConfiguration.NM_FPGA_ALLOWED_DEVICES, allowed); try { fpgaResourceHandler.bootstrap(configuration); } catch (ResourceHandlerException e) { invalidConfiguration = true; } Assert.assertTrue(invalidConfiguration); }