use of org.apache.tez.common.security.JobTokenIdentifier in project hive by apache.
the class GenericUDTFGetSplits method getSplits.
public InputSplit[] getSplits(JobConf job, int numSplits, TezWork work, Schema schema, ApplicationId applicationId) throws IOException {
DAG dag = DAG.create(work.getName());
dag.setCredentials(job.getCredentials());
DagUtils utils = DagUtils.getInstance();
Context ctx = new Context(job);
MapWork mapWork = (MapWork) work.getAllWork().get(0);
// bunch of things get setup in the context based on conf but we need only the MR tmp directory
// for the following method.
JobConf wxConf = utils.initializeVertexConf(job, ctx, mapWork);
// TODO: should we also whitelist input formats here? from mapred.input.format.class
Path scratchDir = utils.createTezDir(ctx.getMRScratchDir(), job);
FileSystem fs = scratchDir.getFileSystem(job);
try {
LocalResource appJarLr = createJarLocalResource(utils.getExecJarPathLocal(ctx.getConf()), utils, job);
LlapCoordinator coordinator = LlapCoordinator.getInstance();
if (coordinator == null) {
throw new IOException("LLAP coordinator is not initialized; must be running in HS2 with " + ConfVars.LLAP_HS2_ENABLE_COORDINATOR.varname + " enabled");
}
// Update the queryId to use the generated applicationId. See comment below about
// why this is done.
HiveConf.setVar(wxConf, HiveConf.ConfVars.HIVEQUERYID, applicationId.toString());
Vertex wx = utils.createVertex(wxConf, mapWork, scratchDir, fs, ctx, false, work, work.getVertexType(mapWork), DagUtils.createTezLrMap(appJarLr, null));
String vertexName = wx.getName();
dag.addVertex(wx);
utils.addCredentials(mapWork, dag);
// we have the dag now proceed to get the splits:
Preconditions.checkState(HiveConf.getBoolVar(wxConf, ConfVars.HIVE_TEZ_GENERATE_CONSISTENT_SPLITS));
Preconditions.checkState(HiveConf.getBoolVar(wxConf, ConfVars.LLAP_CLIENT_CONSISTENT_SPLITS));
HiveSplitGenerator splitGenerator = new HiveSplitGenerator(wxConf, mapWork);
List<Event> eventList = splitGenerator.initialize();
InputSplit[] result = new InputSplit[eventList.size() - 1];
InputConfigureVertexTasksEvent configureEvent = (InputConfigureVertexTasksEvent) eventList.get(0);
List<TaskLocationHint> hints = configureEvent.getLocationHint().getTaskLocationHints();
Preconditions.checkState(hints.size() == eventList.size() - 1);
if (LOG.isDebugEnabled()) {
LOG.debug("NumEvents=" + eventList.size() + ", NumSplits=" + result.length);
}
// This assumes LLAP cluster owner is always the HS2 user.
String llapUser = UserGroupInformation.getLoginUser().getShortUserName();
String queryUser = null;
byte[] tokenBytes = null;
LlapSigner signer = null;
if (UserGroupInformation.isSecurityEnabled()) {
signer = coordinator.getLlapSigner(job);
// 1. Generate the token for query user (applies to all splits).
queryUser = SessionState.getUserFromAuthenticator();
if (queryUser == null) {
queryUser = UserGroupInformation.getCurrentUser().getUserName();
LOG.warn("Cannot determine the session user; using " + queryUser + " instead");
}
LlapTokenLocalClient tokenClient = coordinator.getLocalTokenClient(job, llapUser);
// We put the query user, not LLAP user, into the message and token.
Token<LlapTokenIdentifier> token = tokenClient.createToken(applicationId.toString(), queryUser, true);
LOG.info("Created the token for remote user: {}", token);
bos.reset();
token.write(dos);
tokenBytes = bos.toByteArray();
} else {
queryUser = UserGroupInformation.getCurrentUser().getUserName();
}
// Generate umbilical token (applies to all splits)
Token<JobTokenIdentifier> umbilicalToken = JobTokenCreator.createJobToken(applicationId);
LOG.info("Number of splits: " + (eventList.size() - 1));
SignedMessage signedSvs = null;
for (int i = 0; i < eventList.size() - 1; i++) {
TaskSpec taskSpec = new TaskSpecBuilder().constructTaskSpec(dag, vertexName, eventList.size() - 1, applicationId, i);
// 2. Generate the vertex/submit information for all events.
if (i == 0) {
// The queryId could either be picked up from the current request being processed, or
// generated. The current request isn't exactly correct since the query is 'done' once we
// return the results. Generating a new one has the added benefit of working once this
// is moved out of a UDTF into a proper API.
// Setting this to the generated AppId which is unique.
// Despite the differences in TaskSpec, the vertex spec should be the same.
signedSvs = createSignedVertexSpec(signer, taskSpec, applicationId, queryUser, applicationId.toString());
}
SubmitWorkInfo submitWorkInfo = new SubmitWorkInfo(applicationId, System.currentTimeMillis(), taskSpec.getVertexParallelism(), signedSvs.message, signedSvs.signature, umbilicalToken);
byte[] submitWorkBytes = SubmitWorkInfo.toBytes(submitWorkInfo);
// 3. Generate input event.
SignedMessage eventBytes = makeEventBytes(wx, vertexName, eventList.get(i + 1), signer);
// 4. Make location hints.
SplitLocationInfo[] locations = makeLocationHints(hints.get(i));
result[i] = new LlapInputSplit(i, submitWorkBytes, eventBytes.message, eventBytes.signature, locations, schema, llapUser, tokenBytes);
}
return result;
} catch (Exception e) {
throw new IOException(e);
}
}
use of org.apache.tez.common.security.JobTokenIdentifier in project tez by apache.
the class TestShuffleManager method createShuffleManager.
private ShuffleManagerForTest createShuffleManager(InputContext inputContext, int expectedNumOfPhysicalInputs) throws IOException {
Path outDirBase = new Path(".", "outDir");
String[] outDirs = new String[] { outDirBase.toString() };
doReturn(outDirs).when(inputContext).getWorkDirs();
conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, inputContext.getWorkDirs());
DataOutputBuffer out = new DataOutputBuffer();
Token<JobTokenIdentifier> token = new Token<JobTokenIdentifier>(new JobTokenIdentifier(), new JobTokenSecretManager(null));
token.write(out);
doReturn(ByteBuffer.wrap(out.getData())).when(inputContext).getServiceConsumerMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
return new ShuffleManagerForTest(inputContext, conf, expectedNumOfPhysicalInputs, 1024, false, -1, null, inputAllocator);
}
use of org.apache.tez.common.security.JobTokenIdentifier in project tez by apache.
the class TestShuffle method createTezInputContext.
private InputContext createTezInputContext() throws IOException {
ApplicationId applicationId = ApplicationId.newInstance(1, 1);
InputContext inputContext = mock(InputContext.class);
doReturn(applicationId).when(inputContext).getApplicationId();
doReturn("sourceVertex").when(inputContext).getSourceVertexName();
when(inputContext.getCounters()).thenReturn(new TezCounters());
ExecutionContext executionContext = new ExecutionContextImpl("localhost");
doReturn(executionContext).when(inputContext).getExecutionContext();
ByteBuffer shuffleBuffer = ByteBuffer.allocate(4).putInt(0, 4);
doReturn(shuffleBuffer).when(inputContext).getServiceProviderMetaData(anyString());
Token<JobTokenIdentifier> sessionToken = new Token<JobTokenIdentifier>(new JobTokenIdentifier(new Text("text")), new JobTokenSecretManager());
ByteBuffer tokenBuffer = TezCommonUtils.serializeServiceData(sessionToken);
doReturn(tokenBuffer).when(inputContext).getServiceConsumerMetaData(anyString());
when(inputContext.createTezFrameworkExecutorService(anyInt(), anyString())).thenAnswer(new Answer<ExecutorService>() {
@Override
public ExecutorService answer(InvocationOnMock invocation) throws Throwable {
return sharedExecutor.createExecutorService(invocation.getArgumentAt(0, Integer.class), invocation.getArgumentAt(1, String.class));
}
});
return inputContext;
}
use of org.apache.tez.common.security.JobTokenIdentifier in project tez by apache.
the class TestShuffleInputEventHandlerImpl method createShuffleManager.
private ShuffleManager createShuffleManager(InputContext inputContext) throws IOException {
Path outDirBase = new Path(".", "outDir");
String[] outDirs = new String[] { outDirBase.toString() };
doReturn(outDirs).when(inputContext).getWorkDirs();
conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, inputContext.getWorkDirs());
DataOutputBuffer out = new DataOutputBuffer();
Token<JobTokenIdentifier> token = new Token<JobTokenIdentifier>(new JobTokenIdentifier(), new JobTokenSecretManager(null));
token.write(out);
doReturn(ByteBuffer.wrap(out.getData())).when(inputContext).getServiceConsumerMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
ShuffleManager realShuffleManager = new ShuffleManager(inputContext, conf, 2, 1024, false, -1, null, inputAllocator);
ShuffleManager shuffleManager = spy(realShuffleManager);
return shuffleManager;
}
use of org.apache.tez.common.security.JobTokenIdentifier in project tez by apache.
the class TezClientUtils method createSessionToken.
static void createSessionToken(String tokenIdentifier, JobTokenSecretManager jobTokenSecretManager, Credentials credentials) {
JobTokenIdentifier identifier = new JobTokenIdentifier(new Text(tokenIdentifier));
Token<JobTokenIdentifier> sessionToken = new Token<JobTokenIdentifier>(identifier, jobTokenSecretManager);
sessionToken.setService(identifier.getJobId());
TokenCache.setSessionToken(sessionToken, credentials);
}
Aggregations