feat: pass env_vars in engine

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -1,8 +1,10 @@
+import os
 import inspect
 import logging
 from collections import defaultdict, deque
 from functools import wraps
 from typing import Any, Callable, Dict, List, Set
+from dotenv import load_dotenv
 
 import ray
 import ray.data
@@ -12,6 +14,7 @@
 from graphgen.utils import logger
 from graphgen.common import init_llm
 
+load_dotenv()
 
 class Engine:
     def __init__(
@@ -31,6 +34,16 @@ def __init__(
         ctx.enable_tensor_extension_casting = False
         ctx._metrics_export_port = 0  # Disable metrics exporter to avoid RpcError
 
+        all_env_vars = os.environ.copy()
+        if "runtime_env" not in ray_init_kwargs:
+            ray_init_kwargs["runtime_env"] = {}
+
+        existing_env_vars = ray_init_kwargs["runtime_env"].get("env_vars", {})
+        ray_init_kwargs["runtime_env"]["env_vars"] = {
+            **all_env_vars,
+            **existing_env_vars
+        }
+
         if not ray.is_initialized():
             context = ray.init(
                 ignore_reinit_error=True,

graphgen/run.py

@@ -6,7 +6,6 @@
 
 import ray
 import yaml
-from dotenv import load_dotenv
 from ray.data.block import Block
 from ray.data.datasource.filename_provider import FilenameProvider
 
@@ -16,8 +15,6 @@
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
 
-load_dotenv()
-
 
 def set_working_dir(folder):
     os.makedirs(folder, exist_ok=True)