feat: write results in output folder

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -198,7 +198,7 @@ def _find_leaf_nodes(nodes: List[Node]) -> Set[str]:
             deps_set.update(n.dependencies)
         return all_ids - deps_set
 
-    def execute(self, initial_ds: ray.data.Dataset) -> Dict[str, List[Any]]:
+    def execute(self, initial_ds: ray.data.Dataset) -> Dict[str, ray.data.Dataset]:
         sorted_nodes = self._topo_sort(self.config.nodes)
 
         for node in sorted_nodes:
@@ -210,7 +210,4 @@ def execute(self, initial_ds: ray.data.Dataset) -> Dict[str, List[Any]]:
         def _fetch_result(ds: ray.data.Dataset) -> List[Any]:
             return ds.take_all()
 
-        results = ray.get(
-            [_fetch_result.remote(self.datasets[node_id]) for node_id in leaf_nodes]
-        )
-        return dict(zip(leaf_nodes, results))
+        return {node_id: self.datasets[node_id] for node_id in leaf_nodes}

graphgen/run.py

@@ -1,13 +1,17 @@
 import argparse
 import os
 import time
-from importlib.resources import files
+from importlib import resources
+from typing import Any, Dict
 
+import ray
 import yaml
 from dotenv import load_dotenv
+from ray.data.block import Block
+from ray.data.datasource.filename_provider import FilenameProvider
 
-from graphgen.engine import Context, Engine, collect_ops
-from graphgen.graphgen import GraphGen
+from graphgen.engine import Engine
+from graphgen.operators import operators
 from graphgen.utils import logger, set_logger
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
@@ -28,12 +32,38 @@ def save_config(config_path, global_config):
         )
 
 
+class NodeFilenameProvider(FilenameProvider):
+    def __init__(self, node_id: str):
+        self.node_id = node_id
+
+    def get_filename_for_block(
+        self, block: Block, write_uuid: str, task_index: int, block_index: int
+    ) -> str:
+        # format: {node_id}_{write_uuid}_{task_index:06}_{block_index:06}.json
+        return f"{self.node_id}_{write_uuid}_{task_index:06d}_{block_index:06d}.jsonl"
+
+    def get_filename_for_row(
+        self,
+        row: Dict[str, Any],
+        write_uuid: str,
+        task_index: int,
+        block_index: int,
+        row_index: int,
+    ) -> str:
+        raise NotImplementedError(
+            f"Row-based filenames are not supported by write_json. "
+            f"Node: {self.node_id}, write_uuid: {write_uuid}"
+        )
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--config_file",
         help="Config parameters for GraphGen.",
-        default=files("graphgen").joinpath("configs", "aggregated_config.yaml"),
+        default=resources.files("graphgen")
+        .joinpath("configs")
+        .joinpath("aggregated_config.yaml"),
         type=str,
     )
     parser.add_argument(
@@ -51,6 +81,8 @@ def main():
     with open(args.config_file, "r", encoding="utf-8") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
 
+    engine = Engine(config, operators)
+
     unique_id = int(time.time())
 
     output_path = os.path.join(working_dir, "data", "graphgen", f"{unique_id}")
@@ -65,15 +97,22 @@ def main():
         unique_id,
         os.path.join(working_dir, f"{unique_id}.log"),
     )
-
-    graph_gen = GraphGen(unique_id=unique_id, working_dir=working_dir)
-
-    # share context between different steps
-    ctx = Context(config=config, graph_gen=graph_gen)
-    ops = collect_ops(config, graph_gen)
-
-    # run operations
-    Engine(max_workers=config.get("max_workers", 4)).run(ops, ctx)
+    ds = ray.data.from_items([])
+    results = engine.execute(ds)
+
+    for node_id, dataset in results.items():
+        output_path = os.path.join(output_path, f"{node_id}")
+        os.makedirs(output_path, exist_ok=True)
+        dataset.write_json(
+            output_path,
+            filename_provider=NodeFilenameProvider(node_id),
+            pandas_json_args_fn=lambda: {
+                "force_ascii": False,
+                "orient": "records",
+                "lines": True,
+            },
+        )
+        logger.info("Node %s results saved to %s", node_id, output_path)
 
     save_config(os.path.join(output_path, "config.yaml"), config)
     logger.info("GraphGen completed successfully. Data saved to %s", output_path)