424 lines
14 KiB
Python
424 lines
14 KiB
Python
"""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,
|
|
)
|