feat: Add more configuration options

pyproject.toml

@@ -3,6 +3,7 @@ name = "snakemake-executor-plugin-azure-batch"
 version = "0.3.0"
 description = "A Snakemake executor plugin for submitting jobs to Microsoft Azure Batch."
 authors = [
+  "Simon Ungar Felding <sife@adm.ku.dk>",
   "Jake VanCampen <jake.vancampen7@gmail.com>",
   "Johannes Koester <johannes.koester@uni-due.de>"
 ]
@@ -18,8 +19,8 @@ snakemake-interface-common = "^1.17.2"
 snakemake-interface-executor-plugins = "^9.2.0"
 azure-storage-blob = "^12.20.0"
 azure-batch = "^14.2.0"
-azure-mgmt-batch = "^17.0.0"
-azure-identity = "^1.17.1"
+azure-mgmt-batch = "^18.0.0"
+azure-identity = "^1.25.0"
 msrest = "^0.7.1"
 
 [tool.poetry.group.dev.dependencies]

snakemake_executor_plugin_azure_batch/__init__.py

@@ -1,6 +1,6 @@
-__author__ = "Jake VanCampen, Johannes Köster, Andreas Wilm"
-__copyright__ = "Copyright 2023, Snakemake community"
-__email__ = "jake.vancampen7@gmail.com"
+__author__ = "Simon Ungar Felding, Jake VanCampen, Johannes Köster, Andreas Wilm"
+__copyright__ = "Copyright 2025, Snakemake community"
+__email__ = "sife@adm.ku.dk"
 __license__ = "MIT"
 
 
@@ -62,6 +62,14 @@
 # Omit this class if you don't need any.
 @dataclass
 class ExecutorSettings(ExecutorSettingsBase):
+    accelerated_networking_enabled: bool = field(
+        default=False,
+        metadata={
+            "help": "Enable accelerated networking on the batch pool nodes.",
+            "required": False,
+            "env_var": False,
+        },
+    )
     account_url: Optional[str] = field(
         default=None,
         metadata={
@@ -105,6 +113,14 @@ class ExecutorSettings(ExecutorSettingsBase):
             "env_var": True,
         },
     )
+    container_run_options: Optional[str] = field(
+        default="--rm",
+        metadata={
+            "help": "Additional docker run options for the container execution.",
+            "required": False,
+            "env_var": False,
+        },
+    )
     keep_pool: bool = field(
         default=False,
         metadata={
@@ -138,7 +154,18 @@ class ExecutorSettings(ExecutorSettingsBase):
         metadata={
             "help": "Azure Batch node start task bash script url."
             "This can be any url that hosts your start task bash script. Azure blob SAS"
-            "urls work nicely here",
+            "urls work nicely here."
+            "Cannot be used along with node_start_task.",
+            "required": False,
+            "env_var": False,
+        },
+    )
+    node_start_task: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Azure Batch node start task bash script."
+            "Example: echo Hello World > /mnt/batch/tasks/startup/welcome.txt"
+            "Cannot be used along with node_start_task_url.",
             "required": False,
             "env_var": False,
         },
@@ -152,13 +179,25 @@ class ExecutorSettings(ExecutorSettingsBase):
         },
     )
     node_communication_mode: Optional[str] = field(
-        default=None,
+        default="default",
         metadata={
             "help": "Azure Batch node communication mode.",
             "required": False,
             "env_var": False,
         },
     )
+    pool_mount_configuration: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Azure Batch pool mount configuration in JSON format."
+            "Each mount must be an object in the JSON array."
+            "See https://learn.microsoft.com/en-us/azure/batch/batch-mount-azure-storage?tabs=azure-portal#mount-configuration "
+            "for details on the mount configuration format."
+            "Example: [{'azureBlobFileSystemConfiguration': {'accountName': '123', 'containerName': 'ok', 'relativeMountPath': 'data'}}]",
+            "required": False,
+            "env_var": False,
+        },
+    )
     pool_subnet_id: Optional[str] = field(
         default=None,
         metadata={
@@ -184,15 +223,15 @@ class ExecutorSettings(ExecutorSettingsBase):
         },
     )
     pool_image_sku: str = field(
-        default="20-04-lts",
+        default="22-04-lts",
         metadata={
             "help": "Batch pool image sku.",
             "required": False,
             "env_var": False,
         },
     )
     pool_vm_node_agent_sku_id: str = field(
-        default="batch.node.ubuntu 20.04",
+        default="batch.node.ubuntu 22.04",
         metadata={
             "help": "Azure batch pool vm node agent sku id.",
             "required": False,
@@ -215,6 +254,14 @@ class ExecutorSettings(ExecutorSettingsBase):
             "env_var": False,
         },
     )
+    public_address_provisioning_type: str = field(
+        default="NoPublicIPAddresses",
+        metadata={
+            "help": "Azure Batch public address provisioning type.",
+            "required": False,
+            "env_var": False,
+        },
+    )
     resource_group_name: Optional[str] = field(
         default=None,
         metadata={
@@ -351,7 +398,7 @@ def run_job(self, job: JobExecutorInterface):
         self.logger.debug(f"Remote command: {remote_command}")
 
         task: bm.TaskAddParameter = build.batch_task(
-            job, self.container_image, self.envvars(), remote_command
+            job, self.container_image, self.envvars(), remote_command,  self.settings
         )
 
         job_info = SubmittedJobInfo(job, external_jobid=task.id)

snakemake_executor_plugin_azure_batch/build.py

@@ -1,5 +1,6 @@
 import datetime
 import uuid
+import json
 
 import azure.batch.models as bm
 from azure.mgmt.batch.models import (
@@ -15,6 +16,7 @@
     ElevationLevel,
     FixedScaleSettings,
     ImageReference,
+    MountConfiguration,
     NetworkConfiguration,
     NodeCommunicationMode,
     Pool,
@@ -52,9 +54,10 @@ def batch_pool_params(pool_id: str, settings, container_image: str) -> Pool:
         version="latest",
     )
 
-    network_config = None
-    if settings.pool_subnet_id is not None:
-        network_config = NetworkConfiguration(subnet_id=settings.pool_subnet_id)
+    network_config = NetworkConfiguration(subnet_id=settings.pool_subnet_id,
+                                          public_address_provisioning_type=settings.public_address_provisioning_type,
+                                          accelerated_networking_enabled=settings.accelerated_networking_enabled
+                                          )
 
     # configure batch pool identity
     batch_pool_identity = None
@@ -122,14 +125,14 @@ def batch_pool_params(pool_id: str, settings, container_image: str) -> Pool:
 
     # if configured use start task bash script from url
     # can be SAS url or other accessible url hosting bash script
-    if settings.node_start_task_url is not None:
-        _SIMPLE_TASK_NAME = "start_task.sh"
-        start_task_admin = UserIdentity(
-            auto_user=AutoUserSpecification(
-                elevation_level=ElevationLevel.ADMIN,
-                scope=AutoUserScope.POOL,
-            )
+    start_task_admin = UserIdentity(
+        auto_user=AutoUserSpecification(
+            elevation_level=ElevationLevel.ADMIN,
+            scope=AutoUserScope.POOL,
         )
+    )
+    if settings.node_start_task_url:
+        _SIMPLE_TASK_NAME = "start_task.sh"
         start_task_conf = StartTask(
             command_line=f"bash {_SIMPLE_TASK_NAME}",
             resource_files=[
@@ -140,7 +143,15 @@ def batch_pool_params(pool_id: str, settings, container_image: str) -> Pool:
             ],
             user_identity=start_task_admin,
         )
-
+    if settings.node_start_task:
+        start_task_conf = StartTask(
+            command_line=settings.node_start_task,
+            user_identity=start_task_admin,
+        )
+    if settings.node_start_task and settings.node_start_task_url:
+        raise WorkflowError(
+            "You cannot set both node_start_task and node_start_task_url."
+        )
     # auto scale requires the initial dedicated node count to be zero
     # min allowed interval of five minutes
     if settings.autoscale:
@@ -156,6 +167,24 @@ def batch_pool_params(pool_id: str, settings, container_image: str) -> Pool:
         fixed_scale=FixedScaleSettings(target_dedicated_nodes=settings.pool_node_count)
     )
 
+    if settings.pool_mount_configuration:
+        mount_configuration = MountConfiguration.deserialize(
+            settings.pool_mount_configuration
+        )
+        try:
+            json.loads(settings.pool_mount_configuration)
+        except Exception as e:
+            raise WorkflowError(
+                f"Invalid mount configuration (invalid JSON): "
+                f"{settings.pool_mount_configuration}"
+            ) from e
+        if mount_configuration is None:  # This tests if it parsed correctly by the SDK
+            raise WorkflowError(
+                f"Invalid mount configuration: {settings.pool_mount_configuration}"
+            )
+    else:
+        mount_configuration = None
+
     return Pool(
         identity=batch_pool_identity,
         display_name=pool_id,
@@ -174,7 +203,10 @@ def batch_pool_params(pool_id: str, settings, container_image: str) -> Pool:
         task_scheduling_policy=TaskSchedulingPolicy(
             node_fill_type=settings.node_fill_type
         ),
-        target_node_communication_mode=NodeCommunicationMode.CLASSIC,
+        target_node_communication_mode=NodeCommunicationMode(
+            settings.node_communication_mode.title()
+            ),
+        mount_configuration=mount_configuration,
     )
 
 
@@ -183,6 +215,7 @@ def batch_task(
     container_image: str,
     envvars: dict,
     remote_task_command: str,
+    settings,
 ) -> bm.TaskAddParameter:
     """
     Creates a batch task for executing a remote task in Azure Batch.
@@ -216,7 +249,7 @@ def batch_task(
         ),
         container_settings=bm.TaskContainerSettings(
             image_name=container_image,
-            container_run_options="--rm",
+            container_run_options=settings.container_run_options,
         ),
         environment_settings=env_settings,
     )