ds_bench

#!/usr/bin/env python3

from benchmarks.communication.run_all import main
from benchmarks.communication.constants import *
from benchmarks.communication.utils import *
import argparse
import os
import sys

# Run the same file with deepspeed launcher. This is required since setuptools will auto-detect python files and insert a python shebang for both 'scripts' and 'entry_points', and this benchmarks require the DS launcher
required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
if not all(map(lambda v: v in os.environ, required_env)):
    import subprocess
    subprocess.run("deepspeed $(which ds_bench) " + " ".join(sys.argv[1:]), shell=True)
else:
    args = benchmark_parser().parse_args()
    rank = args.local_rank
    main(args, rank)