Autoscale node resize (#128)

* Added unit tests and documentation updates * Added example autoscaling formula * Accepting pool name as function parameter instead of config parameter * Update documentation * Autosize resize time reduction * Update setup.py version * Update pyproject.toml version * Upload docker image function (#129) * new docker upload image * updated docker image upload functions * included missing : * add print statement * bump version * Added unit tests and documentation updates (#127) * Added unit tests and documentation updates * Added example autoscaling formula * Accepting pool name as function parameter instead of config parameter * Update documentation * change inputs to helpers.update_pool * update parameters for helpers.Update_pool in clients * Validate parameters being passed * Added unit tests * Add unit tests and moved check_autoscale_parameters params to helpers * bump setup.py version * bump pyproject.toml version * Changed scaling mode to parameter * Update setup.py version * Update pyproject.toml version --------- Co-authored-by: Ryan Raasch <[email protected]> * Removed unnecessary tests --------- Co-authored-by: Ryan Raasch <[email protected]>
CDCgov · Sep 26, 2024 · fede4c8 · fede4c8
1 parent dd1474c
commit fede4c8
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 8 deletions.
diff --git a/cfa_azure/helpers.py b/cfa_azure/helpers.py
@@ -270,6 +270,8 @@ def get_batch_pool_json(
     output_container_name: str,
     config: dict,
     autoscale_formula_path: str,
+    autoscale_evaluation_interval: str = "PT5M",
+    fixedscale_resize_timeout: str = "PT15M"
 ):
     """creates a json output with various components needed for batch pool creation
 
@@ -385,7 +387,7 @@ def get_batch_pool_json(
                 #     "resizeTimeout": "PT15M"
                 # }
                 "autoScale": {
-                    "evaluationInterval": "PT5M",
+                    "evaluationInterval": autoscale_evaluation_interval,
                     "formula": get_autoscale_formula(
                         filepath=autoscale_formula_path
                     ),
@@ -394,7 +396,7 @@ def get_batch_pool_json(
             "resizeOperationStatus": {
                 "targetDedicatedNodes": 1,
                 "nodeDeallocationOption": "Requeue",
-                "resizeTimeout": "PT15M",
+                "resizeTimeout": fixedscale_resize_timeout,
                 "startTime": "2023-07-05T13:18:25.7572321Z",
             },
             "currentDedicatedNodes": 1,
@@ -1221,6 +1223,7 @@ def get_pool_parameters(
     config: dict,
     mount_config: list,
     autoscale_formula_path: str = None,
+    autoscale_evaluation_interval: str = 'PT5M',
     timeout: int = 60,
     dedicated_nodes: int = 1,
     low_priority_nodes: int = 0,
@@ -1251,19 +1254,21 @@ def get_pool_parameters(
     logger.debug(
         f"Setting up pool parameters in '{mode}' mode with timeout={timeout} minutes..."
     )
+    fixedscale_resize_timeout = 'PT15M'
     if mode == "fixed":
+        fixedscale_resize_timeout = f"PT{timeout}M"
         scale_settings = {
             "fixedScale": {
                 "targetDedicatedNodes": dedicated_nodes,
                 "targetLowPriorityNodes": low_priority_nodes,
-                "resizeTimeout": f"PT{timeout}M",
+                "resizeTimeout": fixedscale_resize_timeout
             }
         }
         logger.debug("Fixed mode set with scale settings.")
     elif mode == "autoscale" and use_default_autoscale_formula is False:
         scale_settings = {
             "autoScale": {
-                "evaluationInterval": "PT5M",
+                "evaluationInterval": autoscale_evaluation_interval,
                 "formula": get_autoscale_formula(
                     filepath=autoscale_formula_path
                 ),
@@ -1273,7 +1278,7 @@ def get_pool_parameters(
     elif mode == "autoscale" and use_default_autoscale_formula is True:
         scale_settings = {
             "autoScale": {
-                "evaluationInterval": "PT5M",
+                "evaluationInterval": autoscale_evaluation_interval,
                 "formula": generate_autoscale_formula(
                     max_nodes=max_autoscale_nodes
                 ),
@@ -1302,7 +1307,7 @@ def get_pool_parameters(
             "resizeOperationStatus": {
                 "targetDedicatedNodes": 1,
                 "nodeDeallocationOption": "Requeue",
-                "resizeTimeout": "PT15M",
+                "resizeTimeout": fixedscale_resize_timeout,
                 "startTime": "2023-07-05T13:18:25.7572321Z",
             },
             "currentDedicatedNodes": 1,

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cfa_azure"
-version = "0.3.6"
+version = "0.3.7"
 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.6",
+    version="0.3.7",
     description="module for use with Azure and Azure Batch",
     packages=find_packages(exclude=["tests", "venv"]),
     author="Ryan Raasch",

diff --git a/tests/helpers_tests.py b/tests/helpers_tests.py
@@ -84,6 +84,22 @@ def test_get_batch_pool_json(self):
             }
         )
 
+    def test_get_batch_pool_json_custominterval(self):
+        batch_json = cfa_azure.helpers.get_batch_pool_json(
+            input_container_name=FAKE_INPUT_CONTAINER,
+            output_container_name=FAKE_OUTPUT_CONTAINER,
+            config=FAKE_CONFIG,
+            autoscale_formula_path=FAKE_AUTOSCALE_FORMULA,
+            autoscale_evaluation_interval="PT35M",
+            fixedscale_resize_timeout="PT45M"
+        )
+        self.assertEqual(
+            batch_json['pool_parameters']['properties']['scaleSettings']['autoScale']['evaluationInterval'], "PT35M"
+        )
+        self.assertEqual(
+            batch_json['pool_parameters']['properties']['resizeOperationStatus']['resizeTimeout'], "PT45M"
+        )
+
     def test_format_extensions(self):
         extension = 'csv'
         formatted = cfa_azure.helpers.format_extensions(extension)