[SW-393] Allow to configure timeout for clouding up in external mode …

…from sparkling water
h2oai · Apr 6, 2017 · 3591572 · 3591572
1 parent 90adde0
commit 3591572
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git a/core/src/main/scala/org/apache/spark/h2o/backends/external/ExternalBackendConf.scala b/core/src/main/scala/org/apache/spark/h2o/backends/external/ExternalBackendConf.scala
@@ -38,6 +38,7 @@ trait ExternalBackendConf extends SharedBackendConf {
   def clusterStartMode = sparkConf.get(PROP_EXTERNAL_CLUSTER_START_MODE._1, PROP_EXTERNAL_CLUSTER_START_MODE._2)
   def isAutoClusterStartUsed = clusterStartMode == "auto"
   def isManualClusterStartUsed = !isAutoClusterStartUsed
+  def clusterStartTimeout = sparkConf.getInt(PROP_EXTERNAL_CLUSTER_START_TIMEOUT._1, PROP_EXTERNAL_CLUSTER_START_TIMEOUT._2)
 
   /**
     * Sets node and port representing H2O Cluster to which should H2O connect when started in external mode.
@@ -53,6 +54,11 @@ trait ExternalBackendConf extends SharedBackendConf {
     self
   }
 
+  def setClusterStartTimeout(clusterStartTimeout: Int): H2OConf = {
+    sparkConf.set(PROP_EXTERNAL_CLUSTER_START_TIMEOUT._1, clusterStartTimeout.toString)
+    self
+  }
+
   def setH2OCluster(hostPort: String): H2OConf = {
     setExternalClusterMode()
     sparkConf.set(PROP_EXTERNAL_CLUSTER_REPRESENTATIVE._1, hostPort)
@@ -129,6 +135,9 @@ trait ExternalBackendConf extends SharedBackendConf {
 
 object ExternalBackendConf {
 
+  /** Timeout for starting h2o external cluster */
+  val PROP_EXTERNAL_CLUSTER_START_TIMEOUT = ("spark.ext.h2o.cluster.start.timeout", 120)
+
   val PROP_EXTERNAL_CLUSTER_INFO_FILE = ("spark.ext.h2o.cluster.info.name", None)
 
   val PROP_EXTERNAL_H2O_MEMORY = ("spark.ext.h2o.hadoop.memory", "6g")

diff --git a/core/src/main/scala/org/apache/spark/h2o/backends/external/ExternalH2OBackend.scala b/core/src/main/scala/org/apache/spark/h2o/backends/external/ExternalH2OBackend.scala
@@ -66,6 +66,7 @@ class ExternalH2OBackend(val hc: H2OContext) extends SparklingBackend with Exter
       "-jobname", H2O_JOB_NAME.format(sparkAppId),
       "-mapperXmx", conf.mapperXmx,
       "-output", conf.HDFSOutputDir.get,
+      "-timeout", conf.clusterStartTimeout
       "-disown"
     )
 

diff --git a/py/pysparkling/conf.py b/py/pysparkling/conf.py
@@ -91,8 +91,15 @@ def use_manual_cluster_start(self):
         self._jconf.useManualClusterStart()
         return self
 
+    def set_cluster_start_timeout(self, timeout):
+        self._jconf.setClusterStartTimeout(timeout)
+        return self
+
 # getters
 
+    def cluster_start_timeout(self):
+        return self._jconf.clusterStartTimeout()
+
     def h2o_cluster(self):
         return self._get_option(self._jconf.h2oCluster)