#! 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))