@VisibleForTesting public WmTezSession getSession( TezSessionState session, MappingInput input, HiveConf conf) throws Exception { return getSession(session, input, conf, null); }
@Override public WmTezSession getSession( TezSessionState session, MappingInput input, HiveConf conf, final WmContext wmContext) throws Exception { // We want to wait for the iteration to finish and set the cluster fraction. WmTezSession state = super.getSession(session, input, conf, null); ensureWm(); return state; }
@Override public void run() { WmTezSession old = session.get(); session.set(null); if (cdl != null) { cdl.countDown(); } LOG.info("About to call get with " + old); try { session.set((WmTezSession) wm.getSession(old, mappingInput(userName), conf)); LOG.info("Received " + session.get()); } catch (Throwable e) { error.compareAndSet(null, e); } } }
public static TezSessionState getSession(TezSessionState session, HiveConf conf, MappingInput input, boolean isUnmanagedLlapMode, WmContext wmContext) throws Exception { Set<String> desiredCounters = new HashSet<>(); // 1. If WM is not present just go to unmanaged. WorkloadManager wm = WorkloadManager.getInstance(); if (wm == null) { LOG.debug("Using unmanaged session - WM is not initialized"); return getUnmanagedSession(session, conf, desiredCounters, isUnmanagedLlapMode, wmContext); } // 2. We will ask WM for a preliminary mapping. This allows us to escape to the unmanaged path // quickly in the common case. It's still possible that resource plan will be updated and // our preliminary mapping won't work out. We'll handle that below. if (!wm.isManaged(input)) { LOG.info("Using unmanaged session - no mapping for " + input); return getUnmanagedSession(session, conf, desiredCounters, isUnmanagedLlapMode, wmContext); } // 3. Finally, try WM. try { // Note: this may just block to wait for a session based on parallelism. LOG.info("Getting a WM session for " + input); WmTezSession result = wm.getSession(session, input, conf, wmContext); result.setWmContext(wmContext); wm.updateTriggers(result); return result; } catch (WorkloadManager.NoPoolMappingException ex) { LOG.info("NoPoolMappingException thrown. Getting an un-managed session"); return getUnmanagedSession(session, conf, desiredCounters, isUnmanagedLlapMode, wmContext); } }
private static void verifyMapping( WorkloadManager wm, HiveConf conf, MappingInput mi, String result) throws Exception { WmTezSession session = (WmTezSession) wm.getSession(null, mi, conf, null); assertEquals(result, session.getPoolName()); session.returnToSessionManager(); }
@Test(timeout = 10000) public void testReuse() throws Exception { HiveConf conf = createConf(); MockQam qam = new MockQam(); WorkloadManager wm = new WorkloadManagerForTest("test", conf, 1, qam); wm.start(); TezSessionState nonPool = mock(TezSessionState.class); when(nonPool.getConf()).thenReturn(conf); doNothing().when(nonPool).close(anyBoolean()); TezSessionState session = wm.getSession(nonPool, mappingInput("user"), conf); verify(nonPool).close(anyBoolean()); assertNotSame(nonPool, session); session.returnToSessionManager(); TezSessionPoolSession diffPool = mock(TezSessionPoolSession.class); when(diffPool.getConf()).thenReturn(conf); doNothing().when(diffPool).returnToSessionManager(); session = wm.getSession(diffPool, mappingInput("user"), conf); verify(diffPool).returnToSessionManager(); assertNotSame(diffPool, session); TezSessionState session2 = wm.getSession(session, mappingInput("user"), conf); assertSame(session, session2); }
final WorkloadManager wm = new WorkloadManagerForTest("test", conf, qam, plan); wm.start(); WmTezSession sessionA1 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf), sessionA2 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf), sessionB1 = (WmTezSession) wm.getSession(null, mappingInput("B"), conf); final AtomicReference<WmTezSession> sessionA3 = new AtomicReference<>(), sessionA4 = new AtomicReference<>(); checkError(error); WmTezSession sessionB2 = (WmTezSession) wm.getSession(null, mappingInput("B"), conf); sessionB1.returnToSessionManager(); sessionB2.returnToSessionManager();
final WorkloadManager wm = new WorkloadManagerForTest("test", conf, 2, qam); wm.start(); WmTezSession session1 = (WmTezSession) wm.getSession( null, mappingInput("user"), conf); WmTezSession session1a = (WmTezSession) wm.getSession( session1, mappingInput("user"), conf); assertSame(session1, session1a); assertEquals(1.0, session1.getClusterFraction(), EPSILON); WmTezSession session2 = (WmTezSession) wm.getSession( null, mappingInput("user"), conf); WmTezSession session2a = (WmTezSession) wm.getSession( session2, mappingInput("user"), conf); assertSame(session2, session2a);
WmTezSession sessionA1 = (WmTezSession) wm.getSession( null, mappingInput("A", null), conf, null); final AtomicReference<WmTezSession> sessionA2 = new AtomicReference<>(); TezSessionState r = wm.getSession(null, mappingInput("A", null), conf, null); fail("Expected an error but got " + r); } catch (WorkloadManager.NoPoolMappingException ex) { sessionA1 = (WmTezSession) wm.getSession(null, mappingInput("A", null), conf, null); assertEquals("A", sessionA1.getPoolName()); sessionA1.returnToSessionManager();
@Test(timeout = 10000) public void testQueueName() throws Exception { HiveConf conf = createConf(); MockQam qam = new MockQam(); WorkloadManager wm = new WorkloadManagerForTest("test", conf, 1, qam); wm.start(); // The queue should be ignored. conf.set(TezConfiguration.TEZ_QUEUE_NAME, "test2"); TezSessionState session = wm.getSession(null, mappingInput("user"), conf); assertEquals("test", session.getQueueName()); assertEquals("test", conf.get(TezConfiguration.TEZ_QUEUE_NAME)); session.setQueueName("test2"); session = wm.getSession(session, mappingInput("user"), conf); assertEquals("test", session.getQueueName()); }
WmTezSession sessionA1 = (WmTezSession) wm.getSession(null, mappingInput("U"), conf); assertEquals("A", sessionA1.getPoolName()); assertEquals(0.5f, sessionA1.getClusterFraction(), EPSILON); WmTezSession sessionB1 = (WmTezSession) wm.getSession(null, mappingInput("U"), conf); assertEquals("B", sessionB1.getPoolName()); assertEquals(0.4f, sessionB1.getClusterFraction(), EPSILON);
assertEquals(5, wm.getNumSessions()); WmTezSession session05of06 = (WmTezSession) wm.getSession( null, mappingInput("p1"), conf); assertEquals(0.3, session05of06.getClusterFraction(), EPSILON); WmTezSession session03of06 = (WmTezSession) wm.getSession( null, mappingInput("p2"), conf); assertEquals(0.18, session03of06.getClusterFraction(), EPSILON); WmTezSession session03of06_2 = (WmTezSession) wm.getSession( null, mappingInput("p2"), conf); assertEquals(0.09, session03of06.getClusterFraction(), EPSILON); assertEquals(0.09, session03of06_2.getClusterFraction(), EPSILON); WmTezSession session02of06 = (WmTezSession) wm.getSession( null,mappingInput("r1"), conf); assertEquals(0.12, session02of06.getClusterFraction(), EPSILON); WmTezSession session04 = (WmTezSession) wm.getSession( null, mappingInput("r2"), conf); assertEquals(0.4, session04.getClusterFraction(), EPSILON);
@Test(timeout=10000) public void testReuseWithDifferentPool() throws Exception { final HiveConf conf = createConf(); MockQam qam = new MockQam(); WMFullResourcePlan plan = new WMFullResourcePlan(plan(), Lists.newArrayList( pool("A", 2, 0.6f), pool("B", 1, 0.4f))); plan.setMappings(Lists.newArrayList(mapping("A", "A"), mapping("B", "B"))); final WorkloadManager wm = new WorkloadManagerForTest("test", conf, qam, plan); wm.start(); WmTezSession sessionA1 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf), sessionA2 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf); assertEquals("A", sessionA1.getPoolName()); assertEquals(0.3f, sessionA1.getClusterFraction(), EPSILON); assertEquals("A", sessionA2.getPoolName()); assertEquals(0.3f, sessionA2.getClusterFraction(), EPSILON); WmTezSession sessionB1 = (WmTezSession) wm.getSession(sessionA1, mappingInput("B"), conf); assertSame(sessionA1, sessionB1); assertEquals("B", sessionB1.getPoolName()); assertEquals(0.4f, sessionB1.getClusterFraction(), EPSILON); assertEquals(0.6f, sessionA2.getClusterFraction(), EPSILON); // A1 removed from A. // Make sure that we can still get a session from A. WmTezSession sessionA3 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf); assertEquals("A", sessionA3.getPoolName()); assertEquals(0.3f, sessionA3.getClusterFraction(), EPSILON); assertEquals(0.3f, sessionA3.getClusterFraction(), EPSILON); sessionA3.returnToSessionManager(); sessionB1.returnToSessionManager(); sessionA2.returnToSessionManager(); }
WmTezSession sessionA1 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf), sessionB1 = (WmTezSession) wm.getSession(null, mappingInput("B"), conf), sessionB2 = (WmTezSession) wm.getSession(null, mappingInput("B"), conf), sessionC1 = (WmTezSession) wm.getSession(null, mappingInput("C"), conf), sessionD1 = (WmTezSession) wm.getSession(null, mappingInput("D"), conf); final AtomicReference<WmTezSession> sessionA2 = new AtomicReference<>(), sessionD2 = new AtomicReference<>();
@Test(timeout=10000) public void testClusterChange() throws Exception { final HiveConf conf = createConf(); MockQam qam = new MockQam(); WMFullResourcePlan plan = new WMFullResourcePlan(plan(), Lists.newArrayList(pool("A", 2, 1f))); plan.getPlan().setDefaultPoolPath("A"); final WorkloadManager wm = new WorkloadManagerForTest("test", conf, qam, plan); wm.start(); WmTezSession session1 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf), session2 = (WmTezSession) wm.getSession(null, mappingInput("A"), conf); assertEquals(0.5, session1.getClusterFraction(), EPSILON); assertEquals(0.5, session2.getClusterFraction(), EPSILON); qam.assertWasCalledAndReset(); // If cluster info changes, qam should be called with the same fractions. wm.notifyOfClusterStateChange(); assertEquals(0.5, session1.getClusterFraction(), EPSILON); assertEquals(0.5, session2.getClusterFraction(), EPSILON); qam.assertWasCalledAndReset(); session1.returnToSessionManager(); session2.returnToSessionManager(); }
@Test(timeout = 10000) public void testDestroyAndReturn() throws Exception { // Session should not be lost; however the fraction should be discarded. HiveConf conf = createConf(); MockQam qam = new MockQam(); WorkloadManager wm = new WorkloadManagerForTest("test", conf, 2, qam); wm.start(); WmTezSession session = (WmTezSession) wm.getSession(null, mappingInput("user"), conf); assertEquals(1.0, session.getClusterFraction(), EPSILON); qam.assertWasCalledAndReset(); WmTezSession session2 = (WmTezSession) wm.getSession(null, mappingInput("user"), conf); assertEquals(0.5, session.getClusterFraction(), EPSILON); assertEquals(0.5, session2.getClusterFraction(), EPSILON); qam.assertWasCalledAndReset(); assertNotSame(session, session2); session.destroy(); // Destroy before returning to the pool. assertEquals(1.0, session2.getClusterFraction(), EPSILON); assertFalse(session.hasClusterFraction()); qam.assertWasCalledAndReset(); // We never lose pool session, so we should still be able to get. session = (WmTezSession) wm.getSession(null, mappingInput("user"), conf); session.returnToSessionManager(); assertEquals(1.0, session2.getClusterFraction(), EPSILON); assertFalse(session.hasClusterFraction()); qam.assertWasCalledAndReset(); }
@Test(timeout = 10000) public void testReopen() throws Exception { // We should always get a different object, and cluster fraction should be propagated. HiveConf conf = createConf(); MockQam qam = new MockQam(); WorkloadManager wm = new WorkloadManagerForTest("test", conf, 1, qam); wm.start(); WmTezSession session = (WmTezSession) wm.getSession( null, mappingInput("user"), conf); assertEquals(1.0, session.getClusterFraction(), EPSILON); qam.assertWasCalledAndReset(); WmTezSession session2 = (WmTezSession) session.reopen(); assertNotSame(session, session2); wm.addTestEvent().get(); assertEquals(session2.toString(), 1.0, session2.getClusterFraction(), EPSILON); assertFalse(session.hasClusterFraction()); qam.assertWasCalledAndReset(); }
WmTezSession sessionA1 = (WmTezSession) wm.getSession( null, mappingInput("A", null), conf, null), sessionA2 = (WmTezSession) wm.getSession(null, mappingInput("A", null), conf, null); assertEquals(0.5f, sessionA1.getClusterFraction(), EPSILON); assertEquals(0.5f, sessionA2.getClusterFraction(), EPSILON); WmTezSession sessionA3 = (WmTezSession) wm.getSession( null, mappingInput("A", null), conf, null); assertEquals(0f, sessionA3.getClusterFraction(), EPSILON); WmTezSession sessionA4 = (WmTezSession) wm.getSession( sessionA2, mappingInput("A", null), conf, null); assertSame(sessionA2, sessionA4);
private SampleTezSessionState validatePoolAfterCleanup( SampleTezSessionState oldSession, HiveConf conf, WorkloadManager wm, TezSessionPool<WmTezSession> pool, String sessionPoolName) throws Exception { // Make sure the cleanup doesn't leave the pool without a session. SampleTezSessionState theOnlySession = (SampleTezSessionState) pool.getSession(); assertNotNull(theOnlySession); theOnlySession.setWaitForAmRegistryFuture(null); assertNull(oldSession.getPoolName()); assertFalse(oldSession.hasClusterFraction()); pool.returnSession(theOnlySession); // Make sure we can actually get a session still - parallelism/etc. should not be affected. WmTezSession result = (WmTezSession) wm.getSession(null, mappingInput("A"), conf); assertEquals(sessionPoolName, result.getPoolName()); assertEquals(1f, result.getClusterFraction(), EPSILON); result.returnToSessionManager(); return theOnlySession; }
TezSessionState r = wm.getSession( null, mappingInput("u0", groups("g0", "g1"), "u2"), conf); fail("Expected failure, but got " + r); TezSessionState r = wm.getSession( null, mappingInput("u0", groups("g0", "g1"), "zzz"), conf); fail("Expected failure, but got " + r);