diff --git a/docs/cudf/source/cudf_polars/memory_errors.md b/docs/cudf/source/cudf_polars/memory_errors.md index 88e0ec5a5fc..f39a4938d4a 100644 --- a/docs/cudf/source/cudf_polars/memory_errors.md +++ b/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). diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index 6744ad959a9..d4af20497ff 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/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` | diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index b519d2aba9d..e1272475f94 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/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", diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 88d6784b094..fc355addc10 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/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]] diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 44957443e72..80e501d3e00 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/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" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 6aa4f3139d3..47c8a9f6dfd 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/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( diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index daf24990263..f2a11f91bc7 100644 --- a/python/cudf_polars/docs/overview.md +++ b/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. diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index f0adf211f72..781b4166816 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -44,11 +44,13 @@ def test_all_fields_unspecified_by_default(monkeypatch: pytest.MonkeyPatch) -> N monkeypatch.delenv(key, raising=False) opts = StreamingOptions() # Fields with no env var are always UNSPECIFIED. + assert isinstance(opts.max_concurrent_io_tasks, Unspecified) assert isinstance(opts.raise_on_fail, Unspecified) assert isinstance(opts.parquet_options, Unspecified) # Fields whose env vars were cleared are also UNSPECIFIED. assert isinstance(opts.fallback_mode, Unspecified) assert isinstance(opts.log, Unspecified) + assert opts.to_executor_options() == {} # --------------------------------------------------------------------------- @@ -56,10 +58,6 @@ def test_all_fields_unspecified_by_default(monkeypatch: pytest.MonkeyPatch) -> N # --------------------------------------------------------------------------- -def test_executor_options_empty_when_all_unspecified() -> None: - assert StreamingOptions().to_executor_options() == {} - - def test_executor_options_includes_set_fields() -> None: opts = StreamingOptions(fallback_mode="raise", max_rows_per_partition=500_000) result = opts.to_executor_options() @@ -73,9 +71,41 @@ def test_executor_options_num_py_executors() -> None: assert result["num_py_executors"] == 4 -def test_executor_options_max_concurrent_io_tasks() -> None: - result = StreamingOptions(max_concurrent_io_tasks=6).to_executor_options() - assert result["max_concurrent_io_tasks"] == 6 +@pytest.mark.parametrize( + "value", + [6, {"local": 3, "remote": 7}, {"remote": 7}, None], +) +def test_executor_options_max_concurrent_io_tasks( + value: int | dict[str, int] | None, +) -> None: + result = StreamingOptions(max_concurrent_io_tasks=value).to_executor_options() + assert result["max_concurrent_io_tasks"] == value + + +@pytest.mark.parametrize( + "env, expected", + [ + ('{"local": 2, "remote": 7}', {"local": 2, "remote": 7}), + ('{"remote": 7}', {"remote": 7}), + ], +) +def test_executor_options_max_concurrent_io_tasks_env( + monkeypatch: pytest.MonkeyPatch, + env: str, + expected: dict[str, int], +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", env) + opts = StreamingOptions() + assert opts.max_concurrent_io_tasks == expected + assert opts.to_executor_options()["max_concurrent_io_tasks"] == expected + + +def test_executor_options_max_concurrent_io_tasks_env_rejects_auto( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", "auto") + with pytest.raises(json.JSONDecodeError): + StreamingOptions() def test_executor_options_kvikio_nthreads() -> None: @@ -262,8 +292,11 @@ def test_from_dict_maps_known_fields() -> None: def test_from_dict_none_value_is_unspecified() -> None: - opts = StreamingOptions.from_dict({"fallback_mode": None}) + opts = StreamingOptions.from_dict( + {"fallback_mode": None, "max_concurrent_io_tasks": None} + ) assert isinstance(opts.fallback_mode, Unspecified) + assert isinstance(opts.max_concurrent_io_tasks, Unspecified) def test_from_dict_unknown_key_raises() -> None: @@ -390,9 +423,16 @@ def test_to_dict_empty_when_all_unspecified() -> None: def test_to_dict_contains_only_set_fields() -> None: - opts = StreamingOptions(fallback_mode="silent", num_streaming_threads=4) - d = opts.to_dict() - assert d == {"fallback_mode": "silent", "num_streaming_threads": 4} + opts = StreamingOptions( + fallback_mode="silent", + num_streaming_threads=4, + max_concurrent_io_tasks=6, + ) + assert opts.to_dict() == { + "fallback_mode": "silent", + "num_streaming_threads": 4, + "max_concurrent_io_tasks": 6, + } def test_to_dict_roundtrip() -> None: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index b1f4c617fd8..f6c1450ceb1 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -23,6 +23,7 @@ prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.io import resolve_max_concurrent_io_tasks from cudf_polars.streaming.base import ( DataSourceInfo, IOPartitionFlavor, @@ -41,7 +42,11 @@ from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import SMALL_MAX_ROWS_PER_PARTITION from cudf_polars.testing.io import make_partitioned_source -from cudf_polars.utils.config import ConfigOptions, ParquetOptions +from cudf_polars.utils.config import ( + ConfigOptions, + MaxConcurrentIOTasks, + ParquetOptions, +) if TYPE_CHECKING: import concurrent.futures @@ -159,6 +164,59 @@ def test_cached_parquet_info_hybrid_scan_reader_lazy(tmp_path, df) -> None: assert info._hybrid_scan_metadata is not None +@pytest.mark.parametrize( + "paths,expected", + [ + ([], 2), + (["file.parquet"], 2), + (["file.parquet", "s3://bucket/file.parquet"], 8), + (["s3://bucket/file.parquet"], 8), + ], +) +def test_resolve_max_concurrent_io_tasks_default( + paths: list[str], expected: int +) -> None: + assert resolve_max_concurrent_io_tasks(MaxConcurrentIOTasks(), paths) == expected + + +def test_resolve_max_concurrent_io_tasks_explicit() -> None: + assert ( + resolve_max_concurrent_io_tasks( + MaxConcurrentIOTasks(local=6, remote=6), ["s3://bucket/file.parquet"] + ) + == 6 + ) + + +@pytest.mark.parametrize( + "paths,expected", + [ + (["file.parquet"], 3), + (["s3://bucket/file.parquet"], 7), + ], +) +def test_resolve_max_concurrent_io_tasks_local_remote_policy( + paths: list[str], expected: int +) -> None: + assert ( + resolve_max_concurrent_io_tasks(MaxConcurrentIOTasks(local=3, remote=7), paths) + == expected + ) + + +def test_resolve_max_concurrent_io_tasks_partial_override() -> None: + max_concurrent_io_tasks = MaxConcurrentIOTasks(remote=7) + assert ( + resolve_max_concurrent_io_tasks(max_concurrent_io_tasks, ["file.parquet"]) == 2 + ) + assert ( + resolve_max_concurrent_io_tasks( + max_concurrent_io_tasks, ["s3://bucket/file.parquet"] + ) + == 7 + ) + + def test_prefetch_file_metadata_select_fast_count( df: pl.DataFrame, streaming_engine_factory: Callable[..., StreamingEngine], diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 062f543ff76..4f0a953dd6c 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -33,6 +33,7 @@ DynamicPlanningOptions, InMemoryExecutor, JoinFilterPushdownOptions, + MaxConcurrentIOTasks, MemoryResourceConfig, ParquetOptions, StreamingExecutor, @@ -425,10 +426,119 @@ def test_config_option_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.executor.max_rows_per_partition == 42 assert config.executor.target_partition_size == 100 assert config.executor.broadcast_limit == 44 - assert config.executor.max_concurrent_io_tasks == 6 + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks( + local=6, remote=6 + ) assert config.executor.quent_context is not None +def test_max_concurrent_io_tasks_local_remote_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with monkeypatch.context() as m: + m.setenv( + "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", + '{"local": 2, "remote": 7}', + ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks( + local=2, remote=7 + ) + + with monkeypatch.context() as m: + m.setenv("CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", '{"remote": 7}') + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks( + local=2, remote=7 + ) + + +def test_max_concurrent_io_tasks_default_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", raising=False) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks() + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": 6}, + ) + ) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks( + local=6, remote=6 + ) + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "max_concurrent_io_tasks": {"local": 3, "remote": 7}, + }, + ) + ) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks( + local=3, remote=7 + ) + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": None}, + ) + ) + assert config.executor.max_concurrent_io_tasks == MaxConcurrentIOTasks() + + +def test_max_concurrent_io_tasks_accepts_dataclass() -> None: + value = MaxConcurrentIOTasks(local=3, remote=7) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": value}, + ) + ) + assert config.executor.max_concurrent_io_tasks is value + + +@pytest.mark.parametrize( + "value", + [0, -1, {"local": 0}, {"remote": 0}, {"local": -1}, {"remote": -1}], +) +def test_max_concurrent_io_tasks_rejects_non_positive( + value: int | dict[str, int], +) -> None: + with pytest.raises(ValueError, match="must be positive"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": value}, + ) + ) + + +@pytest.mark.parametrize( + "value", + [ + True, + False, + {"local": True}, + {"remote": False}, + {"local": 1.5}, + {"remote": "8"}, + ], +) +def test_max_concurrent_io_tasks_rejects_non_int(value: object) -> None: + with pytest.raises(TypeError, match="must be ints"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": value}, + ) + ) + + def test_quent_context_from_env_disabled(monkeypatch: pytest.MonkeyPatch) -> None: with monkeypatch.context() as m: m.setenv("CUDF_POLARS__EXECUTOR__QUENT_CONTEXT", "0") @@ -583,6 +693,7 @@ def test_parquet_options_unspecified_dict_factory() -> None: assert isinstance(config.parquet_options.prefetch_file_metadata, Unspecified) result = dataclasses.asdict(config, dict_factory=ConfigOptions.dict_factory) assert result["parquet_options"]["prefetch_file_metadata"] is None + assert result["executor"]["max_concurrent_io_tasks"] == {"local": 2, "remote": 8} def test_validate_raise_on_fail() -> None: