[REEF-1900] Add option for number of cores for the Driver in REEF Java
This adds all the plumbing to allow for the configuration of the Driver's CPU
cores in the REEF Java implementation.
Notes:
* Only the YARN and HDInsight runtimes read the new field.
JIRA:
[REEF-1900](https://issues.apache.org/jira/browse/REEF-1900)
Pull Request:
This closes #1388
diff --git a/lang/java/reef-bridge-client/src/main/java/org/apache/reef/bridge/client/YarnJobSubmissionClient.java b/lang/java/reef-bridge-client/src/main/java/org/apache/reef/bridge/client/YarnJobSubmissionClient.java
index 868201b..b62c2e4 100644
--- a/lang/java/reef-bridge-client/src/main/java/org/apache/reef/bridge/client/YarnJobSubmissionClient.java
+++ b/lang/java/reef-bridge-client/src/main/java/org/apache/reef/bridge/client/YarnJobSubmissionClient.java
@@ -162,7 +162,7 @@
.addLocalResource(this.fileNames.getREEFFolderName(), jarFileOnDFS)
.addLocalResource(fileNames.getYarnBootstrapJobParamFilePath(), jobFileOnDFS)
.setApplicationName(yarnSubmission.getJobId())
- .setDriverMemory(yarnSubmission.getDriverMemory())
+ .setDriverResources(yarnSubmission.getDriverMemory(), 1)
.setPriority(yarnSubmission.getPriority())
.setQueue(yarnSubmission.getQueue())
.setMaxApplicationAttempts(yarnSubmission.getMaxApplicationSubmissions())
diff --git a/lang/java/reef-common/src/main/java/org/apache/reef/client/DriverConfiguration.java b/lang/java/reef-common/src/main/java/org/apache/reef/client/DriverConfiguration.java
index 6edb52f..a832316 100644
--- a/lang/java/reef-common/src/main/java/org/apache/reef/client/DriverConfiguration.java
+++ b/lang/java/reef-common/src/main/java/org/apache/reef/client/DriverConfiguration.java
@@ -59,6 +59,11 @@
public static final OptionalParameter<Integer> DRIVER_MEMORY = new OptionalParameter<>();
/**
+ * The number of cores allocated for the Driver.
+ */
+ public static final OptionalParameter<Integer> DRIVER_CPU_CORES = new OptionalParameter<>();
+
+ /**
* Files to be made available on the Driver and all Evaluators.
*/
public static final OptionalParameter<String> GLOBAL_FILES = new OptionalParameter<>();
@@ -214,6 +219,7 @@
public static final ConfigurationModule CONF = new DriverConfiguration().merge(DriverRuntimeConfiguration.CONF)
.bindNamedParameter(DriverIdentifier.class, DRIVER_IDENTIFIER)
.bindNamedParameter(DriverMemory.class, DRIVER_MEMORY)
+ .bindNamedParameter(DriverCPUCores.class, DRIVER_CPU_CORES)
.bindNamedParameter(DriverJobSubmissionDirectory.class, DRIVER_JOB_SUBMISSION_DIRECTORY)
.bindNamedParameter(MaxApplicationSubmissions.class, MAX_APPLICATION_SUBMISSIONS)
.bindSetEntry(JobGlobalFiles.class, GLOBAL_FILES)
diff --git a/lang/java/reef-common/src/main/java/org/apache/reef/driver/parameters/DriverCPUCores.java b/lang/java/reef-common/src/main/java/org/apache/reef/driver/parameters/DriverCPUCores.java
new file mode 100644
index 0000000..c4ad1c7
--- /dev/null
+++ b/lang/java/reef-common/src/main/java/org/apache/reef/driver/parameters/DriverCPUCores.java
@@ -0,0 +1,31 @@
+/*
+ * 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.reef.driver.parameters;
+
+import org.apache.reef.tang.annotations.Name;
+import org.apache.reef.tang.annotations.NamedParameter;
+
+/**
+ * Number of cores assigned to the Driver Container.
+ */
+@NamedParameter(doc = "Number of CPU Cores allocated to the Driver.", default_value = "1")
+public final class DriverCPUCores implements Name<Integer> {
+ private DriverCPUCores() {
+ }
+}
diff --git a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/JobSubmissionHelper.java b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/JobSubmissionHelper.java
index a974ed7..9f439ba 100644
--- a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/JobSubmissionHelper.java
+++ b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/JobSubmissionHelper.java
@@ -83,6 +83,7 @@
final JobSubmissionEventImpl.Builder jbuilder = JobSubmissionEventImpl.newBuilder()
.setIdentifier(returnOrGenerateDriverId(injector.getNamedInstance(DriverIdentifier.class)))
.setDriverMemory(injector.getNamedInstance(DriverMemory.class))
+ .setDriverCpuCores(injector.getNamedInstance(DriverCPUCores.class))
.setUserName(System.getProperty("user.name"))
.setPreserveEvaluators(preserveEvaluators)
.setMaxApplicationSubmissions(maxAppSubmissions)
diff --git a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEvent.java b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEvent.java
index 763bc73..10a8ba1 100644
--- a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEvent.java
+++ b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEvent.java
@@ -70,6 +70,11 @@
Optional<Integer> getDriverMemory();
/**
+ * @return The number of CPU cores to be allocated for the Driver
+ */
+ Optional<Integer> getDriverCPUCores();
+
+ /**
* @return Priority to be given to the Job
*/
Optional<Integer> getPriority();
diff --git a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEventImpl.java b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEventImpl.java
index 5c640ee..860439e 100644
--- a/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEventImpl.java
+++ b/lang/java/reef-common/src/main/java/org/apache/reef/runtime/common/client/api/JobSubmissionEventImpl.java
@@ -38,6 +38,7 @@
private final Set<FileResource> globalFileSet;
private final Set<FileResource> localFileSet;
private final Optional<Integer> driverMemory;
+ private final Optional<Integer> driverCpuCores;
private final Optional<Integer> priority;
private final Optional<String> queue;
private final Optional<Boolean> preserveEvaluators;
@@ -51,6 +52,7 @@
this.globalFileSet = BuilderUtils.notNull(builder.globalFileSet);
this.localFileSet = BuilderUtils.notNull(builder.localFileSet);
this.driverMemory = Optional.ofNullable(builder.driverMemory);
+ this.driverCpuCores = Optional.ofNullable(builder.driverCpuCores);
this.priority = Optional.ofNullable(builder.priority);
this.preserveEvaluators = Optional.ofNullable(builder.preserveEvaluators);
this.queue = Optional.ofNullable(builder.queue);
@@ -102,6 +104,11 @@
}
@Override
+ public Optional<Integer> getDriverCPUCores(){
+ return driverCpuCores;
+ }
+
+ @Override
public Optional<Integer> getPriority() {
return priority;
}
@@ -131,6 +138,7 @@
private Set<FileResource> globalFileSet = new HashSet<>();
private Set<FileResource> localFileSet = new HashSet<>();
private Integer driverMemory;
+ private Integer driverCpuCores;
private Integer priority;
private String queue;
private Boolean preserveEvaluators;
@@ -195,6 +203,14 @@
}
/**
+ * @see JobSubmissionEvent#getDriverCPUCores()
+ */
+ public Builder setDriverCpuCores(final Integer driverCpuCores){
+ this.driverCpuCores = driverCpuCores;
+ return this;
+ }
+
+ /**
* @see JobSubmissionEvent#getPriority()
*/
public Builder setPriority(final Integer priority) {
diff --git a/lang/java/reef-runtime-hdinsight/src/main/java/org/apache/reef/runtime/hdinsight/client/HDInsightJobSubmissionHandler.java b/lang/java/reef-runtime-hdinsight/src/main/java/org/apache/reef/runtime/hdinsight/client/HDInsightJobSubmissionHandler.java
index 6d61a11..72f655e 100644
--- a/lang/java/reef-runtime-hdinsight/src/main/java/org/apache/reef/runtime/hdinsight/client/HDInsightJobSubmissionHandler.java
+++ b/lang/java/reef-runtime-hdinsight/src/main/java/org/apache/reef/runtime/hdinsight/client/HDInsightJobSubmissionHandler.java
@@ -141,7 +141,7 @@
return new Resource()
.setMemory(jobSubmissionEvent.getDriverMemory().get())
- .setvCores(1);
+ .setvCores(jobSubmissionEvent.getDriverCPUCores().get());
}
/**
diff --git a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnJobSubmissionHandler.java b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnJobSubmissionHandler.java
index c79c668..d9f89e6 100644
--- a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnJobSubmissionHandler.java
+++ b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnJobSubmissionHandler.java
@@ -121,7 +121,7 @@
submissionHelper
.addLocalResource(this.fileNames.getREEFFolderName(), driverJarOnDfs)
.setApplicationName(id)
- .setDriverMemory(jobSubmissionEvent.getDriverMemory().get())
+ .setDriverResources(jobSubmissionEvent.getDriverMemory().get(), jobSubmissionEvent.getDriverCPUCores().get())
.setPriority(getPriority(jobSubmissionEvent))
.setQueue(getQueue(jobSubmissionEvent))
.setPreserveEvaluators(getPreserveEvaluators(jobSubmissionEvent))
diff --git a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnSubmissionHelper.java b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnSubmissionHelper.java
index 7eec786..428a7fa 100644
--- a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnSubmissionHelper.java
+++ b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/YarnSubmissionHelper.java
@@ -37,7 +37,10 @@
import org.apache.reef.runtime.yarn.util.YarnTypes;
import java.io.IOException;
-import java.util.*;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;
@@ -137,12 +140,15 @@
}
/**
- * Set the amount of memory to be allocated to the Driver.
- * @param megabytes
- * @return
+ * Set the resources (memory and number of cores) for the Driver.
+ *
+ * @param memoryinMegabytes memory to be allocated for the Driver, in MegaBytes.
+ * @param numberOfCores Number of cores to allocate for the Driver.
+ *
+ * @return this.
*/
- public YarnSubmissionHelper setDriverMemory(final int megabytes) {
- applicationSubmissionContext.setResource(Resource.newInstance(getMemory(megabytes), 1));
+ public YarnSubmissionHelper setDriverResources(final int memoryinMegabytes, final int numberOfCores) {
+ applicationSubmissionContext.setResource(Resource.newInstance(getMemory(memoryinMegabytes), numberOfCores));
return this;
}
