[train] group consecutive workers by IP

[matthewdeng]

Why are these changes needed?

Some libraries (e.g. Deepspeed) may have expectations that consecutive global rank workers will have the same ordering of local ranks.

In general, this should also make debugging things slightly easier.

Logic

For now, I did the simplest implementation which is to retain the original sequencing of IPs (i.e. node rank), and group the IPs together.

group_workers_by_ip()

[2, 3, 2, 4, 2, 1, 3] --> [2, 2, 2, 3, 3, 4, 1]

This extends existing logic which moves existing logic which moves workers from the Driver IP to the front of the list. Based on this existing, I assuming that moving workers around here is generally safe. The "LOCAL_RANK" and (world) "RANK" will be calculated from the final list.

group_workers_by_ip(4)

[2, 3, 2, 4, 2, 1, 3] --> [4, 2, 2, 2, 3, 3, 1]

Note: This logic can be extended in the future if new requirements come up, e.g. if the IPs need to be sorted by a certain order.

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 :(