chore: import upstream snapshot with attribution

python/ray/train/_internal/data_config.py

@@ -0,0 +1,178 @@
+import copy
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+
+from ray.actor import ActorHandle
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator, Dataset, ExecutionOptions, NodeIdStr
+
+
+@PublicAPI(stability="stable")
+class DataConfig:
+    """Class responsible for configuring Train dataset preprocessing.
+
+    For advanced use cases, this class can be subclassed and the `configure()` method
+    overridden for custom data preprocessing.
+    """
+
+    def __init__(
+        self,
+        datasets_to_split: Union[Literal["all"], List[str]] = "all",
+        execution_options: Optional[
+            Union["ExecutionOptions", Dict[str, "ExecutionOptions"]]
+        ] = None,
+        enable_shard_locality: bool = True,
+    ):
+        """Construct a DataConfig.
+
+        Args:
+            datasets_to_split: Specifies which datasets should be split among workers.
+                Can be set to "all" or a list of dataset names. Defaults to "all",
+                i.e. split all datasets.
+            execution_options: The execution options to pass to Ray Data. Can be either:
+                1. A single ExecutionOptions object that is applied to all datasets.
+                2. A dict mapping dataset names to ExecutionOptions. If a dataset name
+                is not in the dict, it defaults to ``DataConfig.default_ingest_options()``.
+                By default, the options are optimized for data ingest. When overriding,
+                base your options off ``DataConfig.default_ingest_options()``.
+            enable_shard_locality: If true, dataset sharding across Train workers will
+                consider locality to minimize cross-node data transfer. Enabled by default.
+        """
+        from ray.data import ExecutionOptions
+
+        if isinstance(datasets_to_split, list) or datasets_to_split == "all":
+            self._datasets_to_split = datasets_to_split
+        else:
+            raise TypeError(
+                "`datasets_to_split` should be a 'all' or a list of strings of "
+                "dataset names. Received "
+                f"{type(datasets_to_split).__name__} with value {datasets_to_split}."
+            )
+
+        default_execution_options = DataConfig.default_ingest_options()
+
+        if isinstance(execution_options, ExecutionOptions):
+            default_execution_options = execution_options
+        # If None, all datasets will use the default ingest options.
+        self._execution_options: Dict[str, "ExecutionOptions"] = defaultdict(
+            lambda: copy.deepcopy(default_execution_options)
+        )
+        if isinstance(execution_options, dict):
+            self._execution_options.update(execution_options)
+
+        self._enable_shard_locality = enable_shard_locality
+
+        self._num_train_cpus = 0.0
+        self._num_train_gpus = 0.0
+
+    def set_train_total_resources(self, num_train_cpus: float, num_train_gpus: float):
+        """Set the total number of CPUs and GPUs used by training.
+
+        If CPU or GPU resource limits are not set, they will be set to the
+        total cluster resources minus the resources used by training.
+        """
+        # TODO: We may also include other resources besides CPU and GPU.
+        self._num_train_cpus = num_train_cpus
+        self._num_train_gpus = num_train_gpus
+
+    def _get_execution_options(self, dataset_name: str) -> "ExecutionOptions":
+        """Return a copy of the configured execution options for a given dataset name."""
+        return copy.deepcopy(self._execution_options[dataset_name])
+
+    @DeveloperAPI
+    def configure(
+        self,
+        datasets: Dict[str, "Dataset"],
+        world_size: int,
+        worker_handles: Optional[List[ActorHandle]],
+        worker_node_ids: Optional[List["NodeIdStr"]],
+        **kwargs,
+    ) -> List[Dict[str, "DataIterator"]]:
+        """Configure how Train datasets should be assigned to workers.
+
+        Args:
+            datasets: The datasets dict passed to Train by the user.
+            world_size: The number of Train workers in total.
+            worker_handles: The actor handles of the Train workers.
+            worker_node_ids: The node ids of the Train workers.
+            **kwargs: Forwards compatibility placeholder.
+
+        Returns:
+            A list of dataset splits for each worker. The size of the list must be
+            equal to `world_size`. Each element of the list contains the assigned
+            `DataIterator` instances by name for the worker.
+        """
+        from ray.data._internal.execution.interfaces.execution_options import (
+            ExecutionResources,
+        )
+
+        output = [{} for _ in range(world_size)]
+
+        for dataset_name, dataset in datasets.items():
+            if dataset.name is None:
+                dataset.set_name(dataset_name)
+
+        if self._datasets_to_split == "all":
+            datasets_to_split = set(datasets.keys())
+        else:
+            datasets_to_split = set(self._datasets_to_split)
+
+        locality_hints = worker_node_ids if self._enable_shard_locality else None
+        for name, ds in datasets.items():
+            execution_options = self._get_execution_options(name)
+
+            if execution_options.is_resource_limits_default():
+                if not self._scaling_policy_reserves_train_resources():
+                    execution_options.exclude_resources = (
+                        execution_options.exclude_resources.add(
+                            ExecutionResources(
+                                cpu=self._num_train_cpus, gpu=self._num_train_gpus
+                            )
+                        )
+                    )
+
+            ds = ds.copy(ds)
+            ds.context.execution_options = execution_options
+
+            if name in datasets_to_split:
+                for i, split in enumerate(
+                    ds.streaming_split(
+                        world_size, equal=True, locality_hints=locality_hints
+                    )
+                ):
+                    output[i][name] = split
+            else:
+                for i in range(world_size):
+                    output[i][name] = ds.iterator()
+
+        return output
+
+    @classmethod
+    def _scaling_policy_reserves_train_resources(cls) -> bool:
+        """True iff Ray Train V2's ScalingPolicy will register training resources
+        with the AutoscalingCoordinator for this run.
+
+        """
+        from ray.train.v2._internal.constants import is_v2_enabled
+
+        return is_v2_enabled()
+
+    @staticmethod
+    def default_ingest_options() -> "ExecutionOptions":
+        """The default Ray Data options used for data ingest.
+
+        By default, configurations are carried over from what is already set
+        in DataContext.
+        """
+        from ray.data import ExecutionOptions
+        from ray.data.context import DataContext
+
+        ctx = DataContext.get_current()
+        return ExecutionOptions(
+            resource_limits=ctx.execution_options.resource_limits,
+            exclude_resources=ctx.execution_options.exclude_resources,
+            preserve_order=ctx.execution_options.preserve_order,
+            verbose_progress=ctx.execution_options.verbose_progress,
+        )