[data][train] Create a deepcopy of the data context on the split coordinator process (#56211)

The main change of this PR is to create a deepcopy of the base dataset's
context before setting the process-global context.

Otherwise, mutations to the base dataset's context during the planning
phase are also propagated to the global context, which can affect future
dataset executions launched from the same process.

Misc. drive-by changes:
* Utility to create a `StorageContext` from the `RunConfig` directly
* Pipe the `DatasetShardMetadata` from the outermost level among other
changes, for easier patching

---------

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

Author: justinvyu

python/ray/data/_internal/iterator/stream_split_iterator.py

@@ -139,9 +139,13 @@ def __init__(
         locality_hints: Optional[List[NodeIdStr]],
     ):
         dataset = dataset_wrapper._dataset
+
         # Set current DataContext.
-        self._data_context = dataset.context
+        # This needs to be a deep copy so that updates to the base dataset's
+        # context does not affect this process's global DataContext.
+        self._data_context = dataset.context.copy()
         ray.data.DataContext._set_current(self._data_context)
+
         if self._data_context.execution_options.locality_with_output is True:
             self._data_context.execution_options.locality_with_output = locality_hints
             logger.info(f"Auto configuring locality_with_output={locality_hints}")

python/ray/data/iterator.py

@@ -158,6 +158,11 @@ def iter_batches(
             local_shuffle_seed=local_shuffle_seed,
         )
 
+    def _create_batch_iterator(
+        self, ref_bundles_iter: Iterator[RefBundle], **kwargs
+    ) -> BatchIterator:
+        return BatchIterator(ref_bundles_iter, **kwargs)
+
     def _iter_batches(
         self,
         *,
@@ -186,7 +191,7 @@ def _create_iterator() -> Iterator[DataBatch]:
 
             dataset_tag = self._get_dataset_tag()
 
-            batch_iterator = BatchIterator(
+            batch_iterator = self._create_batch_iterator(
                 ref_bundles_iterator,
                 stats=stats,
                 dataset_tag=dataset_tag,

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

@@ -7,9 +7,9 @@
 from ray.train.v2._internal.data_integration.interfaces import (
     DatasetShardMetadata,
     DatasetShardProvider,
-    GenDataset,
 )
 from ray.train.v2._internal.execution.callback import WorkerGroupCallback
+from ray.train.v2._internal.execution.context import TrainRunContext
 from ray.train.v2._internal.execution.worker_group.worker_group import (
     Worker,
     WorkerGroup,
@@ -37,15 +37,10 @@ def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> DataIterator:
 class DatasetsSetupCallback(WorkerGroupCallback):
     """The callback to setup Ray Datasets for the worker group."""
 
-    def __init__(
-        self,
-        datasets: Dict[str, GenDataset],
-        data_config: ray.train.DataConfig,
-        scaling_config: ray.train.ScalingConfig,
-    ):
-        self._datasets = datasets
-        self._data_config = data_config
-        self._scaling_config = scaling_config
+    def __init__(self, train_run_context: TrainRunContext):
+        self._datasets = train_run_context.datasets
+        self._data_config = copy.deepcopy(train_run_context.dataset_config)
+        self._scaling_config = train_run_context.scaling_config
 
         # Capture the current DataContext to propagate it to
         # the Train workers later.

python/ray/train/v2/_internal/execution/controller/controller.py

@@ -50,7 +50,6 @@
     ResizeDecision,
     ScalingPolicy,
 )
-from ray.train.v2._internal.execution.storage import StorageContext
 from ray.train.v2._internal.execution.worker_group import (
     WorkerGroup,
     WorkerGroupPollStatus,
@@ -126,11 +125,7 @@ def __init__(
         self._failure_policy = failure_policy
         self._run_config = self._train_run_context.run_config
         self._callbacks = callbacks or []
-        self._storage_context = StorageContext(
-            storage_path=self._run_config.storage_path,
-            experiment_dir_name=self._run_config.name,
-            storage_filesystem=self._run_config.storage_filesystem,
-        )
+        self._storage_context = self._train_run_context.run_config.storage_context
 
         self._checkpoint_manager = CheckpointManager(
             checkpoint_config=self._run_config.checkpoint_config,

python/ray/train/v2/_internal/execution/train_fn_utils.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from ray.data import DataIterator
+from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
 from ray.train.v2._internal.execution import collective_impl
 from ray.train.v2._internal.execution.context import (
     get_train_context as get_internal_train_context,
@@ -68,14 +69,11 @@ def get_all_reported_checkpoints(self) -> List["ReportedCheckpoint"]:
         pass
 
     @abstractmethod
-    def get_dataset_shard(self, dataset_name: str) -> DataIterator:
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> DataIterator:
         """Get the dataset shard for this training process.
 
-        This method is used by the public API function :func:`ray.train.get_dataset_shard`.
-        Users should typically call ``ray.train.get_dataset_shard()`` instead of calling this method directly.
-
         Args:
-            dataset_name: The name of the dataset to get the shard for.
+            dataset_info: The metadata of the dataset to get the shard for.
 
         Returns:
             The DataIterator shard for this worker.
@@ -131,14 +129,8 @@ def report(
     def get_checkpoint(self):
         return get_internal_train_context().get_checkpoint()
 
-    def get_dataset_shard(self, dataset_name: str) -> DataIterator:
-        from ray.train.v2._internal.data_integration.interfaces import (
-            DatasetShardMetadata,
-        )
-
-        return get_internal_train_context().get_dataset_shard(
-            DatasetShardMetadata(dataset_name=dataset_name)
-        )
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> DataIterator:
+        return get_internal_train_context().get_dataset_shard(dataset_info)
 
     def get_context(self) -> DistributedTrainContext:
         return DistributedTrainContext()
@@ -182,7 +174,8 @@ def report(
     def get_checkpoint(self) -> Optional["Checkpoint"]:
         return self._last_checkpoint
 
-    def get_dataset_shard(self, dataset_name: str) -> DataIterator:
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> DataIterator:
+        dataset_name = dataset_info.dataset_name
         assert (
             self._dataset_shards is not None and dataset_name in self._dataset_shards
         ), f"Dataset shard {dataset_name} not found."

python/ray/train/v2/_internal/execution/worker_group/worker_group.py

@@ -37,7 +37,6 @@
 from ray.train.v2._internal.execution.checkpoint.sync_actor import SynchronizationActor
 from ray.train.v2._internal.execution.context import (
     DistributedContext,
-    StorageContext,
     TrainRunContext,
 )
 from ray.train.v2._internal.execution.worker_group.poll import (
@@ -145,11 +144,7 @@ def __init__(
         """
         self._train_run_context = train_run_context
         run_config = self._train_run_context.run_config
-        self._storage_context = StorageContext(
-            storage_path=run_config.storage_path,
-            experiment_dir_name=run_config.name,
-            storage_filesystem=run_config.storage_filesystem,
-        )
+        self._storage_context = run_config.storage_context
 
         self._worker_group_context: WorkerGroupContext = worker_group_context
 

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

@@ -1,5 +1,6 @@
 import logging
 from dataclasses import dataclass
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, Union
 
@@ -12,6 +13,7 @@
 )
 from ray.runtime_env import RuntimeEnv
 from ray.train.v2._internal.constants import _DEPRECATED
+from ray.train.v2._internal.execution.storage import StorageContext
 from ray.train.v2._internal.migration_utils import (
     FAIL_FAST_DEPRECATION_MESSAGE,
     TRAINER_RESOURCES_DEPRECATION_MESSAGE,
@@ -261,3 +263,11 @@ def __post_init__(self):
                 "See this issue for more context: "
                 "https://github.com/ray-project/ray/issues/49454"
             )
+
+    @cached_property
+    def storage_context(self) -> StorageContext:
+        return StorageContext(
+            storage_path=self.storage_path,
+            experiment_dir_name=self.name,
+            storage_filesystem=self.storage_filesystem,
+        )

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

@@ -30,7 +30,6 @@
     TPUReservationCallback,
     WorkingDirectorySetupCallback,
 )
-from ray.train.v2._internal.callbacks.datasets import GenDataset
 from ray.train.v2._internal.callbacks.env_callback import _initialize_env_callbacks
 from ray.train.v2._internal.callbacks.metrics import (
     ControllerMetricsCallback,
@@ -42,6 +41,7 @@
     METRICS_ENABLED_ENV_VAR,
     get_env_vars_to_propagate,
 )
+from ray.train.v2._internal.data_integration.interfaces import GenDataset
 from ray.train.v2._internal.execution.callback import RayTrainCallback
 from ray.train.v2._internal.execution.context import TrainRunContext
 from ray.train.v2._internal.execution.controller import TrainController
@@ -164,9 +164,7 @@ def _create_default_callbacks(self) -> List[RayTrainCallback]:
         )
         backend_setup_callback = BackendSetupCallback(self.backend_config)
         datasets_setup_callback = DatasetsSetupCallback(
-            datasets=self.datasets,
-            data_config=self.data_config,
-            scaling_config=self.scaling_config,
+            train_run_context=self.train_run_context
         )
         tpu_reservation_setup_callback = TPUReservationCallback()
         callbacks.extend(

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

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
+from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
 from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
 from ray.train.v2.api.context import TrainContext
 from ray.util.annotations import PublicAPI
@@ -241,4 +242,6 @@ def train_loop_per_worker(config):
         The ``DataIterator`` shard to use for this worker.
         If no dataset is passed into Trainer, then return None.
     """
-    return get_train_fn_utils().get_dataset_shard(dataset_name)
+    return get_train_fn_utils().get_dataset_shard(
+        DatasetShardMetadata(dataset_name=dataset_name)
+    )

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

@@ -1,5 +1,3 @@
-from unittest.mock import MagicMock
-
 import pytest
 
 import ray.data
@@ -9,12 +7,15 @@
 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.data_integration.interfaces import DatasetShardMetadata
-from ray.train.v2._internal.execution.context import TrainRunContext
 from ray.train.v2._internal.execution.worker_group.worker_group import (
     WorkerGroupContext,
 )
 from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
-from ray.train.v2.tests.util import DummyObjectRefWrapper, DummyWorkerGroup
+from ray.train.v2.tests.util import (
+    DummyObjectRefWrapper,
+    DummyWorkerGroup,
+    create_dummy_run_context,
+)
 
 # TODO(justinvyu): Bring over more tests from ray/air/tests/test_new_dataset_config.py
 
@@ -77,17 +78,18 @@ def test_dataset_setup_callback(ray_start_4_cpus):
         num_workers=scaling_config.num_workers,
         resources_per_worker=scaling_config.resources_per_worker,
     )
+    train_run_context = create_dummy_run_context(
+        datasets={"train": train_ds, "valid": valid_ds},
+        dataset_config=data_config,
+        scaling_config=scaling_config,
+    )
     worker_group = DummyWorkerGroup(
-        train_run_context=MagicMock(spec=TrainRunContext),
+        train_run_context=train_run_context,
         worker_group_context=worker_group_context,
     )
     worker_group._start()
 
-    callback = DatasetsSetupCallback(
-        datasets={"train": train_ds, "valid": valid_ds},
-        data_config=data_config,
-        scaling_config=scaling_config,
-    )
+    callback = DatasetsSetupCallback(train_run_context)
     dataset_manager_for_each_worker = callback.before_init_train_context(
         worker_group.get_workers()
     )["dataset_shard_provider"]