Move collectives module

[rjzamora]

Description

• Follow up to [cudf_polars] Reorganize package layout #22491
• Moves the collectives module under actor_graph to break circular dependencies. The "collectives" are mostly used to build the actor graph anyway.

Note: Before merging this, I'd like to get confirmation that others see circular-import errors locally. E.g.

pytest -v python/cudf_polars/tests/streaming/test_groupby.py

...

E   ImportError: cannot import name 'ShuffleManager' from partially initialized module 'cudf_polars.streaming.collectives.shuffle' (most likely due to a circular import) (/raid/rzamora/rapids-26.06/cudf/python/cudf_polars/cudf_polars/streaming/collectives/shuffle.py)

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1a356e8b-32fb-4569-901e-9289b2f86007

📥 Commits

Reviewing files that changed from the base of the PR and between 2dba3ec and e417b31.

📒 Files selected for processing (19)

• python/cudf_polars/cudf_polars/engine/core.py
• python/cudf_polars/cudf_polars/engine/spmd.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
• python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py
• python/cudf_polars/tests/streaming/test_allgather.py
• python/cudf_polars/tests/streaming/test_groupby.py
• python/cudf_polars/tests/streaming/test_metadata.py
• python/cudf_polars/tests/streaming/test_shuffler.py
• python/cudf_polars/tests/streaming/test_spmd.py

---

📝 Walkthrough

Summary by CodeRabbit

• Chores
  • Reorganized internal module structure for streaming components to improve code organization and maintainability.

Walkthrough

This PR consolidates collective operation utilities into the cudf_polars.streaming.actor_graph.collectives module hierarchy, moving ReserveOpIDs, reserve_op_id, and collective managers (AllGatherManager, ShuffleManager) from streaming.collectives to their new namespace. All dependent imports across engine, actor-graph, benchmark, and test modules are updated accordingly.

Changes

Collective Utilities Reorganization

Layer / File(s)	Summary
Collectives module reorganization cudf_polars/streaming/actor_graph/collectives/__init__.py, cudf_polars/streaming/actor_graph/collectives/common.py, cudf_polars/streaming/actor_graph/collectives/shuffle.py, cudf_polars/streaming/actor_graph/collectives/sort.py	Module docstrings and imports updated to reference actor_graph.collectives as the source for ReserveOpIDs, reserve_op_id, and collective managers. Streaming operator types (StreamingSink, Join, Over, Repartition, Shuffle) are imported from non-experimental cudf_polars.streaming modules.
Actor-graph modules integration cudf_polars/streaming/actor_graph/__init__.py, cudf_polars/streaming/actor_graph/groupby.py, cudf_polars/streaming/actor_graph/join.py, cudf_polars/streaming/actor_graph/over.py, cudf_polars/streaming/actor_graph/repartition.py, cudf_polars/streaming/actor_graph/utils.py	Groupby, join, over, repartition, and utils modules updated to import collective managers (AllGatherManager, ShuffleManager) and reservation utilities from the reorganized actor_graph.collectives hierarchy. Actor-graph __init__.py adds side-effect imports for shuffle and sort registration.
Engine and benchmark updates cudf_polars/engine/core.py, cudf_polars/engine/spmd.py, cudf_polars/streaming/benchmarks/utils.py	Engine modules and benchmark utilities updated to import reserve_op_id from cudf_polars.streaming.actor_graph.collectives.common instead of the prior cudf_polars.streaming.collectives.common.
Test module imports tests/streaming/test_allgather.py, tests/streaming/test_groupby.py, tests/streaming/test_metadata.py, tests/streaming/test_shuffler.py, tests/streaming/test_spmd.py	Test modules updated to import collective managers and reservation utilities from the new cudf_polars.streaming.actor_graph.collectives module paths.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• rapidsai/cudf#22491: Both PRs rewire the same collective operation ID helper (reserve_op_id / ReserveOpIDs) imports in the streaming/engine layer to new cudf_polars.streaming module paths (main PR targets actor_graph.collectives.common, retrieved PR targets the broader streaming layout).
• rapidsai/cudf#22191: Both PRs are connected through the collective-ID reservation plumbing (reserve_op_id/ReserveOpIDs) needed for streaming over() support: the main PR repoints the code to the actor-graph collective reservation utilities, while the retrieved PR extends collective reservation to include the Over IR node when dynamic planning is enabled.

Suggested reviewers

• TomAugspurger
• vyasr
• mroeschke

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title "Move collectives module" accurately summarizes the main change: relocating the collectives module under actor_graph to resolve circular dependencies.
Description check	✅ Passed	The PR description clearly explains the purpose (breaking circular dependencies), references the related PR #22491, and provides context about why the change is needed.
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

---

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

[rjzamora]

/merge