File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 77
88import ray
99import ray .data
10- from dotenv import load_dotenv
1110from ray .data import DataContext
1211
1312from graphgen .bases import Config , Node
1413from graphgen .common import init_llm , init_storage
1514from graphgen .utils import logger
1615
17- load_dotenv ()
18-
1916
2017class Engine :
2118 def __init__ (
Original file line number Diff line number Diff line change 66
77import ray
88import yaml
9+ from dotenv import load_dotenv
910from ray .data .block import Block
1011from ray .data .datasource .filename_provider import FilenameProvider
1112
1213from graphgen .engine import Engine
1314from graphgen .operators import operators
1415from graphgen .utils import CURRENT_LOGGER_VAR , logger , set_logger
1516
17+ load_dotenv ()
18+
1619sys_path = os .path .abspath (os .path .dirname (__file__ ))
1720
1821
Original file line number Diff line number Diff line change @@ -160,6 +160,7 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
160160 "op_name" : "generate" ,
161161 "type" : "map_batch" ,
162162 "dependencies" : ["partition" ],
163+ "save_output" : True ,
163164 "execution_params" : {"replicas" : 1 , "batch_size" : 128 },
164165 "params" : {
165166 "method" : params .mode ,
@@ -168,14 +169,17 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
168169 }
169170 )
170171
171- config = {"global_params" : {"working_dir" : working_dir }, "nodes" : nodes }
172+ config = {
173+ "global_params" : {
174+ "working_dir" : working_dir ,
175+ "graph_backend" : "kuzu" ,
176+ "kv_backend" : "rocksdb" ,
177+ },
178+ "nodes" : nodes ,
179+ }
172180
173181 try :
174182 # 4. Initialize and Run Engine
175- # Initialize Ray if not already running (Engine handles this mostly, but good for safety)
176- if not ray .is_initialized ():
177- ray .init (ignore_reinit_error = True , log_to_driver = True )
178-
179183 engine = Engine (config , operators )
180184
181185 # Start with an empty dataset to kick off the pipeline
You can’t perform that action at this time.
0 commit comments