public MockTaskAttemptListenerImpl(AppContext context, JobTokenSecretManager jobTokenSecretManager, RMHeartbeatHandler rmHeartbeatHandler, TaskHeartbeatHandler hbHandler) { super(context, jobTokenSecretManager, rmHeartbeatHandler, null); this.taskHeartbeatHandler = hbHandler; }
public MockTaskAttemptListenerImpl(AppContext context, JobTokenSecretManager jobTokenSecretManager, RMHeartbeatHandler rmHeartbeatHandler, TaskHeartbeatHandler hbHandler, AMPreemptionPolicy policy) { super(context, jobTokenSecretManager, rmHeartbeatHandler, policy); this.taskHeartbeatHandler = hbHandler; }
public MockTaskAttemptListenerImpl(AppContext context, JobTokenSecretManager jobTokenSecretManager, RMHeartbeatHandler rmHeartbeatHandler, TaskHeartbeatHandler hbHandler) { super(context, jobTokenSecretManager, rmHeartbeatHandler); this.taskHeartbeatHandler = hbHandler; }
protected void registerHeartbeatHandler(Configuration conf) { taskHeartbeatHandler = new TaskHeartbeatHandler(context.getEventHandler(), context.getClock(), conf.getInt(MRJobConfig.MR_AM_TASK_LISTENER_THREAD_COUNT, MRJobConfig.DEFAULT_MR_AM_TASK_LISTENER_THREAD_COUNT)); addService(taskHeartbeatHandler); }
@Test (timeout=5000) public void testGetTask() throws IOException { AppContext appCtx = mock(AppContext.class); JobTokenSecretManager secret = mock(JobTokenSecretManager.class); RMHeartbeatHandler rmHeartbeatHandler = mock(RMHeartbeatHandler.class); TaskHeartbeatHandler hbHandler = mock(TaskHeartbeatHandler.class); MockTaskAttemptListenerImpl listener = new MockTaskAttemptListenerImpl(appCtx, secret, rmHeartbeatHandler, hbHandler); Configuration conf = new Configuration(); listener.init(conf); listener.start(); JVMId id = new JVMId("foo",1, true, 1); WrappedJvmID wid = new WrappedJvmID(id.getJobId(), id.isMap, id.getId()); // Verify ask before registration. //The JVM ID has not been registered yet so we should kill it. JvmContext context = new JvmContext(); context.jvmId = id; JvmTask result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); // Verify ask after registration but before launch. // Don't kill, should be null. TaskAttemptId attemptID = mock(TaskAttemptId.class); Task task = mock(Task.class); //Now put a task with the ID listener.registerPendingTask(task, wid); result = listener.getTask(context); assertNull(result); // Unregister for more testing. listener.unregister(attemptID, wid); // Verify ask after registration and launch //Now put a task with the ID listener.registerPendingTask(task, wid); listener.registerLaunchedTask(attemptID, wid); verify(hbHandler).register(attemptID); result = listener.getTask(context); assertNotNull(result); assertFalse(result.shouldDie); // Don't unregister yet for more testing. //Verify that if we call it again a second time we are told to die. result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.unregister(attemptID, wid); // Verify after unregistration. result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.stop(); // test JVMID JVMId jvmid = JVMId.forName("jvm_001_002_m_004"); assertNotNull(jvmid); try { JVMId.forName("jvm_001_002_m_004_006"); Assert.fail(); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "TaskId string : jvm_001_002_m_004_006 is not properly formed"); } }
@Test (timeout=5000) public void testGetTask() throws IOException { AppContext appCtx = mock(AppContext.class); JobTokenSecretManager secret = mock(JobTokenSecretManager.class); RMHeartbeatHandler rmHeartbeatHandler = mock(RMHeartbeatHandler.class); TaskHeartbeatHandler hbHandler = mock(TaskHeartbeatHandler.class); Dispatcher dispatcher = mock(Dispatcher.class); EventHandler ea = mock(EventHandler.class); when(dispatcher.getEventHandler()).thenReturn(ea); when(appCtx.getEventHandler()).thenReturn(ea); CheckpointAMPreemptionPolicy policy = new CheckpointAMPreemptionPolicy(); policy.init(appCtx); MockTaskAttemptListenerImpl listener = new MockTaskAttemptListenerImpl(appCtx, secret, rmHeartbeatHandler, hbHandler, policy); Configuration conf = new Configuration(); listener.init(conf); listener.start(); JVMId id = new JVMId("foo",1, true, 1); WrappedJvmID wid = new WrappedJvmID(id.getJobId(), id.isMap, id.getId()); // Verify ask before registration. //The JVM ID has not been registered yet so we should kill it. JvmContext context = new JvmContext(); context.jvmId = id; JvmTask result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); // Verify ask after registration but before launch. // Don't kill, should be null. TaskAttemptId attemptID = mock(TaskAttemptId.class); Task task = mock(Task.class); //Now put a task with the ID listener.registerPendingTask(task, wid); result = listener.getTask(context); assertNull(result); // Unregister for more testing. listener.unregister(attemptID, wid); // Verify ask after registration and launch //Now put a task with the ID listener.registerPendingTask(task, wid); listener.registerLaunchedTask(attemptID, wid); verify(hbHandler).register(attemptID); result = listener.getTask(context); assertNotNull(result); assertFalse(result.shouldDie); // Don't unregister yet for more testing. //Verify that if we call it again a second time we are told to die. result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.unregister(attemptID, wid); // Verify after unregistration. result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.stop(); // test JVMID JVMId jvmid = JVMId.forName("jvm_001_002_m_004"); assertNotNull(jvmid); try { JVMId.forName("jvm_001_002_m_004_006"); fail(); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "TaskId string : jvm_001_002_m_004_006 is not properly formed"); } }
@SuppressWarnings("rawtypes") @Test public void testStatusUpdateProgress() throws IOException, InterruptedException { AppContext appCtx = mock(AppContext.class); JobTokenSecretManager secret = mock(JobTokenSecretManager.class); RMHeartbeatHandler rmHeartbeatHandler = mock(RMHeartbeatHandler.class); TaskHeartbeatHandler hbHandler = mock(TaskHeartbeatHandler.class); Dispatcher dispatcher = mock(Dispatcher.class); EventHandler ea = mock(EventHandler.class); when(dispatcher.getEventHandler()).thenReturn(ea); when(appCtx.getEventHandler()).thenReturn(ea); CheckpointAMPreemptionPolicy policy = new CheckpointAMPreemptionPolicy(); policy.init(appCtx); MockTaskAttemptListenerImpl listener = new MockTaskAttemptListenerImpl(appCtx, secret, rmHeartbeatHandler, hbHandler, policy); Configuration conf = new Configuration(); listener.init(conf); listener.start(); JVMId id = new JVMId("foo",1, true, 1); WrappedJvmID wid = new WrappedJvmID(id.getJobId(), id.isMap, id.getId()); TaskAttemptID attemptID = new TaskAttemptID("1", 1, TaskType.MAP, 1, 1); TaskAttemptId attemptId = TypeConverter.toYarn(attemptID); Task task = mock(Task.class); listener.registerPendingTask(task, wid); listener.registerLaunchedTask(attemptId, wid); verify(hbHandler).register(attemptId); // make sure a ping doesn't report progress AMFeedback feedback = listener.statusUpdate(attemptID, null); assertTrue(feedback.getTaskFound()); verify(hbHandler, never()).progressing(eq(attemptId)); // make sure a status update does report progress MapTaskStatus mockStatus = new MapTaskStatus(attemptID, 0.0f, 1, TaskStatus.State.RUNNING, "", "RUNNING", "", TaskStatus.Phase.MAP, new Counters()); feedback = listener.statusUpdate(attemptID, mockStatus); assertTrue(feedback.getTaskFound()); verify(hbHandler).progressing(eq(attemptId)); listener.close(); }