feat: seperate ray logs and service logs

Author: ChenZiHong-Gavin

graphgen/bases/__init__.py

@@ -2,15 +2,11 @@
 from .base_generator import BaseGenerator
 from .base_kg_builder import BaseKGBuilder
 from .base_llm_wrapper import BaseLLMWrapper
+from .base_operator import BaseOperator
 from .base_partitioner import BasePartitioner
 from .base_reader import BaseReader
 from .base_searcher import BaseSearcher
 from .base_splitter import BaseSplitter
-from .base_storage import (
-    BaseGraphStorage,
-    BaseKVStorage,
-    BaseListStorage,
-    StorageNameSpace,
-)
+from .base_storage import BaseGraphStorage, BaseKVStorage, StorageNameSpace
 from .base_tokenizer import BaseTokenizer
 from .datatypes import Chunk, Config, Node, QAPair, Token

graphgen/bases/base_operator.py

@@ -0,0 +1,57 @@
+import inspect
+import os
+from abc import ABC, abstractmethod
+from typing import Iterable, Union
+
+import pandas as pd
+import ray
+
+from graphgen.utils import CURRENT_LOGGER_VAR, set_logger
+
+
+class BaseOperator(ABC):
+    def __init__(self, working_dir: str = "cache", op_name: str = None):
+        log_dir = os.path.join(working_dir, "logs")
+        self.op_name = op_name or self.__class__.__name__
+
+        try:
+            ctx = ray.get_runtime_context()
+            worker_id = ctx.get_actor_id() or ctx.get_worker_id()
+            worker_id_short = worker_id[-6:] if worker_id else "driver"
+        except Exception as e:
+            print(
+                "Warning: Could not get Ray worker ID, defaulting to 'local'. Exception:",
+                e,
+            )
+            worker_id_short = "local"
+
+        # e.g. cache/logs/ChunkService_a1b2c3.log
+        log_file = os.path.join(log_dir, f"{self.op_name}_{worker_id_short}.log")
+
+        self.logger = set_logger(
+            log_file=log_file, name=f"{self.op_name}.{worker_id_short}", force=True
+        )
+
+        self.logger.info(
+            "[%s] Operator initialized on Worker %s", self.op_name, worker_id_short
+        )
+
+    def __call__(
+        self, batch: pd.DataFrame
+    ) -> Union[pd.DataFrame, Iterable[pd.DataFrame]]:
+        logger_token = CURRENT_LOGGER_VAR.set(self.logger)
+        try:
+            result = self.process(batch)
+            if inspect.isgenerator(result):
+                yield from result
+            else:
+                yield result
+        finally:
+            CURRENT_LOGGER_VAR.reset(logger_token)
+
+    @abstractmethod
+    def process(self, batch):
+        raise NotImplementedError("Subclasses must implement the process method.")
+
+    def get_logger(self):
+        return self.logger

graphgen/bases/base_storage.py

@@ -41,6 +41,9 @@ def upsert(self, data: dict[str, T]):
     def drop(self):
         raise NotImplementedError
 
+    def reload(self):
+        raise NotImplementedError
+
 
 class BaseGraphStorage(StorageNameSpace):
     def has_node(self, node_id: str) -> bool:
@@ -88,3 +91,6 @@ def upsert_edge(
 
     def delete_node(self, node_id: str):
         raise NotImplementedError
+
+    def reload(self):
+        raise NotImplementedError