Blame for mindxtrain/deploy/bankml_push.py

"""Serve a trained model through bankml — the verified CPU engine.

`push_to_bankml` is the bankml twin of `ollama_push.push_to_ollama`: merge the LoRA (or take an
already-merged directory) → write a Modelfile that bankml can reproduce → `bankml create`, so the
tag answers on `bankml serve --native` (OpenAI `/v1` + Ollama `/api`, a receipt on every answer).

bankml (https://github.com/cryptoAGI/bankml) is reached **only** as a CLI subprocess here and as
HTTP in `operator/backends/bankml.py`; none of its code is vendored (clean-room policy).

What bankml takes, and what this module therefore refuses *before* running anything (it never
drops an instruction silently — a dropped PARAMETER would serve a model that answers differently
from the one described):

- **Taken:** `FROM` (the merged safetensors directory, which `bankml create` converts to GGUF F16
  byte-identically to llama.cpp b11192, or a GGUF `bankml convert` already pinned), `SYSTEM`,
  `MESSAGE`, `LICENSE`, `REQUIRES`, `PARAMETER` temperature / top_k / top_p / min_p / seed /
  num_ctx / num_predict, and `stop`; from bankml 0.3.6 (its O2) also the penalties —
  repeat_penalty, repeat_last_n, presence_penalty, frequency_penalty — which it applies token
  for token as llama-server b11192 does. A small generation (SmolLM2 / mindx-genN) needs
  `repeat_penalty` to answer in words rather than runs of `,`.
- **Refused:** `ADAPTER` (merge first — this module does it for you given `adapter_dir`), a
  `TEMPLATE` other than the base's own, the penalties on a bankml older than 0.3.6,
  mirostat*, typical_p, resource options (num_gpu, num_thread, num_batch,
  num_keep, draft_num_predict), and any architecture other than Llama (SmolLM2 = mindx-genN):
  bankml converts Llama safetensors only; Qwen3 it serves only as a pinned Q1_0/Q2_0_g64 GGUF.

The merged directory is converted through a link named after the tag: llama.cpp's converter
names a model after the directory it reads (`merged` → general.name "Merged"), and only the
tag's name gives the GGUF byte-identical to the one mindX pins (gen39 → 6b64c748…).

`bankml create` / `bankml convert` arrive in bankml 0.3.5. An older binary is detected from
`bankml --help` (the verbs it lists) and reported as `status="bankml_too_old"`, never raised.

Every public function returns a result; nothing here raises past `BankmlPushResult`.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

from mindxtrain.deploy.modelfile import ModelfileSpec, render_modelfile

BANKML_URL = "https://github.com/cryptoAGI/bankml"

# PARAMETERs bankml reproduces token-for-token (plus `stop`, carried on ModelfileSpec.stop).
BANKML_PARAMS: frozenset[str] = frozenset(
    {"temperature", "top_k", "top_p", "min_p", "seed", "num_ctx", "num_predict"},
)

_PENALTY = frozenset({"repeat_penalty", "presence_penalty", "frequency_penalty", "repeat_last_n"})
PENALTIES_SINCE = (0, 3, 6)  # bankml O2: penalties token-identical to llama-server b11192
_MIROSTAT = frozenset({"mirostat", "mirostat_tau", "mirostat_eta"})
_RESOURCE = frozenset({"num_gpu", "num_thread", "num_batch", "num_keep", "draft_num_predict"})

# Architectures `bankml convert` / `bankml create FROM <dir>` accept (config.json `architectures`).
BANKML_CONVERT_ARCHS: frozenset[str] = frozenset({"LlamaForCausalLM"})

_NAME = re.compile(r"^[a-z0-9_][a-z0-9._-]{0,127}$")
_SHA = r"[0-9a-f]{64}"

PushStatus = Literal[
    "created", "refused", "bankml_missing", "bankml_too_old", "merge_failed", "failed", "error",
]


def version_tuple(version: str) -> tuple[int, ...]:
    """'0.3.6' → (0, 3, 6); '' or unparseable → () (treated as older than any release)."""
    m = re.match(r"(\d+)\.(\d+)\.(\d+)", version or "")
    return tuple(int(g) for g in m.groups()) if m else ()


def penalties_supported(version: str) -> bool:
    return version_tuple(version) >= PENALTIES_SINCE


def _param_refusal(name: str) -> str:
    if name in _PENALTY:
        return (f"PARAMETER {name}: penalties are reproduced from bankml "
                f"{'.'.join(map(str, PENALTIES_SINCE))} (O2); this bankml predates it — upgrade, "
                "or leave it out")
    if name in _MIROSTAT:
        return f"PARAMETER {name}: mirostat is not reproduced by bankml; leave it out"
    if name == "typical_p":
        return "PARAMETER typical_p: locally-typical sampling is not reproduced by bankml"
    if name in _RESOURCE:
        return (f"PARAMETER {name}: a resource option is not part of a bankml model "
                "(bankml's answer does not depend on threads, batch or GPU layers)")
    if name == "stop":
        return "PARAMETER stop: pass stop sequences as ModelfileSpec.stop, not in parameters"
    return f"PARAMETER {name}: not a parameter bankml reproduces"


# ---- capability probe -------------------------------------------------------


@dataclass(frozen=True)
class BankmlCapabilities:
    """What the local `bankml` binary can do. `reason` is set when it cannot do what was asked."""

    binary: str | None
    version: str = ""
    has_create: bool = False
    has_convert: bool = False
    reason: str = ""

    @property
    def found(self) -> bool:
        return self.binary is not None


def bankml_capabilities(bankml_bin: str | None = None, *, timeout_s: float = 30.0) -> BankmlCapabilities:
    """Locate `bankml` and read its version and verbs (`bankml version`, `bankml --help`).

    The verbs listed in the usage text are the truth — an unreleased build may still say 0.3.4
    and already carry `create`. Never raises.
    """
    binary = bankml_bin or shutil.which("bankml")
    if not binary:
        return BankmlCapabilities(
            binary=None,
            reason=(f"`bankml` not found on PATH; install it from {BANKML_URL} "
                    "(or pass --bankml-bin)"),
        )
    try:
        ver = subprocess.run([binary, "version"], capture_output=True, text=True,
                             timeout=timeout_s, check=False)
        usage = subprocess.run([binary, "--help"], capture_output=True, text=True,
                               timeout=timeout_s, check=False)
    except (OSError, subprocess.SubprocessError) as exc:
        return BankmlCapabilities(binary=binary, reason=f"cannot run {binary}: {exc}")
    m = re.search(r"bankml\s+(\d+\.\d+\.\d+\S*)", ver.stdout + ver.stderr)
    text = usage.stdout + usage.stderr
    return BankmlCapabilities(
        binary=binary,
        version=m.group(1) if m else "",
        has_create="bankml create" in text,
        has_convert="bankml convert" in text,
    )


# ---- the Modelfile subset ----------------------------------------------------


@dataclass(frozen=True)
class BankmlSanitizeResult:
    """`ok` → `spec` is exactly what bankml will reproduce; else `refusals` say why, one per line."""

    spec: ModelfileSpec
    refusals: tuple[str, ...] = ()

    @property
    def ok(self) -> bool:
        return not self.refusals


def bankml_sanitize(
    spec: ModelfileSpec, *, base_template: str | None = None, penalties: bool = False,
) -> BankmlSanitizeResult:
    """Check a ModelfileSpec against the subset `bankml create` reproduces.

    Refuses (and names) every instruction bankml would not honour; it never drops one, because a
    silently dropped PARAMETER serves a model that answers differently from the spec. A
    `TEMPLATE` passes only when it equals `base_template` (the base GGUF's own chat template).
    `penalties=True` admits the penalties (bankml 0.3.6+; see `penalties_supported`).
    """
    refusals: list[str] = []
    if spec.adapter:
        refusals.append("ADAPTER: bankml does not merge a LoRA at create time; merge it first "
                        "(push_to_bankml does, given adapter_dir) and FROM the merged directory")
    if spec.template and spec.template != base_template:
        refusals.append("TEMPLATE: bankml renders the base model's own chat template "
                        "(byte-identical to llama.cpp); a different template is not reproduced — "
                        "leave TEMPLATE out")
    allowed = BANKML_PARAMS | _PENALTY if penalties else BANKML_PARAMS
    for name in sorted(spec.parameters):
        if name not in allowed:
            refusals.append(_param_refusal(name))
    for label, value in (("SYSTEM", spec.system), ("TEMPLATE", spec.template),
                         ("LICENSE", spec.license)):
        if '"""' in value:
            refusals.append(f'{label}: contains """, which a Modelfile cannot quote')
    for stop in spec.stop:
        if '"' in stop or "\n" in stop:
            refusals.append(f"stop {stop!r}: a quote or newline cannot be written as PARAMETER stop")
    return BankmlSanitizeResult(spec=spec, refusals=tuple(refusals))


def check_merged_arch(merged_dir: Path) -> str | None:
    """None when `merged_dir/config.json` names an architecture bankml converts (or has none to
    read — then bankml decides); else the refusal reason."""
    try:
        cfg = json.loads((merged_dir / "config.json").read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    archs = cfg.get("architectures") or []
    if not archs or any(a in BANKML_CONVERT_ARCHS for a in archs):
        return None
    return (f"architecture {', '.join(map(str, archs))}: bankml converts Llama-architecture "
            "safetensors only (SmolLM2 / mindx-genN); Qwen3 it serves only as a pinned "
            "Q1_0/Q2_0_g64 GGUF — use `serve --to ollama` or `--to vllm` for this model")


# Families mindXtrain trains that bankml cannot convert (checked by name before a merge is spent).
_UNCONVERTIBLE_FAMILIES = ("qwen", "mistral", "phi", "gemma", "glm", "deepseek", "instella")


def base_family_refusal(base_model: str) -> str | None:
    """A cheap, name-based pre-check (before merging): None unless `base_model` names a family
    bankml cannot convert. The merged `config.json` is still checked afterwards."""
    low = base_model.lower()
    if "smollm" in low or "llama" in low:
        return None
    hit = next((f for f in _UNCONVERTIBLE_FAMILIES if f in low), None)
    if hit is None:
        return None
    return (f"base {base_model}: the {hit} family is not a Llama-architecture model bankml can "
            "convert (bankml converts SmolLM2 / mindx-genN; Qwen3 only as a pinned "
            "Q1_0/Q2_0_g64 GGUF) — use `serve --to ollama` or `--to vllm`")


def bankml_registry_dir(registry_dir: Path | None = None) -> Path:
    """bankml's registry (pins + derived models): `--registry`, else $BANKML_FORKS, else
    ~/.local/share/bankml/forks — the same resolution `bankml create` uses."""
    if registry_dir is not None:
        return Path(registry_dir).expanduser()
    env = os.environ.get("BANKML_FORKS")
    if env:
        return Path(env).expanduser()
    return Path.home() / ".local" / "share" / "bankml" / "forks"


# ---- push --------------------------------------------------------------------


@dataclass(frozen=True)
class BankmlPushResult:
    """Outcome of `push_to_bankml`. `status == "created"` is the only success."""

    status: PushStatus
    tag: str
    reason: str = ""
    refusals: tuple[str, ...] = ()
    bankml_version: str = ""
    merged_dir: Path | None = None
    gguf: Path | None = None
    modelfile: Path | None = None
    model_sha256: str = ""
    digest: str = ""
    output: str = ""
    mindx_fallback_swap: dict[str, str] | None = None
    log: tuple[str, ...] = field(default=())

    @property
    def ok(self) -> bool:
        return self.status == "created"


def _run(cmd: list[str], timeout_s: float) -> tuple[int, str, str]:
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_s, check=False)
    return proc.returncode, proc.stdout or "", proc.stderr or ""


def push_to_bankml(
    base_model: str,
    tag: str,
    *,
    adapter_dir: Path | None = None,
    merged_dir: Path | None = None,
    system: str | None = None,
    params: dict[str, float | int | str] | None = None,
    stop: list[str] | None = None,
    bankml_bin: str | None = None,
    convert: bool = False,
    work_dir: Path | None = None,
    registry_dir: Path | None = None,
    register_with_mindx: bool = False,
    mindx_base_url: str | None = None,
    sink: Callable[[str], None] | None = None,
    timeout_s: float = 1800.0,
) -> BankmlPushResult:
    """Merge (if given an adapter) → Modelfile (bankml subset) → `bankml create <tag>`.

    Give exactly one of `adapter_dir` (a PEFT LoRA over `base_model`; merging needs
    `uv sync --extra ml`) or `merged_dir` (an already-merged HF safetensors directory).

    `convert=False` (default) writes `FROM <merged dir>` and lets `bankml create` convert and pin
    it. `convert=True` runs `bankml convert` first (GGUF + FORK.json into the registry) and
    writes `FROM <gguf>` — the explicit, inspectable two-step.

    The model sha256 recorded is the one bankml prints for the base it pinned and verified.
    With `register_with_mindx`, mindX's fallback model is swapped to `{provider: "bankml",
    model: tag}` (best-effort; a failure is logged, the push still stands).
    """
    lines: list[str] = []

    def emit(line: str) -> None:
        lines.append(line)
        if sink:
            sink(line)

    def done(status: PushStatus, **kw: object) -> BankmlPushResult:
        return BankmlPushResult(status=status, tag=tag, log=tuple(lines), **kw)  # type: ignore[arg-type]

    name = tag.strip().lower().removesuffix(":latest")
    if not _NAME.match(name) or name.endswith(".gguf"):
        return done("refused", reason=(f"tag {tag!r}: bankml names use letters, digits, '.', '_' "
                                       "and '-' (at most 128, no ':tag', not ending in .gguf)"))
    if (adapter_dir is None) == (merged_dir is None):
        return done("error", reason="give exactly one of adapter_dir or merged_dir")

    # 1. the subset — checked before anything expensive runs
    spec = ModelfileSpec(from_model="<merged>", system=system or "", parameters=dict(params or {}),
                         stop=list(stop or []))
    checked = bankml_sanitize(spec, penalties=True)  # penalties are gated on the version below
    if not checked.ok:
        for r in checked.refusals:
            emit(f"[push-bankml] refuse: {r}")
        return done("refused", reason="the Modelfile asks for what bankml does not reproduce",
                    refusals=checked.refusals)

    # 2. the binary, and whether it has the verbs
    caps = bankml_capabilities(bankml_bin)
    if not caps.found:
        return done("bankml_missing", reason=caps.reason)
    if caps.reason:
        return done("error", reason=caps.reason, bankml_version=caps.version)
    missing = [v for v, have in (("create", caps.has_create), ("convert", caps.has_convert))
               if (v == "create" or convert) and not have]
    if missing:
        return done("bankml_too_old", bankml_version=caps.version, reason=(
            f"bankml {caps.version or '(unknown version)'} has no `{'`/`'.join(missing)}` verb "
            f"(they arrive in bankml 0.3.5); upgrade from {BANKML_URL}"))
    asked = sorted(_PENALTY & set(spec.parameters))
    if asked and not penalties_supported(caps.version):
        refusals = tuple(_param_refusal(n) for n in asked)
        for r in refusals:
            emit(f"[push-bankml] refuse: {r}")
        return done("refused", bankml_version=caps.version, refusals=refusals,
                    reason=f"bankml {caps.version or '(unknown version)'} does not reproduce penalties")
    emit(f"[push-bankml] bankml {caps.version} at {caps.binary}")

    work = Path(work_dir) if work_dir else (
        (adapter_dir or merged_dir).parent / "bankml_push")  # type: ignore[union-attr]

    # 3. merge, when given an adapter
    if adapter_dir is not None:
        try:
            from mindxtrain.deploy.ollama_push import merge_lora_adapter

            merged = merge_lora_adapter(base_model, Path(adapter_dir), work / "merged", sink=emit)
        except ImportError as exc:
            return done("merge_failed", bankml_version=caps.version, reason=(
                f"{exc} — merging a LoRA needs `uv sync --extra ml`"))
        except Exception as exc:  # a merge failure is a result, not a crash
            return done("merge_failed", bankml_version=caps.version,
                        reason=f"{type(exc).__name__}: {exc}")
    else:
        merged = Path(merged_dir)  # type: ignore[arg-type]
        if not merged.is_dir():
            return done("error", bankml_version=caps.version,
                        reason=f"merged_dir {merged} is not a directory")

    arch_refusal = check_merged_arch(merged)
    if arch_refusal:
        emit(f"[push-bankml] refuse: {arch_refusal}")
        return done("refused", bankml_version=caps.version, merged_dir=merged,
                    reason=arch_refusal, refusals=(arch_refusal,))

    registry = bankml_registry_dir(registry_dir)
    reg_args = ["--registry", str(registry)] if registry_dir is not None else []
    gguf: Path | None = None
    model_sha = ""
    try:
        # 4. optional explicit conversion — through a link named after the tag, so the GGUF's
        #    general.name (and so its bytes and sha256) is the one mindX pins
        source = str(_named_source(merged, work / name, name))
        if convert:
            registry.mkdir(parents=True, exist_ok=True)
            gguf = registry / f"{name}-base-F16.gguf"
            cmd = [str(caps.binary), "convert", source, "-o", str(gguf),
                   "--fork", str(registry / f"{gguf.name}.FORK.json"), "--source", source]
            emit(f"[push-bankml] $ {' '.join(cmd)}")
            code, out, err = _run(cmd, timeout_s)
            if err.strip():
                emit(err.rstrip())
            if code != 0:
                status: PushStatus = "refused" if code == 2 else "failed"
                return done(status, bankml_version=caps.version, merged_dir=merged,
                            reason=_last_line(err) or f"bankml convert exited {code}",
                            output=(out + err)[-2000:])
            m = re.search(rf"^({_SHA})\s", out, re.MULTILINE)
            model_sha = m.group(1) if m else ""

        # 5. the Modelfile and `bankml create`
        spec = spec.model_copy(update={"from_model": str(gguf) if gguf else source})
        modelfile = work / name / "Modelfile"
        modelfile.parent.mkdir(parents=True, exist_ok=True)
        modelfile.write_text(render_modelfile(spec), encoding="utf-8")
        cmd = [str(caps.binary), "create", name, "-f", str(modelfile), *reg_args]
        emit(f"[push-bankml] $ {' '.join(cmd)}")
        code, out, err = _run(cmd, timeout_s)
    except (OSError, subprocess.SubprocessError) as exc:
        return done("error", bankml_version=caps.version, merged_dir=merged, gguf=gguf,
                    reason=f"{type(exc).__name__}: {exc}")
    for chunk in (out, err):
        if chunk.strip():
            emit(chunk.rstrip())
    if code != 0:
        return done("refused" if code == 2 else "failed", bankml_version=caps.version,
                    merged_dir=merged, gguf=gguf, modelfile=modelfile,
                    reason=_last_line(err) or f"bankml create exited {code}",
                    output=(out + err)[-2000:])

    both = out + "\n" + err
    over = re.search(rf"over (\S+) \(sha256 ({_SHA})\)", both)
    digest = re.search(rf"digest sha256:({_SHA})", both)
    if over:
        model_sha = over.group(2)
        if gguf is None:
            gguf = registry / over.group(1)
    elif not model_sha and gguf is not None:
        code_s, out_s, _ = _run([str(caps.binary), "sha256", str(gguf)], 600.0)
        m = re.match(rf"({_SHA})\s", out_s)
        model_sha = m.group(1) if (code_s == 0 and m) else ""

    swap: dict[str, str] | None = None
    if register_with_mindx:
        try:
            from mindxtrain.deploy.api_client import swap_mindx_fallback_model

            swap = swap_mindx_fallback_model(provider="bankml", model=name, api_url=mindx_base_url)
            emit(f"[push-bankml] mindX swap: {swap.get('previous', '?')} -> "
                 f"{swap.get('current', '?')}")
        except Exception as exc:  # best-effort, as push_to_ollama
            emit(f"[push-bankml] mindX registration failed (push still ok): {exc}")
            swap = None

    emit(f"[push-bankml] created {name} over sha256 {model_sha or '(not reported)'}")
    return done("created", bankml_version=caps.version, merged_dir=merged, gguf=gguf,
                modelfile=modelfile, model_sha256=model_sha,
                digest=digest.group(1) if digest else "", output=both.strip()[-2000:],
                mindx_fallback_swap=swap)


def _named_source(merged: Path, at: Path, name: str) -> Path:
    """A link `at/<name>` → `merged`, unless `merged` is already called `name`."""
    target = merged.resolve()
    if target.name == name:
        return target
    at.mkdir(parents=True, exist_ok=True)
    link = at / name
    if link.is_symlink():
        link.unlink()
    elif link.exists():
        return link  # a real directory by that name: the caller put it there
    link.symlink_to(target, target_is_directory=True)
    return link


def _last_line(text: str) -> str:
    rows = [r.strip() for r in text.strip().splitlines() if r.strip()]
    return rows[-1] if rows else ""


__all__ = [
    "BANKML_CONVERT_ARCHS",
    "BANKML_PARAMS",
    "PENALTIES_SINCE",
    "BankmlCapabilities",
    "BankmlPushResult",
    "BankmlSanitizeResult",
    "bankml_capabilities",
    "bankml_registry_dir",
    "bankml_sanitize",
    "base_family_refusal",
    "check_merged_arch",
    "penalties_supported",
    "push_to_bankml",
    "version_tuple",
]