add first commit
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
[build-system]
|
||||
build-backend = "hatchling.build"
|
||||
requires = ["hatchling"]
|
||||
|
||||
[project]
|
||||
name = "chongming-request"
|
||||
version = "0.1.0"
|
||||
description = "worker direct communication"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"eclipse-zenoh>=1.9.0",
|
||||
]
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/chongming_request"]
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
import zenoh
|
||||
import orjson
|
||||
|
||||
|
||||
_SENTINEL = object()
|
||||
|
||||
class Client:
|
||||
def __init__(self,
|
||||
session: Optional[zenoh.Session] = None,
|
||||
logger: Optional[logging.Logger] = None):
|
||||
if session is None:
|
||||
self._session = zenoh.open(zenoh.Config())
|
||||
self._owns_session = True
|
||||
else:
|
||||
self._session = session
|
||||
self._owns_session = False
|
||||
self._logger = logger or logging.getLogger(__name__)
|
||||
self._worker_cache = {}
|
||||
self._cache_ttl = 15
|
||||
self._cache_timestamps = {}
|
||||
|
||||
async def _get_workers(self, worker_type: str) -> List[str]:
|
||||
now = time.monotonic()
|
||||
if worker_type in self._worker_cache:
|
||||
if now - self._cache_timestamps[worker_type] < self._cache_ttl:
|
||||
return self._worker_cache[worker_type]
|
||||
try:
|
||||
replies = self._session.liveliness().get(f"worker/{worker_type}/*")
|
||||
workers = [str(r.ok.key_expr) for r in replies if r.ok]
|
||||
self._worker_cache[worker_type] = workers
|
||||
self._cache_timestamps[worker_type] = now
|
||||
return workers
|
||||
except Exception as e:
|
||||
if worker_type in self._worker_cache:
|
||||
self._logger.warning("Using expired cache")
|
||||
return self._worker_cache[worker_type]
|
||||
self._logger.error(str(e))
|
||||
raise ValueError(str(e))
|
||||
|
||||
async def call(
|
||||
self,
|
||||
worker_type: str,
|
||||
subject: str,
|
||||
payload: Dict[str, Any],
|
||||
method: str = "GET",
|
||||
timeout: float = 3.0,
|
||||
retries: int = 1,
|
||||
) -> Dict[str, Any]:
|
||||
payload.setdefault("_method", method.upper())
|
||||
full_key = ""
|
||||
for attempt in range(retries + 1):
|
||||
workers = await self._get_workers(worker_type)
|
||||
if not workers:
|
||||
raise RuntimeError(...)
|
||||
target = random.choice(workers)
|
||||
full_key = f"{target}/{subject}"
|
||||
try:
|
||||
reply_iter = self._session.get(
|
||||
full_key,
|
||||
payload=orjson.dumps(payload),
|
||||
consolidation=zenoh.ConsolidationMode.NONE,
|
||||
)
|
||||
def _get_first(iterator):
|
||||
for item in iterator:
|
||||
return item
|
||||
return _SENTINEL
|
||||
first_reply = await asyncio.wait_for(
|
||||
asyncio.to_thread(_get_first, reply_iter),
|
||||
timeout=timeout,
|
||||
)
|
||||
if first_reply is _SENTINEL:
|
||||
self._logger.warning(f"Worker {target} did not respond to '{full_key}', retrying {attempt+1}/{retries}...")
|
||||
continue
|
||||
|
||||
reply = first_reply
|
||||
if reply.ok: # type: ignore
|
||||
payload_bytes = bytes(reply.ok.payload) if reply.ok.payload else b"" # type: ignore
|
||||
if not payload_bytes:
|
||||
result = {}
|
||||
else:
|
||||
result = orjson.loads(payload_bytes)
|
||||
return result
|
||||
else:
|
||||
error_payload = reply.err.payload # type: ignore
|
||||
error_msg = str(error_payload) if error_payload else "unknown error"
|
||||
self._logger.warning(f"Worker Error: {error_msg} retrying {attempt+1}/{retries}...")
|
||||
continue
|
||||
except asyncio.TimeoutError:
|
||||
raise RuntimeError(f"Request to '{full_key}' timed out after {timeout} seconds")
|
||||
raise RuntimeError(f"Failed to get a valid response from '{full_key}' after {retries} retries")
|
||||
|
||||
async def close(self):
|
||||
if self._owns_session and self._session:
|
||||
self._session.close()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
await self.close()
|
||||
|
||||
def call_sync(self, *args, **kwargs):
|
||||
return asyncio.run(self.call(*args, **kwargs))
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
[build-system]
|
||||
build-backend = "hatchling.build"
|
||||
requires = ["hatchling"]
|
||||
|
||||
[project]
|
||||
dependencies = [
|
||||
"eclipse-zenoh>=1.9.0",
|
||||
"tomli>=2.4.1",
|
||||
]
|
||||
description = "chongming worker"
|
||||
name = "chongming-worker"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
version = "0.1.0"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/chongming_worker"]
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
chongming_worker - 通用 Worker 框架
|
||||
|
||||
提供 Zenoh RPC 端点的装饰器定义、生命周期管理和配置管理。
|
||||
"""
|
||||
from chongming_worker.decorators import (
|
||||
RPCMethod,
|
||||
rpc_method,
|
||||
bind_model,
|
||||
auto_bind_models,
|
||||
get_rpc_methods,
|
||||
get_all_rpc_methods,
|
||||
clear_registry,
|
||||
generate_route_configs,
|
||||
)
|
||||
from chongming_worker.config import (
|
||||
load_worker_config,
|
||||
WorkerConfig,
|
||||
OperationConfig,
|
||||
)
|
||||
from chongming_worker.runtime import (
|
||||
create_worker,
|
||||
create_workers,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"RPCMethod",
|
||||
"rpc_method",
|
||||
"bind_model",
|
||||
"auto_bind_models",
|
||||
"get_rpc_methods",
|
||||
"get_all_rpc_methods",
|
||||
"clear_registry",
|
||||
"generate_route_configs",
|
||||
"load_worker_config",
|
||||
"WorkerConfig",
|
||||
"OperationConfig",
|
||||
"create_worker",
|
||||
"create_workers",
|
||||
]
|
||||
@@ -0,0 +1,263 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TypedDict, Literal, Optional, cast, List, Dict, Any
|
||||
|
||||
# 根据Python版本选择合适的TOML解析库和Required类型
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
from typing import Required
|
||||
else:
|
||||
import tomli as tomllib
|
||||
from typing_extensions import Required
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Worker 配置类型
|
||||
# ============================================================
|
||||
|
||||
class WorkerParams(TypedDict, total=False):
|
||||
"""Worker 操作参数定义"""
|
||||
type: str
|
||||
required: bool
|
||||
default: Any
|
||||
description: str
|
||||
|
||||
|
||||
class WorkerResponse(TypedDict, total=False):
|
||||
"""Worker 操作响应定义"""
|
||||
type: str
|
||||
required: bool
|
||||
default: Any
|
||||
|
||||
|
||||
class OperationConfig(TypedDict, total=False):
|
||||
"""
|
||||
Worker 操作配置结构
|
||||
|
||||
对应 config/workers/<name>/*.toml 中的每个操作定义。
|
||||
"""
|
||||
name: Required[str]
|
||||
method: Required[Literal["GET", "POST", "PUT", "DELETE"]]
|
||||
path: str
|
||||
subject: Required[str]
|
||||
summary: str
|
||||
timeout: float
|
||||
params: Dict[str, WorkerParams]
|
||||
response_model: Dict[str, WorkerResponse]
|
||||
|
||||
|
||||
class WorkerConfig(TypedDict, total=False):
|
||||
"""
|
||||
Worker 配置结构
|
||||
|
||||
对应 config/workers/<name>/config.toml 中的顶级配置。
|
||||
"""
|
||||
name: Required[str]
|
||||
type: Required[str]
|
||||
description: str
|
||||
version: str
|
||||
enabled: bool
|
||||
operations: List[OperationConfig]
|
||||
|
||||
|
||||
def load_worker_config(config_dir: Path) -> Dict[str, Any]:
|
||||
"""
|
||||
加载 Worker 配置目录中的所有 TOML 配置文件
|
||||
|
||||
Args:
|
||||
config_dir: Worker 配置目录路径，如 Path("config/workers/example")
|
||||
|
||||
Returns:
|
||||
解析后的配置字典，包含所有操作的配置
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: 当配置目录不存在时抛出异常
|
||||
"""
|
||||
if not config_dir.exists():
|
||||
raise FileNotFoundError(f"Worker 配置目录不存在: {config_dir}")
|
||||
|
||||
config: Dict[str, Any] = {"operations": []}
|
||||
|
||||
# 首先加载 config.toml（如果存在）
|
||||
config_file = config_dir / "config.toml"
|
||||
if config_file.exists():
|
||||
with open(config_file, "rb") as f:
|
||||
base_config = tomllib.load(f)
|
||||
config.update(base_config)
|
||||
|
||||
# 然后加载所有 *.toml 中的操作定义
|
||||
for toml_file in sorted(config_dir.glob("*.toml")):
|
||||
if toml_file.name == "config.toml":
|
||||
continue
|
||||
with open(toml_file, "rb") as f:
|
||||
op_data = tomllib.load(f)
|
||||
# 将文件名（不含扩展名）作为默认操作名
|
||||
op_name = toml_file.stem
|
||||
if "name" not in op_data:
|
||||
op_data["name"] = op_name
|
||||
config["operations"].append(op_data)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 旧的配置系统（与主应用 config.py 兼容）
|
||||
# ============================================================
|
||||
|
||||
# 设置 env 为 development 或 production
|
||||
|
||||
|
||||
class DefaultApp(TypedDict):
|
||||
"""应用默认配置结构"""
|
||||
|
||||
name: str
|
||||
version: str
|
||||
description: str
|
||||
debug: bool
|
||||
|
||||
|
||||
class Default(TypedDict):
|
||||
"""系统默认配置结构，包含应用信息、环境、路径等基础配置"""
|
||||
|
||||
app: DefaultApp
|
||||
env: Literal["production", "development"]
|
||||
prefix: str
|
||||
upload_path: str
|
||||
|
||||
|
||||
class SqliteConnectionArgs(TypedDict, total=False):
|
||||
"""SQLite数据库连接参数配置"""
|
||||
|
||||
timeout: int
|
||||
check_same_thread: bool
|
||||
|
||||
|
||||
class DatabaseSqlite(TypedDict, total=False):
|
||||
"""SQLite数据库连接池及行为配置"""
|
||||
|
||||
echo: bool
|
||||
pool_size: int
|
||||
max_overflow: int
|
||||
pool_timeout: int
|
||||
pool_recycle: int
|
||||
pool_pre_ping: bool
|
||||
connect_args: SqliteConnectionArgs
|
||||
|
||||
|
||||
class Database(TypedDict, total=False):
|
||||
"""数据库配置结构，支持SQLite、MySQL、PostgreSQL三种数据库类型"""
|
||||
|
||||
type: Required[Literal["sqlite", "mysql", "postgresql"]]
|
||||
url: Required[str]
|
||||
generate_schemas: bool
|
||||
sqlite: DatabaseSqlite
|
||||
mysql: dict
|
||||
postgresql: dict
|
||||
database_path: str
|
||||
|
||||
|
||||
class Security(TypedDict):
|
||||
"""安全相关配置，包括JWT令牌和安全策略设置"""
|
||||
|
||||
secret_key: str
|
||||
algorithm: Literal["HS256"]
|
||||
access_token_expire_minutes: int
|
||||
refresh_token_expire_days: int
|
||||
max_sessions_per_user: int
|
||||
|
||||
|
||||
class Server(TypedDict, total=False):
|
||||
"""服务器运行配置，包含主机、端口、工作进程等设置"""
|
||||
|
||||
host: Required[str]
|
||||
port: Required[int]
|
||||
reload: bool
|
||||
workers: int
|
||||
access_log: bool
|
||||
|
||||
|
||||
class Logging(TypedDict):
|
||||
"""日志系统配置，定义日志级别、输出格式和文件管理策略"""
|
||||
|
||||
level: str
|
||||
file: str
|
||||
format: str
|
||||
max_size: int
|
||||
backup_count: int
|
||||
|
||||
|
||||
class Cors(TypedDict, total=False):
|
||||
"""跨域资源共享(CORS)配置"""
|
||||
|
||||
allow_origins: List[str]
|
||||
allow_credentials: bool
|
||||
allow_methods: List[str]
|
||||
allow_headers: List[str]
|
||||
|
||||
|
||||
class ModuleSystem(TypedDict, total=False):
|
||||
"""模块系统配置，定义模块加载类型和路径"""
|
||||
|
||||
type: str
|
||||
path: List[str]
|
||||
|
||||
|
||||
class FileSystem(TypedDict, total=False):
|
||||
"""文件系统配置，定义存储类型和路径"""
|
||||
|
||||
type: str
|
||||
path: str
|
||||
|
||||
|
||||
class Env(TypedDict, total=False):
|
||||
"""环境特定配置结构，用于开发和生产环境的差异化配置"""
|
||||
|
||||
debug: Required[bool]
|
||||
server: Required[Server]
|
||||
logging: Required[Logging]
|
||||
cors: Cors
|
||||
module_system: ModuleSystem
|
||||
file_system: FileSystem
|
||||
|
||||
|
||||
class Scheduler(TypedDict):
|
||||
"""任务调度器配置"""
|
||||
|
||||
job_store_path: str
|
||||
|
||||
|
||||
class Cache(TypedDict):
|
||||
"""缓存系统配置"""
|
||||
|
||||
cache_store_path: str
|
||||
|
||||
|
||||
class Config(TypedDict):
|
||||
"""完整配置结构，整合所有配置项，包含默认配置、数据库、调度器、缓存、安全和环境配置"""
|
||||
|
||||
default: Default
|
||||
database: Database
|
||||
scheduler: Scheduler
|
||||
cache: Cache
|
||||
security: Security
|
||||
development: Env
|
||||
production: Env
|
||||
|
||||
|
||||
def load_config(config_path: Path = Path("config.toml")) -> Config:
|
||||
"""
|
||||
从TOML配置文件加载配置数据
|
||||
|
||||
Args:
|
||||
config_path: 配置文件路径，默认为当前目录下的config.toml
|
||||
|
||||
Returns:
|
||||
Config: 解析后的配置字典对象
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: 当配置文件不存在时抛出异常
|
||||
"""
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"配置文件不存在: {config_path}")
|
||||
with open(config_path, "rb") as f:
|
||||
return cast(Config, tomllib.load(f))
|
||||
@@ -0,0 +1,546 @@
|
||||
"""
|
||||
通用 Working 装饰器框架
|
||||
|
||||
提供 @rpc_method 装饰器，用于将普通函数注册为 Zenoh RPC 端点。
|
||||
装饰器自动处理：
|
||||
- 访问类型（GET/POST/PUT/DELETE）
|
||||
- 接口名称（subject）
|
||||
- 参数解析与类型转换
|
||||
- 响应序列化
|
||||
|
||||
兼容 gen_models 命令生成的 Pydantic 模型：
|
||||
- 通过 request_model / response_model 参数绑定已生成的模型
|
||||
- 自动从插件 models 目录发现匹配的模型
|
||||
- 请求参数自动校验与响应格式自动序列化
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import inspect
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, get_type_hints, Union
|
||||
|
||||
from pydantic import BaseModel, Field, create_model, ValidationError
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 类型系统
|
||||
# ============================================================
|
||||
|
||||
TYPE_MAP: Dict[str, Any] = {
|
||||
"float": float,
|
||||
"int": int,
|
||||
"str": str,
|
||||
"bool": bool,
|
||||
"any": Any,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_type(annotation: Any) -> Any:
|
||||
"""将类型注解解析为可用的 Python 类型"""
|
||||
if annotation is inspect.Parameter.empty:
|
||||
return str
|
||||
origin = getattr(annotation, "__origin__", None)
|
||||
if origin is list:
|
||||
args = getattr(annotation, "__args__", (Any,))
|
||||
return List[args[0]] if args else List[Any] # type: ignore
|
||||
if origin is dict:
|
||||
args = getattr(annotation, "__args__", (str, Any))
|
||||
return Dict[args[0], args[1]] if len(args) == 2 else Dict[str, Any] # type: ignore
|
||||
if origin is tuple:
|
||||
args = getattr(annotation, "__args__", (Any,))
|
||||
return Tuple[tuple(args)] # type: ignore
|
||||
if origin is Optional:
|
||||
args = getattr(annotation, "__args__", (Any,))
|
||||
return Optional[args[0]] # type: ignore
|
||||
return annotation
|
||||
|
||||
|
||||
def _type_to_config_str(py_type: Any) -> str:
|
||||
"""将 Python 类型转换为配置字符串（用于 gen_models 兼容）"""
|
||||
origin = getattr(py_type, "__origin__", None)
|
||||
# 处理 Union 类型（包括 Optional[X]）
|
||||
if origin is Union:
|
||||
args = getattr(py_type, "__args__", (Any,))
|
||||
non_none_args = [a for a in args if a is not type(None)]
|
||||
if len(non_none_args) == 1:
|
||||
return _type_to_config_str(non_none_args[0])
|
||||
return "str"
|
||||
if origin is list:
|
||||
args = getattr(py_type, "__args__", (Any,))
|
||||
return f"list[{_type_to_config_str(args[0])}]"
|
||||
if origin is dict:
|
||||
args = getattr(py_type, "__args__", (str, Any))
|
||||
return f"dict[{_type_to_config_str(args[0])}, {_type_to_config_str(args[1])}]"
|
||||
if py_type is float:
|
||||
return "float"
|
||||
if py_type is int:
|
||||
return "int"
|
||||
if py_type is bool:
|
||||
return "bool"
|
||||
if py_type is str:
|
||||
return "str"
|
||||
return "str"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 装饰器定义
|
||||
# ============================================================
|
||||
|
||||
class RPCMethod:
|
||||
"""
|
||||
装饰后的 RPC 方法元数据
|
||||
|
||||
存储原始函数、访问类型、接口名称、参数信息等，
|
||||
供 Worker 在运行时发现和调用。
|
||||
|
||||
支持绑定 gen_models 生成的 Pydantic 模型：
|
||||
- request_model: 用于请求参数的自动校验和转换
|
||||
- response_model: 用于响应数据的格式化和序列化
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: Callable,
|
||||
subject: str,
|
||||
method: str = "GET",
|
||||
summary: str = "",
|
||||
timeout: float = 3.0,
|
||||
request_model: Optional[Type[BaseModel]] = None,
|
||||
response_model: Optional[Type[BaseModel]] = None,
|
||||
):
|
||||
self.func = func
|
||||
self.subject = subject
|
||||
self.method = method.upper()
|
||||
self.summary = summary or func.__doc__ or ""
|
||||
self.timeout = timeout
|
||||
self.name = func.__name__
|
||||
|
||||
# 绑定的 Pydantic 模型（由 gen_models 生成，或手动传入）
|
||||
self.request_model: Optional[Type[BaseModel]] = request_model
|
||||
self.response_model: Optional[Type[BaseModel]] = response_model
|
||||
|
||||
# 解析函数签名
|
||||
sig = inspect.signature(func)
|
||||
hints = get_type_hints(func)
|
||||
self.params: Dict[str, Any] = {}
|
||||
self.model_param: Optional[str] = None # 如果有参数是 BaseModel，记录参数名
|
||||
self.return_type: Any = str
|
||||
|
||||
for param_name, param in sig.parameters.items():
|
||||
if param_name == "self":
|
||||
continue
|
||||
py_type = _resolve_type(hints.get(param_name, str))
|
||||
# 检查参数类型是否是 BaseModel 子类
|
||||
if isinstance(py_type, type) and issubclass(py_type, BaseModel):
|
||||
self.model_param = param_name
|
||||
# 如果是 BaseModel 参数，直接用它的字段作为 params（用于配置和文档）
|
||||
for field_name, field in py_type.model_fields.items():
|
||||
self.params[field_name] = _resolve_type(field.annotation)
|
||||
else:
|
||||
self.params[param_name] = py_type
|
||||
|
||||
# 返回值类型
|
||||
self.return_type = _resolve_type(hints.get("return", Any))
|
||||
|
||||
def build_pydantic_request_model(self, model_name: str = "AutoRequest") -> Type[BaseModel]:
|
||||
"""根据参数注解生成 Pydantic 请求模型"""
|
||||
fields: Dict[str, Any] = {}
|
||||
for param_name, py_type in self.params.items():
|
||||
fields[param_name] = (py_type, Field(...))
|
||||
return create_model(model_name, **fields)
|
||||
|
||||
def build_pydantic_response_model(self, model_name: str = "AutoResponse") -> Type[BaseModel]:
|
||||
"""根据返回值注解生成 Pydantic 响应模型"""
|
||||
fields: Dict[str, Any] = {}
|
||||
# 如果返回类型是 BaseModel 子类，直接使用
|
||||
if isinstance(self.return_type, type) and issubclass(self.return_type, BaseModel):
|
||||
return self.return_type
|
||||
# 否则包装为 {result: return_type}
|
||||
fields["result"] = (self.return_type, Field(...))
|
||||
return create_model(model_name, **fields)
|
||||
|
||||
def get_request_model(self) -> Optional[Type[BaseModel]]:
|
||||
"""
|
||||
获取有效的请求模型，优先级：
|
||||
1. 绑定的 request_model（由用户或 gen_models 提供）
|
||||
2. 如果函数参数本身就是 Pydantic 模型（model_param），从中提取字段生成模型
|
||||
3. 自动生成的模型（根据函数参数注解，仅 POST/PUT/PATCH）
|
||||
"""
|
||||
if self.request_model is not None:
|
||||
return self.request_model
|
||||
# 如果函数接受了 Pydantic 模型参数，用它的字段生成请求模型
|
||||
if self.model_param is not None:
|
||||
sig = inspect.signature(self.func)
|
||||
hints = get_type_hints(self.func)
|
||||
model_type = hints.get(self.model_param)
|
||||
if isinstance(model_type, type) and issubclass(model_type, BaseModel):
|
||||
return model_type
|
||||
# 对于 POST/PUT/PATCH，根据 params 自动生成
|
||||
if self.params and self.method in ("POST", "PUT", "PATCH"):
|
||||
return self.build_pydantic_request_model(f"{self.name.title()}_Request")
|
||||
return None
|
||||
|
||||
def get_response_model(self) -> Optional[Type[BaseModel]]:
|
||||
"""
|
||||
获取有效的响应模型，优先级：
|
||||
1. 绑定的 response_model（由用户或 gen_models 提供）
|
||||
2. 自动生成的模型（根据函数返回值注解）
|
||||
"""
|
||||
if self.response_model is not None:
|
||||
return self.response_model
|
||||
return self.build_pydantic_response_model(f"{self.name.title()}_Response")
|
||||
|
||||
def validate_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
使用 Pydantic 模型校验请求参数
|
||||
|
||||
Args:
|
||||
payload: 原始请求参数字典
|
||||
|
||||
Returns:
|
||||
校验通过后的参数字典（包含类型转换）
|
||||
"""
|
||||
req_model = self.get_request_model()
|
||||
if req_model is None:
|
||||
return payload
|
||||
try:
|
||||
validated = req_model(**payload)
|
||||
return validated.model_dump()
|
||||
except ValidationError as e:
|
||||
raise ValueError(f"请求参数校验失败: {e.errors()}")
|
||||
|
||||
def serialize_response(self, raw_result: Any) -> Any:
|
||||
"""
|
||||
使用 Pydantic 模型序列化响应数据
|
||||
|
||||
Args:
|
||||
raw_result: 原始函数返回值
|
||||
|
||||
Returns:
|
||||
序列化后的响应数据（dict 或列表等可 JSON 序列化的类型）
|
||||
"""
|
||||
resp_model = self.get_response_model()
|
||||
if resp_model is None:
|
||||
return raw_result
|
||||
if isinstance(raw_result, dict):
|
||||
try:
|
||||
validated = resp_model(**raw_result)
|
||||
return validated.model_dump()
|
||||
except ValidationError:
|
||||
# 如果无法用响应模型校验，返回原始结果
|
||||
return raw_result
|
||||
if isinstance(raw_result, resp_model):
|
||||
return raw_result.model_dump()
|
||||
return raw_result
|
||||
|
||||
async def call(self, payload: Dict[str, Any]) -> Any:
|
||||
"""
|
||||
调用原始函数，自动从 payload 提取参数
|
||||
|
||||
流程：
|
||||
1. 校验请求参数（使用绑定的 request_model 或自动生成）
|
||||
2. 提取函数需要的参数
|
||||
3. 调用函数
|
||||
4. 序列化响应（使用绑定的 response_model 或自动生成）
|
||||
|
||||
支持两种函数签名：
|
||||
- 展开参数: def add(a: float, b: float) -> ... # 传统方式
|
||||
- Pydantic 模型: def add(req: AddGetRequest) -> ... # 工程化方式
|
||||
|
||||
Args:
|
||||
payload: 请求参数字典
|
||||
|
||||
Returns:
|
||||
序列化后的响应数据
|
||||
"""
|
||||
# 1. 校验请求参数
|
||||
validated_payload = self.validate_request(payload)
|
||||
|
||||
# 2. 提取函数需要的参数
|
||||
kwargs = {}
|
||||
|
||||
if self.model_param is not None:
|
||||
# 2a. 函数参数是 Pydantic 模型 — 将整个校验后的 payload 转为模型实例传入
|
||||
req_model = self.get_request_model()
|
||||
if req_model is not None:
|
||||
kwargs[self.model_param] = req_model(**validated_payload)
|
||||
else:
|
||||
# 如果没有 request_model，从函数签名获取模型类型
|
||||
sig = inspect.signature(self.func)
|
||||
hints = get_type_hints(self.func)
|
||||
model_type = hints.get(self.model_param)
|
||||
if isinstance(model_type, type) and issubclass(model_type, BaseModel):
|
||||
kwargs[self.model_param] = model_type(**validated_payload)
|
||||
else:
|
||||
kwargs[self.model_param] = validated_payload
|
||||
else:
|
||||
# 2b. 传统方式：逐字段提取参数
|
||||
for param_name in self.params:
|
||||
if param_name in validated_payload:
|
||||
kwargs[param_name] = validated_payload[param_name]
|
||||
|
||||
# 3. 调用函数
|
||||
if inspect.iscoroutinefunction(self.func):
|
||||
result = await self.func(**kwargs)
|
||||
else:
|
||||
result = self.func(**kwargs)
|
||||
|
||||
# 4. 序列化响应
|
||||
return self.serialize_response(result)
|
||||
|
||||
|
||||
# 全局注册表：worker_type -> {subject -> RPCMethod}
|
||||
_RPC_REGISTRY: Dict[str, Dict[str, RPCMethod]] = {}
|
||||
|
||||
|
||||
def rpc_method(
|
||||
subject: str,
|
||||
method: str = "GET",
|
||||
summary: str = "",
|
||||
timeout: float = 3.0,
|
||||
request_model: Optional[Type[BaseModel]] = None,
|
||||
response_model: Optional[Type[BaseModel]] = None,
|
||||
):
|
||||
"""
|
||||
装饰器：将函数注册为 RPC 端点
|
||||
|
||||
用法:
|
||||
# 基本用法（自动生成模型）
|
||||
@rpc_method(subject="calc.add", method="GET", summary="Add two numbers")
|
||||
def add(a: float, b: float) -> float:
|
||||
return a + b
|
||||
|
||||
# 绑定 gen_models 生成的 Pydantic 模型
|
||||
@rpc_method(
|
||||
subject="calc.add",
|
||||
method="GET",
|
||||
summary="Add two numbers",
|
||||
response_model=AddGetResponse,
|
||||
)
|
||||
def add(a: float, b: float):
|
||||
return a + b
|
||||
|
||||
参数:
|
||||
subject: 接口名称，格式为 "worker_type.operation"
|
||||
method: HTTP 方法 (GET/POST/PUT/DELETE)
|
||||
summary: 接口描述
|
||||
timeout: 超时时间（秒）
|
||||
request_model: 可选的 Pydantic 请求模型（用于验证请求参数）
|
||||
response_model: 可选的 Pydantic 响应模型（用于格式化响应）
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
rpc = RPCMethod(
|
||||
func=func,
|
||||
subject=subject,
|
||||
method=method,
|
||||
summary=summary,
|
||||
timeout=timeout,
|
||||
request_model=request_model,
|
||||
response_model=response_model,
|
||||
)
|
||||
# 提取 worker_type（subject 的第一部分）
|
||||
worker_type = subject.split(".")[0] if "." in subject else "default"
|
||||
if worker_type not in _RPC_REGISTRY:
|
||||
_RPC_REGISTRY[worker_type] = {}
|
||||
# 使用 subject + "#" + method 作为唯一 key，允许多个 HTTP 方法共享同一 subject
|
||||
registry_key = f"{subject}#{method.upper()}"
|
||||
_RPC_REGISTRY[worker_type][registry_key] = rpc
|
||||
|
||||
# 给函数附加元数据，方便外部访问
|
||||
func._rpc_meta = rpc # type: ignore
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def bind_model(
|
||||
subject: str = "",
|
||||
method: str = "",
|
||||
request_model: Optional[Type[BaseModel]] = None,
|
||||
response_model: Optional[Type[BaseModel]] = None,
|
||||
):
|
||||
"""
|
||||
装饰器：为已注册的 RPC 端点绑定 Pydantic 模型
|
||||
|
||||
用于 gen_models 生成模型后，将模型绑定到用 @rpc_method 定义的端点。
|
||||
无需修改原有的 @rpc_method 调用。
|
||||
|
||||
用法:
|
||||
from .models.pydantic import AddGetResponse
|
||||
|
||||
@bind_model(subject="example.add", method="GET", response_model=AddGetResponse)
|
||||
def add():
|
||||
pass # 函数体不重要，仅用于匹配已注册的端点
|
||||
|
||||
参数:
|
||||
subject: 要绑定的接口名称（与 @rpc_method 中的 subject 一致）
|
||||
method: 要绑定的 HTTP 方法
|
||||
request_model: Pydantic 请求模型
|
||||
response_model: Pydantic 响应模型
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
nonlocal subject, method
|
||||
if not subject and hasattr(func, '_rpc_meta'):
|
||||
subject = func._rpc_meta.subject # type: ignore[attr-defined]
|
||||
method = func._rpc_meta.method # type: ignore[attr-defined]
|
||||
|
||||
worker_type = subject.split(".")[0] if "." in subject else "default"
|
||||
method_upper = method.upper()
|
||||
registry = _RPC_REGISTRY.get(worker_type, {})
|
||||
# 用新 key 格式查找: subject#METHOD
|
||||
rpc = registry.get(f"{subject}#{method_upper}")
|
||||
|
||||
if rpc is None:
|
||||
raise ValueError(
|
||||
f"未找到已注册的 RPC 端点: subject='{subject}', method='{method}'. "
|
||||
f"请确保 @rpc_method 装饰器已先于 @bind_model 执行。"
|
||||
)
|
||||
|
||||
# 如果指定了 method，进一步筛选
|
||||
if method and rpc.method != method.upper():
|
||||
return func
|
||||
|
||||
if request_model is not None:
|
||||
rpc.request_model = request_model
|
||||
if response_model is not None:
|
||||
rpc.response_model = response_model
|
||||
|
||||
func._rpc_meta = rpc # type: ignore[attr-defined]
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def auto_bind_models(plugin_name: str):
|
||||
"""
|
||||
自动从插件 models 目录发现并绑定 Pydantic 模型
|
||||
|
||||
扫描 src/plugins/<plugin_name>/models/pydantic.py 中定义的所有模型，
|
||||
尝试按命名约定匹配已注册的 RPC 端点。
|
||||
|
||||
命名约定：
|
||||
- 请求模型: <Operation><Method>Request (e.g., AddPostRequest)
|
||||
- 响应模型: <Operation><Method>Response (e.g., AddGetResponse)
|
||||
|
||||
用法:
|
||||
# 在插件 __init__.py 中调用
|
||||
from ...decorators import auto_bind_models
|
||||
auto_bind_models("example")
|
||||
"""
|
||||
try:
|
||||
models_module = importlib.import_module(f"plugins.{plugin_name}.models.pydantic")
|
||||
except ImportError:
|
||||
try:
|
||||
models_module = importlib.import_module(f"src.plugins.{plugin_name}.models.pydantic")
|
||||
except ImportError:
|
||||
print(f"⚠️ auto_bind_models: 未找到 models 模块 (plugins.{plugin_name}.models.pydantic)")
|
||||
return
|
||||
|
||||
# 获取模块中所有 BaseModel 子类
|
||||
model_classes = {}
|
||||
for name in dir(models_module):
|
||||
obj = getattr(models_module, name)
|
||||
if isinstance(obj, type) and issubclass(obj, BaseModel) and obj is not BaseModel:
|
||||
model_classes[name] = obj
|
||||
|
||||
if not model_classes:
|
||||
print(f"⚠️ auto_bind_models: 未找到任何 Pydantic 模型")
|
||||
return
|
||||
|
||||
registry = _RPC_REGISTRY.get(plugin_name, {})
|
||||
bound_count = 0
|
||||
|
||||
for _key, rpc in registry.items():
|
||||
op_name = rpc.subject.split(".")[-1] if "." in rpc.subject else rpc.subject
|
||||
op_camel = op_name.title().replace("-", "").replace("_", "")
|
||||
method_tag = rpc.method.title()
|
||||
|
||||
# 尝试匹配请求模型: <Op><Method>Request
|
||||
req_model_name = f"{op_camel}{method_tag}Request"
|
||||
if req_model_name in model_classes:
|
||||
rpc.request_model = model_classes[req_model_name]
|
||||
bound_count += 1
|
||||
print(f" ✓ 绑定请求模型: {req_model_name} -> {rpc.subject} [{rpc.method}]")
|
||||
|
||||
# 尝试匹配响应模型: <Op><Method>Response
|
||||
resp_model_name = f"{op_camel}{method_tag}Response"
|
||||
if resp_model_name in model_classes:
|
||||
rpc.response_model = model_classes[resp_model_name]
|
||||
bound_count += 1
|
||||
print(f" ✓ 绑定响应模型: {resp_model_name} -> {rpc.subject} [{rpc.method}]")
|
||||
|
||||
if bound_count > 0:
|
||||
print(f"✅ auto_bind_models: 已绑定 {bound_count} 个模型到 {plugin_name} 的 RPC 端点")
|
||||
else:
|
||||
print(f"⚠️ auto_bind_models: 未找到匹配的模型绑定")
|
||||
|
||||
|
||||
def get_rpc_methods(worker_type: str) -> Dict[str, RPCMethod]:
|
||||
"""获取指定 worker 类型的所有 RPC 方法"""
|
||||
return _RPC_REGISTRY.get(worker_type, {})
|
||||
|
||||
|
||||
def get_all_rpc_methods() -> Dict[str, Dict[str, RPCMethod]]:
|
||||
"""获取所有已注册的 RPC 方法"""
|
||||
return _RPC_REGISTRY
|
||||
|
||||
|
||||
def clear_registry():
|
||||
"""清空注册表（主要用于测试）"""
|
||||
_RPC_REGISTRY.clear()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 路由配置生成器（从装饰器元数据生成 TOML 兼容的配置）
|
||||
# ============================================================
|
||||
|
||||
def generate_route_configs(worker_type: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
从装饰器注册表生成路由配置列表
|
||||
|
||||
返回的配置与 config/workers/<type>/*.toml 格式兼容，
|
||||
可用于动态路由注册或生成配置文件。
|
||||
"""
|
||||
methods = get_rpc_methods(worker_type)
|
||||
configs: List[Dict[str, Any]] = []
|
||||
for _key, rpc in methods.items():
|
||||
# 从 subject 提取操作名（最后一部分）
|
||||
op_name = rpc.subject.split(".")[-1] if "." in rpc.subject else rpc.subject
|
||||
configs.append({
|
||||
"subject": rpc.subject,
|
||||
"path": f"/{op_name}",
|
||||
"method": rpc.method,
|
||||
"params": [f"{name}: {_type_to_config_str(pt)}" for name, pt in rpc.params.items()],
|
||||
"response_model": _build_response_spec(rpc),
|
||||
"summary": rpc.summary,
|
||||
"timeout": rpc.timeout,
|
||||
})
|
||||
return configs
|
||||
|
||||
|
||||
def _build_response_spec(rpc: RPCMethod) -> Dict[str, Any]:
|
||||
"""从 RPC 方法构建 response_model 配置"""
|
||||
spec: Dict[str, Any] = {}
|
||||
|
||||
# 优先使用绑定的 response_model
|
||||
resp_model = rpc.response_model
|
||||
if resp_model is None:
|
||||
resp_model = rpc.build_pydantic_response_model()
|
||||
|
||||
if isinstance(resp_model, type) and issubclass(resp_model, BaseModel):
|
||||
# 遍历 Pydantic 模型字段
|
||||
for field_name, field in resp_model.model_fields.items():
|
||||
field_type = _type_to_config_str(field.annotation)
|
||||
required = field.is_required()
|
||||
default = field.default if not field.is_required() else None
|
||||
spec[field_name] = {
|
||||
"type": field_type,
|
||||
"required": required,
|
||||
}
|
||||
if default is not None:
|
||||
spec[field_name]["default"] = default
|
||||
else:
|
||||
spec["result"] = {"type": _type_to_config_str(rpc.return_type), "required": True}
|
||||
return spec
|
||||
@@ -0,0 +1,179 @@
|
||||
"""
|
||||
Worker 运行时 — 通用的生命周期管理
|
||||
|
||||
提供 create_worker / create_workers 工厂函数，供所有 Worker 插件共享。
|
||||
"""
|
||||
import os
|
||||
import asyncio
|
||||
import json
|
||||
import orjson
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import zenoh
|
||||
|
||||
from chongming_worker.decorators import (
|
||||
RPCMethod,
|
||||
get_rpc_methods,
|
||||
)
|
||||
|
||||
async def create_worker(
|
||||
session: zenoh.Session,
|
||||
worker_type: str,
|
||||
worker_id: str,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
worker_manager=None,
|
||||
):
|
||||
"""
|
||||
创建一个 Worker 协程
|
||||
|
||||
自动发现通过 @rpc_method 注册的端点，为每个 subject 声明 Zenoh Queryable，
|
||||
收到请求后按 _method 字段分派到具体方法。
|
||||
|
||||
Args:
|
||||
session: Zenoh 会话
|
||||
worker_type: Worker 类型标识（匹配 subject 的第一部分）
|
||||
worker_id: Worker 唯一 ID
|
||||
config: 自定义配置参数
|
||||
worker_manager: WorkerManager 实例（用于本地注册）
|
||||
"""
|
||||
print(f"[Worker] {worker_id} 启动 (type={worker_type})")
|
||||
|
||||
# 1. 注册存活令牌
|
||||
token = session.liveliness().declare_token(f"worker/{worker_type}/{worker_id}")
|
||||
|
||||
# 2. 注册到本地 WorkerManager（若有）
|
||||
if worker_manager is not None:
|
||||
worker_manager.register_worker(worker_type, worker_id)
|
||||
|
||||
# 3. 获取该 worker 类型的所有 RPC 方法
|
||||
rpc_methods = get_rpc_methods(worker_type)
|
||||
if not rpc_methods:
|
||||
print(f"[Worker] {worker_id} 警告: 未找到任何 RPC 方法")
|
||||
else:
|
||||
print(f"[Worker] {worker_id} 发现 {len(rpc_methods)} 个端点:")
|
||||
for _key, rpc in rpc_methods.items():
|
||||
print(f" {rpc.method} {rpc.subject} -> {rpc.name}({', '.join(rpc.params.keys())})")
|
||||
|
||||
# 4. 按 subject 分组（多个方法可共享同一 subject，如 GET + POST）
|
||||
subject_methods: Dict[str, List[RPCMethod]] = {}
|
||||
for registry_key, rpc in rpc_methods.items():
|
||||
subject_methods.setdefault(rpc.subject, []).append(rpc)
|
||||
|
||||
# 5. 为每个 subject 创建回调函数并声明 Queryable
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
def make_callback(subject: str, methods: List[RPCMethod]):
|
||||
"""创建处理特定 subject 的回调函数"""
|
||||
def callback(query: zenoh.Query):
|
||||
try:
|
||||
# 解析请求负载
|
||||
# payload_str = str(query.payload) if query.payload else "{}"
|
||||
payload_bytes = bytes(query.payload) if query.payload else b""
|
||||
if not payload_bytes:
|
||||
data = {}
|
||||
else:
|
||||
data = orjson.loads(payload_bytes)
|
||||
http_method = data.get("_method", "GET").upper()
|
||||
|
||||
# 匹配方法（按 HTTP 方法）
|
||||
matched = next((m for m in methods if m.method.upper() == http_method), None)
|
||||
if not matched and methods:
|
||||
matched = methods[0] # 降级到第一个
|
||||
if not matched:
|
||||
query.reply(
|
||||
key_expr=query.key_expr,
|
||||
payload=json.dumps({"error": f"No method for {subject}"}).encode("utf-8")
|
||||
)
|
||||
return
|
||||
|
||||
# 调用 RPC 方法（支持同步和异步）
|
||||
result = asyncio.run(matched.call(data))
|
||||
|
||||
# 回复
|
||||
query.reply(
|
||||
key_expr=query.key_expr,
|
||||
payload=orjson.dumps(result)
|
||||
)
|
||||
print(f"[Worker] {worker_id} 回复 {matched.method} {subject}: {result}")
|
||||
except Exception as e:
|
||||
error_response = {
|
||||
"result": 0.0,
|
||||
"operation": "error",
|
||||
"error": str(e),
|
||||
"worker": worker_id
|
||||
}
|
||||
query.reply(
|
||||
key_expr=query.key_expr,
|
||||
payload=json.dumps(error_response).encode("utf-8")
|
||||
)
|
||||
print(f"[Worker] {worker_id} 异常 {subject}: {e}")
|
||||
return callback
|
||||
|
||||
# 声明所有 Queryable
|
||||
queryables: List[zenoh.Queryable] = []
|
||||
for subject, methods in subject_methods.items():
|
||||
key = f"worker/{worker_type}/{worker_id}/{subject}"
|
||||
# 使用回调方式声明
|
||||
q = session.declare_queryable(key, make_callback(subject, methods))
|
||||
queryables.append(q)
|
||||
print(f"[Worker] {worker_id} Queryable: '{key}' [{', '.join(m.method for m in methods)}]")
|
||||
|
||||
# 6. 保持 worker 存活直到被取消
|
||||
try:
|
||||
await stop_event.wait()
|
||||
except asyncio.CancelledError:
|
||||
print(f"[Worker] {worker_id} 正在关闭...")
|
||||
except Exception as e:
|
||||
print(f"[Worker] {worker_id} 发生异常: {e}")
|
||||
finally:
|
||||
# 清理
|
||||
if worker_manager is not None:
|
||||
worker_manager.unregister_worker(worker_type, worker_id)
|
||||
token.undeclare()
|
||||
for q in queryables:
|
||||
q.undeclare()
|
||||
print(f"[Worker] {worker_id} 已关闭")
|
||||
|
||||
def create_workers(
|
||||
session: zenoh.Session,
|
||||
config: dict,
|
||||
count: int = 1,
|
||||
worker_manager=None,
|
||||
) -> list[asyncio.Task]:
|
||||
"""
|
||||
创建一组 Worker 协程任务
|
||||
|
||||
工厂函数，供 WorkerManager 动态加载使用。
|
||||
|
||||
Args:
|
||||
session: Zenoh 会话
|
||||
config: Worker 配置字典（至少需要 "type" 字段）
|
||||
count: 启动的 Worker 副本数
|
||||
worker_manager: WorkerManager 实例（可选）
|
||||
|
||||
Returns:
|
||||
list[asyncio.Task]: Worker 协程任务列表
|
||||
"""
|
||||
worker_type = config.get("type", "default")
|
||||
pid = os.getpid()
|
||||
tasks = []
|
||||
for i in range(count):
|
||||
worker_id = f"worker-{i}-{pid}"
|
||||
task = asyncio.create_task(
|
||||
create_worker(
|
||||
session=session,
|
||||
worker_type=worker_type,
|
||||
worker_id=worker_id,
|
||||
config=config,
|
||||
worker_manager=worker_manager,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
print(f"[Worker] 已创建任务: {worker_id} (type={worker_type})")
|
||||
return tasks
|
||||
|
||||
|
||||
__all__ = [
|
||||
"create_worker",
|
||||
"create_workers",
|
||||
]
|
||||
Reference in New Issue
Block a user