NvidiaGPUPluginForRuntimeV2 implements DevicePlugin,
    DevicePluginScheduler { @Override public Set<Device> getDevices() throws Exception { shellExecutor.searchBinary(); TreeSet<Device> r = new TreeSet<>(); String output; try { output = shellExecutor.getDeviceInfo(); String[] lines = output.trim().split("\n"); int id = 0; for (String oneLine : lines) { String[] tokensEachLine = oneLine.split(","); if (tokensEachLine.length != 2) { throw new Exception("Cannot parse the output to get device info. " + "Unexpected format in it:" + oneLine); } String minorNumber = tokensEachLine[0].trim(); String busId = tokensEachLine[1].trim(); String majorNumber = getMajorNumber(DEV_NAME_PREFIX + minorNumber); if (majorNumber != null) { r.add(Device.Builder.newInstance() .setId(id) .setMajorNumber(Integer.parseInt(majorNumber)) .setMinorNumber(Integer.parseInt(minorNumber)) .setBusID(busId) .setDevPath("/dev/" + DEV_NAME_PREFIX + minorNumber) .setHealthy(true) .build()); id++; } } lastTimeFoundDevices = r; return r; } catch (IOException e) { LOG.debug("Failed to get output from {}", pathOfGpuBinary); throw new YarnException(e); } }  @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 testTopologySchedulingPerformanceWithPackPolicyWithNVLink() throws Exception { NvidiaGPUPluginForRuntimeV2 plugin = mockEightGPUPlugin(); NvidiaGPUPluginForRuntimeV2 spyPlugin = spy(plugin); Set<Device> allDevices = spyPlugin.getDevices(); Map<String, String> env = new HashMap<>(); env.put(NvidiaGPUPluginForRuntimeV2.TOPOLOGY_POLICY_ENV_KEY, NvidiaGPUPluginForRuntimeV2.TOPOLOGY_POLICY_PACK); ActualPerformanceReport report = new ActualPerformanceReport(); report.readFromFile(); ArrayList<ActualPerformanceReport.DataRecord> dataSet = report.getDataSet(); assertThat(dataSet).hasSize(2952); String[] allModels = {"alexnet", "resnet50", "vgg16", "inception3"}; int[] batchSizes = {32, 64, 128}; int[] gpuCounts = {2, 3, 4, 5, 6, 7}; float totalBoostAgainstMedian = 0; int count = 0; float maxBoostAgainstMedian = 0; float totalBoostAgainstMin = 0; float maxBoostAgainstMin = 0; for (String model : allModels) { float totalBoostAgainstMinCertainModel = 0; float totalBoostAgainstMedianCertainModel = 0; float maxBoostAgainstMinCertainModel = 0; float maxBoostAgainstMedianCertainModel = 0; int countOfEachModel = 0; for (int bs : batchSizes) { for (int gpuCount: gpuCounts) { float bstAgainstMedian = calculatePerformanceBoostAgainstMedian( report, model, bs, gpuCount, plugin, allDevices, env); float bstAgainstMinimum = calculatePerformanceBoostAgainstMinimum( report, model, bs, gpuCount, plugin, allDevices, env); totalBoostAgainstMedian += bstAgainstMedian; totalBoostAgainstMin += bstAgainstMinimum; count++; if (maxBoostAgainstMedian < bstAgainstMedian) { maxBoostAgainstMedian = bstAgainstMedian; } if (maxBoostAgainstMin < bstAgainstMinimum) { maxBoostAgainstMin = bstAgainstMinimum; } totalBoostAgainstMinCertainModel += bstAgainstMinimum; totalBoostAgainstMedianCertainModel += bstAgainstMedian; if (maxBoostAgainstMinCertainModel < bstAgainstMinimum) { maxBoostAgainstMinCertainModel = bstAgainstMinimum; } if (maxBoostAgainstMedianCertainModel < bstAgainstMedian) { maxBoostAgainstMedianCertainModel = bstAgainstMedian; } countOfEachModel++; } } LOG.info("Model:{}, The best performance boost against median value is " + "{}", model, maxBoostAgainstMedianCertainModel); LOG.info("Model:{}, The aggregated average performance boost against " + "median value is {}", model, totalBoostAgainstMedianCertainModel/countOfEachModel); LOG.info("Model:{}, The best performance boost against min value is {}", model, maxBoostAgainstMinCertainModel); LOG.info("Model:{}, The aggregated average performance boost against " + "min value is {}", model, totalBoostAgainstMinCertainModel/countOfEachModel); } LOG.info("For all, the best performance boost against median value is " + maxBoostAgainstMedian); LOG.info("For all, the aggregated average performance boost against median " + "value is " + totalBoostAgainstMedian/count); LOG.info("For all, the best performance boost against min value is " + maxBoostAgainstMin); LOG.info("For all, the aggregated average performance boost against min " + "value is " + totalBoostAgainstMin/count); }