Fix canceling inflight messages for which requests are closed

ucp/_libs/ucx_endpoint.pyx

@@ -22,12 +22,12 @@ cdef void _cancel_inflight_msgs(UCXWorker worker, set inflight_msgs):
     cdef dict req_info
     cdef str name
     for req in list(inflight_msgs):
-        assert not req.closed()
-        req_info = <dict>req._handle.info
-        name = req_info["name"]
-        logger.debug("Future cancelling: %s" % name)
-        # Notice, `request_cancel()` evoke the send/recv callback functions
-        worker.request_cancel(req)
+        if not req.closed():
+            req_info = <dict>req._handle.info
+            name = req_info["name"]
+            logger.debug("Future cancelling: %s" % name)
+            # Notice, `request_cancel()` evoke the send/recv callback functions
+            worker.request_cancel(req)
 
 
 class UCXEndpointCloseCallback():

ucp/_libs/ucx_worker.pyx

@@ -277,16 +277,17 @@ cdef class UCXWorker(UCXObject):
         ucp_worker_print_info(self._handle, text_fd)
         return decode_text_fd(text_fd)
 
-    def query_total_inflight_messages_to_cancel(self):
-        """Query the total of inflight messages scheduled to cancel
+    def cancel_inflight_messages(self):
+        """Cancel inflight messages scheduled for canceling
 
-        While there are messages scheduled for canceling, we need to progress
-        the worker. Therefore, this can be used to query if there are still any
-        such messages and progress while the result is larger than 0.
+        If any messages are scheduled for canceling, we need to trigger their
+        cancelation and return the number of canceled messages, if there are
+        any messages to cancel, the worker needs to progress to complete the
+        cancelation.
 
         Returns
         -------
-        total: The total number of inflight messages scheduled to cancel.
+        total: The total number of inflight messages canceled.
         """
         len_inflight_msgs_to_cancel = len(self._inflight_msgs_to_cancel)
         if len_inflight_msgs_to_cancel > 0:

ucp/continuous_ucx_progress.py

@@ -100,7 +100,7 @@ async def _arm_worker(self):
             # Cancel inflight messages that couldn't be completed. This may
             # happen if the user called ep.recv() but the remote worker
             # errored before sending the message.
-            if worker.query_total_inflight_messages_to_cancel() > 0:
+            if worker.cancel_inflight_messages() > 0:
                 worker.progress()
 
             del worker