[Train] Sort Local Train Workers by GPU id

[woshiyyya]

Why are these changes needed?

Sort the local Ray Train workers according to the GPU device id. This ensures that the allocated GPU == "cuda:{local_rank}". More details in #40803

Related issue number

Close #40803

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

[matthewdeng]

Converting to int won't work in the future if we support UUIDs (e.g for MIG). Maybe we can first try to convert to int and then fallback to string?

[woshiyyya]

Got it. Make sense. I've made take str ID as a backup.

[matthewdeng]

Wondering if we should update the name of this method now that it's no longer grouping but also sorting. At the very least we should update the docstring.

In the future we may want to start generalizing this more (e.g. have it be a generic sort function that takes in a comparator) and define the comparison logic upstream in the caller, since this IP/GPU sorting logic doesn't actually belong to the "worker group", but the "backend" layer on top of it.

[woshiyyya]

I agree. This is a temporary solution for now. We should actually put the ranking allocation logic into backend executor (e.g. here:

ray/python/ray/train/_internal/backend_executor.py

Line 358 in 473d8f5

def _create_rank_world_size_mappings(self) -> List[Dict]:

) and support different allocation strategies in the future. I'll update the name and docstring. What about naming it as sort_workers_by_ip_and_gpu_id?

[matthewdeng]

I think we should still test default sorting behavior when there are no GPU IDs?

[woshiyyya]

When using cpu actors, GPU ids should be empty, then the sorted order will be non-deterministic since all the worker has the same key 0.

[matthewdeng]

Oh I think sort will retain the ordering when the entries have the same sort value. But to your point I don't know if this is guaranteed or just based on the implementation.

entries = [5, 4, 3, 2, 1]
print(entries)

entries.sort(key = lambda x: 0)
print(entries)

entries.sort(key = lambda x: x)
print(entries)

[5, 4, 3, 2, 1]
[5, 4, 3, 2, 1]
[1, 2, 3, 4, 5]

[woshiyyya]

Ah I got you. Seems that the python's sort will maintain the order if multiple entries have identical keys: https://docs.python.org/3.7/howto/sorting.html#sort-stability-and-complex-sorts

I'll add a test for it.