TEZ-3821: Ability to fail fast tasks that write too much to local dis…

…k. (#314) (Ayush Saxena reviewed by Rajesh Balamohan)

tez-api/src/main/java/org/apache/tez/dag/api/TezConfiguration.java

@@ -2304,4 +2304,12 @@ static Set<String> getPropertySet() {
  public static final String TEZ_THREAD_DUMP_INTERVAL = "tez.thread.dump.interval";
  public static final String TEZ_THREAD_DUMP_INTERVAL_DEFAULT = "0ms";
 
+ /**
+ * Limits the amount of data that can be written to LocalFileSystem by a Task.
+ */
+ @ConfigurationScope(Scope.DAG)
+ @ConfigurationProperty(type = "long")
+ public static final String TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES = "tez.task.local-fs.write-limit.bytes";
+ public static final long TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES_DEFAULT = -1;
+
 }

tez-runtime-internals/src/main/java/org/apache/tez/runtime/RuntimeTask.java

@@ -18,6 +18,7 @@
 
 package org.apache.tez.runtime;
 
+import java.io.IOException;
 import java.util.Collection;
 import java.util.EnumSet;
 import java.util.Map;
@@ -26,6 +27,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.LocalFileSystem;
 import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.dag.records.TezTaskAttemptID;
 import org.apache.tez.runtime.api.impl.TaskSpec;
@@ -35,6 +38,11 @@
 import org.apache.tez.runtime.metrics.TaskCounterUpdater;
 
 import com.google.common.collect.Maps;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static org.apache.tez.dag.api.TezConfiguration.TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES;
+import static org.apache.tez.dag.api.TezConfiguration.TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES_DEFAULT;
 
 public abstract class RuntimeTask {
 
@@ -54,6 +62,9 @@ public abstract class RuntimeTask {
  private final TaskStatistics statistics;
  private final AtomicBoolean progressNotified = new AtomicBoolean(false);
 
+ private final long lfsBytesWriteLimit;
+ private static final Logger LOG = LoggerFactory.getLogger(RuntimeTask.class);
+
  protected RuntimeTask(TaskSpec taskSpec, Configuration tezConf,
  TezUmbilical tezUmbilical, String pid, boolean setupSysCounterUpdater) {
  this.taskSpec = taskSpec;
@@ -71,6 +82,8 @@ protected RuntimeTask(TaskSpec taskSpec, Configuration tezConf,
  } else {
  this.counterUpdater = null;
  }
+ this.lfsBytesWriteLimit =
+ tezConf.getLong(TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES, TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES_DEFAULT);
  }
 
  protected enum State {
@@ -182,4 +195,36 @@ protected void setTaskDone() {
  protected final boolean isUpdatingSystemCounters() {
  return counterUpdater != null;
  }
+
+ /**
+ * Check whether the task has exceeded any configured limits.
+ *
+ * @throws LocalWriteLimitException in case the limit is exceeded.
+ */
+ public void checkTaskLimits() throws LocalWriteLimitException {
+ // check the limit for writing to local file system
+ if (lfsBytesWriteLimit >= 0) {
+ Long lfsBytesWritten = null;
+ try {
+ LocalFileSystem localFS = FileSystem.getLocal(tezConf);
+ lfsBytesWritten = FileSystem.getGlobalStorageStatistics().get(localFS.getScheme()).getLong("bytesWritten");
+ } catch (IOException e) {
+ LOG.warn("Could not get LocalFileSystem bytesWritten counter");
+ }
+ if (lfsBytesWritten != null && lfsBytesWritten > lfsBytesWriteLimit) {
+ throw new LocalWriteLimitException(
+ "Too much write to local file system." + " current value is " + lfsBytesWritten + " the limit is "
+ + lfsBytesWriteLimit);
+ }
+ }
+ }
+
+ /**
+ * Exception thrown when the task exceeds some configured limits.
+ */
+ public static class LocalWriteLimitException extends IOException {
+ public LocalWriteLimitException(String str) {
+ super(str);
+ }
+ }
 }

tez-runtime-internals/src/main/java/org/apache/tez/runtime/task/TaskReporter.java

@@ -44,6 +44,7 @@
 import org.apache.tez.dag.api.TezException;
 import org.apache.tez.dag.records.TezTaskAttemptID;
 import org.apache.tez.runtime.RuntimeTask;
+import org.apache.tez.runtime.RuntimeTask.LocalWriteLimitException;
 import org.apache.tez.runtime.api.*;
 import org.apache.tez.runtime.api.events.TaskAttemptCompletedEvent;
 import org.apache.tez.runtime.api.events.TaskAttemptFailedEvent;
@@ -141,6 +142,8 @@ static class HeartbeatCallable implements Callable<Boolean> {
  private static final float LOG_COUNTER_BACKOFF = 1.3f;
  private static final int HEAP_MEMORY_USAGE_UPDATE_INTERVAL = 5000; // 5 seconds
 
+ private static final int LOCAL_FILE_SYSTEM_BYTES_WRITTEN_CHECK_INTERVAL = 10000; // 10 seconds
+
  private final RuntimeTask task;
  private final EventMetaData updateEventMetadata;
 
@@ -165,6 +168,9 @@ static class HeartbeatCallable implements Callable<Boolean> {
  private long usedMemory = 0;
  private long heapMemoryUsageUpdatedTime = System.currentTimeMillis() - HEAP_MEMORY_USAGE_UPDATE_INTERVAL;
 
+ private long localFileSystemBytesWrittenCheckInterval =
+ System.currentTimeMillis() - LOCAL_FILE_SYSTEM_BYTES_WRITTEN_CHECK_INTERVAL;
+
  /*
  * Keeps track of regular timed heartbeats. Is primarily used as a timing mechanism to send /
  * log counters.
@@ -262,6 +268,17 @@ private synchronized ResponseWrapper heartbeat(Collection<TezEvent> eventsArg) t
  sendCounters = true;
  prevCounterSendHeartbeatNum = nonOobHeartbeatCounter.get();
  }
+ try {
+ long now = System.currentTimeMillis();
+ if (now - localFileSystemBytesWrittenCheckInterval > LOCAL_FILE_SYSTEM_BYTES_WRITTEN_CHECK_INTERVAL) {
+ task.checkTaskLimits();
+ localFileSystemBytesWrittenCheckInterval = now;
+ }
+ } catch (LocalWriteLimitException lwle) {
+ LOG.error("Local FileSystem write limit exceeded", lwle);
+ askedToDie.set(true);
+ return new ResponseWrapper(true, 1);
+ }
  updateEvent = new TezEvent(getStatusUpdateEvent(sendCounters), updateEventMetadata);
  events.add(updateEvent);
  }

tez-runtime-internals/src/test/java/org/apache/tez/runtime/task/TestTaskReporter.java

@@ -18,6 +18,7 @@
 
 package org.apache.tez.runtime.task;
 
+import static org.apache.tez.dag.api.TezConfiguration.TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
@@ -27,7 +28,11 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.io.File;
+import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -36,12 +41,21 @@
 
 import com.google.common.collect.Lists;
 
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.LocalFileSystem;
+import org.apache.hadoop.fs.Path;
 import org.apache.tez.common.TezTaskUmbilicalProtocol;
 import org.apache.tez.common.counters.TezCounters;
+import org.apache.tez.dag.api.TezConfiguration;
 import org.apache.tez.dag.api.TezUncheckedException;
 import org.apache.tez.dag.records.TezTaskAttemptID;
 import org.apache.tez.runtime.LogicalIOProcessorRuntimeTask;
+import org.apache.tez.runtime.RuntimeTask.LocalWriteLimitException;
 import org.apache.tez.runtime.api.events.TaskStatusUpdateEvent;
+import org.apache.tez.runtime.api.impl.InputSpec;
+import org.apache.tez.runtime.api.impl.OutputSpec;
+import org.apache.tez.runtime.api.impl.TaskSpec;
 import org.apache.tez.runtime.api.impl.TaskStatistics;
 import org.apache.tez.runtime.api.impl.TezEvent;
 import org.apache.tez.runtime.api.impl.TezHeartbeatRequest;
@@ -55,6 +69,9 @@
 @SuppressWarnings("rawtypes")
 public class TestTaskReporter {
 
+ private static final File TEST_DIR =
+ new File(System.getProperty("test.build.data"), TestTaskReporter.class.getName()).getAbsoluteFile();
+
  @Test(timeout = 10000)
  public void testContinuousHeartbeatsOnMaxEvents() throws Exception {
 
@@ -218,6 +235,38 @@ public void testStatusUpdateAfterInitializationAndCounterFlag() {
 
  }
 
+ @Test
+ public void testLocalFileSystemBytesWrittenLimit() throws IOException {
+ TaskSpec mockSpec = mock(TaskSpec.class);
+ when(mockSpec.getInputs()).thenReturn(Collections.singletonList(mock(InputSpec.class)));
+ when(mockSpec.getOutputs()).thenReturn(Collections.singletonList(mock(OutputSpec.class)));
+ TezConfiguration tezConf = new TezConfiguration();
+ LogicalIOProcessorRuntimeTask lio1 =
+ new LogicalIOProcessorRuntimeTask(mockSpec, 0, tezConf, null, null, null, null, null, null, "", null,
+ Runtime.getRuntime().maxMemory(), true, null, null);
+
+ LocalFileSystem localFS = FileSystem.getLocal(tezConf);
+ FileSystem.clearStatistics();
+ Path tmpPath =
+ new Path(TEST_DIR + "/testLocalFileSystemBytesWrittenLimit" + new Random(System.currentTimeMillis()).nextInt());
+ try (FSDataOutputStream out = localFS.create(tmpPath, true)) {
+ out.write(new byte[1024]);
+ }
+ // Check limits with default shouldn't throw exception.
+ lio1.checkTaskLimits();
+
+ tezConf.setLong(TEZ_TASK_LOCAL_FS_WRITE_LIMIT_BYTES, 10);
+ lio1 = new LogicalIOProcessorRuntimeTask(mockSpec, 0, tezConf, null, null, null, null, null, null, "", null,
+ Runtime.getRuntime().maxMemory(), true, null, null);
+
+ try {
+ lio1.checkTaskLimits();
+ Assert.fail("Expected to throw LocalWriteLimitException");
+ } catch (LocalWriteLimitException localWriteLimitException) {
+ Assert.assertTrue(localWriteLimitException.getMessage().contains("Too much write to local file system"));
+ }
+ }
+
  private List<TezEvent> createEvents(int numEvents) {
  List<TezEvent> list = Lists.newArrayListWithCapacity(numEvents);
  for (int i = 0; i < numEvents; i++) {