neuron-tai/packages/node/meshnet_node/activation_compression.py

"""Policy-driven zstd compression for activation seam bodies.

Policies are intentionally local to a hop condition: a LAN prefill can favour
wire savings while a one-token decode keeps a larger raw fast path.  Environment
overrides make a trace-tuned rollout possible without changing the wire format.
"""

from __future__ import annotations

from dataclasses import dataclass
import os
import time


@dataclass(frozen=True)
class CompressionPolicy:
    """The measurable conditions required before an activation is compressed."""

    min_input_bytes: int
    min_savings_bytes: int = 4096
    min_savings_ratio: float = 0.05
    level: int = 1
    enabled: bool = True


@dataclass(frozen=True)
class CompressionResult:
    body: bytes
    encoding: str | None
    input_bytes: int
    output_bytes: int
    elapsed_seconds: float
    decision: str

    @property
    def compressed(self) -> bool:
        return self.encoding == "zstd"


_DEFAULTS: dict[tuple[str, str], CompressionPolicy] = {
    # Decode activations usually contain one position; keep that hot path raw.
    ("lan", "prefill"): CompressionPolicy(64 * 1024),
    ("lan", "decode"): CompressionPolicy(128 * 1024),
    ("relay", "prefill"): CompressionPolicy(32 * 1024),
    ("relay", "decode"): CompressionPolicy(128 * 1024),
    # The deterministic benchmark can explicitly model either policy family.
    ("benchmark", "prefill"): CompressionPolicy(64 * 1024),
    ("benchmark", "decode"): CompressionPolicy(128 * 1024),
}


class CompressionPolicies:
    """Explicit policies for LAN, relay, and benchmark prefill/decode seams.

    Set ``MESHNET_COMPRESSION_<ROUTE>_<PHASE>_MIN_INPUT_BYTES``,
    ``..._MIN_SAVINGS_BYTES``, ``..._MIN_SAVINGS_RATIO``, or ``..._ENABLED`` to
    tune a condition from production traces.  E.g.
    ``MESHNET_COMPRESSION_RELAY_PREFILL_MIN_INPUT_BYTES=32768``.
    """

    def __init__(self, policies: dict[tuple[str, str], CompressionPolicy] | None = None) -> None:
        self._policies = dict(_DEFAULTS if policies is None else policies)

    def for_condition(self, route: str, phase: str) -> CompressionPolicy:
        key = (route.lower(), phase.lower())
        try:
            policy = self._policies[key]
        except KeyError as exc:
            raise ValueError(f"unknown compression condition {route}/{phase}") from exc
        prefix = f"MESHNET_COMPRESSION_{key[0].upper()}_{key[1].upper()}_"
        return CompressionPolicy(
            min_input_bytes=_env_int(prefix + "MIN_INPUT_BYTES", policy.min_input_bytes),
            min_savings_bytes=_env_int(prefix + "MIN_SAVINGS_BYTES", policy.min_savings_bytes),
            min_savings_ratio=_env_float(prefix + "MIN_SAVINGS_RATIO", policy.min_savings_ratio),
            level=_env_int(prefix + "LEVEL", policy.level),
            enabled=_env_bool(prefix + "ENABLED", policy.enabled),
        )


def compress_activation(body: bytes, policy: CompressionPolicy) -> CompressionResult:
    """Compress only when zstd clears both configured savings thresholds."""
    started = time.monotonic()
    if not policy.enabled:
        return _raw(body, started, "disabled")
    if len(body) < policy.min_input_bytes:
        return _raw(body, started, "below_min_input")
    try:
        import zstandard as zstd

        candidate = zstd.ZstdCompressor(level=policy.level).compress(body)
    except Exception:
        # Compression is an optional transport optimisation, never a reason to
        # reject an otherwise valid activation.
        return _raw(body, started, "unavailable")
    saved = len(body) - len(candidate)
    ratio = saved / max(1, len(body))
    if saved < policy.min_savings_bytes or ratio < policy.min_savings_ratio:
        return _raw(body, started, "below_savings")
    return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed")


def decompress_activation(
    body: bytes,
    encoding: str | None,
    *,
    max_output_bytes: int | None = None,
) -> CompressionResult:
    """Decode a modern zstd body or preserve a legacy raw body with metrics."""
    started = time.monotonic()
    if not encoding:
        return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, "legacy_raw")
    if encoding != "zstd":
        raise ValueError("unsupported X-Meshnet-Encoding")
    try:
        import zstandard as zstd
    except ImportError as exc:
        raise ValueError("zstd support is unavailable") from exc
    if max_output_bytes is not None and max_output_bytes < 0:
        raise ValueError("max_output_bytes must be non-negative")
    try:
        if max_output_bytes is None:
            raw = zstd.ZstdDecompressor().decompress(body)
        else:
            # Cap both decoder window allocation and bytes read. zstandard's
            # max_window_size unit is KiB.
            max_window_kib = max(1024, (max_output_bytes + 1023) // 1024)
            decompressor = zstd.ZstdDecompressor(max_window_size=max_window_kib)
            # `decompress(max_output_size=...)` may trust a frame's advertised
            # content size. A bounded stream read enforces the limit regardless
            # of frame metadata and detects trailing expansion with one byte.
            with decompressor.stream_reader(body) as reader:
                raw = reader.read(max_output_bytes + 1)
            if len(raw) > max_output_bytes:
                raise ValueError("zstd activation body exceeds its output limit")
    except zstd.ZstdError as exc:
        raise ValueError("invalid zstd activation body") from exc
    return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed")


def _raw(body: bytes, started: float, decision: str) -> CompressionResult:
    return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, decision)


def _env_int(name: str, default: int) -> int:
    try:
        return max(0, int(os.getenv(name, str(default))))
    except ValueError:
        return default


def _env_float(name: str, default: float) -> float:
    try:
        return max(0.0, float(os.getenv(name, str(default))))
    except ValueError:
        return default


def _env_bool(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() not in {"0", "false", "no", "off"}