org.apache.reef.driver.evaluator.EvaluatorRequest Java Examples

@Test
public void testSerializeDeserializeNoRacksNorNodesRequest() {
  final EvaluatorRequest originalRequest = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(10)
      .setNumberOfCores(3)
      .build();
  final String serialized = AvroEvaluatorRequestSerializer.toString(originalRequest);
  final EvaluatorRequest deserializedRequest = AvroEvaluatorRequestSerializer.fromString(serialized);

  Assert.assertEquals(originalRequest.getMegaBytes(), deserializedRequest.getMegaBytes());
  Assert.assertEquals(originalRequest.getNumber(), deserializedRequest.getNumber());
  Assert.assertEquals(originalRequest.getNumberOfCores(), deserializedRequest.getNumberOfCores());
  Assert.assertTrue(deserializedRequest.getRackNames().size() == 0);
  Assert.assertTrue(deserializedRequest.getNodeNames().size() == 0);
}
 

@Test
public void testSerializeDeserializeCompleteRequest() {
  final EvaluatorRequest originalRequest = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(10)
      .setNumberOfCores(3)
      .addRackName("rack1")
      .addRackName("rack2")
      .addNodeName("nodename1")
      .addNodeName("nodename2")
      .build();
  final String serializedRequest = AvroEvaluatorRequestSerializer.toString(originalRequest);
  final EvaluatorRequest deserializedRequest = AvroEvaluatorRequestSerializer.fromString(serializedRequest);

  Assert.assertEquals(originalRequest.getMegaBytes(), deserializedRequest.getMegaBytes());
  Assert.assertEquals(originalRequest.getNumber(), deserializedRequest.getNumber());
  Assert.assertEquals(originalRequest.getNumberOfCores(), deserializedRequest.getNumberOfCores());
  Assert.assertEquals(originalRequest.getRackNames(), deserializedRequest.getRackNames());
  Assert.assertEquals(originalRequest.getNodeNames(), deserializedRequest.getNodeNames());
}
 

@Test
public void scheduleHeronWorkersRequestsContainersForPacking() throws Exception {
  Set<PackingPlan.ContainerPlan> containers = new HashSet<>();

  PackingPlan.ContainerPlan container1 = PackingTestUtils.testContainerPlan(1, 1, 2);
  containers.add(container1);
  PackingPlan.ContainerPlan container2 = PackingTestUtils.testContainerPlan(2, 1, 2, 3);
  containers.add(container2);

  PackingPlan packing = new PackingPlan("packingId", containers);
  spyDriver.scheduleHeronWorkers(packing);
  verify(mockRequestor, times(2)).submit(any(EvaluatorRequest.class));
  verify(spyDriver, times(1)).requestContainerForWorker(eq(1), anyHeronWorker());
  verify(spyDriver, times(1)).requestContainerForWorker(eq(2), anyHeronWorker());
  verify(spyDriver, times(1)).createEvaluatorRequest(getCpu(container1), getRam(container1));
  verify(spyDriver, times(1)).createEvaluatorRequest(getCpu(container2), getRam(container2));
}
 

@Test
public void scheduleHeronWorkersAddsContainers() throws Exception {
  PackingPlan packingPlan = PackingTestUtils.testPackingPlan("test", new RoundRobinPacking());
  spyDriver.scheduleHeronWorkers(packingPlan);
  verify(spyDriver, times(1)).requestContainerForWorker(eq(1), anyHeronWorker());
  verify(mockRequestor, times(1)).submit(any(EvaluatorRequest.class));

  Set<PackingPlan.ContainerPlan> toBeAddedContainerPlans = new HashSet<>();
  toBeAddedContainerPlans.add(PackingTestUtils.testContainerPlan(2));
  toBeAddedContainerPlans.add(PackingTestUtils.testContainerPlan(3));

  spyDriver.scheduleHeronWorkers(toBeAddedContainerPlans);
  verify(spyDriver, times(1)).requestContainerForWorker(eq(2), anyHeronWorker());
  verify(spyDriver, times(1)).requestContainerForWorker(eq(3), anyHeronWorker());
  verify(mockRequestor, times(3)).submit(any(EvaluatorRequest.class));
}
 

@Override
public void onNext(final StartTime startTime) {
  LOG.log(Level.INFO, "StartTime: {0} Number of files in the set: {1}",
      new Object[]{startTime, FileResourceTestDriver.this.fileNamesToExpect.size()});

  // Check whether the files made it
  for (final String fileName : FileResourceTestDriver.this.fileNamesToExpect) {
    final File file = new File(localFolder, fileName);
    LOG.log(Level.INFO, "Testing file: " + file.getAbsolutePath());
    if (!file.exists()) {
      throw new DriverSideFailure("Cannot find file: " + fileName);
    } else if (!file.isFile()) {
      throw new DriverSideFailure("Not a file: " + fileName);
    } else if (!file.canRead()) {
      throw new DriverSideFailure("Can't read: " + fileName);
    }
  }

  // Ask for a single evaluator.
  FileResourceTestDriver.this.requestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(1).setMemory(64).setNumberOfCores(1).build());
}
 

@Override
public void onNext(final EvaluatorRequest request) {
  try {

    LOG.log(Level.FINE,
        "Processing a request with count: {0} - Waiting for gate to be released",
        request.getNumber());

    this.resourceRequestGate.await();

    LOG.log(Level.FINE, "Gate released. Submitting request: {0}", request);
    this.resourceRequestGate = new CountDownLatch(1);
    this.requestor.submit(request);

  } catch (final InterruptedException ex) {
    LOG.log(Level.FINEST, "Interrupted", ex);
  }
}
 

private Configuration getDataLoadConfiguration(final String jobName) {
  final EvaluatorRequest computeRequest = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(memory)
      .build();
  final Configuration dataLoadConfiguration = new DataLoadingRequestBuilder()
      .setInputFormatClass(TextInputFormat.class)
      .setInputPath(input)
      .setNumberOfDesiredSplits(numSplits)
      .addComputeRequest(computeRequest)
      .renewFailedEvaluators(false)
      .setDriverConfigurationModule(DriverConfiguration.CONF
          .setMultiple(DriverConfiguration.GLOBAL_LIBRARIES, EnvironmentUtils.getAllClasspathJars())
          .set(DriverConfiguration.DRIVER_MEMORY, Integer.toString(memory))
          .set(DriverConfiguration.ON_CONTEXT_ACTIVE, BGDDriver.ContextActiveHandler.class)
          .set(DriverConfiguration.ON_TASK_RUNNING, BGDDriver.TaskRunningHandler.class)
          .set(DriverConfiguration.ON_TASK_FAILED, BGDDriver.TaskFailedHandler.class)
          .set(DriverConfiguration.ON_TASK_COMPLETED, BGDDriver.TaskCompletedHandler.class)
          .set(DriverConfiguration.DRIVER_IDENTIFIER, jobName))
      .build();
  return dataLoadConfiguration;
}
 

private static AvroEvaluatorRequest toAvro(final EvaluatorRequest request) {
  final List<CharSequence> nodeNames = new ArrayList<>();
  for (final String nodeName : request.getNodeNames()) {
    nodeNames.add(nodeName);
  }

  final List<CharSequence> rackNames = new ArrayList<>();
  for (final String rackName : request.getRackNames()) {
    rackNames.add(rackName);
  }

  return AvroEvaluatorRequest.newBuilder()
      .setCores(request.getNumberOfCores())
      .setMegaBytes(request.getMegaBytes())
      .setNumber(request.getNumber())
      .setNodeNames(nodeNames)
      .setRackNames(rackNames)
      .build();
}
 

/**
 * Expect an IllegalArgumentException when a non-positive evaluator count is passed.
 */
@Test(expected = IllegalArgumentException.class)
public void testIllegalCount() {
  final int memory = 128;
  final int count = 0;
  final DummyRequestHandler requestHandler = new DummyRequestHandler();
  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, requestHandler, loggingScopeFactory);
  evaluatorRequestor.submit(EvaluatorRequest.newBuilder().setMemory(memory).setNumberOfCores(1).setNumber(count)
      .build());
}
 

@Override
public void onNext(final StartTime startTime) {
  this.requestor.submit(EvaluatorRequest.newBuilder()
      .setMemory(64)
      .setNumber(1)
      .setNumberOfCores(1)
      .build());
}
 

@Override
public void onNext(final StartTime time) {

  LOG.log(Level.FINE, "StartTime: {0} :: request {1} evaluators",
      new Object[]{time, NUM_EVALUATORS});

  requestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(NUM_EVALUATORS).setMemory(64).setNumberOfCores(1).build());
}
 

@Override
public void onNext(final StartTime startTime) {

  // Check that the parameters were set
  checkParameter(DriverParameter.class, DriverParameter.TEST_VALUE);
  checkParameter(DriverAndEvaluatorParameter.class, DriverAndEvaluatorParameter.TEST_VALUE);

  // If we got here, the parameters were set correctly and we can proceed.
  evaluatorRequestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(32)
      .build());
}
 

@Override
public void onNext(final StartTime startTime) {
  EvaluatorSizeTestDriver.this.evaluatorRequestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(EvaluatorSizeTestDriver.this.memorySize)
      .setNumberOfCores(1)
      .build());
}
 

/**
 * If only memory, no count is given, 1 evaluator should be requested.
 */
@Test
public void testMemoryOnly() {
  final int memory = 777;
  final DummyRequestHandler requestHandler = new DummyRequestHandler();
  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, requestHandler, loggingScopeFactory);
  evaluatorRequestor.submit(EvaluatorRequest.newBuilder().setMemory(memory).build());
  Assert.assertEquals("Memory request did not make it",
      memory, requestHandler.get().getMemorySize().get().intValue());
  Assert.assertEquals("Number of requests did not make it", 1, requestHandler.get().getResourceCount());
}
 

/**
 * Checks whether memory and count make it correctly.
 */
@Test
public void testMemoryAndCount() {
  final int memory = 777;
  final int count = 9;
  final DummyRequestHandler requestHandler = new DummyRequestHandler();
  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, requestHandler, loggingScopeFactory);
  evaluatorRequestor.submit(EvaluatorRequest.newBuilder().setMemory(memory).setNumber(count).build());
  Assert.assertEquals("Memory request did not make it",
      memory, requestHandler.get().getMemorySize().get().intValue());
  Assert.assertEquals("Number of requests did not make it", count, requestHandler.get().getResourceCount());
}
 

/**
 * Expect an IllegalArgumentException when a non-positive memory amount is passed.
 */
@Test(expected = IllegalArgumentException.class)
public void testIllegalMemory() {
  final int memory = 0;
  final int count = 1;
  final DummyRequestHandler requestHandler = new DummyRequestHandler();
  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, requestHandler, loggingScopeFactory);
  evaluatorRequestor.submit(EvaluatorRequest.newBuilder().setMemory(memory).setNumberOfCores(1).setNumber(count)
      .build());
}
 

/**
 * Requests containers/evaluators with the given specifications.
 * @param numToRequest number of containers to request
 * @param resourceSpecification containing the specifications of
 */
public synchronized void requestContainer(final int numToRequest,
                                          final ResourceSpecification resourceSpecification) {
  if (numToRequest > 0) {
    // Create a list of executor specifications to be used when containers are allocated.
    final List<ResourceSpecification> resourceSpecificationList = new ArrayList<>(numToRequest);
    for (int i = 0; i < numToRequest; i++) {
      resourceSpecificationList.add(resourceSpecification);
    }
    executorsByContainerType.putIfAbsent(resourceSpecification.getContainerType(), new ArrayList<>(numToRequest));

    // Mark the request as pending with the given specifications.
    pendingContainerRequestsByContainerType.putIfAbsent(resourceSpecification.getContainerType(), new ArrayList<>());
    pendingContainerRequestsByContainerType.get(resourceSpecification.getContainerType())
        .addAll(resourceSpecificationList);

    requestLatchByResourceSpecId.put(resourceSpecification.getResourceSpecId(),
        new CountDownLatch(numToRequest));

    // Request the evaluators
    evaluatorRequestor.submit(EvaluatorRequest.newBuilder()
        .setNumber(numToRequest)
        .setMemory(resourceSpecification.getMemory())
        .setNumberOfCores(resourceSpecification.getCapacity())
        .build());
  } else {
    LOG.info("Request {} containers", numToRequest);
  }
}
 

@Override
public void onNext(final StartTime startTime) {
  requestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(numEvaluators)
      .setMemory(64)
      .setNumberOfCores(1)
      .build());
}
 

@Override
public void onEvaluatorRequest(final EvaluatorRequest evaluatorRequest) {
  this.serviceStub.requestResources(
      ResourceRequest.newBuilder()
          .setCores(evaluatorRequest.getNumberOfCores())
          .setMemorySize(evaluatorRequest.getMegaBytes())
          .setRelaxLocality(evaluatorRequest.getRelaxLocality())
          .setResourceCount(evaluatorRequest.getNumber())
          .setRuntimeName(evaluatorRequest.getRuntimeName())
          .addAllRackNameList(evaluatorRequest.getRackNames())
          .addAllNodeNameList(evaluatorRequest.getNodeNames())
          .build());
}
 

@Override
public void onNext(final StartTime startTime) {
  LOG.log(Level.INFO, "StartTime: ", startTime);
  HelloDriver.this.requestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(nCLRTasks + nJVMTasks)
      .setMemory(128)
      .setNumberOfCores(1)
      .build());
}
 

/**
 * Tests whether the amount of memory is transferred correctly.
 */
@Test
public void testDifferentMemory() throws InjectionException {

  final LoggingScopeFactory loggingScopeFactory =
      Tang.Factory.getTang().newInjector().getInstance(LoggingScopeFactory.class);

  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, resourceRequestHandler, loggingScopeFactory);

  final EvaluatorRequest requestOne = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(64)
      .setNumberOfCores(1)
      .build();

  final EvaluatorRequest requestTwo = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(128)
      .setNumberOfCores(2)
      .build();

  evaluatorRequestor.submit(requestOne);
  Assert.assertEquals("Request in REEF and YARN form should have the same amount of memory",
      requestOne.getMegaBytes(),
      containerRequestHandler.getRequests()[0].getCapability().getMemory());

  evaluatorRequestor.submit(requestTwo);
  Assert.assertEquals("Request in REEF and YARN form should have the same amount of memory",
      requestTwo.getMegaBytes(),
      containerRequestHandler.getRequests()[0].getCapability().getMemory());

  evaluatorRequestor.submit(requestOne);
  Assert.assertNotEquals("Second YARN request should have different amount of memory",
      requestTwo.getMegaBytes(),
      containerRequestHandler.getRequests()[0].getCapability().getMemory());
}
 

@Test
public void testEvaluatorCount() throws InjectionException {

  final LoggingScopeFactory loggingScopeFactory =
      Tang.Factory.getTang().newInjector().getInstance(LoggingScopeFactory.class);

  final EvaluatorRequestor evaluatorRequestor =
      new EvaluatorRequestorImpl(resourceCatalog, resourceRequestHandler, loggingScopeFactory);

  final EvaluatorRequest requestOne = EvaluatorRequest.newBuilder()
      .setNumber(1)
      .setMemory(64)
      .setNumberOfCores(1)
      .build();

  final EvaluatorRequest requestTwo = EvaluatorRequest.newBuilder()
      .setNumber(2)
      .setMemory(128)
      .setNumberOfCores(2)
      .build();

  evaluatorRequestor.submit(requestOne);
  Assert.assertEquals("Request in REEF and YARN form should have the same number of Evaluators",
      requestOne.getNumber(),
      containerRequestHandler.getRequests().length);

  evaluatorRequestor.submit(requestTwo);
  Assert.assertEquals("Request in REEF and YARN form should have the same number of Evaluators",
      requestTwo.getNumber(),
      containerRequestHandler.getRequests().length);

  evaluatorRequestor.submit(requestTwo);
  Assert.assertNotEquals("Second YARN request should have different number of Evaluators",
      requestOne.getNumber(),
      containerRequestHandler.getRequests().length);
}
 

@Override
public void submit(final EvaluatorRequest req) {
  if (this.clock.get().isClosed()) {
    throw new IllegalStateException("clock closed");
  }
  final NodeDescriptor nodeDescriptor = new MockNodeDescriptor();
  final MockEvaluatorDescriptor evaluatorDescriptor = new MockEvaluatorDescriptor(nodeDescriptor);
  for (int i = 0; i < req.getNumber(); i++) {
    final MockAllocatedEvaluator mockEvaluator = new MockAllocatedEvaluator(
        this.mockRuntimeDriver.get(), UUID.randomUUID().toString(), evaluatorDescriptor);
    this.mockRuntimeDriver.get().add(new AllocateEvaluator(mockEvaluator));
  }
}
 

public void submit(final int evaluatorsNumber,
                   final int memory,
                   final int virtualCore,
                   final boolean relaxLocality,
                   final String rack,
                   final String runtimeName,
                   final ArrayList<String> nodeNames,
                   final String nodeLabelExpression) {
  if (this.isBlocked) {
    throw new RuntimeException("Cannot request additional Evaluator, this is probably because " +
        "the Driver has crashed and restarted, and cannot ask for new container due to YARN-2433.");
  }

  if (rack != null && !rack.isEmpty()) {
    LOG.log(Level.WARNING, "Ignoring rack preference.");
  }

  try (LoggingScope ls = loggingScopeFactory.evaluatorRequestSubmitToJavaDriver(evaluatorsNumber)) {
    clrEvaluatorsNumber += evaluatorsNumber;

    final EvaluatorRequest request = EvaluatorRequest.newBuilder()
        .setNumber(evaluatorsNumber)
        .setMemory(memory)
        .setNumberOfCores(virtualCore)
        .setRuntimeName(runtimeName)
        .setRelaxLocality(relaxLocality)
        .addNodeNames(nodeNames)
        .setNodeLabelExpression(nodeLabelExpression)
        .build();

    LOG.log(Level.FINE, "submitting evaluator request {0}", request);
    jevaluatorRequestor.submit(request);
  }
}
 

/**
 * Adds the requests to the data requests list.
 *
 * @param dataRequests
 *          the data requests to add
 * @return this
 */
@SuppressWarnings("checkstyle:hiddenfield")
public DataLoadingRequestBuilder addDataRequests(final List<EvaluatorRequest> dataRequests) {
  for (final EvaluatorRequest dataRequest : dataRequests) {
    addDataRequest(dataRequest);
  }
  return this;
}
 

/**
 * Requests containers/evaluators with the given specifications.
 *
 * @param numToRequest          number of containers to request
 * @param resourceSpecification containing the specifications of
 */
public void requestContainer(final int numToRequest, final ResourceSpecification resourceSpecification) {
  if (isTerminated) {
    LOG.info("ContainerManager is terminated, ignoring {}", resourceSpecification.toString());
    return;
  }

  if (numToRequest > 0) {
    // Create a list of executor specifications to be used when containers are allocated.
    final List<ResourceSpecification> resourceSpecificationList = new ArrayList<>(numToRequest);
    for (int i = 0; i < numToRequest; i++) {
      resourceSpecificationList.add(resourceSpecification);
    }

    // Mark the request as pending with the given specifications.
    pendingContainerRequestsByContainerType.putIfAbsent(resourceSpecification.getContainerType(), new ArrayList<>());
    pendingContainerRequestsByContainerType.get(resourceSpecification.getContainerType())
      .addAll(resourceSpecificationList);

    requestLatchByResourceSpecId.put(resourceSpecification.getResourceSpecId(),
      new CountDownLatch(numToRequest));

    // Request the evaluators
    evaluatorRequestor.submit(EvaluatorRequest.newBuilder()
      .setNumber(numToRequest)
      .setMemory(resourceSpecification.getMemory())
      .setNumberOfCores(resourceSpecification.getCapacity())
      .build());
  } else {
    LOG.info("Request {} containers", numToRequest);
  }
}
 

@VisibleForTesting
void requestContainerForWorker(int id, final HeronWorker worker) {
  int cpu = worker.cores;
  ByteAmount mem = worker.mem;
  EvaluatorRequest evaluatorRequest = createEvaluatorRequest(cpu, mem);
  LOG.info(String.format("Requesting container for worker: %d, RAM: %s, CPU: %d", id, mem, cpu));
  requestor.submit(evaluatorRequest);
}
 

@VisibleForTesting
EvaluatorRequest createEvaluatorRequest(int cpu, ByteAmount mem) {
  return EvaluatorRequest
      .newBuilder()
      .setNumber(1)
      .setMemory(((Long) mem.asMegabytes()).intValue())
      .setNumberOfCores(cpu)
      .build();
}
 

@Test
public void requestContainerForWorkerSubmitsValidRequest() {
  ByteAmount memory = ByteAmount.fromMegabytes(786);
  EvaluatorRequest request = spyDriver.createEvaluatorRequest(5, memory);
  doReturn(request).when(spyDriver).createEvaluatorRequest(5, memory);
  HeronMasterDriver.HeronWorker worker = new HeronMasterDriver.HeronWorker(3, 5, memory);
  spyDriver.requestContainerForWorker(3, worker);
  verify(mockRequestor, times(1)).submit(request);
}
 

@Test(expected = HeronMasterDriver.ContainerAllocationException.class)
public void scheduleHeronWorkersFailsOnDuplicateRequest() throws Exception {
  PackingPlan packingPlan = PackingTestUtils.testPackingPlan("test", new RoundRobinPacking());
  spyDriver.scheduleHeronWorkers(packingPlan);
  verify(spyDriver, times(1)).requestContainerForWorker(eq(1), anyHeronWorker());
  verify(mockRequestor, times(1)).submit(any(EvaluatorRequest.class));

  PackingPlan.ContainerPlan duplicatePlan = PackingTestUtils.testContainerPlan(1);
  Set<PackingPlan.ContainerPlan> toBeAddedContainerPlans = new HashSet<>();
  toBeAddedContainerPlans.add(duplicatePlan);
  spyDriver.scheduleHeronWorkers(toBeAddedContainerPlans);
}