Taskslots (#126)

* task slots update * bump the version
CDCgov · Oct 9, 2024 · bdfcfb1 · bdfcfb1
1 parent 97b21b9
commit bdfcfb1
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 15 deletions.
diff --git a/cfa_azure/clients.py b/cfa_azure/clients.py
@@ -108,6 +108,7 @@ def set_pool_info(
         dedicated_nodes=1,
         low_priority_nodes=0,
         cache_blobfuse: bool = True,
+        task_slots_per_node: int = 1
     ) -> None:
         """Sets the scaling mode of the client, either "fixed" or "autoscale".
         If "fixed" is selected, debug must be turned off.
@@ -122,6 +123,7 @@ def set_pool_info(
             dedicated_nodes (int, optional): number of dedicated nodes for the pool. Defaults to 1.
             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.
         """
         # check if debug and scaling mode match, otherwise alert the user
         if self.debug is True and mode == "autoscale":
@@ -163,18 +165,19 @@ def set_pool_info(
             self.low_priority_nodes = low_priority_nodes
             # create batch_json with fixed
             self.pool_parameters = helpers.get_pool_parameters(
-                mode,
-                self.container_image_name,
-                self.registry_url,
-                self.container_registry_server,
-                self.config,
-                self.mount_config,
-                autoscale_formula_path,
-                timeout,
-                dedicated_nodes,
-                low_priority_nodes,
-                use_default_autoscale_formula,
-                max_autoscale_nodes,
+                mode = mode,
+                container_image_name= self.container_image_name,
+                container_registry_url=self.registry_url,
+                container_registry_server=self.container_registry_server,
+                config=self.config,
+                mount_config=self.mount_config,
+                autoscale_formula_path=autoscale_formula_path,
+                timeout=timeout,
+                dedicated_nodes=dedicated_nodes,
+                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
             )
             logger.debug("pool parameters generated")
         else:

diff --git a/cfa_azure/helpers.py b/cfa_azure/helpers.py
@@ -1202,6 +1202,7 @@ def get_pool_parameters(
     low_priority_nodes: int = 0,
     use_default_autoscale_formula: bool = False,
     max_autoscale_nodes: int = 3,
+    task_slots_per_node: int = 1
 ):
     """creates a pool parameter dictionary to be used with pool creation.
 
@@ -1217,6 +1218,8 @@ def get_pool_parameters(
         dedicated_nodes (int, optional): number of dedicated nodes. Defaults to 1.
         low_priority_nodes (int, optional): number of low priority nodes. Defaults to 0.
         use_default_autoscale_formula (bool, optional)
+        max_autoscale_nodes (int): maximum number of nodes to use with autoscaling. Default 3.
+        task_slots_per_node (int): number of task slots per node. Default is 1.
 
     Returns:
         dict: dict of pool parameters for pool creation
@@ -1262,7 +1265,7 @@ def get_pool_parameters(
         "properties": {
             "vmSize": config["Batch"]["pool_vm_size"],
             "interNodeCommunication": "Disabled",
-            "taskSlotsPerNode": 1,
+            "taskSlotsPerNode": task_slots_per_node,
             "taskSchedulingPolicy": {"nodeFillType": "Spread"},
             "deploymentConfiguration": get_deployment_config(
                 container_image_name,

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cfa_azure"
-version = "0.3.3"
+version = "0.3.4"
 description = "module for use with Azure and Azure Batch"
 authors = ["Ryan Raasch <[email protected]>"]
 readme = "README.md"

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="cfa_azure",
-    version="0.3.3",
+    version="0.3.4",
     description="module for use with Azure and Azure Batch",
     packages=find_packages(exclude=["tests", "venv"]),
     author="Ryan Raasch",