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"

CONCURRENCY = 32
ATTENTION_DP_SIZE = 8
PREFILL_CHUNK_TOKENS = 4096 * ATTENTION_DP_SIZE
MODEL_REVISION = "a7e62ac04ecb2c0a54d736dc46601c5606cf10a6"

TOKENIZER_PATH = (
    "/cache/hub/models--deepseek-ai--DeepSeek-V3.2/snapshots/"
    f"{MODEL_REVISION}"
)

chute = build_sglang_chute(
    tee=True,
    username="chutes",
    readme="deepseek-ai/DeepSeek-V3.2",
    model_name="deepseek-ai/DeepSeek-V3.2",
    name="deepseek-ai/DeepSeek-V3.2-TEE",
    revision=MODEL_REVISION,
    image="chutes/sglang:0.5.21",
    official_package=True,
    concurrency=CONCURRENCY,
    node_selector=NodeSelector(
        gpu_count=8,
        include=["h200"],
    ),
    engine_args=(
        f"--tokenizer-path {TOKENIZER_PATH} "
        "--tp-size 8 "
        f"--dp-size {ATTENTION_DP_SIZE} "
        "--enable-dp-attention "
        "--tool-call-parser deepseekv32 "
        "--reasoning-parser deepseek-v3 "
        "--trust-remote-code "
        "--speculative-algorithm EAGLE "
        "--speculative-num-steps 3 "
        "--speculative-eagle-topk 1 "
        "--speculative-num-draft-tokens 4 "
        "--speculative-use-rejection-sampling "
        f"--cuda-graph-max-bs-decode {CONCURRENCY} "
        "--cuda-graph-backend-prefill disabled "
        f"--chunked-prefill-size {PREFILL_CHUNK_TOKENS} "
        "--mem-fraction-static 0.85 "
        "--context-length 131072 "
        "--dsa-topk-backend sgl-kernel "
        "--enforce-disable-flashinfer-allreduce-fusion "
        "--disable-custom-all-reduce"
    ),
)