swactor/apps/mvp-node/tinygrad_worker.py

#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import json
import os
import sys
import time
import traceback
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

Tensor: Any = None
dtypes: Any = None
model: Any = None
tokenizer: Any = None
role: dict[str, Any] = {}
loaded: dict[str, Any] = {}


def control(**event: Any) -> None:
    print(json.dumps(event, separators=(",", ":")), flush=True)


def log(message: str) -> None:
    print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True)


def fatal(reason: str, **fields: Any) -> None:
    control(type="WorkerFatal", reason=reason, **fields)
    raise SystemExit(1)


def test_mode() -> bool:
    return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"}



def initialize(cmd: dict[str, Any]) -> None:
    global Tensor, dtypes
    if int(cmd.get("helper_abi_version", 1)) != 1:
        fatal("UnsupportedHelperAbi", helper_abi_version=cmd.get("helper_abi_version"))
    device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA")
    os.environ["DEV"] = device
    started = time.monotonic()
    control(type="TinygradImportStarted", device=device)
    from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes

    Tensor = TinyTensor
    dtypes = tiny_dtypes
    value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist()
    control(
        type="WorkerReady",
        pid=os.getpid(),
        backend={"device": device},
        cuda_probe=value,
        elapsed_ms=int((time.monotonic() - started) * 1000),
    )


def configure_role(cmd: dict[str, Any]) -> None:
    config = cmd.get("config", {})
    role.clear()
    role.update(
        role_id=int(cmd.get("role_id", 1)),
        run_id=int(config.get("run_id", 1)),
        stage_index=int(config.get("stage_index", 0)),
        layer_start=int(config.get("layer_start", 0)),
        layer_end_exclusive=int(config.get("layer_end_exclusive", 0)),
    )
    control(type="RoleConfigured", role_id=role["role_id"], stage_index=role["stage_index"])


def cache_root() -> Path:
    raw = os.environ.get("MVP_MODEL_CACHE_DIR", "").strip()
    root = Path(raw).expanduser() if raw else Path.home() / ".cache" / "mvp-node"
    root.mkdir(parents=True, exist_ok=True)
    return root


def hf_url(repo: str, file: str, revision: str | None) -> str:
    encoded_file = "/".join(urllib.parse.quote(part) for part in file.split("/"))
    return f"https://huggingface.co/{repo}/resolve/{revision or 'main'}/{encoded_file}"


def source_url(source: dict[str, Any]) -> str | None:
    if "HuggingFaceGguf" not in source:
        return None
    hf = source["HuggingFaceGguf"]
    return hf_url(str(hf["repo"]), str(hf["file"]), hf.get("revision"))


def source_path(source: dict[str, Any]) -> Path | None:
    if "LocalPath" not in source:
        return None
    return Path(str(source["LocalPath"])).expanduser()


def cache_path_for(url: str) -> Path:
    parsed = urllib.parse.urlparse(url)
    basename = Path(parsed.path).name or "model.gguf"
    digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16]
    return cache_root() / f"{digest}-{basename}"


def request_headers() -> dict[str, str]:
    headers = {"User-Agent": "swactor-mvp-node/0.1"}
    token = os.environ.get("HF_TOKEN", "").strip()
    if token:
        headers["Authorization"] = f"Bearer {token}"
    return headers


def fetch_whole(source: dict[str, Any]) -> Path:
    local = source_path(source)
    if local is not None:
        if not local.is_file():
            fatal("GgufLocalPathMissing", path=str(local))
        control(type="GgufCacheReady", path=str(local), cache_hit=True, source="local")
        return local

    url = source_url(source)
    if not url:
        fatal("UnsupportedGgufSource", source=source)
    target = cache_path_for(url)
    if target.is_file() and target.stat().st_size > 0:
        control(type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=True, url=url)
        return target

    partial = target.with_name(target.name + ".partial")
    started = time.monotonic()
    req = urllib.request.Request(url, headers=request_headers())
    control(type="GgufDownloadStarted", url=url, path=str(target))
    try:
        with urllib.request.urlopen(req, timeout=60) as response, partial.open("wb") as out:
            total = int(response.headers.get("Content-Length") or 0)
            done = 0
            last_event = 0.0
            while True:
                chunk = response.read(1024 * 1024)
                if not chunk:
                    break
                out.write(chunk)
                done += len(chunk)
                now = time.monotonic()
                if now - last_event >= float(os.environ.get("MVP_DOWNLOAD_PROGRESS_SECS", "5")):
                    control(
                        type="GgufDownloadProgress",
                        bytes_done=done,
                        bytes_total=total,
                        elapsed_ms=int((now - started) * 1000),
                    )
                    last_event = now
        partial.replace(target)
    except Exception as exc:
        try:
            partial.unlink(missing_ok=True)
        except Exception:
            pass
        fatal("GgufDownloadFailed", url=url, error=str(exc))
    control(
        type="GgufCacheReady",
        path=str(target),
        bytes=target.stat().st_size,
        cache_hit=False,
        elapsed_ms=int((time.monotonic() - started) * 1000),
        url=url,
    )
    return target


def require_tinygrad() -> Any:
    if Tensor is None:
        fatal("BackendNotInitialized")
    return Tensor


def load_weights(cmd: dict[str, Any]) -> None:
    global model, tokenizer
    TensorCls = require_tinygrad()
    started = time.monotonic()
    model_id = str(cmd["model_id"])
    if test_mode():
        model = {"test_mode": True}
        tokenizer = {"test_mode": True}
        loaded.clear()
        loaded.update(
            model_id=model_id,
            path="mvp-tinygrad-test-mode",
            layer_start=int(cmd.get("layer_start", 0)),
            layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)),
        )
        control(
            type="WeightsLoaded",
            model_id=model_id,
            path=loaded["path"],
            test_mode=True,
            elapsed_ms=int((time.monotonic() - started) * 1000),
        )
        return
    source = cmd["gguf_source"]
    path = fetch_whole(source)
    try:
        from tinygrad.apps.llm import SimpleTokenizer, Transformer

        max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512")
        max_context = int(max_context_raw) if max_context_raw else 512
        model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True)
        tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None})
        if "EmbeddedGguf" in tok_src:
            tokenizer = SimpleTokenizer.from_gguf_kv(kv)
        else:
            fatal("UnsupportedTokenizerSource", tokenizer=tok_src)
    except SystemExit:
        raise
    except Exception as exc:
        tb = traceback.format_exc()
        print(tb, file=sys.stderr, flush=True)
        fatal("ModelLoadFailed", error=str(exc), traceback=tb)
    loaded.clear()
    loaded.update(
        model_id=model_id,
        path=str(path),
        layer_start=int(cmd.get("layer_start", 0)),
        layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)),
    )
    control(
        type="WeightsLoaded",
        model_id=model_id,
        path=str(path),
        elapsed_ms=int((time.monotonic() - started) * 1000),
    )


def infer_prompt(cmd: dict[str, Any]) -> None:
    if model is None or tokenizer is None:
        fatal("WeightsNotLoaded")
    prompt = str(cmd.get("prompt", ""))
    max_tokens = int(cmd.get("max_tokens", 1))
    started = time.monotonic()
    if test_mode():
        text = f"mvp-test response: {prompt}"
        control(
            type="PromptCompleted",
            model_id=loaded.get("model_id"),
            prompt_tokens=[],
            generated_tokens=list(range(min(max_tokens, 3))),
            text=text,
            test_mode=True,
            elapsed_ms=int((time.monotonic() - started) * 1000),
        )
        return
    prompt_tokens = tokenizer.encode(prompt)
    generated: list[int] = []
    for token in model.generate(list(prompt_tokens)):
        generated.append(int(token))
        if len(generated) >= max_tokens:
            break
    text = tokenizer.decode(generated) if generated else ""
    control(
        type="PromptCompleted",
        model_id=loaded.get("model_id"),
        prompt_tokens=prompt_tokens,
        generated_tokens=generated,
        text=text,
        elapsed_ms=int((time.monotonic() - started) * 1000),
    )


def shutdown_worker(_: dict[str, Any]) -> None:
    control(type="WorkerStopped", reason="Graceful")
    raise SystemExit(0)


HANDLERS = {
    "InitializeWorker": initialize,
    "ConfigureRole": configure_role,
    "LoadWeights": load_weights,
    "InferPrompt": infer_prompt,
    "ShutdownWorker": shutdown_worker,
}

for raw in sys.stdin:
    if not raw.strip():
        continue
    try:
        command = json.loads(raw)
        handler = HANDLERS.get(command.get("type"))
        if handler is None:
            fatal("UnknownCommand", command=command.get("type"))
        handler(command)
    except SystemExit:
        raise
    except Exception as exc:
        tb = traceback.format_exc()
        print(tb, file=sys.stderr, flush=True)
        fatal("UnhandledWorkerException", error=str(exc), traceback=tb)