[train] Sort workers by node ID rather than by node IP

[KyleKoon]

Why are these changes needed?

We should not sort workers by IP, since it may be the case that multiple nodes are placed on the same IP host. If multiple nodes are collocated, the local rank and node rank of the workers will be incorrect. We should instead use a node's unique ID to map to a set of workers.

Example:

Worker 0: node 0, ip 0
Worker 1: node 0, ip 0
Worker 2: node 1, ip 0
Worker 3: node 1, ip 0

Current Behavior:

worker 0: local_rank 0, node_rank 0
worker 1: local_rank 1, node_rank 0
worker 2: local_rank 2, node_rank 0
worker 3: local_rank 3, node_rank 0

Expected Behavior:

worker 0: local_rank 0, node_rank 0
worker 1: local_rank 1, node_rank 0
worker 2: local_rank 0, node_rank 1
worker 3: local_rank 1, node_rank 1

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[woshiyyya]

Thank you @KyleKoon, in general it looks good. Can we add some tests to reproduce the failure case we want to resolve (multiple workers nodes shares the same IP)?

You can refer to this example to construct virtual clusters:

ray/python/ray/train/tests/conftest.py

Line 45 in 450df60

def ray_2_node_2_gpu():

[woshiyyya]

The 2 virtual nodes will have the same node_ip. Can we also have a test without patching the add_worker method?

[KyleKoon]

Done

[woshiyyya]

It looks good to me. Left one comment

[matthewdeng]

Thanks!