[BUG] pending requests should be cancelled when requesting tasks complete

[abellina]

Given: #1177, a better approach is to cancel pending requests for any task that has completed.

This means we need:

1. Out of band message to the server that says "cancel transfer request for buffer ids x,y,z".
2. We would need to make sure we can "un-send" from UCX, using the cancel function, which means remove the tag send request.
3. After this, we need to release bounce buffers.

[abellina]

So there are four categories of shuffle block requests we could be cancelling:

Client requests:

• Waiting for bounce buffers: These were queued, and have not been matched with a bounce buffers.
• With bounce buffers: Bounce buffer acquired, we may have sent a TransferRequest

Server requests:

• Waiting for bounce buffers: Queued. The client is waiting, however, with a receive.
• With bounce buffers: Bounce buffer acquired, we may have issued a send.

The lowest hanging fruit are client request that are waiting for bounce buffers: No resources have been acquired, and also no risk of a server having pending sends. So I think this part is doable today. This could cover a very large number of messages as well. Just thinking in terms of bounce buffer lengths.. we could have several bounce buffer lengths not ready to be received in the queue.

For the rest of the situations I think we need that out-of-band message. The number of requests in these situations is going to be those that fit in the bounce buffers we have. The server is at a disadvantage because it doesn't know when tasks completed, so it does have # receive bounce buffers * # of peer requests to deal with.

1. A client with bounce buffers would see the cancellation.
2. The client should keep all pending TransferRequests around, and iterate through them calling UCX.cancel for each receive issued as part of a TransferRequest. We need to make sure this part works correctly: before or during a transfer.
3. We also (concurrently with 2) need to issue a TransferRequestCancel message. Currently we assign each TransferRequest a tag, so we could use this as an id. The server would keep around all TransferRequestss pending, and match on the id. We'd need to cancel all pending sends, like the receives. Again making sure all of this works correctly with UCX. It may be we make this better by first issuing a TransferRequestCancel for any receive that hasn't been issued, and on ack we know that there aren't any sends either, but we allow the receives that have been issue to carry to completion.
4. As we cancel we need to unacquire buffers from the spill store (so they can be spilled/removed), we also need to give back bounce buffers, and free temporary buffer memory allocated in BufferReceiveState for that buffer that was larger than the window.

[sameerz]

@abellina will this go into 21.08 or should we move it out of the release?

[abellina]

So there are four categories of shuffle block requests we could be cancelling:

Client requests:

1. Waiting for bounce buffers: These were queued, and have not been matched with a bounce buffers.
2. With bounce buffers: Bounce buffer acquired, we may have sent a TransferRequest

Server requests:

3. Waiting for bounce buffers: Queued. The client is waiting, however, with a receive.
4. With bounce buffers: Bounce buffer acquired, we may have issued a send.

Out of the original categories we handle (1) on task complete. This was introduced here #2553. This was the "low hanging fruit" as identified in a later comment.

The rest of the categories are not handled yet. Changes for active messages in 21.08 help with this, since the receive side is not actively receiving, instead it has an active message waiting which receives dynamically. We are also now using a single id for the BufferReceiveState in the headers of active messages, so it's simpler to cancel, and the server is already tracking these.

So the main task pending is to issue a TransferRequestCancel and unwind (2) in the client. The server should handle it and do (3) and (4). I think this more doable with 21.08 than before. I'll move to 21.10 because it is still a rather large task, and need to discuss priorities.

[mattahrens]

Closing as won't fix for now