flopscope-server 0.3.0__py3-none-any.whl

flopscope_server/_server.py

@@ -0,0 +1,375 @@
+"""FlopscopeServer -- ZMQ REP loop that dispatches requests to a Session + RequestHandler."""
+
+from __future__ import annotations
+
+import sys
+from time import monotonic, perf_counter_ns
+
+import msgpack
+import zmq
+
+import flopscope
+from flopscope_server._protocol import (
+    InvalidRequestError,
+    decode_request,
+    encode_error_response,
+    encode_response,
+    validate_request,
+)
+from flopscope_server._request_handler import RequestHandler
+from flopscope_server._session import Session
+
+
+class FlopscopeServer:
+    """ZMQ REP server for the flopscope budget-controlled compute service.
+
+    Parameters
+    ----------
+    url:
+        ZMQ endpoint to bind (e.g. ``"ipc:///tmp/flopscope.sock"`` or
+        ``"tcp://127.0.0.1:15555"``).
+    session_timeout_s:
+        Seconds of inactivity after which an open session is reaped.
+    """
+
+    def __init__(
+        self,
+        url: str = "ipc:///tmp/flopscope.sock",
+        session_timeout_s: float = 60.0,
+    ) -> None:
+        self._url = url
+        self._session_timeout_s = session_timeout_s
+        self._running = False
+
+        # Session state (at most one session at a time)
+        self._session: Session | None = None
+        self._handler: RequestHandler | None = None
+        self._last_activity: float = 0.0
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def run(self) -> None:
+        """Start the blocking ZMQ REP loop.
+
+        Creates a ZMQ Context and REP socket, binds to ``self._url``, and
+        loops until :meth:`stop` is called.
+        """
+        ctx = zmq.Context()
+        sock = ctx.socket(zmq.REP)
+        sock.setsockopt(zmq.RCVTIMEO, 1000)  # 1 s poll interval
+        sock.bind(self._url)
+        self._running = True
+
+        try:
+            while self._running:
+                try:
+                    raw = sock.recv()
+                except zmq.Again:
+                    self._reap_session_if_stale()
+                    continue
+
+                response = self._process_request(raw)
+                sock.send(response)
+        finally:
+            sock.close(linger=0)
+            ctx.term()
+
+    def stop(self) -> None:
+        """Signal the run loop to exit."""
+        self._running = False
+
+    # ------------------------------------------------------------------
+    # Request processing
+    # ------------------------------------------------------------------
+
+    def _process_request(self, raw: bytes) -> bytes:
+        """Decode, dispatch, encode a single request.
+
+        Returns msgpack bytes ready to send back to the client.
+        """
+        t0 = perf_counter_ns()
+
+        # --- Decode ---
+        try:
+            msg = decode_request(raw)
+        except InvalidRequestError as exc:
+            return encode_error_response("InvalidRequestError", str(exc))
+
+        # Normalise remaining bytes→str fields that decode_request leaves as
+        # raw bytes (it only converts _STRING_FIELDS: op, dtype, request_id).
+        _normalize_msg(msg)
+
+        t1 = perf_counter_ns()
+        op = msg["op"]
+
+        # --- Pre-session ops (no active session required) ---
+        if op == "hello":
+            return self._handle_hello(msg)
+
+        # --- Session lifecycle ops ---
+        if op == "budget_open":
+            return self._handle_budget_open(msg, t0, t1)
+
+        if op == "budget_close":
+            return self._handle_budget_close(t0, t1)
+
+        # --- Require active session for everything else ---
+        if self._session is None:
+            return encode_error_response(
+                "NoBudgetContextError",
+                "no active session -- send budget_open first",
+            )
+
+        # --- Validate (whitelist check) ---
+        try:
+            validate_request(msg)
+        except InvalidRequestError as exc:
+            return encode_error_response("InvalidRequestError", str(exc))
+
+        # --- Dispatch ---
+        assert self._handler is not None
+        result = self._handler.handle(msg)
+
+        t2 = perf_counter_ns()
+
+        # --- Encode response ---
+        is_fetch = op in ("fetch", "fetch_slice")
+        response_bytes: bytes
+
+        try:
+            if result.get("status") == "error":
+                response_bytes = encode_error_response(
+                    result["error_type"], result["message"]
+                )
+            elif "data" in result:
+                # fetch / fetch_slice -- use raw response packing
+                payload = {
+                    "status": "ok",
+                    "data": result["data"],
+                    "shape": result["shape"],
+                    "dtype": result["dtype"],
+                    "comms_overhead_ns": 0,  # placeholder, updated below
+                }
+                response_bytes = msgpack.packb(payload, use_bin_type=True)
+            else:
+                response_bytes = encode_response(
+                    result.get("result"),
+                    result.get(
+                        "budget", self._session.budget_remaining if self._session else 0
+                    ),
+                    comms_overhead_ns=0,  # placeholder
+                )
+        except Exception as enc_err:
+            # Response serialization failed -- return error instead of crashing
+            response_bytes = encode_error_response(
+                "FlopscopeServerError",
+                f"failed to serialize response: {type(enc_err).__name__}",
+            )
+
+        t3 = perf_counter_ns()
+
+        # --- Record comms overhead ---
+        comms_ns = (t1 - t0) + (t3 - t2)
+        compute_ns = t2 - t1
+
+        self._session.comms_tracker.record_request(
+            bytes_received=len(raw),
+            bytes_sent=len(response_bytes),
+            comms_overhead_ns=comms_ns,
+            compute_time_ns=compute_ns,
+            is_fetch=is_fetch,
+        )
+
+        self._last_activity = monotonic()
+        return response_bytes
+
+    # ------------------------------------------------------------------
+    # Budget lifecycle handlers
+    # ------------------------------------------------------------------
+
+    def _handle_budget_open(self, msg: dict, t0: int, t1: int) -> bytes:
+        """Open a new session; error if one is already open."""
+        if self._session is not None and self._session.is_open:
+            return encode_error_response(
+                "RuntimeError",
+                "session already open -- send budget_close first",
+            )
+
+        # Support both top-level and kwargs-based flop_budget
+        flop_budget = msg.get("flop_budget")
+        flop_multiplier = msg.get("flop_multiplier")
+        if flop_budget is None:
+            kwargs = msg.get("kwargs") or {}
+            flop_budget = kwargs.get("flop_budget", 1_000_000)
+            if flop_multiplier is None:
+                flop_multiplier = kwargs.get("flop_multiplier", 1.0)
+        if flop_multiplier is None:
+            flop_multiplier = 1.0
+        self._session = Session(
+            flop_budget=flop_budget, flop_multiplier=flop_multiplier
+        )
+        self._handler = RequestHandler(self._session)
+        self._last_activity = monotonic()
+
+        t2 = perf_counter_ns()
+        response_bytes = encode_response(
+            {"session": "opened", "flop_budget": flop_budget},
+            budget=flop_budget,
+            comms_overhead_ns=0,
+        )
+        t3 = perf_counter_ns()
+
+        comms_ns = (t1 - t0) + (t3 - t2)
+        compute_ns = t2 - t1
+        self._session.comms_tracker.record_request(
+            bytes_received=0,
+            bytes_sent=len(response_bytes),
+            comms_overhead_ns=comms_ns,
+            compute_time_ns=compute_ns,
+            is_fetch=False,
+        )
+        return response_bytes
+
+    def _handle_budget_close(self, t0: int, t1: int) -> bytes:
+        """Close the active session and return a summary."""
+        if self._session is None or not self._session.is_open:
+            return encode_error_response(
+                "NoBudgetContextError",
+                "no active session to close",
+            )
+
+        summary = self._session.close()
+        self._session = None
+        self._handler = None
+
+        t2 = perf_counter_ns()
+        response_bytes = encode_response(summary, budget=0, comms_overhead_ns=0)
+        t3 = perf_counter_ns()
+        # No session to record to (already closed); that's fine.
+        return response_bytes
+
+    # ------------------------------------------------------------------
+    # Handshake
+    # ------------------------------------------------------------------
+
+    def _handle_hello(self, msg: dict) -> bytes:
+        """Compare client and server flopscope versions.
+
+        Returns an ok response carrying the server's leading X.Y.Z when
+        the client reports a matching version, or a ``VersionMismatch``
+        error otherwise. The handshake is intentionally session-free so
+        clients can detect mismatch before opening a budget.
+        """
+        kwargs = msg.get("kwargs") or {}
+        client_version = (
+            kwargs.get("client_version") if isinstance(kwargs, dict) else None
+        )
+        server_xyz = flopscope.__version__.split("+", 1)[0]
+
+        if not isinstance(client_version, str) or not client_version:
+            return encode_error_response(
+                "VersionMismatch",
+                f"client did not report a flopscope version; server is {server_xyz}",
+            )
+
+        if client_version != server_xyz:
+            return encode_error_response(
+                "VersionMismatch",
+                (
+                    f"flopscope-client {client_version} cannot talk to "
+                    f"flopscope-server {server_xyz}: versions must match. "
+                    "Install matched versions: see "
+                    "https://pypi.org/project/flopscope/ for the latest."
+                ),
+            )
+
+        return msgpack.packb(
+            {"status": "ok", "server_version": server_xyz},
+            use_bin_type=True,
+        )
+
+    # ------------------------------------------------------------------
+    # Session reaping
+    # ------------------------------------------------------------------
+
+    def _reap_session_if_stale(self) -> None:
+        """Close and discard the session if it has been idle too long."""
+        if self._session is None or not self._session.is_open:
+            return
+        if monotonic() - self._last_activity > self._session_timeout_s:
+            print(
+                "[flopscope-server] session timed out -- reaping",
+                file=sys.stderr,
+            )
+            self._session.close()
+            self._session = None
+            self._handler = None
+
+
+# ---------------------------------------------------------------------------
+# Message normalisation helper
+# ---------------------------------------------------------------------------
+
+
+def _decode_if_bytes(v: object) -> object:
+    """Decode bytes to str if possible, otherwise return as-is."""
+    if isinstance(v, bytes):
+        try:
+            return v.decode("utf-8")
+        except UnicodeDecodeError:
+            return v
+    return v
+
+
+def _normalize_arg(a: object) -> object:
+    """Normalize a single arg: decode short bytes to str, normalize dict keys/values.
+
+    Binary data payloads (array bytes) must stay as bytes.  Heuristic: only
+    decode if short AND contains no null bytes (handles, dtype strings, etc.
+    are always short ASCII).
+    """
+    if isinstance(a, bytes):
+        # Only decode bytes that are short AND contain only ASCII printable
+        # characters (no bytes > 127).  This catches handle IDs and dtype
+        # strings but never touches binary array data (which often contains
+        # high bytes even if it happens to be valid UTF-8).
+        if len(a) > 0 and len(a) <= 32 and all(32 <= b < 128 for b in a):
+            return a.decode("ascii")
+        return a  # keep as raw bytes (likely binary payload)
+    if isinstance(a, dict):
+        return {_decode_if_bytes(k): _normalize_arg(v) for k, v in a.items()}
+    if isinstance(a, list):
+        return [_normalize_arg(x) for x in a]
+    return a
+
+
+def _normalize_msg(msg: dict) -> None:
+    """In-place normalise a decoded request dict.
+
+    ``decode_request`` only converts a small set of known string fields
+    (op, dtype, request_id) from bytes to str.  Other string-valued
+    fields -- ``id``, ``ids``, and items inside ``args`` -- may still be
+    bytes after msgpack decoding with ``raw=True``.  This helper converts
+    them so the downstream RequestHandler receives plain strings.
+    """
+    # id (used by fetch, fetch_slice, free)
+    if "id" in msg:
+        msg["id"] = _decode_if_bytes(msg["id"])
+
+    # ids (used by free)
+    if "ids" in msg and isinstance(msg["ids"], list):
+        msg["ids"] = [_decode_if_bytes(x) for x in msg["ids"]]
+
+    # args -- handle IDs are short ASCII strings, may also contain dicts
+    if "args" in msg and isinstance(msg["args"], list):
+        msg["args"] = [_normalize_arg(a) for a in msg["args"]]
+
+    # kwargs keys are already str (decode_request normalises all keys).
+    # kwargs *values* may contain handles, dicts, or lists that need full
+    # normalization (not just _decode_if_bytes).
+    if "kwargs" in msg and isinstance(msg["kwargs"], dict):
+        msg["kwargs"] = {
+            _decode_if_bytes(k): _normalize_arg(v) for k, v in msg["kwargs"].items()
+        }

flopscope_server/_session.py

@@ -0,0 +1,171 @@
+"""Session — ties together ArrayStore, BudgetContext, and CommsTracker for a single participant session."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+import flopscope as flops
+from flopscope_server._array_store import ArrayStore
+from flopscope_server._comms_tracker import CommsTracker
+
+
+class Session:
+    """A single participant session combining ArrayStore, BudgetContext, and CommsTracker.
+
+    Parameters
+    ----------
+    flop_budget : int
+        Maximum number of FLOPs allowed for this session.
+    flop_multiplier : float
+        Multiplier applied to each operation's raw FLOP cost.
+    """
+
+    def __init__(self, flop_budget: int, flop_multiplier: float = 1.0) -> None:
+        self._store = ArrayStore()
+        self._comms_tracker = CommsTracker()
+        self._budget_ctx = flops.BudgetContext(
+            flop_budget=flop_budget,
+            flop_multiplier=flop_multiplier,
+            quiet=True,
+        )
+        self._budget_ctx.__enter__()
+        self._is_open = True
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def is_open(self) -> bool:
+        """True if this session is still active (not yet closed)."""
+        return self._is_open
+
+    @property
+    def budget_remaining(self) -> int:
+        """FLOPs remaining in the current budget."""
+        return self._budget_ctx.flops_remaining
+
+    @property
+    def budget_context(self) -> flops.BudgetContext:
+        """The active BudgetContext for this session.
+
+        Raises
+        ------
+        RuntimeError
+            If the session is already closed.
+        """
+        if not self._is_open:
+            raise RuntimeError(
+                "Session is closed; BudgetContext is no longer available."
+            )
+        return self._budget_ctx
+
+    @property
+    def comms_tracker(self) -> CommsTracker:
+        """The CommsTracker for this session."""
+        return self._comms_tracker
+
+    # ------------------------------------------------------------------
+    # Array operations (delegate to ArrayStore)
+    # ------------------------------------------------------------------
+
+    def store_array(self, arr: Any) -> str:
+        """Store *arr* and return its handle ID.
+
+        Delegates to :meth:`ArrayStore.put`.
+        """
+        return self._store.put(arr)
+
+    def get_array(self, handle: str) -> Any:
+        """Return the array for *handle*.
+
+        Delegates to :meth:`ArrayStore.get`.
+
+        Raises
+        ------
+        KeyError
+            If *handle* is not in the store.
+        """
+        return self._store.get(handle)
+
+    def array_metadata(self, handle: str) -> dict:
+        """Return metadata dict for *handle*.
+
+        Delegates to :meth:`ArrayStore.metadata`.
+
+        Raises
+        ------
+        KeyError
+            If *handle* is not in the store.
+        """
+        return self._store.metadata(handle)
+
+    def free_arrays(self, handles: list) -> None:
+        """Remove arrays by handle; silently ignore unknown handles.
+
+        Delegates to :meth:`ArrayStore.free`.
+        """
+        self._store.free(handles)
+
+    # ------------------------------------------------------------------
+    # Budget
+    # ------------------------------------------------------------------
+
+    def budget_status(self) -> dict:
+        """Return the current FLOP budget status.
+
+        Returns
+        -------
+        dict with keys:
+            flop_budget: total budget
+            flops_used: FLOPs consumed so far
+            flops_remaining: budget minus used
+        """
+        return {
+            "flop_budget": self._budget_ctx.flop_budget,
+            "flops_used": self._budget_ctx.flops_used,
+            "flops_remaining": self._budget_ctx.flops_remaining,
+        }
+
+    # ------------------------------------------------------------------
+    # Session lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> dict:
+        """Close the session, exiting the BudgetContext and clearing the ArrayStore.
+
+        Returns
+        -------
+        dict with keys:
+            budget_summary: str — human-readable FLOP budget summary,
+                including a namespace section when labeled ops were recorded
+            budget_breakdown: dict — machine-readable summary data with
+                ``by_namespace`` buckets for direct ingestion
+            comms_summary: dict — CommsTracker summary
+
+        Raises
+        ------
+        RuntimeError
+            If the session is already closed.
+        """
+        if not self._is_open:
+            raise RuntimeError("Session is already closed.")
+
+        self._budget_ctx.__exit__(None, None, None)
+        budget_breakdown = self._budget_ctx.summary_dict(by_namespace=True)
+        show_namespaces = any(
+            namespace is not None
+            for namespace in budget_breakdown.get("by_namespace", {})
+        )
+        budget_summary = self._budget_ctx.summary(by_namespace=show_namespaces)
+        comms_summary = self._comms_tracker.summary()
+        self._store.clear()
+        self._is_open = False
+
+        return {
+            "budget_summary": budget_summary,
+            "budget_breakdown": budget_breakdown,
+            "comms_summary": comms_summary,
+        }

flopscope_server-0.3.0.dist-info/METADATA

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: flopscope-server
+Version: 0.3.0
+Summary: Backend server for flopscope client-server architecture
+Requires-Python: >=3.10
+Requires-Dist: flopscope==0.3.0
+Requires-Dist: msgpack>=1.0.0
+Requires-Dist: numpy<2.5.0,>=2.0.0
+Requires-Dist: pyzmq>=26.0.0

flopscope_server-0.3.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+flopscope_server/__init__.py,sha256=AVnTf0x1li_DqBv19AszabChlmxAN-bqCQqJx7H5uIY,113
+flopscope_server/__main__.py,sha256=3dhGzDNLgrni3LsTjYvA5aeNz_Ny2bVXpZ2Wrj-ABr4,1056
+flopscope_server/_array_store.py,sha256=iIoDuBQcdnk3UPUdKBee8mczOWhbo5yuO77k4fPzrj4,2947
+flopscope_server/_comms_tracker.py,sha256=T8HC00xbMbFOhQ_gK9mLyCCJhJJpZMTbPUSmKfgT4uc,2833
+flopscope_server/_protocol.py,sha256=jiwx-lJd9xik0HBR72h_vuUPWkWmgs9m-4xp4SBuL6g,6588
+flopscope_server/_request_handler.py,sha256=UXBwrxvB1zM3hfCGDk4nlaXALA8Rs8Jm6nsvDK5RE0M,17025
+flopscope_server/_server.py,sha256=WZ9bYiJjYFavyaHLfYG9h1BQrFHJTbDDxK46u6RcQME,13466
+flopscope_server/_session.py,sha256=23wf19c9uJs_Kiq9mgBrdHR2W48_3jpWLoQE5pa7YJE,5401
+flopscope_server-0.3.0.dist-info/METADATA,sha256=zya9tXmhH1ie2HlzQrh_KBA3uJbqPTXUpNFwaNmR4uM,275
+flopscope_server-0.3.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+flopscope_server-0.3.0.dist-info/entry_points.txt,sha256=_fVtS1QpIrRT22oEsh9PEuucsyUfH8v2FHctWdYAzU4,68
+flopscope_server-0.3.0.dist-info/RECORD,,

flopscope_server-0.3.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

flopscope_server-0.3.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+flopscope-server = flopscope_server.__main__:main