[data] Cache cluster_resources

python/ray/data/_internal/execution/streaming_executor.py

@@ -4,7 +4,6 @@
 import uuid
 from typing import Dict, Iterator, List, Optional
 
-import ray
 from ray.data._internal.dataset_logger import DatasetLogger
 from ray.data._internal.execution.autoscaling_requester import (
     get_or_create_autoscaling_requester_actor,
@@ -35,6 +34,7 @@
 )
 from ray.data._internal.progress_bar import ProgressBar
 from ray.data._internal.stats import DatasetStats, StatsManager
+from ray.data._internal.util import cluster_resources
 from ray.data.context import DataContext
 
 logger = DatasetLogger(__name__)
@@ -335,7 +335,7 @@ def _get_or_refresh_resource_limits(self) -> ExecutionResources:
         """
         base = self._options.resource_limits
         exclude = self._options.exclude_resources
-        cluster = ray.cluster_resources()
+        cluster = cluster_resources()
 
         cpu = base.cpu
         if cpu is None:

python/ray/data/_internal/util.py

@@ -13,6 +13,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Dict,
     Iterable,
     Iterator,
     List,
@@ -55,6 +56,25 @@
 _pyarrow_dataset: LazyModule = None
 
 
+cached_cluster_resources = {}
+cluster_resources_last_fetch_time = 0
+CLUSTER_RESOURCES_FETCH_INTERVAL_SECONDS = 10
+
+
+def cluster_resources() -> Dict[str, float]:
+    """Fetch Ray cluster resources with cache."""
+    global cached_cluster_resources
+    global cluster_resources_last_fetch_time
+    now = time.time()
+    if (
+        now - cluster_resources_last_fetch_time
+        > CLUSTER_RESOURCES_FETCH_INTERVAL_SECONDS
+    ):
+        cached_cluster_resources = ray.cluster_resources()
+        cluster_resources_last_fetch_time = now
+    return cached_cluster_resources
+
+
 def _lazy_import_pyarrow_dataset() -> LazyModule:
     global _pyarrow_dataset
     if _pyarrow_dataset is None:

python/ray/data/tests/test_util.py

@@ -1,4 +1,6 @@
+import time
 from typing import Any, Dict, Optional
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -111,6 +113,32 @@ def get_parquet_read_logical_op(
     return read_op
 
 
+def test_cluster_resources():
+    """Test ray.data._internal.util.cluster_resources()."""
+    resources = {"CPU": 4, "GPU": 1}
+    cache_interval_s = 0.1
+    with patch.object(
+        ray.data._internal.util,
+        "CLUSTER_RESOURCES_FETCH_INTERVAL_SECONDS",
+        cache_interval_s,
+    ):
+        with patch(
+            "ray.cluster_resources",
+            return_value=resources,
+        ) as ray_cluster_resources:
+            # The first call should call ray.cluster_resources().
+            assert ray.data._internal.util.cluster_resources() == resources
+            assert ray_cluster_resources.call_count == 1
+            # The second call should return the cached value.
+            assert ray.data._internal.util.cluster_resources() == resources
+            assert ray_cluster_resources.call_count == 1
+            time.sleep(cache_interval_s)
+            # After the cache interval, the third call should call
+            # ray.cluster_resources() again.
+            assert ray.data._internal.util.cluster_resources() == resources
+            assert ray_cluster_resources.call_count == 2
+
+
 if __name__ == "__main__":
     import sys