Examples with RunTaskCommand alluxio.grpc.RunTaskCommand used on opensource projects

Example 1 with RunTaskCommand

use of alluxio.grpc.RunTaskCommand in project alluxio by Alluxio.

the class TaskExecutorTest method runCompletion.

@Test
public void runCompletion() throws Exception {
    long jobId = 1;
    long taskId = 2;
    JobConfig jobConfig = mock(JobConfig.class);
    Serializable taskArgs = Lists.newArrayList(1);
    RunTaskContext context = mock(RunTaskContext.class);
    Integer taskResult = 1;
    @SuppressWarnings("unchecked") PlanDefinition<JobConfig, Serializable, Serializable> planDefinition = mock(PlanDefinition.class);
    when(mRegistry.getJobDefinition(any(JobConfig.class))).thenReturn(planDefinition);
    when(planDefinition.runTask(any(JobConfig.class), eq(taskArgs), any(RunTaskContext.class))).thenReturn(taskResult);
    RunTaskCommand command = RunTaskCommand.newBuilder().setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig))).setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs))).build();
    TaskExecutor executor = new TaskExecutor(jobId, taskId, command, context, mTaskExecutorManager);
    executor.run();
    verify(planDefinition).runTask(any(JobConfig.class), eq(taskArgs), eq(context));
    verify(mTaskExecutorManager).notifyTaskCompletion(jobId, taskId, taskResult);
}

Example 2 with RunTaskCommand

use of alluxio.grpc.RunTaskCommand in project alluxio by Alluxio.

the class TaskExecutorTest method runFailure.

@Test
public void runFailure() throws Exception {
    long jobId = 1;
    long taskId = 2;
    JobConfig jobConfig = new SleepJobConfig(10);
    Serializable taskArgs = Lists.newArrayList(1);
    RunTaskContext context = mock(RunTaskContext.class);
    @SuppressWarnings("unchecked") PlanDefinition<JobConfig, Serializable, Serializable> planDefinition = mock(PlanDefinition.class);
    when(mRegistry.getJobDefinition(eq(jobConfig))).thenReturn(planDefinition);
    when(planDefinition.runTask(eq(jobConfig), any(Serializable.class), any(RunTaskContext.class))).thenThrow(new UnsupportedOperationException("failure"));
    RunTaskCommand command = RunTaskCommand.newBuilder().setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig))).setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs))).build();
    TaskExecutor executor = new TaskExecutor(jobId, taskId, command, context, mTaskExecutorManager);
    executor.run();
    verify(mTaskExecutorManager).notifyTaskFailure(eq(jobId), eq(taskId), any());
}

Example 3 with RunTaskCommand

use of alluxio.grpc.RunTaskCommand in project alluxio by Alluxio.

the class CommandHandlingExecutorTest method heartbeat.

@Test
public void heartbeat() throws Exception {
    JobCommand.Builder command = JobCommand.newBuilder();
    RunTaskCommand.Builder runTaskCommand = RunTaskCommand.newBuilder();
    long jobId = 1;
    runTaskCommand.setJobId(jobId);
    long taskId = 2;
    runTaskCommand.setTaskId(taskId);
    JobConfig jobConfig = new TestPlanConfig("/test");
    runTaskCommand.setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig)));
    Serializable taskArgs = Lists.newArrayList(1);
    runTaskCommand.setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs)));
    command.setRunTaskCommand(runTaskCommand);
    Mockito.when(mJobMasterClient.heartbeat(any(JobWorkerHealth.class), eq(Lists.newArrayList()))).thenReturn(Lists.newArrayList(command.build()));
    mCommandHandlingExecutor.heartbeat();
    ExecutorService executorService = Whitebox.getInternalState(mCommandHandlingExecutor, "mCommandHandlingService");
    executorService.shutdown();
    Assert.assertTrue(executorService.awaitTermination(5000, TimeUnit.MILLISECONDS));
    Mockito.verify(mTaskExecutorManager).getAndClearTaskUpdates();
    Mockito.verify(mTaskExecutorManager).executeTask(Mockito.eq(jobId), Mockito.eq(taskId), Mockito.eq(runTaskCommand.build()), Mockito.any(RunTaskContext.class));
}