/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.tez.mapreduce.examples;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.util.ToolRunner;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.apache.hadoop.yarn.api.records.LocalResourceType;
import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.apache.tez.client.TezClient;
import org.apache.tez.dag.api.DAG;
import org.apache.tez.dag.api.ProcessorDescriptor;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.api.TezException;
import org.apache.tez.dag.api.UserPayload;
import org.apache.tez.dag.api.Vertex;
import org.apache.tez.examples.TezExampleBase;
import org.apache.tez.runtime.api.ProcessorContext;
import org.apache.tez.runtime.library.processor.SimpleProcessor;
import org.apache.tez.util.StopWatch;
public class RPCLoadGen extends TezExampleBase {
private static final Logger LOG = LoggerFactory.getLogger(RPCLoadGen.class);
private static final String VIA_RPC = "viaRpc";
private static final byte VIA_RPC_BYTE = (byte) 0x00;
private static final String VIA_HDFS_DIST_CACHE = "viaHdfsDistCache";
private static final byte VIA_HDFS_DIST_CACHE_BYTE = (byte) 0x01;
private static final String VIA_HDFS_DIRECT_READ = "viaHdfsDirectRead";
private static final byte VIA_HDFS_DIRECT_READ_BYTE = (byte) 0x02;
private static final Random random = new Random();
private static final String DISK_PAYLOAD_NAME = RPCLoadGen.class.getSimpleName() + "_payload";
private FileSystem fs;
private Path resourcePath;
@Override
protected final int runJob(String[] args, TezConfiguration tezConf, TezClient tezClient) throws
TezException, InterruptedException, IOException {
LOG.info("Running: " +
this.getClass().getSimpleName());
String mode = VIA_RPC;
if (args.length == 4) {
if (args[3].equals(VIA_RPC) || args[3].equals(VIA_HDFS_DIRECT_READ) ||
args[3].equals(VIA_HDFS_DIST_CACHE)) {
mode = args[3];
} else {
printUsage();
return 2;
}
}
int numTasks = Integer.parseInt(args[0]);
int maxSleepTimeMillis = Integer.parseInt(args[1]);
int payloadSizeBytes = Integer.parseInt(args[2]);
LOG.info("Parameters: numTasks=" + numTasks + ", maxSleepTime(ms)=" + maxSleepTimeMillis +
", payloadSize(bytes)=" + payloadSizeBytes + ", mode=" + mode);
DAG dag = createDAG(tezConf, numTasks, maxSleepTimeMillis, payloadSizeBytes, mode);
try {
return runDag(dag, false, LOG);
} finally {
if (fs != null) {
if (resourcePath != null) {
fs.delete(resourcePath, false);
}
}
}
}
@Override
protected void printUsage() {
System.err.println(
"Usage: " + "RPCLoadGen <numTasks> <max_sleep_time_millis> <get_task_payload_size> [" +
"<" + VIA_RPC + ">|" + VIA_HDFS_DIST_CACHE + "|" + VIA_HDFS_DIRECT_READ + "]");
ToolRunner.printGenericCommandUsage(System.err);
}
@Override
protected final int validateArgs(String[] otherArgs) {
return (otherArgs.length >=3 && otherArgs.length <=4) ? 0 : 2;
}
private DAG createDAG(TezConfiguration conf, int numTasks, int maxSleepTimeMillis,
int payloadSize, String mode) throws IOException {
Map<String, LocalResource> localResourceMap = new HashMap<String, LocalResource>();
UserPayload payload =
createUserPayload(conf, maxSleepTimeMillis, payloadSize, mode, localResourceMap);
Vertex vertex = Vertex.create("RPCLoadVertex",
ProcessorDescriptor.create(RPCSleepProcessor.class.getName()).setUserPayload(
payload), numTasks).addTaskLocalFiles(localResourceMap);
return DAG.create("RPCLoadGen").addVertex(vertex);
}
private UserPayload createUserPayload(TezConfiguration conf, int maxSleepTimeMillis,
int payloadSize, String mode,
Map<String, LocalResource> localResources) throws
IOException {
ByteBuffer payload;
if (mode.equals(VIA_RPC)) {
if (payloadSize < 5) {
payloadSize = 5; // To Configure the processor
}
byte[] payloadBytes = new byte[payloadSize];
random.nextBytes(payloadBytes);
payload = ByteBuffer.wrap(payloadBytes);
payload.put(4, VIA_RPC_BYTE); // ViaRPC
} else {
// Actual payload
byte[] payloadBytes = new byte[5];
payload = ByteBuffer.wrap(payloadBytes);
// Disk payload
byte[] diskPayload = new byte[payloadSize];
random.nextBytes(diskPayload);
fs = FileSystem.get(conf);
resourcePath = new Path(Path.SEPARATOR + "tmp", DISK_PAYLOAD_NAME);
resourcePath = fs.makeQualified(resourcePath);
FSDataOutputStream dataOut = fs.create(resourcePath, true);
dataOut.write(diskPayload);
dataOut.close();
fs.setReplication(resourcePath, (short)10);
FileStatus fileStatus = fs.getFileStatus(resourcePath);
if (mode.equals(VIA_HDFS_DIST_CACHE)) {
LocalResource lr = LocalResource.newInstance(ConverterUtils.getYarnUrlFromPath(resourcePath),
LocalResourceType.ARCHIVE.FILE, LocalResourceVisibility.PRIVATE, fileStatus.getLen(),
fileStatus.getModificationTime());
localResources.put(DISK_PAYLOAD_NAME, lr);
payload.put(4, VIA_HDFS_DIST_CACHE_BYTE); // ViaRPC
} else if (mode.equals(VIA_HDFS_DIRECT_READ)) {
payload.put(4, VIA_HDFS_DIRECT_READ_BYTE); // ViaRPC
}
}
payload.putInt(0, maxSleepTimeMillis);
return UserPayload.create(payload);
}
public static class RPCSleepProcessor extends SimpleProcessor {
private final int sleepTimeMax;
private final byte modeByte;
public RPCSleepProcessor(ProcessorContext context) {
super(context);
sleepTimeMax = getContext().getUserPayload().getPayload().getInt(0);
modeByte = getContext().getUserPayload().getPayload().get(4);
}
@Override
public void run() throws Exception {
StopWatch sw = new StopWatch().start();
long sleepTime = random.nextInt(sleepTimeMax);
if (modeByte == VIA_RPC_BYTE) {
LOG.info("Received via RPC.");
} else if (modeByte == VIA_HDFS_DIST_CACHE_BYTE) {
LOG.info("Reading from local filesystem");
FileSystem localFs = FileSystem.getLocal(new Configuration());
FSDataInputStream is = localFs.open(new Path(DISK_PAYLOAD_NAME));
IOUtils.toByteArray(is);
} else if (modeByte == VIA_HDFS_DIRECT_READ_BYTE) {
LOG.info("Reading from HDFS");
FileSystem fs = FileSystem.get(new Configuration());
FSDataInputStream is = fs.open(new Path(Path.SEPARATOR + "tmp", DISK_PAYLOAD_NAME));
IOUtils.toByteArray(is);
} else {
throw new IllegalArgumentException("Unknown execution mode: [" + modeByte + "]");
}
LOG.info("TimeTakenToAccessPayload=" + sw.stop().now(TimeUnit.MILLISECONDS));
LOG.info("Sleeping for: " + sleepTime);
Thread.sleep(sleepTime);
}
}
public static void main(String[] args) throws Exception {
int res = ToolRunner.run(new Configuration(), new RPCLoadGen(), args);
System.exit(res);
}
}