tez-tests/src/main/java/org/apache/tez/mapreduce/examples/RPCLoadGen.java - tez - Git at Google

 /*
  * 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);
   }
 }
	/*
	* 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);
	}
	}