[AIRFLOW-6359] Make Spark status_poll_interval explicit (#6978)

(cherry picked from commit 6cc1123)
apache · Jan 21, 2020 · edb4f62 · edb4f62
1 parent ec9d891
commit edb4f62
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/airflow/contrib/hooks/spark_submit_hook.py b/airflow/contrib/hooks/spark_submit_hook.py
@@ -86,6 +86,9 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
     :type name: str
     :param num_executors: Number of executors to launch
     :type num_executors: int
+    :param status_poll_interval: Seconds to wait between polls of driver status in cluster
+        mode (Default: 1)
+    :type status_poll_interval: int
     :param application_args: Arguments for the application being submitted
     :type application_args: list
     :param env_vars: Environment variables for spark-submit. It
@@ -118,6 +121,7 @@ def __init__(self,
                  proxy_user=None,
                  name='default-name',
                  num_executors=None,
+                 status_poll_interval=1,
                  application_args=None,
                  env_vars=None,
                  verbose=False,
@@ -142,6 +146,7 @@ def __init__(self,
         self._proxy_user = proxy_user
         self._name = name
         self._num_executors = num_executors
+        self._status_poll_interval = status_poll_interval
         self._application_args = application_args
         self._env_vars = env_vars
         self._verbose = verbose

diff --git a/airflow/contrib/operators/spark_submit_operator.py b/airflow/contrib/operators/spark_submit_operator.py
@@ -76,6 +76,9 @@ class SparkSubmitOperator(BaseOperator):
     :type name: str
     :param num_executors: Number of executors to launch
     :type num_executors: int
+    :param status_poll_interval: Seconds to wait between polls of driver status in cluster
+        mode (Default: 1)
+    :type status_poll_interval: int
     :param application_args: Arguments for the application being submitted (templated)
     :type application_args: list
     :param env_vars: Environment variables for spark-submit. It supports yarn and k8s mode too. (templated)
@@ -114,6 +117,7 @@ def __init__(self,
                  proxy_user=None,
                  name='airflow-spark',
                  num_executors=None,
+                 status_poll_interval=1,
                  application_args=None,
                  env_vars=None,
                  verbose=False,
@@ -141,6 +145,7 @@ def __init__(self,
         self._proxy_user = proxy_user
         self._name = name
         self._num_executors = num_executors
+        self._status_poll_interval = status_poll_interval
         self._application_args = application_args
         self._env_vars = env_vars
         self._verbose = verbose
@@ -173,6 +178,7 @@ def execute(self, context):
             proxy_user=self._proxy_user,
             name=self._name,
             num_executors=self._num_executors,
+            status_poll_interval=self._status_poll_interval,
             application_args=self._application_args,
             env_vars=self._env_vars,
             verbose=self._verbose,