File size: 6,031 Bytes
66f749a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
Ray cluster initialization and utilities
"""
import logging
import os
import ray
logger = logging.getLogger(__name__)
def _clean_env(name: str, default: str = "") -> str:
"""Read env var and trim surrounding whitespace/newlines."""
return (os.getenv(name, default) or default).strip()
def init_ray_cluster(
num_cpus: int = None,
num_gpus: int = 0,
address: str = None
):
"""
Initialize Ray cluster
Args:
num_cpus: Number of CPUs to use (None = auto-detect)
num_gpus: Number of GPUs to use
address: Existing cluster address to connect to
Returns:
Ray context info
"""
if ray.is_initialized():
logger.info("Ray already initialized, shutting down first")
ray.shutdown()
# Get API key(s) to pass to workers via runtime_env
api_key = _clean_env("GEMINI_API_KEY", "")
api_keys = _clean_env("GEMINI_API_KEYS", "")
if api_keys:
key_count = len([k for k in api_keys.split(",") if k.strip()])
logger.info(f"GEMINI_API_KEYS found ({key_count} keys for rotation)")
elif api_key:
logger.info(f"GEMINI_API_KEY found ({len(api_key)} chars, starts with {api_key[:8]}...)")
else:
logger.error("GEMINI_API_KEY is NOT set! LLM calls will fail.")
qwen_api_url = _clean_env("QWEN_API_URL", "")
if not qwen_api_url:
logger.warning("QWEN_API_URL is not set, using default HuggingFace endpoint")
qwen_api_url = "https://vish85521-qwen.hf.space/api/generate"
qwen_model_name = _clean_env("QWEN_MODEL_NAME", "qwen3.5:397b-cloud")
runtime_env = {
"env_vars": {
"GEMINI_API_KEY": api_key,
"GEMINI_API_KEYS": api_keys,
"QWEN_API_URL": qwen_api_url,
"QWEN_MODEL_NAME": qwen_model_name,
}
}
if address:
# Connect to existing cluster
logger.info(f"Connecting to Ray cluster at {address}")
context = ray.init(address=address, runtime_env=runtime_env)
else:
# Start local cluster
logger.info("Initializing local Ray cluster")
# Workaround: Ray on Windows breaks if project path contains spaces.
# Ray's runtime_env/context.py replaces spaces with "\ " in worker
# script paths, turning e.g. "CGP 1" into "CGP\ 1" which is invalid
# on Windows. Fix: use a temp dir without spaces for Ray's internal
# files AND copy the worker scripts there so the path has no spaces.
import tempfile
import shutil
ray_temp = os.path.join(tempfile.gettempdir(), "ray_agentsociety")
os.makedirs(ray_temp, exist_ok=True)
logger.info(f"Using Ray temp dir: {ray_temp}")
# Copy default_worker.py to space-free path if needed
extra_init_kwargs = {}
ray_private_dir = os.path.join(
os.path.dirname(ray.__file__), "_private"
)
worker_src = os.path.join(
ray_private_dir, "workers", "default_worker.py"
)
if " " in worker_src and os.path.exists(worker_src):
workers_dest = os.path.join(ray_temp, "workers")
os.makedirs(workers_dest, exist_ok=True)
worker_dst = os.path.join(workers_dest, "default_worker.py")
shutil.copy2(worker_src, worker_dst)
# Also copy setup_worker.py if it exists
setup_src = os.path.join(
ray_private_dir, "workers", "setup_worker.py"
)
if os.path.exists(setup_src):
shutil.copy2(
setup_src,
os.path.join(workers_dest, "setup_worker.py"),
)
logger.info(
f"Copied Ray workers to space-free path: {workers_dest}"
)
# Tell Ray to use the space-free worker path
extra_init_kwargs["_system_config"] = {
"worker_register_timeout_seconds": 120,
}
# Monkey-patch the default worker path before ray.init()
import ray._private.parameter as ray_parameter
original_update = ray_parameter.RayParams.update_if_absent
_patched_worker_path = worker_dst
_patched_setup_path = os.path.join(
workers_dest, "setup_worker.py"
)
def patched_update_if_absent(self, **kwargs):
if "worker_path" in kwargs:
kwargs["worker_path"] = _patched_worker_path
if "setup_worker_path" in kwargs and os.path.exists(
_patched_setup_path
):
kwargs["setup_worker_path"] = _patched_setup_path
return original_update(self, **kwargs)
ray_parameter.RayParams.update_if_absent = patched_update_if_absent
logger.info("Patched Ray worker path to avoid spaces-in-path bug")
context = ray.init(
num_cpus=num_cpus,
num_gpus=num_gpus,
include_dashboard=False,
logging_level=logging.WARNING,
ignore_reinit_error=True,
log_to_driver=True,
configure_logging=False,
runtime_env=runtime_env,
_temp_dir=ray_temp,
**extra_init_kwargs,
)
resources = ray.available_resources()
logger.info(f"Ray initialized with resources: {resources}")
return context
def shutdown_ray():
"""Shutdown Ray cluster"""
if ray.is_initialized():
ray.shutdown()
logger.info("Ray cluster shut down")
def get_cluster_info() -> dict:
"""Get current cluster information"""
if not ray.is_initialized():
return {"status": "not_initialized"}
return {
"status": "running",
"resources": ray.available_resources(),
"nodes": len(ray.nodes())
}
|