Allow Using Availability Zones for Pools

README.md

@@ -46,6 +46,17 @@ For example, if we would like to persist stdout and stderr to the blob container
 client.add_job("persisting_test", save_logs_to_blob = "input-test")
 ```
 
+**Availability Zones**  
+To make use of Azure's availability zone functionality there is a parameter available in the `set_pool_info()` method called `availability_zones`. To use availability zones when building a pool, set this parameter to True. If you want to stick with the default Regional configuration, this parameter can be left out or set to False. Turn availability zone on like the following:
+```
+client.set_pool_info(
+  ...
+  availability_zones = True,
+  ...
+)
+```
+
+
 ### Functions
 - create_pool: creates a new Azure batch pool using default autoscale mode   
   Example:

cfa_azure/clients.py

@@ -110,7 +110,8 @@ def set_pool_info(
         dedicated_nodes=0,
         low_priority_nodes=1,
         cache_blobfuse: bool = True,
-        task_slots_per_node: int = 1
+        task_slots_per_node: int = 1,
+        availability_zones = False
     ) -> None:
         """Sets the scaling mode of the client, either "fixed" or "autoscale".
         If "fixed" is selected, debug must be turned off.
@@ -126,6 +127,7 @@ def set_pool_info(
             low_priority_nodes (int, optional): number of low priority nodes for the pool. Defaults to 0.
             cache_blobfuse (bool): True to use blobfuse caching, False to download data from blobfuse every time. Defaults to True.
             task_slots_per_node (int): number of task slots per node. Default 1.
+            availability_zones (bool): whether to use availability zones for the pool. True to use Availability Zones. False to stay Regional. Default False.
         """
         # check if debug and scaling mode match, otherwise alert the user
         if self.debug is True and mode == "autoscale":
@@ -179,7 +181,8 @@ def set_pool_info(
                 low_priority_nodes=low_priority_nodes,
                 use_default_autoscale_formula=use_default_autoscale_formula,
                 max_autoscale_nodes=max_autoscale_nodes,
-                task_slots_per_node= task_slots_per_node
+                task_slots_per_node= task_slots_per_node,
+                availability_zones = availability_zones
             )
             logger.debug("pool parameters generated")
         else:

cfa_azure/helpers.py

@@ -1145,6 +1145,7 @@ def get_deployment_config(
     container_registry_url: str,
     container_registry_server: str,
     config: str,
+    availability_zones: bool = False
 ):
     """gets the deployment config based on the config information
 
@@ -1157,6 +1158,12 @@ def get_deployment_config(
     Returns:
         dict: dictionary containing info for container deployment. Uses ubuntu server with info obtained from config file.
     """
+    logger.debug("setting availability zone info")
+    if availability_zones:
+        policy = "Zonal"
+    else:
+        policy = "Regional"
+
     logger.debug("Getting deployment config.")
     deployment_config = {
         "virtualMachineConfiguration": {
@@ -1166,6 +1173,9 @@ def get_deployment_config(
                 "sku": "20-04-lts",
                 "version": "latest",
             },
+            "nodePlacementConfiguration": {
+                "policy": policy
+            },
             "nodeAgentSkuId": "batch.node.ubuntu 20.04",
             "containerConfiguration": {
                 "type": "dockercompatible",
@@ -1262,7 +1272,8 @@ def get_pool_parameters(
     low_priority_nodes: int = 1,
     use_default_autoscale_formula: bool = False,
     max_autoscale_nodes: int = 3,
-    task_slots_per_node: int = 1
+    task_slots_per_node: int = 1,
+    availability_zones: bool = False
 ):
     """creates a pool parameter dictionary to be used with pool creation.
 
@@ -1334,6 +1345,7 @@ def get_pool_parameters(
                 container_registry_url,
                 container_registry_server,
                 config,
+                availability_zones
             ),
             "networkConfiguration": get_network_config(config),
             "scaleSettings": scale_settings,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cfa_azure"
-version = "1.0.6"
+version = "1.0.7"
 description = "module for use with Azure and Azure Batch"
 authors = ["Ryan Raasch <[email protected]>"]
 readme = "README.md"

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="cfa_azure",
-    version="1.0.6",
+    version="1.0.7",
     description="module for use with Azure and Azure Batch",
     packages=find_packages(exclude=["tests", "venv"]),
     author="Ryan Raasch",