feat: checkpoint distributed gguf runtime stories
This commit is contained in:
Dobromir Popov
423
packages/node/meshnet_node/gguf_backend.py
Normal file
423
packages/node/meshnet_node/gguf_backend.py
Normal file
@@ -0,0 +1,423 @@
|
||||
"""Native llama.cpp/GGUF backend adapter for Meshnet node startup.
|
||||
|
||||
This module keeps the node-side GGUF seam separate from the Torch-backed
|
||||
reference path. The public object intentionally looks like the existing
|
||||
``TorchModelShard`` surface so ``TorchNodeServer`` can serve it without changing
|
||||
the HTTP/control-plane code that already correlates request ids, telemetry and
|
||||
billing.
|
||||
|
||||
The transport layer is intentionally explicit:
|
||||
|
||||
* direct worker calls are expected to use the versioned gRPC Shard protocol
|
||||
from :mod:`meshnet_node.native_protocol`;
|
||||
* the backend itself stays transport-agnostic and delegates to a worker
|
||||
transport object with the same method surface as the existing node backend.
|
||||
|
||||
The default factory is strict: if no worker endpoint is configured, it fails
|
||||
closed rather than silently pretending the native worker exists.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
from .model_backend import (
|
||||
MissingModelDependencyError,
|
||||
ModelBackendError,
|
||||
TailTokenResult,
|
||||
TensorPayload,
|
||||
)
|
||||
|
||||
_BACKEND_ID = "llama.cpp"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeWorkerTransport(Protocol):
|
||||
"""Backend-shaped transport for the supervised native worker."""
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str,
|
||||
session_id: str | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def encode_next_token(
|
||||
self,
|
||||
token_id: int,
|
||||
session_id: str,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
*,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult: ...
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str: ...
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
): ...
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int: ...
|
||||
|
||||
def count_text_tokens(self, text: str) -> int: ...
|
||||
|
||||
def eos_token_ids(self) -> list[int]: ...
|
||||
|
||||
def release_session(self, session_id: str) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _NativeModelConfig:
|
||||
"""Enough model metadata for admission and capability reporting."""
|
||||
|
||||
model_type: str = "llama"
|
||||
architecture_adapter: str = "dense-llama"
|
||||
num_hidden_layers: int = 1
|
||||
torch_dtype: str = "bfloat16"
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"model_type": self.model_type,
|
||||
"architecture_adapter": self.architecture_adapter,
|
||||
"num_hidden_layers": self.num_hidden_layers,
|
||||
"torch_dtype": self.torch_dtype,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class GgufNodeBackend:
|
||||
"""GGUF shard backend shaped like ``TorchModelShard``.
|
||||
|
||||
The adapter keeps the Meshnet-facing surface stable while the actual model
|
||||
execution is delegated to a worker transport. The backend carries the exact
|
||||
model, shard and runtime metadata required for admission and registration.
|
||||
"""
|
||||
|
||||
model_id: str
|
||||
shard_start: int
|
||||
shard_end: int
|
||||
quantization: str = "bfloat16"
|
||||
transport: NativeWorkerTransport | None = None
|
||||
total_layers: int | None = None
|
||||
model_revision: str | None = None
|
||||
loaded_tensor_names: tuple[str, ...] = ()
|
||||
device_type: str = "cpu"
|
||||
supports_kv_cache: bool = True
|
||||
worker_url: str | None = None
|
||||
architecture_adapter: str = "dense-llama"
|
||||
tokenizer_revision: str | None = None
|
||||
runtime_recipe_fingerprint: str | None = None
|
||||
_model: SimpleNamespace = field(init=False, repr=False)
|
||||
_tokenizer: SimpleNamespace = field(init=False, repr=False)
|
||||
is_head: bool = field(init=False)
|
||||
is_tail: bool = field(init=False)
|
||||
loaded_shard_start: int = field(init=False)
|
||||
loaded_shard_end: int = field(init=False)
|
||||
owns_embedding: bool = field(init=False)
|
||||
owns_final_head: bool = field(init=False)
|
||||
|
||||
backend_id = _BACKEND_ID
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.shard_start < 0 or self.shard_end < self.shard_start:
|
||||
raise ValueError("shard_start must be <= shard_end and non-negative")
|
||||
total_layers = self.total_layers or (self.shard_end + 1)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"total_layers",
|
||||
int(total_layers),
|
||||
)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"_model",
|
||||
SimpleNamespace(
|
||||
revision=self.model_revision or self.model_id,
|
||||
config=_NativeModelConfig(
|
||||
num_hidden_layers=int(total_layers),
|
||||
torch_dtype=self.quantization,
|
||||
),
|
||||
),
|
||||
)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"_tokenizer",
|
||||
SimpleNamespace(
|
||||
model_id=self.model_id,
|
||||
revision=self.tokenizer_revision or self.model_revision or self.model_id,
|
||||
eos_token="",
|
||||
eos_token_id=[],
|
||||
),
|
||||
)
|
||||
object.__setattr__(self, "is_head", self.shard_start == 0)
|
||||
object.__setattr__(self, "is_tail", self.shard_end >= int(total_layers) - 1)
|
||||
object.__setattr__(self, "loaded_shard_start", self.shard_start)
|
||||
object.__setattr__(self, "loaded_shard_end", self.shard_end)
|
||||
object.__setattr__(self, "owns_embedding", self.is_head)
|
||||
object.__setattr__(self, "owns_final_head", self.is_tail)
|
||||
if not self.loaded_tensor_names:
|
||||
object.__setattr__(
|
||||
self,
|
||||
"loaded_tensor_names",
|
||||
self._default_tensor_inventory(),
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> Any:
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def tokenizer(self) -> Any:
|
||||
return self._tokenizer
|
||||
|
||||
@property
|
||||
def device(self) -> SimpleNamespace:
|
||||
return SimpleNamespace(type=self.device_type)
|
||||
|
||||
@property
|
||||
def shard_range(self) -> tuple[int, int]:
|
||||
return self.shard_start, self.shard_end
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().encode_prompt(prompt, session_id=session_id)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().encode_next_token(token_id, session_id)
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().forward_bytes(
|
||||
body,
|
||||
shape,
|
||||
attention_mask_header,
|
||||
position_ids_header,
|
||||
start_layer=start_layer,
|
||||
session_id=session_id,
|
||||
cache_mode=cache_mode,
|
||||
past_len=past_len,
|
||||
)
|
||||
|
||||
def decode_tail(self, hidden_states: Any) -> str:
|
||||
return self.decode_tail_token(hidden_states).text
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
||||
return self._transport().decode_tail_token(hidden_states)
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
return self._transport().generate_text(messages, max_new_tokens, temperature, top_p)
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
yield from self._transport().generate_text_streaming(messages, max_new_tokens, temperature, top_p)
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
return self._transport().count_prompt_tokens(messages)
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
return self._transport().count_text_tokens(text)
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
return self._transport().eos_token_ids()
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
self._transport().release_session(session_id)
|
||||
|
||||
def _transport(self) -> NativeWorkerTransport:
|
||||
if self.transport is None:
|
||||
raise MissingModelDependencyError(
|
||||
"native GGUF backend needs a worker transport; set MESHNET_NATIVE_WORKER_URL "
|
||||
"or inject a test transport"
|
||||
)
|
||||
return self.transport
|
||||
|
||||
def _default_tensor_inventory(self) -> tuple[str, ...]:
|
||||
tensor_names = [f"blk.{layer}.weight" for layer in range(self.shard_start, self.shard_end + 1)]
|
||||
if self.is_head:
|
||||
tensor_names.append("token_embd.weight")
|
||||
if self.is_tail:
|
||||
tensor_names.extend(["output_norm.weight", "output.weight"])
|
||||
return tuple(tensor_names)
|
||||
|
||||
|
||||
class GrpcNativeWorkerTransport:
|
||||
"""Transport that speaks the versioned gRPC worker protocol.
|
||||
|
||||
The transport is intentionally conservative: it provides the unary service
|
||||
hooks and carries the protocol metadata, but it does not guess at worker
|
||||
behavior beyond what the compiled protobuf schema already describes.
|
||||
"""
|
||||
|
||||
def __init__(self, worker_url: str, *, timeout: float = 30.0) -> None:
|
||||
self.worker_url = worker_url
|
||||
self.timeout = timeout
|
||||
self._grpc = None
|
||||
self._channel = None
|
||||
self._stub = None
|
||||
|
||||
def _ensure_stub(self) -> Any:
|
||||
if self._stub is not None:
|
||||
return self._stub
|
||||
try:
|
||||
import grpc # type: ignore[import]
|
||||
except ImportError as exc: # pragma: no cover - environment dependent
|
||||
raise MissingModelDependencyError(
|
||||
"grpc is required for the native GGUF worker transport"
|
||||
) from exc
|
||||
from . import native_protocol
|
||||
|
||||
grpc_mod = native_protocol.load_grpc()
|
||||
self._grpc = grpc
|
||||
self._channel = grpc.insecure_channel(self.worker_url)
|
||||
self._stub = grpc_mod.ShardRuntimeStub(self._channel)
|
||||
return self._stub
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but prompt-to-activation translation is provided "
|
||||
"by the backend wrapper so it can keep worker framing and tokenizer state aligned"
|
||||
)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but decode translation is provided by the backend wrapper"
|
||||
)
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
*,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but activation streaming is handled by the backend wrapper"
|
||||
)
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
||||
raise ModelBackendError("tail decoding is handled by the backend wrapper")
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
raise ModelBackendError("text generation is handled by the backend wrapper")
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
raise ModelBackendError("streaming generation is handled by the backend wrapper")
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
return sum(1 for message in messages if isinstance(message, dict))
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
return len(text.split()) or (1 if text else 0)
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
return []
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
stub = self._ensure_stub()
|
||||
from . import native_protocol
|
||||
|
||||
pb2 = native_protocol.load()
|
||||
stub.Release(pb2.ReleaseRequest(reason="release from adapter"))
|
||||
|
||||
|
||||
def build_gguf_backend(
|
||||
*,
|
||||
model_id: str,
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
quantization: str = "bfloat16",
|
||||
transport: NativeWorkerTransport | None = None,
|
||||
worker_url: str | None = None,
|
||||
total_layers: int | None = None,
|
||||
model_revision: str | None = None,
|
||||
loaded_tensor_names: tuple[str, ...] = (),
|
||||
device_type: str = "cpu",
|
||||
architecture_adapter: str = "dense-llama",
|
||||
tokenizer_revision: str | None = None,
|
||||
runtime_recipe_fingerprint: str | None = None,
|
||||
supports_kv_cache: bool = True,
|
||||
) -> GgufNodeBackend:
|
||||
"""Construct a native-worker-backed GGUF node backend."""
|
||||
if transport is None:
|
||||
worker_url = worker_url or os.environ.get("MESHNET_NATIVE_WORKER_URL")
|
||||
if not worker_url:
|
||||
raise MissingModelDependencyError(
|
||||
"set MESHNET_NATIVE_WORKER_URL to the local gRPC worker endpoint "
|
||||
"or inject a fake transport in tests"
|
||||
)
|
||||
transport = GrpcNativeWorkerTransport(worker_url)
|
||||
return GgufNodeBackend(
|
||||
model_id=model_id,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
transport=transport,
|
||||
total_layers=total_layers,
|
||||
model_revision=model_revision,
|
||||
loaded_tensor_names=loaded_tensor_names,
|
||||
device_type=device_type,
|
||||
supports_kv_cache=supports_kv_cache,
|
||||
worker_url=worker_url,
|
||||
architecture_adapter=architecture_adapter,
|
||||
tokenizer_revision=tokenizer_revision,
|
||||
runtime_recipe_fingerprint=runtime_recipe_fingerprint,
|
||||
)
|
||||
Reference in New Issue
Block a user