neuron-tai/scripts/refresh_glm_target_manifest.py

#!/usr/bin/env python3
"""Refresh (or check) the pinned GLM-5.2 target manifest and architecture snapshot.

Resolves revisions, shard sizes, and LFS SHA-256 digests from the Hugging Face
metadata API. It never downloads a weight payload: sizes and digests come from
``/api/models/.../paths-info``, which returns the LFS pointer metadata, and the only
files fetched in full are the small config/tokenizer/chat-template documents whose
bytes the snapshot hashes.

Usage::

    python scripts/refresh_glm_target_manifest.py --check    # CI: pinned bytes still resolve?
    python scripts/refresh_glm_target_manifest.py --write    # re-pin HEAD after human review

``--check`` validates the already locked revisions through Hugging Face's
revision-specific API. A normal new upstream commit does not mutate the target.
``--write`` intentionally follows moving HEAD and is therefore a reviewed target
change, never an automatic refresh.

This script requires network access and is not part of the default test suite.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import urllib.request
from pathlib import Path
from typing import Any

DATA_DIR = Path(__file__).resolve().parent.parent / "packages/node/meshnet_node/glm_alpha/data"
MANIFEST_PATH = DATA_DIR / "target-manifest.json"
SNAPSHOT_PATH = DATA_DIR / "architecture-snapshot.json"

SOURCE_REPO = "zai-org/GLM-5.2"
GGUF_REPO = "unsloth/GLM-5.2-GGUF"
QUANT = "UD-IQ1_S"
FALLBACK_QUANT = "UD-IQ1_M"
SHARD_COUNT = 6

SNAPSHOT_FILES = (
    "config.json",
    "chat_template.jinja",
    "generation_config.json",
    "tokenizer_config.json",
)

TIMEOUT = 60


def _get(url: str) -> bytes:
    with urllib.request.urlopen(url, timeout=TIMEOUT) as response:  # noqa: S310 - fixed HTTPS host
        return response.read()


def _get_json(url: str) -> Any:
    return json.loads(_get(url))


def _post_json(url: str, payload: dict) -> Any:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=TIMEOUT) as response:  # noqa: S310
        return json.loads(response.read())


def _shard_paths(quant: str) -> list[str]:
    return [
        f"{quant}/GLM-5.2-{quant}-{index:05d}-of-{SHARD_COUNT:05d}.gguf"
        for index in range(1, SHARD_COUNT + 1)
    ]


def _resolve_shards(revision: str, quant: str) -> list[dict]:
    info = _post_json(
        f"https://huggingface.co/api/models/{GGUF_REPO}/paths-info/{revision}",
        {"paths": _shard_paths(quant)},
    )
    by_path = {entry["path"]: entry for entry in info}
    shards = []
    for index, path in enumerate(_shard_paths(quant), start=1):
        entry = by_path.get(path)
        if entry is None:
            raise SystemExit(f"upstream is missing shard {path} at revision {revision}")
        lfs = entry.get("lfs") or {}
        oid = lfs.get("oid")
        if not oid:
            raise SystemExit(
                f"{path} has no LFS oid; a non-LFS shard is not the published artifact"
            )
        shards.append(
            {
                "index": index,
                "path": path,
                "size_bytes": int(entry["size"]),
                "sha256": oid,
                "url": f"https://huggingface.co/{GGUF_REPO}/resolve/{revision}/{path}",
            }
        )
    return shards


def build_documents(
    *,
    source_revision: str | None = None,
    gguf_revision: str | None = None,
) -> tuple[dict, dict]:
    """Resolve documents at explicit pins, or at current HEAD for reviewed re-pinning."""

    source_api = f"https://huggingface.co/api/models/{SOURCE_REPO}"
    gguf_api = f"https://huggingface.co/api/models/{GGUF_REPO}"
    if source_revision is not None:
        source_api += f"/revision/{source_revision}"
    if gguf_revision is not None:
        gguf_api += f"/revision/{gguf_revision}"
    source_info = _get_json(source_api)
    gguf_info = _get_json(gguf_api)
    source_rev = source_info["sha"]
    gguf_rev = gguf_info["sha"]
    if source_revision is not None and source_rev != source_revision:
        raise SystemExit(
            f"source revision endpoint returned {source_rev}, expected {source_revision}"
        )
    if gguf_revision is not None and gguf_rev != gguf_revision:
        raise SystemExit(f"GGUF revision endpoint returned {gguf_rev}, expected {gguf_revision}")

    shards = _resolve_shards(gguf_rev, QUANT)
    total = sum(shard["size_bytes"] for shard in shards)
    fallback = _resolve_shards(gguf_rev, FALLBACK_QUANT)
    fallback_total = sum(shard["size_bytes"] for shard in fallback)

    manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
    manifest["source_model"]["revision"] = source_rev
    manifest["source_model"]["last_modified"] = source_info.get("lastModified")
    manifest["source_model"]["revision_url"] = (
        f"https://huggingface.co/{SOURCE_REPO}/tree/{source_rev}"
    )
    manifest["gguf_artifact"]["revision"] = gguf_rev
    manifest["gguf_artifact"]["last_modified"] = gguf_info.get("lastModified")
    manifest["gguf_artifact"]["revision_url"] = (
        f"https://huggingface.co/{GGUF_REPO}/tree/{gguf_rev}"
    )
    manifest["gguf_artifact"]["shards"] = shards
    manifest["gguf_artifact"]["total_bytes"] = total
    manifest["gguf_artifact"]["total_gib"] = round(total / 1024**3, 3)
    manifest["gguf_artifact"]["total_gb"] = round(total / 1000**3, 3)
    manifest["diagnostic_fallback"]["total_bytes"] = fallback_total
    manifest["diagnostic_fallback"]["total_gib"] = round(fallback_total / 1024**3, 3)
    manifest["diagnostic_fallback"]["total_gb"] = round(fallback_total / 1000**3, 3)

    snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8"))
    snapshot["source_revision"] = source_rev
    source_files = []
    config: dict[str, Any] = {}
    for name in SNAPSHOT_FILES:
        url = f"https://huggingface.co/{SOURCE_REPO}/resolve/{source_rev}/{name}"
        body = _get(url)
        source_files.append(
            {
                "path": name,
                "size_bytes": len(body),
                "sha256": hashlib.sha256(body).hexdigest(),
                "url": url,
            }
        )
        if name == "config.json":
            config = json.loads(body)
    snapshot["source_files"] = source_files

    indexer_types = config["indexer_types"]
    arch = snapshot["architecture"]
    arch["num_hidden_layers"] = config["num_hidden_layers"]
    arch["num_nextn_predict_layers"] = config["num_nextn_predict_layers"]
    arch["total_artifact_layers"] = config["num_hidden_layers"] + config["num_nextn_predict_layers"]
    arch["hidden_size"] = config["hidden_size"]
    arch["n_routed_experts"] = config["n_routed_experts"]
    arch["num_experts_per_tok"] = config["num_experts_per_tok"]
    arch["n_shared_experts"] = config["n_shared_experts"]
    arch["index_topk"] = config["index_topk"]
    arch["index_head_dim"] = config["index_head_dim"]
    arch["kv_lora_rank"] = config["kv_lora_rank"]
    arch["qk_rope_head_dim"] = config["qk_rope_head_dim"]
    arch["mla_cached_values_per_token_per_layer"] = (
        config["kv_lora_rank"] + config["qk_rope_head_dim"]
    )
    arch["indexer_full_layers"] = sum(1 for role in indexer_types if role == "full")
    arch["indexer_shared_layers"] = sum(1 for role in indexer_types if role == "shared")
    arch["indexer_types_sha256"] = hashlib.sha256(
        json.dumps(indexer_types, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    arch["max_position_embeddings"] = config["max_position_embeddings"]
    arch["vocab_size"] = config["vocab_size"]
    arch["first_k_dense_replace"] = config["first_k_dense_replace"]
    arch["dense_layers"] = config["first_k_dense_replace"]
    arch["sparse_moe_layers"] = config["num_hidden_layers"] - config["first_k_dense_replace"]

    return manifest, snapshot


def _dump(document: dict) -> str:
    return json.dumps(document, indent=2, ensure_ascii=False) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--check", action="store_true", help="fail if the pins have drifted")
    group.add_argument("--write", action="store_true", help="re-pin from upstream")
    args = parser.parse_args()

    if args.check:
        pinned = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
        manifest, snapshot = build_documents(
            source_revision=pinned["source_model"]["revision"],
            gguf_revision=pinned["gguf_artifact"]["revision"],
        )
    else:
        # --write intentionally follows current HEAD and therefore requires review.
        manifest, snapshot = build_documents()

    if args.write:
        MANIFEST_PATH.write_text(_dump(manifest), encoding="utf-8")
        SNAPSHOT_PATH.write_text(_dump(snapshot), encoding="utf-8")
        print(f"wrote {MANIFEST_PATH}")
        print(f"wrote {SNAPSHOT_PATH}")
        print("Re-pinning changes the alpha target. Update the alpha contract under review.")
        return 0

    drifted = False
    for path, fresh in ((MANIFEST_PATH, manifest), (SNAPSHOT_PATH, snapshot)):
        current = json.loads(path.read_text(encoding="utf-8"))
        if current != fresh:
            drifted = True
            print(f"DRIFT: {path.name} no longer matches upstream", file=sys.stderr)
    if drifted:
        print(
            "\nPinned revision metadata no longer matches the locked files. Treat this as "
            "artifact-integrity drift; do not heal or re-pin without human review.",
            file=sys.stderr,
        )
        return 1

    print("target manifest and architecture snapshot match upstream")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())