fix: update webui

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -289,7 +289,9 @@ def func_wrapper(row_or_batch: Dict[str, Any]) -> Dict[str, Any]:
                     f"Unsupported node type {node.type} for node {node.id}"
                 )
 
-    def execute(self, initial_ds: ray.data.Dataset, output_dir: str):
+    def execute(
+        self, initial_ds: ray.data.Dataset, output_dir: str
+    ) -> Dict[str, ray.data.Dataset]:
         sorted_nodes = self._topo_sort(self.config.nodes)
 
         for node in sorted_nodes:
@@ -314,3 +316,5 @@ def execute(self, initial_ds: ray.data.Dataset, output_dir: str):
 
                 # ray will lazy read the dataset
                 self.datasets[node.id] = ray.data.read_json(node_output_path)
+
+        return self.datasets

webui/app.py

@@ -1,7 +1,7 @@
+import gc
 import json
 import os
 import sys
-import gc
 import tempfile
 from importlib.resources import files
 
@@ -188,7 +188,7 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
         ds = ray.data.from_items([])
 
         # Execute pipeline
-        results = engine.execute(ds)
+        results = engine.execute(ds, output_dir=working_dir)
 
         # 5. Process Output
         # Extract the result from the 'generate' node