feat: Add support for customer compute service account

.gitignore

@@ -1,3 +1,4 @@
 __pycache__
 env
 .snakemake
+poetry.lock

docs/further.md

@@ -240,6 +240,20 @@ rule hello_world:
         "..."
 ```
 
+#### googlebatch_service_account
+
+The email of custom compute service account to be used by Batch (e.g., `[email protected]`)
+
+```console
+rule hello_world:
+	output:
+		"...",
+	resources: 
+		googlebatch_service_account="[email protected]"
+	shell:
+        "..."
+```
+
 #### googlebatch_cpu_milli
 
 This will define the milliseconds per cpu-second for a particular step, overriding the default from the command line.

pyproject.toml

@@ -28,7 +28,7 @@ black = "^24.4.0"
 flake8 = "^6.1.0"
 coverage = "^7.3.1"
 pytest = "^7.4.2"
-snakemake = {git = "https://github.com/snakemake/snakemake.git", branch="main"}
+snakemake = "^8.18.0"
 snakemake-storage-plugin-s3 = "^0.2.10"
 
 [tool.coverage.run]

snakemake_executor_plugin_googlebatch/__init__.py

@@ -155,6 +155,15 @@ class ExecutorSettings(ExecutorSettingsBase):
         },
     )
 
+    service_account: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The email of a customer compute service account",
+            "env_var": True,
+            "required": False,
+        },
+    )
+
     # local SSD uses type "local-ssd".
     # Also "pd-balanced", "pd-extreme", "pd-ssd", "pd-standard"
     boot_disk_type: Optional[str] = field(

snakemake_executor_plugin_googlebatch/executor.py

@@ -359,6 +359,12 @@ def get_allocation_policy(self, job):
         network_policy = self.get_network_policy(job)
         if network_policy is not None:
             allocation_policy.network = network_policy
+            # Add custom compute service account
+        service_account = self.get_service_account(job)
+
+        if service_account is not None:
+            allocation_policy.service_account = service_account
+
         return allocation_policy
 
     def get_network_policy(self, job):
@@ -379,6 +385,16 @@ def get_network_policy(self, job):
         policy.network_interfaces = [interface]
         return policy
 
+    def get_service_account(self, job):
+        """
+        Givena job request, get the service account
+        """
+        service_account_email = self.get_param(job, "service_account")
+        service_account = batch_v1.ServiceAccount()
+        if service_account_email is not None:
+            service_account.email = service_account_email
+        return service_account
+
     def get_boot_disk(self, job):
         """
         Given a job request, add a customized boot disk.