|
@@ -0,0 +1,209 @@
|
|
|
+/**
|
|
|
+ * 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.hadoop.yarn.server.nodemanager.nodelabels;
|
|
|
+
|
|
|
+import java.io.File;
|
|
|
+import java.io.FileOutputStream;
|
|
|
+import java.io.IOException;
|
|
|
+import java.io.PrintWriter;
|
|
|
+import java.util.TimerTask;
|
|
|
+
|
|
|
+import org.apache.hadoop.conf.Configuration;
|
|
|
+import org.apache.hadoop.fs.FileContext;
|
|
|
+import org.apache.hadoop.fs.FileUtil;
|
|
|
+import org.apache.hadoop.fs.Path;
|
|
|
+import org.apache.hadoop.service.ServiceStateException;
|
|
|
+import org.apache.hadoop.util.Shell;
|
|
|
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
|
|
|
+import org.apache.hadoop.yarn.nodelabels.NodeLabelTestBase;
|
|
|
+import org.junit.After;
|
|
|
+import org.junit.Assert;
|
|
|
+import org.junit.Before;
|
|
|
+import org.junit.Test;
|
|
|
+
|
|
|
+public class TestScriptBasedNodeLabelsProvider extends NodeLabelTestBase {
|
|
|
+
|
|
|
+ protected static File testRootDir = new File("target",
|
|
|
+ TestScriptBasedNodeLabelsProvider.class.getName() + "-localDir")
|
|
|
+ .getAbsoluteFile();
|
|
|
+
|
|
|
+ private final File nodeLabelsScriptFile =
|
|
|
+ new File(testRootDir, Shell.appendScriptExtension("failingscript"));
|
|
|
+
|
|
|
+ private ScriptBasedNodeLabelsProvider nodeLabelsProvider;
|
|
|
+
|
|
|
+ @Before
|
|
|
+ public void setup() {
|
|
|
+ testRootDir.mkdirs();
|
|
|
+ nodeLabelsProvider = new ScriptBasedNodeLabelsProvider();
|
|
|
+ }
|
|
|
+
|
|
|
+ @After
|
|
|
+ public void tearDown() throws Exception {
|
|
|
+ if (testRootDir.exists()) {
|
|
|
+ FileContext.getLocalFSFileContext()
|
|
|
+ .delete(new Path(testRootDir.getAbsolutePath()), true);
|
|
|
+ }
|
|
|
+ if (nodeLabelsProvider != null) {
|
|
|
+ nodeLabelsProvider.stop();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ private Configuration getConfForNodeLabelScript() {
|
|
|
+ Configuration conf = new Configuration();
|
|
|
+ conf.set(YarnConfiguration.NM_SCRIPT_BASED_NODE_LABELS_PROVIDER_PATH,
|
|
|
+ nodeLabelsScriptFile.getAbsolutePath());
|
|
|
+ // set bigger interval so that test cases can be run
|
|
|
+ conf.setLong(YarnConfiguration.NM_NODE_LABELS_PROVIDER_FETCH_INTERVAL_MS,
|
|
|
+ 1 * 60 * 60 * 1000l);
|
|
|
+ conf.setLong(YarnConfiguration.NM_NODE_LABELS_PROVIDER_FETCH_TIMEOUT_MS,
|
|
|
+ 1000);
|
|
|
+ return conf;
|
|
|
+ }
|
|
|
+
|
|
|
+ private void writeNodeLabelsScriptFile(String scriptStr,
|
|
|
+ boolean setExecutable) throws IOException {
|
|
|
+ PrintWriter pw = null;
|
|
|
+ try {
|
|
|
+ FileUtil.setWritable(nodeLabelsScriptFile, true);
|
|
|
+ FileUtil.setReadable(nodeLabelsScriptFile, true);
|
|
|
+ pw = new PrintWriter(new FileOutputStream(nodeLabelsScriptFile));
|
|
|
+ pw.println(scriptStr);
|
|
|
+ pw.flush();
|
|
|
+ } catch (Exception e) {
|
|
|
+ e.printStackTrace();
|
|
|
+ Assert.fail();
|
|
|
+ } finally {
|
|
|
+ if (null != pw) {
|
|
|
+ pw.close();
|
|
|
+ }
|
|
|
+ }
|
|
|
+ FileUtil.setExecutable(nodeLabelsScriptFile, setExecutable);
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ public void testNodeLabelsScriptRunnerCreation() throws IOException {
|
|
|
+ // If no script configured then initialization of service should fail
|
|
|
+ ScriptBasedNodeLabelsProvider nodeLabelsProvider =
|
|
|
+ new ScriptBasedNodeLabelsProvider();
|
|
|
+ initilizeServiceFailTest(
|
|
|
+ "Expected to fail fast when no script is configured and "
|
|
|
+ + "ScriptBasedNodeLabelsProvider service is inited",
|
|
|
+ nodeLabelsProvider);
|
|
|
+
|
|
|
+ // If script configured is blank then initialization of service should fail
|
|
|
+ nodeLabelsProvider = new ScriptBasedNodeLabelsProvider();
|
|
|
+ Configuration conf = new Configuration();
|
|
|
+ conf.set(YarnConfiguration.NM_SCRIPT_BASED_NODE_LABELS_PROVIDER_PATH, "");
|
|
|
+ initilizeServiceFailTest(
|
|
|
+ "Expected to fail fast when script path configuration is blank"
|
|
|
+ + "and ScriptBasedNodeLabelsProvider service is inited.",
|
|
|
+ nodeLabelsProvider);
|
|
|
+
|
|
|
+ // If script configured is not executable then no timertask /
|
|
|
+ // NodeLabelsScriptRunner initialized
|
|
|
+ nodeLabelsProvider = new ScriptBasedNodeLabelsProvider();
|
|
|
+ writeNodeLabelsScriptFile("", false);
|
|
|
+ initilizeServiceFailTest(
|
|
|
+ "Expected to fail fast when script is not executable"
|
|
|
+ + "and ScriptBasedNodeLabelsProvider service is inited.",
|
|
|
+ nodeLabelsProvider);
|
|
|
+
|
|
|
+ // If configured script is executable then timertask /
|
|
|
+ // NodeLabelsScriptRunner should be initialized
|
|
|
+ nodeLabelsProvider = new ScriptBasedNodeLabelsProvider();
|
|
|
+ writeNodeLabelsScriptFile("", true);
|
|
|
+ nodeLabelsProvider.init(getConfForNodeLabelScript());
|
|
|
+ nodeLabelsProvider.start();
|
|
|
+ Assert
|
|
|
+ .assertNotNull("Node Label Script runner should be started when script"
|
|
|
+ + " is executable", nodeLabelsProvider.getTimerTask());
|
|
|
+ nodeLabelsProvider.stop();
|
|
|
+ }
|
|
|
+
|
|
|
+ private void initilizeServiceFailTest(String message,
|
|
|
+ ScriptBasedNodeLabelsProvider nodeLabelsProvider) {
|
|
|
+ try {
|
|
|
+ nodeLabelsProvider.init(new Configuration());
|
|
|
+ Assert.fail(message);
|
|
|
+ } catch (ServiceStateException ex) {
|
|
|
+ Assert.assertEquals("IOException was expected", IOException.class,
|
|
|
+ ex.getCause().getClass());
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ public void testConfigForNoTimer() throws Exception {
|
|
|
+ Configuration conf = getConfForNodeLabelScript();
|
|
|
+ conf.setLong(YarnConfiguration.NM_NODE_LABELS_PROVIDER_FETCH_INTERVAL_MS,
|
|
|
+ AbstractNodeLabelsProvider.DISABLE_NODE_LABELS_PROVIDER_FETCH_TIMER);
|
|
|
+ String normalScript = "echo NODE_PARTITION:X86";
|
|
|
+ writeNodeLabelsScriptFile(normalScript, true);
|
|
|
+ nodeLabelsProvider.init(conf);
|
|
|
+ nodeLabelsProvider.start();
|
|
|
+ Assert.assertNull(
|
|
|
+ "Timer is not expected to be created when interval is configured as -1",
|
|
|
+ nodeLabelsProvider.nodeLabelsScheduler);
|
|
|
+ // Ensure that even though timer is not run script is run at least once so
|
|
|
+ // that NM registers/updates Labels with RM
|
|
|
+ assertNLCollectionEquals(toNodeLabelSet("X86"),
|
|
|
+ nodeLabelsProvider.getNodeLabels());
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ public void testNodeLabelsScript() throws Exception {
|
|
|
+ String scriptWithoutLabels = "";
|
|
|
+ String normalScript = "echo NODE_PARTITION:Windows";
|
|
|
+ String scrptWithMultipleLinesHavingNodeLabels =
|
|
|
+ "echo NODE_PARTITION:RAM\n echo NODE_PARTITION:JDK1_6";
|
|
|
+ String timeOutScript = Shell.WINDOWS
|
|
|
+ ? "@echo off\nping -n 4 127.0.0.1 >nul\n" + "echo NODE_PARTITION:ALL"
|
|
|
+ : "sleep 4\necho NODE_PARTITION:ALL";
|
|
|
+
|
|
|
+ writeNodeLabelsScriptFile(scriptWithoutLabels, true);
|
|
|
+ nodeLabelsProvider.init(getConfForNodeLabelScript());
|
|
|
+ nodeLabelsProvider.start();
|
|
|
+ Thread.sleep(500l);
|
|
|
+ TimerTask timerTask = nodeLabelsProvider.getTimerTask();
|
|
|
+ timerTask.run();
|
|
|
+ Assert.assertNull(
|
|
|
+ "Node Label Script runner should return null when script doesnt "
|
|
|
+ + "give any Labels output",
|
|
|
+ nodeLabelsProvider.getNodeLabels());
|
|
|
+
|
|
|
+ writeNodeLabelsScriptFile(normalScript, true);
|
|
|
+ timerTask.run();
|
|
|
+ assertNLCollectionEquals(toNodeLabelSet("Windows"),
|
|
|
+ nodeLabelsProvider.getNodeLabels());
|
|
|
+
|
|
|
+ // multiple lines with partition tag then the last line's partition info
|
|
|
+ // needs to be taken.
|
|
|
+ writeNodeLabelsScriptFile(scrptWithMultipleLinesHavingNodeLabels, true);
|
|
|
+ timerTask.run();
|
|
|
+ assertNLCollectionEquals(toNodeLabelSet("JDK1_6"),
|
|
|
+ nodeLabelsProvider.getNodeLabels());
|
|
|
+
|
|
|
+ // timeout script.
|
|
|
+ writeNodeLabelsScriptFile(timeOutScript, true);
|
|
|
+ timerTask.run();
|
|
|
+
|
|
|
+ Assert.assertNotEquals("Node Labels should not be set after timeout ",
|
|
|
+ toNodeLabelSet("ALL"), nodeLabelsProvider.getNodeLabels());
|
|
|
+ }
|
|
|
+}
|