flopscope-server 0.3.0__py3-none-any.whl

flopscope_server/__init__.py

@@ -0,0 +1,3 @@
+"""Flopscope backend server — executes numpy operations on behalf of remote clients."""
+
+__version__ = "0.3.0"

flopscope_server/__main__.py

@@ -0,0 +1,42 @@
+"""Entry point for ``python -m flopscope_server``."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from flopscope_server._server import FlopscopeServer
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Flopscope budget-controlled compute server",
+    )
+    parser.add_argument(
+        "--url",
+        default="ipc:///tmp/flopscope.sock",
+        help="ZMQ endpoint to bind (default: ipc:///tmp/flopscope.sock)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=60.0,
+        help="Session inactivity timeout in seconds (default: 60.0)",
+    )
+    args = parser.parse_args()
+
+    print(
+        f"[flopscope-server] binding to {args.url}  (timeout={args.timeout}s)",
+        file=sys.stderr,
+    )
+
+    server = FlopscopeServer(url=args.url, session_timeout_s=args.timeout)
+    try:
+        server.run()
+    except KeyboardInterrupt:
+        print("\n[flopscope-server] shutting down", file=sys.stderr)
+        server.stop()
+
+
+if __name__ == "__main__":
+    main()

flopscope_server/_array_store.py

@@ -0,0 +1,97 @@
+"""ArrayStore — in-process dict-based mapping from handle IDs to numpy arrays."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+
+#: Maximum number of arrays allowed in a single store (configurable via env var).
+MAX_ARRAY_COUNT = int(os.environ.get("FLOPSCOPE_MAX_ARRAY_COUNT", "100000"))
+
+
+class ArrayStore:
+    """Simple store that maps string handle IDs to numpy arrays.
+
+    Handle IDs are monotonically increasing strings of the form "a0", "a1",
+    "a2", … The counter never resets, so IDs remain unique across put/free
+    cycles.
+    """
+
+    def __init__(self) -> None:
+        self._arrays: dict[str, Any] = {}
+        self._counter: int = 0
+
+    # ------------------------------------------------------------------
+    # Core operations
+    # ------------------------------------------------------------------
+
+    def put(self, arr: Any) -> str:
+        """Store *arr* and return its handle ID.
+
+        Raises
+        ------
+        MemoryError
+            If the store already contains :data:`MAX_ARRAY_COUNT` arrays.
+        """
+        if len(self._arrays) >= MAX_ARRAY_COUNT:
+            raise MemoryError(f"array store limit reached: {MAX_ARRAY_COUNT} arrays")
+        handle = f"a{self._counter}"
+        self._arrays[handle] = arr
+        self._counter += 1
+        return handle
+
+    def get(self, handle: str) -> Any:
+        """Return the array for *handle*.
+
+        Raises
+        ------
+        KeyError
+            If *handle* is not in the store.
+        """
+        if handle not in self._arrays:
+            raise KeyError(f"Array handle {handle!r} not found in store")
+        return self._arrays[handle]
+
+    def metadata(self, handle: str) -> dict:
+        """Return metadata dict for *handle*.
+
+        Returns
+        -------
+        dict
+            ``{"id": handle, "shape": list[int], "dtype": str}``
+
+        Raises
+        ------
+        KeyError
+            If *handle* is not in the store.
+        """
+        arr = self.get(handle)  # propagates KeyError with helpful message
+        meta = {
+            "id": handle,
+            "shape": list(arr.shape),
+            "dtype": str(arr.dtype),
+        }
+        symmetry = getattr(arr, "symmetry", None)
+        if symmetry is not None:
+            meta["symmetry"] = symmetry.to_payload()
+        return meta
+
+    def free(self, handles: list[str]) -> None:
+        """Remove arrays by handle; silently ignore unknown handles."""
+        for handle in handles:
+            self._arrays.pop(handle, None)
+
+    def clear(self) -> None:
+        """Remove all arrays from the store."""
+        self._arrays.clear()
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def count(self) -> int:
+        """Number of arrays currently in the store."""
+        return len(self._arrays)

flopscope_server/_comms_tracker.py

@@ -0,0 +1,77 @@
+"""CommsTracker — accumulates per-request timing and byte counts for a session."""
+
+from __future__ import annotations
+
+
+class CommsTracker:
+    """Accumulates communication and compute statistics across requests in a session."""
+
+    def __init__(self) -> None:
+        self._request_count: int = 0
+        self._fetch_count: int = 0
+        self._total_bytes_sent: int = 0
+        self._total_bytes_received: int = 0
+        self._total_comms_overhead_ns: int = 0
+        self._total_compute_time_ns: int = 0
+
+    def record_request(
+        self,
+        *,
+        bytes_received: int,
+        bytes_sent: int,
+        comms_overhead_ns: int,
+        compute_time_ns: int,
+        is_fetch: bool,
+    ) -> None:
+        """Accumulate statistics for a single request.
+
+        Parameters
+        ----------
+        bytes_received:
+            Number of bytes received in this request.
+        bytes_sent:
+            Number of bytes sent in this request.
+        comms_overhead_ns:
+            Communications overhead for this request, in nanoseconds.
+        compute_time_ns:
+            Compute time for this request, in nanoseconds.
+        is_fetch:
+            Whether this request is a fetch (array retrieval) request.
+        """
+        self._request_count += 1
+        if is_fetch:
+            self._fetch_count += 1
+        self._total_bytes_sent += bytes_sent
+        self._total_bytes_received += bytes_received
+        self._total_comms_overhead_ns += comms_overhead_ns
+        self._total_compute_time_ns += compute_time_ns
+
+    def summary(self) -> dict:
+        """Return a summary of accumulated statistics.
+
+        Returns
+        -------
+        dict with keys:
+            request_count: total number of requests recorded
+            fetch_count: number of fetch requests recorded
+            total_bytes_sent: total bytes sent across all requests
+            total_bytes_received: total bytes received across all requests
+            total_comms_overhead_ns: total communications overhead in nanoseconds
+            total_compute_time_ns: total compute time in nanoseconds
+            overhead_ratio: comms / (comms + compute); 0.0 if total is 0
+        """
+        total_ns = self._total_comms_overhead_ns + self._total_compute_time_ns
+        if total_ns == 0:
+            overhead_ratio = 0.0
+        else:
+            overhead_ratio = self._total_comms_overhead_ns / total_ns
+
+        return {
+            "request_count": self._request_count,
+            "fetch_count": self._fetch_count,
+            "total_bytes_sent": self._total_bytes_sent,
+            "total_bytes_received": self._total_bytes_received,
+            "total_comms_overhead_ns": self._total_comms_overhead_ns,
+            "total_compute_time_ns": self._total_compute_time_ns,
+            "overhead_ratio": overhead_ratio,
+        }

flopscope_server/_protocol.py

@@ -0,0 +1,233 @@
+"""Protocol layer for flopscope server — message encoding/decoding.
+
+Message format: msgpack over ZMQ.
+
+All messages are msgpack-encoded dicts. By default, msgpack returns bytes keys
+when raw=False is not used, so decode_request handles normalisation carefully:
+  - top-level string fields (op, dtype, request_id) are decoded to str
+  - binary payload fields (data) are kept as bytes
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import msgpack
+
+from flopscope._registry import REGISTRY
+
+# ---------------------------------------------------------------------------
+# Whitelist
+# ---------------------------------------------------------------------------
+
+#: Protocol-level ops not in the numpy REGISTRY.
+_PROTOCOL_OPS: frozenset[str] = frozenset(
+    {
+        "hello",
+        "budget_open",
+        "budget_close",
+        "budget_status",
+        "fetch",
+        "fetch_slice",
+        "free",
+        "create_from_data",
+        "__getitem__",
+        "astype",
+    }
+)
+
+#: Full set of permitted op names.
+WHITELIST: frozenset[str] = frozenset(REGISTRY.keys()) | _PROTOCOL_OPS
+
+# ---------------------------------------------------------------------------
+# String fields that should always be decoded from bytes -> str
+# ---------------------------------------------------------------------------
+
+_STRING_FIELDS: frozenset[str] = frozenset({"op", "dtype", "request_id"})
+
+
+# ---------------------------------------------------------------------------
+# Exceptions
+# ---------------------------------------------------------------------------
+
+
+class InvalidRequestError(Exception):
+    """Raised when a client request cannot be decoded or is not permitted."""
+
+
+# ---------------------------------------------------------------------------
+# decode_request
+# ---------------------------------------------------------------------------
+
+
+def decode_request(raw: bytes) -> dict:
+    """Decode a msgpack-encoded request from the client.
+
+    Parameters
+    ----------
+    raw:
+        Raw msgpack bytes received from the client.
+
+    Returns
+    -------
+    dict
+        Decoded request with top-level keys as strings.
+        Binary payload fields (e.g. ``data``) are kept as :class:`bytes`.
+
+    Raises
+    ------
+    InvalidRequestError
+        If *raw* is empty, cannot be unpacked, or is missing the ``op`` field.
+    """
+    if not raw:
+        raise InvalidRequestError("malformed request: empty bytes")
+
+    try:
+        # Use raw=True so that we get bytes keys and can selectively decode.
+        msg_raw = msgpack.unpackb(raw, raw=True)
+    except Exception as exc:
+        raise InvalidRequestError(f"malformed request: {exc}") from exc
+
+    if not isinstance(msg_raw, dict):
+        raise InvalidRequestError("malformed request: top-level value must be a dict")
+
+    # Normalise keys and selected string values.
+    msg: dict[str, Any] = {}
+    for k, v in msg_raw.items():
+        # Decode key from bytes -> str
+        key: str = k.decode("utf-8") if isinstance(k, bytes) else str(k)
+
+        # Selectively decode known string-valued fields
+        if key in _STRING_FIELDS and isinstance(v, bytes):
+            value: Any = v.decode("utf-8")
+        else:
+            value = v
+
+        msg[key] = value
+
+    if "op" not in msg:
+        raise InvalidRequestError("malformed request: missing 'op' field")
+
+    return msg
+
+
+# ---------------------------------------------------------------------------
+# validate_request
+# ---------------------------------------------------------------------------
+
+
+def validate_request(msg: dict) -> None:
+    """Check that the op in *msg* is on the permitted whitelist.
+
+    Parameters
+    ----------
+    msg:
+        Decoded request dict (as returned by :func:`decode_request`).
+
+    Raises
+    ------
+    InvalidRequestError
+        If the op name is not in :data:`WHITELIST`.
+    """
+    op = msg.get("op", "")
+    if op not in WHITELIST:
+        raise InvalidRequestError(f"unknown op: {op!r}")
+
+
+# ---------------------------------------------------------------------------
+# encode_response
+# ---------------------------------------------------------------------------
+
+
+def encode_response(result: Any, budget: int, comms_overhead_ns: int) -> bytes:
+    """Encode a successful operation response.
+
+    Parameters
+    ----------
+    result:
+        The value returned by the operation (must be msgpack-serialisable).
+    budget:
+        Remaining budget after the operation.
+    comms_overhead_ns:
+        Round-trip communications overhead in nanoseconds.
+
+    Returns
+    -------
+    bytes
+        msgpack-encoded response dict with ``status="ok"``.
+    """
+    payload = {
+        "status": "ok",
+        "result": result,
+        "budget": budget,
+        "comms_overhead_ns": comms_overhead_ns,
+    }
+    return msgpack.packb(payload, use_bin_type=True)
+
+
+# ---------------------------------------------------------------------------
+# encode_error_response
+# ---------------------------------------------------------------------------
+
+
+def encode_error_response(error_type: str, message: str) -> bytes:
+    """Encode an error response.
+
+    Parameters
+    ----------
+    error_type:
+        Name of the exception class (e.g. ``"InvalidRequestError"``).
+    message:
+        Human-readable error description.
+
+    Returns
+    -------
+    bytes
+        msgpack-encoded response dict with ``status="error"``.
+    """
+    payload = {
+        "status": "error",
+        "error_type": error_type,
+        "message": message,
+    }
+    return msgpack.packb(payload, use_bin_type=True)
+
+
+# ---------------------------------------------------------------------------
+# encode_fetch_response
+# ---------------------------------------------------------------------------
+
+
+def encode_fetch_response(
+    data: bytes,
+    shape: tuple[int, ...],
+    dtype: str,
+    comms_overhead_ns: int,
+) -> bytes:
+    """Encode a fetch response carrying raw array bytes.
+
+    Parameters
+    ----------
+    data:
+        Raw array bytes (e.g. ``array.tobytes()``).
+    shape:
+        Array shape as a tuple/list of ints.
+    dtype:
+        NumPy dtype string (e.g. ``"float32"``).
+    comms_overhead_ns:
+        Round-trip communications overhead in nanoseconds.
+
+    Returns
+    -------
+    bytes
+        msgpack-encoded response dict with ``status="ok"`` and ``data`` as
+        raw bytes.
+    """
+    payload = {
+        "status": "ok",
+        "data": data,
+        "shape": list(shape),
+        "dtype": dtype,
+        "comms_overhead_ns": comms_overhead_ns,
+    }
+    return msgpack.packb(payload, use_bin_type=True)