[FLINK-2954] Add config parameter for passing environment variables to YARN

This closes #1409

docs/setup/config.md

@@ -415,6 +415,16 @@ so that the Flink client is able to pick those details up. This configuration pa
 changing the default location of that file (for example for environments sharing a Flink 
 installation between users)
 
+- `yarn.application-master.env.`*ENV_VAR1=value* Configuration values prefixed with `yarn.application-master.env.`
+will be passed as environment variables to the ApplicationMaster/JobManager process.
+For example for passing `LD_LIBRARY_PATH` as an env variable to the ApplicationMaster, set:
+
+      yarn.application-master.env.LD_LIBRARY_PATH: "/usr/lib/native"
+
+
+- `yarn.taskmanager.env.` Similar to the configuration prefix about, this prefix allows setting custom
+environment variables for the TaskManager processes.
+
 ## High Availability Mode
 
 - `recovery.mode`: (Default 'standalone') Defines the recovery mode used for the cluster execution. Currently,

flink-core/src/main/java/org/apache/flink/configuration/ConfigConstants.java

@@ -243,6 +243,20 @@ public final class ConfigConstants {
 	 */
 	public static final String YARN_PROPERTIES_FILE_LOCATION = "yarn.properties-file.location";
 
+	/**
+	 * Prefix for passing custom environment variables to Flink's ApplicationMaster (JobManager).
+	 * For example for passing LD_LIBRARY_PATH as an env variable to the AppMaster, set:
+	 * 	yarn.application-master.env.LD_LIBRARY_PATH: "/usr/lib/native"
+	 * in the flink-conf.yaml.
+	 */
+	public static final String YARN_APPLICATION_MASTER_ENV_PREFIX = "yarn.application-master.env.";
+
+	/**
+	 * Similar to the {@see YARN_APPLICATION_MASTER_ENV_PREFIX}, this configuration prefix allows
+	 * setting custom environment variables.
+	 */
+	public static final String YARN_TASK_MANAGER_ENV_PREFIX = "yarn.taskmanager.env.";
+
 
 	// ------------------------ Hadoop Configuration ------------------------
 

flink-dist/pom.xml

@@ -107,12 +107,6 @@ under the License.
 			<version>${project.version}</version>
 		</dependency>
 
-		<dependency>
-			<groupId>org.apache.flink</groupId>
-			<artifactId>flink-language-binding-generic</artifactId>
-			<version>${project.version}</version>
-		</dependency>
-
 		<dependency>
 			<groupId>org.apache.flink</groupId>
 			<artifactId>flink-python</artifactId>

flink-yarn-tests/pom.xml

@@ -156,6 +156,7 @@ under the License.
 							</filters>
 							<transformers>
 								<transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+								<transformer implementation="org.apache.maven.plugins.shade.resource.ApacheNoticeResourceTransformer"/>
 							</transformers>
 							<relocations>
 								<relocation>
@@ -174,9 +175,6 @@ under the License.
 									<shadedPattern>org.apache.flink.hadoop.shaded.org.jboss.netty</shadedPattern>
 								</relocation>
 							</relocations>
-							<transformers>
-								<transformer implementation="org.apache.maven.plugins.shade.resource.ApacheNoticeResourceTransformer"/>
-							</transformers>
 						</configuration>
 					</execution>
 				</executions>

flink-yarn-tests/src/main/java/org/apache/flink/yarn/UtilsTest.java

@@ -31,6 +31,7 @@ import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 public class UtilsTest {
 	private static final Logger LOG = LoggerFactory.getLogger(UtilsTest.class);
@@ -63,7 +64,7 @@ public class UtilsTest {
 		Assert.assertEquals(8500, Utils.calculateHeapSize(10000, conf) );
 
 		// test different configuration
-		Assert.assertEquals(3400, Utils.calculateHeapSize(4000, conf) );
+		Assert.assertEquals(3400, Utils.calculateHeapSize(4000, conf));
 
 		conf.setString(ConfigConstants.YARN_HEAP_CUTOFF_MIN, "1000");
 		conf.setString(ConfigConstants.YARN_HEAP_CUTOFF_RATIO, "0.1");
@@ -97,6 +98,28 @@ public class UtilsTest {
 		Assert.assertEquals(0, Utils.calculateHeapSize(4000, conf));
 	}
 
+	@Test
+	public void testGetEnvironmentVariables() {
+		Configuration testConf = new Configuration();
+		testConf.setString("yarn.application-master.env.LD_LIBRARY_PATH", "/usr/lib/native");
+
+		Map<String, String> res = Utils.getEnvironmentVariables("yarn.application-master.env.", testConf);
+
+		Assert.assertEquals(1, res.size());
+		Map.Entry<String, String> entry = res.entrySet().iterator().next();
+		Assert.assertEquals("LD_LIBRARY_PATH", entry.getKey());
+		Assert.assertEquals("/usr/lib/native", entry.getValue());
+	}
+
+	@Test
+	public void testGetEnvironmentVariablesErroneous() {
+		Configuration testConf = new Configuration();
+		testConf.setString("yarn.application-master.env.", "/usr/lib/native");
+
+		Map<String, String> res = Utils.getEnvironmentVariables("yarn.application-master.env.", testConf);
+
+		Assert.assertEquals(0, res.size());
+	}
 
 	//
 	// --------------- Tools to test if a certain string has been logged with Log4j. -------------

flink-yarn-tests/src/main/java/org/apache/flink/yarn/YARNSessionFIFOITCase.java

@@ -479,7 +479,7 @@ public class YARNSessionFIFOITCase extends YarnTestBase {
 		}
 
 		// get temporary file for reading input data for wordcount example
-		File tmpInFile = null;
+		File tmpInFile;
 		try{
 			tmpInFile = tmp.newFile();
 			FileUtils.writeStringToFile(tmpInFile,WordCountData.TEXT);

flink-yarn/src/main/java/org/apache/flink/yarn/FlinkYarnClientBase.java

@@ -587,8 +587,11 @@ public abstract class FlinkYarnClientBase extends AbstractFlinkYarnClient {
 
 		// Setup CLASSPATH for ApplicationMaster
 		Map<String, String> appMasterEnv = new HashMap<String, String>();
+		// set user specified app master environment variables
+		appMasterEnv.putAll(Utils.getEnvironmentVariables(ConfigConstants.YARN_APPLICATION_MASTER_ENV_PREFIX, flinkConfiguration));
+		// set classpath from YARN configuration
 		Utils.setupEnv(conf, appMasterEnv);
-		// set configuration values
+		// set Flink on YARN internal configuration values
 		appMasterEnv.put(FlinkYarnClient.ENV_TM_COUNT, String.valueOf(taskManagerCount));
 		appMasterEnv.put(FlinkYarnClient.ENV_TM_MEMORY, String.valueOf(taskManagerMemoryMb));
 		appMasterEnv.put(FlinkYarnClient.FLINK_JAR_PATH, remotePathJar.toString() );

flink-yarn/src/main/java/org/apache/flink/yarn/Utils.java

@@ -18,11 +18,11 @@
 package org.apache.flink.yarn;
 
 import java.io.File;
-import java.io.FilenameFilter;
 import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.nio.ByteBuffer;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Map;
 
 import org.slf4j.Logger;
@@ -98,7 +98,7 @@ public final class Utils {
 		
 		Path dst = new Path(homedir, suffix);
 		
-		LOG.info("Copying from " + localRsrcPath + " to " + dst );
+		LOG.info("Copying from " + localRsrcPath + " to " + dst);
 		fs.copyFromLocalFile(localRsrcPath, dst);
 		registerLocalResource(fs, dst, appMasterJar);
 		return dst;
@@ -187,17 +187,6 @@ public final class Utils {
 		}
 	}
 
-	public static void logFilesInCurrentDirectory(final Logger logger) {
-		new File(".").list(new FilenameFilter() {
-			
-			@Override
-			public boolean accept(File dir, String name) {
-				logger.info(dir.getAbsolutePath() + "/" + name);
-				return true;
-			}
-		});
-	}
-	
 	/**
 	 * Copied method from org.apache.hadoop.yarn.util.Apps
 	 * It was broken by YARN-1824 (2.4.0) and fixed for 2.4.1
@@ -221,4 +210,22 @@ public final class Utils {
 	private Utils() {
 		throw new RuntimeException();
 	}
+
+	/**
+	 * Method to extract environment variables from the flinkConfiguration based on the given prefix String.
+	 *
+	 * @param envPrefix Prefix for the environment variables key
+	 * @param flinkConfiguration The Flink config to get the environment variable defintion from
+	 */
+	public static Map<String, String> getEnvironmentVariables(String envPrefix, org.apache.flink.configuration.Configuration flinkConfiguration) {
+		Map<String, String> result  = new HashMap<>();
+		for(Map.Entry<String, String> entry: flinkConfiguration.toMap().entrySet()) {
+			if(entry.getKey().startsWith(envPrefix) && entry.getKey().length() > envPrefix.length()) {
+				// remove prefix
+				String key = entry.getKey().substring(envPrefix.length());
+				result.put(key, entry.getValue());
+			}
+		}
+		return result;
+	}
 }

flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterBase.scala

@@ -117,7 +117,7 @@ abstract class ApplicationMasterBase {
 
       // if a web monitor shall be started, set the port to random binding
       if (config.getInteger(ConfigConstants.JOB_MANAGER_WEB_PORT_KEY, 0) >= 0) {
-        config.setInteger(ConfigConstants.JOB_MANAGER_WEB_PORT_KEY, 0);
+        config.setInteger(ConfigConstants.JOB_MANAGER_WEB_PORT_KEY, 0)
       }
 
       val (actorSystem, jmActor, archiveActor, webMonitor) =
@@ -147,7 +147,7 @@ abstract class ApplicationMasterBase {
         jobManagerPort, webServerPort, slots, taskManagerCount,
         dynamicPropertiesEncodedString)
 
-      val hadoopConfig = new YarnConfiguration();
+      val hadoopConfig = new YarnConfiguration()
 
       // send "start yarn session" message to YarnJobManager.
       log.info("Starting YARN session on Job Manager.")

flink-yarn/src/main/scala/org/apache/flink/yarn/YarnJobManager.scala

@@ -707,8 +707,12 @@ class YarnJobManager(
 
     ctx.setLocalResources(taskManagerLocalResources.asJava)
 
-    // Setup classpath for container ( = TaskManager )
+    // Setup classpath and environment variables for container ( = TaskManager )
     val containerEnv = new java.util.HashMap[String, String]()
+    // user defined TaskManager environment variables
+    containerEnv.putAll(Utils.getEnvironmentVariables(ConfigConstants.YARN_TASK_MANAGER_ENV_PREFIX,
+      flinkConfiguration))
+    // YARN classpath
     Utils.setupEnv(yarnConf, containerEnv)
     containerEnv.put(FlinkYarnClientBase.ENV_CLIENT_USERNAME, yarnClientUsername)
     ctx.setEnvironment(containerEnv)