import os
from chutes.chute import NodeSelector
from chutes.chute.template.sglang import build_sglang_chute

os.environ["SGLANG_OPT_USE_TOPK_V2"] = "1"
os.environ["SGLANG_ENABLE_MOE_DEFERRED_FINALIZE"] = "1"

chute = build_sglang_chute(
    tee=True,
    username="chutes",
    readme="nvidia/GLM-5.2-NVFP4",
    model_name="nvidia/GLM-5.2-NVFP4",
    name="zai-org/GLM-5.2-TEE",
    image="chutes/sglang:0.5.15-post1a",
    official_package=True,
    concurrency=32,
    revision="aec724e8c7b8ee9db3b48c01c320f63f9cdaf8aa",
    node_selector=NodeSelector(
        gpu_count=8,
        include=["b200", "b300"],
    ),
    engine_args=(
        "--tensor-parallel-size 8 "
        "--ep-size 8 "
        "--quantization modelopt_fp4 "
        "--max-prefill-tokens 8192 "
        "--chunked-prefill-size 8192 "
        "--cuda-graph-max-bs-decode 32 "
        "--mem-fraction-static 0.87 "
        "--trust-remote-code "
        "--kv-cache-dtype fp8_e4m3 "
        "--bf16-gemm-backend cutedsl "
        "--reasoning-parser glm45 "
        "--tool-call-parser glm47 "
        "--speculative-algorithm EAGLE "
        "--speculative-num-steps 5 "
        "--speculative-eagle-topk 1 "
        "--speculative-num-draft-tokens 6 "
        "--enable-cache-report "
        "--disable-custom-all-reduce"
    ),
)