chore: import upstream snapshot with attribution

release/train_tests/benchmark/ray_dataloader_factory.py

@@ -0,0 +1,281 @@
+import logging
+import threading
+import time
+from abc import abstractmethod
+from typing import Any, Dict, Optional
+
+import ray
+import ray.train
+from ray._private.internal_api import get_memory_info_reply, get_state_from_address
+from ray.data import Dataset
+from ray.data.collate_fn import CollateFn
+
+from constants import DatasetKey
+from config import BenchmarkConfig, RayDataConfig
+from dataloader_factory import BaseDataLoaderFactory
+
+logger = logging.getLogger(__name__)
+
+SPILL_MONITOR_ACTOR_NAME = "spill_metrics_monitor"
+SPILL_MONITOR_ACTOR_NAMESPACE = "_spill_metrics_monitor"
+
+
+@ray.remote(num_cpus=0)
+class SpillMetricsMonitor:
+    """Actor that periodically polls object store spill metrics
+    to compute peak and average spilling rates (GB/min).
+
+    GB/min is used instead of GB/s because object store spilling rates are
+    typically small fractions of a GB per second, making GB/s values hard to
+    read and compare (e.g. 0.0021 GB/s vs 0.13 GB/min). GB/min produces
+    more human-friendly numbers while still matching the 60-second poll
+    interval naturally.
+
+    A single instance is shared across all workers via a named actor.
+    """
+
+    def __init__(self, poll_interval_s: float = 60.0):
+        self._poll_interval_s = poll_interval_s
+        self._count = 0
+        self._sum_gb_min = 0.0
+        self._max_gb_min = 0.0
+        self._lock = threading.Lock()
+
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def _get_spilled_bytes(self) -> int:
+        memory_info = get_memory_info_reply(
+            get_state_from_address(ray.get_runtime_context().gcs_address)
+        )
+        return memory_info.store_stats.spilled_bytes_total
+
+    def _poll_loop(self) -> None:
+        prev_spilled_bytes: Optional[int] = None
+        prev_time: Optional[float] = None
+
+        while True:
+            time.sleep(self._poll_interval_s)
+            try:
+                current_bytes = self._get_spilled_bytes()
+                current_time = time.monotonic()
+
+                if prev_spilled_bytes is not None and prev_time is not None:
+                    delta_bytes = current_bytes - prev_spilled_bytes
+                    delta_time = current_time - prev_time
+
+                    if delta_time > 0 and delta_bytes >= 0:
+                        rate_gb_min = (delta_bytes / (1024**3)) / delta_time * 60
+                        with self._lock:
+                            self._count += 1
+                            self._sum_gb_min += rate_gb_min
+                            self._max_gb_min = max(self._max_gb_min, rate_gb_min)
+
+                prev_spilled_bytes = current_bytes
+                prev_time = current_time
+            except Exception as e:
+                logger.warning(f"SpillMetricsMonitor: poll failed: {e}")
+
+    def get_metrics(self) -> Dict[str, float]:
+        with self._lock:
+            count = self._count
+            sum_gb_min = self._sum_gb_min
+            max_gb_min = self._max_gb_min
+
+        if count == 0:
+            return {}
+
+        return {
+            "object_store_spilling_peak_gb_min": round(max_gb_min, 4),
+            "object_store_spilling_avg_gb_min": round(sum_gb_min / count, 4),
+        }
+
+
+def get_or_create_spill_metrics_monitor(
+    poll_interval_s: float = 60.0,
+) -> ray.actor.ActorHandle:
+    return SpillMetricsMonitor.options(
+        name=SPILL_MONITOR_ACTOR_NAME,
+        namespace=SPILL_MONITOR_ACTOR_NAMESPACE,
+        get_if_exists=True,
+        lifetime="detached",
+    ).remote(poll_interval_s=poll_interval_s)
+
+
+class RayDataLoaderFactory(BaseDataLoaderFactory):
+    def __init__(self, benchmark_config: BenchmarkConfig) -> None:
+        super().__init__(benchmark_config)
+        self._ray_ds_iterators = {}
+        self._spill_monitor: Optional[ray.actor.ActorHandle] = None
+
+        dataloader_config = self.get_dataloader_config()
+        assert isinstance(dataloader_config, RayDataConfig), type(dataloader_config)
+
+        # Configure Ray Data settings.
+        data_context = ray.data.DataContext.get_current()
+        data_context.enable_operator_progress_bars = (
+            dataloader_config.enable_operator_progress_bars
+        )
+        # Retry transient S3 errors that sometimes show up due to
+        # throttling during read operations.
+        data_context.retried_io_errors.append("AWS Error ACCESS_DENIED")
+        data_context.retried_io_errors.append("AWS Error UNKNOWN (HTTP status 500)")
+
+        data_context.execution_options.locality_with_output = (
+            dataloader_config.locality_with_output
+        )
+        data_context.execution_options.actor_locality_enabled = (
+            dataloader_config.actor_locality_enabled
+        )
+        data_context.execution_options.preserve_order = dataloader_config.preserve_order
+
+    @abstractmethod
+    def get_ray_datasets(self) -> Dict[str, Dataset]:
+        """Get Ray datasets."""
+        raise NotImplementedError
+
+    def _get_collate_fn(self) -> Optional[CollateFn]:
+        """Return the collate function for the dataloader."""
+        return None
+
+    def get_ray_data_config(self) -> ray.train.DataConfig:
+        return ray.train.DataConfig(
+            enable_shard_locality=self.get_dataloader_config().enable_shard_locality,
+        )
+
+    def get_train_dataloader(self):
+        """Get the training dataloader.
+
+        Returns:
+            Iterator of training batches
+        """
+        # Get or create the shared spill monitor actor on first call.
+        if self._spill_monitor is None:
+            self._spill_monitor = get_or_create_spill_metrics_monitor()
+
+        ds_iterator = ray.train.get_dataset_shard(DatasetKey.TRAIN)
+        self._ray_ds_iterators[DatasetKey.TRAIN] = ds_iterator
+
+        dataloader_config = self.get_dataloader_config()
+        return iter(
+            ds_iterator.iter_torch_batches(
+                batch_size=dataloader_config.train_batch_size,
+                local_shuffle_buffer_size=(
+                    dataloader_config.local_buffer_shuffle_size
+                    if dataloader_config.local_buffer_shuffle_size > 0
+                    else None
+                ),
+                collate_fn=self._get_collate_fn(),
+                prefetch_batches=dataloader_config.ray_data_prefetch_batches,
+                drop_last=True,
+                pin_memory=dataloader_config.ray_data_pin_memory,
+            )
+        )
+
+    def get_val_dataloader(self):
+        """Get the validation dataloader.
+
+        Returns:
+            Iterator of validation batches
+        """
+        ds_iterator = ray.train.get_dataset_shard(DatasetKey.VALID)
+        self._ray_ds_iterators[DatasetKey.VALID] = ds_iterator
+
+        dataloader_config = self.get_dataloader_config()
+        return iter(
+            ds_iterator.iter_torch_batches(
+                batch_size=dataloader_config.validation_batch_size,
+                collate_fn=self._get_collate_fn(),
+                prefetch_batches=dataloader_config.ray_data_prefetch_batches,
+                drop_last=True,
+            )
+        )
+
+    def get_metrics(self) -> Dict[str, Any]:
+        metrics = {}
+        for ds_key, ds_iterator in self._ray_ds_iterators.items():
+            stats = ray.get(ds_iterator._coord_actor.stats.remote())
+            summary = stats.to_summary()
+            summary.iter_stats = ds_iterator._iter_stats.to_summary().iter_stats
+            summary.iter_stats.streaming_split_coord_time.add(
+                stats.streaming_split_coordinator_s.get()
+            )
+
+            if not summary.parents:
+                continue
+
+            # The split() operator has no metrics, so pull the stats
+            # from the final dataset stage.
+            ds_output_summary = summary.parents[0]
+            ds_throughput = (
+                ds_output_summary.operators_stats[-1].output_num_rows.sum
+                / ds_output_summary.get_total_wall_time()
+            )
+
+            iter_stats = summary.iter_stats
+
+            metrics[f"dataloader/{ds_key}"] = {
+                "producer_throughput": ds_throughput,
+                "iter_stats": {
+                    "prefetch_block-avg": iter_stats.wait_time.avg(),
+                    "prefetch_block-min": iter_stats.wait_time.min(),
+                    "prefetch_block-max": iter_stats.wait_time.max(),
+                    "prefetch_block-total": iter_stats.wait_time.get(),
+                    "get_ref_bundles-avg": iter_stats.get_ref_bundles_time.avg(),
+                    "get_ref_bundles-min": iter_stats.get_ref_bundles_time.min(),
+                    "get_ref_bundles-max": iter_stats.get_ref_bundles_time.max(),
+                    "get_ref_bundles-total": iter_stats.get_ref_bundles_time.get(),
+                    "fetch_block-avg": iter_stats.get_time.avg(),
+                    "fetch_block-min": iter_stats.get_time.min(),
+                    "fetch_block-max": iter_stats.get_time.max(),
+                    "fetch_block-total": iter_stats.get_time.get(),
+                    "block_to_batch-avg": iter_stats.next_time.avg(),
+                    "block_to_batch-min": iter_stats.next_time.min(),
+                    "block_to_batch-max": iter_stats.next_time.max(),
+                    "block_to_batch-total": iter_stats.next_time.get(),
+                    "format_batch-avg": iter_stats.format_time.avg(),
+                    "format_batch-min": iter_stats.format_time.min(),
+                    "format_batch-max": iter_stats.format_time.max(),
+                    "format_batch-total": iter_stats.format_time.get(),
+                    "collate-avg": iter_stats.collate_time.avg(),
+                    "collate-min": iter_stats.collate_time.min(),
+                    "collate-max": iter_stats.collate_time.max(),
+                    "collate-total": iter_stats.collate_time.get(),
+                    "finalize-avg": iter_stats.finalize_batch_time.avg(),
+                    "finalize-min": iter_stats.finalize_batch_time.min(),
+                    "finalize-max": iter_stats.finalize_batch_time.max(),
+                    "finalize-total": iter_stats.finalize_batch_time.get(),
+                    "time_spent_blocked-avg": iter_stats.block_time.avg(),
+                    "time_spent_blocked-min": iter_stats.block_time.min(),
+                    "time_spent_blocked-max": iter_stats.block_time.max(),
+                    "time_spent_blocked-total": iter_stats.block_time.get(),
+                    "time_spent_training-avg": iter_stats.user_time.avg(),
+                    "time_spent_training-min": iter_stats.user_time.min(),
+                    "time_spent_training-max": iter_stats.user_time.max(),
+                    "time_spent_training-total": iter_stats.user_time.get(),
+                },
+            }
+
+        # Collect object store spill metrics.
+        try:
+            memory_info = get_memory_info_reply(
+                get_state_from_address(ray.get_runtime_context().gcs_address)
+            )
+            spilled_bytes_total = memory_info.store_stats.spilled_bytes_total
+            metrics["object_store_spilled_total_gb"] = round(
+                spilled_bytes_total / (1024**3), 4
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to collect object_store_spilled_total_gb metric: {e}"
+            )
+
+        # Collect peak and average spilling rate from the background monitor.
+        if self._spill_monitor is not None:
+            try:
+                spill_metrics = ray.get(self._spill_monitor.get_metrics.remote())
+                metrics.update(spill_metrics)
+            except Exception as e:
+                logger.warning(f"Failed to collect spill rate metrics: {e}")
+
+        return metrics