From f39656fd4bfe0a841955ec861d18ff98186ee8bd Mon Sep 17 00:00:00 2001
From: Ryan Raasch <xng3@cdc.gov>
Date: Wed, 11 Sep 2024 09:45:15 -0700
Subject: [PATCH 1/2] task slots update

---
 cfa_azure/clients.py | 27 +++++++++++++++------------
 cfa_azure/helpers.py |  5 ++++-
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/cfa_azure/clients.py b/cfa_azure/clients.py
index d4a5c13..226b291 100644
--- 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
index 11adead..07319d9 100644
--- a/cfa_azure/helpers.py
+++ b/cfa_azure/helpers.py
@@ -1200,6 +1200,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.
 
@@ -1215,6 +1216,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
@@ -1260,7 +1263,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,

From 3143863fcc8333d658b5ab2a2bfa677a36e7c374 Mon Sep 17 00:00:00 2001
From: Ryan Raasch <xng3@cdc.gov>
Date: Wed, 11 Sep 2024 10:21:27 -0700
Subject: [PATCH 2/2] bump the version

---
 pyproject.toml | 2 +-
 setup.py       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 586ee73..7c32361 100644
--- 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 <xng3@cdc.gov>"]
 readme = "README.md"
diff --git a/setup.py b/setup.py
index d9c8961..5893b50 100644
--- 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",