commit 0a17b3367cb6a9cfc05a673f7482cff224f355ce Author: chakcy <947105045@qq.com> Date: Tue Jul 21 15:46:02 2026 +0800 add first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9989abd --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +# Python-generated files +__pycache__/ +*.py[oc] +build/ +build_nuitka/ +dist/ +wheels/ +*.egg-info + +# Virtual environments +.venv \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..292f9e3 --- /dev/null +++ b/pyproject.toml @@ -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"] diff --git a/src/chongming_worker/__init__.py b/src/chongming_worker/__init__.py new file mode 100644 index 0000000..3f2aca2 --- /dev/null +++ b/src/chongming_worker/__init__.py @@ -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", +] diff --git a/src/chongming_worker/config.py b/src/chongming_worker/config.py new file mode 100644 index 0000000..240fc54 --- /dev/null +++ b/src/chongming_worker/config.py @@ -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//*.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//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)) diff --git a/src/chongming_worker/decorators.py b/src/chongming_worker/decorators.py new file mode 100644 index 0000000..d943cc2 --- /dev/null +++ b/src/chongming_worker/decorators.py @@ -0,0 +1,537 @@ +""" +通用 Working 装饰器框架 + +提供 @rpc_method 装饰器,用于将普通函数注册为 Zenoh RPC 端点。 +装饰器自动处理: +- 访问类型(GET/POST/PUT/DELETE) +- 接口名称(subject) +- 参数解析与类型转换 +- 响应序列化 + +兼容 gen_models 命令生成的 Pydantic 模型: +- 通过 request_model / response_model 参数绑定已生成的模型 +- 自动从插件 models 目录发现匹配的模型 +- 请求参数自动校验与响应格式自动序列化 +""" +import inspect +import importlib +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], **extra_kwargs) -> Any: + """ + 调用原始函数,支持从 payload 和 extra_kwargs 提取参数。 + + 支持两种函数签名: + - 展开参数: def add(a: float, b: float) -> ... + - Pydantic 模型: def add(req: AddGetRequest) -> ... + + Args: + payload: 请求参数字典 + **extra_kwargs: 额外关键字参数(如 session, worker_id) + + Returns: + 序列化后的响应数据 + """ + validated_payload = self.validate_request(payload) + + sig = inspect.signature(self.func) + param_names = set(sig.parameters.keys()) - {'self'} + kwargs = {} + + # 处理模型参数 + if self.model_param is not None: + req_model = self.get_request_model() + if req_model is not None: + kwargs[self.model_param] = req_model(**validated_payload) + else: + model_type = sig.parameters[self.model_param].annotation + 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 + param_names.discard(self.model_param) + + # 普通参数:优先 validated_payload,其次 extra_kwargs + for name in param_names: + if name in validated_payload: + kwargs[name] = validated_payload[name] + elif name in extra_kwargs: + kwargs[name] = extra_kwargs[name] + # 否则函数使用默认值(如果有) + + # 调用函数 + if inspect.iscoroutinefunction(self.func): + result = await self.func(**kwargs) + else: + result = self.func(**kwargs) + + 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//models/pydantic.py 中定义的所有模型, + 尝试按命名约定匹配已注册的 RPC 端点。 + + 命名约定: + - 请求模型: Request (e.g., AddPostRequest) + - 响应模型: 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() + + # 尝试匹配请求模型: 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}]") + + # 尝试匹配响应模型: 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//*.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 diff --git a/src/chongming_worker/runtime.py b/src/chongming_worker/runtime.py new file mode 100644 index 0000000..0d496b5 --- /dev/null +++ b/src/chongming_worker/runtime.py @@ -0,0 +1,187 @@ +""" +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 的回调函数""" + loop = asyncio.get_running_loop() + async def _handle_query(query: zenoh.Query): + try: + # 解析请求负载 + # payload_str = str(query.payload) if query.payload else "{}" + payload_bytes = bytes(query.payload) if query.payload else b"" + data = orjson.loads(payload_bytes) if payload_bytes else {} + 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 = await matched.call( + data, + session=session, + worker_id=worker_id + ) + + # 回复 + 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}") + + def callback(query: zenoh.Query): + # 当前 create_worker 运行在事件循环中,可直接 create_task + asyncio.run_coroutine_threadsafe(_handle_query(query), loop) + + 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 存活直到被取消 + stop_event = asyncio.Event() + 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", +]