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()
@@ -488,9 +476,18 @@ def get_env(
             # if this changes and we do need it - please add a comment about *why* we need it!
             # XXX: update PAASTA_RESOURCE_* env vars to use the correct value from spark_args and set
             # these to the correct values for the executors as part of the driver commandline
+            env["KUBECONFIG"] = system_paasta_config.get_spark_kubeconfig()
 
         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 = {}
@@ -896,7 +893,7 @@ def format_master_config(master_config, default_volumes, dockercfg_location):
 def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = False):
     """Generate a dict of tronfig for an action, from the TronActionConfig.
 
-    :param job_config: TronActionConfig
+    :param action_config: TronActionConfig
     """
     executor = action_config.get_executor()
     result = {
@@ -912,7 +909,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 +918,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
@@ -997,20 +994,25 @@ def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = Fal
             )
 
         if executor == "spark":
+            system_paasta_config = load_system_paasta_config()
             # this service account will only be used by Spark drivers since executors don't
             # need Kubernetes access permissions
             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
+            result["extra_volumes"] = [{
+                "container_path": system_paasta_config.get_spark_kubeconfig(),
+                "host_path": system_paasta_config.get_spark_kubeconfig(),
+                "mode": "RO",
+            }]
+            # spark, unlike normal batches, needs to expose several ports for things like the spark
             # ui and for executor->driver communication
             result["ports"] = list(
                 set(
                     _get_spark_ports(
-                        system_paasta_config=load_system_paasta_config()
+                        system_paasta_config=system_paasta_config
                     ).values()
                 )
             )
@@ -1034,7 +1036,10 @@ def format_tron_action_dict(action_config: TronActionConfig, use_k8s: bool = Fal
         result["cpus"] = action_config.get_cpus()
         result["mem"] = action_config.get_mem()
         result["disk"] = action_config.get_disk()
-        result["extra_volumes"] = format_volumes(action_config.get_extra_volumes())
+        if "extra_volumes" in result:
+            result["extra_volumes"] = result["extra_volumes"] + format_volumes(action_config.get_extra_volumes())
+        else:
+            result["extra_volumes"] = format_volumes(action_config.get_extra_volumes())
         result["docker_image"] = action_config.get_docker_url()
 
     # Only pass non-None values, so Tron will use defaults for others

paasta_tools/utils.py

@@ -1887,6 +1887,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):
@@ -2777,6 +2778,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)