[Train] Add local mode support to Ray Train v2 (num_workers=0)

[xinyuangui2]

This PR implements local mode support for Ray Train v2, allowing training functions to run in the same process when num_workers=0 is specified in the ScalingConfig. This feature provides a lightweight alternative for development, debugging, and single-node training scenarios.

Key Changes:

Core Infrastructure

• Abstract TrainContext: Refactored TrainContext to be an abstract base class with concrete implementations for distributed and local modes
  • Created LocalTrainContext and DistributedTrainContext implementations
• TrainFnUtils Implementations:
  • DistributedTrainFnUtils: Handles distributed training scenarios
  • LocalTrainFnUtils: Provides local mode implementation with simulated distributed training APIs
• LocalController: New controller class for managing local mode training execution

API Enhancements

• ScalingConfig: Added support for num_workers=0 with informational logging
• DataParallelTrainer: Enhanced to detect local mode and route to appropriate controller

Framework Integration

• PyTorch: Updated device detection to work correctly in local mode
• Other Frameworks: Full support for TensorFlow, Lightning, LightGBM, JAX, Xgboost, and HuggingFace Transformers

Testing

• Comprehensive Test Suite: Added extensive tests covering all supported frameworks in local mode
• Edge Cases: Validation of metrics reporting, checkpointing, and dataset handling

Usage Example:

from ray.train import ScalingConfig
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer

def train_fn():
    # Your training code here
    ray.train.report({"accuracy": 0.95})

trainer = DataParallelTrainer(
    train_fn, 
    scaling_config=ScalingConfig(num_workers=0)  # Local mode
)
result = trainer.fit()

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

[gemini-code-assist]

Summary of Changes

Hello @xinyuangui2, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the flexibility of Ray's TorchTrainer by enabling it to operate in a standalone mode, independent of the full Ray Train distributed execution framework. This change is particularly beneficial for local development, debugging, or scenarios where the overhead of a distributed Ray cluster is not required. It introduces a new backend to manage local training execution, including integration with torchrun, and refactors key internal APIs to support both distributed and local training contexts seamlessly. The overall impact is a more versatile TorchTrainer that can be adopted for a wider range of PyTorch training workflows.

Highlights

• Standalone TorchTrainer Execution: This PR introduces the capability for TorchTrainer to run training loops without requiring the full Ray Train distributed execution environment. This is achieved by adding a new running_without_ray_train flag to the TorchTrainer constructor.
• New Local Training Backend: A new backend, TorchBackendWithoutRayTrain, has been implemented to manage the execution flow when TorchTrainer is used in this standalone mode. This backend handles local process group initialization (e.g., for torchrun) and executes the user-defined training function directly.
• API Abstraction and Context Management: The core TrainContext and TrainFnUtils classes have been refactored into Abstract Base Classes (ABCs). New concrete implementations, TrainContextWithoutRayTrain and TorchWithoutRayTrainTrainFnUtils, provide the necessary context and utility functions tailored for non-Ray Train execution, ensuring consistent API usage.
• Device Management Utilities: New get_device and get_devices utility functions have been added to ray.train.torch.train_loop_utils to correctly identify and return the appropriate torch.device (CPU or CUDA) based on whether the training is managed by Ray Train or running locally.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a significant and valuable feature: the ability to run TorchTrainer without a full Ray Train environment. This is excellent for local development, debugging, and single-node training. The overall approach, which involves abstracting TrainContext and TrainFnUtils, is well-designed. I've identified a couple of critical bugs related to an undefined attribute and a missing return value, along with some suggestions to improve the design and maintainability. Once these issues are addressed, this will be a great addition to Ray Train.