Minimize Dask resource acquisition in cudf_polars tests

[mroeschke]

Description

The motivation is to help alleviate potential CI issues due to Dask/Ray resource spin-up in cudf_polars tests, starting with test_dask.py

• Removes test_scan, test_filter, test_group_by, test_join, and test_empty_dataframe as they already have coverage in other test where we parameterize over engine
• For tests needing an instantiated DaskEngine, uses the session-scoped DaskEngine already created in conftest.py
• For tests exercising DaskEngine construction, uses a module-scoped dask_client with a LocalCluster (tests that use this are just testing engine properties so just uses 1 worker)

Overall, this PR reduces the 5 Dask clusters allocated to just 2

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• Tests
  • Improved Dask-related test setup with a dedicated, shared engine fixture and a separate local cluster/client fixture to improve isolation and reliability.
  • Tests now gracefully skip when the optional cluster dependency is unavailable; test coverage refocused on engine construction, cluster info, resource limits, run/reset, and shutdown behavior.

Walkthrough

Adds a session-scoped configured dask_engine fixture in tests/conftest.py and refactors python/cudf_polars/tests/streaming/test_dask.py to provide a module-scoped dask_client fixture for explicit Client injection; tests are updated to use dask_client for construction and dask_engine for runtime assertions.

Changes

Dask Engine Fixture Refactoring

Layer / File(s)	Summary
Shared dask_engine fixture in conftest python/cudf_polars/tests/conftest.py	Imports DaskEngine, adds dask_engine fixture that configures and returns a shared DaskEngine using configure_streaming_engine, and updates pytest_generate_tests to parametrize _engine_param only for the "dask" variant when dask_engine is requested.
test_dask.py fixture and test rewiring python/cudf_polars/tests/streaming/test_dask.py	Replaces module-scoped engine with dask_client fixture that creates LocalCluster + Client (silencing logs); updates constructor-focused tests to accept dask_client and call DaskEngine.from_options(..., dask_client=...); migrates runtime/GPU tests to the shared dask_engine; updates reset_engine and test_reset_after_shutdown_raises to use dask_client; removes several earlier query-behavior tests.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• rapidsai/cudf#22493: Refactors engine parametrization and fixture wiring machinery, including _engine_param and engine-selection behavior in conftest.py that closely relates to this change.

Suggested reviewers

• pentschev
• madsbk

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly and concisely summarizes the main change: minimizing Dask resource acquisition in tests by reducing cluster allocations from five to two.
Description check	✅ Passed	The description is related to the changeset, explaining the motivation (reducing CI issues), listing removed tests, and describing the new fixture approach to reduce resource usage.
Docstring Coverage	✅ Passed	Docstring coverage is 86.67% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

Review ran into problems

🔥 Problems

Stopped waiting for pipeline failures after 30000ms. One of your pipelines takes longer than our 30000ms fetch window to run, so review may not consider pipeline-failure results for inline comments if any failures occurred after the fetch window. Increase the timeout if you want to wait longer or run a @coderabbit review after the pipeline has finished.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[wence-]

Just import distributed in the TYPE_CHECKING block, no?

[mroeschke]

I ended up using a try/except for import distributed instead of distributed = pytest.importorskip in 01688c0 since the latter was probably preventing mypy from using that alias as a module

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

python/cudf_polars/tests/streaming/test_dask.py (1)

171-184: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Guarantee engine cleanup when assertions fail

test_reset_after_shutdown_raises does manual teardown, but if an assertion fails before the last shutdown(), the engine may not be cleaned up. Wrap the assertions in try/finally so cleanup is unconditional.

Proposed fix

 def test_reset_after_shutdown_raises(dask_client: distributed.Client) -> None:
     """``shutdown`` is idempotent; ``_reset`` after shutdown raises every time."""
     engine = DaskEngine(
         dask_client=dask_client,
         executor_options={"max_rows_per_partition": 10},
     )
-    engine.shutdown()
-    engine.shutdown()  # idempotent
-    with pytest.raises(RuntimeError, match="shut-down"):
-        engine._reset()
-    with pytest.raises(RuntimeError, match="shut-down"):
-        engine._reset()  # still raises on a second attempt
-    engine.shutdown()  # still safe after a failed _reset
+    try:
+        engine.shutdown()
+        engine.shutdown()  # idempotent
+        with pytest.raises(RuntimeError, match="shut-down"):
+            engine._reset()
+        with pytest.raises(RuntimeError, match="shut-down"):
+            engine._reset()  # still raises on a second attempt
+    finally:
+        engine.shutdown()  # always clean up even on assertion failure

As per coding guidelines, "Ensure proper cleanup in del and context managers to prevent GPU memory leaks."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf_polars/tests/streaming/test_dask.py` around lines 171 - 184, The
test test_reset_after_shutdown_raises creates a DaskEngine and calls shutdown
multiple times but doesn't guarantee engine.shutdown() runs if an assertion
fails; wrap the assertions that call engine._reset() in a try/finally so
engine.shutdown() is always invoked for cleanup. Concretely, after constructing
engine (DaskEngine(...)) and the initial shutdown calls, put the two with
pytest.raises(...) blocks inside a try block and call engine.shutdown() in the
finally block to ensure cleanup even on failure.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@python/cudf_polars/tests/streaming/test_dask.py`:
- Around line 171-184: The test test_reset_after_shutdown_raises creates a
DaskEngine and calls shutdown multiple times but doesn't guarantee
engine.shutdown() runs if an assertion fails; wrap the assertions that call
engine._reset() in a try/finally so engine.shutdown() is always invoked for
cleanup. Concretely, after constructing engine (DaskEngine(...)) and the initial
shutdown calls, put the two with pytest.raises(...) blocks inside a try block
and call engine.shutdown() in the finally block to ensure cleanup even on
failure.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d84be9bd-4cbf-4bc6-8d52-a0cdec7ade48

📥 Commits

Reviewing files that changed from the base of the PR and between a0c3541 and 1253c33.

📒 Files selected for processing (1)

• python/cudf_polars/tests/streaming/test_dask.py

[mroeschke]

/merge