feat: checkpoint distributed gguf runtime stories
This commit is contained in:
Dobromir Popov
300
packages/node/meshnet_node/native_protocol/__init__.py
Normal file
300
packages/node/meshnet_node/native_protocol/__init__.py
Normal file
@@ -0,0 +1,300 @@
|
||||
"""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",
|
||||
]
|
||||
Reference in New Issue
Block a user