runtime/src/resynthesis/release_model.py

"""Portable, inference-only Release 188 model assembly.

The public repository is one standard safetensors weight map.  This module
combines its four runtime roles without consulting a training checkpoint,
private manifest, optimizer, or host-specific path:

* the immutable lexical/projection parent;
* the dense ResynthesisRBO graph;
* the resident NoNE routers (including their learned trauma state); and
* lazily materialized page tensors selected by those routers.

Generation remains uncapped at the host boundary.  The trained additive
completion surfaces inside :func:`resynthesis_rbo_generate` own stopping.
"""

from __future__ import annotations

import contextlib
import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Final, cast

import torch
from safetensors import safe_open

from resynthesis.config import ResynthesisConfig
from resynthesis.none_migration import (
    NONE_V2_RESIDENT_RUNTIME_SCHEMA,
    build_checkpoint_direct_science_stack_boundary,
)
from resynthesis.none_paging import (
    NoNEImmutablePageStore,
    NoNEPagedExpertRuntime,
)
from resynthesis.rbo import (
    COMPOSED_ADDITIVE_LINEAGE_SCHEMA,
    RBOEmissionTracePacket,
    RBOGenerationResult,
    ResynthesisRBO,
    ResynthesisRBOConfig,
    resynthesis_rbo_generate,
)
from resynthesis.release_page_store import (
    RELEASE_188_GENERATION,
    Release188SafetensorsPageStore,
)
from resynthesis.release_parent import (
    RELEASE_188_LEXICAL_ROWS,
    Release188ParentAdapter,
)
from resynthesis.science_layers import ResynthesisScienceLayerConfig

RELEASE_188_RUNTIME_SCHEMA: Final[str] = (
    "nnf.resynthesis.release_runtime.v1"
)
RELEASE_188_SCIENCE_LAYERS: Final[int] = 18
RELEASE_188_SCIENCE_EXPERTS: Final[int] = 8
RELEASE_188_DEVICE_POLICY: Final[str] = (
    "single_best_free_gpu_else_cpu"
)
RELEASE_188_DTYPE_POLICY: Final[str] = "auto"
RELEASE_188_MINIMUM_GPU_HEADROOM_BYTES: Final[int] = 48 * 1024**3
_DENSE_FILENAME: Final[str] = (
    "weights/safetensors/resynthesis.safetensors"
)
_RESIDENT_FILENAME: Final[str] = (
    "weights/safetensors/resident-runtime.safetensors"
)
_SPEC_FIELDS: Final[frozenset[str]] = frozenset(
    (
        "schema",
        "release",
        "model_index",
        "dtype",
        "device_policy",
        "science_layers",
        "science_experts",
    )
)


@dataclass(frozen=True)
class Release188RuntimeSpec:
    """Minimal public configuration resolved from ``runtime/model.json``."""

    source_path: Path
    model_index_path: Path
    dtype_policy: str
    device_policy: str
    science_layers: int
    science_experts: int

    @property
    def release_root(self) -> Path:
        """Return the public repository root containing the weight map."""

        return self.model_index_path.parent


def release_188_runtime_payload() -> dict[str, object]:
    """Return the complete portable ``runtime/model.json`` payload."""

    return {
        "schema": RELEASE_188_RUNTIME_SCHEMA,
        "release": RELEASE_188_GENERATION,
        "model_index": "../model.safetensors.index.json",
        "dtype": RELEASE_188_DTYPE_POLICY,
        "device_policy": RELEASE_188_DEVICE_POLICY,
        "science_layers": RELEASE_188_SCIENCE_LAYERS,
        "science_experts": RELEASE_188_SCIENCE_EXPERTS,
    }


def load_release_188_runtime_spec(
    spec_path: str | Path,
) -> Release188RuntimeSpec:
    """Load the small public runtime contract without accepting local paths."""

    resolved_spec = Path(spec_path).expanduser().resolve()
    if not resolved_spec.is_file():
        raise FileNotFoundError(
            f"Release 188 runtime specification is absent: {resolved_spec}"
        )
    with resolved_spec.open(encoding="utf-8") as handle:
        payload: object = json.load(handle)
    if not isinstance(payload, dict) or set(payload) != _SPEC_FIELDS:
        raise RuntimeError("Release 188 runtime specification fields differ")
    if (
        payload.get("schema") != RELEASE_188_RUNTIME_SCHEMA
        or payload.get("release") != RELEASE_188_GENERATION
        or payload.get("dtype") != RELEASE_188_DTYPE_POLICY
        or payload.get("device_policy") != RELEASE_188_DEVICE_POLICY
        or payload.get("science_layers") != RELEASE_188_SCIENCE_LAYERS
        or payload.get("science_experts") != RELEASE_188_SCIENCE_EXPERTS
    ):
        raise RuntimeError("Release 188 runtime specification differs")
    model_index_value = payload.get("model_index")
    if not isinstance(model_index_value, str):
        raise RuntimeError("Release 188 model index path is malformed")
    relative_index = PurePosixPath(model_index_value)
    if relative_index.is_absolute():
        raise RuntimeError("Release 188 model index must be repository-relative")
    model_index_path = resolved_spec.parent.joinpath(
        *relative_index.parts
    ).resolve()
    expected_index_path = (
        resolved_spec.parent.parent / "model.safetensors.index.json"
    ).resolve()
    if (
        model_index_path != expected_index_path
        or model_index_path.name != "model.safetensors.index.json"
        or not model_index_path.is_file()
    ):
        raise RuntimeError(
            "Release 188 runtime does not point to the repository weight map"
        )
    return Release188RuntimeSpec(
        source_path=resolved_spec,
        model_index_path=model_index_path,
        dtype_policy=cast(str, payload["dtype"]),
        device_policy=cast(str, payload["device_policy"]),
        science_layers=cast(int, payload["science_layers"]),
        science_experts=cast(int, payload["science_experts"]),
    )


def select_release_188_device_boundary(
    *,
    minimum_free_bytes: int = RELEASE_188_MINIMUM_GPU_HEADROOM_BYTES,
) -> torch.device:
    """Choose one clean accelerator for placement, otherwise use CPU.

    This external resource boundary does not select model routes, experts, or
    pages.  It only chooses where the complete model executes.
    """

    if (
        isinstance(minimum_free_bytes, bool)
        or minimum_free_bytes < 1
    ):
        raise ValueError("Release 188 minimum GPU headroom must be positive")
    if not torch.cuda.is_available():
        return torch.device("cpu")
    candidates: list[tuple[int, int]] = []
    for index in range(torch.cuda.device_count()):
        free_bytes, _total_bytes = torch.cuda.mem_get_info(index)
        candidates.append((free_bytes, index))
    if not candidates:
        return torch.device("cpu")
    free_bytes, index = max(candidates)
    if free_bytes < minimum_free_bytes:
        return torch.device("cpu")
    return torch.device("cuda", index)


def _release_dtype_boundary(
    device: torch.device,
    dtype_policy: str,
) -> torch.dtype:
    if dtype_policy != RELEASE_188_DTYPE_POLICY:
        raise RuntimeError("Release 188 dtype policy differs")
    return torch.bfloat16 if device.type == "cuda" else torch.float32


def _science_config(
    cfg: ResynthesisConfig,
) -> ResynthesisScienceLayerConfig:
    return ResynthesisScienceLayerConfig(
        hidden_size=cfg.hidden_size,
        knowledge_transfer_dim=cfg.knowledge_transfer_dim,
        num_layers=cfg.num_layers,
        num_experts=cfg.num_experts,
        expert_hidden_size=cfg.expert_hidden_size,
        memory_slots=cfg.memory_slots,
        attention_heads=cfg.attention_heads,
        mhc_heads=cfg.mhc_heads,
        recursive_steps=cfg.recursive_steps,
        residual_init=cfg.residual_init,
        logit_residual_init=cfg.logit_residual_init,
        kl_anchor_weight=cfg.kl_anchor_weight,
        kl_anchor_warmup_steps=cfg.kl_anchor_warmup_steps,
        glyph_input_dim=cfg.glyph_input_dim,
    )


def _load_safetensors_boundary(
    path: Path,
    *,
    role: str,
    schema: str,
) -> dict[str, torch.Tensor]:
    """Read one public safetensors role at the external load boundary."""

    if not path.is_file():
        raise FileNotFoundError(f"Release 188 safetensors is absent: {path}")
    with safe_open(  # type: ignore[no-untyped-call]
        str(path),
        framework="pt",
        device="cpu",
    ) as handle:
        metadata = handle.metadata()
        if (
            metadata is None
            or metadata.get("release") != str(RELEASE_188_GENERATION)
            or metadata.get("role") != role
            or metadata.get("schema") != schema
            or metadata.get("optimizer_state") != "excluded"
        ):
            raise RuntimeError(
                f"Release 188 {role} safetensors metadata differs"
            )
        state = {
            name: handle.get_tensor(name)
            for name in handle.keys()
        }
    if not state:
        raise RuntimeError(f"Release 188 {role} safetensors is empty")
    return state


def _resident_layer_states_boundary(
    state: Mapping[str, torch.Tensor],
    *,
    layer_count: int,
) -> tuple[dict[str, torch.Tensor], ...]:
    layers: list[dict[str, torch.Tensor]] = [
        {} for _ in range(layer_count)
    ]
    for name, value in state.items():
        layer_stem, separator, local_name = name.partition(".")
        layer_digits = layer_stem.removeprefix("layer_")
        if (
            not separator
            or not layer_stem.startswith("layer_")
            or not layer_digits.isdigit()
            or not local_name
        ):
            raise RuntimeError(
                "Release 188 resident runtime tensor name differs"
            )
        layer_id = int(layer_digits)
        if layer_id < 0 or layer_id >= layer_count:
            raise RuntimeError(
                "Release 188 resident runtime layer is outside the graph"
            )
        layers[layer_id][local_name] = value
    if any(not layer for layer in layers):
        raise RuntimeError(
            "Release 188 resident runtime omits a science layer"
        )
    return tuple(layers)


def _resident_geometry_boundary(
    state: Mapping[str, torch.Tensor],
    *,
    expected_hidden_size: int,
) -> tuple[int, int, torch.Tensor]:
    route_keys_t = state.get("router.page_route_keys")
    hidden_projection_t = state.get("router.hidden_projection.weight")
    action_projection_t = state.get("router.action_projection.weight")
    page_catalog_ids_t = state.get("router.page_catalog_ids_t")
    if (
        not isinstance(route_keys_t, torch.Tensor)
        or route_keys_t.ndim != 2
        or not isinstance(hidden_projection_t, torch.Tensor)
        or hidden_projection_t.ndim != 2
        or hidden_projection_t.shape[1] != expected_hidden_size
        or not isinstance(action_projection_t, torch.Tensor)
        or action_projection_t.ndim != 2
        or not isinstance(page_catalog_ids_t, torch.Tensor)
        or page_catalog_ids_t.ndim != 1
        or page_catalog_ids_t.dtype != torch.long
        or route_keys_t.shape[0] != page_catalog_ids_t.shape[0]
        or route_keys_t.shape[1] != hidden_projection_t.shape[0]
        or route_keys_t.shape[1] != action_projection_t.shape[0]
        or page_catalog_ids_t.numel() < 1
    ):
        raise RuntimeError(
            "Release 188 resident runtime geometry differs"
        )
    return (
        int(action_projection_t.shape[1]),
        int(route_keys_t.shape[1]),
        page_catalog_ids_t.detach().cpu().long(),
    )


def _construct_release_188_model(
    spec: Release188RuntimeSpec,
    *,
    device: torch.device,
    dtype: torch.dtype,
    session_id_t: torch.Tensor,
) -> Release188Model:
    release_root = spec.release_root
    dense_state = _load_safetensors_boundary(
        release_root / _DENSE_FILENAME,
        role="resynthesis-additive-graph",
        schema=COMPOSED_ADDITIVE_LINEAGE_SCHEMA,
    )
    resident_state = _load_safetensors_boundary(
        release_root / _RESIDENT_FILENAME,
        role="paged-resident-runtime",
        schema=NONE_V2_RESIDENT_RUNTIME_SCHEMA,
    )
    page_store = Release188SafetensorsPageStore(spec.model_index_path)
    parent = Release188ParentAdapter(
        release_root,
        device_map={"": device},
        dtype=dtype,
    )
    parent.load_weights()

    cfg = ResynthesisConfig(
        num_layers=spec.science_layers,
        num_experts=spec.science_experts,
    )
    construction_context = (
        torch.device(device)
        if device.type != "cpu"
        else contextlib.nullcontext()
    )
    with construction_context:
        science_stack = build_checkpoint_direct_science_stack_boundary(
            _science_config(cfg),
            dense_state,
            device=device,
        )
        rbo = ResynthesisRBO(
            base=parent,
            science_stack=science_stack,
            cfg=ResynthesisRBOConfig(),
            model_cfg=cfg,
        )
    rbo.ensure_parent_capability_attachment()

    # The direct CUDA builder has already assigned every science tensor in its
    # exact first-forward BF16 representation.  Reuse those assigned tensors
    # for the full-graph strict load rather than widening them back through a
    # redundant host round trip.
    active_dense_state = dict(dense_state)
    if device.type == "cuda":
        active_dense_state.update(
            {
                f"science_stack.{name}": value
                for name, value in science_stack.state_dict().items()
            }
        )
    rbo.load_trainable_state_dict(active_dense_state)
    rbo.to(device=device, dtype=dtype)
    rbo.preserve_trainable_control_precision()

    layer_states = _resident_layer_states_boundary(
        resident_state,
        layer_count=spec.science_layers,
    )
    runtimes: list[NoNEPagedExpertRuntime] = []
    expected_catalogs: list[torch.Tensor] = []
    for layer_id, layer_state in enumerate(layer_states):
        action_size, router_size, page_catalog_ids_t = (
            _resident_geometry_boundary(
                layer_state,
                expected_hidden_size=cfg.hidden_size,
            )
        )
        with construction_context:
            runtime = NoNEPagedExpertRuntime(
                hidden_size=cfg.hidden_size,
                action_size=action_size,
                router_size=router_size,
                page_count=int(page_catalog_ids_t.shape[0]),
                layer_id=layer_id,
                page_catalog_ids_t=page_catalog_ids_t,
            )
        runtime.load_seed_resident_state_boundary(layer_state)
        runtime.to(device=device, dtype=dtype)
        runtime.preserve_training_proof_precision()
        generation_t = runtime.bind_store_boundary(
            cast(NoNEImmutablePageStore, page_store),
            session_id_t,
        )
        if not torch.equal(
            generation_t.detach().cpu().long().reshape(()),
            torch.tensor(RELEASE_188_GENERATION, dtype=torch.long),
        ):
            raise RuntimeError(
                "Release 188 resident runtime generation differs"
            )
        science_stack.attach_paged_expert_runtime(layer_id, runtime)
        runtimes.append(runtime)
        expected_catalogs.append(page_catalog_ids_t)

    # This is the manifest-free equivalent of the internal composition
    # boundary: the router tensors themselves are the layer/page authority.
    rbo._paged_none_store_boundary = page_store
    rbo._paged_none_expected_layer_ids_boundary = tuple(
        range(spec.science_layers)
    )
    rbo._paged_none_expected_layer_catalog_ids_t_boundary = tuple(
        expected_catalogs
    )
    rbo._paged_none_reasoning_source_layer_count = spec.science_layers
    rbo._apply_paged_none_model_wide_residency_budget_boundary(
        tuple(runtimes)
    )
    validated_generation_t = (
        rbo.validate_paged_none_graph_frontier_boundary()
    )
    if not torch.equal(
        validated_generation_t.detach().cpu().long().reshape(()),
        torch.tensor(RELEASE_188_GENERATION, dtype=torch.long),
    ):
        raise RuntimeError("Release 188 page frontier generation differs")

    rbo.requires_grad_(False)
    rbo.eval()
    return Release188Model(
        spec=spec,
        parent=parent,
        rbo=rbo,
        page_store=page_store,
        device=device,
        dtype=dtype,
        session_id_t=session_id_t,
    )


class Release188Model:
    """Loaded Release 188 graph with tensor-native and text generation seams."""

    def __init__(
        self,
        *,
        spec: Release188RuntimeSpec,
        parent: Release188ParentAdapter,
        rbo: ResynthesisRBO,
        page_store: Release188SafetensorsPageStore,
        device: torch.device,
        dtype: torch.dtype,
        session_id_t: torch.Tensor,
    ) -> None:
        self.spec = spec
        self.parent = parent
        self.rbo = rbo
        self.page_store = page_store
        self.device = device
        self.dtype = dtype
        self.session_id_t = (
            session_id_t.detach().cpu().long().reshape(-1).clone()
        )

    @classmethod
    def load(
        cls,
        spec_path: str | Path,
        *,
        device: str | torch.device | None = None,
        session_id_t: torch.Tensor | None = None,
    ) -> Release188Model:
        """Load one inference-only Release 188 runtime."""

        spec = load_release_188_runtime_spec(spec_path)
        resolved_device = (
            select_release_188_device_boundary()
            if device is None
            else torch.device(device)
        )
        dtype = _release_dtype_boundary(
            resolved_device,
            spec.dtype_policy,
        )
        resolved_session_id_t = (
            torch.randint(
                0,
                torch.iinfo(torch.int32).max,
                (4,),
                dtype=torch.long,
            )
            if session_id_t is None
            else session_id_t.detach().cpu().long().reshape(-1).clone()
        )
        if (
            resolved_session_id_t.shape != (4,)
            or torch.unique(resolved_session_id_t).numel() < 2
        ):
            raise ValueError(
                "Release 188 inference session identity is malformed"
            )
        return _construct_release_188_model(
            spec,
            device=resolved_device,
            dtype=dtype,
            session_id_t=resolved_session_id_t,
        )

    def generate_token_ids_boundary(
        self,
        prompt_ids_t: torch.Tensor,
        *,
        emission_observer: (
            Callable[[RBOEmissionTracePacket], object] | None
        ) = None,
    ) -> RBOGenerationResult:
        """Generate through the uncapped model-owned additive stop boundary."""

        if (
            prompt_ids_t.dtype != torch.long
            or prompt_ids_t.ndim != 2
            or prompt_ids_t.shape[0] != 1
            or prompt_ids_t.shape[1] < 1
        ):
            raise ValueError(
                "Release 188 prompt IDs must be one nonempty int64 row"
            )
        if bool(
            torch.logical_or(
                prompt_ids_t.lt(0),
                prompt_ids_t.ge(RELEASE_188_LEXICAL_ROWS),
            ).any()
        ):
            raise ValueError(
                "Release 188 prompt contains a projection-only token ID"
            )
        self.rbo.begin_session()
        return resynthesis_rbo_generate(
            self.rbo,
            prompt_ids_t.to(device=self.device),
            emission_observer=emission_observer,
        )

    def _encode_prompt_boundary(self, prompt: str) -> torch.Tensor:
        if not isinstance(prompt, str) or not prompt.strip():
            raise ValueError("Release 188 prompt must be nonempty text")
        tokenizer = self.parent.tokenizer
        backend = tokenizer.backend
        if backend.chat_template:
            encoded = backend.apply_chat_template(
                [{"role": "user", "content": prompt}],
                tokenize=True,
                add_generation_prompt=True,
                return_tensors="pt",
            )
            if not isinstance(encoded, torch.Tensor):
                raise RuntimeError(
                    "Release 188 chat template returned no tensor"
                )
            return encoded.long()
        encoded_batch = tokenizer(
            prompt,
            return_tensors="pt",
            add_special_tokens=True,
        )
        encoded_ids = encoded_batch.get("input_ids")
        if not isinstance(encoded_ids, torch.Tensor):
            raise RuntimeError("Release 188 tokenizer returned no input IDs")
        return encoded_ids.long()

    def solve(
        self,
        prompt: str,
        *,
        emission_observer: (
            Callable[[RBOEmissionTracePacket], object] | None
        ) = None,
    ) -> str:
        """Encode one prompt, traverse Release 188, and decode its own surface."""

        prompt_ids_t = self._encode_prompt_boundary(prompt)
        generated = self.generate_token_ids_boundary(
            prompt_ids_t,
            emission_observer=emission_observer,
        )
        completed_row = generated.select_batch_row_boundary(
            0,
            prompt_width=prompt_ids_t.shape[1],
        )
        generated_ids_t = completed_row.token_ids[
            0,
            prompt_ids_t.shape[1] :,
        ]
        return self.parent.tokenizer.decode(
            generated_ids_t,
            skip_special_tokens=True,
        )


def load_release_188_model(
    spec_path: str | Path,
    *,
    device: str | torch.device | None = None,
    session_id_t: torch.Tensor | None = None,
) -> Release188Model:
    """Public convenience loader for ``runtime/model.json``."""

    return Release188Model.load(
        spec_path,
        device=device,
        session_id_t=session_id_t,
    )