Fix missing parameter for _add_instances_for_nodes

Fix missing parameter assign_node_batch_size for _add_instances_for_nodes Signed-off-by: Luca Carrogu <[email protected]>
aws · Sep 26, 2023 · ea58da0 · ea58da0
1 parent 5bb429d
commit ea58da0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/src/slurm_plugin/instance_manager.py b/src/slurm_plugin/instance_manager.py
@@ -637,6 +637,7 @@ def _scaling_for_jobs_single_node(
                 self._add_instances_for_nodes(
                     node_list=single_nodes,
                     launch_batch_size=launch_batch_size,
+                    assign_node_batch_size=assign_node_batch_size,
                     update_node_address=update_node_address,
                     all_or_nothing_batch=False,
                 )

diff --git a/tests/slurm_plugin/test_instance_manager.py b/tests/slurm_plugin/test_instance_manager.py
@@ -3376,6 +3376,7 @@ def test_scaling_for_jobs_single_node(
             instance_manager._add_instances_for_nodes.assert_called_once_with(
                 node_list=expected_single_nodes_no_oversubscribe,
                 launch_batch_size=launch_batch_size,
+                assign_node_batch_size=assign_node_batch_size,
                 update_node_address=update_node_address,
                 all_or_nothing_batch=False,
             )