Use different max_concurrent_io_tasks defaults for local and remote paths (#23847)

Add path-aware automatic defaults for `max_concurrent_io_tasks`.

By default, Scan actors use lower IO concurrency for local paths and higher concurrency for scans with remote URIs. Users can override this with an integer for all scans, or a `{"local": ..., "remote": ...}` dict. The environment variable also accepts `auto` or a JSON dict.

Closes #23552. Once this PR merges we won't be trying to set the value based on any heuristic, but we will at least be defining it based on observed benchmark results.

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Peter Andreas Entschev (https://github.com/pentschev)
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: #23847

Author: rjzamora

docs/cudf/source/cudf_polars/memory_errors.md

@@ -59,6 +59,14 @@ enters the pipeline at once. For formats that do not support partial reads, such
 the engine must load an entire file before it can begin processing, which may produce
 chunks much larger than `target_partition_size`.
 
+### Concurrent file reads
+
+Each scan node may read more than one input chunk at a time. At least two IO
+producer tasks help overlap IO with GPU compute. Larger
+`max_concurrent_io_tasks` values may improve high-latency IO throughput but
+increase memory use per Scan actor. See
+{class}`~cudf_polars.engine.options.StreamingOptions` for current defaults.
+
 ## Spilling to host memory
 
 When GPU memory pressure rises above a configurable threshold
@@ -95,10 +103,11 @@ constructing the GPU engine for queries.
 | Option | Default | Effect |
 |---|---|---|
 | `target_partition_size` (executor option or `CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE`) | 1.5 GB or 2.5% of smallest GPU | Target chunk size in bytes. Smaller values reduce peak memory at some cost to compute efficiency. |
+| `max_concurrent_io_tasks` (executor option or `CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`) | auto | Number of concurrent IO producer tasks for each scan node. Larger values may improve high-latency IO throughput but increase memory pressure. |
 | `RAPIDSMPF_SPILL_DEVICE_LIMIT` | `80%` | GPU memory fraction at which spilling begins. Lower values give more headroom for peaks. |
 | `RAPIDSMPF_PINNED_MEMORY` | disabled | Set to `true` to enable pinned host memory for spill buffers. |
 | `RAPIDSMPF_PINNED_INITIAL_POOL_SIZE` | (none) | Size of the pinned memory pool to pre-allocate (e.g. `32GB`). |
 
-For the full list of engine configuration options, including `target_partition_size`,
-see {doc}`options`. For the full list of memory and spill configuration options see the
-[RapidsMPF configuration reference](https://docs.rapids.ai/api/rapidsmpf/stable/configuration/#general).
+For the full list of engine configuration options, including `target_partition_size`
+and `max_concurrent_io_tasks`, see {doc}`options`. For the full list of memory
+and spill configuration options see the [RapidsMPF configuration reference](https://docs.rapids.ai/api/rapidsmpf/stable/configuration/#general).

docs/cudf/source/cudf_polars/options.md

@@ -107,6 +107,7 @@ Environment variables follow these patterns:
 | `max_rows_per_partition` | Maximum number of rows per partition. Only used for in-memory `DataFrame` sources, never for disk IO or dynamic planning.                           | `1_000_000` |
 | `broadcast_limit`        | Maximum number of bytes for broadcast joins.                                                                                                        | auto        |
 | `target_partition_size`  | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto.                                                                   | auto        |
+| `max_concurrent_io_tasks` | Number of concurrent IO producer tasks for each scan node. Tune with an integer or a `{"local": ..., "remote": ...}` dict. | auto        |
 | `dynamic_planning`       | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables.                                 | enabled     |
 | `join_filter_pushdown`   | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables.        | enabled     |
 | `sink_to_directory`      | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True`      |

python/cudf_polars/cudf_polars/engine/options.py

@@ -21,6 +21,7 @@
 from cudf_polars.utils.config import (
     UNSPECIFIED,
     DynamicPlanningOptions,
+    MaxConcurrentIOTasks,
     MemoryResourceConfig,
     Unspecified,
 )
@@ -210,7 +211,10 @@ class StreamingOptions:
     max_concurrent_io_tasks
         Maximum concurrent IO tasks for each scan node.
         Env: ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS``.
-        Default: ``2``.
+        Default: automatic, resolved separately for each scan based on its paths.
+        Python and config values may be an ``int``, a dict with ``local``
+        and/or ``remote`` keys, or omitted/``None`` for the default policy.
+        The environment variable accepts an int or a JSON dict.
         Category: executor.
     fallback_mode
         Fallback behavior (``"warn"``, ``"raise"``, ``"silent"``).
@@ -338,8 +342,10 @@ class StreamingOptions:
     kvikio_statistics: bool | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS", parse_boolean
     )
-    max_concurrent_io_tasks: int | Unspecified = _opt(
-        "executor", "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", int
+    max_concurrent_io_tasks: int | dict[str, int] | Unspecified | None = _opt(
+        "executor",
+        "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS",
+        MaxConcurrentIOTasks.parse_env,
     )
     fallback_mode: str | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__FALLBACK_MODE"
@@ -723,7 +729,7 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None:
             help=textwrap.dedent("""\
                 Maximum concurrent IO tasks for each scan node.
                 Env: CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS.
-                Built-in default: 2."""),
+                Omit to use the path-dependent default."""),
         )
         g.add_argument(
             "--raise-on-fail",

python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py

@@ -21,7 +21,11 @@
         PartitionInfo,
         StatsCollector,
     )
-    from cudf_polars.utils.config import ConfigOptions, StreamingExecutor
+    from cudf_polars.utils.config import (
+        ConfigOptions,
+        MaxConcurrentIOTasks,
+        StreamingExecutor,
+    )
 
 
 class FanoutInfo(NamedTuple):
@@ -52,7 +56,7 @@ class GenState(TypedDict):
     ir_context
         The execution context for the IR node.
     max_concurrent_io_tasks
-        The maximum number of concurrent IO tasks to use for a single IO node.
+        The local and remote IO task limits to use for scan nodes.
     stats
         Statistics collector.
     collective_id_map
@@ -65,7 +69,7 @@ class GenState(TypedDict):
     partition_info: MutableMapping[IR, PartitionInfo]
     fanout_nodes: dict[IR, FanoutInfo]
     ir_context: IRExecutionContext
-    max_concurrent_io_tasks: int
+    max_concurrent_io_tasks: MaxConcurrentIOTasks
     stats: StatsCollector
     collective_id_map: dict[IR, list[int]]
 

python/cudf_polars/cudf_polars/streaming/actor_graph/io.py

@@ -13,6 +13,7 @@
 
 import polars as pl
 
+import pylibcudf as plc
 from cudf_streaming.channel_metadata import ChannelMetadata
 from cudf_streaming.table_chunk import TableChunk
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
@@ -43,7 +44,7 @@
 from cudf_polars.streaming.rank_aware_source import RankAwareSource
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
+    from collections.abc import Callable, Iterable, Sequence
 
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
@@ -57,6 +58,17 @@
         PartitionInfo,
     )
     from cudf_polars.streaming.io import FusedScan, SplitScan
+    from cudf_polars.utils.config import MaxConcurrentIOTasks
+
+
+def resolve_max_concurrent_io_tasks(
+    max_concurrent_io_tasks: MaxConcurrentIOTasks,
+    paths: Iterable[str],
+) -> int:
+    """Resolve the scan-local IO producer count."""
+    if any(plc.io.SourceInfo._is_remote_uri(path) for path in paths):
+        return max_concurrent_io_tasks.remote
+    return max_concurrent_io_tasks.local
 
 
 class Lineariser:
@@ -280,7 +292,9 @@ def _(
 ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]:
     config_options = rec.state["config_options"]
     rows_per_partition = config_options.executor.max_rows_per_partition
-    num_producers = rec.state["max_concurrent_io_tasks"]
+    num_producers = resolve_max_concurrent_io_tasks(
+        rec.state["max_concurrent_io_tasks"], ()
+    )
     # Use target_partition_size as the estimated chunk size
     estimated_chunk_bytes = config_options.executor.target_partition_size
 
@@ -674,7 +688,10 @@ def _(
     config_options = rec.state["config_options"]
     executor = config_options.executor
     partition_info = rec.state["partition_info"][ir]
-    num_producers = rec.state["max_concurrent_io_tasks"]
+    num_producers = resolve_max_concurrent_io_tasks(
+        rec.state["max_concurrent_io_tasks"],
+        ir.base_scan.paths,
+    )
     channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])}
 
     assert partition_info.io_plan is not None, "Scan node must have a partition plan"

python/cudf_polars/cudf_polars/utils/config.py

@@ -61,6 +61,7 @@
     "DynamicPlanningOptions",
     "InMemoryExecutor",
     "JoinFilterPushdownOptions",
+    "MaxConcurrentIOTasks",
     "ParquetOptions",
     "RayContext",
     "SPMDContext",
@@ -105,6 +106,47 @@ def __repr__(self) -> str:
 """
 
 
+@dataclasses.dataclass(frozen=True)
+class MaxConcurrentIOTasks:
+    """Concurrent IO task defaults for local and remote scan paths."""
+
+    local: int = 2
+    remote: int = 8
+
+    @staticmethod
+    def parse_env(raw: str) -> int | dict[str, int] | None:
+        """Parse an environment-variable value."""
+        raw = raw.strip()
+        try:
+            return int(raw)
+        except ValueError:
+            value = json.loads(raw)
+            MaxConcurrentIOTasks.from_config(value)
+            return value
+
+    @classmethod
+    def from_config(
+        cls, value: int | dict[str, int] | MaxConcurrentIOTasks | None
+    ) -> MaxConcurrentIOTasks:
+        """Construct from the supported configuration shapes."""
+        if value is None:
+            return cls()
+        if isinstance(value, int):
+            return cls(local=value, remote=value)
+        if isinstance(value, MaxConcurrentIOTasks):
+            return value
+        if not isinstance(value, dict):
+            raise TypeError("max_concurrent_io_tasks must be an int, dict, or None")
+        return cls(**value)
+
+    def __post_init__(self) -> None:
+        """Validate local and remote values."""
+        if type(self.local) is not int or type(self.remote) is not int:
+            raise TypeError("max_concurrent_io_tasks values must be ints")
+        if self.local < 1 or self.remote < 1:
+            raise ValueError("max_concurrent_io_tasks values must be positive")
+
+
 def _env_get_int(name: str, default: int) -> int:
     try:
         return int(os.getenv(name, default))
@@ -790,11 +832,16 @@ class StreamingExecutor:
         Enable through environment variables with
         ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1``.
     max_concurrent_io_tasks
-        Maximum number of concurrent IO tasks for each scan node. Default is 2.
-        This can be set via
+        Maximum number of concurrent IO tasks for each scan node. The default
+        uses ``2`` for local paths and ``8`` for scans with remote URIs.
+        Passing an ``int`` uses the same value for all scans. Passing a dict
+        with ``local`` and/or ``remote`` keys tunes local and remote paths
+        separately. Omit the option, or pass ``None``, to use the default
+        policy. This can be set via
 
         - ``executor_options`` passed to ``polars.GPUEngine``
-        - the ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`` environment variable
+        - the ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`` environment
+          variable, as an int or JSON dict
     num_py_executors
         Maximum number of workers for the Python ThreadPoolExecutor.
         Default is 8.
@@ -881,9 +928,13 @@ class StreamingExecutor:
     join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field(
         default_factory=JoinFilterPushdownOptions
     )
-    max_concurrent_io_tasks: int = dataclasses.field(
+    max_concurrent_io_tasks: MaxConcurrentIOTasks = dataclasses.field(
         default_factory=_make_default_factory(
-            f"{_env_prefix}__MAX_CONCURRENT_IO_TASKS", int, default=2
+            f"{_env_prefix}__MAX_CONCURRENT_IO_TASKS",
+            lambda raw: MaxConcurrentIOTasks.from_config(
+                MaxConcurrentIOTasks.parse_env(raw)
+            ),
+            default=MaxConcurrentIOTasks(),
         )
     )
     num_py_executors: int = dataclasses.field(
@@ -967,7 +1018,6 @@ def __post_init__(self) -> None:  # noqa: D105
             object.__setattr__(self, "sink_to_directory", True)
         elif self.sink_to_directory is None:
             object.__setattr__(self, "sink_to_directory", False)
-
         # Type / value check everything else
         if not isinstance(self.max_rows_per_partition, int):
             raise TypeError("max_rows_per_partition must be an int")
@@ -979,8 +1029,6 @@ def __post_init__(self) -> None:  # noqa: D105
             raise TypeError("sink_to_directory must be bool")
         if not isinstance(self.client_device_threshold, float):
             raise TypeError("client_device_threshold must be a float")
-        if not isinstance(self.max_concurrent_io_tasks, int):
-            raise TypeError("max_concurrent_io_tasks must be an int")
         if not isinstance(self.num_py_executors, int):
             raise TypeError("num_py_executors must be an int")
         if not isinstance(self.kvikio_nthreads, int):
@@ -994,6 +1042,9 @@ def __hash__(self) -> int:  # noqa: D105
         d = dataclasses.asdict(self)
         d["dynamic_planning"] = json.dumps(d["dynamic_planning"])
         d["join_filter_pushdown"] = json.dumps(d["join_filter_pushdown"])
+        d["max_concurrent_io_tasks"] = json.dumps(
+            d["max_concurrent_io_tasks"], sort_keys=True
+        )
 
         # Hash the quent context UUIDs as ints
         quent_context = d["quent_context"]
@@ -1189,6 +1240,12 @@ def from_polars_engine(
                 user_executor_options = user_executor_options.copy()
                 if "min_device_size" not in user_executor_options:
                     user_executor_options["min_device_size"] = get_total_device_memory()
+                if "max_concurrent_io_tasks" in user_executor_options:
+                    user_executor_options["max_concurrent_io_tasks"] = (
+                        MaxConcurrentIOTasks.from_config(
+                            user_executor_options["max_concurrent_io_tasks"]
+                        )
+                    )
 
                 # Handle dynamic_planning: check user config, then env var
                 user_dynamic_planning = user_executor_options.get(

python/cudf_polars/docs/overview.md

@@ -410,19 +410,24 @@ engine = pl.GPUEngine(
 )
 ```
 
-Each scan node may run up to `max_concurrent_io_tasks` reads concurrently. The
-limit applies independently to each scan node, each corresponding to a
-single `pl.scan_parquet` call in the query. Configure it through
-`executor_options` or
+Each scan node may run up to `max_concurrent_io_tasks` reads concurrently. By
+default, the streaming executor chooses this limit automatically based on the
+scan's paths. The limit applies independently to each scan node, each
+corresponding to a single `pl.scan_parquet` call in the query. Configure it
+explicitly through `executor_options` or
 `CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`:
 
 ```python
 engine = pl.GPUEngine(
     executor="streaming",
-    executor_options={"max_concurrent_io_tasks": 8},
+    executor_options={"max_concurrent_io_tasks": 4},
 )
 ```
 
+Passing an integer uses the same limit for all scans. Pass a
+`{"local": ..., "remote": ...}` dict, or set the environment variable to a
+JSON value like `{"remote": 16}`, to tune local and remote scans separately.
+
 Before each read is submitted, it waits for a device-memory reservation.
 This makes aggregate read concurrency respond to memory pressure across all
 scan nodes on the rank.