[Train] Split overloaded ray.train.torch.get_device into another get_devices API for multi-GPU worker setup

[woshiyyya]

Why are these changes needed?

The original ray.train.torch.get_device behaves inconsistently depending on the number of available devices on a ray train worker:

• Single Device: Returns a single torch.device object.
• Multiple Devices: Returns a list of torch.device objects.

The proposal involves two key changes:

• Modification of ray.train.torch.get_device():
  • New Behavior: This function will always return a single torch.device object. If multiple devices are available, it will return the device with smaller id.
  • Rationale: Ensures consistent return type, simplifying the usage and handling in user code.
• Introduction of ray.train.torch.get_devices():
  • Behavior: This new function will return a list of torch.device objects, representing all available devices.
  • Rationale: Provides a clear and explicit way to retrieve all devices for multi-gpu worker scenario.

Example:

Single-GPU workers:

def train_func():
    device = ray.train.torch.get_device() # torch.device("cuda:0") 
    devices = ray.train.torch.get_devices() # [torch.device("cuda:0")]

Multi-GPU workers (e.g. 2 gpus per worker):

def train_func():
    device = ray.train.torch.get_device() # torch.device("cuda:0")
    devices = ray.train.torch.get_devices() # [torch.device("cuda:0"), torch.device("cuda:1")]

Related issue number

Closes #42003, #38115

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

[justinvyu]

One more round of suggestions for docs, then I think it's good.

1. Add a small subsection called "Assigning multiple GPUs to a worker" that shows a small tested example recommending get_devices
   • https://anyscale-ray--42314.com.readthedocs.build/en/42314/train/user-guides/using-gpus.html#using-gpus-in-the-training-function
2. Add get_devices to the API reference here: https://anyscale-ray--42314.com.readthedocs.build/en/42314/train/api/api.html#pytorch

[matthewdeng]

This should have a new TRAIN_TORCH_GET_DEVICES key.

[woshiyyya]

Ah previously Justin mentioned we can use a single Key for these two APIs. But now I think it makes more sense to have a separate one to get more accurate telemetry data.

[matthewdeng]

Nice, this wasn't actually working as expected before.

[woshiyyya]

@ArturNiederfahrenhorst Can you take a look? The only change to rllib is in rllib/core/learner/torch/torch_learner.py

[ArturNiederfahrenhorst]

Yep, sorry!

[ArturNiederfahrenhorst]

Approved for RLlib changes

[justinvyu]

Very nice!

[pcmoritz]

For usage.proto