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

/**
 * Uses the AllocatedEvaluator to launch a JVM task.
 *
 * @param allocatedEvaluator
 */
void onNextJVM(final AllocatedEvaluator allocatedEvaluator) {
  try {
    final Configuration contextConfiguration = ContextConfiguration.CONF
        .set(ContextConfiguration.IDENTIFIER, "HelloREEFContext")
        .build();

    final Configuration taskConfiguration = TaskConfiguration.CONF
        .set(TaskConfiguration.IDENTIFIER, "HelloREEFTask")
        .set(TaskConfiguration.TASK, HelloTask.class)
        .build();

    allocatedEvaluator.submitContextAndTask(contextConfiguration, taskConfiguration);
  } catch (final BindException ex) {
    final String message = "Unable to setup Task or Context configuration.";
    LOG.log(Level.SEVERE, message, ex);
    throw new RuntimeException(message, ex);
  }
}
 

@Override
public void fail(final AllocatedEvaluator evaluator) {
  if (this.clock.get().isClosed()) {
    throw new IllegalStateException("clock is closed");
  }
  if (this.allocatedEvaluatorMap.remove(evaluator.getId()) == null) {
    throw new IllegalStateException("unknown evaluator " + evaluator);
  }
  FailedTask failedTask = null;
  if (this.runningTasks.containsKey(evaluator.getId())) {
    final RunningTask task = this.runningTasks.remove(evaluator.getId());
    failedTask = new FailedTask(
        task.getId(),
        "mock",
        Optional.<String>empty(),
        Optional.<Throwable>empty(),
        Optional.<byte[]>empty(),
        Optional.<ActiveContext>of(task.getActiveContext()));
  }
  final List<FailedContext> failedContexts = new ArrayList<>();
  for (final MockActiveContext context : this.allocatedContextsMap.get(evaluator.getId())) {
    failedContexts.add(new MockFailedContext(context));
  }
  post(this.failedEvaluatorHandlers, new MockFailedEvaluator(
      evaluator.getId(), failedContexts, Optional.ofNullable(failedTask)));
}
 

@Override
public void onNext(final AllocatedEvaluator eval) {
  try {

    LOG.log(Level.INFO, "Allocated Evaluator: {0}", eval.getId());

    final Configuration thisContextConfiguration = ContextConfiguration.CONF.set(
        ContextConfiguration.IDENTIFIER, eval.getId() + "_context").build();

    eval.submitContext(Tang.Factory.getTang()
        .newConfigurationBuilder(thisContextConfiguration, contextConfig).build());

  } catch (final BindException ex) {
    throw new RuntimeException(ex);
  }
}
 

@Override
public void onNext(final AllocatedEvaluator eb) {
  try {
    final Configuration contextConfiguration = ContextConfiguration.CONF
        .set(ContextConfiguration.IDENTIFIER, "StatePassingContext")
        .build();

    final Configuration serviceConfiguration = ServiceConfiguration.CONF
        .set(ServiceConfiguration.SERVICES, Counter.class)
        .build();

    eb.submitContextAndService(contextConfiguration, serviceConfiguration);
  } catch (final BindException e) {
    throw new RuntimeException(e);
  }
}
 

private AllocatedEvaluator simulateContainerAllocation(String evaluatorId,
                                                       int cores,
                                                       ByteAmount ram,
                                                       int workerId) {
  AllocatedEvaluator evaluator = createMockEvaluator(evaluatorId, cores, ram);
  HeronMasterDriver.HeronWorker worker = new HeronMasterDriver.HeronWorker(workerId, cores, ram);

  Set<HeronMasterDriver.HeronWorker> workers = new HashSet<>();
  workers.add(worker);
  doReturn(workers).when(spyDriver).getWorkersAwaitingAllocation();

  doReturn(Optional.of(worker)).when(spyDriver)
      .findLargestFittingWorker(eq(evaluator), eq(workers), eq(false));

  spyDriver.new ContainerAllocationHandler().onNext(evaluator);
  return evaluator;
}
 

@Test
public void testMockFailures() {
  // make sure we're running
  assertTrue("mock application received start event", this.mockApplication.isRunning());

  // allocate an evaluator and get root context
  this.mockApplication.requestEvaluators(1);
  this.mockRuntime.succeed(this.mockRuntime.getNextProcessRequest());
  final AllocatedEvaluator evaluator = this.mockRuntime.getCurrentAllocatedEvaluators().iterator().next();
  this.mockApplication.submitContext(evaluator, "FOO");
  // fail evaluator
  this.mockRuntime.fail(evaluator);
  assertEquals("evaluator failed", 1, this.mockApplication.getFailedEvaluators().size());

  // both contexts should be failed
  assertEquals("root and child contexts failed", 2,
      this.mockApplication.getFailedContext().size());
}
 

/**
 * Fires the EvaluatorAllocatedEvent to the handlers. Can only be done once.
 */
public synchronized void fireEvaluatorAllocatedEvent() {

  if (this.stateManager.isAllocated() && this.allocationNotFired) {

    final AllocatedEvaluator allocatedEvaluator =
        new AllocatedEvaluatorImpl(this,
            this.remoteManager.getMyIdentifier(),
            this.configurationSerializer,
            getJobIdentifier(),
            this.loggingScopeFactory,
            this.evaluatorConfigurationProviders);

    LOG.log(Level.FINEST, "Firing AllocatedEvaluator event for Evaluator with ID [{0}]", this.evaluatorId);

    this.messageDispatcher.onEvaluatorAllocated(allocatedEvaluator);
    this.allocationNotFired = false;

  } else {
    LOG.log(Level.WARNING, "AllocatedEvaluator event fired more than once.");
  }
}
 

public void restartWorker(int id) throws ContainerAllocationException {
  LOG.log(Level.INFO, "Find & restart container for id={0}", id);

  Optional<HeronWorker> worker = multiKeyWorkerMap.lookupByWorkerId(id);
  if (!worker.isPresent()) {
    LOG.log(Level.WARNING, "Requesting a new container for: {0}", id);
    ContainerPlan containerPlan = containerPlans.get(id);
    if (containerPlan == null) {
      throw new IllegalArgumentException(
          String.format("There is no container for %s in packing plan.", id));
    }
    worker = Optional.of(new HeronWorker(id, containerPlan.getRequiredResource()));
  } else {
    AllocatedEvaluator evaluator = multiKeyWorkerMap.detachEvaluatorAndRemove(worker.get());
    LOG.log(Level.INFO, "Shutting down container {0}", evaluator.getId());
    evaluator.close();
  }

  requestContainerForWorker(worker.get().workerId, worker.get());
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {

  LOG.log(Level.FINE, "Submitting root context");

  try {

    final Configuration contextConfiguration = ContextConfiguration.CONF
        .set(ContextConfiguration.ON_CONTEXT_STARTED, ContextStartHandler1.class)
        .set(ContextConfiguration.ON_CONTEXT_STOP, ContextStopHandler1.class)
        .set(ContextConfiguration.IDENTIFIER, CONTEXT_1_IDENTIFIER)
        .build();

    allocatedEvaluator.submitContext(contextConfiguration);

    synchronized (SubContextDriver.this) {
      SubContextDriver.this.state = State.CONTEXT_1_SUBMITTED;
    }

  } catch (final BindException e) {
    throw new RuntimeException(e);
  }
}
 

@Test
public void onNextFailedEvaluatorRestartsContainer() throws Exception {
  int numContainers = 3;
  AllocatedEvaluator[] mockEvaluators = createApplicationWithContainers(numContainers);

  FailedEvaluator mockFailedContainer = mock(FailedEvaluator.class);
  when(mockFailedContainer.getId()).thenReturn("e1");
  verify(spyDriver, never()).requestContainerForWorker(anyInt(), anyHeronWorker());
  spyDriver.new FailedContainerHandler().onNext(mockFailedContainer);

  for (int id = 0; id < numContainers; id++) {
    if (id == 1) {
      verify(spyDriver, times(1)).requestContainerForWorker(eq(id), anyHeronWorker());
      assertFalse(spyDriver.lookupByEvaluatorId("e" + id).isPresent());
      continue;
    }
    verify(mockEvaluators[id], never()).close();
    assertEquals(Integer.valueOf(id), spyDriver.lookupByEvaluatorId("e" + id).get());
  }
}
 

@Test
public void onKillClosesContainersKillsTMaster() throws Exception {
  HeronMasterDriver.TMaster mockTMaster = mock(HeronMasterDriver.TMaster.class);
  when(spyDriver.buildTMaster(any(ExecutorService.class))).thenReturn(mockTMaster);

  int numContainers = 3;
  AllocatedEvaluator[] mockEvaluators = createApplicationWithContainers(numContainers);
  spyDriver.launchTMaster();

  spyDriver.killTopology();

  for (int id = 0; id < numContainers; id++) {
    Mockito.verify(mockEvaluators[id]).close();
    assertFalse(spyDriver.lookupByEvaluatorId("e" + id).isPresent());
  }

  verify(mockTMaster, times(1)).killTMaster();
}
 

@Test
public void restartWorkerRestartsSpecificWorker() throws Exception {
  int numContainers = 3;
  AllocatedEvaluator[] mockEvaluators = createApplicationWithContainers(numContainers);

  verify(spyDriver, never()).requestContainerForWorker(anyInt(), anyHeronWorker());
  spyDriver.restartWorker(1);

  for (int id = 0; id < numContainers; id++) {
    if (id == 1) {
      verify(spyDriver, times(1)).requestContainerForWorker(eq(id), anyHeronWorker());
      Mockito.verify(mockEvaluators[1]).close();
      assertFalse(spyDriver.lookupByEvaluatorId("e" + id).isPresent());
      continue;
    }
    verify(mockEvaluators[id], never()).close();
    assertEquals(Integer.valueOf(id), spyDriver.lookupByEvaluatorId("e" + id).get());
  }
}
 

@Override
public void onNext(final AllocatedEvaluator eval) {
  synchronized (HttpShellJobDriver.this) {
    LOG.log(Level.INFO, "Allocated Evaluator: {0} expect {1} running {2}",
        new Object[]{eval.getId(), HttpShellJobDriver.this.expectCount, HttpShellJobDriver.this.contexts.size()});
    assert HttpShellJobDriver.this.state == State.WAIT_EVALUATORS;
    try {
      eval.submitContext(ContextConfiguration.CONF.set(
          ContextConfiguration.IDENTIFIER, eval.getId() + "_context").build());
    } catch (final BindException ex) {
      LOG.log(Level.SEVERE, "Failed to submit a context to evaluator: " + eval.getId(), ex);
      throw new RuntimeException(ex);
    }
  }
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  final Configuration taskConfiguration = TaskConfiguration.CONF
      .set(TaskConfiguration.IDENTIFIER, "RogueThreadTestTask")
      .set(TaskConfiguration.TASK, RogueThreadTask.class)
      .build();
  allocatedEvaluator.submitTask(taskConfiguration);
}
 

@Test
public void testSuccessRequests() throws Exception {
  assertTrue("mock application received start event", this.mockApplication.isRunning());

  this.mockApplication.requestEvaluators(1);
  assertTrue("check for process event", this.mockRuntime.hasProcessRequest());
  final ProcessRequest allocateEvaluatorRequest = this.mockRuntime.getNextProcessRequest();
  assertEquals("allocate evaluator request", ProcessRequest.Type.ALLOCATE_EVALUATOR,
      allocateEvaluatorRequest.getType());
  final AllocatedEvaluator evaluator =
      ((ProcessRequestInternal<AllocatedEvaluator, Object>)allocateEvaluatorRequest)
          .getSuccessEvent();
  this.mockRuntime.succeed(allocateEvaluatorRequest);
  assertTrue("evaluator allocation succeeded",
      this.mockApplication.getAllocatedEvaluators().contains(evaluator));
  final String contextId = "foo";
  this.mockApplication.submitContext(evaluator, contextId);
  final ActiveContext rootContext = ((MockAllocatedEvaluator) evaluator).getRootContext();
  assertTrue("root context", rootContext != null);


  // submit a task
  this.mockApplication.submitTask(rootContext, "test-task");
  assertTrue("create task queued", this.mockRuntime.hasProcessRequest());
  final ProcessRequest createTaskRequest = this.mockRuntime.getNextProcessRequest();
  assertEquals("create task request", ProcessRequest.Type.CREATE_TASK,
      createTaskRequest.getType());
  final RunningTask task = (RunningTask) ((ProcessRequestInternal)createTaskRequest).getSuccessEvent();
  this.mockRuntime.succeed(createTaskRequest);
  assertTrue("task running", this.mockApplication.getRunningTasks().contains(task));

  // check task auto complete
  assertTrue("check for request", this.mockRuntime.hasProcessRequest());
  final ProcessRequestInternal completedTask =
      (ProcessRequestInternal) this.mockRuntime.getNextProcessRequest();
  assertEquals("complete task request", ProcessRequest.Type.COMPLETE_TASK,
      completedTask.getType());
  this.mockRuntime.succeed(completedTask);
  assertEquals("no running tasks", 0, this.mockApplication.getRunningTasks().size());
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  LOG.log(Level.INFO, "Submitting HelloREEF task to AllocatedEvaluator: {0}", allocatedEvaluator);
  final Configuration taskConfiguration = TaskConfiguration.CONF
      .set(TaskConfiguration.IDENTIFIER, "HelloREEFTask")
      .set(TaskConfiguration.TASK, HelloTask.class)
      .build();
  allocatedEvaluator.submitTask(taskConfiguration);
}
 

/**
 * This constructor should only be called by the AllocatedEvaluatorBridgeFactory.
 */
AllocatedEvaluatorBridge(final AllocatedEvaluator allocatedEvaluator,
                         final String serverInfo) {
  this.jallocatedEvaluator = allocatedEvaluator;
  this.evaluatorId = allocatedEvaluator.getId();
  this.nameServerInfo = serverInfo;
}
 

@Override
public void onNext(final AllocatedEvaluator eval) {
  final String taskId = "Task_" + eval.getId();
  LOG.log(Level.INFO, "Submit task: {0}", taskId);

  final Configuration taskConfig = TaskConfiguration.CONF
      .set(TaskConfiguration.IDENTIFIER, taskId)
      .set(TaskConfiguration.TASK, TaskMessagingTask.class)
      .set(TaskConfiguration.ON_MESSAGE, TaskMessagingTask.DriverMessageHandler.class)
      .set(TaskConfiguration.ON_SEND_MESSAGE, TaskMessagingTask.class)
      .build();
  eval.submitTask(taskConfig);
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  allocatedEvaluator.submitTask(TaskConfiguration.CONF
      .set(TaskConfiguration.TASK, ConfigurationProviderTestTask.class)
      .set(TaskConfiguration.IDENTIFIER, "ConfigurationProviderTestTask")
      .build());
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  final Configuration taskConfiguration = TaskConfiguration.CONF
      .set(TaskConfiguration.IDENTIFIER, "EvaluatorExitTestTask")
      .set(TaskConfiguration.TASK, EvaluatorExitTestTask.class)
      .build();
  allocatedEvaluator.submitTask(taskConfiguration);
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  LOG.log(Level.INFO, "Submitting an id context to AllocatedEvaluator: {0}", allocatedEvaluator);
  final Configuration contextConfiguration = ContextConfiguration.CONF
      .set(ContextConfiguration.IDENTIFIER, "BroadcastContext-" +
          BroadcastDriver.this.numberOfAllocatedEvaluators.getAndDecrement())
      .build();
  allocatedEvaluator.submitContext(contextConfiguration);
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  LOG.log(Level.INFO, "Submitting HelloREEF task to AllocatedEvaluator: {0}", allocatedEvaluator);
  final Configuration taskConfiguration = TaskConfiguration.CONF
      .set(TaskConfiguration.IDENTIFIER, "HelloREEFTask")
      .set(TaskConfiguration.TASK, HelloTask.class)
      .build();
  allocatedEvaluator.submitTask(taskConfiguration);
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  synchronized (HelloDriver.this) {
    if (HelloDriver.this.nJVMTasks > 0) {
      HelloDriver.this.onNextJVM(allocatedEvaluator);
      HelloDriver.this.nJVMTasks -= 1;
    } else if (HelloDriver.this.nCLRTasks > 0) {
      HelloDriver.this.onNextCLR(allocatedEvaluator);
      HelloDriver.this.nCLRTasks -= 1;
    }
  }
}
 

@Override
public void onNext(final AllocatedEvaluator evaluator) {
  LOG.log(Level.INFO, "Evaluator is ready");
  synchronized (SchedulerDriver.this) {
    nActiveEval++;
    nRequestedEval--;
  }

  evaluator.submitContext(ContextConfiguration.CONF
      .set(ContextConfiguration.IDENTIFIER, "SchedulerContext")
      .build());
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {

  final String actual =
      allocatedEvaluator.getEvaluatorDescriptor().getNodeDescriptor().getRackDescriptor().getName();
  if (!expectedRackName.equals(actual)) {
    throw new DriverSideFailure("The rack received is different that the expected one, received " + actual
        + " expected " + expectedRackName);
  }
  allocatedEvaluator.close();
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  synchronized (expectedRunningTaskIds) {
    final String taskId = "Task-" + numberOfTaskSubmissions.getAndDecrement();
    final Configuration taskConfiguration = getTaskConfiguration(taskId);
    allocatedEvaluator.submitTask(taskConfiguration);
    expectedRunningTaskIds.add(taskId);
  }
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  final Configuration contextConfiguration = ContextConfiguration.CONF
      .set(ContextConfiguration.IDENTIFIER, "FailingEvaluator")
      .set(ContextConfiguration.ON_CONTEXT_STARTED, FailureSchedulingContextStartHandler.class)
      .build();
  allocatedEvaluator.submitContext(contextConfiguration);
}
 

@Override
public Collection<AllocatedEvaluator> getCurrentAllocatedEvaluators() {
  if (this.clock.get().isClosed()) {
    throw new IllegalStateException("clock is closed");
  }
  return new ArrayList<AllocatedEvaluator>(this.allocatedEvaluatorMap.values());
}
 

@Override
public void onNext(final AllocatedEvaluator allocatedEvaluator) {
  if (isFirstEvaluator.compareAndSet(true, false)) {
    allocatedEvaluator.submitContext(getFailedContextConfiguration());
  } else {
    allocatedEvaluator.submitContext(ContextConfiguration.CONF
        .set(ContextConfiguration.IDENTIFIER, ROOT_CONTEXT_ID)
        .build());
  }
}
 

@Override
public Configuration getServiceConfiguration(final AllocatedEvaluator allocatedEvaluator) {

  try {

    final NumberedSplit<InputSplit> numberedSplit =
        this.evaluatorToPartitionStrategy.getInputSplit(
            allocatedEvaluator.getEvaluatorDescriptor().getNodeDescriptor(),
            allocatedEvaluator.getId());

    final Configuration serviceConfiguration = ServiceConfiguration.CONF
        .set(ServiceConfiguration.SERVICES,
            this.inMemory ? InMemoryInputFormatDataSet.class : InputFormatDataSet.class)
        .build();

    return Tang.Factory.getTang().newConfigurationBuilder(serviceConfiguration)
        .bindImplementation(
            DataSet.class,
            this.inMemory ? InMemoryInputFormatDataSet.class : InputFormatDataSet.class)
        .bindNamedParameter(JobConfExternalConstructor.InputFormatClass.class, inputFormatClass)
        .bindNamedParameter(JobConfExternalConstructor.InputPath.class, numberedSplit.getPath())
        .bindNamedParameter(
            InputSplitExternalConstructor.SerializedInputSplit.class,
            WritableSerializer.serialize(numberedSplit.getEntry()))
        .bindConstructor(InputSplit.class, InputSplitExternalConstructor.class)
        .bindConstructor(JobConf.class, JobConfExternalConstructor.class)
        .build();

  } catch (final BindException ex) {
    final String evalId = allocatedEvaluator.getId();
    final String msg = "Unable to create configuration for evaluator " + evalId;
    LOG.log(Level.WARNING, msg, ex);
    throw new RuntimeException(msg, ex);
  }
}