301 lines
10 KiB
Python
301 lines
10 KiB
Python
"""Loader and helpers for the versioned gRPC Shard protocol (ADR-0024, DGR-002).
|
|
|
|
The ``.proto`` schema at ``packages/node/native/proto/shard_runtime.proto`` is the
|
|
single source of truth. Rather than commit generated stubs (which pin a protobuf
|
|
runtime version and drift from the schema), this package generates the Python
|
|
stubs on demand into a gitignored build directory and imports them. Generation is
|
|
reproducible: it shells out to the pinned ``grpc_tools.protoc`` with the exact
|
|
same flags as ``packages/node/native/scripts/generate_python.py``.
|
|
|
|
Typical use::
|
|
|
|
from meshnet_node import native_protocol as proto
|
|
pb2 = proto.load()
|
|
header = pb2.MessageHeader(work_id="w1", route_session_id="s1")
|
|
|
|
The checksum/fragment helpers encode the bounded-fragment tensor-bundle semantics
|
|
so callers (and DGR-008/DGR-009) do not re-derive them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib
|
|
import importlib.util
|
|
import pathlib
|
|
import sys
|
|
import threading
|
|
import types
|
|
import zlib
|
|
|
|
# The wire schema version this build targets. Keep in sync with the
|
|
# ``SCHEMA_VERSION_1`` enum member in the .proto.
|
|
SCHEMA_VERSION = 1
|
|
|
|
_NATIVE_ROOT = pathlib.Path(__file__).resolve().parents[2] / "native"
|
|
PROTO_DIR = _NATIVE_ROOT / "proto"
|
|
PROTO_FILE = PROTO_DIR / "shard_runtime.proto"
|
|
# ``build/`` is globally gitignored, so generated stubs never enter version control.
|
|
GEN_DIR = _NATIVE_ROOT / "build" / "python"
|
|
|
|
_PB2_MODULE = "shard_runtime_pb2"
|
|
_GRPC_MODULE = "shard_runtime_pb2_grpc"
|
|
|
|
# Reentrant: load_grpc() holds the lock and calls load(), which re-acquires it.
|
|
_lock = threading.RLock()
|
|
_cached_pb2: types.ModuleType | None = None
|
|
_cached_grpc: types.ModuleType | None = None
|
|
|
|
|
|
class ProtocGenerationError(RuntimeError):
|
|
"""Raised when the protobuf stubs cannot be generated from the schema."""
|
|
|
|
|
|
def _needs_regen(target: pathlib.Path) -> bool:
|
|
if not target.exists():
|
|
return True
|
|
try:
|
|
return PROTO_FILE.stat().st_mtime > target.stat().st_mtime
|
|
except OSError:
|
|
return True
|
|
|
|
|
|
def generate(*, force: bool = False) -> pathlib.Path:
|
|
"""Generate ``shard_runtime_pb2{,_grpc}.py`` into :data:`GEN_DIR`.
|
|
|
|
Returns the output directory. Reproducible and idempotent: regenerates only
|
|
when the schema is newer than the stubs (or ``force`` is set). Requires the
|
|
pinned ``grpc_tools`` (available in the project ``.venv``).
|
|
"""
|
|
if not PROTO_FILE.exists():
|
|
raise ProtocGenerationError(f"schema not found: {PROTO_FILE}")
|
|
|
|
pb2_path = GEN_DIR / f"{_PB2_MODULE}.py"
|
|
if not force and not _needs_regen(pb2_path):
|
|
return GEN_DIR
|
|
|
|
try:
|
|
from grpc_tools import protoc
|
|
except ImportError as exc: # pragma: no cover - environment-dependent
|
|
raise ProtocGenerationError(
|
|
"grpc_tools is required to generate the Shard protocol stubs; "
|
|
"install grpcio-tools (present in the project .venv)."
|
|
) from exc
|
|
|
|
GEN_DIR.mkdir(parents=True, exist_ok=True)
|
|
well_known = _well_known_include()
|
|
args = [
|
|
"grpc_tools.protoc",
|
|
f"-I{PROTO_DIR}",
|
|
*([f"-I{well_known}"] if well_known else []),
|
|
f"--python_out={GEN_DIR}",
|
|
f"--grpc_python_out={GEN_DIR}",
|
|
str(PROTO_FILE.name),
|
|
]
|
|
# protoc resolves the proto by name relative to -I, so run with PROTO_DIR
|
|
# semantics by passing the bare filename plus the include path above.
|
|
rc = protoc.main([a for a in args])
|
|
if rc != 0:
|
|
raise ProtocGenerationError(
|
|
f"grpc_tools.protoc exited with status {rc} for {PROTO_FILE}"
|
|
)
|
|
if not pb2_path.exists(): # pragma: no cover - defensive
|
|
raise ProtocGenerationError(f"protoc did not produce {pb2_path}")
|
|
return GEN_DIR
|
|
|
|
|
|
def _well_known_include() -> str | None:
|
|
"""Bundled well-known .proto include dir shipped with grpc_tools, if any."""
|
|
try:
|
|
import grpc_tools
|
|
|
|
candidate = pathlib.Path(grpc_tools.__file__).parent / "_proto"
|
|
return str(candidate) if candidate.is_dir() else None
|
|
except Exception: # pragma: no cover - defensive
|
|
return None
|
|
|
|
|
|
def _import_generated(module_name: str) -> types.ModuleType:
|
|
gen_dir = str(GEN_DIR)
|
|
if gen_dir not in sys.path:
|
|
sys.path.insert(0, gen_dir)
|
|
if module_name in sys.modules:
|
|
return sys.modules[module_name]
|
|
return importlib.import_module(module_name)
|
|
|
|
|
|
def load(*, force: bool = False) -> types.ModuleType:
|
|
"""Return the generated ``shard_runtime_pb2`` module (messages only).
|
|
|
|
Generates the stubs on first use. Thread-safe and cached. Does not import
|
|
grpc; message serialization/round-trip needs only this module.
|
|
"""
|
|
global _cached_pb2
|
|
with _lock:
|
|
if _cached_pb2 is not None and not force:
|
|
return _cached_pb2
|
|
generate(force=force)
|
|
_cached_pb2 = _import_generated(_PB2_MODULE)
|
|
return _cached_pb2
|
|
|
|
|
|
def load_grpc(*, force: bool = False) -> types.ModuleType:
|
|
"""Return the generated ``shard_runtime_pb2_grpc`` module (service stubs).
|
|
|
|
Requires the ``grpc`` runtime. Use for building the C++/Python worker; the
|
|
round-trip/compat tests only need :func:`load`.
|
|
"""
|
|
global _cached_grpc
|
|
with _lock:
|
|
if _cached_grpc is not None and not force:
|
|
return _cached_grpc
|
|
generate(force=force)
|
|
load() # ensure the _pb2 module the grpc stub imports is present
|
|
_cached_grpc = _import_generated(_GRPC_MODULE)
|
|
return _cached_grpc
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Checksum + bounded-fragment helpers (shared bundle semantics)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Algorithm-name strings mirror the ChecksumAlgorithm enum members without
|
|
# importing the generated module (so this table is usable before load()).
|
|
_CHECKSUM_CRC32C = "CHECKSUM_CRC32C"
|
|
_CHECKSUM_CRC32 = "CHECKSUM_CRC32"
|
|
_CHECKSUM_SHA256 = "CHECKSUM_SHA256"
|
|
_CHECKSUM_NONE = "CHECKSUM_NONE"
|
|
|
|
|
|
def _crc32c(data: bytes) -> int:
|
|
"""Castagnoli CRC32C (software table). Deterministic, no external deps."""
|
|
crc = 0xFFFFFFFF
|
|
for byte in data:
|
|
crc ^= byte
|
|
for _ in range(8):
|
|
crc = (crc >> 1) ^ (0x82F63B78 & -(crc & 1))
|
|
return crc ^ 0xFFFFFFFF
|
|
|
|
|
|
def compute_checksum(algorithm: int, data: bytes):
|
|
"""Build a ``Checksum`` message for ``data`` under the given enum value.
|
|
|
|
``algorithm`` is a ``ChecksumAlgorithm`` enum int from the generated module.
|
|
Uses only the standard library (crc32c software table, zlib.crc32, hashlib).
|
|
"""
|
|
pb2 = load()
|
|
name = pb2.ChecksumAlgorithm.Name(algorithm)
|
|
if name == _CHECKSUM_SHA256:
|
|
value = hashlib.sha256(data).digest()
|
|
elif name == _CHECKSUM_CRC32C:
|
|
value = _crc32c(data).to_bytes(4, "big")
|
|
elif name == _CHECKSUM_CRC32:
|
|
value = (zlib.crc32(data) & 0xFFFFFFFF).to_bytes(4, "big")
|
|
elif name == _CHECKSUM_NONE:
|
|
value = b""
|
|
else:
|
|
raise ValueError(f"unsupported checksum algorithm: {name}")
|
|
return pb2.Checksum(algorithm=algorithm, value=value)
|
|
|
|
|
|
def verify_checksum(checksum, data: bytes) -> bool:
|
|
"""True if ``checksum`` matches ``data`` (CHECKSUM_NONE always verifies)."""
|
|
pb2 = load()
|
|
if checksum.algorithm in (0, pb2.CHECKSUM_NONE):
|
|
return True
|
|
return compute_checksum(checksum.algorithm, data).value == checksum.value
|
|
|
|
|
|
def fragment_tensor(
|
|
*,
|
|
name: str,
|
|
shape,
|
|
dtype: int,
|
|
payload: bytes,
|
|
byte_order: int | None = None,
|
|
max_fragment_bytes: int = 1 << 20,
|
|
compression: int | None = None,
|
|
checksum_algorithm: int | None = None,
|
|
):
|
|
"""Build a :class:`NamedTensor` splitting ``payload`` into bounded fragments.
|
|
|
|
Fragments are ordered by ``byte_offset`` and each carries an optional
|
|
per-fragment checksum. ``payload`` is treated as already compressed if
|
|
``compression`` is set; this helper does not compress (that is the seam's
|
|
policy in ``activation_compression``), it only frames.
|
|
"""
|
|
if max_fragment_bytes <= 0:
|
|
raise ValueError("max_fragment_bytes must be positive")
|
|
pb2 = load()
|
|
if byte_order is None:
|
|
byte_order = pb2.BYTE_ORDER_LITTLE_ENDIAN
|
|
if compression is None:
|
|
compression = pb2.COMPRESSION_NONE
|
|
|
|
chunks = [
|
|
payload[i : i + max_fragment_bytes]
|
|
for i in range(0, len(payload), max_fragment_bytes)
|
|
] or [b""]
|
|
fragments = []
|
|
offset = 0
|
|
for index, chunk in enumerate(chunks):
|
|
frag = pb2.TensorFragment(
|
|
fragment_index=index,
|
|
fragment_count=len(chunks),
|
|
byte_offset=offset,
|
|
data=chunk,
|
|
)
|
|
if checksum_algorithm is not None:
|
|
frag.checksum.CopyFrom(compute_checksum(checksum_algorithm, chunk))
|
|
fragments.append(frag)
|
|
offset += len(chunk)
|
|
return pb2.NamedTensor(
|
|
name=name,
|
|
shape=list(shape),
|
|
dtype=dtype,
|
|
byte_order=byte_order,
|
|
total_byte_length=len(payload),
|
|
compression=compression,
|
|
fragments=fragments,
|
|
)
|
|
|
|
|
|
def reassemble_tensor(named_tensor) -> bytes:
|
|
"""Concatenate a :class:`NamedTensor`'s fragments back into the full payload.
|
|
|
|
Validates fragment ordering, total length, and any per-fragment checksums.
|
|
"""
|
|
fragments = sorted(named_tensor.fragments, key=lambda f: f.byte_offset)
|
|
out = bytearray()
|
|
for frag in fragments:
|
|
if frag.byte_offset != len(out):
|
|
raise ValueError(
|
|
f"non-contiguous fragment at offset {frag.byte_offset} "
|
|
f"(expected {len(out)})"
|
|
)
|
|
if frag.HasField("checksum") and not verify_checksum(frag.checksum, frag.data):
|
|
raise ValueError(f"fragment {frag.fragment_index} checksum mismatch")
|
|
out.extend(frag.data)
|
|
if named_tensor.total_byte_length and len(out) != named_tensor.total_byte_length:
|
|
raise ValueError(
|
|
f"reassembled length {len(out)} != declared "
|
|
f"{named_tensor.total_byte_length}"
|
|
)
|
|
return bytes(out)
|
|
|
|
|
|
__all__ = [
|
|
"SCHEMA_VERSION",
|
|
"PROTO_FILE",
|
|
"PROTO_DIR",
|
|
"GEN_DIR",
|
|
"ProtocGenerationError",
|
|
"generate",
|
|
"load",
|
|
"load_grpc",
|
|
"compute_checksum",
|
|
"verify_checksum",
|
|
"fragment_tensor",
|
|
"reassemble_tensor",
|
|
]
|