fix ut

intelligent-machine-learning · Jan 22, 2025 · 22c0da2 · 22c0da2
1 parent bd904a7
commit 22c0da2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/dlrover/python/master/node/event_callback.py b/dlrover/python/master/node/event_callback.py
@@ -272,7 +272,7 @@ def on_node_succeeded(self, node: Node, cluster_context: ClusterContext):
     @NodeEventCallback.log_callback_exception
     def on_node_failed(self, node: Node, cluster_context):
         node.finish_time = datetime.now()  # type: ignore
-        self._job_context.report_failed_node()
+        self._job_context.report_failed_node(node.id)
         self._stop_job_if_needed(node)
         if node.is_unrecoverable_failure():
             self._master.speed_monitor.reduce_target_worker_num(

diff --git a/dlrover/python/tests/test_event_callback.py b/dlrover/python/tests/test_event_callback.py
@@ -72,7 +72,7 @@ def test_on_node_failed(self):
         worker.relaunch_count = 1
         worker.exit_reason = NodeExitReason.FATAL_ERROR
         self.event_cb.on_node_failed(worker, None)
-        self.assertEqual(self.event_cb._failed_worker_count, 1)
+        self.assertEqual(self.event_cb._job_context.get_failed_node_cnt(), 1)
         self.assertTrue(self.master._stop_requested)
         self.master._stop_requested = False
         _dlrover_ctx.relaunch_always = True