pybind/mgr/balancer/module.py: assign weight-sets to all buckets before balancing

Add an additional check to make sure that the choose_args section has the same
number of buckets as the crushmap. If not, ensure that
get_compat_weight_set_weights assigns weight-sets to all buckets.

Without this change, if we end up with an orig_ws, which has fewer buckets
than the crushmap, the mgr will crash due a KeyError in do_crush_compat().

Fixes: https://tracker.ceph.com/issues/49576
Signed-off-by: Neha Ojha <nojha@redhat.com>

diff --git a/src/pybind/mgr/balancer/module.py b/src/pybind/mgr/balancer/module.py
index ea857bac66ffc563d667864c6ab8ab30649f422a..d313bc4de848dc801a79ac9582a4dd2276210926 100644 (file)
--- a/src/pybind/mgr/balancer/module.py
+++ b/src/pybind/mgr/balancer/module.py
@@ -1220,8 +1220,13 @@ class Module(MgrModule):
                                 'change balancer mode and retry might help'
 
     def get_compat_weight_set_weights(self, ms: MappingState):
-        if not CRUSHMap.have_default_choose_args(ms.crush_dump):
+        have_choose_args = CRUSHMap.have_default_choose_args(ms.crush_dump)
+        if have_choose_args:
+            # get number of buckets in choose_args
+            choose_args_len = len(CRUSHMap.get_default_choose_args(ms.crush_dump))
+        if not have_choose_args or choose_args_len != len(ms.crush_dump['buckets']):
             # enable compat weight-set first
+            self.log.debug('no choose_args or all buckets do not have weight-sets')
             self.log.debug('ceph osd crush weight-set create-compat')
             result = CommandResult('')
             self.send_command(result, 'mon', '', json.dumps({