File tree Expand file tree Collapse file tree 3 files changed +12
-8
lines changed
Expand file tree Collapse file tree 3 files changed +12
-8
lines changed 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