[feat] add Serializer for rpc server

Author: CormickKneey

areal/scheduler/rpc/app.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import argparse
+import json
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from flask import Flask
+from flask_jsonrpc import JSONRPC
+
+from areal.api.cli_args import NameResolveConfig
+from areal.utils import logging, name_resolve
+from areal.utils.dynamic_import import import_from_string
+from areal.utils.seeding import set_random_seed
+from areal.utils.stats_tracker import export_all
+
+logger = logging.getLogger(__name__)
+
+
+def _ensure_jsonable(value: Any) -> Any:
+    """Convert Python objects into JSON-serialisable structures."""
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    if isinstance(value, Mapping):
+        return {str(k): _ensure_jsonable(v) for k, v in value.items()}
+    if isinstance(value, Iterable) and not isinstance(value, (bytes, bytearray)):
+        return [_ensure_jsonable(v) for v in value]
+    try:
+        json.dumps(value)
+    except TypeError:
+        return repr(value)
+    return value
+
+
+class EngineRPCServer:
+    engine = None
+
+    def create_engine(
+        self,
+        payload: dict[str, Any],
+    ) -> dict[str, Any]:
+        if not isinstance(payload, Mapping):
+            raise ValueError("payload must be a mapping object")
+
+        cls_path = payload.get("engine_class")
+        if not cls_path:
+            raise ValueError("engine_class is required")
+
+        cls_args = payload.get("engine_args", [])
+        cls_kwargs = payload.get("engine_kwargs", {})
+        init_args = payload.get("initialize_args", [])
+        init_kwargs = payload.get("initialize_kwargs", {})
+
+        try:
+            engine_cls = import_from_string(cls_path)
+        except (ImportError, AttributeError, ValueError) as exc:
+            raise ValueError(
+                f"Failed to import engine class '{cls_path}': {exc}"
+            ) from exc
+
+        try:
+            engine = engine_cls(*cls_args, **cls_kwargs)
+        except Exception as exc:  # pylint: disable=broad-except
+            raise RuntimeError(f"Failed to instantiate engine: {exc}") from exc
+
+        try:
+            init_result = engine.initialize(*init_args, **init_kwargs)
+        except Exception as exc:  # pylint: disable=broad-except
+            raise RuntimeError(f"Engine initialise failed: {exc}") from exc
+
+        self.engine = engine
+        return {
+            "status": "initialized",
+            "initialize_result": _ensure_jsonable(init_result),
+        }
+
+    def call_engine(self, payload: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(payload, Mapping):
+            raise ValueError("payload must be a mapping object")
+
+        method_name = payload.get("method")
+        if method_name is None:
+            raise ValueError("method is required")
+
+        if self.engine is None:
+            raise RuntimeError("Engine not found")
+
+        engine = self.engine
+        args = payload.get("args", [])
+        kwargs = payload.get("kwargs", {})
+
+        try:
+            target = getattr(engine, method_name)
+        except AttributeError as exc:
+            raise ValueError(f"Engine has no method '{method_name}'") from exc
+
+        try:
+            result = target(*args, **kwargs)
+        except Exception as exc:  # pylint: disable=broad-except
+            raise RuntimeError(f"Engine call failed: {exc}") from exc
+
+        return {
+            "method": method_name,
+            "result": _ensure_jsonable(result),
+        }
+
+    def configure(self, payload: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(payload, Mapping):
+            raise ValueError("payload must be a mapping object")
+
+        seed_cfg = payload.get("seed")
+        if seed_cfg:
+            base_seed = seed_cfg.get("base_seed")
+            key = seed_cfg.get("key", "default")
+            if base_seed is None:
+                raise ValueError("seed.base_seed is required")
+            set_random_seed(int(base_seed), str(key))
+            logger.info("Random seed configured: base=%s key=%s", base_seed, key)
+
+        name_resolve_cfg = payload.get("name_resolve")
+        if name_resolve_cfg:
+            try:
+                config = NameResolveConfig(**name_resolve_cfg)
+            except TypeError as exc:
+                raise ValueError(f"Invalid name_resolve payload: {exc}") from exc
+            name_resolve.reconfigure(config)
+            logger.info("Name resolver reconfigured: type=%s", config.type)
+
+        return {"status": "configured"}
+
+    def export_stats(self, reset: bool = True) -> dict[str, Any]:
+        try:
+            stats = export_all(reset=reset)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.exception("Failed to export stats")
+            raise RuntimeError(f"Failed to export stats: {exc}") from exc
+
+        return {"stats": _ensure_jsonable(stats), "reset": reset}
+
+    def health(self) -> dict[str, str]:
+        return {"status": "ok"}
+
+
+def create_app():
+    app = Flask(__name__)
+    engine_rpc_server = EngineRPCServer()
+    jsonrpc = JSONRPC(app, "/api", enable_web_browsable_api=False)
+
+    @jsonrpc.method("App.create_engine")
+    def rpc_create_engine(payload: dict[str, Any]):  # type: ignore[override]
+        return engine_rpc_server.create_engine(payload)
+
+    @jsonrpc.method("App.call_engine")
+    def rpc_call_engine(payload: dict[str, Any]):  # type: ignore[override]
+        return engine_rpc_server.call_engine(payload)
+
+    @jsonrpc.method("App.configure")
+    def rpc_configure(payload: dict[str, Any]):  # type: ignore[override]
+        return engine_rpc_server.configure(payload)
+
+    @jsonrpc.method("App.health")
+    def rpc_health():  # type: ignore[override]
+        return engine_rpc_server.health()
+
+    @jsonrpc.method("App.export_stats")
+    def rpc_export_stats(reset: bool = True):  # type: ignore[override]
+        return engine_rpc_server.export_stats(reset)
+
+    return app
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", default=5000, type=int, required=False)
+    args = parser.parse_args()
+
+    app = create_app()
+    app.run(host=args.host, port=args.port)