Define AllGather.__enter/exit__ for insert_finished#1043
Conversation
| with nogil: | ||
| deref(self._handle).insert_finished() | ||
|
|
||
| def __enter__(self): |
There was a problem hiding this comment.
I think we have to be careful in the streaming case because of
rapidsmpf/cpp/include/rapidsmpf/streaming/coll/allgather.hpp
Lines 23 to 30 in d8e53ff
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.
There was a problem hiding this comment.
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 selfand could move Inserter.insert onto AllGatherManager.insert if desired
There was a problem hiding this comment.
Oh ok I think that work, thanks.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
Yeah it's kinda nice to keep an API like inserting to signal that "after exiting the cm this calls insert_finished"
In cudf_polars we often have to use a
finally: AllGather.insert_finished()pattern to ensure that method is called afterAllGather.inserte.g. rapidsai/cudf#22516.This PR defines
.__enter/exit__like we do onContextto simplify this convention when we need to callAllGather.insert