[train] Improve error message if users call training function utils outside of a Ray Train worker

[justinvyu]

Description

Introduce a decorator to mark functions that require running inside a worker process spawned by Ray Train.

Additional information

Before:

>>> import ray.train
>>> ray.train.get_context()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/justin/Developer/ray/python/ray/train/v2/api/train_fn_utils.py", line 153, in get_context
    return get_train_fn_utils().get_context()
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/justin/Developer/ray/python/ray/train/v2/_internal/execution/train_fn_utils.py", line 264, in get_train_fn_utils
    raise RuntimeError("TrainFnUtils has not been initialized.")
RuntimeError: TrainFnUtils has not been initialized.

After:

>>> import ray.train
>>> ray.train.get_context()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/justin/Developer/ray/python/ray/train/v2/_internal/util.py", line 281, in _wrapped_fn
    raise RuntimeError(
RuntimeError: `get_context` cannot be used outside of a Ray Train training function. You are calling this API from the driver or another non-training process. These utilities are only available within a function launched by `trainer.fit()`.

[gemini-code-assist]

Code Review

This pull request introduces a decorator requires_train_worker to enforce that certain functions are only called within a Ray Train worker process. This enhances error messaging and prevents misuse of training utilities outside of the intended environment. The changes include adding the decorator to relevant functions in collective.py, train_fn_utils.py, and torch/train_loop_utils.py, and updating a test case in test_api_migrations.py to reflect the new error handling.

[TimothySeah]

lgtm with a few nits

[TimothySeah]

Is this supposed to be cuda:2? The other examples make it seem like it should be cuda:<get_gpu_ids()>

[justinvyu]

cuda:0 is the "logical cuda device", which points to the 0th index in the CUDA_VISIBLE_DEVICES string: CUDA_VISIBLE_DEVICES[0] == "2"