"""Isolated stdin/stdout worker for the CP-SAT assignment protocol.""" from __future__ import annotations import contextlib import hashlib import importlib.metadata import importlib.util import io import json import os import platform import site import sys import traceback from collections.abc import Mapping from pathlib import Path from typing import Any # ``python -I /abs/path/solver_worker.py`` does not trust cwd/PYTHONPATH. _REPO_ROOT = Path(__file__).resolve().parents[2] if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) from server.engines.solver_process import ( ERROR_MARKER, PROTOCOL_VERSION, SUCCESS_MARKER, _canonical_json, _json_value, _request_digest, ) _MAX_REQUEST_BYTES = 64 * 1024 * 1024 _DIAGNOSTIC_LIMIT = 8_000 _REQUIRED_NATIVE_PACKAGES = ("ortools", "numpy", "pandas") def _resolved(path: str | os.PathLike[str]) -> Path: return Path(path).expanduser().resolve() def _under(path: Path, root: Path) -> bool: try: path.relative_to(root) return True except ValueError: return False def _package_identity(name: str, allowed_roots: list[Path]) -> tuple[dict[str, Any], list[str]]: reasons: list[str] = [] spec = importlib.util.find_spec(name) origin = None if spec is None else spec.origin try: version = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: version = None identity: dict[str, Any] = {"version": version, "origin": origin} if spec is None or not origin: reasons.append(f"{name}:missing") return identity, reasons origin_path = _resolved(origin) if not any(_under(origin_path, root) for root in allowed_roots): reasons.append(f"{name}:outside-runtime-prefix") user_site = site.getusersitepackages() if user_site and _under(origin_path, _resolved(user_site)): reasons.append(f"{name}:user-site") return identity, reasons def runtime_identity() -> dict[str, Any]: frozen = bool(getattr(sys, "frozen", False)) roots = [_resolved(sys.prefix), _resolved(sys.exec_prefix)] meipass = getattr(sys, "_MEIPASS", None) if meipass: roots.append(_resolved(meipass)) roots = list(dict.fromkeys(roots)) reasons: list[str] = [] if platform.python_implementation() != "CPython": reasons.append("implementation:not-cpython") if not frozen and not bool(sys.flags.isolated): reasons.append("python:not-isolated") if bool(site.ENABLE_USER_SITE): reasons.append("python:user-site-enabled") packages: dict[str, Any] = {} for package in _REQUIRED_NATIVE_PACKAGES: packages[package], package_reasons = _package_identity(package, roots) reasons.extend(package_reasons) return { "implementation": platform.python_implementation(), "pythonVersion": platform.python_version(), "executable": str(_resolved(sys.executable)), "baseExecutable": str(_resolved(getattr(sys, "_base_executable", sys.executable))), "prefix": str(_resolved(sys.prefix)), "basePrefix": str(_resolved(sys.base_prefix)), "isolated": bool(sys.flags.isolated), "noUserSite": bool(sys.flags.no_user_site), "enableUserSite": bool(site.ENABLE_USER_SITE), "frozen": frozen, "allowedRoots": [str(root) for root in roots], "packages": packages, "safe": not reasons, "reasons": sorted(set(reasons)), } def _trim(value: str) -> str: return value if len(value) <= _DIAGNOSTIC_LIMIT else value[:_DIAGNOSTIC_LIMIT] + "\n...[truncated]" def _emit(payload: Mapping[str, Any]) -> None: sys.stdout.write(_canonical_json(_json_value(payload))) sys.stdout.flush() def _emit_error( code: str, message: str, *, request_id: str = "", details: Mapping[str, Any] | None = None, ) -> None: _emit( { "protocolVersion": PROTOCOL_VERSION, "marker": ERROR_MARKER, "ok": False, "requestId": request_id, "error": {"code": code, "message": message, "details": dict(details or {})}, } ) def _read_request() -> dict[str, Any]: raw = sys.stdin.buffer.read(_MAX_REQUEST_BYTES + 1) if len(raw) > _MAX_REQUEST_BYTES: raise ValueError("request exceeds size limit") try: payload = json.loads(raw.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ValueError(f"invalid request JSON: {exc}") from exc if not isinstance(payload, dict): raise TypeError("request must be an object") return payload def _validate_request(request: Mapping[str, Any]) -> tuple[str, dict[str, Any], list[dict], dict, list[dict] | None, str]: request_id = request.get("requestId") if not isinstance(request_id, str) or not request_id: raise ValueError("requestId is required") if request.get("protocolVersion") != PROTOCOL_VERSION: raise RuntimeError("protocol mismatch") body = {key: value for key, value in request.items() if key != "requestId"} if _request_digest(body) != request_id: raise ValueError("requestId digest mismatch") if request.get("operation") != "optimize_line_assignment": raise ValueError("unsupported operation") world = request.get("world") entries = request.get("entries") params = request.get("params") warm_start = request.get("warmStart") pipeline_label = request.get("pipelineLabel") if not isinstance(world, dict): raise TypeError("world must be an object") if not isinstance(entries, list) or not all(isinstance(item, dict) for item in entries): raise TypeError("entries must be an object array") if not isinstance(params, dict): raise TypeError("params must be an object") if warm_start is not None and ( not isinstance(warm_start, list) or not all(isinstance(item, dict) for item in warm_start) ): raise TypeError("warmStart must be an object array") if not isinstance(pipeline_label, str) or not pipeline_label.strip(): raise ValueError("pipelineLabel is required") return request_id, world, entries, params, warm_start, pipeline_label def main() -> int: request_id = "" try: request = _read_request() request_id = str(request.get("requestId") or "") except ValueError as exc: _emit_error("SOLVER_REQUEST_INVALID", "求解请求非法", details={"error": str(exc)}) return 0 if request.get("protocolVersion") != PROTOCOL_VERSION: _emit_error( "SOLVER_PROTOCOL_MISMATCH", "求解请求协议版本不一致", request_id=request_id, details={"expected": PROTOCOL_VERSION, "actual": request.get("protocolVersion")}, ) return 0 try: request_id, world, entries, params_payload, warm_start, pipeline_label = _validate_request(request) except RuntimeError: _emit_error( "SOLVER_PROTOCOL_MISMATCH", "求解请求协议版本不一致", request_id=request_id, ) return 0 except (TypeError, ValueError) as exc: _emit_error( "SOLVER_REQUEST_INVALID", "求解请求非法", request_id=request_id, details={"error": str(exc)}, ) return 0 identity = runtime_identity() if identity.get("safe") is not True: _emit_error( "SOLVER_RUNTIME_UNSAFE", "求解子进程运行时身份不安全", request_id=request_id, details={"runtimeIdentity": identity}, ) return 0 captured_stdout = io.StringIO() captured_stderr = io.StringIO() try: with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr): from server.engines.base import EngineParams from server.engines.cp_engine import optimize_line_assignment params = EngineParams.model_validate(params_payload) ordered, solver_meta = optimize_line_assignment( world, entries, params, warm_start=warm_start, pipeline_label=pipeline_label, ) if not isinstance(ordered, list) or not all(isinstance(item, dict) for item in ordered): raise TypeError("optimizer entries result is not an object array") if not isinstance(solver_meta, dict): raise TypeError("optimizer solverMeta result is not an object") except Exception as exc: # noqa: BLE001 - child must translate every Python-level failure _emit_error( "SOLVER_EXECUTION_FAILED", "CP-SAT 子进程执行失败", request_id=request_id, details={ "exceptionType": type(exc).__name__, "error": str(exc), "stdout": _trim(captured_stdout.getvalue()), "stderr": _trim(captured_stderr.getvalue()), "traceback": _trim(traceback.format_exc()), "runtimeIdentity": identity, }, ) return 0 response_body = { "protocolVersion": PROTOCOL_VERSION, "marker": SUCCESS_MARKER, "ok": True, "requestId": request_id, "runtimeIdentity": identity, "result": {"entries": ordered, "solverMeta": solver_meta}, } response_body["responseDigest"] = hashlib.sha256( _canonical_json(_json_value(response_body)).encode("utf-8") ).hexdigest() _emit(response_body) return 0 if __name__ == "__main__": raise SystemExit(main())