GCP Cloud Run中Ray客户端重复连接报错及连接复用方案

问题背景 #

在GCP Cloud Run容器中运行的Flask应用，负责接收Firestore事件并处理任务。每次请求执行ray.init()时都会触发错误：ray.init() called, but ray client is already connected，需要实现Ray连接的复用而非每次请求重新初始化。

核心原因 #

1. 当前代码在每次请求中先强制断开Ray连接，再尝试重新初始化，导致多线程环境下（Gunicorn配置threads=8）连接状态混乱
2. Ray客户端并非线程安全，多线程共享全局连接会引发冲突
3. finally块中每次请求后断开连接，完全破坏了连接复用的可能

解决方案 #

修改思路 #

• 使用线程本地存储为每个线程维护独立的Ray连接，避免多线程冲突
• 检查当前线程的Ray连接状态，仅在未连接时执行初始化
• 移除每次请求后的强制断开操作，让连接随线程复用保留
• 优化连接初始化逻辑，避免无意义的断开尝试

修改后的完整代码 #

import os
import sys
import logging
import json
import time
import ray
import threading
import google.protobuf.json_format
from flask import Flask, request, jsonify
from google.events.cloud.firestore import DocumentEventData

# 配置日志
logging.basicConfig(
    stream=sys.stdout,
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 线程本地存储，为每个线程保存独立的Ray连接状态
thread_local = threading.local()

def get_or_create_ray_connection():
    """获取或创建当前线程的Ray连接"""
    try:
        # 检查当前线程是否已连接Ray
        if hasattr(thread_local, 'ray_connected') and thread_local.ray_connected:
            logger.info("✅ 复用当前线程已有的Ray连接")
            return True

        # 未连接时执行初始化
        creds_path = "./ray-gcs-sa.json"
        if not os.path.exists(creds_path):
            logger.error(f"❌ 未找到服务账号密钥: {creds_path}")
            return False
            
        with open(creds_path, "r") as f:
            credentials_content = f.read()
        
        logger.info(f"🔄 当前线程初始化Ray连接: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        worker_id = f"worker_{os.getpid()}_{threading.get_ident()}_{time.time_ns()}"
        
        ray.init(
            address="ray://xxxxxxxxxxxxxxxx",
            namespace=worker_id,
            runtime_env={
                "pip": ["google-cloud-storage>=2.0.0"],
                "env_vars": {
                    "GOOGLE_APPLICATION_CREDENTIALS_CONTENT": credentials_content
                }
            },
            logging_level=logging.INFO
        )
        
        # 测试连接并标记当前线程已连接
        @ray.remote
        def ping():
            return "Connection successful"
        
        result = ray.get(ping.remote(), timeout=5)
        logger.info(f"✨ Ray连接验证成功: {result}")
        thread_local.ray_connected = True
        return True

    except ray.exceptions.RayClientAlreadyConnectedError:
        # 捕获已连接错误，标记状态并返回成功
        logger.info("✅ Ray客户端已连接，复用现有连接")
        thread_local.ray_connected = True
        return True
    except Exception as e:
        logger.error(f"❌ Ray连接错误: {str(e)}")
        thread_local.ray_connected = False
        return False

def extract_user_query(doc_json):
    """从Firestore文档中提取用户查询"""
    try:
        doc_dict = json.loads(doc_json)
        value_dict = doc_dict.get('value', {})
        fields = value_dict.get('fields', {})
        user_query = fields.get('user_query', {}).get('stringValue', '')
        logger.info(f"🔍 提取到用户查询: {user_query}")
        return user_query
    except Exception as e:
        logger.error(f"❌ 提取用户查询失败: {e}", exc_info=True)
        raise

app = Flask(__name__)

@app.route("/", methods=["GET"])
def home():
    """健康检查端点"""
    logger.info("🏥 健康检查端点被调用")
    return "Firestore Listener is running!", 200

@app.route("/trigger", methods=["POST"])
def trigger():
    """接收Firestore事件并处理"""
    try:
        logger.info("🔹 收到/trigger请求")
        raw_body = request.get_data()
        logger.info(f"📜 Protobuf原始数据长度: {len(raw_body)} bytes")
        
        if not raw_body:
            logger.warning("⚠️ 请求体为空")
            return jsonify({"status": "error", "message": "Empty request body"}), 400

        # 解析Protobuf消息
        doc_event = DocumentEventData()
        doc_event._pb.ParseFromString(raw_body)
        doc_json = google.protobuf.json_format.MessageToJson(doc_event._pb)
        logger.info(f"✅ 解析DocumentEventData为JSON: {doc_json}")

        # 提取用户查询
        new_user_query = extract_user_query(doc_json)
        logger.info(f"📝 新用户查询: {new_user_query}")

        # 获取或创建Ray连接
        logger.info("🔄 检查Ray连接状态...")
        connection_active = get_or_create_ray_connection()
        
        response_data = {
            "status": "success",
            "message": "事件已接收并处理",
            "user_query": new_user_query,
            "ray_connected": connection_active
        }

        logger.info(f"✨ 请求处理完成，Ray连接状态: {connection_active}")
        return jsonify(response_data), 200

    except Exception as e:
        error_msg = f"❌ 请求处理失败: {str(e)}"
        logger.error(error_msg, exc_info=True)
        return jsonify({
            "status": "error",
            "message": error_msg
        }), 500
    # 移除finally中的强制断开操作，保留连接供线程复用

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8081))
    logger.info(f"🚀 Flask应用启动，端口: {port}")
    app.run(host="0.0.0.0", port=port)

关键修改点说明 #

1. 线程本地存储：新增thread_local对象，为每个线程独立保存ray_connected状态，避免多线程间的连接状态干扰
2. 连接状态检查：在get_or_create_ray_connection中优先检查当前线程是否已连接，仅在未连接时执行初始化
3. 异常捕获：专门捕获RayClientAlreadyConnectedError，处理意外的已连接情况并标记状态
4. 移除强制断开：删除finally块中的ray.disconnect()，让连接随线程复用保留，避免频繁创建销毁连接
5. 命名空间优化：在worker_id中加入线程ID，确保每个线程的Ray命名空间唯一

额外注意事项 #

• 保持Gunicorn的workers=1配置：如果使用多worker进程，每个进程会独立维护自己的线程连接池，无需额外修改
• Ray服务端配置：确保Ray集群允许足够的并发连接数，适配Cloud Run的请求并发量
• 连接泄漏处理：如果应用长时间运行，可考虑定期检查线程状态并清理闲置连接（可选）

内容的提问来源于stack exchange，提问作者Juan Lozano