feat: add support for Kubernetes service account name spec

snakemake/__init__.py

@@ -174,6 +174,7 @@ def snakemake(
     kubernetes=None,
     container_image=None,
     k8s_cpu_scalar=1.0,
+    k8s_service_account_name=None,
     flux=False,
     tibanna=False,
     tibanna_sfn=None,
@@ -312,6 +313,7 @@ def snakemake(
         kubernetes (str):           submit jobs to Kubernetes, using the given namespace.
         container_image (str):      Docker image to use, e.g., for Kubernetes.
         k8s_cpu_scalar (float):     What proportion of each k8s node's CPUs are availabe to snakemake?
+        k8s_service_account_name:   Custom k8s service account, needed for workload identity.
         flux (bool):                Launch workflow to flux cluster.
         default_remote_provider (str): default remote provider to use instead of local files (e.g. S3, GS)
         default_remote_prefix (str): prefix for default remote provider (e.g. name of the bucket).
@@ -726,6 +728,7 @@ def snakemake(
                     kubernetes=kubernetes,
                     container_image=container_image,
                     k8s_cpu_scalar=k8s_cpu_scalar,
+                    k8s_service_account_name=k8s_service_account_name,
                     conda_create_envs_only=conda_create_envs_only,
                     default_remote_provider=default_remote_provider,
                     default_remote_prefix=default_remote_prefix,
@@ -786,6 +789,7 @@ def snakemake(
                     kubernetes=kubernetes,
                     container_image=container_image,
                     k8s_cpu_scalar=k8s_cpu_scalar,
+                    k8s_service_account_name=k8s_service_account_name,
                     tibanna=tibanna,
                     tibanna_sfn=tibanna_sfn,
                     google_lifesciences=google_lifesciences,
@@ -2409,6 +2413,16 @@ def get_argument_parser(profile=None):
         "see the original value, i.e. as the value substituted in {threads}.",
     )
 
+    group_kubernetes.add_argument(
+        "--k8s-service-account-name",
+        metavar="SERVICEACCOUNTNAME",
+        default=None,
+        help="This argument allows the use of customer service accounts for "
+        "kubernetes pods. If specified serviceAccountName will be added to the "
+        "pod specs. This is needed when using workload identity which is enforced "
+        "when using Google Cloud GKE Autopilot.",
+    )
+
     group_tibanna.add_argument(
         "--tibanna",
         action="store_true",
@@ -3009,6 +3023,7 @@ def open_browser():
             kubernetes=args.kubernetes,
             container_image=args.container_image,
             k8s_cpu_scalar=args.k8s_cpu_scalar,
+            k8s_service_account_name=args.k8s_service_account_name,
             flux=args.flux,
             tibanna=args.tibanna,
             tibanna_sfn=args.tibanna_sfn,

snakemake/executors/__init__.py

@@ -1679,6 +1679,7 @@ def __init__(
         namespace,
         container_image=None,
         k8s_cpu_scalar=1.0,
+        k8s_service_account_name=None,
         jobname="{rulename}.{jobid}",
         printreason=False,
         quiet=False,
@@ -1720,6 +1721,7 @@ def __init__(
         import kubernetes.client
 
         self.k8s_cpu_scalar = k8s_cpu_scalar
+        self.k8s_service_account_name = k8s_service_account_name
         self.kubeapi = kubernetes.client.CoreV1Api()
         self.batchapi = kubernetes.client.BatchV1Api()
         self.namespace = namespace
@@ -1901,6 +1903,10 @@ def run(self, job, callback=None, submit_callback=None, error_callback=None):
         body.spec = kubernetes.client.V1PodSpec(
             containers=[container], node_selector=node_selector
         )
+        #Add service account name if provided
+        if self.k8s_service_account_name:
+            body.spec.service_account_name = self.k8s_service_account_name
+
         # fail on first error
         body.spec.restart_policy = "Never"
 

snakemake/scheduler.py

@@ -87,6 +87,7 @@ def __init__(
         env_modules=None,
         kubernetes=None,
         k8s_cpu_scalar=1.0,
+        k8s_service_account_name=None,
         container_image=None,
         flux=None,
         tibanna=None,
@@ -317,6 +318,7 @@ def __init__(
                 kubernetes,
                 container_image=container_image,
                 k8s_cpu_scalar=k8s_cpu_scalar,
+                k8s_service_account_name=k8s_service_account_name,
                 printreason=printreason,
                 quiet=quiet,
                 printshellcmds=printshellcmds,

snakemake/workflow.py

@@ -530,6 +530,7 @@ def execute(
         drmaa_log_dir=None,
         kubernetes=None,
         k8s_cpu_scalar=1.0,
+        k8s_service_account_name=None,
         flux=None,
         tibanna=None,
         tibanna_sfn=None,
@@ -983,6 +984,7 @@ def files(items):
             drmaa_log_dir=drmaa_log_dir,
             kubernetes=kubernetes,
             k8s_cpu_scalar=k8s_cpu_scalar,
+            k8s_service_account_name=k8s_service_account_name,
             flux=flux,
             tibanna=tibanna,
             tibanna_sfn=tibanna_sfn,