flopscope-server 0.4.0__tar.gz → 0.4.2__tar.gz

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: flopscope-server
-Version: 0.4.0
+Version: 0.4.2
 Summary: Backend server for flopscope client-server architecture
 License-Expression: MIT
 Requires-Python: >=3.10
-Requires-Dist: flopscope==0.4.0
+Requires-Dist: flopscope==0.4.2
 Requires-Dist: msgpack>=1.0.0
 Requires-Dist: numpy<2.5.0,>=2.0.0
 Requires-Dist: pyzmq>=26.0.0

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flopscope-server"
-version = "0.4.0"
+version = "0.4.2"
 description = "Backend server for flopscope client-server architecture"
 readme = "README.md"
 requires-python = ">=3.10"
 license = "MIT"
 dependencies = [
-    "flopscope==0.4.0",
+    "flopscope==0.4.2",
     "numpy>=2.0.0,<2.5.0",
     "pyzmq>=26.0.0",
     "msgpack>=1.0.0",

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/src/flopscope_server/__init__.py

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

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/src/flopscope_server/_protocol.py

@@ -130,6 +130,10 @@ def validate_request(msg: dict) -> None:
         If the op name is not in :data:`WHITELIST`.
     """
     op = msg.get("op", "")
+    # Generator.<method> ops are dispatched to a server-side RNG handle and are
+    # method-whitelisted in the request handler (_ALLOWED_GEN_METHODS).
+    if op.startswith("Generator."):
+        return
     if op not in WHITELIST:
         raise InvalidRequestError(f"unknown op: {op!r}")
 

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/src/flopscope_server/_request_handler.py

@@ -14,6 +14,27 @@ from flopscope_server._session import Session
 
 _HANDLE_RE = re.compile(r"^a\d+$")
 
+# Generator methods the server will dispatch to a server-side numpy Generator
+# resolved from a ``{"__gen__": handle}`` argument. Array-returning samplers
+# only (``shuffle`` mutates in place and is intentionally excluded).
+_ALLOWED_GEN_METHODS = frozenset(
+    {
+        "uniform",
+        "standard_normal",
+        "normal",
+        "integers",
+        "random",
+        "standard_exponential",
+        "exponential",
+        "poisson",
+        "binomial",
+        "beta",
+        "gamma",
+        "choice",
+        "permutation",
+    }
+)
+
 
 def _make_serializable(obj):
     """Convert a nested structure to be msgpack-safe (no numpy types)."""
@@ -221,6 +242,24 @@ class RequestHandler:
             result = arr.astype(dtype)
             return self._pack_result(result)
 
+        # Generator method calls: op is "Generator.<method>" with the remote
+        # generator handle as the first arg. Resolve it and call the method
+        # server-side, so the RNG state lives + advances on the server (the
+        # stream stays deterministic per seed and FLOP-counted).
+        if op.startswith("Generator."):
+            method = op[len("Generator.") :]
+            if method not in _ALLOWED_GEN_METHODS:
+                return {
+                    "status": "error",
+                    "error_type": "UnsupportedFunctionError",
+                    "message": f"Generator.{method} is not supported by the flopscope server",
+                }
+            gen = self._resolve_arg(raw_args[0])
+            rest = [self._resolve_arg(a) for a in raw_args[1:]]
+            resolved_kwargs = {k: self._resolve_arg(v) for k, v in kwargs.items()}
+            result = getattr(gen, method)(*rest, **resolved_kwargs)
+            return self._pack_result(result)
+
         func = _get_flopscope_func(op)
         resolved_args = [self._resolve_arg(a) for a in raw_args]
         resolved_kwargs = {k: self._resolve_arg(v) for k, v in kwargs.items()}
@@ -247,6 +286,13 @@ class RequestHandler:
                 if isinstance(handle, bytes):
                     handle = handle.decode("utf-8")
                 return self._session.get_array(handle)
+            gen_handle = arg.get("__gen__")
+            if gen_handle is None:
+                gen_handle = arg.get(b"__gen__")
+            if gen_handle is not None:
+                if isinstance(gen_handle, bytes):
+                    gen_handle = gen_handle.decode("utf-8")
+                return self._session.get_generator(gen_handle)
             # SymmetryGroup wire format
             pg_data = arg.get("__symmetry_group__") or arg.get(b"__symmetry_group__")
             if pg_data is not None:
@@ -370,6 +416,10 @@ class RequestHandler:
                 "result": {"value": str(result), "dtype": "str"},
                 "budget": budget,
             }
+        if isinstance(result, np.random.Generator):
+            handle = self._session.store_generator(result)
+            return {"status": "ok", "result": {"gen_id": handle}, "budget": budget}
+
         # Fallback: try to make it serializable
         try:
             serializable = _make_serializable(result)

{flopscope_server-0.4.0 → flopscope_server-0.4.2}/src/flopscope_server/_session.py

@@ -24,6 +24,8 @@ class Session:
 
     def __init__(self, flop_budget: int, flop_multiplier: float = 1.0) -> None:
         self._store = ArrayStore()
+        self._generators: dict[str, Any] = {}
+        self._gen_counter: int = 0
         self._comms_tracker = CommsTracker()
         self._budget_ctx = flops.BudgetContext(
             flop_budget=flop_budget,
@@ -109,6 +111,29 @@ class Session:
         """
         self._store.free(handles)
 
+    # ------------------------------------------------------------------
+    # Generator operations (server-side RNG handles)
+    # ------------------------------------------------------------------
+
+    def store_generator(self, gen: Any) -> str:
+        """Store an RNG ``Generator`` and return its handle ID (``g0``, ``g1`` …)."""
+        handle = f"g{self._gen_counter}"
+        self._generators[handle] = gen
+        self._gen_counter += 1
+        return handle
+
+    def get_generator(self, handle: str) -> Any:
+        """Return the ``Generator`` for *handle*.
+
+        Raises
+        ------
+        KeyError
+            If *handle* is not a known generator handle.
+        """
+        if handle not in self._generators:
+            raise KeyError(f"Generator handle {handle!r} not found in store")
+        return self._generators[handle]
+
     # ------------------------------------------------------------------
     # Budget
     # ------------------------------------------------------------------
@@ -162,6 +187,7 @@ class Session:
         budget_summary = self._budget_ctx.summary(by_namespace=show_namespaces)
         comms_summary = self._comms_tracker.summary()
         self._store.clear()
+        self._generators.clear()
         self._is_open = False
 
         return {