279 lines
9.5 KiB
Python
279 lines
9.5 KiB
Python
"""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())
|