Define AllGather.__enter/exit__ for insert_finished - #1043
Conversation
mroeschke
added
improvement
| 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"
pentschev
left a comment
pentschev
left a comment
There was a problem hiding this comment.
Left a suggestion that we keep testing both cases, otherwise LGTM.
| with gather as ag: | ||
| while (msg := await ch_in.recv(context)) is not None: | ||
| chunk = PackedDataChunk.from_message(msg, br=context.br()).to_packed_data() | ||
| ag.insert(msg.sequence_number, chunk) |
There was a problem hiding this comment.
Could we actually test both the context manager and the "raw" use with insert_finished()?
| # Insert data from this rank and mark as finished | ||
| with allgather as ag: | ||
| for i in range(n_inserts): | ||
| packed_data = generate_packed_data( | ||
| n_elements, gen_offset(i, this_rank), stream, br | ||
| ) | ||
| ag.insert(i, packed_data) |
|
You need to ban |
…xercise context and non-context paths
…ke/rapidsmpf into feat/python/allgather_cm
Done in da751a3 |
|
/merge |
4 participants
In cudf_polars we often have to use a
finally: AllGather.insert_finished()pattern to ensure that method is called afterAllGather.inserte.g. NVIDIA/cudf#22516.This PR defines
.__enter/exit__like we do onContextto simplify this convention when we need to callAllGather.insert