diff --git a/.gitignore b/.gitignore index a598507..a72c935 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__ env .snakemake +poetry.lock \ No newline at end of file diff --git a/docs/further.md b/docs/further.md index daad1e9..fe36d0c 100644 --- a/docs/further.md +++ b/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., `snakemake-sa@projectid.iam.gserviceaccount.com`) + +```console +rule hello_world: + output: + "...", + resources: + googlebatch_service_account="snakemake-sa@projectid.iam.gserviceaccount.com" + shell: + "..." +``` + #### googlebatch_cpu_milli This will define the milliseconds per cpu-second for a particular step, overriding the default from the command line. diff --git a/pyproject.toml b/pyproject.toml index 92edd0d..96c98cd 100644 --- a/pyproject.toml +++ b/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] diff --git a/snakemake_executor_plugin_googlebatch/__init__.py b/snakemake_executor_plugin_googlebatch/__init__.py index 1cb1f31..f8cd049 100644 --- a/snakemake_executor_plugin_googlebatch/__init__.py +++ b/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( diff --git a/snakemake_executor_plugin_googlebatch/executor.py b/snakemake_executor_plugin_googlebatch/executor.py index 509b1d7..e9adb01 100644 --- a/snakemake_executor_plugin_googlebatch/executor.py +++ b/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.