Skip to content

Define AllGather.__enter/exit__ for insert_finished#1043

Open
mroeschke wants to merge 1 commit into
rapidsai:mainfrom
mroeschke:feat/python/allgather_cm
Open

Define AllGather.__enter/exit__ for insert_finished#1043
mroeschke wants to merge 1 commit into
rapidsai:mainfrom
mroeschke:feat/python/allgather_cm

Conversation

@mroeschke
Copy link
Copy Markdown
Contributor

In cudf_polars we often have to use a finally: AllGather.insert_finished() pattern to ensure that method is called after AllGather.insert e.g. rapidsai/cudf#22516.

This PR defines .__enter/exit__ like we do on Context to simplify this convention when we need to call AllGather.insert

@mroeschke mroeschke self-assigned this May 15, 2026
@mroeschke mroeschke requested a review from a team as a code owner May 15, 2026 18:00
@mroeschke mroeschke added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels May 15, 2026
with nogil:
deref(self._handle).insert_finished()

def __enter__(self):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we have to be careful in the streaming case because of

/**
* @brief Asynchronous (coroutine) interface to `coll::AllGather`.
*
* Once the AllGather is created, many tasks may insert data into it. If multiple tasks
* insert data, the user is responsible for arranging that `insert_finished` is only
* called after all `insert`ions have completed. A single consumer task should extract
* data.
*/

It's the callers responsibility to do this. We probably should only do the "bulk" case in this PR. And keep using AllGatherManager.Inserter in cudf-polars.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure we can still use the AllGatherManager.Inserter API in cudf_polars, but this would help simplify that implementation too, with this inserting could become

@contextmanager
def inserting(self):
    with self.allgather:
        yield self

and could move Inserter.insert onto AllGatherManager.insert if desired

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh ok I think that work, thanks.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think we can remove the Inserter class but we have to keep inserting cm, so the caller can control when insert_finished is called?

Copy link
Copy Markdown
Contributor Author

@mroeschke mroeschke May 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it's kinda nice to keep an API like inserting to signal that "after exiting the cm this calls insert_finished"

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants