list.py 500 Bytes
Newer Older
liangjing's avatar
liangjing committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import hostlist
import os
import json
def get_list():
  worker = []
  nodelist = os.environ["SLURM_JOB_NODELIST"]
  nodelist = hostlist.expand_hostlist(nodelist)
  num_nodes = int(os.getenv("SLURM_JOB_NUM_NODES"))
  port_number = 60001
  worker_nodes = [node for i, node in enumerate(nodelist) if i >= 0 ]
  for node in worker_nodes:
      for index in range(4):
          worker_sockets = ":".join([node, str(port_number + index )])
          worker.append(worker_sockets)
  return ','.join(worker)