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

os.environ["NCCL_NVLS_ENABLE"] = "0"
os.environ["NCCL_CUMEM_ENABLE"] = "1"
os.environ["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "0"

chute = build_sglang_chute(
    tee=True,
    username="chutes",
    readme="zai-org/GLM-5.1-FP8",
    model_name="zai-org/GLM-5.1-FP8",
    name="zai-org/GLM-5.1-TEE",
    image="chutes/sglang:0.5.21",
    official_package=True,
    concurrency=32,
    revision="f396cf805182f4ca10fa675e1a99815b3ca384db",
    node_selector=NodeSelector(
        gpu_count=8,
        include=["h200"],
    ),
    engine_args=(
        "--tp-size 8 "
        "--reasoning-parser glm45 "
        "--tool-call-parser glm47 "
        f"--cuda-graph-max-bs-decode 32 "
        "--mem-fraction-static 0.87 "
        "--dsa-topk-backend sgl-kernel "
        "--disable-custom-all-reduce"
    ),
)