<commit_before>/*
 * 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.twill.yarn;

import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicates;
import com.google.common.base.Throwables;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import com.google.common.io.ByteStreams;
import com.google.common.io.OutputSupplier;
import com.google.common.reflect.TypeToken;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.twill.api.ClassAcceptor;
import org.apache.twill.api.EventHandlerSpecification;
import org.apache.twill.api.LocalFile;
import org.apache.twill.api.RunId;
import org.apache.twill.api.RuntimeSpecification;
import org.apache.twill.api.SecureStore;
import org.apache.twill.api.TwillController;
import org.apache.twill.api.TwillPreparer;
import org.apache.twill.api.TwillSpecification;
import org.apache.twill.api.logging.LogEntry;
import org.apache.twill.api.logging.LogHandler;
import org.apache.twill.filesystem.Location;
import org.apache.twill.filesystem.LocationFactory;
import org.apache.twill.internal.ApplicationBundler;
import org.apache.twill.internal.Arguments;
import org.apache.twill.internal.Configs;
import org.apache.twill.internal.Constants;
import org.apache.twill.internal.DefaultLocalFile;
import org.apache.twill.internal.DefaultRuntimeSpecification;
import org.apache.twill.internal.DefaultTwillSpecification;
import org.apache.twill.internal.EnvKeys;
import org.apache.twill.internal.JvmOptions;
import org.apache.twill.internal.LogOnlyEventHandler;
import org.apache.twill.internal.ProcessController;
import org.apache.twill.internal.ProcessLauncher;
import org.apache.twill.internal.RunIds;
import org.apache.twill.internal.appmaster.ApplicationMasterInfo;
import org.apache.twill.internal.appmaster.ApplicationMasterMain;
import org.apache.twill.internal.container.TwillContainerMain;
import org.apache.twill.internal.json.ArgumentsCodec;
import org.apache.twill.internal.json.JvmOptionsCodec;
import org.apache.twill.internal.json.LocalFileCodec;
import org.apache.twill.internal.json.TwillSpecificationAdapter;
import org.apache.twill.internal.utils.Dependencies;
import org.apache.twill.internal.utils.Paths;
import org.apache.twill.internal.utils.Resources;
import org.apache.twill.internal.yarn.YarnAppClient;
import org.apache.twill.internal.yarn.YarnApplicationReport;
import org.apache.twill.internal.yarn.YarnUtils;
import org.apache.twill.launcher.FindFreePort;
import org.apache.twill.launcher.TwillLauncher;
import org.apache.twill.zookeeper.ZKClient;
import org.apache.twill.zookeeper.ZKClients;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;

/**
 * Implementation for {@link TwillPreparer} to prepare and launch distributed application on Hadoop YARN.
 */
final class YarnTwillPreparer implements TwillPreparer {

  private static final Logger LOG = LoggerFactory.getLogger(YarnTwillPreparer.class);

  private final YarnConfiguration yarnConfig;
  private final TwillSpecification twillSpec;
  private final YarnAppClient yarnAppClient;
  private final ZKClient zkClient;
  private final LocationFactory locationFactory;
  private final YarnTwillControllerFactory controllerFactory;
  private final RunId runId;

  private final List<LogHandler> logHandlers = Lists.newArrayList();
  private final List<String> arguments = Lists.newArrayList();
  private final Set<Class<?>> dependencies = Sets.newIdentityHashSet();
  private final List<URI> resources = Lists.newArrayList();
  private final List<String> classPaths = Lists.newArrayList();
  private final ListMultimap<String, String> runnableArgs = ArrayListMultimap.create();
  private final Map<String, Map<String, String>> environments = new HashMap<>();
  private final List<String> applicationClassPaths = Lists.newArrayList();
  private final Credentials credentials;
  private final int reservedMemory;
  private String user;
  private String schedulerQueue;
  private String extraOptions;
  private JvmOptions.DebugOptions debugOptions = JvmOptions.DebugOptions.NO_DEBUG;
  private ClassAcceptor classAcceptor;
  private LogEntry.Level logLevel;

  YarnTwillPreparer(YarnConfiguration yarnConfig, TwillSpecification twillSpec,
                    YarnAppClient yarnAppClient, ZKClient zkClient,
                    LocationFactory locationFactory, String extraOptions, LogEntry.Level logLevel,
                    YarnTwillControllerFactory controllerFactory) {
    this.yarnConfig = yarnConfig;
    this.twillSpec = twillSpec;
    this.yarnAppClient = yarnAppClient;
    this.zkClient = ZKClients.namespace(zkClient, "/" + twillSpec.getName());
    this.locationFactory = locationFactory;
    this.controllerFactory = controllerFactory;
    this.runId = RunIds.generate();
    this.credentials = createCredentials();
    this.reservedMemory = yarnConfig.getInt(Configs.Keys.JAVA_RESERVED_MEMORY_MB,
                                            Configs.Defaults.JAVA_RESERVED_MEMORY_MB);
    this.user = System.getProperty("user.name");
    this.extraOptions = extraOptions;
    this.logLevel = logLevel;
    this.classAcceptor = new ClassAcceptor();
  }

  @Override
  public TwillPreparer addLogHandler(LogHandler handler) {
    logHandlers.add(handler);
    return this;
  }

  @Override
  public TwillPreparer setUser(String user) {
    this.user = user;
    return this;
  }

  @Override
  public TwillPreparer setSchedulerQueue(String name) {
    this.schedulerQueue = name;
    return this;
  }

  @Override
  public TwillPreparer setJVMOptions(String options) {
    this.extraOptions = options;
    return this;
  }

  @Override
  public TwillPreparer addJVMOptions(String options) {
    this.extraOptions = extraOptions == null ? options : extraOptions + " " + options;
    return this;
  }

  @Override
  public TwillPreparer enableDebugging(String... runnables) {
    return enableDebugging(false, runnables);
  }

  @Override
  public TwillPreparer enableDebugging(boolean doSuspend, String... runnables) {
    this.debugOptions = new JvmOptions.DebugOptions(true, doSuspend, ImmutableSet.copyOf(runnables));
    return this;
  }

  @Override
  public TwillPreparer withApplicationArguments(String... args) {
    return withApplicationArguments(ImmutableList.copyOf(args));
  }

  @Override
  public TwillPreparer withApplicationArguments(Iterable<String> args) {
    Iterables.addAll(arguments, args);
    return this;
  }

  @Override
  public TwillPreparer withArguments(String runnableName, String... args) {
    return withArguments(runnableName, ImmutableList.copyOf(args));
  }

  @Override
  public TwillPreparer withArguments(String runnableName, Iterable<String> args) {
    Preconditions.checkArgument(twillSpec.getRunnables().containsKey(runnableName),
                                "Runnable %s is not defined in the application.", runnableName);
    runnableArgs.putAll(runnableName, args);
    return this;
  }

  @Override
  public TwillPreparer withDependencies(Class<?>... classes) {
    return withDependencies(ImmutableList.copyOf(classes));
  }

  @Override
  public TwillPreparer withDependencies(Iterable<Class<?>> classes) {
    Iterables.addAll(dependencies, classes);
    return this;
  }

  @Override
  public TwillPreparer withResources(URI... resources) {
    return withResources(ImmutableList.copyOf(resources));
  }

  @Override
  public TwillPreparer withResources(Iterable<URI> resources) {
    Iterables.addAll(this.resources, resources);
    return this;
  }

  @Override
  public TwillPreparer withClassPaths(String... classPaths) {
    return withClassPaths(ImmutableList.copyOf(classPaths));
  }

  @Override
  public TwillPreparer withClassPaths(Iterable<String> classPaths) {
    Iterables.addAll(this.classPaths, classPaths);
    return this;
  }

  @Override
  public TwillPreparer withEnv(Map<String, String> env) {
    // Add the given environments to all runnables
    for (String runnableName : twillSpec.getRunnables().keySet()) {
      setEnv(runnableName, env, false);
    }
    return this;
  }

  @Override
  public TwillPreparer withEnv(String runnableName, Map<String, String> env) {
    Preconditions.checkArgument(twillSpec.getRunnables().containsKey(runnableName),
                                "Runnable %s is not defined in the application.", runnableName);
    setEnv(runnableName, env, true);
    return this;
  }

  @Override
  public TwillPreparer withApplicationClassPaths(String... classPaths) {
    return withApplicationClassPaths(ImmutableList.copyOf(classPaths));
  }

  @Override
  public TwillPreparer withApplicationClassPaths(Iterable<String> classPaths) {
    Iterables.addAll(this.applicationClassPaths, classPaths);
    return this;
  }

  @Override
  public TwillPreparer withBundlerClassAcceptor(ClassAcceptor classAcceptor) {
    this.classAcceptor = classAcceptor;
    return this;
  }

  @Override
  public TwillPreparer addSecureStore(SecureStore secureStore) {
    Object store = secureStore.getStore();
    Preconditions.checkArgument(store instanceof Credentials, "Only Hadoop Credentials is supported.");
    this.credentials.mergeAll((Credentials) store);
    return this;
  }

  @Override
  public TwillPreparer setLogLevel(LogEntry.Level logLevel) {
    this.logLevel = logLevel;
    return this;
  }

  @Override
  public TwillController start() {
    try {
      final ProcessLauncher<ApplicationMasterInfo> launcher = yarnAppClient.createLauncher(twillSpec, schedulerQueue);
      final ApplicationMasterInfo appMasterInfo = launcher.getContainerInfo();
      Callable<ProcessController<YarnApplicationReport>> submitTask =
        new Callable<ProcessController<YarnApplicationReport>>() {
        @Override
        public ProcessController<YarnApplicationReport> call() throws Exception {
          String fsUser = locationFactory.getHomeLocation().getName();

          // Local files needed by AM
          Map<String, LocalFile> localFiles = Maps.newHashMap();
          // Local files declared by runnables
          Multimap<String, LocalFile> runnableLocalFiles = HashMultimap.create();

          createAppMasterJar(createBundler(), localFiles);
          createContainerJar(createBundler(), localFiles);
          populateRunnableLocalFiles(twillSpec, runnableLocalFiles);
          saveSpecification(twillSpec, runnableLocalFiles, localFiles);
          saveLogback(localFiles);
          saveLauncher(localFiles);
          saveJvmOptions(localFiles);
          saveArguments(new Arguments(arguments, runnableArgs), localFiles);
          saveEnvironments(localFiles);
          saveLocalFiles(localFiles, ImmutableSet.of(Constants.Files.TWILL_SPEC,
                                                     Constants.Files.LOGBACK_TEMPLATE,
                                                     Constants.Files.CONTAINER_JAR,
                                                     Constants.Files.LAUNCHER_JAR,
                                                     Constants.Files.ARGUMENTS));

          LOG.debug("Submit AM container spec: {}", appMasterInfo);
          // java -Djava.io.tmpdir=tmp -cp launcher.jar:$HADOOP_CONF_DIR -XmxMemory
          //     org.apache.twill.internal.TwillLauncher
          //     appMaster.jar
          //     org.apache.twill.internal.appmaster.ApplicationMasterMain
          //     false
          ImmutableMap.Builder<String, String> builder = ImmutableMap.<String, String>builder()
            .put(EnvKeys.TWILL_FS_USER, fsUser)
            .put(EnvKeys.TWILL_APP_DIR, getAppLocation().toURI().toASCIIString())
            .put(EnvKeys.TWILL_ZK_CONNECT, zkClient.getConnectString())
            .put(EnvKeys.TWILL_RUN_ID, runId.getId())
            .put(EnvKeys.TWILL_RESERVED_MEMORY_MB, Integer.toString(reservedMemory))
            .put(EnvKeys.TWILL_APP_NAME, twillSpec.getName())
            .put(EnvKeys.YARN_RM_SCHEDULER_ADDRESS, yarnConfig.get(YarnConfiguration.RM_SCHEDULER_ADDRESS));

          if (logLevel != null) {
            LOG.debug("Log level is set to {} for the Twill application.", logLevel);
            builder.put(EnvKeys.TWILL_APP_LOG_LEVEL, logLevel.toString());
          }

          int memory = Resources.computeMaxHeapSize(appMasterInfo.getMemoryMB(),
                                                    Constants.APP_MASTER_RESERVED_MEMORY_MB, Constants.HEAP_MIN_RATIO);
          return launcher.prepareLaunch(builder.build(), localFiles.values(), credentials)
            .addCommand(
              "$JAVA_HOME/bin/java",
              "-Djava.io.tmpdir=tmp",
              "-Dyarn.appId=$" + EnvKeys.YARN_APP_ID_STR,
              "-Dtwill.app=$" + EnvKeys.TWILL_APP_NAME,
              "-cp", Constants.Files.LAUNCHER_JAR + ":$HADOOP_CONF_DIR",
              "-Xmx" + memory + "m",
              extraOptions == null ? "" : extraOptions,
              TwillLauncher.class.getName(),
              Constants.Files.APP_MASTER_JAR,
              ApplicationMasterMain.class.getName(),
              Boolean.FALSE.toString())<commit_msg>Remove the if-check in the YarnTwillPreparer class for log level when setting env key.<commit_after>/*
 * 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.twill.yarn;

import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicates;
import com.google.common.base.Throwables;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import com.google.common.io.ByteStreams;
import com.google.common.io.OutputSupplier;
import com.google.common.reflect.TypeToken;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.twill.api.ClassAcceptor;
import org.apache.twill.api.EventHandlerSpecification;
import org.apache.twill.api.LocalFile;
import org.apache.twill.api.RunId;
import org.apache.twill.api.RuntimeSpecification;
import org.apache.twill.api.SecureStore;
import org.apache.twill.api.TwillController;
import org.apache.twill.api.TwillPreparer;
import org.apache.twill.api.TwillSpecification;
import org.apache.twill.api.logging.LogEntry;
import org.apache.twill.api.logging.LogHandler;
import org.apache.twill.filesystem.Location;
import org.apache.twill.filesystem.LocationFactory;
import org.apache.twill.internal.ApplicationBundler;
import org.apache.twill.internal.Arguments;
import org.apache.twill.internal.Configs;
import org.apache.twill.internal.Constants;
import org.apache.twill.internal.DefaultLocalFile;
import org.apache.twill.internal.DefaultRuntimeSpecification;
import org.apache.twill.internal.DefaultTwillSpecification;
import org.apache.twill.internal.EnvKeys;
import org.apache.twill.internal.JvmOptions;
import org.apache.twill.internal.LogOnlyEventHandler;
import org.apache.twill.internal.ProcessController;
import org.apache.twill.internal.ProcessLauncher;
import org.apache.twill.internal.RunIds;
import org.apache.twill.internal.appmaster.ApplicationMasterInfo;
import org.apache.twill.internal.appmaster.ApplicationMasterMain;
import org.apache.twill.internal.container.TwillContainerMain;
import org.apache.twill.internal.json.ArgumentsCodec;
import org.apache.twill.internal.json.JvmOptionsCodec;
import org.apache.twill.internal.json.LocalFileCodec;
import org.apache.twill.internal.json.TwillSpecificationAdapter;
import org.apache.twill.internal.utils.Dependencies;
import org.apache.twill.internal.utils.Paths;
import org.apache.twill.internal.utils.Resources;
import org.apache.twill.internal.yarn.YarnAppClient;
import org.apache.twill.internal.yarn.YarnApplicationReport;
import org.apache.twill.internal.yarn.YarnUtils;
import org.apache.twill.launcher.FindFreePort;
import org.apache.twill.launcher.TwillLauncher;
import org.apache.twill.zookeeper.ZKClient;
import org.apache.twill.zookeeper.ZKClients;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;

/**
 * Implementation for {@link TwillPreparer} to prepare and launch distributed application on Hadoop YARN.
 */
final class YarnTwillPreparer implements TwillPreparer {

  private static final Logger LOG = LoggerFactory.getLogger(YarnTwillPreparer.class);

  private final YarnConfiguration yarnConfig;
  private final TwillSpecification twillSpec;
  private final YarnAppClient yarnAppClient;
  private final ZKClient zkClient;
  private final LocationFactory locationFactory;
  private final YarnTwillControllerFactory controllerFactory;
  private final RunId runId;

  private final List<LogHandler> logHandlers = Lists.newArrayList();
  private final List<String> arguments = Lists.newArrayList();
  private final Set<Class<?>> dependencies = Sets.newIdentityHashSet();
  private final List<URI> resources = Lists.newArrayList();
  private final List<String> classPaths = Lists.newArrayList();
  private final ListMultimap<String, String> runnableArgs = ArrayListMultimap.create();
  private final Map<String, Map<String, String>> environments = new HashMap<>();
  private final List<String> applicationClassPaths = Lists.newArrayList();
  private final Credentials credentials;
  private final int reservedMemory;
  private String user;
  private String schedulerQueue;
  private String extraOptions;
  private JvmOptions.DebugOptions debugOptions = JvmOptions.DebugOptions.NO_DEBUG;
  private ClassAcceptor classAcceptor;
  private LogEntry.Level logLevel;

  YarnTwillPreparer(YarnConfiguration yarnConfig, TwillSpecification twillSpec,
                    YarnAppClient yarnAppClient, ZKClient zkClient,
                    LocationFactory locationFactory, String extraOptions, LogEntry.Level logLevel,
                    YarnTwillControllerFactory controllerFactory) {
    this.yarnConfig = yarnConfig;
    this.twillSpec = twillSpec;
    this.yarnAppClient = yarnAppClient;
    this.zkClient = ZKClients.namespace(zkClient, "/" + twillSpec.getName());
    this.locationFactory = locationFactory;
    this.controllerFactory = controllerFactory;
    this.runId = RunIds.generate();
    this.credentials = createCredentials();
    this.reservedMemory = yarnConfig.getInt(Configs.Keys.JAVA_RESERVED_MEMORY_MB,
                                            Configs.Defaults.JAVA_RESERVED_MEMORY_MB);
    this.user = System.getProperty("user.name");
    this.extraOptions = extraOptions;
    this.logLevel = logLevel;
    this.classAcceptor = new ClassAcceptor();
  }

  @Override
  public TwillPreparer addLogHandler(LogHandler handler) {
    logHandlers.add(handler);
    return this;
  }

  @Override
  public TwillPreparer setUser(String user) {
    this.user = user;
    return this;
  }

  @Override
  public TwillPreparer setSchedulerQueue(String name) {
    this.schedulerQueue = name;
    return this;
  }

  @Override
  public TwillPreparer setJVMOptions(String options) {
    this.extraOptions = options;
    return this;
  }

  @Override
  public TwillPreparer addJVMOptions(String options) {
    this.extraOptions = extraOptions == null ? options : extraOptions + " " + options;
    return this;
  }

  @Override
  public TwillPreparer enableDebugging(String... runnables) {
    return enableDebugging(false, runnables);
  }

  @Override
  public TwillPreparer enableDebugging(boolean doSuspend, String... runnables) {
    this.debugOptions = new JvmOptions.DebugOptions(true, doSuspend, ImmutableSet.copyOf(runnables));
    return this;
  }

  @Override
  public TwillPreparer withApplicationArguments(String... args) {
    return withApplicationArguments(ImmutableList.copyOf(args));
  }

  @Override
  public TwillPreparer withApplicationArguments(Iterable<String> args) {
    Iterables.addAll(arguments, args);
    return this;
  }

  @Override
  public TwillPreparer withArguments(String runnableName, String... args) {
    return withArguments(runnableName, ImmutableList.copyOf(args));
  }

  @Override
  public TwillPreparer withArguments(String runnableName, Iterable<String> args) {
    Preconditions.checkArgument(twillSpec.getRunnables().containsKey(runnableName),
                                "Runnable %s is not defined in the application.", runnableName);
    runnableArgs.putAll(runnableName, args);
    return this;
  }

  @Override
  public TwillPreparer withDependencies(Class<?>... classes) {
    return withDependencies(ImmutableList.copyOf(classes));
  }

  @Override
  public TwillPreparer withDependencies(Iterable<Class<?>> classes) {
    Iterables.addAll(dependencies, classes);
    return this;
  }

  @Override
  public TwillPreparer withResources(URI... resources) {
    return withResources(ImmutableList.copyOf(resources));
  }

  @Override
  public TwillPreparer withResources(Iterable<URI> resources) {
    Iterables.addAll(this.resources, resources);
    return this;
  }

  @Override
  public TwillPreparer withClassPaths(String... classPaths) {
    return withClassPaths(ImmutableList.copyOf(classPaths));
  }

  @Override
  public TwillPreparer withClassPaths(Iterable<String> classPaths) {
    Iterables.addAll(this.classPaths, classPaths);
    return this;
  }

  @Override
  public TwillPreparer withEnv(Map<String, String> env) {
    // Add the given environments to all runnables
    for (String runnableName : twillSpec.getRunnables().keySet()) {
      setEnv(runnableName, env, false);
    }
    return this;
  }

  @Override
  public TwillPreparer withEnv(String runnableName, Map<String, String> env) {
    Preconditions.checkArgument(twillSpec.getRunnables().containsKey(runnableName),
                                "Runnable %s is not defined in the application.", runnableName);
    setEnv(runnableName, env, true);
    return this;
  }

  @Override
  public TwillPreparer withApplicationClassPaths(String... classPaths) {
    return withApplicationClassPaths(ImmutableList.copyOf(classPaths));
  }

  @Override
  public TwillPreparer withApplicationClassPaths(Iterable<String> classPaths) {
    Iterables.addAll(this.applicationClassPaths, classPaths);
    return this;
  }

  @Override
  public TwillPreparer withBundlerClassAcceptor(ClassAcceptor classAcceptor) {
    this.classAcceptor = classAcceptor;
    return this;
  }

  @Override
  public TwillPreparer addSecureStore(SecureStore secureStore) {
    Object store = secureStore.getStore();
    Preconditions.checkArgument(store instanceof Credentials, "Only Hadoop Credentials is supported.");
    this.credentials.mergeAll((Credentials) store);
    return this;
  }

  @Override
  public TwillPreparer setLogLevel(LogEntry.Level logLevel) {
    Preconditions.checkNotNull(logLevel);
    this.logLevel = logLevel;
    return this;
  }

  @Override
  public TwillController start() {
    try {
      final ProcessLauncher<ApplicationMasterInfo> launcher = yarnAppClient.createLauncher(twillSpec, schedulerQueue);
      final ApplicationMasterInfo appMasterInfo = launcher.getContainerInfo();
      Callable<ProcessController<YarnApplicationReport>> submitTask =
        new Callable<ProcessController<YarnApplicationReport>>() {
        @Override
        public ProcessController<YarnApplicationReport> call() throws Exception {
          String fsUser = locationFactory.getHomeLocation().getName();

          // Local files needed by AM
          Map<String, LocalFile> localFiles = Maps.newHashMap();
          // Local files declared by runnables
          Multimap<String, LocalFile> runnableLocalFiles = HashMultimap.create();

          createAppMasterJar(createBundler(), localFiles);
          createContainerJar(createBundler(), localFiles);
          populateRunnableLocalFiles(twillSpec, runnableLocalFiles);
          saveSpecification(twillSpec, runnableLocalFiles, localFiles);
          saveLogback(localFiles);
          saveLauncher(localFiles);
          saveJvmOptions(localFiles);
          saveArguments(new Arguments(arguments, runnableArgs), localFiles);
          saveEnvironments(localFiles);
          saveLocalFiles(localFiles, ImmutableSet.of(Constants.Files.TWILL_SPEC,
                                                     Constants.Files.LOGBACK_TEMPLATE,
                                                     Constants.Files.CONTAINER_JAR,
                                                     Constants.Files.LAUNCHER_JAR,
                                                     Constants.Files.ARGUMENTS));

          LOG.debug("Submit AM container spec: {}", appMasterInfo);
          // java -Djava.io.tmpdir=tmp -cp launcher.jar:$HADOOP_CONF_DIR -XmxMemory
          //     org.apache.twill.internal.TwillLauncher
          //     appMaster.jar
          //     org.apache.twill.internal.appmaster.ApplicationMasterMain
          //     false
          ImmutableMap.Builder<String, String> builder = ImmutableMap.<String, String>builder()
            .put(EnvKeys.TWILL_FS_USER, fsUser)
            .put(EnvKeys.TWILL_APP_DIR, getAppLocation().toURI().toASCIIString())
            .put(EnvKeys.TWILL_ZK_CONNECT, zkClient.getConnectString())
            .put(EnvKeys.TWILL_RUN_ID, runId.getId())
            .put(EnvKeys.TWILL_RESERVED_MEMORY_MB, Integer.toString(reservedMemory))
            .put(EnvKeys.TWILL_APP_NAME, twillSpec.getName())
            .put(EnvKeys.YARN_RM_SCHEDULER_ADDRESS, yarnConfig.get(YarnConfiguration.RM_SCHEDULER_ADDRESS));

          LOG.debug("Log level is set to {} for the Twill application.", logLevel);
          builder.put(EnvKeys.TWILL_APP_LOG_LEVEL, logLevel.toString());

          int memory = Resources.computeMaxHeapSize(appMasterInfo.getMemoryMB(),
                                                    Constants.APP_MASTER_RESERVED_MEMORY_MB, Constants.HEAP_MIN_RATIO);
          return launcher.prepareLaunch(builder.build(), localFiles.values(), credentials)
            .addCommand(
              "$JAVA_HOME/bin/java",
              "-Djava.io.tmpdir=tmp",
              "-Dyarn.appId=$" + EnvKeys.YARN_APP_ID_STR,
              "-Dtwill.app=$" + EnvKeys.TWILL_APP_NAME,
              "-cp", Constants.Files.LAUNCHER_JAR + ":$HADOOP_CONF_DIR",
              "-Xmx" + memory + "m",
              extraOptions == null ? "" : extraOptions,
              TwillLauncher.class.getName(),
              Constants.Files.APP_MASTER_JAR,
              ApplicationMasterMain.class.getName(),
              Boolean.FALSE.toString())
{'fraction_non_alphanumeric': 0.09282363509332268, 'fraction_numerical': 0.00039924144126160294, 'mean_word_length': 3.8371419375603475, 'pattern_counts': {'":': 2, '<': 65, '<?xml version=': 0, '>': 65, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4155034', 'n_tokens_mistral': 8754, 'n_tokens_neox': 8546, 'n_words': 1742}