flopscope-client 0.4.0__py3-none-any.whl

flopscope/__init__.py

@@ -0,0 +1,335 @@
+"""flopscope — transparent proxy to a remote flopscope server.
+
+This module exposes a numpy-like API where every operation is dispatched
+to a remote server over ZMQ.  Participants use it as::
+
+    import flopscope as flops
+    import flopscope.numpy as fnp
+
+    with flops.BudgetContext(flop_budget=1_000_000) as ctx:
+        a = fnp.array([[1.0, 2.0], [3.0, 4.0]])
+        b = fnp.zeros((2, 2))
+        c = fnp.add(a, b)
+"""
+
+from __future__ import annotations
+
+import builtins
+import struct
+from typing import Any
+
+__version__ = "0.4.0"
+
+# ---------------------------------------------------------------------------
+# Errors
+# ---------------------------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Budget
+# ---------------------------------------------------------------------------
+from flopscope._budget import (  # noqa: E402
+    BudgetContext,
+    OpRecord,
+    budget,
+    budget_summary_dict,
+)
+from flopscope._display import budget_live, budget_summary  # noqa: E402
+from flopscope._math_compat import e, inf, nan, pi  # noqa: E402
+from flopscope._perm_group import SymmetryGroup  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Remote types
+# ---------------------------------------------------------------------------
+from flopscope._remote_array import (  # noqa: E402
+    _DTYPE_INFO,
+    RemoteArray,
+    RemoteScalar,
+    _encode_arg,
+    _result_from_response,
+)
+from flopscope.errors import (  # noqa: E402
+    BudgetExhaustedError,
+    FlopscopeError,
+    FlopscopeServerError,
+    FlopscopeWarning,
+    NoBudgetContextError,
+    SymmetryError,
+)
+
+# Alias: ``fnp.ndarray`` refers to the RemoteArray class.
+ndarray = RemoteArray
+
+# ---------------------------------------------------------------------------
+# Connection / protocol (private)
+# ---------------------------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Submodules (imported so ``fnp.linalg``, ``fnp.random``, ``fnp.fft`` work)
+# ---------------------------------------------------------------------------
+from flopscope import (
+    fft,  # noqa: E402, F401
+    flops,  # noqa: E402, F401
+    linalg,  # noqa: E402, F401
+    random,  # noqa: E402, F401
+    stats,  # noqa: E402, F401
+)
+from flopscope._connection import get_connection  # noqa: E402
+from flopscope._protocol import (  # noqa: E402
+    encode_create_from_data,
+    encode_request,
+)
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+from flopscope._registry import (  # noqa: E402
+    BLACKLISTED,
+    FUNCTION_CATEGORIES,
+    get_category,
+    is_valid_op,
+    iter_proxyable,
+)
+from flopscope._registry_data import FUNCTION_CATEGORIES as _FC  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Constants (no server round-trip needed)
+# ---------------------------------------------------------------------------
+
+pi: float = pi
+e: float = e
+inf: float = inf
+nan: float = nan
+newaxis = None
+
+# ---------------------------------------------------------------------------
+# Dtype strings (mirror numpy dtype names as plain strings)
+# ---------------------------------------------------------------------------
+
+float16: str = "float16"
+float32: str = "float32"
+float64: str = "float64"
+int8: str = "int8"
+int16: str = "int16"
+int32: str = "int32"
+int64: str = "int64"
+uint8: str = "uint8"
+bool_: str = "bool"
+complex64: str = "complex64"
+complex128: str = "complex128"
+
+# ---------------------------------------------------------------------------
+# Proxy factory
+# ---------------------------------------------------------------------------
+
+
+def _make_proxy(op_name: str):
+    """Create a proxy function that dispatches *op_name* to the server."""
+
+    def proxy(*args: Any, **kwargs: Any):
+        conn = get_connection()
+        encoded_args = [_encode_arg(a) for a in args]
+        encoded_kwargs = {k: _encode_arg(v) for k, v in kwargs.items()}
+        resp = conn.send_recv(
+            encode_request(op_name, args=encoded_args, kwargs=encoded_kwargs)
+        )
+        return _result_from_response(resp)
+
+    proxy.__name__ = op_name
+    proxy.__qualname__ = op_name
+    return proxy
+
+
+# ---------------------------------------------------------------------------
+# Special-case: array()
+# ---------------------------------------------------------------------------
+
+
+def _flatten(obj):
+    """Recursively flatten a nested list/tuple and return ``(flat, shape)``."""
+    if not isinstance(obj, (list, tuple)):
+        return [obj], ()
+    if len(obj) == 0:
+        return [], (0,)
+    first_flat, inner_shape = _flatten(obj[0])
+    flat = list(first_flat)
+    for item in obj[1:]:
+        item_flat, item_shape = _flatten(item)
+        if item_shape != inner_shape:
+            raise ValueError(
+                f"Inhomogeneous shape: expected inner shape {inner_shape}, "
+                f"got {item_shape}"
+            )
+        flat.extend(item_flat)
+    return flat, (len(obj),) + inner_shape
+
+
+def _infer_dtype(values):
+    """Infer a dtype string from a list of Python scalars."""
+    # Use builtins.any/all to avoid collision with the proxy functions
+    # that shadow these names at module level.
+    _any = builtins.any
+    _all = builtins.all
+    has_float = _any(isinstance(v, float) for v in values)
+    has_complex = _any(isinstance(v, complex) for v in values)
+    if has_complex:
+        return "complex128"
+    if has_float:
+        return "float64"
+    if _all(isinstance(v, bool) for v in values):
+        return "bool"
+    if _all(isinstance(v, int) for v in values):
+        return "int64"
+    return "float64"  # mixed or float values
+
+
+def array(object, dtype=None, **kwargs):
+    """Create a remote array from a Python list, tuple, or existing RemoteArray.
+
+    Parameters
+    ----------
+    object:
+        Data to create the array from.  May be a nested list/tuple of
+        numbers or an existing :class:`RemoteArray`.
+    dtype:
+        Optional dtype string (e.g. ``"float64"``).  Inferred from data
+        if not given.
+
+    Returns
+    -------
+    RemoteArray
+        A new remote array on the server.
+    """
+    if isinstance(object, RemoteArray):
+        if dtype is None:
+            return object
+        # dtype cast: dispatch to server
+        conn = get_connection()
+        resp = conn.send_recv(
+            encode_request("astype", args=[{"__handle__": object.handle_id}, dtype])
+        )
+        return _result_from_response(resp)
+
+    if isinstance(object, (list, tuple)):
+        flat, shape = _flatten(object)
+        if not flat:
+            # Empty array
+            dtype_str = dtype if isinstance(dtype, str) else "float64"
+            conn = get_connection()
+            resp = conn.send_recv(encode_create_from_data(b"", list(shape), dtype_str))
+            return _result_from_response(resp)
+
+        dtype_str = dtype if isinstance(dtype, str) else (dtype or _infer_dtype(flat))
+        info = _DTYPE_INFO.get(dtype_str)
+        if info is None:
+            raise TypeError(f"Unsupported dtype: {dtype_str!r}")
+        fmt_char, _ = info
+
+        # Complex types: split each value into (real, imag) pairs
+        if dtype_str in ("complex64", "complex128"):
+            expanded = []
+            for v in flat:
+                c = complex(v)
+                expanded.extend([c.real, c.imag])
+            flat = expanded
+            fmt_char = "f" if dtype_str == "complex64" else "d"
+            data = struct.pack(f"<{len(flat)}{fmt_char}", *flat)
+        else:
+            data = struct.pack(f"<{len(flat)}{fmt_char}", *flat)
+
+        conn = get_connection()
+        resp = conn.send_recv(encode_create_from_data(data, list(shape), dtype_str))
+        return _result_from_response(resp)
+
+    if isinstance(object, (int, float, complex)):
+        # Scalar -> 0-d array
+        if isinstance(object, complex) and dtype is None:
+            dtype_str = "complex128"
+        else:
+            dtype_str = dtype if isinstance(dtype, str) else "float64"
+        info = _DTYPE_INFO.get(dtype_str)
+        if info is None:
+            raise TypeError(f"Unsupported dtype: {dtype_str!r}")
+        fmt_char, _ = info
+
+        if dtype_str in ("complex64", "complex128"):
+            c = complex(object)
+            pack_fmt = "f" if dtype_str == "complex64" else "d"
+            data = struct.pack(f"<2{pack_fmt}", c.real, c.imag)
+        else:
+            data = struct.pack(f"<1{fmt_char}", object)
+        conn = get_connection()
+        resp = conn.send_recv(encode_create_from_data(data, [], dtype_str))
+        return _result_from_response(resp)
+
+    raise TypeError(
+        f"Cannot create array from {type(object).__name__}. "
+        f"Expected list, tuple, int, float, or RemoteArray."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Special-case: einsum()
+# ---------------------------------------------------------------------------
+
+
+def einsum(subscripts, *operands, **kwargs):
+    """Einstein summation on remote arrays.
+
+    Parameters
+    ----------
+    subscripts:
+        Subscript string (e.g. ``"ij,jk->ik"``).
+    *operands:
+        Input :class:`RemoteArray` objects.
+    **kwargs:
+        Additional keyword arguments forwarded to the server.
+
+    Returns
+    -------
+    RemoteArray
+        Result of the einsum operation.
+    """
+    conn = get_connection()
+    encoded_args = [subscripts] + [_encode_arg(op) for op in operands]
+    encoded_kwargs = {k: _encode_arg(v) for k, v in kwargs.items()}
+    resp = conn.send_recv(
+        encode_request("einsum", args=encoded_args, kwargs=encoded_kwargs)
+    )
+    return _result_from_response(resp)
+
+
+# ---------------------------------------------------------------------------
+# Auto-generate proxy functions for all non-blacklisted top-level ops
+# ---------------------------------------------------------------------------
+
+# Functions that are special-cased above and should not be overwritten.
+_SPECIAL_CASED = frozenset({"array", "einsum"})
+
+# Functions that belong to submodules (contain a dot) are handled by the
+# submodule packages themselves.
+_generated_proxies: list[str] = []
+for _op_name in iter_proxyable():
+    if "." in _op_name:
+        continue  # submodule function
+    if _op_name in _SPECIAL_CASED:
+        continue
+    globals()[_op_name] = _make_proxy(_op_name)
+    _generated_proxies.append(_op_name)
+
+del _op_name  # clean up loop variable
+
+
+# ---------------------------------------------------------------------------
+# Module-level __getattr__ for blacklisted / unknown names
+# ---------------------------------------------------------------------------
+
+# We import the factory but define the function inline so we can also
+# check against names that are already defined in the module namespace.
+
+from flopscope._getattr import make_module_getattr as _make_module_getattr  # noqa: E402
+
+_module_getattr = _make_module_getattr("", "flopscope")
+
+
+def __getattr__(name: str):
+    return _module_getattr(name)

flopscope/_budget.py

@@ -0,0 +1,324 @@
+"""Client-side BudgetContext proxy that delegates to the flopscope server."""
+
+from __future__ import annotations
+
+from flopscope._connection import get_connection
+from flopscope._protocol import (
+    encode_budget_close,
+    encode_budget_open,
+    encode_budget_status,
+)
+
+# Module-level guard: only one BudgetContext can be active at a time.
+_active_context = None
+
+
+class OpRecord:
+    """Record of a single operation's FLOP cost.
+
+    Parameters
+    ----------
+    op_name:
+        Name of the operation (e.g. ``"dot"``).
+    flop_cost:
+        FLOPs charged for this operation.
+    cumulative:
+        Total FLOPs used after this operation.
+    """
+
+    def __init__(self, op_name: str, flop_cost: int, cumulative: int) -> None:
+        self.op_name = op_name
+        self.flop_cost = flop_cost
+        self.cumulative = cumulative
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return (
+            f"OpRecord(op_name={self.op_name!r}, "
+            f"flop_cost={self.flop_cost}, cumulative={self.cumulative})"
+        )
+
+
+class BudgetContext:
+    """Context manager that opens/closes a FLOP budget on the server.
+
+    Parameters
+    ----------
+    flop_budget:
+        Maximum FLOPs allowed within this context.
+    flop_multiplier:
+        Scaling factor applied to each operation's raw FLOP count before
+        it is charged against the budget.  Defaults to ``1.0``.
+    quiet:
+        If ``True``, suppress informational output.  Defaults to ``False``.
+    namespace:
+        Optional label for grouping budget records.
+
+    Example
+    -------
+    >>> with BudgetContext(flop_budget=1_000_000) as ctx:
+    ...     result = flopscope.dot(a, b)
+    ...     print(ctx.summary())
+    """
+
+    def __init__(
+        self,
+        flop_budget: int,
+        flop_multiplier: float = 1.0,
+        quiet: bool = False,
+        namespace: str | None = None,
+    ) -> None:
+        self._flop_budget = flop_budget
+        self._flop_multiplier = flop_multiplier
+        self._quiet = quiet
+        self._namespace = namespace
+        self._flops_used: int = 0
+        self._close_summary: str | None = None
+        self._is_open: bool = False
+        self._previous_context = None
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def flop_budget(self) -> int:
+        """Maximum FLOP allowance for this context."""
+        return self._flop_budget
+
+    @property
+    def flops_used(self) -> int:
+        """FLOPs consumed so far (cached locally, updated from server responses)."""
+        return self._flops_used
+
+    @property
+    def flops_remaining(self) -> int:
+        """FLOPs remaining in the budget (``budget - used``)."""
+        return self._flop_budget - self._flops_used
+
+    @property
+    def flop_multiplier(self) -> float:
+        """FLOP scaling multiplier."""
+        return self._flop_multiplier
+
+    @property
+    def quiet(self) -> bool:
+        """Whether informational output is suppressed."""
+        return self._quiet
+
+    @property
+    def namespace(self) -> str | None:
+        """Optional namespace label for this context."""
+        return self._namespace
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _update_budget(self, budget_info: dict) -> None:
+        """Update the local ``flops_used`` cache from a server-response dict.
+
+        Parameters
+        ----------
+        budget_info:
+            Dict that may contain a ``"flops_used"`` key.  Missing key is
+            silently ignored.
+        """
+        if "flops_used" in budget_info:
+            self._flops_used = int(budget_info["flops_used"])
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def summary(self) -> str:
+        """Query the server for current budget status and return a formatted string.
+
+        Also updates the local ``flops_used`` cache.
+
+        Returns
+        -------
+        str
+            Human-readable summary of budget usage.
+        """
+        conn = get_connection()
+        response = conn.send_recv(encode_budget_status())
+        # Budget status is nested inside "result" key
+        result = response.get("result", {})
+        self._update_budget(result)
+        budget = result.get("flop_budget", self._flop_budget)
+        used = self._flops_used
+        remaining = int(budget) - used
+        return f"BudgetContext: {used}/{budget} FLOPs used ({remaining} remaining)"
+
+    # ------------------------------------------------------------------
+    # Decorator support
+    # ------------------------------------------------------------------
+
+    def __call__(self, func):
+        """Use BudgetContext as a decorator."""
+        import functools
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    # ------------------------------------------------------------------
+    # Context manager protocol
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> BudgetContext:
+        """Open the budget on the server and update the local cache."""
+        global _active_context
+        if _active_context is not None and _active_context is not _global_default:
+            raise RuntimeError(
+                "Nested BudgetContext is not supported. "
+                "Only one context can be active at a time."
+            )
+        self._previous_context = _active_context
+        conn = get_connection()
+        response = conn.send_recv(
+            encode_budget_open(self._flop_budget, self._flop_multiplier)
+        )
+        self._update_budget(response)
+        self._is_open = True
+        _active_context = self
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        """Close the budget on the server and store the close summary."""
+        global _active_context
+        if self._is_open:
+            conn = get_connection()
+            response = conn.send_recv(encode_budget_close())
+            self._update_budget(response)
+            self._close_summary = (
+                f"BudgetContext closed: {self._flops_used}/{self._flop_budget} "
+                f"FLOPs used"
+            )
+            self._is_open = False
+            _accumulator.record(self)
+        _active_context = self._previous_context
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return (
+            f"BudgetContext(flop_budget={self._flop_budget}, "
+            f"flops_used={self._flops_used}, "
+            f"flop_multiplier={self._flop_multiplier})"
+        )
+
+
+# ------------------------------------------------------------------
+# Accumulator
+# ------------------------------------------------------------------
+
+
+class NamespaceRecord:
+    """Snapshot of a BudgetContext's state at close time."""
+
+    def __init__(self, namespace, flop_budget, flops_used):
+        self.namespace = namespace
+        self.flop_budget = flop_budget
+        self.flops_used = flops_used
+
+
+class BudgetAccumulator:
+    """Collects budget records across multiple BudgetContext sessions."""
+
+    def __init__(self):
+        self._records = []
+
+    def record(self, ctx):
+        self._records.append(
+            NamespaceRecord(
+                namespace=ctx.namespace,
+                flop_budget=ctx.flop_budget,
+                flops_used=ctx.flops_used,
+            )
+        )
+
+    def get_data(self, by_namespace=False):
+        total_budget = sum(r.flop_budget for r in self._records)
+        total_used = sum(r.flops_used for r in self._records)
+        result = {
+            "flop_budget": total_budget,
+            "flops_used": total_used,
+            "flops_remaining": total_budget - total_used,
+            "operations": {},
+        }
+        if by_namespace:
+            by_ns = {}
+            for r in self._records:
+                ns = r.namespace
+                if ns not in by_ns:
+                    by_ns[ns] = {"flop_budget": 0, "flops_used": 0, "operations": {}}
+                by_ns[ns]["flop_budget"] += r.flop_budget
+                by_ns[ns]["flops_used"] += r.flops_used
+            result["by_namespace"] = by_ns
+        return result
+
+    def reset(self):
+        self._records.clear()
+
+
+_accumulator = BudgetAccumulator()
+
+
+def budget(flop_budget, flop_multiplier=1.0, quiet=False, namespace=None):
+    """Create a BudgetContext usable as both a context manager and decorator."""
+    return BudgetContext(
+        flop_budget=flop_budget,
+        flop_multiplier=flop_multiplier,
+        quiet=quiet,
+        namespace=namespace,
+    )
+
+
+def budget_summary_dict(by_namespace=False):
+    """Return aggregated budget data across all recorded contexts."""
+    return _accumulator.get_data(by_namespace=by_namespace)
+
+
+# Note: No budget_reset() in the client — participants must not clear usage.
+
+
+_global_default = None
+
+
+def _get_default_budget_amount():
+    import os
+
+    raw = os.environ.get("FLOPSCOPE_DEFAULT_BUDGET")
+    if raw is not None:
+        return int(float(raw))
+    return int(1e15)
+
+
+def _get_global_default():
+    global _global_default, _active_context
+    if _global_default is None:
+        _global_default = BudgetContext(
+            flop_budget=_get_default_budget_amount(),
+            quiet=True,
+            namespace=None,
+        )
+        # Open it on the server
+        conn = get_connection()
+        response = conn.send_recv(
+            encode_budget_open(
+                _global_default._flop_budget, _global_default._flop_multiplier
+            )
+        )
+        _global_default._update_budget(response)
+        _global_default._is_open = True
+        _active_context = _global_default
+    return _global_default
+
+
+def _reset_global_default():
+    global _global_default, _active_context
+    if _global_default is not None and _active_context is _global_default:
+        _active_context = None
+    _global_default = None

flopscope/_comms_tracker.py

@@ -0,0 +1,46 @@
+"""Client-side communications tracker."""
+
+from __future__ import annotations
+
+
+class ClientCommsTracker:
+    """Accumulates network statistics for flopscope client requests."""
+
+    def __init__(self) -> None:
+        self._request_count: int = 0
+        self._total_round_trip_ns: int = 0
+        self._total_request_bytes: int = 0
+        self._total_response_bytes: int = 0
+
+    def record(
+        self,
+        *,
+        round_trip_ns: int,
+        request_bytes: int,
+        response_bytes: int,
+    ) -> None:
+        """Record statistics for a single request/response round trip."""
+        self._request_count += 1
+        self._total_round_trip_ns += round_trip_ns
+        self._total_request_bytes += request_bytes
+        self._total_response_bytes += response_bytes
+
+    @property
+    def request_count(self) -> int:
+        """Total number of requests recorded."""
+        return self._request_count
+
+    @property
+    def total_round_trip_ns(self) -> int:
+        """Sum of all round-trip times in nanoseconds."""
+        return self._total_round_trip_ns
+
+    @property
+    def total_request_bytes(self) -> int:
+        """Sum of all request payload sizes in bytes."""
+        return self._total_request_bytes
+
+    @property
+    def total_response_bytes(self) -> int:
+        """Sum of all response payload sizes in bytes."""
+        return self._total_response_bytes