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

os.environ["SGLANG_ENABLE_OVERLAP_PLAN_STREAM"] = "1"

chute = build_sglang_chute(
    tee=True,
    username="chutes",
    readme="Qwen/Qwen3.5-397B-A17B-FP8",
    model_name="Qwen/Qwen3.5-397B-A17B-FP8",
    name="Qwen/Qwen3.5-397B-A17B-TEE",
    image="chutes/sglang:0.5.17-qwen38",
    official_package=True,
    concurrency=40,
    revision="9f1f3de9a3a48cfd340fd73ca98c02625b7afb3b",
    node_selector=NodeSelector(
        gpu_count=8,
        include=["h200"],
    ),
    engine_args=(
        "--mem-fraction-static 0.85 "
        "--context-length 262144 "
        "--expert-parallel-size 8 "
        "--reasoning-parser qwen3 "
        "--tool-call-parser qwen3_coder "
        "--mamba-radix-cache-strategy extra_buffer "
        "--max-running-requests 40 "
        "--chunked-prefill-size 16384 "
        "--enable-flashinfer-allreduce-fusion "
        "--attention-backend flashinfer "
        "--mamba-ssm-dtype bfloat16 "
        "--cuda-graph-max-bs 48 "
        "--disable-custom-all-reduce"
    ),
)