File size: 846 Bytes
e6f9ec5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | #! coding: utf-8
import sys
import os
import json
def main(args):
env_name = args[1]
cluster_spec = json.loads(os.environ["AFO_SPEC"])
workers = cluster_spec["cluster"]["worker"]
master = workers[0]
master_addr, master_ports = master.split(":")
master_port = master_ports.split(",")[0]
nproc_per_node = os.popen("nvidia-smi --list-gpus | wc -l").read().strip()
nnodes = len(workers)
node_rank = cluster_spec["taskId"]
if env_name == 'nproc_per_node':
return nproc_per_node
elif env_name == 'master_addr':
return master_addr
elif env_name == 'master_port':
return master_port
elif env_name == 'nnodes':
return nnodes
elif env_name == 'node_rank':
return node_rank
raise "WRONG ARGUMENT"
if __name__ == '__main__':
print(main(sys.argv))
|