NvidiaGPUPluginForRuntimeV2 implements DevicePlugin,
    DevicePluginScheduler { @Override public DeviceRuntimeSpec onDevicesAllocated(Set<Device> allocatedDevices, YarnRuntimeType yarnRuntime) throws Exception { LOG.debug("Generating runtime spec for allocated devices: {}, {}", allocatedDevices, yarnRuntime.getName()); if (yarnRuntime == YarnRuntimeType.RUNTIME_DOCKER) { String nvidiaRuntime = "nvidia"; String nvidiaVisibleDevices = "NVIDIA_VISIBLE_DEVICES"; StringBuffer gpuMinorNumbersSB = new StringBuffer(); for (Device device : allocatedDevices) { gpuMinorNumbersSB.append(device.getMinorNumber() + ","); } String minorNumbers = gpuMinorNumbersSB.toString(); LOG.info("Nvidia Docker v2 assigned GPU: " + minorNumbers); return DeviceRuntimeSpec.Builder.newInstance() .addEnv(nvidiaVisibleDevices, minorNumbers.substring(0, minorNumbers.length() - 1)) .setContainerRuntime(nvidiaRuntime) .build(); } return null; }  @Override DeviceRegisterRequest getRegisterRequestInfo(); @Override Set<Device> getDevices(); @Override DeviceRuntimeSpec onDevicesAllocated(Set<Device> allocatedDevices,
      YarnRuntimeType yarnRuntime); @Override void onDevicesReleased(Set<Device> releasedDevices); @Override Set<Device> allocateDevices(Set<Device> availableDevices, int count,
      Map<String, String> envs); @VisibleForTesting void initCostTable(); @VisibleForTesting int computeCostOfDevices(Device[] devices); @VisibleForTesting void topologyAwareSchedule(Set<Device> allocation, int count,
      Map<String, String> envs,
      Set<Device> availableDevices,
      Map<Integer, List<Map.Entry<Set<Device>, Integer>>> cTable); @VisibleForTesting void basicSchedule(Set<Device> allocation, int count,
      Set<Device> availableDevices); void parseTopo(String topo,
      Map<String, Integer> deviceLinkToWeight); @VisibleForTesting void setPathOfGpuBinary(String pOfGpuBinary); @VisibleForTesting void setShellExecutor(
      NvidiaCommandExecutor shellExecutor); @VisibleForTesting boolean isTopoInitialized(); @VisibleForTesting Map<Integer, List<Map.Entry<Set<Device>, Integer>>> getCostTable(); @VisibleForTesting Map<String, Integer> getDevicePairToWeight(); static final Logger LOG; static final String NV_RESOURCE_NAME; static final String TOPOLOGY_POLICY_ENV_KEY; static final String TOPOLOGY_POLICY_PACK; static final String TOPOLOGY_POLICY_SPREAD; }
@Test public void testOnDeviceAllocated() throws Exception { NvidiaGPUPluginForRuntimeV2 plugin = new NvidiaGPUPluginForRuntimeV2(); Set<Device> allocatedDevices = new TreeSet<>(); DeviceRuntimeSpec spec = plugin.onDevicesAllocated(allocatedDevices, YarnRuntimeType.RUNTIME_DEFAULT); Assert.assertNull(spec); allocatedDevices.add(Device.Builder.newInstance() .setId(0).setHealthy(true) .setBusID("00000000:04:00.0") .setDevPath("/dev/nvidia0") .setMajorNumber(195) .setMinorNumber(0).build()); spec = plugin.onDevicesAllocated(allocatedDevices, YarnRuntimeType.RUNTIME_DOCKER); Assert.assertEquals("nvidia", spec.getContainerRuntime()); Assert.assertEquals("0", spec.getEnvs().get("NVIDIA_VISIBLE_DEVICES")); allocatedDevices.add(Device.Builder.newInstance() .setId(0).setHealthy(true) .setBusID("00000000:82:00.0") .setDevPath("/dev/nvidia1") .setMajorNumber(195) .setMinorNumber(1).build()); spec = plugin.onDevicesAllocated(allocatedDevices, YarnRuntimeType.RUNTIME_DOCKER); Assert.assertEquals("nvidia", spec.getContainerRuntime()); Assert.assertEquals("0,1", spec.getEnvs().get("NVIDIA_VISIBLE_DEVICES")); }