feat: Add support for customer compute service account (#51, @hnawar) (#52)

Test for PR #51.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Added support for specifying a custom compute service account in the
`hello_world` rule, enhancing flexibility for cloud interactions.
- Introduced a new method to retrieve service account information,
providing better control over job allocations.
  
- **Bug Fixes**
- Updated `.gitignore` to exclude the `poetry.lock` file, ensuring
cleaner version control.
  
- **Documentation**
- Enhanced documentation for the `hello_world` rule to include the new
`googlebatch_service_account` parameter, improving user guidance.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Hatem Nawar <hnawar@google.com>
Co-authored-by: Hatem <hnawar@google.com>
Co-authored-by: hatem.nawar <hatem.nawar@gmail.com>

Author: 3 people

.gitignore

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

docs/further.md

@@ -119,11 +119,11 @@ This will define the machine type for a particular step, overriding the default
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_machine_type="c3-standard-112"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_machine_type="c3-standard-112"
+    shell:
         "..."
 ```
 
@@ -136,11 +136,11 @@ This will define the image family for a particular step, overriding the default
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_image_family="hpc-centos-7"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_image_family="hpc-centos-7"
+    shell:
         "..."
 ```
 
@@ -160,11 +160,11 @@ This will define the image project for a particular step, overriding the default
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_image_project="cloud-hpc-image-public"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_image_project="cloud-hpc-image-public"
+    shell:
         "..."
 ```
 
@@ -175,11 +175,11 @@ This will define the bucket for a particular step, overriding the default from t
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_bucket="my-snakemake-batch-bucket"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_bucket="my-snakemake-batch-bucket"
+    shell:
         "..."
 ```
 
@@ -189,11 +189,11 @@ This will define the mount path for a bucket for a particular step, overriding t
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_mount_path="/mnt/workflow"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_mount_path="/mnt/workflow"
+    shell:
         "..."
 ```
 
@@ -204,11 +204,11 @@ This will define the work tasks for a particular step, overriding the default fr
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_work_tasks=1
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_work_tasks=1
+    shell:
         "..."
 ```
 
@@ -218,11 +218,11 @@ The URL of an existing network resource (e.g., `projects/{project}/global/networ
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_network="projects/{project}/global/networks/{network}"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_network="projects/{project}/global/networks/{network}"
+    shell:
         "..."
 ```
 
@@ -232,11 +232,25 @@ The URL of an existing subnetwork resource (e.g., `projects/{project}/regions/{r
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_subnetwork="projects/{project}/regions/{region}/subnetworks/{subnetwork}"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_subnetwork="projects/{project}/regions/{region}/subnetworks/{subnetwork}"
+    shell:
+        "..."
+```
+
+#### googlebatch_service_account
+
+The email of custom compute service account to be used by Batch (e.g., `snakemake-sa@projectid.iam.gserviceaccount.com`)
+
+```console
+rule hello_world:
+    output:
+        "...",
+    resources: 
+        googlebatch_service_account="snakemake-sa@projectid.iam.gserviceaccount.com"
+    shell:
         "..."
 ```
 
@@ -246,11 +260,11 @@ This will define the milliseconds per cpu-second for a particular step, overridi
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_cpu_mulli=2000
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_cpu_mulli=2000
+    shell:
         "..."
 ```
 
@@ -260,11 +274,11 @@ This will define the work tasks per node (Google batch calls these tasks) for a
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_work_tasks_per_node=2
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_work_tasks_per_node=2
+    shell:
         "..."
 ```
 
@@ -274,11 +288,11 @@ This will define the memory for a particular step as an integer in MiB, overridi
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_memory=2000
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_memory=2000
+    shell:
         "..."
 ```
 
@@ -289,11 +303,11 @@ This is the [boot disk type](https://cloud.google.com/compute/docs/disks#pdspecs
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_boot_disk_type="pd-standard"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_boot_disk_type="pd-standard"
+    shell:
         "..."
 ```
 
@@ -304,11 +318,11 @@ This is the boot disk [image](https://github.com/googleapis/googleapis/blob/2fd7
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_boot_disk_image="batch-centos"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_boot_disk_image="batch-centos"
+    shell:
         "..."
 ```
 
@@ -319,11 +333,11 @@ The [size of the boot disk](https://github.com/googleapis/googleapis/blob/2fd762
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_boot_disk_gb=40
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_boot_disk_gb=40
+    shell:
         "..."
 ```
 
@@ -333,11 +347,11 @@ This will define the retry times for a step overriding the default from the comm
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_retry_count=2
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_retry_count=2
+    shell:
         "..."
 ```
 
@@ -347,11 +361,11 @@ This will define the max run duration for a step overriding the default from the
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_max_run_duration="3600s"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_max_run_duration="3600s"
+    shell:
         "..."
 ```
 
@@ -361,11 +375,11 @@ This will define the extra labels to add to the Google Batch job.
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_labels="model=c3,stage=test"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_labels="model=c3,stage=test"
+    shell:
         "..."
 ```
 
@@ -376,11 +390,11 @@ A container to use only with `image_family` set to batch-cos* (see [here](https:
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_container="ghcr.io/rse-ops/atacseq:app-latest"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_container="ghcr.io/rse-ops/atacseq:app-latest"
+    shell:
         "..."
 ```
 
@@ -391,10 +405,10 @@ One or more named (or file-derived) snippets to add to setup.
 
 ```console
 rule hello_world:
-	output:
-		"...",
-	resources: 
-		googlebatch_snippets="mpi,myscript.sh"
-	shell:
+    output:
+        "...",
+    resources: 
+        googlebatch_snippets="mpi,myscript.sh"
+    shell:
         "..."
 ```

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: JobExecutorInterface) -> batch_v1.ServiceAccount:
+        """
+        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.