[train] Rename DatasetsSetupCallback to DatasetsCallback (ray-project#59423)

PR ray-project#58325 adds shutdown and abort hooks to enhance resource-cleanup
logic in DatasetsSetupCallback, the callback’s responsibilities have
expanded beyond initial setup. Accordingly, this PR renames it to
DatasetsCallback for better alignment with its behavior.

Signed-off-by: JasonLi1909 <jasli1909@gmail.com>
Signed-off-by: peterxcli <peterxcli@gmail.com>

Author: JasonLi1909

python/ray/train/v2/_internal/callbacks/__init__.py

@@ -1,14 +1,14 @@
 from .accelerators import AcceleratorSetupCallback
 from .backend_setup import BackendSetupCallback
-from .datasets import DatasetsSetupCallback
+from .datasets import DatasetsCallback
 from .state_manager import StateManagerCallback
 from .tpu_reservation_callback import TPUReservationCallback
 from .working_dir_setup import WorkingDirectorySetupCallback
 
 __all__ = [
     "AcceleratorSetupCallback",
     "BackendSetupCallback",
-    "DatasetsSetupCallback",
+    "DatasetsCallback",
     "StateManagerCallback",
     "TPUReservationCallback",
     "WorkingDirectorySetupCallback",

python/ray/train/v2/_internal/callbacks/datasets.py

@@ -46,7 +46,7 @@ def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> DataIterator:
         return self._dataset_iterators[dataset_info.dataset_name]
 
 
-class DatasetsSetupCallback(WorkerGroupCallback, ControllerCallback):
+class DatasetsCallback(WorkerGroupCallback, ControllerCallback):
     """A callback for managing Ray Datasets for the worker group."""
 
     def __init__(self, train_run_context: TrainRunContext):

python/ray/train/v2/api/data_parallel_trainer.py

@@ -27,7 +27,7 @@
 from ray.train.v2._internal.callbacks import (
     AcceleratorSetupCallback,
     BackendSetupCallback,
-    DatasetsSetupCallback,
+    DatasetsCallback,
     TPUReservationCallback,
     WorkingDirectorySetupCallback,
 )
@@ -203,9 +203,7 @@ def _create_default_callbacks(self) -> List[RayTrainCallback]:
             self.backend_config, self.scaling_config
         )
         backend_setup_callback = BackendSetupCallback(self.backend_config)
-        datasets_callback = DatasetsSetupCallback(
-            train_run_context=self.train_run_context
-        )
+        datasets_callback = DatasetsCallback(train_run_context=self.train_run_context)
         tpu_reservation_setup_callback = TPUReservationCallback()
         placement_group_cleaner_callback = PlacementGroupCleanerCallback()
         callbacks.extend(

python/ray/train/v2/tests/test_data_integration.py

@@ -16,7 +16,7 @@
 )
 from ray.data._internal.iterator.stream_split_iterator import StreamSplitDataIterator
 from ray.data.tests.conftest import restore_data_context  # noqa: F401
-from ray.train.v2._internal.callbacks.datasets import DatasetsSetupCallback
+from ray.train.v2._internal.callbacks.datasets import DatasetsCallback
 from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
 from ray.train.v2._internal.execution.worker_group.worker_group import (
     WorkerGroupContext,
@@ -92,7 +92,7 @@ def test_data_config_validation():
 
 
 def test_datasets_callback(ray_start_4_cpus):
-    """Check that the `DatasetsSetupCallback` correctly configures the
+    """Check that the `DatasetsCallback` correctly configures the
     dataset shards and execution options."""
     NUM_WORKERS = 2
 
@@ -121,7 +121,7 @@ def test_datasets_callback(ray_start_4_cpus):
     )
     worker_group._start()
 
-    callback = DatasetsSetupCallback(train_run_context)
+    callback = DatasetsCallback(train_run_context)
     dataset_manager_for_each_worker = callback.before_init_train_context(
         worker_group.get_workers()
     )["dataset_shard_provider"]

python/ray/train/v2/tests/test_data_resource_cleanup.py

@@ -12,7 +12,7 @@
     SplitCoordinator,
     _DatasetWrapper,
 )
-from ray.train.v2._internal.callbacks.datasets import DatasetsSetupCallback
+from ray.train.v2._internal.callbacks.datasets import DatasetsCallback
 from ray.train.v2._internal.execution.worker_group import (
     WorkerGroup,
     WorkerGroupContext,
@@ -23,7 +23,7 @@
 
 
 def test_datasets_callback_multiple_datasets(ray_start_4_cpus):
-    """Test that the DatasetsSetupCallback properly collects the coordinator actors for multiple datasets"""
+    """Test that the DatasetsCallback properly collects the coordinator actors for multiple datasets"""
     # Start worker group
     worker_group_context = WorkerGroupContext(
         run_attempt_id="test",
@@ -49,7 +49,7 @@ def test_datasets_callback_multiple_datasets(ray_start_4_cpus):
         datasets=datasets, dataset_config=dataset_config
     )
 
-    callback = DatasetsSetupCallback(train_run_context)
+    callback = DatasetsCallback(train_run_context)
     callback.before_init_train_context(wg.get_workers())
 
     # Two coordinator actors, one for each sharded dataset
@@ -58,7 +58,7 @@ def test_datasets_callback_multiple_datasets(ray_start_4_cpus):
 
 
 def test_after_worker_group_abort():
-    callback = DatasetsSetupCallback(create_dummy_run_context())
+    callback = DatasetsCallback(create_dummy_run_context())
 
     # Mock SplitCoordinator shutdown_executor method
     coord_mock = create_autospec(SplitCoordinator)
@@ -79,7 +79,7 @@ def test_after_worker_group_abort():
 
 
 def test_after_worker_group_shutdown():
-    callback = DatasetsSetupCallback(create_dummy_run_context())
+    callback = DatasetsCallback(create_dummy_run_context())
 
     # Mock SplitCoordinator shutdown_executor method
     coord_mock = create_autospec(SplitCoordinator)