refactored

BalaBalaYi · BalaBalaYi · commit aa44f9cfe6b2 · 2024-10-10T11:48:07.000+08:00
diff --git a/dlrover/python/common/grpc.py b/dlrover/python/common/grpc.py
@@ -489,10 +489,9 @@ class NodeCheckpointState(Message):
 
 
 @dataclass
-class WorkerDiagnosisData(Message):
-    type: str = ""
-    timestamp: int = 0
-    content: str = ""
+class DiagnosisReportData(Message):
+    data_cls: str = ""
+    data_content: str = ""
     node_rank: int = -1
 
 
diff --git a/dlrover/python/diagnosis/common/diagnosis_data.py b/dlrover/python/diagnosis/common/diagnosis_data.py
@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from abc import ABCMeta
 from datetime import datetime
 from typing import List
@@ -20,18 +21,36 @@
 
 
 class DiagnosisData(metaclass=ABCMeta):
+    """
+    Basic definition of diagnosis data.
+
+    Args:
+        timestamp (datetime): Timestamp of diagnosis data.
+        data_type (str): Type of metric. Defaults to "GENERIC".
+        data_content (str): Content of the metric. Defaults to "".
+        node_id (int): Node ID. Defaults to -1.
+        node_type (str): Node type. Defaults to "".
+        node_rank (int): Node rank. Defaults to -1.
+    """
+
     def __init__(
         self,
         timestamp: int = 0,
         data_type: str = DiagnosisDataType.GENERIC,
         data_content: str = "",
+        node_id: int = -1,
+        node_type: str = "",
+        node_rank: int = -1,
     ):
         if timestamp == 0:
             self._timestamp = int(round(datetime.now().timestamp()))
         else:
             self._timestamp = timestamp
         self._data_type = data_type
         self._data_content = data_content
+        self._node_id = node_id
+        self._node_type = node_type
+        self._node_rank = node_rank
 
     @property
     def data_type(self) -> str:
@@ -45,33 +64,6 @@ def timestamp(self) -> int:
     def data_content(self) -> str:
         return self._data_content
 
-
-class WorkerDiagnosisData(DiagnosisData):
-    def __init__(
-        self,
-        timestamp: int = 0,
-        data_type: str = DiagnosisDataType.GENERIC,
-        data_content: str = "",
-        node_id: int = -1,
-        node_type: str = "",
-        node_rank: int = -1,
-    ):
-        """
-        General metric
-
-        Args:
-            data_type (str): Type of metric. Defaults to "GENERIC".
-            data_content (str): Content of the metric. Defaults to "".
-            node_id (int): Node ID. Defaults to -1.
-            node_type (str): Node type. Defaults to "".
-            node_rank (int): Node rank. Defaults to -1.
-        """
-
-        super().__init__(timestamp, data_type, data_content)
-        self._node_id = node_id
-        self._node_type = node_type
-        self._node_rank = node_rank
-
     @property
     def node_id(self):
         return self._node_id
@@ -84,35 +76,47 @@ def node_type(self):
     def node_rank(self):
         return self._node_rank
 
+    def to_json(self):
+        data = {k.lstrip("_"): v for k, v in self.__dict__.items()}
+        return json.dumps(data)
+
+    @classmethod
+    def from_json(cls, json_data):
+        return cls(**json.loads(json_data))
+
+    def is_from_worker(self):
+        return self._node_id != -1
+
+
+class WorkerTrainingMetric(DiagnosisData):
+    """
+    Diagnosis data for worker training metric.
+
+    Args:
+        timestamp (datetime): Timestamp of diagnosis data.
+        data_type (str): Type of metric. Defaults to "GENERIC".
+        data_content (str): Content of the metric. Defaults to "".
+        node_id (int): Node ID. Defaults to -1.
+        node_type (str): Node type. Defaults to "".
+        node_rank (int): Node rank. Defaults to -1.
+        is_final_result (bool, optional): Whether the metric is final result.
+            Defaults to False.
+        need_report (bool, optional): Whether the metric needs report.
+            Defaults to False.
+    """
 
-class WorkerTrainingMetric(WorkerDiagnosisData):
     def __init__(
         self,
         timestamp: int = 0,
         data_type: str = DiagnosisDataType.GENERIC,
         data_content: str = "",
-        node_id: int = -1,
-        node_type: str = "",
-        node_rank: int = -1,
+        node_id=env_utils.get_node_id(),
+        node_type=env_utils.get_node_type(),
+        node_rank=env_utils.get_node_rank(),
         is_final_result=False,
         need_report=False,
     ):
-        """
-        General metric
-
-        Args:
-            data_type (str): Type of metric. Defaults to "GENERIC".
-            data_content (str): Content of the metric. Defaults to "".
-            is_final_result (bool, optional): Whether the metric is final
-                result or not. Defaults to False.
-            need_report (bool, optional): Whether the metric needs
-                report(to Brain). Defaults to False.
-            node_id (int): Node ID. Defaults to -1.
-            node_type (str): Node type. Defaults to "".
-            node_rank (int): Node rank. Defaults to -1.
-        """
-
-        super().__init__(
+        super(WorkerTrainingMetric, self).__init__(
             timestamp, data_type, data_content, node_id, node_type, node_rank
         )
         self._is_final_result = is_final_result
@@ -133,8 +137,26 @@ def is_resolvable(self):
         return False
 
 
-class TrainingLog(WorkerDiagnosisData):
-    def __init__(self, timestamp: int = 0, logs: List[str] = None):
+class TrainingLog(DiagnosisData):
+    """
+    Worker's training log.
+
+    Args:
+        timestamp (datetime): Timestamp of diagnosis data.
+        logs (list): Log content in list format.
+        node_id (int): Node ID. Defaults to -1.
+        node_type (str): Node type. Defaults to "".
+        node_rank (int): Node rank. Defaults to -1.
+    """
+
+    def __init__(
+        self,
+        timestamp: int = 0,
+        logs: List[str] = None,
+        node_id=env_utils.get_node_id(),
+        node_type=env_utils.get_node_type(),
+        node_rank=env_utils.get_node_rank(),
+    ):
         if logs is None:
             data_content = ""
         else:
@@ -144,9 +166,9 @@ def __init__(self, timestamp: int = 0, logs: List[str] = None):
             timestamp,
             DiagnosisDataType.TRAINING_LOG,
             data_content,
-            env_utils.get_node_id(),
-            env_utils.get_node_type(),
-            env_utils.get_node_rank(),
+            node_id,
+            node_type,
+            node_rank,
         )
 
     @property
diff --git a/dlrover/python/elastic_agent/master_client.py b/dlrover/python/elastic_agent/master_client.py
@@ -23,7 +23,7 @@
 from dlrover.python.common.constants import NetworkFailureReason, NodeEnv
 from dlrover.python.common.log import default_logger as logger
 from dlrover.python.common.singleton import Singleton
-from dlrover.python.diagnosis.common.diagnosis_data import WorkerDiagnosisData
+from dlrover.python.diagnosis.common.diagnosis_data import DiagnosisData
 
 
 def retry_grpc_request(func):
@@ -382,11 +382,10 @@ def report_failures(self, error_data, restart_count=-1, level=""):
     def report_paral_config(self, config: grpc.ParallelConfig):
         self._report(config)
 
-    def report_diagnosis_agent_metrics(self, data: WorkerDiagnosisData):
-        message = grpc.WorkerDiagnosisData(
-            data.data_type,
-            data.timestamp,
-            data.data_content,
+    def report_diagnosis_agent_metrics(self, data: DiagnosisData):
+        message = grpc.DiagnosisReportData(
+            data.__class__.__name__,
+            data.to_json(),
             data.node_rank,
         )
         self._report(message)
diff --git a/dlrover/python/master/servicer.py b/dlrover/python/master/servicer.py
@@ -14,7 +14,7 @@
 import threading
 import time
 from concurrent import futures
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import grpc as grpc_lib
 
@@ -32,7 +32,7 @@
 )
 from dlrover.python.common.global_context import Context
 from dlrover.python.common.log import default_logger as logger
-from dlrover.python.diagnosis.common.diagnosis_data import WorkerTrainingMetric
+from dlrover.python.diagnosis.common.diagnosis_data import DiagnosisData
 from dlrover.python.master.diagnosis.diagnosis import DiagnosisManager
 from dlrover.python.master.elastic_training.kv_store_service import (
     KVStoreService,
@@ -354,10 +354,8 @@ def report(self, request, _):
             success = self._report_heartbeat(node_type, node_id, message)
         elif isinstance(message, grpc.NodeCheckpointState):
             success = self._sync_checkpoint(node_type, node_id, message)
-        elif isinstance(message, grpc.WorkerDiagnosisData):
-            success = self._report_worker_diagnosis_data(
-                node_type, node_id, message
-            )
+        elif isinstance(message, grpc.DiagnosisReportData):
+            success = self._report_worker_diagnosis_data(message)
 
         response.success = success
         return response
@@ -613,19 +611,17 @@ def _sync_checkpoint(
         rdzv_manager = self._rdzv_managers[RendezvousName.ELASTIC_TRAINING]
         return rdzv_manager.sync_ckpt_nodes(node_id, message.step)
 
-    def _report_worker_diagnosis_data(
-        self, node_type, node_id, message: grpc.WorkerDiagnosisData
-    ):
+    def _report_worker_diagnosis_data(self, message: grpc.DiagnosisReportData):
         if self._diagnosis_manager:
-            data = WorkerTrainingMetric(
-                timestamp=message.timestamp,
-                data_type=message.type,
-                data_content=message.content,
-                node_id=node_id,
-                node_type=node_type,
-                node_rank=message.node_rank,
-            )
-            self._diagnosis_manager.collect_diagnosis_data(data)
+            data_cls: Optional[DiagnosisData] = globals().get(message.data_cls)
+            if data_cls is None:
+                logger.warning(
+                    "Invalid diagnosis report "
+                    f"data type: {message.data_cls}"
+                )
+                return False
+            data_obj = data_cls.from_json(message.data_content)
+            self._diagnosis_manager.collect_diagnosis_data(data_obj)
         return True
 
     def _sync_training_ports(
diff --git a/dlrover/python/tests/test_diagnosis_agent.py b/dlrover/python/tests/test_diagnosis_agent.py
@@ -166,6 +166,36 @@ def test_xpu_timer_metric_collect(self):
         self.assertEqual(agent_xpu_metric.node_rank, 1)
         self.assertTrue(agent_xpu_metric.timestamp > 0)
 
+    def test_worker_training_metric(self):
+        test = WorkerTrainingMetric(
+            data_content="test123",
+            node_id=env_utils.get_node_id(),
+            node_type=env_utils.get_node_type(),
+            node_rank=env_utils.get_node_rank(),
+            is_final_result=True,
+        )
+
+        test_str = test.to_json()
+        self.assertTrue('"data_content": "test123"' in test_str)
+
+        test_new = WorkerTrainingMetric.from_json(test_str)
+        self.assertEqual(test_new.timestamp, test.timestamp)
+        self.assertEqual(test_new.data_content, test.data_content)
+        self.assertEqual(test_new.data_type, test.data_type)
+        self.assertEqual(test_new.is_final_result, test.is_final_result)
+
+        test_new = globals().get("WorkerTrainingMetric").from_json(test_str)
+        self.assertEqual(test_new.timestamp, test.timestamp)
+        self.assertEqual(test_new.data_content, test.data_content)
+        self.assertEqual(test_new.data_type, test.data_type)
+        self.assertEqual(test_new.is_final_result, test.is_final_result)
+
+        test_new = globals().get(test.__class__.__name__).from_json(test_str)
+        self.assertEqual(test_new.timestamp, test.timestamp)
+        self.assertEqual(test_new.data_content, test.data_content)
+        self.assertEqual(test_new.data_type, test.data_type)
+        self.assertEqual(test_new.is_final_result, test.is_final_result)
+
 
 if __name__ == "__main__":
     unittest.main()
