From 1ec2a36443d737f3ec0eee41fb2bc3edeffcc50d Mon Sep 17 00:00:00 2001 From: Eddy Mwiti Date: Fri, 6 Oct 2023 20:06:22 +0300 Subject: [PATCH] Add support for greedy-all-or-nothing Signed-off-by: Eddy Mwiti --- src/slurm_plugin/common.py | 1 + src/slurm_plugin/instance_manager.py | 2 +- tests/slurm_plugin/test_instance_manager.py | 29 +++++++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/slurm_plugin/common.py b/src/slurm_plugin/common.py index 0647052d2..fa8dc0a12 100644 --- a/src/slurm_plugin/common.py +++ b/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): diff --git a/src/slurm_plugin/instance_manager.py b/src/slurm_plugin/instance_manager.py index 80db55b58..4a251357d 100644 --- a/src/slurm_plugin/instance_manager.py +++ b/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, diff --git a/tests/slurm_plugin/test_instance_manager.py b/tests/slurm_plugin/test_instance_manager.py index 4f1bfd372..0e8a0e07f 100644 --- a/tests/slurm_plugin/test_instance_manager.py +++ b/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",