Remove CUDAStreamPolicy enum and simplify CUDA stream policy (rapidsai#22086)

Currently it is possible to produce streams for the rapidsmpf execution path in cudf to get streams in a few different ways. You can...

1) use the CUDA default stream
2) construct a new stream before pipeline evaluation,
3) get streams from the rapidsmpf stream pool. 

We have no real need for the second option, and it is a potential performance footgun for conflicting with the stream pool. This PR removes that ability, and in the process simplifies much of the stream policy handling leaving us with just the stream pool config.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: rapidsai#22086

Author: vyasr

docs/cudf/source/cudf_polars/api.md

@@ -8,7 +8,7 @@ For the most part, the public API of `cudf-polars` is the polars API.
    :members:
       Cluster,
       ConfigOptions,
-      CUDAStreamPolicy,
+      CUDAStreamPoolConfig,
       DynamicPlanningOptions,
       ExecutorType,
       InMemoryExecutor,

python/cudf_polars/cudf_polars/callback.py

@@ -301,7 +301,7 @@ def _callback(
         ),
     ):
         if config_options.executor.name == "in-memory":
-            context = IRExecutionContext.from_config_options(config_options)
+            context = IRExecutionContext()
             df = ir.evaluate(cache={}, timer=timer, context=context).to_polars()
             if timer is None:
                 return df

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -43,11 +43,9 @@
     range_window_bounds,
 )
 from cudf_polars.utils import dtypes
-from cudf_polars.utils.config import CUDAStreamPolicy
 from cudf_polars.utils.cuda_stream import (
     get_cuda_stream,
     get_joined_cuda_stream,
-    get_new_cuda_stream,
     join_cuda_streams,
 )
 from cudf_polars.utils.versions import (
@@ -68,7 +66,7 @@
 
     from cudf_polars.containers.dataframe import NamedColumn
     from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice
-    from cudf_polars.utils.config import ConfigOptions, ParquetOptions
+    from cudf_polars.utils.config import ParquetOptions
     from cudf_polars.utils.timer import Timer
 
 __all__ = [
@@ -114,25 +112,9 @@ class IRExecutionContext:
         A zero-argument callable that returns a CUDA stream.
     """
 
-    get_cuda_stream: Callable[[], Stream]
+    get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream)
     query_id: uuid.UUID = field(default_factory=uuid.uuid4)
 
-    @classmethod
-    def from_config_options(
-        cls, config_options: ConfigOptions, query_id: uuid.UUID | None = None
-    ) -> IRExecutionContext:
-        """Create an IRExecutionContext from ConfigOptions."""
-        query_id = query_id or uuid.uuid4()
-        match config_options.cuda_stream_policy:
-            case CUDAStreamPolicy.DEFAULT:
-                return cls(get_cuda_stream=get_cuda_stream, query_id=query_id)
-            case CUDAStreamPolicy.NEW:
-                return cls(get_cuda_stream=get_new_cuda_stream, query_id=query_id)
-            case _:  # pragma: no cover
-                raise ValueError(
-                    f"Invalid CUDA stream policy: {config_options.cuda_stream_policy}"
-                )
-
     @contextlib.contextmanager
     def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None]:
         """

python/cudf_polars/cudf_polars/experimental/benchmarks/utils_legacy.py

@@ -812,9 +812,7 @@ def execute_query(
             if args.debug:
                 translator = Translator(q._ldf.visit(), engine)
                 ir = translator.translate_ir()
-                context = IRExecutionContext.from_config_options(
-                    translator.config_options
-                )
+                context = IRExecutionContext()
                 if run_config.executor == "in-memory":
                     t0 = time.monotonic()
                     result = ir.evaluate(

python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py

@@ -713,9 +713,7 @@ def execute_query(
             if args.debug:
                 translator = Translator(q._ldf.visit(), engine)
                 ir = translator.translate_ir()
-                context = IRExecutionContext.from_config_options(
-                    translator.config_options
-                )
+                context = IRExecutionContext()
                 if run_config.executor == "in-memory":
                     t0 = time.monotonic()
                     result = ir.evaluate(

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -141,7 +141,7 @@ def task_graph(
     --------
     generate_ir_tasks
     """
-    context = IRExecutionContext.from_config_options(config_options)
+    context = IRExecutionContext()
     graph = reduce(
         operator.or_,
         (

python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py

@@ -312,9 +312,7 @@ def evaluate_pipeline(
                 get_cuda_stream=rmpf_context.get_stream_from_pool, query_id=query_id
             )
         else:
-            ir_context = IRExecutionContext.from_config_options(
-                config_options, query_id=query_id
-            )
+            ir_context = IRExecutionContext(query_id=query_id)
 
         # Generate network nodes
         assert rmpf_context is not None, "RapidsMPF context must defined."

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py

@@ -242,7 +242,7 @@ class StreamingOptions:
         Env: ``CUDF_POLARS__MEMORY_RESOURCE_CONFIG__*``.
         Category: engine.
     cuda_stream_policy
-        CUDA stream policy (``"default"``, ``"new"``, ``"pool"`` or config dict).
+        CUDA stream policy (``"default"``, ``"pool"`` or config dict).
         Env: ``CUDF_POLARS__CUDA_STREAM_POLICY``.
         Category: engine.
 
@@ -307,9 +307,9 @@ class StreamingOptions:
     raise_on_fail: bool | Unspecified = _opt("engine")
     parquet_options: dict[str, Any] | ParquetOptions | Unspecified = _opt("engine")
     memory_resource_config: MemoryResourceConfig | Unspecified = _opt("engine")
-    cuda_stream_policy: (
-        Literal["default", "new", "pool"] | dict[str, Any] | Unspecified
-    ) = _opt("engine", "CUDF_POLARS__CUDA_STREAM_POLICY")
+    cuda_stream_policy: Literal["default", "pool"] | dict[str, Any] | Unspecified = (
+        _opt("engine", "CUDF_POLARS__CUDA_STREAM_POLICY")
+    )
 
     # ------------------------------------------------------------------
     # Conversion helpers used by the engines

python/cudf_polars/cudf_polars/utils/config.py

@@ -947,24 +947,12 @@ def build(self) -> CudaStreamPool:
         )
 
 
-class CUDAStreamPolicy(enum.StrEnum):
-    """
-    The policy to use for acquiring new CUDA streams.
-
-    * ``CUDAStreamPolicy.DEFAULT`` : Use the default CUDA stream.
-    * ``CUDAStreamPolicy.NEW`` : Create a new CUDA stream.
-    """
-
-    DEFAULT = "default"
-    NEW = "new"
-
-
 def _convert_cuda_stream_policy(
     user_cuda_stream_policy: dict | str,
-) -> CUDAStreamPolicy | CUDAStreamPoolConfig:
+) -> CUDAStreamPoolConfig | None:
     match user_cuda_stream_policy:
-        case "default" | "new":
-            return CUDAStreamPolicy(user_cuda_stream_policy)
+        case "default":
+            return None
         case "pool":
             return CUDAStreamPoolConfig()
         case dict():
@@ -997,6 +985,13 @@ def _convert_cuda_stream_policy(
                         ) from None
 
 
+def _default_cuda_stream_policy() -> CUDAStreamPoolConfig | None:
+    v = os.environ.get("CUDF_POLARS__CUDA_STREAM_POLICY")
+    if v is None:
+        return None
+    return _convert_cuda_stream_policy(v)
+
+
 @dataclasses.dataclass(frozen=True, eq=True)
 class ConfigOptions(Generic[ExecutorType]):
     """
@@ -1017,7 +1012,9 @@ class ConfigOptions(Generic[ExecutorType]):
         The GPU used to run the query. If not provided, the
         query uses the current CUDA device.
     cuda_stream_policy
-        The policy to use for acquiring new CUDA streams. See :class:`~cudf_polars.utils.config.CUDAStreamPolicy` for more.
+        The policy to use for CUDA streams. ``None`` (the default) uses the
+        default CUDA stream. A :class:`~cudf_polars.utils.config.CUDAStreamPoolConfig`
+        can be used to configure a stream pool.
     """
 
     raise_on_fail: bool = False
@@ -1029,12 +1026,8 @@ class ConfigOptions(Generic[ExecutorType]):
     )
     device: int | None = None
     memory_resource_config: MemoryResourceConfig | None = None
-    cuda_stream_policy: CUDAStreamPolicy | CUDAStreamPoolConfig = dataclasses.field(
-        default_factory=_make_default_factory(
-            "CUDF_POLARS__CUDA_STREAM_POLICY",
-            CUDAStreamPolicy.__call__,
-            default=CUDAStreamPolicy.DEFAULT,
-        )
+    cuda_stream_policy: CUDAStreamPoolConfig | None = dataclasses.field(
+        default_factory=_default_cuda_stream_policy
     )
 
     @classmethod
@@ -1137,7 +1130,7 @@ def from_polars_engine(
             "cuda_stream_policy", None
         ) or os.environ.get("CUDF_POLARS__CUDA_STREAM_POLICY", None)
 
-        cuda_stream_policy: CUDAStreamPolicy | CUDAStreamPoolConfig
+        cuda_stream_policy: CUDAStreamPoolConfig | None
 
         if user_cuda_stream_policy is None:
             if (
@@ -1147,7 +1140,7 @@ def from_polars_engine(
                 cuda_stream_policy = CUDAStreamPoolConfig()
             else:
                 # everything else defaults to the default stream
-                cuda_stream_policy = CUDAStreamPolicy.DEFAULT
+                cuda_stream_policy = None
         else:
             cuda_stream_policy = _convert_cuda_stream_policy(user_cuda_stream_policy)
 
@@ -1157,7 +1150,7 @@ def from_polars_engine(
             or (executor.name == "streaming" and executor.runtime != Runtime.RAPIDSMPF)
         ):
             raise ValueError(
-                "CUDAStreamPolicy.POOL is only supported by the rapidsmpf runtime."
+                "A stream pool is only supported by the rapidsmpf runtime."
             )
 
         kwargs["cuda_stream_policy"] = cuda_stream_policy

python/cudf_polars/cudf_polars/utils/cuda_stream.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """CUDA stream utilities."""
@@ -8,11 +8,13 @@
 from typing import TYPE_CHECKING
 
 import pylibcudf as plc
-from rmm.pylibrmm.stream import DEFAULT_STREAM, Stream
+from rmm.pylibrmm.stream import DEFAULT_STREAM
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
+    from rmm.pylibrmm.stream import Stream
+
 
 def get_dask_cuda_stream() -> Stream:
     """Get the default CUDA stream for Dask."""
@@ -24,11 +26,6 @@ def get_cuda_stream() -> Stream:
     return DEFAULT_STREAM
 
 
-def get_new_cuda_stream() -> Stream:
-    """Get a new CUDA stream for the current thread."""
-    return Stream()
-
-
 def join_cuda_streams(
     *, downstreams: Sequence[Stream], upstreams: Sequence[Stream]
 ) -> None: