94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
|
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
from fastapi import FastAPI, HTTPException
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
from typing import Dict, Any
|
|||
|
|
import uvicorn
|
|||
|
|
|
|||
|
|
# 确保 pytrace 在 Python 路径中
|
|||
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|||
|
|
|
|||
|
|
from pytrace.model.graph import WorkflowGraph
|
|||
|
|
from pytrace.runtime.context import TraceContext
|
|||
|
|
from pytrace.runtime.executor import WorkflowExecutor
|
|||
|
|
from pytrace.runtime.repository import NodeRepository
|
|||
|
|
# 导入项目中已有的 TraceServer
|
|||
|
|
from app.services import trace_server
|
|||
|
|
# 假定数学节点套件已在某处定义并可以导入，以便在图中本地使用
|
|||
|
|
from pytrace.nodes.test_math_suite import MathSuite
|
|||
|
|
|
|||
|
|
|
|||
|
|
# --- 1. 数据模型 ---
|
|||
|
|
|
|||
|
|
class GraphExecutionRequest(BaseModel):
|
|||
|
|
graph_data: Dict[str, Any]
|
|||
|
|
initial_inputs: Dict[str, Any] = Field(default_factory=dict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# --- 2. FastAPI 应用定义 ---
|
|||
|
|
|
|||
|
|
app = FastAPI(
|
|||
|
|
title="TraceStudio - Remote Test Server",
|
|||
|
|
description="一个用于测试远程节点注册和执行的服务器。它同时开启HTTP和TCP服务。",
|
|||
|
|
version="2.0.0",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@app.on_event("startup")
|
|||
|
|
async def startup_event():
|
|||
|
|
# 1. 在后台启动TCP服务器，用于监听客户端的连接和注册
|
|||
|
|
trace_server.start_background()
|
|||
|
|
|
|||
|
|
# 2. 将TCP服务器的远程节点注册表配置给全局的NodeRepository
|
|||
|
|
NodeRepository.set_remote_registry(trace_server.remote_registry)
|
|||
|
|
|
|||
|
|
print("--- Remote Test Server ---")
|
|||
|
|
print(f"TCP RPC Server started on port {trace_server.port}. Waiting for clients to connect and register.")
|
|||
|
|
print("HTTP Server ready to execute graphs at /execute_graph")
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.post("/execute_graph", tags=["Execution"])
|
|||
|
|
async def execute_graph(request: GraphExecutionRequest):
|
|||
|
|
"""
|
|||
|
|
接收图的JSON并执行，可能会调用已通过TCP连接注册的远程节点。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
print("\n--- New Graph Execution Request ---")
|
|||
|
|
# 1. 从请求数据构建 WorkflowGraph
|
|||
|
|
graph = WorkflowGraph.from_dict(request.graph_data)
|
|||
|
|
print(f"Executing graph: '{graph.name}'")
|
|||
|
|
|
|||
|
|
# 2. 准备执行上下文，注入 trace_server 实例作为RPC服务
|
|||
|
|
# RemoteProxyNode 将通过这个服务来回调客户端
|
|||
|
|
context = TraceContext(services={"rpc": trace_server})
|
|||
|
|
|
|||
|
|
# 3. 实例化执行器并运行
|
|||
|
|
executor = WorkflowExecutor(graph, context)
|
|||
|
|
executor.execute_batch(request.initial_inputs)
|
|||
|
|
|
|||
|
|
# 4. 获取并返回最终输出
|
|||
|
|
final_outputs = context.get_outputs()
|
|||
|
|
print(f"Graph execution finished. Final outputs: {final_outputs}")
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"outputs": final_outputs
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
import traceback
|
|||
|
|
print(f"ERROR: Graph execution failed: {e}\n{traceback.format_exc()}")
|
|||
|
|
raise HTTPException(
|
|||
|
|
status_code=500,
|
|||
|
|
detail={
|
|||
|
|
"error": str(e),
|
|||
|
|
"type": type(e).__name__,
|
|||
|
|
"traceback": traceback.format_exc()
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# --- 3. 启动服务器 ---
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 使用 uvicorn 启动HTTP服务器。TCP服务器会在startup事件中被启动。
|
|||
|
|
# 您可以通过命令行启动: uvicorn server.remote_test_server:app --reload --port 8000
|
|||
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|