Merge pull request #115 from Yelp/luisp/COMPINFRA-2505-support-use-eks

Add support for --use-eks in PaaSTA

Author: bobokun

service_configuration_lib/spark_config.py

@@ -918,15 +918,17 @@ def _get_k8s_spark_env(
     volumes: Optional[List[Mapping[str, str]]],
     paasta_pool: str,
     service_account_name: Optional[str] = None,
+    include_self_managed_configs: bool = True,
+    k8s_server_address: Optional[str] = None,
 ) -> Dict[str, str]:
     # RFC 1123: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-label-names
     # technically only paasta instance can be longer than 63 chars. But we apply the normalization regardless.
     # NOTE: this affects only k8s labels, not the pod names.
     _paasta_cluster = _get_k8s_resource_name_limit_size_with_hash(paasta_cluster)
     _paasta_service = _get_k8s_resource_name_limit_size_with_hash(paasta_service)
     _paasta_instance = _get_k8s_resource_name_limit_size_with_hash(paasta_instance)
+
     spark_env = {
-        'spark.master': f'k8s://https://k8s.{paasta_cluster}.paasta:6443',
         'spark.executorEnv.PAASTA_SERVICE': paasta_service,
         'spark.executorEnv.PAASTA_INSTANCE': paasta_instance,
         'spark.executorEnv.PAASTA_CLUSTER': paasta_cluster,
@@ -935,7 +937,6 @@ def _get_k8s_spark_env(
         'spark.kubernetes.pyspark.pythonVersion': '3',
         'spark.kubernetes.container.image': docker_img,
         'spark.kubernetes.namespace': 'paasta-spark',
-        'spark.kubernetes.container.image.pullPolicy': 'Always',
         'spark.kubernetes.executor.label.yelp.com/paasta_service': _paasta_service,
         'spark.kubernetes.executor.label.yelp.com/paasta_instance': _paasta_instance,
         'spark.kubernetes.executor.label.yelp.com/paasta_cluster': _paasta_cluster,
@@ -954,14 +955,20 @@ def _get_k8s_spark_env(
                 'spark.kubernetes.authenticate.serviceAccountName': service_account_name,
             },
         )
-    else:
+    elif not include_self_managed_configs:
+        spark_env.update({
+            'spark.master': f'k8s://{k8s_server_address}',
+        })
+    elif include_self_managed_configs:
         spark_env.update(
             {
+                'spark.master': f'k8s://https://k8s.{paasta_cluster}.paasta:6443',
                 'spark.kubernetes.authenticate.caCertFile': f'{K8S_AUTH_FOLDER}/{paasta_cluster}-ca.crt',
                 'spark.kubernetes.authenticate.clientKeyFile': f'{K8S_AUTH_FOLDER}/{paasta_cluster}-client.key',
                 'spark.kubernetes.authenticate.clientCertFile': f'{K8S_AUTH_FOLDER}/{paasta_cluster}-client.crt',
             },
         )
+
     return spark_env
 
 
@@ -1094,6 +1101,8 @@ def get_spark_conf(
     docker_img: str,
     aws_creds: Tuple[Optional[str], Optional[str], Optional[str]],
     extra_volumes: Optional[List[Mapping[str, str]]] = None,
+    use_eks: bool = False,
+    k8s_server_address: Optional[str] = None,
     # the follow arguments only being used for mesos
     extra_docker_params: Optional[MutableMapping[str, str]] = None,
     with_secret: bool = True,
@@ -1201,6 +1210,8 @@ def get_spark_conf(
             extra_volumes,
             paasta_pool,
             service_account_name=service_account_name,
+            include_self_managed_configs=not use_eks,
+            k8s_server_address=k8s_server_address,
         ))
     elif cluster_manager == 'local':
         spark_conf.update(_get_local_spark_env(

tests/spark_config_test.py

@@ -1252,7 +1252,6 @@ def assert_kubernetes_conf(self, base_volumes):
             'spark.kubernetes.authenticate.clientCertFile': (
                 f'{spark_config.K8S_AUTH_FOLDER}/{self.cluster}-client.crt'
             ),
-            'spark.kubernetes.container.image.pullPolicy': 'Always',
             'spark.kubernetes.executor.label.yelp.com/paasta_service': self.service,
             'spark.kubernetes.executor.label.yelp.com/paasta_instance': self.instance,
             'spark.kubernetes.executor.label.yelp.com/paasta_cluster': self.cluster,