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

"""Authoritative identity boundary for a native GGUF Shard backend.

The native loader owns the facts about the mapped artifact.  This module turns
that immutable report and separately pinned deployment inputs into the one
``ShardIdentity`` DGR-003 permits a native worker to emit.  It is intentionally
not an adapter for the legacy Transformers backend: that backend has no
authoritative immutable GGUF artifact pin and must remain identity-free.
"""

from __future__ import annotations

from dataclasses import dataclass

from .native_protocol import BUNDLE_VERSION, SCHEMA_VERSION, pb
from .runtime_recipe import (
    ArtifactIdentity,
    DerivativeBinding,
    RecipeIdentityError,
    RuntimeRecipe,
    ShardIdentity,
    check_session_open,
    handshake_error,
)


@dataclass(frozen=True)
class NativeLoadedArtifactReport:
    """Immutable GGUF facts returned by the loaded native model.

    The report is copied from ``llama_model_meshnet_range_report`` plus the
    parsed GGUF metadata while the model is live.  Byte counts are operational
    evidence rather than compatibility axes, but keeping them beside the range
    prevents a caller from substituting an unverified range declaration.
    """

    owned_start_layer: int
    owned_end_layer: int
    mapped_bytes: int
    resident_bytes: int
    registered_bytes: int
    architecture: str
    architecture_digest: str
    layer_count: int

    def __post_init__(self) -> None:
        if self.owned_start_layer < 0 or self.owned_end_layer <= self.owned_start_layer:
            raise RecipeIdentityError("native report has an invalid owned layer range")
        if self.layer_count < 1 or self.owned_end_layer > self.layer_count:
            raise RecipeIdentityError("native report range is outside GGUF layer metadata")
        if min(self.mapped_bytes, self.resident_bytes, self.registered_bytes) < 0:
            raise RecipeIdentityError("native report byte counts must be non-negative")


@dataclass(frozen=True)
class ImmutableArtifactPin:
    """Deployment-supplied immutable bytes pin, never inferred from a model name."""

    artifact_id: str
    revision: str
    content_digest: str
    derived_from: DerivativeBinding | None = None


@dataclass(frozen=True)
class NativeNumericalRecipe:
    """Immutable numerical inputs selected for this native worker instance."""

    weight_quantization: str
    activation_dtype: str
    compute_dtype: str
    kv_dtype: str
    kv_layout: str
    architecture_adapter: str
    backend_id: str
    runtime_version: str
    recipe_id: str
    recipe_version: str
    catalogue_version: str
    boundary_schema_version: int = BUNDLE_VERSION
    protocol_schema_version: int = int(SCHEMA_VERSION)


@dataclass(frozen=True)
class NativeIdentityInputs:
    """Everything a native backend needs to emit one exact identity."""

    loaded_artifact: NativeLoadedArtifactReport
    artifact_pin: ImmutableArtifactPin
    tokenizer_revision: str
    numerical_recipe: NativeNumericalRecipe


def shard_identity_from_native_report(inputs: NativeIdentityInputs) -> ShardIdentity:
    """Derive identity only from the native report and immutable pinned inputs."""
    report = inputs.loaded_artifact
    pin = inputs.artifact_pin
    recipe = inputs.numerical_recipe
    artifact = ArtifactIdentity(
        artifact_id=pin.artifact_id,
        revision=pin.revision,
        content_digest=pin.content_digest,
        architecture=report.architecture,
        architecture_digest=report.architecture_digest,
        layer_count=report.layer_count,
        derived_from=pin.derived_from,
    )
    return ShardIdentity(
        artifact=artifact,
        recipe=RuntimeRecipe(
            weight_quantization=recipe.weight_quantization,
            activation_dtype=recipe.activation_dtype,
            compute_dtype=recipe.compute_dtype,
            kv_dtype=recipe.kv_dtype,
            kv_layout=recipe.kv_layout,
            tokenizer_revision=inputs.tokenizer_revision,
            architecture_adapter=recipe.architecture_adapter,
            backend_id=recipe.backend_id,
            runtime_version=recipe.runtime_version,
            boundary_schema_version=recipe.boundary_schema_version,
            protocol_schema_version=recipe.protocol_schema_version,
            recipe_id=recipe.recipe_id,
            recipe_version=recipe.recipe_version,
            catalogue_version=recipe.catalogue_version,
        ),
        shard_start=report.owned_start_layer,
        shard_end=report.owned_end_layer,
    )


class NativeSessionRejected(RecipeIdentityError):
    """A native worker refused a ``SessionOpen`` before allocating session state."""

    def __init__(self, error: "pb.ShardError") -> None:
        super().__init__(f"native SessionOpen rejected: {error.code}")
        self.error = error


class NativeWorkerBackendAdapter:
    """Small backend-facing adapter around the native loaded-artifact seam."""

    def __init__(self, identity_inputs: NativeIdentityInputs) -> None:
        self.identity_inputs = identity_inputs
        self.identity = shard_identity_from_native_report(identity_inputs)

    @property
    def loaded_artifact_report(self) -> NativeLoadedArtifactReport:
        return self.identity_inputs.loaded_artifact

    def check_session_open(
        self,
        opened: "pb.SessionOpen",
        *,
        expected_route_session_id: str | None = None,
        expected_route_epoch: int | None = None,
    ) -> None:
        """Reject incompatible/stale opens at the native worker boundary."""
        mismatches = check_session_open(
            self.identity,
            opened,
            expected_route_session_id=expected_route_session_id,
            expected_route_epoch=expected_route_epoch,
        )
        error = handshake_error(mismatches)
        if error is not None:
            raise NativeSessionRejected(error)

    def on_session_open(
        self,
        opened: "pb.SessionOpen",
        *,
        expected_route_session_id: str | None = None,
        expected_route_epoch: int | None = None,
    ) -> "pb.SessionAccepted":
        """The native worker's SessionOpen boundary, before session allocation."""
        self.check_session_open(
            opened,
            expected_route_session_id=expected_route_session_id,
            expected_route_epoch=expected_route_epoch,
        )
        return pb.SessionAccepted(
            schema_version=SCHEMA_VERSION,
            route_session_id=opened.route_session_id,
            route_epoch=opened.route_epoch,
            fingerprint=self.identity.fingerprint.to_proto(),
        )