MLCOMPUTE-967 | changes for Spark drivers on k8s

paasta_tools/tron_tools.py

@@ -98,7 +98,7 @@
 DEFAULT_AWS_REGION = "us-west-2"
 EXECUTOR_TYPE_TO_NAMESPACE = {
     "paasta": "tron",
-    "spark": "paasta-spark",
+    "spark": "tron",
 }
 DEFAULT_TZ = "US/Pacific"
 clusterman_metrics, _ = get_clusterman_metrics()
@@ -291,16 +291,6 @@ def __init__(
         self.for_validation = for_validation
 
     def _build_spark_config(self) -> Dict[str, str]:
-        # this service account will either be used solely for k8s access or for
-        # k8s + AWS access - the executor will have its own account though: and
-        # only if users opt into using pod identity
-        driver_service_account = create_or_find_service_account_name(
-            iam_role=self.get_iam_role(),
-            namespace=EXECUTOR_TYPE_TO_NAMESPACE[self.get_executor()],
-            k8s_role=_spark_k8s_role(),
-            dry_run=self.for_validation,
-        )
-
         # this is pretty similar to what's done in service_configuration_lib, but it doesn't seem
         # worth refactoring things there when it's unlikely that other users of that library will
         # have the k8s creds to deal with service account creation (or the role bindings for the
@@ -322,9 +312,6 @@ def _build_spark_config(self) -> Dict[str, str]:
             "spark.executorEnv.PAASTA_CLUSTER": self.get_cluster(),
             "spark.executorEnv.PAASTA_INSTANCE_TYPE": "spark",
             "spark.executorEnv.SPARK_EXECUTOR_DIRS": "/tmp",
-            # XXX: do we need to set this for the driver if the pod running the driver was created with
-            # this SA already?
-            "spark.kubernetes.authenticate.driver.serviceAccountName": driver_service_account,
             "spark.kubernetes.pyspark.pythonVersion": "3",
             "spark.kubernetes.container.image": self.get_docker_url(
                 system_paasta_config
@@ -415,7 +402,8 @@ def get_job_name(self):
     def get_action_name(self):
         return self.action
 
-    # mypy does not like the SecretVolume -> TronSecretVolume conversion, because TypedDict inheritence is broken. Until this is fixed, let's ignore this issue.
+    # mypy does not like the SecretVolume -> TronSecretVolume conversion, because TypedDict inheritence is broken.
+    # Until this is fixed, let's ignore this issue.
     def get_secret_volumes(self) -> List[TronSecretVolume]:  # type: ignore
         """Adds the secret_volume_name to the objet so tron/task_processing can load it downstream without replicating code."""
         secret_volumes = super().get_secret_volumes()
@@ -491,6 +479,14 @@ def get_env(
 
         return env
 
+    def get_iam_role(self) -> str:
+        iam_role = super().get_iam_role()
+
+        if not iam_role and self.get_executor() == "spark":
+            iam_role = load_system_paasta_config().get_spark_driver_iam_role()
+
+        return iam_role
+
     def get_secret_env(self) -> Mapping[str, dict]:
         base_env = self.config_dict.get("env", {})
         secret_env = {}
@@ -912,7 +908,7 @@ def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = Fal
         "triggered_by": action_config.get_triggered_by(),
         "on_upstream_rerun": action_config.get_on_upstream_rerun(),
         "trigger_timeout": action_config.get_trigger_timeout(),
-        # outside of Spark usescases, we also allow users to specify an expected-to-exist Service Account name
+        # outside of Spark use-cases, we also allow users to specify an expected-to-exist Service Account name
         # in the Tron namespace in case an action needs specific k8s permissions (e.g., a Jolt batch may need
         # k8s permissions to list Jolt pods in the jolt namespace to do science™ to them).
         # if the provided Service Account does not exist, Tron should simply fail to create the Podspec and report
@@ -921,7 +917,7 @@ def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = Fal
         "service_account_name": action_config.get_service_account_name(),
     }
 
-    # while we're tranisitioning, we want to be able to cleanly fallback to Mesos
+    # while we're transitioning, we want to be able to cleanly fall back to Mesos,
     # so we'll default to Mesos unless k8s usage is enabled for both the cluster
     # and job.
     # there are slight differences between k8s and Mesos configs, so we'll translate
@@ -1002,7 +998,6 @@ def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = Fal
             result["service_account_name"] = create_or_find_service_account_name(
                 iam_role=action_config.get_iam_role(),
                 namespace=EXECUTOR_TYPE_TO_NAMESPACE[executor],
-                k8s_role=_spark_k8s_role(),
                 dry_run=action_config.for_validation,
             )
             # spark, unlike normal batches, needs to expose  several ports for things like the spark

paasta_tools/utils.py

@@ -1881,6 +1881,7 @@ class RemoteRunConfig(TypedDict, total=False):
 class SparkRunConfig(TypedDict, total=False):
     default_cluster: str
     default_pool: str
+    default_spark_driver_iam_role: str
 
 
 class PaastaNativeConfig(TypedDict, total=False):
@@ -2771,6 +2772,9 @@ def get_eks_cluster_aliases(self) -> Dict[str, str]:
     def get_cluster_pools(self) -> Dict[str, List[str]]:
         return self.config_dict.get("allowed_pools", {})
 
+    def get_spark_driver_iam_role(self) -> str:
+        return self.get_spark_run_config().get("default_spark_driver_iam_role", "")
+
     def get_hacheck_match_initial_delay(self) -> bool:
         return self.config_dict.get("hacheck_match_initial_delay", False)