BitTransformerLM/mcp_server.py

import io
import os
import gzip
import uuid
import traceback
from concurrent.futures import ThreadPoolExecutor
from flask import Flask, request, jsonify, send_file
import matplotlib.pyplot as plt
import torch

from bit_transformer.dashboard_app import ModelManager
from bit_transformer.dashboard import plot_telemetry
from bit_transformer.hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
from bit_transformer.optimization import configure_optimizer
from bit_transformer.bit_io import text_to_bits

app = Flask(__name__)
manager = ModelManager()

# background job management
executor = ThreadPoolExecutor(max_workers=4)
jobs: dict[str, dict] = {}


def _submit_job(fn, *args, **kwargs) -> str:
    """Schedule a function for background execution and return a job id."""
    job_id = str(uuid.uuid4())
    jobs[job_id] = {"status": "queued", "result": None, "error": None, "logs": []}

    def wrapper():
        jobs[job_id]["status"] = "running"
        try:
            jobs[job_id]["result"] = fn(*args, **kwargs)
            jobs[job_id]["status"] = "completed"
        except Exception as err:  # pragma: no cover - captured for client
            jobs[job_id]["status"] = "error"
            jobs[job_id]["error"] = str(err)
            jobs[job_id]["trace"] = traceback.format_exc()

    executor.submit(wrapper)
    return job_id


@app.errorhandler(Exception)
def handle_exception(err):
    """Return JSON error responses with stack traces."""
    return (
        jsonify({"error": str(err), "trace": traceback.format_exc()}),
        getattr(err, "code", 500),
    )


@app.route("/init", methods=["POST"])
def init_model():
    data = request.json or {}
    int_fields = {
        "d_model",
        "nhead",
        "num_layers",
        "dim_feedforward",
        "max_seq_len",
        "chunk_size",
        "overlap",
    }
    float_fields = {"act_threshold"}
    bool_fields = {"reversible", "use_checkpoint"}
    params = {}
    for k, v in data.items():
        if v is None:
            params[k] = None
        elif k in int_fields:
            params[k] = int(v)
        elif k in float_fields:
            params[k] = float(v)
        elif k in bool_fields:
            params[k] = bool(v)
        else:
            params[k] = v
    manager.init_model(params)
    return jsonify({"status": "initialized", "params": params})

@app.route("/train", methods=["POST"])
def train_model():
    bits = request.json["bits"]

    def task():
        tensor = torch.tensor(bits, dtype=torch.long)
        loss, ratio = manager.train_step(tensor)
        return {"loss": loss, "ratio": ratio}

    job_id = _submit_job(task)
    return jsonify({"job_id": job_id})


@app.route("/train_epochs", methods=["POST"])
def train_epochs_route():
    data = request.json
    bits = data["bits"]
    epochs = int(data.get("epochs", 1))
    compress_prob = float(data.get("compress_prob", 0.5))
    direct_prob = float(data.get("direct_prob", 0.0))

    def task():
        tensor = torch.tensor(bits, dtype=torch.long)
        metrics = manager.train_epochs(
            tensor,
            epochs=epochs,
            compress_prob=compress_prob,
            direct_prob=direct_prob,
        )
        return {"metrics": metrics}

    job_id = _submit_job(task)
    return jsonify({"job_id": job_id})

@app.route("/scale_up", methods=["POST"])
def scale_up():
    width_mult = float(request.json.get("width_mult", 1.0))

    def task():
        manager.scale_up(width_mult)
        return {
            "status": "scaled",
            "layers": manager.model.num_layers,
            "d_model": manager.model.d_model,
        }

    job_id = _submit_job(task)
    return jsonify({"job_id": job_id})

@app.route("/collapse", methods=["POST"])
def collapse_model():
    cluster_bits = request.json["clusters"]
    params = {k: int(v) for k, v in request.json["params"].items()}
    width_scale = float(request.json.get("width_scale", 1.0))

    def task():
        manager.collapse(cluster_bits, params, width_scale)
        return {"status": "collapsed"}

    job_id = _submit_job(task)
    return jsonify({"job_id": job_id})


@app.route("/job/<job_id>", methods=["GET"])
def get_job(job_id: str):
    job = jobs.get(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify(job)


@app.route("/jobs", methods=["GET"])
def list_jobs():
    return jsonify(jobs)

@app.route("/lambdas", methods=["GET", "POST"])
def update_lambdas():
    if request.method == "POST":
        data = request.json
        manager.set_lambdas(float(data["lambda_K"]), float(data["lambda_C"]), float(data["lambda_S"]))
        return jsonify({"status": "updated"})
    else:
        return jsonify({
            "lambda_K": manager.lambda_K,
            "lambda_C": manager.lambda_C,
            "lambda_S": manager.lambda_S,
        })

@app.route("/diffusion", methods=["GET", "POST"])
def update_diffusion():
    if request.method == "POST":
        manager.set_diffusion(bool(request.json.get("diffusion", False)))
        return jsonify({"status": "updated"})
    return jsonify({"diffusion": manager.diffusion})


@app.route("/qat", methods=["GET", "POST"])
def update_qat():
    if request.method == "POST":
        manager.set_qat(bool(request.json.get("qat", False)))
        return jsonify({"status": "updated"})
    return jsonify({"qat": manager.qat})


@app.route("/gpu", methods=["GET", "POST"])
def update_gpu():
    if request.method == "POST":
        manager.set_gpu(bool(request.json.get("use_gpu", False)))
        return jsonify({"status": "updated"})
    return jsonify({"use_gpu": manager.use_gpu})

@app.route("/infer", methods=["POST"])
def inference():
    bits = torch.tensor(request.json["bits"], dtype=torch.long)
    result = manager.infer(bits)
    return jsonify(result)


@app.route("/infer_long", methods=["POST"])
def inference_long():
    bits = torch.tensor(request.json["bits"], dtype=torch.long)
    ctx = int(request.json.get("ctx_bits", 4096))
    overlap = int(request.json.get("overlap", 256))
    result = manager.infer_long(bits, ctx_bits=ctx, overlap=overlap)
    return jsonify(result)

@app.route("/infer_text", methods=["POST"])
def inference_text():
    text = request.json.get("text", "")
    result = manager.infer_text(text)
    return jsonify(result)

@app.route("/status", methods=["GET"])
def status():
    return jsonify(manager.get_status())


@app.route("/model_config", methods=["GET"])
def model_config():
    return jsonify(manager.get_model_config())


@app.route("/metrics", methods=["GET"])
def metrics():
    return jsonify(manager.get_metrics())


@app.route("/save_checkpoint", methods=["POST"])
def save_checkpoint_route():
    repo_id = request.json.get("repo_id")
    token = request.json.get("token") or os.getenv("HF_TOKEN")
    if manager.model is None:
        return jsonify({"error": "model not initialized"}), 400
    if token:
        hf_login(token=token)
    save_checkpoint(manager.model, repo_id=repo_id)
    return jsonify({"status": "saved"})


@app.route("/download_checkpoint", methods=["POST"])
def download_checkpoint_route():
    repo_id = request.json.get("repo_id")
    token = request.json.get("token") or os.getenv("HF_TOKEN")
    if token:
        hf_login(token=token)
    dest = manager.weights_path + ".gz"
    ok = download_checkpoint(dest, repo_id=repo_id)
    if not ok:
        return jsonify({"status": "failed"}), 500
    if manager.model is None:
        return jsonify({"status": "downloaded", "loaded": False})
    with gzip.open(dest, "rb") as f:
        state = torch.load(f, map_location="cpu")
    manager.model.load_state_dict(state)
    manager.optimizer, manager.scheduler = configure_optimizer(
        manager.model, lr=1e-3, total_steps=manager.total_steps
    )
    manager._apply_device()
    manager._save_state()
    return jsonify({"status": "downloaded", "loaded": True})

@app.route("/plot.png")
def plot_png():
    fig, _ = plot_telemetry(manager.metrics)
    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)
    return send_file(buf, mimetype="image/png")


@app.route("/text_to_bits", methods=["POST"])
def text_to_bits_route():
    text = request.json.get("text", "")
    if len(text) > 100_000:
        return jsonify({"error": "text too large"}), 413
    return jsonify({"bits": text_to_bits(text)})


@app.route("/dataset", methods=["GET"])
def dataset_route():
    name = request.args.get("name", "")
    split = request.args.get("split", "train")
    size = int(request.args.get("size", 1))
    seq_len = int(request.args.get("seq_len", 64))
    if size * seq_len > 1_000_000:
        return jsonify({"error": "dataset too large"}), 413
    if name == "wikitext2":
        try:
            from datasets import load_dataset

            ds = load_dataset("wikitext", "wikitext-2-raw-v1", split=split)
            lines = [t for t in ds["text"] if t.strip()][:size]
        except Exception:
            bits = torch.randint(0, 2, (size, seq_len), dtype=torch.long)
            return jsonify({"bits": bits.tolist()})
        bits_list = []
        for text in lines:
            b = text_to_bits(text)[:seq_len]
            if len(b) < seq_len:
                b.extend([0] * (seq_len - len(b)))
            bits_list.append(b)
        if len(bits_list) < size:
            pad = size - len(bits_list)
            bits_list.extend(torch.randint(0, 2, (pad, seq_len), dtype=torch.long).tolist())
        return jsonify({"bits": bits_list})
    return jsonify({"error": "unknown dataset"}), 400


@app.route("/health")
def health_check():
    return jsonify({"status": "ok"})


def run_mcp_server(host: str = "0.0.0.0", port: int = 7000) -> None:
    app.run(host=host, port=port, debug=True)


if __name__ == "__main__":
    import torch
    run_mcp_server()