Add support for greedy-all-or-nothing

Signed-off-by: Eddy Mwiti <[email protected]>

src/slurm_plugin/common.py

@@ -38,6 +38,7 @@
 class ScalingStrategy(Enum):
     ALL_OR_NOTHING = "all-or-nothing"
     BEST_EFFORT = "best-effort"
+    GREEDY_ALL_OR_NOTHING = "greedy-all-or-nothing"
 
     @classmethod
     def _missing_(cls, strategy):

src/slurm_plugin/instance_manager.py

@@ -884,7 +884,7 @@ def _add_instances_for_nodes(
                 successful_launched_nodes += slurm_node_list[:q_cr_instances_launched_length]
                 failed_launch_nodes += slurm_node_list[q_cr_instances_launched_length:]
 
-        if scaling_strategy == ScalingStrategy.ALL_OR_NOTHING:
+        if scaling_strategy in [ScalingStrategy.ALL_OR_NOTHING, ScalingStrategy.GREEDY_ALL_OR_NOTHING]:
             self.all_or_nothing_node_assignment(
                 assign_node_batch_size=assign_node_batch_size,
                 instances_launched=instances_launched,

tests/slurm_plugin/test_instance_manager.py

@@ -2995,6 +2995,35 @@ def test_add_instances_for_nodes(
             instance_manager._assign_instances_to_nodes.assert_not_called()
             assert_that(instance_manager.failed_nodes).is_equal_to(expected_failed_nodes)
 
+    @pytest.mark.parametrize(
+        "scaling_strategy, expect_all_or_nothing_node_assignment, expect_best_effort_node_assignment",
+        [
+            (ScalingStrategy.BEST_EFFORT, False, True),
+            (ScalingStrategy.ALL_OR_NOTHING, True, False),
+            (ScalingStrategy.GREEDY_ALL_OR_NOTHING, True, False),
+        ],
+    )
+    def test_node_assignment_by_scaling_strategy(
+        self,
+        mocker,
+        instance_manager,
+        scaling_strategy,
+        expect_all_or_nothing_node_assignment,
+        expect_best_effort_node_assignment,
+    ):
+        instance_manager.all_or_nothing_node_assignment = mocker.MagicMock()
+        instance_manager.best_effort_node_assignment = mocker.MagicMock()
+        instance_manager._add_instances_for_nodes(
+            node_list=[],
+            launch_batch_size=1,
+            assign_node_batch_size=2,
+            scaling_strategy=scaling_strategy,
+        )
+        if expect_all_or_nothing_node_assignment:
+            instance_manager.all_or_nothing_node_assignment.assert_called_once()
+        if expect_best_effort_node_assignment:
+            instance_manager.best_effort_node_assignment.assert_called_once()
+
     @pytest.mark.parametrize(
         "job, nodes_to_launch, launch_batch_size, unused_launched_instances, launched_instances, "
         "expected_instances_launched, expected_failed_nodes",