# SPDX-License-Identifier: Apache-2.0
"""OpenAI-compatible API server for the OFFICIAL native vLLM MiniMax-M3 (vLLM
0.23.1) + our AutoRound 3.2bit quantization == the LONG-CONTEXT port (②).

Unlike m3_api_server.py (the out-of-tree M1 model, <=2048 ctx, vision), this
serves the official native M3 text backbone with the MSA lightning indexer, so
it supports long context (~46K tokens of KV on 2x RTX PRO 6000). All the ② fixes
(shared expert: n_shared_experts + down_proj/gate_up quant naming) live in
serve_m3_official.py / m3_official_loader.py / m3_quant.py and are reused here.

Engine args mirror serve_m3_official.py's LLM kwargs; we reuse its
_override_quant_method (architecture force + the FORCE config dict incl.
n_shared_experts=1). The OpenAI server's CLI-parsed engine args are replaced with
ours (same trick as m3_api_server.py) since the CLI path mis-resolves the config.

Start:  M3_OFFICIAL_PORT is set automatically.
  python m3_official_api_server.py
Test :  curl http://localhost:8003/v1/models
        curl http://localhost:8003/v1/chat/completions -H 'Content-Type: application/json' \
              -d '{"model":"minimax-m3-long","messages":[{"role":"user","content":"日本の首都は?"}],"max_tokens":30}'
"""
import os
import sys
import asyncio

os.environ["M3_OFFICIAL_PORT"] = "1"   # gate m3_quant's gate_up/qkv handling
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ.setdefault("MIXED_MOE_GROUPED", "1")
os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

CKPT = os.environ.get("M3_CKPT", os.path.dirname(os.path.abspath(__file__)))
PORT = os.environ.get("M3_PORT", "8003")             # M1 keeps :8002
SERVED = os.environ.get("M3_SERVED", "minimax-m3-long")
MAXLEN = int(os.environ.get("M3_MAXLEN", "40960"))   # long context (KV ~46K)

# Importing serve_m3_official runs the registrations (m3_quant + the key-
# translating CausalLM loader) and defines the shared config override; the LLM
# build there is guarded by __main__, so the import is side-effect-safe.
import serve_m3_official  # noqa: E402
_ov = serve_m3_official._override_quant_method

from vllm.engine.arg_utils import AsyncEngineArgs  # noqa: E402

# M3_CUDAGRAPH_MODE=PIECEWISE (etc.) -> override the cudagraph mode. The default
# FULL_AND_PIECEWISE hits an illegal-memory-access RACE during the decode-FULL
# capture on ② (CUDA_LAUNCH_BLOCKING=1 avoids it but is too slow); PIECEWISE-only
# may dodge that path.
_cgmode = os.environ.get("M3_CUDAGRAPH_MODE")
_comp_cfg = {"cudagraph_mode": _cgmode} if _cgmode else {}  # {} = engine defaults

_ENGINE_ARGS = AsyncEngineArgs(
    compilation_config=_comp_cfg,
    model=CKPT,
    quantization="autoround_mixed",
    hf_overrides=_ov,
    trust_remote_code=True,
    tensor_parallel_size=2,
    pipeline_parallel_size=1,
    enable_expert_parallel=True,
    distributed_executor_backend="mp",
    block_size=128,                                  # mandatory for MSA sparse cache
    attention_backend=os.environ.get("M3_ATTN_BACKEND", "TRITON_ATTN"),
    max_model_len=MAXLEN,
    max_num_batched_tokens=2048,
    gpu_memory_utilization=0.97,
    max_num_seqs=int(os.environ.get("M3_MAX_SEQS", "4")),
    # M3_EAGER=0 -> cudagraphs (faster decode; M1 runs the same M3 weights this
    # way). The M3 sparse attention supports UNIFORM_BATCH cudagraphs and breaks
    # out the non-capturable split-K kernels via @eager_break_during_capture.
    enforce_eager=os.environ.get("M3_EAGER", "1") == "1",
    disable_custom_all_reduce=True,                  # Blackwell sm_120
    dtype="bfloat16",
)
# Make the OpenAI server use OUR engine args (the CLI path mis-resolves the config).
AsyncEngineArgs.from_cli_args = classmethod(lambda cls, args: _ENGINE_ARGS)

from vllm.entrypoints.openai.api_server import run_server  # noqa: E402
from vllm.entrypoints.openai.cli_args import (  # noqa: E402
    make_arg_parser,
    validate_parsed_serve_args,
)
from vllm.utils.argparse_utils import FlexibleArgumentParser  # noqa: E402

if __name__ == "__main__":
    parser = make_arg_parser(FlexibleArgumentParser())
    args = parser.parse_args([
        CKPT,
        "--served-model-name", SERVED,
        "--host", "0.0.0.0", "--port", PORT,
    ])
    validate_parsed_serve_args(args)
    print(f"[m3_official_api] starting OpenAI server on :{PORT} "
          f"(model={SERVED}, maxlen={MAXLEN})", flush=True)
    asyncio.run(run_server(args))