From f2a09b690c39c850b8f66bb06513a73ef07bece0 Mon Sep 17 00:00:00 2001
From: lijinghan <lijinghan1@huawei.com>
Date: Thu, 5 Jun 2025 22:36:19 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90taskd=E3=80=91network=E5=88=9D?=
 =?UTF-8?q?=E5=A7=8B=E5=8C=96=E6=88=90=E5=8A=9F=E5=90=8E=E5=8F=91=E9=80=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../agent/base_agent/agent_network.py         | 32 +++++++++++++++++++
 .../framework/agent/base_agent/base_agent.py  |  5 +--
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/component/taskd/taskd/python/framework/agent/base_agent/agent_network.py b/component/taskd/taskd/python/framework/agent/base_agent/agent_network.py
index 6accd3bb4..e565b5519 100644
--- a/component/taskd/taskd/python/framework/agent/base_agent/agent_network.py
+++ b/component/taskd/taskd/python/framework/agent/base_agent/agent_network.py
@@ -121,6 +121,12 @@ class AgentMessageManager():
                 self.lib.DestroyNetwork(self._network_instance)
                 return
 
+    def get_network_instance(self):
+        """
+        Get network instance.
+        """
+        return self._network_instance
+
     def _parse_msg(self, msg_json) -> MsgBody:
         """
         Parse message from taskd manager.
@@ -170,6 +176,19 @@ def init_message_manager(network_config, msg_queue):
         run_log.error("network_config is None!")
         raise Exception("network_config is None!")
     msg_manager = AgentMessageManager(network_config, msg_queue)
+
+    time_use = 0
+    while True:
+        if time_use > 60:
+            run_log.error("init message manager failed!")
+            return
+        if msg_manager.get_network_instance() is not None:
+            run_log.info("init message manager success!")
+            break
+        time.sleep(1)
+        time_use += 1
+        run_log.info("wait get_network_instance")
+
     msg_manager.register(network_config.pos.server_rank)
     msg_manager.receive_message()
 
@@ -185,4 +204,17 @@ def network_send_message(msg :MessageInfo):
     Send message to taskd manager.
     """
     msg_manager = get_message_manager()
+    if msg_manager.get_network_instance() is None:
+        run_log.warning("network instance is None!")
+        return
     msg_manager.send_message(msg)
+
+
+def get_msg_network_instance():
+    """
+    Get network instance.
+    """
+    msg_manager = get_message_manager()
+    if msg_manager is None:
+        return None
+    return msg_manager.get_network_instance()
diff --git a/component/taskd/taskd/python/framework/agent/base_agent/base_agent.py b/component/taskd/taskd/python/framework/agent/base_agent/base_agent.py
index 4ec66b92e..f024db661 100644
--- a/component/taskd/taskd/python/framework/agent/base_agent/base_agent.py
+++ b/component/taskd/taskd/python/framework/agent/base_agent/base_agent.py
@@ -22,7 +22,8 @@ import uuid
 from dataclasses import asdict
 from taskd.python.utils.log import run_log
 from taskd.python.toolkit.fault_checker.fault_check import grace_exit_pids, stop_pids
-from taskd.python.framework.agent.base_agent.agent_network import get_message_manager, network_send_message
+from taskd.python.framework.agent.base_agent.agent_network import get_message_manager, network_send_message,\
+    get_msg_network_instance
 from taskd.python.framework.common.type import MsgBody, MessageInfo
 
 
@@ -115,7 +116,7 @@ class BaseAgent:
             if time_cost > 60:
                 run_log.error('waiting for message manager timeout')
                 raise ValueError("failed to initialized agent network, initialization message_manager timeout")
-            if get_message_manager() is None:
+            if get_msg_network_instance() is None:
                 run_log.info('waiting for message manager')
                 time.sleep(1)
                 time_cost += 1
-- 
Gitee