Java Code Examples for org.apache.hadoop.yarn.api.records.Resource#setGpuCores()

/**
 * Test RegisterNodeManagerRequestPBImpl.
 */

@Test
public void testRegisterNodeManagerRequestPBImpl() {
  RegisterNodeManagerRequestPBImpl original = new RegisterNodeManagerRequestPBImpl();
  original.setHttpPort(8080);
  original.setNodeId(getNodeId());
  Resource resource = recordFactory.newRecordInstance(Resource.class);
  resource.setMemory(10000);
  resource.setVirtualCores(2);
  resource.setGpuCores(3);
  original.setResource(resource);
  RegisterNodeManagerRequestPBImpl copy = new RegisterNodeManagerRequestPBImpl(
      original.getProto());

  assertEquals(8080, copy.getHttpPort());
  assertEquals(9090, copy.getNodeId().getPort());
  assertEquals(10000, copy.getResource().getMemory());
  assertEquals(2, copy.getResource().getVirtualCores());
  assertEquals(3, copy.getResource().getGpuCores());

}
 

@Test(timeout = 15000)
public void testSubmitApplicationOnHA() throws Exception {
  ApplicationSubmissionContext appContext =
      Records.newRecord(ApplicationSubmissionContext.class);
  appContext.setApplicationId(cluster.createFakeAppId());
  ContainerLaunchContext amContainer =
      Records.newRecord(ContainerLaunchContext.class);
  appContext.setAMContainerSpec(amContainer);
  Resource capability = Records.newRecord(Resource.class);
  capability.setMemory(10);
  capability.setVirtualCores(1);
  capability.setGpuCores(1);
  appContext.setResource(capability);
  ApplicationId appId = client.submitApplication(appContext);
  Assert.assertTrue(getActiveRM().getRMContext().getRMApps()
      .containsKey(appId));
}
 

public static Resource createResource(int memory, int cores, int gcores) {
  Resource resource = Records.newRecord(Resource.class);
  resource.setMemory(memory);
  resource.setVirtualCores(cores);
  resource.setGpuCores(gcores);
  return resource;
}
 

public static Resource multiplyAndRoundDown(Resource lhs, double by) {
  Resource out = clone(lhs);
  out.setMemory((int)(lhs.getMemory() * by));
  out.setVirtualCores((int)(lhs.getVirtualCores() * by));
  out.setGpuCores((int)(lhs.getGpuCores() * by));
  return out;
}
 

public static Resource newResource(int memory, int vCores) {
  Resource resource = recordFactory.newRecordInstance(Resource.class);
  resource.setMemory(memory);
  resource.setVirtualCores(vCores);
  resource.setGpuCores(0);
  return resource;
}
 

public static Resource newResource(int memory, int vCores, int gCores) {
  Resource resource = recordFactory.newRecordInstance(Resource.class);
  resource.setMemory(memory);
  resource.setVirtualCores(vCores);
  resource.setGpuCores(gCores);
  return resource;
}
 

private static void setResourceValue(int val, Resource resource, ResourceType type) {
  switch (type) {
  case MEMORY:
    resource.setMemory(val);
    break;
  case CPU:
    resource.setVirtualCores(val);
    break;
  case GPU:
    resource.setGpuCores(val);
    break;
  default:
    throw new IllegalArgumentException("Invalid resource");
  }
}
 

public static Resource addTo(Resource lhs, Resource rhs) {
  lhs.setMemory(lhs.getMemory() + rhs.getMemory());
  lhs.setVirtualCores(lhs.getVirtualCores() + rhs.getVirtualCores());
  lhs.setGpuCores(lhs.getGpuCores() + rhs.getGpuCores());
  return lhs;
}
 

public static Resource subtractFrom(Resource lhs, Resource rhs) {
  lhs.setMemory(lhs.getMemory() - rhs.getMemory());
  lhs.setVirtualCores(lhs.getVirtualCores() - rhs.getVirtualCores());
  lhs.setGpuCores(lhs.getGpuCores() - rhs.getGpuCores());
  return lhs;
}
 

public static Resource multiplyTo(Resource lhs, double by) {
  lhs.setMemory((int)(lhs.getMemory() * by));
  lhs.setVirtualCores((int)(lhs.getVirtualCores() * by));
  lhs.setGpuCores((int) (lhs.getGpuCores() * by));
  return lhs;
}
 

@Test public void testNames() {
  DefaultMetricsSystem.initialize("NodeManager");
  NodeManagerMetrics metrics = NodeManagerMetrics.create();
  Resource total = Records.newRecord(Resource.class);
  total.setMemory(8*GiB);
  total.setVirtualCores(16);
  total.setGpuCores(16);
  Resource resource = Records.newRecord(Resource.class);
  resource.setMemory(512); //512MiB
  resource.setVirtualCores(2);
  resource.setGpuCores(1);


  metrics.addResource(total);

  for (int i = 10; i-- > 0;) {
    // allocate 10 containers(allocatedGB: 5GiB, availableGB: 3GiB)
    metrics.launchedContainer();
    metrics.allocateContainer(resource);
  }

  metrics.initingContainer();
  metrics.endInitingContainer();
  metrics.runningContainer();
  metrics.endRunningContainer();
  // Releasing 3 containers(allocatedGB: 3.5GiB, availableGB: 4.5GiB)
  metrics.completedContainer();
  metrics.releaseContainer(resource);

  metrics.failedContainer();
  metrics.releaseContainer(resource);

  metrics.killedContainer();
  metrics.releaseContainer(resource);

  metrics.initingContainer();
  metrics.runningContainer();

  Assert.assertTrue(!metrics.containerLaunchDuration.changed());
  metrics.addContainerLaunchDuration(1);
  Assert.assertTrue(metrics.containerLaunchDuration.changed());

  // availableGB is expected to be floored,
  // while allocatedGB is expected to be ceiled.
  // allocatedGB: 3.5GB allocated memory is shown as 4GB
  // availableGB: 4.5GB available memory is shown as 4GB
  checkMetrics(10, 1, 1, 1, 1, 1, 4, 7, 4, 14, 2, 7, 9);
}
 

@Test
public void testNodeRegistrationWithMinimumAllocations() throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, "2048");
  conf.set(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, "4");
  conf.set(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_GCORES, "4");
  rm = new MockRM(conf);
  rm.start();

  ResourceTrackerService resourceTrackerService
    = rm.getResourceTrackerService();
  RegisterNodeManagerRequest req = Records.newRecord(
      RegisterNodeManagerRequest.class);
  NodeId nodeId = BuilderUtils.newNodeId("host", 1234);
  req.setNodeId(nodeId);

  Resource capability = BuilderUtils.newResource(1024, 1, 1);
  req.setResource(capability);
  RegisterNodeManagerResponse response1 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response1.getNodeAction());
  
  capability.setMemory(2048);
  capability.setVirtualCores(1);
  capability.setGpuCores(1);
  req.setResource(capability);
  RegisterNodeManagerResponse response2 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response2.getNodeAction());
  
  capability.setMemory(1024);
  capability.setVirtualCores(4);
  capability.setGpuCores(4);
  req.setResource(capability);
  RegisterNodeManagerResponse response3 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response3.getNodeAction());
  
  capability.setMemory(2048);
  capability.setVirtualCores(4);
  capability.setGpuCores(4);
  req.setResource(capability);
  RegisterNodeManagerResponse response4 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.NORMAL,response4.getNodeAction());
}