import os, sys, subprocess, argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True)
    parser.add_argument('--tensor_parallel_size', type=int, default=1)
    parser.add_argument('--port', type=int, default=8000)
    parser.add_argument('--max_model_len', type=int, default=8192)
    parser.add_argument('--max_num_seqs', type=int, default=256)
    parser.add_argument('--dtype', default='bfloat16')
    parser.add_argument('--num_scheduler_steps', type=int, default=10)
    parser.add_argument('--enable_chunked_prefill', action='store_true', default=False)
    args = parser.parse_args()

    os.environ.setdefault('VLLM_USE_TRITON_FLASH_ATTN', '0')
    os.environ.setdefault('NCCL_MIN_NCHANNELS', '112')
    os.environ.setdefault('PYTORCH_TUNABLEOP_ENABLED', '1')
    os.environ.setdefault('TORCHINDUCTOR_MAX_AUTOTUNE', '1')

    cmd = [sys.executable, '-m', 'vllm.entrypoints.openai.api_server',
           '--model', args.model, '--tensor-parallel-size', str(args.tensor_parallel_size),
           '--port', str(args.port), '--max-model-len', str(args.max_model_len),
           '--max-num-seqs', str(args.max_num_seqs), '--dtype', args.dtype,
           '--num-scheduler-steps', str(args.num_scheduler_steps)]
    if not args.enable_chunked_prefill:
        cmd.append('--enable-chunked-prefill=False')

    print('Launching vLLM:', ' '.join(cmd))
    subprocess.run(cmd)

if __name__ == "__main__": main()