Fix rank iter implement. (#1356)

* fix rank iter * fix rank iter
intelligent-machine-learning · Nov 22, 2024 · c639de1 · c639de1
1 parent 1efc2f5
commit c639de1
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/dlrover/python/master/node/training_node.py b/dlrover/python/master/node/training_node.py
@@ -238,20 +238,16 @@ def _get_nodes(self):
     def _update_node(self, node: Node):
         self._job_context.update_job_node(node)
 
-    def update_nodes_iter(self):
+    def update_nodes_iter(self, update_rank_iter=True):
         nodes = self._job_context.job_nodes_by_type(self._node_type)
-
-        # update everytime
         self._node_id_iter = itertools.count(
             max(nodes.keys()) + 1 if len(nodes) > 0 else 0
         )
-
-        # update once
-        if not self._node_rank_iter:
+        if update_rank_iter:
             self._node_rank_iter = itertools.count(len(nodes))
 
     def get_next_node_id(self):
-        self.update_nodes_iter()
+        self.update_nodes_iter(update_rank_iter=False)
         return next(self._node_id_iter)
 
     def remove_node(self, node_id):