diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index b4b2ecb69e07..57ccf6302c5b 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -5,9 +5,11 @@ ARG PYTHON_PACKAGE_MANAGER=conda FROM ${BASE} as pip-base +# libnuma-dev is required for pip devcontainers for cucascade from rapidsmpf RUN apt update -y \ && DEBIAN_FRONTEND=noninteractive apt install -y \ librdkafka-dev \ + libnuma-dev \ && rm -rf /tmp/* /var/tmp/* /var/cache/apt/* /var/lib/apt/lists/*; ENV DEFAULT_VIRTUAL_ENV=rapids diff --git a/.devcontainer/README.md b/.devcontainer/README.md index 91ee7ef85f78..cc5fac22fde2 100644 --- a/.devcontainer/README.md +++ b/.devcontainer/README.md @@ -20,6 +20,7 @@ This container is a turnkey development environment for building and testing the By default, the following directories are bind-mounted into the devcontainer: * `${repo}:/home/coder/cudf` +* `${repo}/../rapidsmpf:/home/coder/rapidsmpf` * `${repo}/../.aws:/home/coder/.aws` * `${repo}/../.local:/home/coder/.local` * `${repo}/../.cache:/home/coder/.cache` @@ -28,6 +29,13 @@ By default, the following directories are bind-mounted into the devcontainer: This ensures caches, configurations, dependencies, and your commits are persisted on the host across container runs. +The [rapidsmpf](https://github.com/rapidsai/rapidsmpf) repository is a required dependency of `cudf_polars` (that also requires `libcudf`) and must be cloned as a sibling directory to the cudf repo before launching the devcontainer: + +``` +# from the parent directory of your cudf clone +git clone https://github.com/rapidsai/rapidsmpf.git +``` + ## Launch a Dev Container To launch a devcontainer from VSCode, open the cuDF repo and select the "Reopen in Container" button in the bottom right:
diff --git a/.devcontainer/cuda12.9-conda/devcontainer.json b/.devcontainer/cuda12.9-conda/devcontainer.json index 9d672bdbde8f..272007e7c953 100644 --- a/.devcontainer/cuda12.9-conda/devcontainer.json +++ b/.devcontainer/cuda12.9-conda/devcontainer.json @@ -49,7 +49,7 @@ "initializeCommand": [ "/bin/bash", "-c", - "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.9-envs}" + "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.9-envs} ${localWorkspaceFolder}/../rapidsmpf" ], "postAttachCommand": [ "/bin/bash", @@ -63,7 +63,8 @@ "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent", - "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.9-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent" + "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.9-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent" ], "customizations": { "vscode": { diff --git a/.devcontainer/cuda12.9-pip/devcontainer.json b/.devcontainer/cuda12.9-pip/devcontainer.json index f99cc4ce5dc3..5012dcfa9797 100644 --- a/.devcontainer/cuda12.9-pip/devcontainer.json +++ b/.devcontainer/cuda12.9-pip/devcontainer.json @@ -5,7 +5,7 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:26.06-cpp-cuda12.9" + "BASE": "rapidsai/devcontainers:26.06-cpp-cuda12.9-ucx1.19.0-openmpi5.0.10" } }, "runArgs": [ @@ -27,7 +27,7 @@ "initializeCommand": [ "/bin/bash", "-c", - "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs}" + "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs} ${localWorkspaceFolder}/../rapidsmpf" ], "postAttachCommand": [ "/bin/bash", @@ -40,7 +40,8 @@ "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", - "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent" + "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent" ], "customizations": { "vscode": { diff --git a/.devcontainer/cuda13.1-conda/devcontainer.json b/.devcontainer/cuda13.1-conda/devcontainer.json index a73953b19894..785302c3c1d0 100644 --- a/.devcontainer/cuda13.1-conda/devcontainer.json +++ b/.devcontainer/cuda13.1-conda/devcontainer.json @@ -49,7 +49,7 @@ "initializeCommand": [ "/bin/bash", "-c", - "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda13.1-envs}" + "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda13.1-envs} ${localWorkspaceFolder}/../rapidsmpf" ], "postAttachCommand": [ "/bin/bash", @@ -63,7 +63,8 @@ "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent", - "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda13.1-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent" + "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda13.1-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent" ], "customizations": { "vscode": { diff --git a/.devcontainer/cuda13.1-pip/devcontainer.json b/.devcontainer/cuda13.1-pip/devcontainer.json index 8596ff6b5035..730b1c1e8cae 100644 --- a/.devcontainer/cuda13.1-pip/devcontainer.json +++ b/.devcontainer/cuda13.1-pip/devcontainer.json @@ -5,7 +5,7 @@ "args": { "CUDA": "13.1", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:26.06-cpp-cuda13.1" + "BASE": "rapidsai/devcontainers:26.06-cpp-cuda13.1-ucx1.19.0-openmpi5.0.10" } }, "runArgs": [ @@ -27,7 +27,7 @@ "initializeCommand": [ "/bin/bash", "-c", - "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda13.1-venvs}" + "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda13.1-venvs} ${localWorkspaceFolder}/../rapidsmpf" ], "postAttachCommand": [ "/bin/bash", @@ -40,7 +40,8 @@ "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", - "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda13.1-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent" + "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda13.1-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent" ], "customizations": { "vscode": { diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 4e6935c38225..41a4c734f21c 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -550,9 +550,16 @@ jobs: SCCACHE_DIST_MAX_RETRIES=inf SCCACHE_SERVER_LOG=sccache=debug SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false + # clone-rapidsmpf to not use the rapidsmpf wheels from cudf_polars dependency + # librapidsmpf-cu13 wheels brings in a hardcoded libnuma-dev cmake target: https://github.com/NVIDIA/cuCascade/issues/118 + # -DBUILD_TESTS=OFF to match rapidsmpf https://github.com/rapidsai/rapidsmpf/blob/main/.github/workflows/pr.yaml#L351 (leads to compilation errors) + # -DCUDF_BUILD_TESTUTIL=OFF to avoid IMPORTED_GLOBAL promotion errors when cuCascade's find_package(cudf) loads cudf-config.cmake from a CPM subdirectory build_command: | sccache --zero-stats; - build-all -j0 -DBUILD_BENCHMARKS=ON --verbose 2>&1 | tee telemetry-artifacts/build.log; + clone-rapidsmpf -j$(nproc) -v -q --branch "$(cat ~/cudf/RAPIDS_BRANCH)" --clone-upstream --depth 1 --single-branch --shallow-submodules; + if [ "$PYTHON_PACKAGE_MANAGER" = "pip" ]; then rapids-make-pip-env --force; elif [ "$PYTHON_PACKAGE_MANAGER" = "conda" ]; then rapids-make-conda-env --force; fi; + rapids-generate-scripts; + build-all -j0 -DBUILD_BENCHMARKS=OFF -DBUILD_NUMA_SUPPORT=OFF -DBUILD_TESTS=OFF -DCUDF_BUILD_TESTUTIL=OFF --verbose 2>&1 | tee telemetry-artifacts/build.log; sccache --show-adv-stats | tee telemetry-artifacts/sccache-stats.txt; unit-tests-cudf-pandas: needs: [wheel-build-cudf, changed-files] diff --git a/ci/test_cudf_polars_experimental.sh b/ci/test_cudf_polars_experimental.sh index 02eab86c0dd5..aa3abd66254b 100755 --- a/ci/test_cudf_polars_experimental.sh +++ b/ci/test_cudf_polars_experimental.sh @@ -28,7 +28,7 @@ rapids-pip-retry install \ -v \ --prefer-binary \ --constraint "${PIP_CONSTRAINT}" \ - "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,experimental,rapidsmpf]" \ + "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,experimental]" \ "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" diff --git a/ci/test_cudf_polars_polars_tests.sh b/ci/test_cudf_polars_polars_tests.sh index 802110b18ac5..4e19464a8951 100755 --- a/ci/test_cudf_polars_polars_tests.sh +++ b/ci/test_cudf_polars_polars_tests.sh @@ -27,7 +27,7 @@ rapids-logger "Install libcudf, pylibcudf and cudf_polars" rapids-pip-retry install \ -v \ --constraint "${PIP_CONSTRAINT}" \ - "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,rapidsmpf]" \ + "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]" \ "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 35939809d53d..3bf006400bce 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -85,6 +85,7 @@ dependencies: - rapids-build-backend>=0.4.0,<0.5.0 - rapids-dask-dependency==26.6.*,>=0.0.0a0 - rapids-logger==0.2.*,>=0.0.0a0 +- rapidsmpf==26.6.*,>=0.0.0a0 - rich - rmm==26.6.*,>=0.0.0a0 - s3fs>=2022.3.0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index d88d82e49990..d206963fc362 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -85,6 +85,7 @@ dependencies: - rapids-build-backend>=0.4.0,<0.5.0 - rapids-dask-dependency==26.6.*,>=0.0.0a0 - rapids-logger==0.2.*,>=0.0.0a0 +- rapidsmpf==26.6.*,>=0.0.0a0 - rich - rmm==26.6.*,>=0.0.0a0 - s3fs>=2022.3.0 diff --git a/conda/environments/all_cuda-131_arch-aarch64.yaml b/conda/environments/all_cuda-131_arch-aarch64.yaml index f423c19b51dc..c73cea89f296 100644 --- a/conda/environments/all_cuda-131_arch-aarch64.yaml +++ b/conda/environments/all_cuda-131_arch-aarch64.yaml @@ -85,6 +85,7 @@ dependencies: - rapids-build-backend>=0.4.0,<0.5.0 - rapids-dask-dependency==26.6.*,>=0.0.0a0 - rapids-logger==0.2.*,>=0.0.0a0 +- rapidsmpf==26.6.*,>=0.0.0a0 - rich - rmm==26.6.*,>=0.0.0a0 - s3fs>=2022.3.0 diff --git a/conda/environments/all_cuda-131_arch-x86_64.yaml b/conda/environments/all_cuda-131_arch-x86_64.yaml index 3c9bffa4738c..f4d888666905 100644 --- a/conda/environments/all_cuda-131_arch-x86_64.yaml +++ b/conda/environments/all_cuda-131_arch-x86_64.yaml @@ -85,6 +85,7 @@ dependencies: - rapids-build-backend>=0.4.0,<0.5.0 - rapids-dask-dependency==26.6.*,>=0.0.0a0 - rapids-logger==0.2.*,>=0.0.0a0 +- rapidsmpf==26.6.*,>=0.0.0a0 - rich - rmm==26.6.*,>=0.0.0a0 - s3fs>=2022.3.0 diff --git a/conda/recipes/cudf-polars/recipe.yaml b/conda/recipes/cudf-polars/recipe.yaml index e3a21aa1afd8..52ac74c7c8bc 100644 --- a/conda/recipes/cudf-polars/recipe.yaml +++ b/conda/recipes/cudf-polars/recipe.yaml @@ -36,6 +36,7 @@ requirements: - nvidia-ml-py>=12 - python - pylibcudf =${{ version }} + - rapidsmpf =${{ minor_version }} - polars>=1.30,<1.39 - packaging - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} diff --git a/dependencies.yaml b/dependencies.yaml index 9728b3ea6679..b1eb276befb7 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -8,6 +8,51 @@ files: matrix: cuda: ["12.9", "13.1"] arch: [x86_64, aarch64] + includes: + - build_base + - build_all + - build_cpp + - build_python_common + - clang + - cuda + - cuda_version + - cudf_polars_trace + - depends_on_cupy + - depends_on_cuda_python + - depends_on_dask_cuda + - depends_on_libkvikio + - depends_on_librmm + - depends_on_libnvcomp + - depends_on_numba_cuda + - depends_on_rapids_logger + - depends_on_rapidsmpf + - depends_on_rmm + - develop + - docs + - iwyu + - notebooks + - numpy_run + - py_version + - pyarrow_run + - rapids_build_skbuild + - rapids_build_setuptools + - run_common + - run_cudf + - run_cudf_polars + - run_pylibcudf + - run_dask_cudf + - run_custreamz + - test_cpp + - test_python_common + - test_python_cudf + - test_python_cudf_common + - test_python_pylibcudf + - test_python_cudf_pandas + - test_python_cudf_polars + - test_python_s3 + devcontainers: + output: none + # Same as "all", excluding depends_on_rapidsmpf (which is built from source) includes: - build_base - build_all @@ -328,6 +373,7 @@ files: table: project includes: - run_cudf_polars + - depends_on_rapidsmpf - depends_on_pylibcudf - depends_on_cuda_python py_run_cudf_polars_experimental: @@ -338,15 +384,6 @@ files: key: experimental includes: - run_cudf_polars_experimental - py_run_cudf_polars_rapidsmpf: - output: pyproject - pyproject_dir: python/cudf_polars - extras: - table: project.optional-dependencies - key: rapidsmpf - includes: - - depends_on_rapidsmpf - - pyarrow_run py_test_cudf_polars: output: pyproject pyproject_dir: python/cudf_polars @@ -360,7 +397,6 @@ files: - test_python_common - test_python_cudf_polars - cudf_polars_trace - - depends_on_rapidsmpf py_trace_cudf_polars: output: pyproject pyproject_dir: python/cudf_polars diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 741b2f6f758c..823954a3b087 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -13,7 +13,6 @@ For the most part, the public API of `cudf-polars` is the polars API. ExecutorType, InMemoryExecutor, ParquetOptions, - ShuffleMethod, StreamingExecutor, StreamingFallbackMode, ``` diff --git a/docs/cudf/source/cudf_polars/engine_options.md b/docs/cudf/source/cudf_polars/engine_options.md index 67e601467d9f..ba6085275b87 100644 --- a/docs/cudf/source/cudf_polars/engine_options.md +++ b/docs/cudf/source/cudf_polars/engine_options.md @@ -52,7 +52,7 @@ For example, the environment variable `max_rows_per_partition` to use if it isn't overridden through `executor_options`. -For boolean options, like `rapidsmpf_spill`, the values `{"1", "true", "yes", "y"}` +For boolean options, like `sink_to_directory`, the values `{"1", "true", "yes", "y"}` are considered `True` and `{"0", "false", "no", "n"}` are considered `False`. See [Configuration Reference](#cudf-polars-api) for a full list of options, and diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index fb915784f964..acd0452ae1bd 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -159,13 +159,12 @@ def set_memory_resource( """ previous = rmm.mr.get_current_device_resource() if mr is None: - # Use cuda async by default with the rapidsmpf runtime. + # Use cuda async by default with the streaming executor. if ( memory_resource_config is None and executor.name == "streaming" - and executor.runtime == "rapidsmpf" and (device_size := get_total_device_memory()) is not None - ): # pragma: no cover; Requires rapidsmpf runtime. + ): # pragma: no cover memory_resource_config = MemoryResourceConfig( qualname="rmm.mr.CudaAsyncMemoryResource", options={ diff --git a/python/cudf_polars/cudf_polars/dsl/expr.py b/python/cudf_polars/cudf_polars/dsl/expr.py index 5dd8898bde21..b21485ac41e8 100644 --- a/python/cudf_polars/cudf_polars/dsl/expr.py +++ b/python/cudf_polars/cudf_polars/dsl/expr.py @@ -15,7 +15,6 @@ from cudf_polars.dsl.expressions.aggregation import Agg from cudf_polars.dsl.expressions.base import ( - AggInfo, Col, ColRef, ErrorExpr, @@ -37,7 +36,6 @@ __all__ = [ "Agg", - "AggInfo", "BinOp", "BooleanFunction", "Cast", diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/base.py b/python/cudf_polars/cudf_polars/dsl/expressions/base.py index 3336ea3fd7ca..b97be71b771f 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/base.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/base.py @@ -8,7 +8,7 @@ import enum from enum import IntEnum -from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple +from typing import TYPE_CHECKING, Any, ClassVar import pylibcudf as plc @@ -20,11 +20,7 @@ from cudf_polars.containers import Column, DataFrame, DataType -__all__ = ["AggInfo", "Col", "ColRef", "ExecutionContext", "Expr", "NamedExpr"] - - -class AggInfo(NamedTuple): - requests: list[tuple[Expr | None, plc.aggregation.Aggregation, Expr]] +__all__ = ["Col", "ColRef", "ExecutionContext", "Expr", "NamedExpr"] class ExecutionContext(IntEnum): diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index bee30183e1c8..1c48f70bb114 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -819,11 +819,13 @@ def read_csv_header( # TODO: Nested column names names = chunk.column_names(include_children=False) concatenated_columns = chunk.tbl.columns() - while reader.has_next(): + while reader.has_next(): # pragma: no cover columns = reader.read_chunk().tbl.columns() # Discard columns while concatenating to reduce memory footprint. # Reverse order to avoid O(n^2) list popping cost. - for i in range(len(concatenated_columns) - 1, -1, -1): + for i in range( # pragma: no cover + len(concatenated_columns) - 1, -1, -1 + ): concatenated_columns[i] = plc.concatenate.concatenate( [concatenated_columns[i], columns.pop()], stream=stream ) @@ -840,7 +842,7 @@ def read_csv_header( num_rows=num_rows, ) if include_file_paths is not None: - df = Scan.add_file_paths( + df = Scan.add_file_paths( # pragma: no cover include_file_paths, paths, chunk.num_rows_per_source, df ) else: @@ -1164,7 +1166,7 @@ def _write_parquet( | plc.io.parquet.ParquetWriterOptionsBuilder ) - if ( + if ( # pragma: no cover parquet_options.chunked and parquet_options.n_output_chunks != 1 and df.table.num_rows() != 0 diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 73ed9b3dbe1e..80ff0dfacbdf 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -11,11 +11,10 @@ from cudf_polars.dsl.traversal import traversal if TYPE_CHECKING: - from collections.abc import Generator, Iterator + from collections.abc import Generator from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR - from cudf_polars.dsl.nodebase import Node class PartitionInfo: @@ -40,22 +39,12 @@ def __init__( self.partitioned_on = partitioned_on self.io_plan = io_plan - def keys(self, node: Node) -> Iterator[tuple[str, int]]: - """Return the partitioned keys for a given node.""" - name = get_key_name(node) - yield from ((name, i) for i in range(self.count)) - def __rich_repr__(self) -> Generator[Any, None, None]: """Formatting for rich.pretty.pprint.""" yield "count", self.count yield "partitioned_on", self.partitioned_on -def get_key_name(node: Node) -> str: - """Generate the key name for a Node.""" - return f"{type(node).__name__.lower()}-{hash(node)}" - - class SerializedDataSourceInfo(TypedDict): """The serialized form of DataSourceInfo.""" diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py index d514d4c44e9e..743869937377 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py @@ -636,26 +636,8 @@ def get_executor_options( executor_options: dict[str, Any] = ( run_config.streaming_options.to_executor_options() ) - executor_options["runtime"] = "rapidsmpf" executor_options["max_io_threads"] = run_config.max_io_threads - # PDSHQueries: inject unique_fraction when dynamic planning is explicitly disabled - if ( - benchmark - and benchmark.__name__ == "PDSHQueries" - and run_config.executor == "streaming" - and run_config.streaming_options.dynamic_planning is None - ): - executor_options.setdefault( - "unique_fraction", - { - "c_custkey": 0.05, - "l_orderkey": 1.0, - "l_partkey": 0.1, - "o_custkey": 0.25, - }, - ) - return executor_options @@ -1110,8 +1092,7 @@ def run_polars_spmd( from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine executor_options = get_executor_options(run_config, benchmark=benchmark) - # "runtime" and "cluster" are reserved — SPMDEngine sets them - executor_options.pop("runtime", None) + # "cluster" is reserved — SPMDEngine sets it executor_options.pop("cluster", None) engine_options = { **run_config.streaming_options.to_engine_options(), @@ -1168,8 +1149,7 @@ def run_polars_ray( from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine executor_options = get_executor_options(run_config, benchmark=benchmark) - # "runtime", "cluster" are reserved — RayEngine sets them - executor_options.pop("runtime", None) + # "cluster" is reserved — RayEngine sets it executor_options.pop("cluster", None) engine_options: dict[str, Any] = { **run_config.streaming_options.to_engine_options(), @@ -1218,8 +1198,7 @@ def run_polars_dask( from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine executor_options = get_executor_options(run_config, benchmark=benchmark) - # "runtime", "cluster" are reserved — DaskEngine sets them - executor_options.pop("runtime", None) + # "cluster" is reserved — DaskEngine sets it executor_options.pop("cluster", None) engine_options: dict[str, Any] = { **run_config.streaming_options.to_engine_options(), diff --git a/python/cudf_polars/cudf_polars/experimental/dispatch.py b/python/cudf_polars/cudf_polars/experimental/dispatch.py index 3ac67b6af463..9ff0cc3156b9 100644 --- a/python/cudf_polars/cudf_polars/experimental/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/dispatch.py @@ -5,7 +5,7 @@ from __future__ import annotations from functools import singledispatch -from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict +from typing import TYPE_CHECKING, TypeAlias, TypedDict from cudf_polars.typing import GenericTransformer @@ -13,7 +13,7 @@ from collections.abc import MutableMapping from cudf_polars.dsl import ir - from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import ( PartitionInfo, StatsCollector, @@ -72,38 +72,3 @@ def lower_ir_node( lower_ir_graph """ raise AssertionError(f"Unhandled type {type(ir)}") # pragma: no cover - - -@singledispatch -def generate_ir_tasks( - ir: IR, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - """ - Generate a task graph for evaluation of an IR node. - - Parameters - ---------- - ir - IR node to generate tasks for. - partition_info - Partitioning information, obtained from :func:`lower_ir_graph`. - context - Runtime context for IR node execution. - - Returns - ------- - mapping - A (partial) dask task graph for the evaluation of an ir node. - - Notes - ----- - Task generation should only produce the tasks for the current node, - referring to child tasks by name. - - See Also - -------- - task_graph - """ - raise AssertionError(f"Unhandled type {type(ir)}") # pragma: no cover diff --git a/python/cudf_polars/cudf_polars/experimental/distinct.py b/python/cudf_polars/cudf_polars/experimental/distinct.py index 9ae148f77d30..564fe570919d 100644 --- a/python/cudf_polars/cudf_polars/experimental/distinct.py +++ b/python/cudf_polars/cudf_polars/experimental/distinct.py @@ -17,8 +17,6 @@ from cudf_polars.experimental.shuffle import Shuffle from cudf_polars.experimental.utils import ( _dynamic_planning_on, - _fallback_inform, - _get_unique_fractions, _lower_ir_fallback, ) @@ -35,8 +33,6 @@ def lower_distinct( child: IR, partition_info: MutableMapping[IR, PartitionInfo], config_options: ConfigOptions[StreamingExecutor], - *, - unique_fraction: float | None = None, ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: """ Lower a Distinct IR into partition-wise stages. @@ -56,9 +52,6 @@ def lower_distinct( associated partitioning information. config_options GPUEngine configuration options. - unique_fraction - Fraction of unique values to total values. Used for algorithm selection. - A value of `1.0` means the column is unique. Returns ------- @@ -68,69 +61,24 @@ def lower_distinct( A mapping from unique nodes in the new graph to associated partitioning information. """ - subset: frozenset[str] = ir.subset or frozenset(ir.schema) - distinct_keys = tuple( - NamedExpr(name, Col(ir.schema[name], name)) - for name in ir.schema - if name in subset - ) - child_count = partition_info[child].count - shuffled = partition_info[child].partitioned_on == distinct_keys - # Check for ordering requirements (shuffle is not stable) - require_tree_reduction = ir.stable or ir.keep in ( - plc.stream_compaction.DuplicateKeepOption.KEEP_FIRST, - plc.stream_compaction.DuplicateKeepOption.KEEP_LAST, - ) - - output_count = 1 - n_ary = 32 # Arbitrary default (for now) + n_ary = 32 if ir.zlice is not None and ir.zlice[1] is not None: - # Head/tail slice operation has been pushed into Distinct - # (caller ensures only simple slices reach here) n_ary = max(1_000_000 // ir.zlice[1], 2) - elif unique_fraction is not None: - # Use unique_fraction to determine partitioning - n_ary = min(max(int(1.0 / unique_fraction), 2), child_count) - output_count = max(int(unique_fraction * child_count), 1) - - if output_count > 1 and require_tree_reduction: - # Need to reduce down to a single partition even - # if the unique_fraction is large. - output_count = 1 - _fallback_inform( - "Unsupported unique options for multiple partitions.", - config_options, - ) # Partition-wise unique count = child_count new_node: IR = ir.reconstruct([child]) partition_info[new_node] = PartitionInfo(count=count) - if shuffled or output_count == 1: - # Tree reduction - while count > output_count: - new_node = Repartition(new_node.schema, new_node) - count = max(math.ceil(count / n_ary), output_count) - partition_info[new_node] = PartitionInfo(count=count) - new_node = ir.reconstruct([new_node]) - partition_info[new_node] = PartitionInfo(count=count) - else: - # Shuffle - new_node = Shuffle( - new_node.schema, - distinct_keys, - config_options.executor.shuffle_method, - new_node, - ) - partition_info[new_node] = PartitionInfo(count=output_count) + # Tree reduction + while count > 1: + new_node = Repartition(new_node.schema, new_node) + count = max(math.ceil(count / n_ary), 1) + partition_info[new_node] = PartitionInfo(count=count) new_node = ir.reconstruct([new_node]) - partition_info[new_node] = PartitionInfo( - count=output_count, - partitioned_on=distinct_keys, - ) + partition_info[new_node] = PartitionInfo(count=count) return new_node, partition_info @@ -172,7 +120,6 @@ def _( child = Shuffle( child.schema, distinct_keys, - config_options.executor.shuffle_method, child, ) partition_info[child] = PartitionInfo( @@ -202,19 +149,9 @@ def _( ) return dynamic_node, partition_info - # Non-dynamic planning: use unique_fraction heuristics - unique_fraction_dict = _get_unique_fractions( - tuple(subset), - config_options.executor.unique_fraction, - ) - unique_fraction = ( - max(unique_fraction_dict.values()) if unique_fraction_dict else None - ) - return lower_distinct( ir, child, partition_info, config_options, - unique_fraction=unique_fraction, ) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 82f023b229c5..d50d9fae0aed 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -288,10 +288,7 @@ def _(ir: GroupBy) -> dict[str, Serializable]: @_serialize_properties.register def _(ir: Shuffle) -> dict[str, Serializable]: - return { - "keys": [ne.name for ne in ir.keys], - "shuffle_method": ir.shuffle_method.value, - } + return {"keys": [ne.name for ne in ir.keys]} @_serialize_properties.register diff --git a/python/cudf_polars/cudf_polars/experimental/expressions.py b/python/cudf_polars/cudf_polars/experimental/expressions.py index d2a0070d0090..d6df4cae8f9c 100644 --- a/python/cudf_polars/cudf_polars/experimental/expressions.py +++ b/python/cudf_polars/cudf_polars/experimental/expressions.py @@ -41,22 +41,18 @@ from cudf_polars.containers import DataType from cudf_polars.dsl.expressions.aggregation import Agg -from cudf_polars.dsl.expressions.base import Col, ExecutionContext, Expr, NamedExpr +from cudf_polars.dsl.expressions.base import Col, ExecutionContext, NamedExpr from cudf_polars.dsl.expressions.binaryop import BinOp from cudf_polars.dsl.expressions.literal import Literal from cudf_polars.dsl.expressions.ternary import Ternary from cudf_polars.dsl.expressions.unary import Cast, Len, UnaryFunction -from cudf_polars.dsl.ir import IR, Distinct, Empty, HConcat, Select +from cudf_polars.dsl.ir import Distinct, Empty, HConcat, Select from cudf_polars.dsl.traversal import ( CachingVisitor, ) from cudf_polars.experimental.base import PartitionInfo from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import ( - _dynamic_planning_on, - _get_unique_fractions, - _leaf_column_names, -) +from cudf_polars.experimental.utils import _dynamic_planning_on if TYPE_CHECKING: from collections.abc import Generator, MutableMapping, Sequence @@ -197,15 +193,6 @@ def _decompose_unique( ) (column,) = columns - unique_fraction_dict = _get_unique_fractions( - _leaf_column_names(child), - config_options.executor.unique_fraction, - ) - - unique_fraction = ( - max(unique_fraction_dict.values()) if unique_fraction_dict else None - ) - input_ir, partition_info = lower_distinct( Distinct( {column.name: column.dtype}, @@ -218,7 +205,6 @@ def _decompose_unique( input_ir, partition_info, config_options, - unique_fraction=unique_fraction, ) return column, input_ir, partition_info @@ -344,7 +330,6 @@ def _decompose_agg_node( input_ir = Shuffle( input_ir.schema, shuffle_on, - config_options.executor.shuffle_method, input_ir, ) partition_info[input_ir] = PartitionInfo( diff --git a/python/cudf_polars/cudf_polars/experimental/groupby.py b/python/cudf_polars/cudf_polars/experimental/groupby.py index 898dfdbf03f1..6a17b56bfc53 100644 --- a/python/cudf_polars/cudf_polars/experimental/groupby.py +++ b/python/cudf_polars/cudf_polars/experimental/groupby.py @@ -36,7 +36,6 @@ from cudf_polars.experimental.shuffle import Shuffle from cudf_polars.experimental.utils import ( _dynamic_planning_on, - _get_unique_fractions, _lower_ir_fallback, ) @@ -390,7 +389,6 @@ def _( # Check if we are dealing with any high-cardinality columns post_aggregation_count = 1 # Default tree reduction - groupby_key_columns = [ne.name for ne in ir.keys] shuffled = partition_info[child].partitioned_on == ir.keys child_count = partition_info[child].count @@ -421,7 +419,6 @@ def _( child = Shuffle( child.schema, ir.keys, - config_options.executor.shuffle_method, child, ) partition_info[child] = PartitionInfo( @@ -441,14 +438,6 @@ def _( ) return dynamic_node, partition_info - if unique_fraction_dict := _get_unique_fractions( - groupby_key_columns, - config_options.executor.unique_fraction, - ): - # Use unique_fraction to determine output partitioning - unique_fraction = max(unique_fraction_dict.values()) - post_aggregation_count = max(int(unique_fraction * child_count), 1) - # Partition-wise groupby operation pwise_schema = {k.name: k.value.dtype for k in ir.keys} | { k.name: k.value.dtype for k in piecewise_exprs @@ -465,46 +454,28 @@ def _( partition_info[gb_pwise] = PartitionInfo(count=child_count) grouped_keys = tuple(NamedExpr(k.name, Col(k.value.dtype, k.name)) for k in ir.keys) - # Reduction - gb_inter: GroupBy | Repartition | Shuffle + # N-ary tree reduction + gb_inter: GroupBy | Repartition reduction_schema = {k.name: k.value.dtype for k in grouped_keys} | { k.name: k.value.dtype for k in reduction_exprs } - if not shuffled and post_aggregation_count > 1: - # Shuffle reduction - if ir.maintain_order: # pragma: no cover - return _lower_ir_fallback( - ir, - rec, - msg="maintain_order not supported for multiple output partitions.", + n_ary = 32 + count = child_count + gb_inter = gb_pwise + while count > post_aggregation_count: + gb_inter = Repartition(gb_inter.schema, gb_inter) + count = max(math.ceil(count / n_ary), post_aggregation_count) + partition_info[gb_inter] = PartitionInfo(count=count) + if count > post_aggregation_count: + gb_inter = GroupBy( + reduction_schema, + grouped_keys, + reduction_exprs, + ir.maintain_order, + None, + gb_inter, ) - - gb_inter = Shuffle( - gb_pwise.schema, - grouped_keys, - config_options.executor.shuffle_method, - gb_pwise, - ) - partition_info[gb_inter] = PartitionInfo(count=post_aggregation_count) - else: - # N-ary tree reduction - n_ary = config_options.executor.groupby_n_ary - count = child_count - gb_inter = gb_pwise - while count > post_aggregation_count: - gb_inter = Repartition(gb_inter.schema, gb_inter) - count = max(math.ceil(count / n_ary), post_aggregation_count) partition_info[gb_inter] = PartitionInfo(count=count) - if count > post_aggregation_count: - gb_inter = GroupBy( - reduction_schema, - grouped_keys, - reduction_exprs, - ir.maintain_order, - None, - gb_inter, - ) - partition_info[gb_inter] = PartitionInfo(count=count) # Final aggregation gb_reduce = GroupBy( diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/experimental/io.py index f45baa054dd8..2cea0274ee66 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/python/cudf_polars/cudf_polars/experimental/io.py @@ -4,13 +4,11 @@ from __future__ import annotations -import dataclasses import functools import itertools import math import statistics from collections import defaultdict -from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload @@ -24,16 +22,14 @@ Empty, Scan, Sink, - Union, ) from cudf_polars.experimental.base import ( IOPartitionFlavor, IOPartitionPlan, PartitionInfo, SerializedDataSourceInfo, - get_key_name, ) -from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node +from cudf_polars.experimental.dispatch import lower_ir_node from cudf_polars.utils.config import Cluster from cudf_polars.utils.cuda_stream import get_cuda_stream from cudf_polars.utils.versions import POLARS_VERSION_LT_137 @@ -62,36 +58,9 @@ def _( ir: DataFrameScan, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - config_options = rec.state["config_options"] + from cudf_polars.experimental.rapidsmpf.io import lower_dataframescan_rapidsmpf - # RapidsMPF runtime: Use rapidsmpf-specific lowering - if ( - config_options.executor.runtime == "rapidsmpf" - ): # pragma: no cover; Requires rapidsmpf runtime - from cudf_polars.experimental.rapidsmpf.io import lower_dataframescan_rapidsmpf - - return lower_dataframescan_rapidsmpf(ir, rec) - - rows_per_partition = config_options.executor.max_rows_per_partition - nrows = max(ir.df.shape()[0], 1) - count = math.ceil(nrows / rows_per_partition) - - if count > 1: - length = math.ceil(nrows / count) - slices = [ - DataFrameScan( - ir.schema, - ir.df.slice(offset, length), - ir.projection, - ) - for offset in range(0, nrows, length) - ] - new_node = Union(ir.schema, None, *slices) - return new_node, {slice: PartitionInfo(count=1) for slice in slices} | { - new_node: PartitionInfo(count=count) - } - - return ir, {ir: PartitionInfo(count=1)} + return lower_dataframescan_rapidsmpf(ir, rec) def scan_partition_plan( @@ -285,84 +254,9 @@ def _( def _( ir: Scan, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - partition_info: MutableMapping[IR, PartitionInfo] - config_options = rec.state["config_options"] - - # RapidsMPF runtime: Use rapidsmpf-specific lowering - if ( - config_options.executor.name == "streaming" - and config_options.executor.runtime == "rapidsmpf" - ): # pragma: no cover; Requires rapidsmpf runtime - from cudf_polars.experimental.rapidsmpf.io import lower_scan_rapidsmpf + from cudf_polars.experimental.rapidsmpf.io import lower_scan_rapidsmpf - return lower_scan_rapidsmpf(ir, rec) - - if ( - ir.typ in ("csv", "parquet", "ndjson") - and ir.n_rows == -1 - and ir.skip_rows == 0 - and ir.row_index is None - ): - plan = scan_partition_plan(ir, rec.state["stats"], config_options) - paths = list(ir.paths) - if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - # Disable chunked reader when splitting files - parquet_options = dataclasses.replace( - config_options.parquet_options, - chunked=False, - ) - - slices: list[SplitScan] = [] - for path in paths: - base_scan = Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - [path], - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) - slices.extend( - SplitScan( - ir.schema, base_scan, sindex, plan.factor, parquet_options - ) - for sindex in range(plan.factor) - ) - new_node = Union(ir.schema, None, *slices) - partition_info = {slice: PartitionInfo(count=1) for slice in slices} | { - new_node: PartitionInfo(count=len(slices)) - } - else: - groups: list[Scan] = [ - Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - paths[i : i + plan.factor], - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - config_options.parquet_options, - ) - for i in range(0, len(paths), plan.factor) - ] - new_node = Union(ir.schema, None, *groups) - partition_info = {group: PartitionInfo(count=1) for group in groups} | { - new_node: PartitionInfo(count=len(groups)) - } - return new_node, partition_info - - return ir, {ir: PartitionInfo(count=1)} # pragma: no cover + return lower_scan_rapidsmpf(ir, rec) class StreamingSink(IR): @@ -441,22 +335,6 @@ def _prepare_sink_directory(path: str) -> None: Path(path).mkdir(parents=True, exist_ok=True) -def _sink_to_directory( - schema: Schema, - kind: str, - path: str, - parquet_options: ParquetOptions, - options: dict[str, Any], - df: DataFrame, - ready: None, - context: IRExecutionContext, -) -> DataFrame: - """Sink a partition to a new file.""" - return Sink.do_evaluate( - schema, kind, path, parquet_options, options, df, context=context - ) - - def _sink_to_parquet_file( path: str, options: dict[str, Any], @@ -545,106 +423,6 @@ def _sink_to_file( return True -def _finalize_file_sink( - kind: str, - writer_state: Any, - df: DataFrame, -) -> DataFrame: - """Finalize the file sink by closing the writer.""" - if kind == "Parquet" and writer_state is not None: - writer_state.close([]) - return df.slice((0, 0)) - - -def _file_sink_graph( - ir: StreamingSink, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - """Sink to a single file.""" - name = get_key_name(ir) - count = partition_info[ir].count - child_name = get_key_name(ir.children[0]) - sink = ir.sink - if count == 1: - return { - (name, 0): ( - partial(sink.do_evaluate, context=context), - *sink._non_child_args, - (child_name, 0), - ) - } - - sink_name = get_key_name(sink) - graph: MutableMapping[Any, Any] = { - (sink_name, i): ( - _sink_to_file, - sink.kind, - sink.path, - sink.options, - None if i == 0 else (sink_name, i - 1), # Writer state - (child_name, i), - ) - for i in range(count) - } - - # Finalize task closes the writer after all chunks are written - graph[(sink_name, "finalize")] = ( - _finalize_file_sink, - sink.kind, - (sink_name, count - 1), # Writer state from last task - (child_name, count - 1), # Last source df for creating empty result - ) - - # Make sure final tasks point to finalize task - graph.update({(name, i): (sink_name, "finalize") for i in range(count)}) - return graph - - -def _directory_sink_graph( - ir: StreamingSink, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - """Sink to a directory of files.""" - name = get_key_name(ir) - count = partition_info[ir].count - child_name = get_key_name(ir.children[0]) - sink = ir.sink - - setup_name = f"setup-{name}" - suffix = sink.kind.lower() - width = math.ceil(math.log10(count)) - graph: MutableMapping[Any, Any] = { - (name, i): ( - _sink_to_directory, - sink.schema, - sink.kind, - f"{sink.path}/part.{str(i).zfill(width)}.{suffix}", - sink.parquet_options, - sink.options, - (child_name, i), - setup_name, - context, - ) - for i in range(count) - } - graph[setup_name] = (_prepare_sink_directory, sink.path) - return graph - - -@generate_ir_tasks.register(StreamingSink) -def _( - ir: StreamingSink, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - if ir.sink_to_directory: - return _directory_sink_graph(ir, partition_info, context=context) - else: - return _file_sink_graph(ir, partition_info, context=context) - - class ParquetMetadata: """ Parquet metadata container. diff --git a/python/cudf_polars/cudf_polars/experimental/join.py b/python/cudf_polars/cudf_polars/experimental/join.py index 47d0ad90d8e2..cd5c514b45a0 100644 --- a/python/cudf_polars/cudf_polars/experimental/join.py +++ b/python/cudf_polars/cudf_polars/experimental/join.py @@ -5,16 +5,15 @@ from __future__ import annotations import operator -from functools import partial, reduce -from typing import TYPE_CHECKING, Any +from functools import reduce +from typing import TYPE_CHECKING from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice -from cudf_polars.experimental.base import PartitionInfo, get_key_name -from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node +from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.experimental.dispatch import lower_ir_node from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import Shuffle, _hash_partition_dataframe +from cudf_polars.experimental.shuffle import Shuffle from cudf_polars.experimental.utils import ( - _concat, _dynamic_planning_on, _fallback_inform, _lower_ir_fallback, @@ -24,16 +23,14 @@ from collections.abc import MutableMapping from cudf_polars.dsl.expr import NamedExpr - from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.dsl.ir import IR from cudf_polars.experimental.parallel import LowerIRTransformer - from cudf_polars.utils.config import ShuffleMethod def _maybe_shuffle_frame( frame: IR, on: tuple[NamedExpr, ...], partition_info: MutableMapping[IR, PartitionInfo], - shuffle_method: ShuffleMethod, output_count: int, ) -> IR: # Shuffle `frame` if it isn't already shuffled. @@ -48,7 +45,6 @@ def _maybe_shuffle_frame( frame = Shuffle( frame.schema, on, - shuffle_method, frame, ) partition_info[frame] = PartitionInfo( @@ -64,21 +60,18 @@ def _make_hash_join( partition_info: MutableMapping[IR, PartitionInfo], left: IR, right: IR, - shuffle_method: ShuffleMethod, ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: # Shuffle left and right dataframes (if necessary) left = _maybe_shuffle_frame( left, ir.left_on, partition_info, - shuffle_method, output_count, ) right = _maybe_shuffle_frame( right, ir.right_on, partition_info, - shuffle_method, output_count, ) # Always reconstruct in case children contain Cache nodes @@ -146,45 +139,7 @@ def _make_bcast_join( partition_info: MutableMapping[IR, PartitionInfo], left: IR, right: IR, - shuffle_method: ShuffleMethod, - *, - streaming_runtime: str, ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - if ir.options[0] != "Inner": - left_count = partition_info[left].count - right_count = partition_info[right].count - - # Shuffle the smaller table (if necessary) - Notes: - # - We need to shuffle the smaller table if - # (1) we are not doing an "inner" join, - # and (2) the small table contains multiple - # partitions. - # - We cannot simply join a large-table partition - # to each small-table partition, and then - # concatenate the partial-join results, because - # a non-"inner" join does NOT commute with - # concatenation. - # - In some cases, we can perform the partial joins - # sequentially. However, we are starting with a - # catch-all algorithm that works for all cases. - if streaming_runtime == "tasks": - if left_count >= right_count: - right = _maybe_shuffle_frame( - right, - ir.right_on, - partition_info, - shuffle_method, - right_count, - ) - else: - left = _maybe_shuffle_frame( - left, - ir.left_on, - partition_info, - shuffle_method, - left_count, - ) - new_node = ir.reconstruct([left, right]) partition_info[new_node] = PartitionInfo(count=output_count) return new_node, partition_info @@ -301,8 +256,6 @@ def _( partition_info, left, right, - config_options.executor.shuffle_method, - streaming_runtime=config_options.executor.runtime, ) else: # Create a hash join @@ -312,109 +265,4 @@ def _( partition_info, left, right, - config_options.executor.shuffle_method, ) - - -@generate_ir_tasks.register(Join) -def _( - ir: Join, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - left, right = ir.children - output_count = partition_info[ir].count - - left_partitioned = ( - partition_info[left].partitioned_on == ir.left_on - and partition_info[left].count == output_count - ) - right_partitioned = ( - partition_info[right].partitioned_on == ir.right_on - and partition_info[right].count == output_count - ) - - if output_count == 1 or (left_partitioned and right_partitioned): - # Partition-wise join - left_name = get_key_name(left) - right_name = get_key_name(right) - return { - key: ( - partial(ir.do_evaluate, context=context), - *ir._non_child_args, - (left_name, i), - (right_name, i), - ) - for i, key in enumerate(partition_info[ir].keys(ir)) - } - else: - # Broadcast join - left_parts = partition_info[left] - right_parts = partition_info[right] - if left_parts.count >= right_parts.count: - small_side = "Right" - small_name = get_key_name(right) - small_size = partition_info[right].count - large_name = get_key_name(left) - large_on = ir.left_on - else: - small_side = "Left" - small_name = get_key_name(left) - small_size = partition_info[left].count - large_name = get_key_name(right) - large_on = ir.right_on - - graph: MutableMapping[Any, Any] = {} - - out_name = get_key_name(ir) - out_size = partition_info[ir].count - split_name = f"split-{out_name}" - getit_name = f"getit-{out_name}" - inter_name = f"inter-{out_name}" - - # Split each large partition if we have - # multiple small partitions (unless this - # is an inner join) - split_large = ir.options[0] != "Inner" and small_size > 1 - - for part_out in range(out_size): - if split_large: - graph[(split_name, part_out)] = ( - _hash_partition_dataframe, - (large_name, part_out), - part_out, - small_size, - None, - large_on, - ) - - _concat_list = [] - for j in range(small_size): - left_key: tuple[str, int] | tuple[str, int, int] - if split_large: - left_key = (getit_name, part_out, j) - graph[left_key] = (operator.getitem, (split_name, part_out), j) - else: - left_key = (large_name, part_out) - join_children = [left_key, (small_name, j)] - if small_side == "Left": - join_children.reverse() - - inter_key = (inter_name, part_out, j) - graph[(inter_name, part_out, j)] = ( - partial(ir.do_evaluate, context=context), - ir.left_on, - ir.right_on, - ir.options, - *join_children, - ) - _concat_list.append(inter_key) - if len(_concat_list) == 1: - graph[(out_name, part_out)] = graph.pop(_concat_list[0]) - else: - graph[(out_name, part_out)] = ( - partial(_concat, context=context), - *_concat_list, - ) - - return graph diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index f77e923bce03..ab5d3b5bd90b 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -4,10 +4,9 @@ from __future__ import annotations -import itertools import operator from functools import partial, reduce -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import polars as pl @@ -26,7 +25,6 @@ Filter, HConcat, HStack, - IRExecutionContext, MapFunction, Projection, Select, @@ -35,16 +33,11 @@ ) from cudf_polars.dsl.traversal import CachingVisitor, traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.base import PartitionInfo, get_key_name -from cudf_polars.experimental.dispatch import ( - generate_ir_tasks, - lower_ir_node, -) +from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.experimental.dispatch import lower_ir_node from cudf_polars.experimental.io import _clear_source_info_cache from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.statistics import collect_statistics from cudf_polars.experimental.utils import ( - _concat, _contains_over, _dynamic_planning_on, _lower_ir_fallback, @@ -52,7 +45,6 @@ if TYPE_CHECKING: from collections.abc import MutableMapping - from typing import Any from cudf_polars.experimental.base import StatsCollector from cudf_polars.experimental.dispatch import LowerIRTransformer, State @@ -109,63 +101,6 @@ def lower_ir_graph( return mapper(ir) -def task_graph( - ir: IR, - partition_info: MutableMapping[IR, PartitionInfo], -) -> tuple[MutableMapping[Any, Any], str | tuple[str, int]]: - """ - Construct a task graph for evaluation of an IR graph. - - Parameters - ---------- - ir - Root of the graph to rewrite. - partition_info - A mapping from all unique IR nodes to the - associated partitioning information. - - Returns - ------- - graph - A task graph for the entire IR graph with root `ir`, - in dict-of-tuples form consumed by - :func:`~cudf_polars.experimental.scheduler.synchronous_scheduler`. - - Notes - ----- - This function traverses the unique nodes of the - graph with root `ir`, and extracts the tasks for - each node with :func:`generate_ir_tasks`. - - See Also - -------- - generate_ir_tasks - """ - context = IRExecutionContext() - graph = reduce( - operator.or_, - ( - generate_ir_tasks(node, partition_info, context=context) - for node in traversal([ir]) - ), - ) - - key_name = get_key_name(ir) - partition_count = partition_info[ir].count - - key: str | tuple[str, int] - if partition_count > 1: - graph[key_name] = ( - partial(_concat, context=context), - *partition_info[ir].keys(ir), - ) - key = key_name - else: - key = (key_name, 0) - - return graph, key - - def evaluate_rapidsmpf( ir: IR, config_options: ConfigOptions[StreamingExecutor], @@ -211,44 +146,7 @@ def evaluate_streaming( # Clear source info cache in case data was overwritten _clear_source_info_cache() - if ( - config_options.executor.runtime == "rapidsmpf" - ): # pragma: no cover; rapidsmpf runtime not tested in CI yet - # Using the RapidsMPF streaming runtime. - return evaluate_rapidsmpf(ir, config_options) - else: - # Using the default task engine. - from cudf_polars.experimental.scheduler import synchronous_scheduler - - stats = collect_statistics(ir, config_options) - ir, partition_info = lower_ir_graph(ir, config_options, stats) - - graph, key = task_graph(ir, partition_info) - - return synchronous_scheduler(graph, key).to_polars() - - -@generate_ir_tasks.register(IR) -def _( - ir: IR, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - # Generate pointwise (embarrassingly-parallel) tasks by default - child_names = [get_key_name(c) for c in ir.children] - bcast_child = [partition_info[c].count == 1 for c in ir.children] - - return { - key: ( - partial(ir.do_evaluate, context=context), - *ir._non_child_args, - *[ - (child_name, 0 if bcast_child[j] else i) - for j, child_name in enumerate(child_names) - ], - ) - for i, key in enumerate(partition_info[ir].keys(ir)) - } + return evaluate_rapidsmpf(ir, config_options) @lower_ir_node.register(Union) @@ -278,21 +176,6 @@ def _( return new_node, partition_info -@generate_ir_tasks.register(Union) -def _( - ir: Union, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - key_name = get_key_name(ir) - partition = itertools.count() - return { - (key_name, next(partition)): child_key - for child in ir.children - for child_key in partition_info[child].keys(child) - } - - @lower_ir_node.register(MapFunction) def _( ir: MapFunction, rec: LowerIRTransformer diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 478c0a33bebb..97168f0b02dc 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -99,8 +99,6 @@ def evaluate_logical_plan( ------- The output DataFrame and metadata collector. """ - assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf" - query_id = uuid.uuid4() with cudf_polars.dsl.tracing.bound_contextvars( @@ -202,8 +200,6 @@ def evaluate_pipeline( ------- The output DataFrame and metadata collector. """ - assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf" - _original_mr: Any = None use_stream_pool = False if rmpf_context is not None: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py index 7bc8dabddecb..26ad95198f6d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py @@ -436,7 +436,7 @@ def execute_ir_on_rank( _RESERVED_EXECUTOR_KEYS: frozenset[str] = frozenset( - {"runtime", "cluster", "spmd_context", "ray_context", "dask_context"} + {"cluster", "spmd_context", "ray_context", "dask_context"} ) _RESERVED_ENGINE_KEYS: frozenset[str] = frozenset({"memory_resource", "executor"}) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py index 49810e998fd2..b4300346132d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py @@ -688,7 +688,6 @@ def __init__( nranks=nranks, executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "dask", "dask_context": dask_ctx, }, @@ -736,7 +735,6 @@ def _reset( nranks=self._nranks, executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "dask", "dask_context": ctx, }, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py index d8464aa7426a..c7650bff513e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py @@ -234,11 +234,6 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING``. Default: enabled. Category: executor. - unique_fraction - Per-column uniqueness estimate (0-1). Defaults to ``1.0``. - Env: ``CUDF_POLARS__EXECUTOR__UNIQUE_FRACTION``. - Default: ``{}``. - Category: executor. sink_to_directory Whether multi-partition sink operations should write to a directory rather than a single file. The ``spmd``/``ray``/``dask`` engines @@ -332,9 +327,6 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) - unique_fraction: dict[str, float] | Unspecified = _opt( - "executor", "CUDF_POLARS__EXECUTOR__UNIQUE_FRACTION", json.loads - ) sink_to_directory: bool | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean ) @@ -515,7 +507,6 @@ def _get(attr: str) -> Any: broadcast_join_limit=_get("broadcast_join_limit"), target_partition_size=target_partition_size, dynamic_planning=dynamic_planning, - unique_fraction=_get("unique_fraction"), raise_on_fail=_get("raise_on_fail"), parquet_options=_get("parquet_options"), memory_resource_config=_get("memory_resource_config"), @@ -711,15 +702,6 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None: Enable dynamic planning. Use --no-dynamic-planning to disable. Env: CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING. Built-in default: enabled."""), ) - g.add_argument( - "--unique-fraction", - dest="unique_fraction", - default=None, - type=json.loads, - help=textwrap.dedent("""\ - Per-column uniqueness estimate as a JSON object (e.g. '{"col": 0.5}'). - Env: CUDF_POLARS__EXECUTOR__UNIQUE_FRACTION. Built-in default: {}."""), - ) g.add_argument( "--stream-policy", dest="stream_policy", diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py index 1ba92de3e493..efbb1db9ad4a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py @@ -90,13 +90,9 @@ def evaluate_pipeline_ray_mode( Raises ------ - RuntimeError - If the configured executor runtime is not ``"rapidsmpf"``. RuntimeError If ``config_options.executor.ray_context`` is not set. """ - if config_options.executor.runtime != "rapidsmpf": - raise RuntimeError("Runtime must be rapidsmpf") if config_options.executor.ray_context is None: raise RuntimeError("ray_context must be set when cluster='ray'") rank_actors = config_options.executor.ray_context.rank_actors @@ -586,7 +582,6 @@ def __init__( nranks=nranks, executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "ray", "ray_context": RayContext(rank_actors), }, @@ -641,7 +636,6 @@ def _reset( nranks=len(self._rank_actors), executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "ray", "ray_context": RayContext(self._rank_actors), }, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py index 65e3eb8b1e7f..7e1bde808cd6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py @@ -23,12 +23,11 @@ from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context -import polars as pl - import pylibcudf as plc import rmm.mr from pylibcudf.contiguous_split import pack +from cudf_polars.containers import DataFrame, DataType from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id from cudf_polars.experimental.rapidsmpf.frontend.core import ( ClusterInfo, @@ -53,6 +52,8 @@ from rapidsmpf.config import Options from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata + import polars as pl + from cudf_polars.dsl.ir import IR from cudf_polars.experimental.parallel import ConfigOptions from cudf_polars.experimental.rapidsmpf.frontend.core import T @@ -98,8 +99,6 @@ def evaluate_pipeline_spmd_mode( The concatenated output DataFrame and, if ``collect_metadata`` is True, the list of channel metadata objects; otherwise ``None``. """ - if config_options.executor.runtime != "rapidsmpf": - raise RuntimeError("Runtime must be rapidsmpf") if config_options.executor.spmd_context is None: raise RuntimeError("spmd_context must be set for SPMD mode") comm = config_options.executor.spmd_context.comm @@ -155,8 +154,9 @@ def allgather_polars_dataframe( ctx = engine.context stream = ctx.get_stream_from_pool() col_names = local_df.columns + dtypes = [DataType(dtype) for dtype in local_df.dtypes] - plc_table = plc.Table.from_arrow(local_df.to_arrow()) + plc_table = plc.Table.from_arrow(local_df, stream=stream) packed_data = PackedData.from_cudf_packed_columns( pack(plc_table, stream), @@ -176,9 +176,12 @@ def allgather_polars_dataframe( plc_result = unpack_and_concat(results, stream, ctx.br()) # pylibcudf Table -> pl.DataFrame (restore column names) - ret = pl.from_arrow(plc_result.to_arrow(col_names)) - assert isinstance(ret, pl.DataFrame) - return ret + return DataFrame.from_table( + plc_result, + col_names, + dtypes, + stream, + ).to_polars() class SPMDEngine(StreamingEngine): @@ -389,7 +392,6 @@ def __init__( nranks=comm.nranks, executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "spmd", "spmd_context": SPMDContext( comm=comm, context=ctx, py_executor=self._py_executor @@ -494,7 +496,6 @@ def _reset( nranks=self._comm.nranks, executor_options={ **executor_options, - "runtime": "rapidsmpf", "cluster": "spmd", "spmd_context": SPMDContext( comm=self._comm, diff --git a/python/cudf_polars/cudf_polars/experimental/repartition.py b/python/cudf_polars/cudf_polars/experimental/repartition.py index 92d89a5f44cb..84c39d930ca9 100644 --- a/python/cudf_polars/cudf_polars/experimental/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/repartition.py @@ -4,20 +4,11 @@ from __future__ import annotations -import itertools -from functools import partial -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from cudf_polars.dsl.ir import IR -from cudf_polars.experimental.base import get_key_name -from cudf_polars.experimental.dispatch import generate_ir_tasks -from cudf_polars.experimental.utils import _concat if TYPE_CHECKING: - from collections.abc import MutableMapping - - from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.parallel import PartitionInfo from cudf_polars.typing import Schema @@ -43,35 +34,3 @@ def __init__(self, schema: Schema, df: IR): self.schema = schema self._non_child_args = () self.children = (df,) - - -@generate_ir_tasks.register(Repartition) -def _( - ir: Repartition, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - # Repartition an IR node. - # Only supports rapartitioning to fewer (for now). - - (child,) = ir.children - count_in = partition_info[child].count - count_out = partition_info[ir].count - - if count_out > count_in: # pragma: no cover - raise NotImplementedError( - f"Repartition {count_in} -> {count_out} not supported." - ) - - key_name = get_key_name(ir) - n, remainder = divmod(count_in, count_out) - # Spread remainder evenly over the partitions. - offsets = [0, *itertools.accumulate(n + (i < remainder) for i in range(count_out))] - child_keys = tuple(partition_info[child].keys(child)) - return { - (key_name, i): ( - partial(_concat, context=context), - *child_keys[offsets[i] : offsets[i + 1]], - ) - for i in range(count_out) - } diff --git a/python/cudf_polars/cudf_polars/experimental/scheduler.py b/python/cudf_polars/cudf_polars/experimental/scheduler.py deleted file mode 100644 index 97eae6ab3782..000000000000 --- a/python/cudf_polars/cudf_polars/experimental/scheduler.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Synchronous task scheduler.""" - -from __future__ import annotations - -from collections import Counter -from collections.abc import MutableMapping -from itertools import chain -from typing import TYPE_CHECKING, Any, TypeVar, Unpack - -if TYPE_CHECKING: - from collections.abc import Mapping - from typing import TypeAlias - - -Key: TypeAlias = str | tuple[str, Unpack[tuple[int, ...]]] -Graph: TypeAlias = MutableMapping[Key, Any] -T_ = TypeVar("T_") - - -# NOTE: This is a slimmed-down version of the single-threaded -# (synchronous) scheduler in `dask.core`. -# -# Key Differences: -# * We do not allow a task to contain a list of key names. -# Keys must be distinct elements of the task. -# * We do not support nested tasks. - - -def istask(x: Any) -> bool: - """Check if x is a callable task.""" - return isinstance(x, tuple) and bool(x) and callable(x[0]) - - -def is_hashable(x: Any) -> bool: - """Check if x is hashable.""" - try: - hash(x) - except BaseException: - return False - else: - return True - - -def _execute_task(arg: Any, cache: Mapping) -> Any: - """Execute a compute task.""" - if istask(arg): - return arg[0](*(_execute_task(a, cache) for a in arg[1:])) - elif is_hashable(arg): - return cache.get(arg, arg) - else: - return arg - - -def required_keys(key: Key, graph: Graph) -> list[Key]: - """ - Return the dependencies to extract a key from the graph. - - Parameters - ---------- - key - Root key we want to extract. - graph - The full task graph. - - Returns - ------- - List of other keys needed to extract ``key``. - """ - maybe_task = graph[key] - return [ - k - for k in ( - maybe_task[1:] - if istask(maybe_task) - else [maybe_task] # maybe_task might be a key - ) - if is_hashable(k) and k in graph - ] - - -def toposort(graph: Graph, dependencies: Mapping[Key, list[Key]]) -> list[Key]: - """Return a list of task keys sorted in topological order.""" - # Stack-based depth-first search traversal. This is based on Tarjan's - # algorithm for strongly-connected components - # (https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm) - ordered: list[Key] = [] - completed: set[Key] = set() - - for key in graph: - if key in completed: - continue - nodes = [key] - while nodes: - # Keep current node on the stack until all descendants are visited - current = nodes[-1] - if current in completed: # pragma: no cover - # Already fully traversed descendants of current - nodes.pop() - continue - - # Add direct descendants of current to nodes stack - next_nodes = set(dependencies[current]) - completed - if next_nodes: - nodes.extend(next_nodes) - else: - # Current has no more descendants to explore - ordered.append(current) - completed.add(current) - nodes.pop() - - return ordered - - -def synchronous_scheduler( - graph: Graph, - key: Key, - *, - cache: MutableMapping | None = None, -) -> Any: - """ - Execute the task graph for a given key. - - Parameters - ---------- - graph - The task graph to execute. - key - The final output key to extract from the graph. - cache - Intermediate-data cache. - - Returns - ------- - Executed task-graph result for ``key``. - """ - if key not in graph: # pragma: no cover - raise KeyError(f"{key} is not a key in the graph") - if cache is None: - cache = {} - - dependencies = {k: required_keys(k, graph) for k in graph} - refcount = Counter(chain.from_iterable(dependencies.values())) - - for k in toposort(graph, dependencies): - cache[k] = _execute_task(graph[k], cache) - for dep in dependencies[k]: - refcount[dep] -= 1 - if refcount[dep] == 0 and dep != key: - del cache[dep] - - return cache[key] diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/experimental/shuffle.py index 8e24dd83fe6c..9381126775f9 100644 --- a/python/cudf_polars/cudf_polars/experimental/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/shuffle.py @@ -4,112 +4,22 @@ from __future__ import annotations -import operator -from functools import partial -from typing import TYPE_CHECKING, Any, Concatenate, TypeVar, TypedDict +from typing import TYPE_CHECKING -import pylibcudf as plc -from rmm.pylibrmm.stream import DEFAULT_STREAM - -from cudf_polars.containers import DataFrame -from cudf_polars.dsl.expr import Col from cudf_polars.dsl.ir import IR -from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars -from cudf_polars.experimental.base import get_key_name -from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node -from cudf_polars.experimental.utils import _concat, _dynamic_planning_on -from cudf_polars.utils.cuda_stream import get_dask_cuda_stream +from cudf_polars.dsl.tracing import log_do_evaluate +from cudf_polars.experimental.dispatch import lower_ir_node +from cudf_polars.experimental.utils import _dynamic_planning_on if TYPE_CHECKING: - from collections.abc import Callable, MutableMapping, Sequence + from collections.abc import MutableMapping - from cudf_polars.containers import DataType + from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.experimental.dispatch import LowerIRTransformer from cudf_polars.experimental.parallel import PartitionInfo from cudf_polars.typing import Schema - from cudf_polars.utils.config import ShuffleMethod - - -# Supported shuffle methods -_SHUFFLE_METHODS = ("rapidsmpf", "tasks") - - -class ShuffleOptions(TypedDict): - """RapidsMPF shuffling options.""" - - on: Sequence[str] - column_names: Sequence[str] - dtypes: Sequence[DataType] - - -# Experimental rapidsmpf shuffler integration -class RMPFIntegration: # pragma: no cover - """cuDF-Polars protocol for rapidsmpf shuffler.""" - - @staticmethod - @nvtx_annotate_cudf_polars(message="RMPFIntegration.insert_partition") - def insert_partition( - df: DataFrame, - partition_id: int, # Not currently used - partition_count: int, - shuffler: Any, - options: ShuffleOptions, - *other: Any, - ) -> None: - """Add cudf-polars DataFrame chunks to an RMP shuffler.""" - from rapidsmpf.integrations.cudf.partition import partition_and_pack - from rapidsmpf.integrations.single import get_worker_context - - context = get_worker_context() - - on = options["on"] - assert not other, f"Unexpected arguments: {other}" - columns_to_hash = tuple(df.column_names.index(val) for val in on) - packed_inputs = partition_and_pack( - df.table, - columns_to_hash=columns_to_hash, - num_partitions=partition_count, - br=context.br, - stream=DEFAULT_STREAM, - ) - - shuffler.insert_chunks(packed_inputs) - - @staticmethod - @nvtx_annotate_cudf_polars(message="RMPFIntegration.extract_partition") - def extract_partition( - partition_id: int, - shuffler: Any, - options: ShuffleOptions, - ) -> DataFrame: - """Extract a finished partition from the RMP shuffler.""" - from rapidsmpf.integrations.cudf.partition import ( - unpack_and_concat, - unspill_partitions, - ) - from rapidsmpf.integrations.single import get_worker_context - - context = get_worker_context() - - shuffler.wait() - column_names = options["column_names"] - dtypes = options["dtypes"] - return DataFrame.from_table( - unpack_and_concat( - unspill_partitions( - shuffler.extract(partition_id), - br=context.br, - allow_overbooking=True, - ), - br=context.br, - stream=DEFAULT_STREAM, - ), - column_names, - dtypes, - get_dask_cuda_stream(), - ) class Shuffle(IR): @@ -118,29 +28,27 @@ class Shuffle(IR): Notes ----- - Only hash-based partitioning is supported (for now). See - `ShuffleSorted` for sorting-based shuffling. + Only hash-based partitioning is supported (for now). """ - __slots__ = ("keys", "shuffle_method") - _non_child = ("schema", "keys", "shuffle_method") - _n_non_child_args = 3 + __slots__ = ("keys",) + _non_child = ( + "schema", + "keys", + ) + _n_non_child_args = 2 keys: tuple[NamedExpr, ...] """Keys to shuffle on.""" - shuffle_method: ShuffleMethod - """Shuffle method to use.""" def __init__( self, schema: Schema, keys: tuple[NamedExpr, ...], - shuffle_method: ShuffleMethod, df: IR, ): self.schema = schema self.keys = keys - self.shuffle_method = shuffle_method - self._non_child_args = (schema, keys, shuffle_method) + self._non_child_args = (schema, keys) self.children = (df,) # the type-ignore is for @@ -153,7 +61,6 @@ def do_evaluate( cls, schema: Schema, keys: tuple[NamedExpr, ...], - shuffle_method: ShuffleMethod, df: DataFrame, *, context: IRExecutionContext, @@ -163,120 +70,6 @@ def do_evaluate( return df -@nvtx_annotate_cudf_polars(message="Shuffle") -def _hash_partition_dataframe( - df: DataFrame, - partition_id: int, # Used only by sorted shuffling - partition_count: int, - options: MutableMapping[str, Any] | None, # No options required - on: tuple[NamedExpr, ...], -) -> dict[int, DataFrame]: - """ - Partition an input DataFrame for hash-based shuffling. - - Parameters - ---------- - df - DataFrame to partition. - partition_id - Partition index (unused for hash partitioning). - partition_count - Total number of output partitions. - options - Options (unused for hash partitioning). - on - Expressions used for the hash partitioning. - - Returns - ------- - A dictionary mapping between int partition indices and - DataFrame fragments. - """ - assert not options, f"Expected no options, got: {options}" - - if df.num_rows == 0: - # Fast path for empty DataFrame - return dict.fromkeys(range(partition_count), df) - - # Hash the specified keys to calculate the output - # partition for each row - partition_map = plc.binaryop.binary_operation( - plc.hashing.murmurhash3_x86_32( - DataFrame([expr.evaluate(df) for expr in on], stream=df.stream).table, - stream=df.stream, - ), - plc.Scalar.from_py( - partition_count, plc.DataType(plc.TypeId.UINT32), stream=df.stream - ), - plc.binaryop.BinaryOperator.PYMOD, - plc.types.DataType(plc.types.TypeId.UINT32), - stream=df.stream, - ) - - # Apply partitioning - t, offsets = plc.partitioning.partition( - df.table, - partition_map, - partition_count, - stream=df.stream, - ) - splits = offsets[1:-1] - - # Split and return the partitioned result - return { - i: DataFrame.from_table( - split, - df.column_names, - df.dtypes, - df.stream, - ) - for i, split in enumerate(plc.copying.split(t, splits, stream=df.stream)) - } - - -# When dropping Python 3.10, can use _simple_shuffle_graph[OPT_T](...) -OPT_T = TypeVar("OPT_T") - - -def _simple_shuffle_graph( - name_in: str, - name_out: str, - count_in: int, - count_out: int, - _partition_dataframe_func: Callable[ - Concatenate[DataFrame, int, int, OPT_T, ...], - MutableMapping[int, DataFrame], - ], - options: OPT_T, - *other: Any, - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - """Make a simple all-to-all shuffle graph.""" - split_name = f"split-{name_out}" - inter_name = f"inter-{name_out}" - - graph: MutableMapping[Any, Any] = {} - for part_out in range(count_out): - _concat_list = [] - for part_in in range(count_in): - graph[(split_name, part_in)] = ( - _partition_dataframe_func, - (name_in, part_in), - part_in, - count_out, - options, - *other, - ) - _concat_list.append((inter_name, part_out, part_in)) - graph[_concat_list[-1]] = ( - operator.getitem, - (split_name, part_in), - part_out, - ) - graph[(name_out, part_out)] = (partial(_concat, context=context), *_concat_list) - return graph - - @lower_ir_node.register(Shuffle) def _( ir: Shuffle, rec: LowerIRTransformer @@ -306,47 +99,3 @@ def _( partitioned_on=ir.keys, ) return new_node, pi - - -@generate_ir_tasks.register(Shuffle) -def _( - ir: Shuffle, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - # Extract "shuffle_method" configuration - shuffle_method = ir.shuffle_method - - # Try using rapidsmpf shuffler if we have "simple" shuffle - # keys, and the "shuffle_method" config is set to "rapidsmpf-single". - _keys: list[Col] - if shuffle_method == "rapidsmpf-single" and len( - _keys := [ne.value for ne in ir.keys if isinstance(ne.value, Col)] - ) == len(ir.keys): # pragma: no cover - from rapidsmpf.integrations.single import rapidsmpf_shuffle_graph - - shuffle_on = [k.name for k in _keys] - - return rapidsmpf_shuffle_graph( - get_key_name(ir.children[0]), - get_key_name(ir), - partition_info[ir.children[0]].count, - partition_info[ir].count, - RMPFIntegration, - { - "on": shuffle_on, - "column_names": list(ir.schema.keys()), - "dtypes": list(ir.schema.values()), - }, - ) - - # Simple task-based fall-back - return partial(_simple_shuffle_graph, context=context)( - get_key_name(ir.children[0]), - get_key_name(ir), - partition_info[ir.children[0]].count, - partition_info[ir].count, - _hash_partition_dataframe, - None, - ir.keys, - ) diff --git a/python/cudf_polars/cudf_polars/experimental/sort.py b/python/cudf_polars/cudf_polars/experimental/sort.py index 6800fb4ab749..fa610324c2dc 100644 --- a/python/cudf_polars/cudf_polars/experimental/sort.py +++ b/python/cudf_polars/cudf_polars/experimental/sort.py @@ -4,47 +4,30 @@ from __future__ import annotations -from functools import partial -from typing import TYPE_CHECKING, Any, TypedDict +from typing import TYPE_CHECKING import polars as pl import pylibcudf as plc -from rmm.pylibrmm.stream import DEFAULT_STREAM from cudf_polars.containers import Column, DataFrame, DataType from cudf_polars.dsl.expr import Col -from cudf_polars.dsl.ir import IR, Slice, Sort +from cudf_polars.dsl.ir import Slice, Sort from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.base import PartitionInfo, get_key_name -from cudf_polars.experimental.dispatch import ( - generate_ir_tasks, - lower_ir_node, -) -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import _simple_shuffle_graph +from cudf_polars.experimental.dispatch import lower_ir_node from cudf_polars.experimental.utils import ( - _concat, - _fallback_inform, _lower_ir_fallback, ) -from cudf_polars.utils.config import ShuffleMethod -from cudf_polars.utils.cuda_stream import ( - get_dask_cuda_stream, - get_joined_cuda_stream, - join_cuda_streams, -) if TYPE_CHECKING: from collections.abc import MutableMapping, Sequence from rmm.pylibrmm.stream import Stream - from cudf_polars.dsl.expr import NamedExpr - from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.dsl.ir import IR + from cudf_polars.experimental.base import PartitionInfo from cudf_polars.experimental.dispatch import LowerIRTransformer - from cudf_polars.typing import Schema def find_sort_splits( @@ -251,248 +234,6 @@ def _get_final_sort_boundaries( ) -def _sort_boundaries_graph( - name_in: str, - by: Sequence[str], - column_order: Sequence[plc.types.Order], - null_order: Sequence[plc.types.NullOrder], - count: int, - context: IRExecutionContext, -) -> tuple[str, MutableMapping[Any, Any]]: - """Graph to get the boundaries from all partitions.""" - local_boundaries_name = f"sort-boundaries_local-{name_in}" - concat_boundaries_name = f"sort-boundaries-concat-{name_in}" - global_boundaries_name = f"sort-boundaries-{name_in}" - graph: MutableMapping[Any, Any] = {} - - _concat_list = [] - for part_id in range(count): - graph[(local_boundaries_name, part_id)] = ( - _select_local_split_candidates, - (name_in, part_id), - by, - count, - part_id, - ) - _concat_list.append((local_boundaries_name, part_id)) - - graph[concat_boundaries_name] = (partial(_concat, context=context), *_concat_list) - graph[global_boundaries_name] = ( - _get_final_sort_boundaries, - concat_boundaries_name, - column_order, - null_order, - count, - ) - return global_boundaries_name, graph - - -class SortedShuffleOptions(TypedDict): - """RapidsMPF shuffling options.""" - - by: Sequence[str] - order: Sequence[plc.types.Order] - null_order: Sequence[plc.types.NullOrder] - column_names: Sequence[str] - column_dtypes: Sequence[DataType] - - -# Experimental rapidsmpf shuffler integration -class RMPFIntegrationSortedShuffle: # pragma: no cover - """cuDF-Polars protocol for rapidsmpf shuffler.""" - - @staticmethod - def insert_partition( - df: DataFrame, - partition_id: int, - partition_count: int, - shuffler: Any, - options: SortedShuffleOptions, - sort_boundaries: DataFrame, - ) -> None: - """Add cudf-polars DataFrame chunks to an RMP shuffler.""" - from rapidsmpf.integrations.cudf.partition import split_and_pack - from rapidsmpf.integrations.single import get_worker_context - - context = get_worker_context() - - by = options["by"] - data_streams = [ - df.stream, - sort_boundaries.stream, - ] - stream = get_joined_cuda_stream(get_dask_cuda_stream, upstreams=data_streams) - - splits = find_sort_splits( - df.select(by).table, - sort_boundaries.table, - partition_id, - options["order"], - options["null_order"], - stream=stream, - ) - packed_inputs = split_and_pack( - df.table, - splits=splits, - br=context.br, - stream=stream, - ) - # TODO: figure out handoff with rapidsmpf - # https://github.com/rapidsai/cudf/issues/20337 - shuffler.insert_chunks(packed_inputs) - - join_cuda_streams(downstreams=data_streams, upstreams=[stream]) - - @staticmethod - def extract_partition( - partition_id: int, - shuffler: Any, - options: SortedShuffleOptions, - ) -> DataFrame: - """Extract a finished partition from the RMP shuffler.""" - from rapidsmpf.integrations.cudf.partition import ( - unpack_and_concat, - unspill_partitions, - ) - from rapidsmpf.integrations.single import get_worker_context - - context = get_worker_context() - - shuffler.wait() - column_names = options["column_names"] - column_dtypes = options["column_dtypes"] - - stream = DEFAULT_STREAM - - # TODO: When sorting, this step should finalize with a merge (unless we - # require stability, as cudf merge is not stable). - # TODO: figure out handoff with rapidsmpf - # https://github.com/rapidsai/cudf/issues/20337 - return DataFrame.from_table( - unpack_and_concat( - unspill_partitions( - shuffler.extract(partition_id), - br=context.br, - allow_overbooking=True, - ), - br=context.br, - stream=stream, - ), - column_names, - column_dtypes, - stream=stream, - ) - - -def _sort_partition_dataframe( - df: DataFrame, - partition_id: int, # Not currently used - partition_count: int, - options: MutableMapping[str, Any], - sort_boundaries: DataFrame, -) -> MutableMapping[int, DataFrame]: - """ - Partition a sorted DataFrame for shuffling. - - Parameters - ---------- - df - The DataFrame to partition. - partition_id - The partition id of the current partition. - partition_count - The total number of partitions. - options - The sort options ``(by, order, null_order)``. - sort_boundaries - The global sort boundary candidates used to decide where to split. - """ - if df.num_rows == 0: # pragma: no cover - # Fast path for empty DataFrame - return dict.fromkeys(range(partition_count), df) - - stream = get_joined_cuda_stream( - get_dask_cuda_stream, upstreams=(df.stream, sort_boundaries.stream) - ) - - splits = find_sort_splits( - df.select(options["by"]).table, - sort_boundaries.table, - partition_id, - options["order"], - options["null_order"], - stream=stream, - ) - - # Split and return the partitioned result - return { - i: DataFrame.from_table( - split, - df.column_names, - df.dtypes, - stream=df.stream, - ) - for i, split in enumerate(plc.copying.split(df.table, splits, stream=stream)) - } - - -class ShuffleSorted(IR): - """ - Shuffle already locally sorted multi-partition data. - - Shuffling is performed by extracting sort boundary candidates from all partitions, - sharing them all-to-all and then exchanging data accordingly. - The sorting information is required to be passed in identically to the already - performed local sort and as of now the final result needs to be sorted again to - merge the partitions. - """ - - __slots__ = ("by", "null_order", "order", "shuffle_method") - _non_child = ("schema", "by", "order", "null_order", "shuffle_method") - _n_non_child_args = 5 - by: tuple[NamedExpr, ...] - """Keys by which the data was sorted.""" - order: tuple[plc.types.Order, ...] - """Sort order if sorted.""" - null_order: tuple[plc.types.NullOrder, ...] - """Null precedence if sorted.""" - shuffle_method: ShuffleMethod - """Shuffle method to use.""" - - def __init__( - self, - schema: Schema, - by: tuple[NamedExpr, ...], - order: tuple[plc.types.Order, ...], - null_order: tuple[plc.types.NullOrder, ...], - shuffle_method: ShuffleMethod, - df: IR, - ): - self.schema = schema - self.by = by - self.order = order - self.null_order = null_order - self.shuffle_method = shuffle_method - self._non_child_args = (schema, by, order, null_order, shuffle_method) - self.children = (df,) - - @classmethod - def do_evaluate( - cls, - schema: Schema, - by: tuple[NamedExpr, ...], - order: tuple[plc.types.Order, ...], - null_order: tuple[plc.types.NullOrder, ...], - shuffle_method: ShuffleMethod, - df: DataFrame, - *, - context: IRExecutionContext, - ) -> DataFrame: # pragma: no cover - """Evaluate and return a dataframe.""" - # Single-partition ShuffleSorted evaluation is a no-op - return df - - def _has_simple_zlice(zlice: tuple[int, int | None] | None) -> bool: """Check if a zlice is a simple top-k/bottom-k operation.""" if zlice is None: @@ -517,26 +258,7 @@ def _( msg="sort currently only supports column names as `by` keys.", ) - config_options = rec.state["config_options"] - executor = config_options.executor - runtime = executor.runtime - - # Special handling for slicing - # (May be a top- or bottom-k operation) - simple_zlice = _has_simple_zlice(ir.zlice) - if simple_zlice and runtime == "tasks": - from cudf_polars.experimental.parallel import _lower_ir_pwise - - new_node, partition_info = _lower_ir_pwise(ir, rec) - if partition_info[new_node].count > 1: - # Collapse down to single partition - inter = Repartition(new_node.schema, new_node) - partition_info[inter] = PartitionInfo(count=1) - # Sort reduced partition - new_node = ir.reconstruct([inter]) - partition_info[new_node] = PartitionInfo(count=1) - return new_node, partition_info - elif ir.zlice is not None and not simple_zlice: + if ir.zlice is not None and not _has_simple_zlice(ir.zlice): # Pull "complex" slices out of the Sort node altogether. return rec( Slice( @@ -557,112 +279,6 @@ def _( # Extract child partitioning child, partition_info = rec(ir.children[0]) - # The "rapidsmpf" runtime uses the sort_actor to handle everything else - if runtime == "rapidsmpf": - sort_node = ir.reconstruct([child]) - partition_info[sort_node] = partition_info[child] - return sort_node, partition_info - - # TODO: Remove everything below here when "tasks" is removed. - - # Avoid rapidsmpf shuffle with maintain_order=True (for now) - shuffle_method = ( - ShuffleMethod("tasks") if ir.stable else config_options.executor.shuffle_method - ) - if ( - shuffle_method != config_options.executor.shuffle_method - ): # pragma: no cover; Requires rapidsmpf - _fallback_inform( - f"shuffle_method={config_options.executor.shuffle_method} does not support maintain_order=True. " - f"Falling back to shuffle_method={shuffle_method}.", - config_options, - ) - - if partition_info[child].count == 1: - single_part_node = ir.reconstruct([child]) - partition_info[single_part_node] = partition_info[child] - return single_part_node, partition_info - - local_sort_node = ir.reconstruct([child]) - partition_info[local_sort_node] = partition_info[child] - - shuffle = ShuffleSorted( - ir.schema, - ir.by, - ir.order, - ir.null_order, - shuffle_method, - local_sort_node, - ) - partition_info[shuffle] = partition_info[child] - - # We sort again locally. - assert ir.zlice is None # zlice handling would be incorrect without adjustment - final_sort_node = ir.reconstruct([shuffle]) - partition_info[final_sort_node] = partition_info[shuffle] - - return final_sort_node, partition_info - - -@generate_ir_tasks.register(ShuffleSorted) -def _( - ir: ShuffleSorted, - partition_info: MutableMapping[IR, PartitionInfo], - context: IRExecutionContext, -) -> MutableMapping[Any, Any]: - by = [ne.value.name for ne in ir.by if isinstance(ne.value, Col)] - if len(by) != len(ir.by): # pragma: no cover - # We should not reach here as this is checked in the lower_ir_node - raise NotImplementedError("Sorting columns must be column names.") - - (child,) = ir.children - - sort_boundaries_name, graph = _sort_boundaries_graph( - get_key_name(child), - by, - ir.order, - ir.null_order, - partition_info[child].count, - context, - ) - - options = { - "by": by, - "order": ir.order, - "null_order": ir.null_order, - "column_names": list(ir.schema.keys()), - "column_dtypes": list(ir.schema.values()), - } - - # Try using rapidsmpf shuffler if we have "simple" shuffle - # keys, and the "shuffle_method" config is set to "rapidsmpf-single". - shuffle_method = ir.shuffle_method - if shuffle_method == "rapidsmpf-single": # pragma: no cover - from rapidsmpf.integrations.single import rapidsmpf_shuffle_graph - - graph.update( - rapidsmpf_shuffle_graph( - get_key_name(child), - get_key_name(ir), - partition_info[child].count, - partition_info[ir].count, - RMPFIntegrationSortedShuffle, - options, - sort_boundaries_name, - ) - ) - return graph - - # Simple task-based fall-back - graph.update( - partial(_simple_shuffle_graph, context=context)( - get_key_name(child), - get_key_name(ir), - partition_info[child].count, - partition_info[ir].count, - _sort_partition_dataframe, - options, - sort_boundaries_name, - ) - ) - return graph + sort_node = ir.reconstruct([child]) + partition_info[sort_node] = partition_info[child] + return sort_node, partition_info diff --git a/python/cudf_polars/cudf_polars/experimental/utils.py b/python/cudf_polars/cudf_polars/experimental/utils.py index 24ce606d41bb..848a4d447598 100644 --- a/python/cudf_polars/cudf_polars/experimental/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/utils.py @@ -10,7 +10,7 @@ from itertools import chain from typing import TYPE_CHECKING -from cudf_polars.dsl.expr import Col, Expr, GroupedWindow, UnaryFunction +from cudf_polars.dsl.expr import Col, GroupedWindow, UnaryFunction from cudf_polars.dsl.ir import Union from cudf_polars.dsl.traversal import traversal from cudf_polars.experimental.base import PartitionInfo @@ -49,11 +49,8 @@ def _fallback_inform( def _dynamic_planning_on(config_options: ConfigOptions[StreamingExecutor]) -> bool: - """Check if dynamic planning is enabled for rapidsmpf runtime.""" - return ( - config_options.executor.runtime == "rapidsmpf" - and config_options.executor.dynamic_planning is not None - ) + """Check if dynamic planning is enabled.""" + return config_options.executor.dynamic_planning is not None def _lower_ir_fallback( @@ -68,9 +65,6 @@ def _lower_ir_fallback( from cudf_polars.experimental.repartition import Repartition from cudf_polars.experimental.select import _inline_hstack_false - config_options = rec.state["config_options"] - rapidsmpf_engine = config_options.executor.runtime == "rapidsmpf" - # Make sure we avoid mixed-length columns in intermediate TableChunks. ir = _inline_hstack_false(ir) @@ -82,13 +76,10 @@ def _lower_ir_fallback( children = [] inform = False for c in lowered_children: - child = c - if multi_partitioned := partition_info[c].count > 1: + if partition_info[c].count > 1: inform = True - if multi_partitioned or rapidsmpf_engine: - # Fall-back logic - child = Repartition(child.schema, child) - partition_info[child] = PartitionInfo(count=1) + child = Repartition(c.schema, c) + partition_info[child] = PartitionInfo(count=1) children.append(child) if inform and msg: @@ -114,32 +105,6 @@ def _leaf_column_names(expr: Expr) -> tuple[str, ...]: return () -def _get_unique_fractions( - column_names: Sequence[str], - user_unique_fractions: dict[str, float], -) -> dict[str, float]: - """ - Return unique-fraction statistics subset. - - Parameters - ---------- - column_names - The column names to get unique-fractions for. - user_unique_fractions - The user-provided unique-fraction dictionary. - - Returns - ------- - unique_fractions - The final unique-fraction dictionary filtered to column_names. - """ - return { - c: max(min(f, 1.0), 0.00001) - for c, f in user_unique_fractions.items() - if c in column_names - } - - def _contains_over(exprs: Sequence[Expr]) -> bool: """Return True if any expression contains a window expression.""" return any(isinstance(e, GroupedWindow) for e in traversal(exprs)) diff --git a/python/cudf_polars/cudf_polars/testing/asserts.py b/python/cudf_polars/cudf_polars/testing/asserts.py index 9f0953cd4dff..5611f8c3e704 100644 --- a/python/cudf_polars/cudf_polars/testing/asserts.py +++ b/python/cudf_polars/cudf_polars/testing/asserts.py @@ -30,7 +30,6 @@ # Will be overriden by `conftest.py` with the value from the `--executor` # and `--cluster` command-line arguments DEFAULT_EXECUTOR = "in-memory" -DEFAULT_RUNTIME = "tasks" DEFAULT_CLUSTER = "single" @@ -200,7 +199,6 @@ def get_default_engine( executor = executor or DEFAULT_EXECUTOR if executor == "streaming": executor_options["cluster"] = DEFAULT_CLUSTER - executor_options["runtime"] = DEFAULT_RUNTIME return GPUEngine( raise_on_fail=True, @@ -290,7 +288,8 @@ def assert_collect_raises( if polars_except != (): raise AssertionError(f"CPU execution DID NOT RAISE {polars_except}") - engine = GPUEngine(raise_on_fail=True) + # TODO: https://github.com/rapidsai/cudf/issues/22346 + engine = GPUEngine(executor="in-memory", raise_on_fail=True) try: lazydf.collect(**final_cudf_collect_kwargs, engine=engine) # type: ignore[misc, call-overload] except cudf_except: diff --git a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py index 6fe2de4d1543..7cfb62c414ee 100644 --- a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py +++ b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py @@ -30,6 +30,7 @@ def pytest_addoption(parser: pytest.Parser) -> None: choices=("in-memory", "spmd"), help="Which GPU engine variant to inject globally.", ) + # TODO: We never run with --inject-gpu-engine-blocksize in ci/run_cudf_polars_polars_tests.sh. Remove? group.addoption( "--inject-gpu-engine-blocksize", action="store", @@ -134,6 +135,7 @@ def pytest_report_header(config: pytest.Config) -> str: return f"injected GPU engine: {cls.__module__}.{cls.__name__}" +# TODO: This is just Mapping[str, str]? EXPECTED_FAILURES: Mapping[str, str | tuple[str, bool]] = { "tests/unit/io/test_csv.py::test_read_csv_only_loads_selected_columns": "Memory usage won't be correct due to GPU", "tests/unit/io/test_delta.py::test_scan_delta_version": "Need to expose hive partitioning", @@ -305,7 +307,6 @@ def pytest_report_header(config: pytest.Config) -> str: # Generally skip for: # 1) Tests that are too slow with --inject-gpu-engine-blocksize=small due to many small partitions for large data -# 2) Tests that fail during cudf_polars execution and segfaults later due to https://github.com/rapidsai/cudf/issues/22138 STREAMING_ENGINE_TESTS_TO_SKIP: Mapping[str, str] = { "tests/unit/operations/aggregation/test_aggregations.py::test_boolean_aggs": "float difference in std/var in the unit of least precision", "tests/benchmark/test_group_by.py::test_groupby_h2oai_q1": "Too slow with --inject-gpu-engine-blocksize=small", diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index a6bbd73929b4..7b5fb5c940c3 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -55,9 +55,7 @@ "InMemoryExecutor", "ParquetOptions", "RayContext", - "Runtime", "SPMDContext", - "ShuffleMethod", "StreamingExecutor", "StreamingFallbackMode", ] @@ -112,15 +110,6 @@ def get_total_device_memory() -> int | None: return None -@functools.cache -def rapidsmpf_single_available() -> bool: # pragma: no cover - """Query whether rapidsmpf is available as a single-process shuffle method.""" - try: - return importlib.util.find_spec("rapidsmpf.integrations.single") is not None - except (ImportError, ValueError): - return False - - class StreamingFallbackMode(enum.StrEnum): """ How the streaming executor handles operations that don't support multiple partitions. @@ -138,20 +127,6 @@ class StreamingFallbackMode(enum.StrEnum): SILENT = "silent" -class Runtime(enum.StrEnum): - """ - The runtime to use for the streaming executor. - - * ``Runtime.TASKS`` : Use the task-based runtime. - This is the default runtime. - * ``Runtime.RAPIDSMPF`` : Use the coroutine-based streaming runtime (rapidsmpf). - This runtime is experimental. - """ - - TASKS = "tasks" - RAPIDSMPF = "rapidsmpf" - - class Cluster(enum.StrEnum): """ The cluster configuration for the streaming executor. @@ -172,27 +147,6 @@ class Cluster(enum.StrEnum): DASK = "dask" -class ShuffleMethod(enum.StrEnum): - """ - The method to use for shuffling data between workers with the streaming executor. - - * ``ShuffleMethod.TASKS`` : Use the task-based shuffler. - * ``ShuffleMethod.RAPIDSMPF`` : Use the rapidsmpf shuffler. - * ``ShuffleMethod._RAPIDSMPF_SINGLE`` : Use the single-process rapidsmpf shuffler. - - With :class:`cudf_polars.utils.config.StreamingExecutor`, the default of ``None`` - resolves to ``ShuffleMethod.TASKS``. - - The user should **not** specify ``ShuffleMethod._RAPIDSMPF_SINGLE`` directly. - A setting of ``ShuffleMethod.RAPIDSMPF`` will be converted to the single-process - shuffler automatically when using single-GPU execution. - """ - - TASKS = "tasks" - RAPIDSMPF = "rapidsmpf" - _RAPIDSMPF_SINGLE = "rapidsmpf-single" - - T = TypeVar("T") @@ -254,7 +208,7 @@ class ParquetOptions: will also be skipped if ``max_footer_samples`` is 0. use_rapidsmpf_native Whether to use the native rapidsmpf node for parquet reading. - This option is only used when the rapidsmpf runtime is enabled. + This option is only used by the streaming executor. Default is False. """ @@ -315,49 +269,32 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("use_rapidsmpf_native must be a bool") -def default_target_partition_size(cluster: str, runtime: str) -> int: +@functools.cache +def default_target_partition_size() -> int: """Return the default blocksize.""" if (device_size := get_total_device_memory()) is None: # pragma: no cover # System doesn't have proper "GPU memory". # Fall back to a conservative 1GB default. return 1_000_000_000 - if ( - cluster == "single" - and runtime == "tasks" - and _env_get_int("POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY", default=1) == 1 - ): - # We can use a larger blocksize when UVM is enabled - blocksize = int(device_size * 0.0625) - else: - # Otherwise, use a conservative default - blocksize = int(device_size * 0.025) + blocksize = int(device_size * 0.025) # Use lower and upper bounds of 1GB and 10GB return min(max(blocksize, 1_000_000_000), 10_000_000_000) -def default_broadcast_join_limit(cluster: str, runtime: str) -> int: +@functools.cache +def default_broadcast_join_limit() -> int: """Return the default broadcast join limit.""" if (device_size := get_total_device_memory()) is None: # pragma: no cover # System doesn't have proper "GPU memory". # We probably want to broadcast in most cases. return 32 - if runtime == "rapidsmpf": - # Target about 12.5% of the device memory when - # default_target_partition_size is used to set the - # target partition size (i.e. 5x the 2.5% default). - return min(5, int(max(1, (device_size * 0.125) // 1e9))) - elif _env_get_int("POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY", default=1) == 1: - # The "tasks" runtime always runs single-GPU; we can lean on UVM - # to support most broadcast joins. - return 32 - else: - # Extra-conservative default for the "tasks" runtime without UVM. - # We cannot spill outside a rapidsmpf shuffle within this runtime, - # so shuffling is usually preferred. - return 2 + # Target about 12.5% of the device memory when + # default_target_partition_size is used to set the + # target partition size (i.e. 5x the 2.5% default). + return min(5, int(max(1, (device_size * 0.125) // 1e9))) @dataclasses.dataclass(frozen=True) @@ -599,17 +536,14 @@ class StreamingExecutor: Parameters ---------- - runtime - The runtime to use for the streaming executor. - ``Runtime.TASKS`` by default. cluster The cluster configuration for the streaming executor. ``Cluster.SINGLE`` by default. * ``Cluster.SINGLE``: Single-GPU execution - * ``Cluster.SPMD``: Multi-GPU SPMD execution (rapidsmpf runtime) - * ``Cluster.RAY``: Multi-GPU Ray execution (rapidsmpf runtime) - * ``Cluster.DASK``: Multi-GPU Dask execution (rapidsmpf runtime) + * ``Cluster.SPMD``: Multi-GPU SPMD execution + * ``Cluster.RAY``: Multi-GPU Ray execution + * ``Cluster.DASK``: Multi-GPU Dask execution fallback_mode How to handle errors when the GPU engine fails to execute a query. @@ -621,13 +555,6 @@ class StreamingExecutor: The maximum number of rows to process per partition. 1_000_000 by default. When the number of rows exceeds this value, the query will be split into multiple partitions and executed in parallel. - unique_fraction - A dictionary mapping column names to floats between 0 and 1 (inclusive - on the right). - - Each factor estimates the fractional number of unique values in the - column. By default, ``1.0`` is used for any column not included in - ``unique_fraction``. target_partition_size Target partition size, in bytes, for IO tasks. This configuration currently controls how large parquet files are split into multiple partitions. @@ -639,11 +566,8 @@ class StreamingExecutor: - keyword argument to ``polars.GPUEngine`` - the ``CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE`` environment variable - By default, cudf-polars uses a target partition size that's a fraction - of the device memory, where the fraction depends on the cluster and runtime: - - - rapidsmpf runtime: 1/40th of the device memory - - single cluster and tasks runtime: 1/16th of the device memory + By default, cudf-polars uses a target partition size of 1/40th of the + device memory. The pynvml library is used to query the total device memory on the first visible GPU. If the device size is not available, the default target @@ -651,26 +575,14 @@ class StreamingExecutor: NOTE: If this configuration is changed manually, it is recommended to set `broadcast_join_limit` manually as well. - groupby_n_ary - The factor by which the number of partitions is decreased when performing - a groupby on a partitioned column. For example, if a column has 64 partitions, - it will first be reduced to ``ceil(64 / 32) = 2`` partitions. - - This is useful when the absolute number of partitions is large. broadcast_join_limit The maximum number of partitions to allow for the smaller table in a broadcast join. For example, if the target partition size is 1GB and the broadcast join limit is 5, then the smaller table will be broadcasted - if it is smaller than 5GB (within the "rapidsmpf" runtime) or contains - fewer than 5 partitions (within the "tasks" runtime). The default depends - on the cluster and runtime. - shuffle_method - The method to use for shuffling data between workers. Defaults to - 'tasks' for the single-GPU cluster. + if it is smaller than 5GB. client_device_threshold - Threshold for spilling data from device memory in rapidsmpf. + Threshold for spilling data from device memory. Default is 50% of device memory on the client process. - This argument is only used by the "rapidsmpf" runtime. sink_to_directory Whether multi-partition sink operations write to a directory rather than a single file. For the spmd, ray, and dask clusters this is @@ -680,7 +592,7 @@ class StreamingExecutor: Options controlling dynamic shuffle planning. See :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. max_io_threads - Maximum number of IO threads for the rapidsmpf runtime. Default is 4. + Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. spill_to_pinned_memory Whether RapidsMPF should spill to pinned host memory when available, @@ -688,8 +600,8 @@ class StreamingExecutor: bandwidth and lower latency for device to host transfers compared to regular pageable host memory. num_py_executors - Maximum number of workers for the Python ThreadPoolExecutor used by - the rapidsmpf runtime. Default is 8. + Maximum number of workers for the Python ThreadPoolExecutor. + Default is 8. Notes ----- @@ -700,13 +612,6 @@ class StreamingExecutor: _env_prefix = "CUDF_POLARS__EXECUTOR" name: Literal["streaming"] = dataclasses.field(default="streaming", init=False) - runtime: Runtime = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__RUNTIME", - Runtime.__call__, - default=Runtime.TASKS, - ) - ) cluster: Cluster | None = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__CLUSTER", @@ -726,33 +631,16 @@ class StreamingExecutor: f"{_env_prefix}__MAX_ROWS_PER_PARTITION", int, default=1_000_000 ) ) - unique_fraction: dict[str, float] = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__UNIQUE_FRACTION", json.loads, default={} - ) - ) target_partition_size: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__TARGET_PARTITION_SIZE", int, default=0 ) ) - groupby_n_ary: int = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__GROUPBY_N_ARY", int, default=32 - ) - ) broadcast_join_limit: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__BROADCAST_JOIN_LIMIT", int, default=0 ) ) - shuffle_method: ShuffleMethod = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__SHUFFLE_METHOD", - ShuffleMethod.__call__, - default=ShuffleMethod.TASKS, - ) - ) client_device_threshold: float = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__CLIENT_DEVICE_THRESHOLD", float, default=0.5 @@ -786,33 +674,10 @@ class StreamingExecutor: dask_context: DaskContext | None = None def __post_init__(self) -> None: # noqa: D105 - # Check for rapidsmpf runtime - if self.runtime == "rapidsmpf": # pragma: no cover; requires rapidsmpf runtime - if not rapidsmpf_single_available(): - raise ValueError("The rapidsmpf streaming engine requires rapidsmpf.") - object.__setattr__(self, "shuffle_method", "rapidsmpf") - if self.cluster is None: object.__setattr__(self, "cluster", Cluster.SINGLE) assert self.cluster is not None, "Expected cluster to be set." - # Handle shuffle_method defaults for streaming executor - if self.shuffle_method is None: - # Use task-based shuffle by default. - # TODO: Evaluate single-process shuffle by default. - object.__setattr__(self, "shuffle_method", "tasks") - elif self.shuffle_method == "rapidsmpf-single": - # The user should NOT specify "rapidsmpf-single" directly. - raise ValueError("rapidsmpf-single is not a supported shuffle method.") - elif self.shuffle_method == "rapidsmpf": - if self.cluster == "single" and not rapidsmpf_single_available(): - raise ValueError( - "rapidsmpf shuffle method requested, but rapidsmpf is not installed." - ) - # Select "rapidsmpf-single" for single-GPU - if self.cluster == "single": - object.__setattr__(self, "shuffle_method", "rapidsmpf-single") - # frozen dataclass, so use object.__setattr__ object.__setattr__( self, "fallback_mode", StreamingFallbackMode(self.fallback_mode) @@ -821,16 +686,15 @@ def __post_init__(self) -> None: # noqa: D105 object.__setattr__( self, "target_partition_size", - default_target_partition_size(self.cluster, self.runtime), + default_target_partition_size(), ) if self.broadcast_join_limit == 0: object.__setattr__( self, "broadcast_join_limit", - default_broadcast_join_limit(self.cluster, self.runtime), + default_broadcast_join_limit(), ) object.__setattr__(self, "cluster", Cluster(self.cluster)) - object.__setattr__(self, "shuffle_method", ShuffleMethod(self.shuffle_method)) # Handle dynamic_planning. # Can be None, dict, or DynamicPlanningOptions @@ -853,12 +717,8 @@ def __post_init__(self) -> None: # noqa: D105 # Type / value check everything else if not isinstance(self.max_rows_per_partition, int): raise TypeError("max_rows_per_partition must be an int") - if not isinstance(self.unique_fraction, dict): - raise TypeError("unique_fraction must be a dict of column name to float") if not isinstance(self.target_partition_size, int): raise TypeError("target_partition_size must be an int") - if not isinstance(self.groupby_n_ary, int): - raise TypeError("groupby_n_ary must be an int") if not isinstance(self.broadcast_join_limit, int): raise TypeError("broadcast_join_limit must be an int") if not isinstance(self.sink_to_directory, bool): @@ -873,10 +733,9 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("num_py_executors must be an int") def __hash__(self) -> int: # noqa: D105 - # cardinality factory, a dict, isn't natively hashable. We'll dump it + # dynamic_planning factory, a dataclass, isn't natively hashable. We'll dump it # to json and hash that. d = dataclasses.asdict(self) - d["unique_fraction"] = json.dumps(d["unique_fraction"]) d["dynamic_planning"] = json.dumps(d["dynamic_planning"]) return hash(tuple(sorted(d.items()))) @@ -1059,19 +918,6 @@ def from_polars_engine( executor = InMemoryExecutor(**user_executor_options) case "streaming": user_executor_options = user_executor_options.copy() - # Handle the interaction between the default shuffle method, the - # cluster, and whether rapidsmpf is available. - env_shuffle_method = os.environ.get( - "CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD", None - ) - if env_shuffle_method is not None: - shuffle_method_default = ShuffleMethod(env_shuffle_method) - else: - shuffle_method_default = None - - user_executor_options.setdefault( - "shuffle_method", shuffle_method_default - ) # Handle dynamic_planning: check user config, then env var user_dynamic_planning = user_executor_options.get( @@ -1097,7 +943,7 @@ def from_polars_engine( } # Handle "cuda-stream-policy". - # The default will depend on the runtime and executor. + # The default will depend on the executor. user_cuda_stream_policy = engine.config.get( "cuda_stream_policy", None ) or os.environ.get("CUDF_POLARS__CUDA_STREAM_POLICY", None) @@ -1105,24 +951,18 @@ def from_polars_engine( cuda_stream_policy: CUDAStreamPoolConfig | None if user_cuda_stream_policy is None: - if ( - executor.name == "streaming" and executor.runtime == Runtime.RAPIDSMPF - ): # pragma: no cover; requires rapidsmpf runtime - # the rapidsmpf runtime defaults to using a stream pool + if executor.name == "streaming": cuda_stream_policy = CUDAStreamPoolConfig() else: - # everything else defaults to the default stream cuda_stream_policy = None else: cuda_stream_policy = _convert_cuda_stream_policy(user_cuda_stream_policy) - # Pool policy is only supported by the rapidsmpf runtime. if isinstance(cuda_stream_policy, CUDAStreamPoolConfig) and ( - (executor.name != "streaming") - or (executor.name == "streaming" and executor.runtime != Runtime.RAPIDSMPF) + executor.name != "streaming" ): raise ValueError( - "A stream pool is only supported by the rapidsmpf runtime." + "A stream pool is only supported by the streaming executor." ) kwargs["cuda_stream_policy"] = cuda_stream_policy diff --git a/python/cudf_polars/cudf_polars/utils/cuda_stream.py b/python/cudf_polars/cudf_polars/utils/cuda_stream.py index c0708d3bea8a..22022ee3401c 100644 --- a/python/cudf_polars/cudf_polars/utils/cuda_stream.py +++ b/python/cudf_polars/cudf_polars/utils/cuda_stream.py @@ -17,11 +17,6 @@ from rmm.pylibrmm.stream import Stream -def get_dask_cuda_stream() -> Stream: - """Get the default CUDA stream for Dask.""" - return DEFAULT_STREAM - - def get_cuda_stream() -> Stream: """Get the default CUDA stream for the current thread.""" return DEFAULT_STREAM diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index d48793f05417..47633e42364e 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "packaging", "polars>=1.30,<1.39", "pylibcudf==26.6.*,>=0.0.0a0", + "rapidsmpf==26.6.*,>=0.0.0a0", "typing_extensions>=4.0.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ @@ -47,7 +48,6 @@ test = [ "pytest-cov", "pytest-httpserver", "pytest-xdist", - "rapidsmpf==26.6.*,>=0.0.0a0", "rich", "structlog", "zstandard", diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 7f00684638f0..b3d83b36d36b 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import importlib.util from typing import TYPE_CHECKING import pytest @@ -54,13 +53,6 @@ def clear_memory_resource_cache(): @pytest.fixture(autouse=True) def _skip_unless_spmd(request: pytest.FixtureRequest) -> None: """Skip tests in SPMD multi-rank mode unless marked with ``pytest.mark.spmd``.""" - # Do not use `pytest.importorskip` here: this fixture is autouse, so an - # import-based skip would skip every test in the suite on environments - # without rapidsmpf (e.g. the coverage CI job), masking real coverage. - # We only want to gate the nranks>1 check on rapidsmpf being available. - if importlib.util.find_spec("rapidsmpf") is None: - return - from rapidsmpf.bootstrap import get_nranks, is_running_with_rrun if ( @@ -79,7 +71,6 @@ def streaming_engines() -> Generator[StreamingEngines, None, None]: name to a single shared engine instance, which is reused across the entire test session. """ - pytest.importorskip("rapidsmpf") from rapidsmpf import bootstrap from rapidsmpf.communicator.single import new_communicator as single_communicator from rapidsmpf.config import Options, get_environment_variables @@ -228,7 +219,8 @@ def engine_raise_on_fail() -> pl.GPUEngine: from ``.collect()``. Uses the in-memory executor so errors are not wrapped by a streaming task group. """ - return pl.GPUEngine(raise_on_fail=True) + # TODO: We should be testing will all supported engine variants + return pl.GPUEngine(executor="in-memory", raise_on_fail=True) def pytest_addoption(parser): @@ -240,14 +232,6 @@ def pytest_addoption(parser): help="Executor to use for GPUEngine.", ) - parser.addoption( - "--runtime", - action="store", - default="tasks", - choices=("tasks", "rapidsmpf"), - help="Runtime to use for the 'streaming' executor.", - ) - parser.addoption( "--cluster", action="store", @@ -278,17 +262,7 @@ def pytest_configure(config): # apply globally rather than per-module. config.addinivalue_line("filterwarnings", "ignore::ResourceWarning") - if config.getoption("--runtime") == "rapidsmpf": - if config.getoption("--executor") == "in-memory": - raise pytest.UsageError("Rapidsmpf runtime requires --executor='streaming'") - - if importlib.util.find_spec("rapidsmpf") is None: - raise pytest.UsageError( - "Rapidsmpf runtime requires the 'rapidsmpf' package" - ) - cudf_polars.testing.asserts.DEFAULT_EXECUTOR = config.getoption("--executor") - cudf_polars.testing.asserts.DEFAULT_RUNTIME = config.getoption("--runtime") cudf_polars.testing.asserts.DEFAULT_CLUSTER = config.getoption("--cluster") diff --git a/python/cudf_polars/tests/experimental/test_dask.py b/python/cudf_polars/tests/experimental/test_dask.py index 5ccdde864ef6..93ef4318490e 100644 --- a/python/cudf_polars/tests/experimental/test_dask.py +++ b/python/cudf_polars/tests/experimental/test_dask.py @@ -64,7 +64,6 @@ def test_yields_engine(engine: DaskEngine) -> None: def test_executor_options_forwarded(engine: DaskEngine) -> None: """Reserved executor_options keys are injected into the engine config.""" opts = engine.config["executor_options"] - assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "dask" assert isinstance(opts["dask_context"], DaskContext) @@ -196,7 +195,6 @@ def test_reset_updates_executor_options(reset_engine: DaskEngine) -> None: opts = reset_engine.config["executor_options"] assert opts["max_rows_per_partition"] == 42 # Reserved keys are still injected by ``_reset``. - assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "dask" assert isinstance(opts["dask_context"], DaskContext) diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index fecd4ba4d03d..7f19e3187783 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -540,8 +540,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None): engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) dag = serialize_query(q, engine) - # walk Union -> Scan - node = dag.nodes[dag.nodes[dag.roots[0]].children[0]] + node = dag.nodes[dag.roots[0]] assert node.type == "Scan" assert node.properties == expected_properties @@ -673,7 +672,6 @@ def test_dynamic_planning_adds_repartition(df, op): executor="streaming", raise_on_fail=True, executor_options={ - "runtime": "rapidsmpf", "dynamic_planning": {}, "max_rows_per_partition": 1_000_000, }, diff --git a/python/cudf_polars/tests/experimental/test_groupby.py b/python/cudf_polars/tests/experimental/test_groupby.py index 8d6ac5927e92..03d87fe23e93 100644 --- a/python/cudf_polars/tests/experimental/test_groupby.py +++ b/python/cudf_polars/tests/experimental/test_groupby.py @@ -270,10 +270,7 @@ def test_groupby_literal_key(df, streaming_engine): @pytest.mark.parametrize("keys", [("y",), ("y", "z")]) def test_groupby_agg_config_options(df, op, keys, streaming_engine_factory): streaming_engine = streaming_engine_factory( - StreamingOptions( - max_rows_per_partition=4, - unique_fraction={"z": 0.5}, - ), + StreamingOptions(max_rows_per_partition=4), ) agg = getattr(pl.col("x"), op)() if op in ("sum", "mean"): diff --git a/python/cudf_polars/tests/experimental/test_hstack.py b/python/cudf_polars/tests/experimental/test_hstack.py index 9bbb4b7aa33a..0c21678f7e25 100644 --- a/python/cudf_polars/tests/experimental/test_hstack.py +++ b/python/cudf_polars/tests/experimental/test_hstack.py @@ -20,7 +20,6 @@ from cudf_polars.experimental.statistics import collect_statistics from cudf_polars.testing.asserts import ( DEFAULT_CLUSTER, - DEFAULT_RUNTIME, assert_gpu_result_equal, ) from cudf_polars.utils.config import ConfigOptions @@ -34,7 +33,6 @@ def engine(): executor_options={ "max_rows_per_partition": 3, "cluster": DEFAULT_CLUSTER, - "runtime": DEFAULT_RUNTIME, }, ) diff --git a/python/cudf_polars/tests/experimental/test_options.py b/python/cudf_polars/tests/experimental/test_options.py index 291cbda7589c..eb64cb97ed8c 100644 --- a/python/cudf_polars/tests/experimental/test_options.py +++ b/python/cudf_polars/tests/experimental/test_options.py @@ -66,11 +66,6 @@ def test_executor_options_includes_set_fields() -> None: assert "log" not in result -def test_executor_options_unique_fraction() -> None: - result = StreamingOptions(unique_fraction={"col_a": 0.5}).to_executor_options() - assert result["unique_fraction"] == {"col_a": 0.5} - - def test_executor_options_num_py_executors() -> None: result = StreamingOptions(num_py_executors=4).to_executor_options() assert result["num_py_executors"] == 4 diff --git a/python/cudf_polars/tests/experimental/test_parallel.py b/python/cudf_polars/tests/experimental/test_parallel.py index 42365a113e26..67fc372e2e44 100644 --- a/python/cudf_polars/tests/experimental/test_parallel.py +++ b/python/cudf_polars/tests/experimental/test_parallel.py @@ -12,13 +12,9 @@ from polars.testing import assert_frame_equal from cudf_polars import Translator -from cudf_polars.dsl.expressions.base import Col, NamedExpr from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.parallel import lower_ir_graph from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal -from cudf_polars.utils.config import ConfigOptions @pytest.mark.parametrize("column", ["a", "b"]) @@ -134,40 +130,3 @@ def test_pickle_conditional_join_args(): ir = Translator(q._ldf.visit(), GPUEngine()).translate_ir() for node in traversal([ir]): pickle.loads(pickle.dumps(node._non_child_args)) - - -def test_preserve_partitioning(streaming_engine_factory): - streaming_engine = streaming_engine_factory( - StreamingOptions( - max_rows_per_partition=2, - broadcast_join_limit=2, - unique_fraction={"a": 1.0}, - ), - ) - left = pl.LazyFrame({"a": [1, 2, 3, 4] * 5, "b": range(20)}) - right = pl.LazyFrame({"a": [3, 4, 5, 6, 7] * 4, "c": range(20)}) - q = ( - left.join(right, on="a") - .filter(pl.col("a") == 2) - .group_by(pl.col("a")) - .mean() - .select(pl.col("a"), pl.col("c")) - ) - _engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - executor_options={ - "max_rows_per_partition": 2, - "broadcast_join_limit": 2, - "unique_fraction": {"a": 1.0}, - }, - ) - config_options = ConfigOptions.from_polars_engine(_engine) - ir = Translator(q._ldf.visit(), _engine).translate_ir() - ir, partition_info = lower_ir_graph( - ir, config_options, collect_statistics(ir, config_options) - ) - expect_dtype = ir.schema["a"] - expect_expr = (NamedExpr("a", Col(expect_dtype, "a")),) - assert partition_info[ir].partitioned_on == expect_expr - assert_gpu_result_equal(q, engine=streaming_engine) diff --git a/python/cudf_polars/tests/experimental/test_ray.py b/python/cudf_polars/tests/experimental/test_ray.py index ded4903c5940..f62c3e3b8316 100644 --- a/python/cudf_polars/tests/experimental/test_ray.py +++ b/python/cudf_polars/tests/experimental/test_ray.py @@ -53,7 +53,7 @@ def engine() -> Iterator[RayEngine]: def test_reserved_executor_keys() -> None: """executor_options rejects reserved keys.""" - for key in ("runtime", "cluster", "spmd_context", "ray_context"): + for key in ("cluster", "spmd_context", "ray_context"): with pytest.raises(TypeError, match="reserved"): RayEngine(executor_options={key: "anything"}) @@ -109,7 +109,6 @@ def test_executor_options_forwarded( ) -> None: """Reserved executor_options keys are injected into the engine config.""" opts = engine.config["executor_options"] - assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "ray" assert isinstance(opts["ray_context"], RayContext) assert engine.rank_actors == opts["ray_context"].rank_actors @@ -258,7 +257,6 @@ def test_reset_updates_executor_options(reset_engine: RayEngine) -> None: opts = reset_engine.config["executor_options"] assert opts["max_rows_per_partition"] == 42 # Reserved keys are still injected by ``_reset``. - assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "ray" assert isinstance(opts["ray_context"], RayContext) assert opts["ray_context"].rank_actors == reset_engine.rank_actors diff --git a/python/cudf_polars/tests/experimental/test_sort.py b/python/cudf_polars/tests/experimental/test_sort.py index 194686acf6b0..f0abf5caadee 100644 --- a/python/cudf_polars/tests/experimental/test_sort.py +++ b/python/cudf_polars/tests/experimental/test_sort.py @@ -9,7 +9,6 @@ from cudf_polars.testing.asserts import ( DEFAULT_CLUSTER, - DEFAULT_RUNTIME, assert_gpu_result_equal, ) @@ -22,7 +21,6 @@ def engine(): executor_options={ "max_rows_per_partition": 3, "cluster": DEFAULT_CLUSTER, - "runtime": DEFAULT_RUNTIME, "fallback_mode": "raise", }, ) @@ -36,7 +34,6 @@ def engine_large(): executor_options={ "max_rows_per_partition": 2_100, "cluster": DEFAULT_CLUSTER, - "runtime": DEFAULT_RUNTIME, "fallback_mode": "raise", }, ) @@ -139,7 +136,6 @@ def test_sort_after_sparse_join(): executor="streaming", executor_options={ "cluster": DEFAULT_CLUSTER, - "runtime": DEFAULT_RUNTIME, "max_rows_per_partition": 4, }, ) diff --git a/python/cudf_polars/tests/experimental/test_spmd.py b/python/cudf_polars/tests/experimental/test_spmd.py index 9fef0e003504..96ec5eab9320 100644 --- a/python/cudf_polars/tests/experimental/test_spmd.py +++ b/python/cudf_polars/tests/experimental/test_spmd.py @@ -66,7 +66,7 @@ def test_single_communicator_outside_rrun() -> None: def test_reserved_keys() -> None: """executor_options rejects reserved keys.""" - for key in ("runtime", "cluster", "spmd_context"): + for key in ("cluster", "spmd_context"): with ( pytest.raises(TypeError, match="reserved"), SPMDEngine(executor_options={key: "anything"}), @@ -320,7 +320,6 @@ def test_reset_updates_executor_options(comm: Communicator) -> None: opts = engine.config["executor_options"] assert opts["max_rows_per_partition"] == 42 # Reserved keys are still injected by ``_reset``. - assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "spmd" assert isinstance(opts["spmd_context"], SPMDContext) diff --git a/python/cudf_polars/tests/experimental/test_unique.py b/python/cudf_polars/tests/experimental/test_unique.py index 49d2b5803008..6bb30624cb66 100644 --- a/python/cudf_polars/tests/experimental/test_unique.py +++ b/python/cudf_polars/tests/experimental/test_unique.py @@ -34,12 +34,9 @@ def df(): @pytest.mark.parametrize("subset", [None, ("y",), ("y", "z")]) @pytest.mark.parametrize("keep", ["first", "last", "any", "none"]) @pytest.mark.parametrize("maintain_order", [True, False]) -@pytest.mark.parametrize("cardinality", [{}, {"y": 0.7}]) -def test_unique( - df, streaming_engine_factory, keep, subset, maintain_order, cardinality -): +def test_unique(df, streaming_engine_factory, keep, subset, maintain_order): engine = streaming_engine_factory( - StreamingOptions(unique_fraction=cardinality, fallback_mode="warn"), + StreamingOptions(fallback_mode="warn"), ) q = df.unique(subset=subset, keep=keep, maintain_order=maintain_order) check_row_order = maintain_order @@ -50,40 +47,16 @@ def test_unique( assert_gpu_result_equal(q, engine=engine, check_row_order=check_row_order) -def test_unique_fallback(df, streaming_engine_factory): - engine = streaming_engine_factory( - StreamingOptions( - unique_fraction={"y": 1.0}, - fallback_mode="raise", - dynamic_planning=None, - ), - ) - q = df.unique(keep="first", maintain_order=True) - with pytest.raises( - NotImplementedError, - match="Unsupported unique options", - ): - assert_gpu_result_equal(q, engine=engine) - - @pytest.mark.parametrize("maintain_order", [True, False]) -@pytest.mark.parametrize("cardinality", [{}, {"y": 0.5}]) -def test_unique_select(df, streaming_engine_factory, maintain_order, cardinality): +def test_unique_select(df, streaming_engine_factory, maintain_order): engine = streaming_engine_factory( StreamingOptions( max_rows_per_partition=4, - unique_fraction=cardinality, fallback_mode="warn", ), ) q = df.select(pl.col("y").unique(maintain_order=maintain_order)) - if cardinality == {"y": 0.5} and maintain_order: - with pytest.warns( - UserWarning, match="Unsupported unique options for multiple partitions." - ): - assert_gpu_result_equal(q, engine=engine, check_row_order=False) - else: - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize("keep", ["first", "last", "any"]) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 3cd66bc527d6..6004c5eef40c 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -4,7 +4,7 @@ from __future__ import annotations import sys -from typing import Any, cast +from typing import cast import pytest @@ -35,20 +35,7 @@ StreamingExecutor, _default_cuda_stream_policy, ) -from cudf_polars.utils.cuda_stream import ( - get_cuda_stream, - get_dask_cuda_stream, -) - - -@pytest.fixture(params=[False, True], ids=["norapidsmpf.single", "rapidsmpf.single"]) -def rapidsmpf_single_available(request, monkeypatch): - monkeypatch.setattr( - cudf_polars.utils.config, - "rapidsmpf_single_available", - lambda: request.param, - ) - return request.param +from cudf_polars.utils.cuda_stream import get_cuda_stream def test_polars_verbose_warns(monkeypatch): @@ -232,47 +219,6 @@ def test_parquet_options_from_none() -> None: assert config.parquet_options.chunked is True -def test_validate_streaming_executor_shuffle_method( - *, rapidsmpf_single_available: bool -) -> None: - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"shuffle_method": "tasks"}, - ) - ) - assert config.executor.name == "streaming" - assert config.executor.shuffle_method == "tasks" - - # rapidsmpf with single cluster - engine = pl.GPUEngine( - executor="streaming", - executor_options={"shuffle_method": "rapidsmpf", "cluster": "single"}, - ) - - if rapidsmpf_single_available: - config = ConfigOptions.from_polars_engine(engine) - assert config.executor.name == "streaming" - assert config.executor.shuffle_method == "rapidsmpf-single" - else: - with pytest.raises(ValueError, match="rapidsmpf is not installed"): - ConfigOptions.from_polars_engine(engine) - - -def test_join_rapidsmpf_single_private_config() -> None: - # The user may not specify "rapidsmpf-single" directly - engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - executor_options={ - "shuffle_method": "rapidsmpf-single", - "runtime": "tasks", - }, - ) - with pytest.raises(ValueError, match="not a supported shuffle method"): - ConfigOptions.from_polars_engine(engine) - - @pytest.mark.parametrize("executor", ["in-memory", "streaming"]) def test_hashable(executor: str) -> None: config = ConfigOptions.from_polars_engine( @@ -319,31 +265,11 @@ def test_validate_cluster() -> None: ) -def test_validate_shuffle_method_defaults() -> None: - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - ) - ) - assert config.executor.name == "streaming" - assert config.executor.shuffle_method == "tasks" # Default for single cluster - - with pytest.raises(ValueError, match="'foo' is not a valid ShuffleMethod"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"shuffle_method": "foo"}, - ) - ) - - @pytest.mark.parametrize( "option", [ "max_rows_per_partition", - "unique_fraction", "target_partition_size", - "groupby_n_ary", "broadcast_join_limit", "sink_to_directory", "client_device_threshold", @@ -409,11 +335,8 @@ def test_config_option_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__EXECUTOR__CLUSTER", "single") m.setenv("CUDF_POLARS__EXECUTOR__FALLBACK_MODE", "silent") m.setenv("CUDF_POLARS__EXECUTOR__MAX_ROWS_PER_PARTITION", "42") - m.setenv("CUDF_POLARS__EXECUTOR__UNIQUE_FRACTION", '{"a": 0.5}') m.setenv("CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE", "100") - m.setenv("CUDF_POLARS__EXECUTOR__GROUPBY_N_ARY", "43") m.setenv("CUDF_POLARS__EXECUTOR__BROADCAST_JOIN_LIMIT", "44") - m.setenv("CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD", "tasks") m.setenv("CUDF_POLARS__CUDA_STREAM_POLICY", "default") engine = pl.GPUEngine() @@ -422,11 +345,8 @@ def test_config_option_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.executor.cluster == "single" assert config.executor.fallback_mode == "silent" assert config.executor.max_rows_per_partition == 42 - assert config.executor.unique_fraction == {"a": 0.5} assert config.executor.target_partition_size == 100 - assert config.executor.groupby_n_ary == 43 assert config.executor.broadcast_join_limit == 44 - assert config.executor.shuffle_method == "tasks" assert config.cuda_stream_policy is None @@ -498,12 +418,6 @@ def test_default_executor() -> None: assert config.executor.name == "streaming" -def test_default_runtime() -> None: - config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert config.executor.name == "streaming" - assert config.executor.runtime == "tasks" - - @pytest.mark.parametrize( "memory_resource, memory_resource_config", [ @@ -537,10 +451,7 @@ def test_memory_resource(memory_resource, memory_resource_config) -> None: if memory_resource is None and memory_resource_config is None: # The default case: We make a new RMM MR, whose type depends on the GPU's features. - if _is_concurrent_managed_access_supported(): - assert isinstance(result, rmm.mr.PrefetchResourceAdaptor) - else: - assert isinstance(result, rmm.mr.CudaAsyncMemoryResource) + assert isinstance(result, rmm.mr.CudaAsyncMemoryResource) elif memory_resource is None: # Configured through memory_resource_config @@ -608,21 +519,17 @@ def test_cuda_stream_pool(): def test_cuda_stream_policy_default(monkeypatch: pytest.MonkeyPatch) -> None: # Default from engine config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert config.cuda_stream_policy is None + assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) - config = ConfigOptions.from_polars_engine( - pl.GPUEngine(executor_options={"runtime": "tasks"}) - ) - assert config.cuda_stream_policy is None + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) # Default from env monkeypatch.setenv("CUDF_POLARS__CUDA_STREAM_POLICY", "default") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.cuda_stream_policy is None - config = ConfigOptions.from_polars_engine( - pl.GPUEngine(executor_options={"runtime": "tasks"}) - ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) assert config.cuda_stream_policy is None @@ -635,26 +542,19 @@ def test_default_cuda_stream_policy(monkeypatch: pytest.MonkeyPatch) -> None: assert isinstance(result, CUDAStreamPoolConfig) -def test_cuda_stream_policy_from_config(*, rapidsmpf_single_available: bool) -> None: +def test_cuda_stream_policy_from_config() -> None: engine = pl.GPUEngine( executor="streaming", - executor_options={"runtime": "rapidsmpf"}, cuda_stream_policy={ "pool_size": 32, "flags": rmm.pylibrmm.CudaStreamFlags.NON_BLOCKING, }, ) - if rapidsmpf_single_available: - config = ConfigOptions.from_polars_engine(engine) - assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) - assert config.cuda_stream_policy.pool_size == 32 - assert ( - config.cuda_stream_policy.flags == rmm.pylibrmm.CudaStreamFlags.NON_BLOCKING - ) - config.cuda_stream_policy.build().get_stream() # no exception - else: - with pytest.raises(ValueError, match="The rapidsmpf streaming engine"): - ConfigOptions.from_polars_engine(engine) + config = ConfigOptions.from_polars_engine(engine) + assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) + assert config.cuda_stream_policy.pool_size == 32 + assert config.cuda_stream_policy.flags == rmm.pylibrmm.CudaStreamFlags.NON_BLOCKING + config.cuda_stream_policy.build().get_stream() # no exception @pytest.mark.parametrize( @@ -667,26 +567,19 @@ def test_cuda_stream_policy_from_config(*, rapidsmpf_single_available: bool) -> '{"pool_size": 32}', ], ) -def test_cuda_stream_policy_from_env( - monkeypatch: pytest.MonkeyPatch, env: str, *, rapidsmpf_single_available: bool -) -> None: +def test_cuda_stream_policy_from_env(monkeypatch: pytest.MonkeyPatch, env: str) -> None: monkeypatch.setenv("CUDF_POLARS__CUDA_STREAM_POLICY", env) - runtime = "tasks" if env == "default" else "rapidsmpf" - engine = pl.GPUEngine(executor="streaming", executor_options={"runtime": runtime}) - if runtime == "rapidsmpf" and rapidsmpf_single_available: - config = ConfigOptions.from_polars_engine(engine) + engine = pl.GPUEngine(executor="streaming") + config = ConfigOptions.from_polars_engine(engine) + if env == "default": + assert config.cuda_stream_policy is None + else: assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) if env == "pool": assert config.cuda_stream_policy.pool_size == 16 assert config.cuda_stream_policy.flags == CudaStreamFlags.NON_BLOCKING else: assert config.cuda_stream_policy.pool_size == 32 - elif runtime == "rapidsmpf": - with pytest.raises(ValueError, match="The rapidsmpf streaming engine"): - ConfigOptions.from_polars_engine(engine) - else: - config = ConfigOptions.from_polars_engine(engine) - assert config.cuda_stream_policy is None def test_cuda_stream_policy_from_env_invalid(monkeypatch: pytest.MonkeyPatch): @@ -696,41 +589,26 @@ def test_cuda_stream_policy_from_env_invalid(monkeypatch: pytest.MonkeyPatch): def test_cuda_stream_policy_default_rapidsmpf(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("rapidsmpf") - # Default from engine - config = ConfigOptions.from_polars_engine( - pl.GPUEngine(executor_options={"runtime": "rapidsmpf"}) - ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) assert isinstance(config.cuda_stream_policy, CUDAStreamPoolConfig) assert config.cuda_stream_policy.pool_size == 16 assert config.cuda_stream_policy.flags == rmm.pylibrmm.CudaStreamFlags.NON_BLOCKING # "default" user argument overrides pool default monkeypatch.setenv("CUDF_POLARS__CUDA_STREAM_POLICY", "default") - config = ConfigOptions.from_polars_engine( - pl.GPUEngine(executor_options={"runtime": "rapidsmpf"}) - ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) assert config.cuda_stream_policy is None -@pytest.mark.parametrize( - "polars_kwargs", - [ - {"executor": "in-memory"}, - {"executor": "streaming", "executor_options": {"runtime": "tasks"}}, - ], -) -def test_cuda_stream_policy_pool_only_supported_by_rapidsmpf( - polars_kwargs: dict[str, Any], -) -> None: +def test_cuda_stream_policy_pool_in_memory_unsupported() -> None: with pytest.raises( ValueError, - match="A stream pool is only supported by the rapidsmpf runtime.", + match="A stream pool is only supported by the streaming executor.", ): ConfigOptions.from_polars_engine( pl.GPUEngine( - **polars_kwargs, + executor="in-memory", cuda_stream_policy={"pool_size": 32, "flags": "NON_BLOCKING"}, ) ) @@ -903,8 +781,3 @@ def test_dask_sink_to_directory_false_raises() -> None: ValueError, match="The dask cluster requires sink_to_directory=True" ): StreamingExecutor(cluster=Cluster.DASK, sink_to_directory=False) - - -def test_get_dask_cuda_stream() -> None: - stream = get_dask_cuda_stream() - assert stream is not None diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index e3e788f2866f..a655efbe422a 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -718,7 +718,7 @@ def test_scan_parquet_zero_width_with_limit( ): request.applymarker( pytest.mark.xfail( - is_streaming_engine(engine) and custom_engine is None, + is_streaming_engine(engine) or custom_engine is not None, reason="https://github.com/rapidsai/cudf/issues/21644", ) ) diff --git a/python/cudf_polars/tests/test_sink.py b/python/cudf_polars/tests/test_sink.py index 7b69f6904b4f..d23559d21340 100644 --- a/python/cudf_polars/tests/test_sink.py +++ b/python/cudf_polars/tests/test_sink.py @@ -157,6 +157,7 @@ def test_chunked_sink_empty_table_to_parquet(tmp_path): pl.LazyFrame(), tmp_path / "out.parquet", engine=pl.GPUEngine( + executor="in-memory", raise_on_fail=True, parquet_options={"chunked": True, "n_output_chunks": 2}, ), diff --git a/python/cudf_polars/tests/test_tracing.py b/python/cudf_polars/tests/test_tracing.py index 184c0a77d38b..283ca361682b 100644 --- a/python/cudf_polars/tests/test_tracing.py +++ b/python/cudf_polars/tests/test_tracing.py @@ -55,9 +55,10 @@ def test_trace_basic( assert b"frames_input" in result assert b"total_bytes_output" in result assert b"total_bytes_input" in result - assert b"rmm_total_bytes_output" in result - assert b"rmm_total_bytes_input" in result - assert b"rmm_current_bytes_output" in result + # TODO: With rapidsmpf are the rmm fields not supposed to be logged? + assert b"rmm_total_bytes_output" not in result + assert b"rmm_total_bytes_input" not in result + assert b"rmm_current_bytes_output" not in result assert b"overhead_duration" in result @@ -79,10 +80,6 @@ def test_import_without_structlog() -> None: subprocess.check_call([sys.executable, "-c", code]) -@pytest.mark.skipif( - cudf_polars.testing.asserts.DEFAULT_RUNTIME != "rapidsmpf", - reason="Requires 'rapidsmpf' runtime.", -) def test_log_query_plan() -> None: """Test that log_query_plan emits a Query Plan event.""" import os @@ -98,7 +95,6 @@ def test_log_query_plan() -> None: executor="streaming", executor_options={ "cluster": "single", - "runtime": "rapidsmpf", "max_rows_per_partition": 5, }, memory_resource=rmm.mr.ManagedMemoryResource(), @@ -126,7 +122,6 @@ def test_log_query_plan() -> None: reason="Requires CUDF_POLARS_LOG_TRACES=1.", ) def test_sets_cudf_polars_query_id(): - pytest.importorskip("rapidsmpf") left = pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) right = pl.LazyFrame({"a": [1, 2, 3], "c": [7, 8, 9]}) @@ -136,7 +131,6 @@ def test_sets_cudf_polars_query_id(): engine = pl.GPUEngine( executor="streaming", raise_on_fail=True, - executor_options={"runtime": "rapidsmpf"}, ) with structlog.testing.capture_logs( diff --git a/python/cudf_polars/tests/testing/test_engine_utils.py b/python/cudf_polars/tests/testing/test_engine_utils.py index faf113502d6d..346a11acf2e7 100644 --- a/python/cudf_polars/tests/testing/test_engine_utils.py +++ b/python/cudf_polars/tests/testing/test_engine_utils.py @@ -3,8 +3,6 @@ from __future__ import annotations -import pytest - from cudf_polars.testing.engine_utils import ( EngineFixtureParam, create_streaming_options, @@ -30,7 +28,6 @@ def test_engine_fixture_param_small_blocksize(): def test_create_streaming_options_medium(): - pytest.importorskip("rapidsmpf") opts = create_streaming_options("medium") assert opts.max_rows_per_partition == 50 assert opts.target_partition_size == 1_000_000 @@ -38,7 +35,6 @@ def test_create_streaming_options_medium(): def test_create_streaming_options_small(): - pytest.importorskip("rapidsmpf") opts = create_streaming_options("small") assert opts.max_rows_per_partition == 4 assert opts.target_partition_size == 10 @@ -46,7 +42,6 @@ def test_create_streaming_options_small(): def test_create_streaming_options_overrides_merge(): """Overrides take precedence over the blocksize baseline.""" - pytest.importorskip("rapidsmpf") from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions overrides = StreamingOptions(max_rows_per_partition=999)