Fix ApplicationContext reference cycle and unhandled write exceptions

cpp/include/ucxx/inflight_requests.h

@@ -96,14 +96,22 @@ class InflightRequests {
   void remove(const std::shared_ptr<Request>& request);
 
   /**
-   * @brief Issue cancelation of all inflight requests and clear the internal container.
+   * @brief Issue cancelation of inflight requests and clear the internal container.
    *
-   * Issue cancelation of all inflight requests known to this object and clear the
-   * internal container. The total number of canceled requests is returned.
+   * Issue cancelation of inflight requests known to this object.  The total number of
+   * canceled requests is returned.
    *
+   * When `workerOnly` is `true`, only worker-scoped operations
+   * (`Request::isWorkerOperation() == true`, i.e. receive variants) are cancelled and
+   * removed; endpoint-scoped operations are left in the container untouched, on the
+   * assumption that the caller will follow up with `ucp_ep_close_nbx(FORCE)` which
+   * handles their UCT-level cleanup atomically.  When `workerOnly` is `false` (the
+   * default), all inflight requests are cancelled.
+   *
+   * @param[in] workerOnly  if `true`, cancel only worker-scoped requests.
    * @returns The total number of canceled requests.
    */
-  size_t cancelAll();
+  size_t cancelAll(bool workerOnly = false);
 
   /**
    * @brief Releases the internally-tracked containers.

cpp/include/ucxx/request.h

@@ -127,6 +127,27 @@ class Request : public Component {
    */
   virtual void cancel();
 
+  /**
+   * @brief Whether this request is worker-scoped (vs endpoint-scoped).
+   *
+   * Returns `true` for operations that are tracked by the UCP worker rather than a
+   * specific UCP endpoint, currently the receive variants (`TagReceive`,
+   * `TagReceiveWithHandle`, `AmReceive`, `StreamReceive`, `TagMultiReceive`).
+   *
+   * `Endpoint::closeBlocking()` uses this to decide which requests must be cancelled
+   * explicitly via `ucp_request_cancel`: worker-scoped requests are not cancelled by
+   * `ucp_ep_close_nbx(UCP_EP_CLOSE_FLAG_FORCE)` (which only tears down endpoint-bound
+   * state) and would otherwise hang forever. Endpoint-scoped requests are left to
+   * UCX's FORCE close, which handles their UCT-level cleanup atomically, calling
+   * `ucp_request_cancel` on them in addition to FORCE close has been observed to
+   * leave UCT-level pending queue entries referencing freed staging buffers, which
+   * the next `ucp_worker_progress()` then crashes dispatching (see
+   * test_send_recv_multi.py CUDA segfault).
+   *
+   * @returns `true` if this request is worker-scoped, `false` if endpoint-scoped.
+   */
+  [[nodiscard]] bool isWorkerOperation() const;
+
   /**
    * @brief Return the status of the request.
    *

cpp/src/endpoint.cpp

@@ -269,13 +269,6 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
 {
   if (_closing.exchange(true) || _handle == nullptr) return;
 
-  size_t canceled = cancelInflightRequestsBlocking(3000000000 /* 3s */, 3);
-  ucxx_debug("ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, canceled %lu requests",
-             __func__,
-             this,
-             _handle,
-             canceled);
-
   ucp_request_param_t param{};
   if (_endpointErrorHandling)
     param = {.op_attr_mask = UCP_OP_ATTR_FIELD_FLAGS, .flags = UCP_EP_CLOSE_FLAG_FORCE};
@@ -288,8 +281,45 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
     bool submitted    = false;
     for (uint64_t i = 0; i < maxAttempts && !closeSuccess; ++i) {
       if (!submitted) {
+        // Cancel WORKER-SCOPED inflight requests (tag_recv & friends) and submit
+        // FORCE close ATOMICALLY in a single pre-callback no
+        // ucp_worker_progress() between them, and no ucp_request_cancel() on
+        // endpoint-scoped operations.
+        //
+        // Why cancel anything at all (UCX FORCE close handles endpoint state):
+        //   Receive operations (`tag_recv`, `am_recv`, ...) post requests on the
+        //   UCP worker via `ucp_*_recv_nbx(worker, ...)`. ucp_ep_close_nbx(FORCE)
+        //   tears down the UCP endpoint but leaves worker-scoped requests
+        //   pending forever, an `await ep.close()` racing with an outstanding
+        //   `await ep.recv()` would hang. See
+        //   test_shutdown.py::test_{server,client}_shutdown.
+        //
+        // Why ONLY worker-scoped (not endpoint-scoped sends/RMA):
+        //   Calling ucp_request_cancel() on an endpoint-scoped tag_send and then
+        //   FORCE-closing the endpoint has been observed to leave UCT-level TCP
+        //   pending queue entries pointing at freed staging buffers; the next
+        //   ucp_worker_progress() then crashed dispatching them
+        //   (uct_tcp_pending_queue_dispatch -> uct_cuda_copy_ep_get_short ->
+        //   cuMemcpyAsync -> SIGSEGV), see test_send_recv_multi.py CUDA segfault.
+        //   FORCE close handles endpoint-scoped requests' UCT cleanup atomically
+        //   on its own, so we leave them alone.
+        //
+        // Why atomic with FORCE close (not as a separate pre-callback):
+        //   When cancelAll and FORCE close were separate pre-callbacks (the
+        //   old cancelInflightRequestsBlocking path), a full ucp_worker_progress()
+        //   ran between them, see prior commit "Cancel inflight requests and
+        //   submit force-close atomically in a single pre-callback".
         if (!worker->registerGenericPre(
-              [this, &status, &param]() { status = ucp_ep_close_nbx(_handle, &param); }, period))
+              [this, &status, &param]() {
+                _inflightRequests->cancelAll(/*workerOnly=*/true);
+                status = ucp_ep_close_nbx(_handle, &param);
+                // Invalidate _handle synchronously immediately, to prevent
+                // time window where _handle` points to freed UCP memory, usually
+                // observed in `populateDelayedSubmission()`.
+                _originalHandle = _handle;
+                _handle         = nullptr;
+              },
+              period))
           continue;
         submitted = true;
       }
@@ -306,7 +336,7 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
                     "endpoint: %s",
                     __func__,
                     this,
-                    _handle,
+                    _originalHandle,
                     ucs_status_string(UCS_PTR_STATUS(status)));
                 }
               },
@@ -323,10 +353,16 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
         "ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, all attempts to close timed out",
         __func__,
         this,
-        _handle);
+        _originalHandle != nullptr ? _originalHandle : _handle);
     }
   } else {
-    status = ucp_ep_close_nbx(_handle, &param);
+    // No progress thread: cancel inflight + FORCE close back-to-back, then
+    // drive progress here. Same atomicity reasoning as the progress-thread
+    // path above (no ucp_worker_progress() between cancel and FORCE close).
+    _inflightRequests->cancelAll(/*workerOnly=*/true);
+    status          = ucp_ep_close_nbx(_handle, &param);
+    _originalHandle = _handle;
+    _handle         = nullptr;
     if (UCS_PTR_IS_PTR(status)) {
       ucs_status_t s;
       while ((s = ucp_request_check_status(status)) == UCS_INPROGRESS)
@@ -337,11 +373,12 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
         "ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, Error while closing endpoint: %s",
         __func__,
         this,
-        _handle,
+        _originalHandle,
         ucs_status_string(UCS_PTR_STATUS(status)));
     }
   }
-  ucxx_trace("ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, closed", __func__, this, _handle);
+  ucxx_trace(
+    "ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, closed", __func__, this, _originalHandle);
 
   if (UCS_PTR_IS_PTR(status)) ucp_request_free(status);
 
@@ -351,14 +388,12 @@ void Endpoint::closeBlocking(uint64_t period, uint64_t maxAttempts)
       ucxx_debug("ucxx::Endpoint::%s, Endpoint: %p, UCP handle: %p, calling user close callback",
                  __func__,
                  this,
-                 _handle);
+                 _originalHandle);
       _closeCallback(_status, _closeCallbackArg);
       _closeCallback    = nullptr;
       _closeCallbackArg = nullptr;
     }
   }
-
-  std::swap(_handle, _originalHandle);
 }
 
 ucp_ep_h Endpoint::getHandle() { return _handle; }
@@ -448,7 +483,7 @@ size_t Endpoint::cancelInflightRequestsBlocking(uint64_t period, uint64_t maxAtt
         "cancel inflight requests failed",
         __func__,
         this,
-        _handle);
+        _originalHandle != nullptr ? _originalHandle : _handle);
   } else {
     canceled = _inflightRequests->cancelAll();
   }

cpp/src/inflight_requests.cpp

@@ -41,18 +41,34 @@ void InflightRequests::merge(TrackedRequests&& trackedRequests)
     if (r) _canceling.insert(std::move(r));
 }
 
-size_t InflightRequests::cancelAll()
+size_t InflightRequests::cancelAll(bool workerOnly)
 {
   decltype(_inflight) toCancel;
+  decltype(_inflight) toKeep;
   {
     std::lock_guard<std::mutex> lock(_mutex);
-    toCancel = std::exchange(_inflight, {});
+    if (workerOnly) {
+      // Partition: receive (worker-scoped) requests get explicit ucp_request_cancel,
+      // endpoint-scoped requests stay in `_inflight` to be cleaned up by FORCE close.
+      for (auto& r : _inflight) {
+        if (r && r->isWorkerOperation())
+          toCancel.insert(r);
+        else
+          toKeep.insert(r);
+      }
+      _inflight = std::move(toKeep);
+    } else {
+      toCancel = std::exchange(_inflight, {});
+    }
   }
 
   size_t total = toCancel.size();
   if (total == 0) return 0;
 
-  ucxx_debug("ucxx::InflightRequests::%s, canceling %lu requests", __func__, total);
+  ucxx_debug("ucxx::InflightRequests::%s, canceling %lu requests (workerOnly=%d)",
+             __func__,
+             total,
+             workerOnly);
 
   for (auto& r : toCancel) {
     if (r) r->cancel();

cpp/src/request.cpp

@@ -71,6 +71,24 @@ Request::~Request()
   ucxx_trace("ucxx::Request destroyed (%s): %p", _operationName.c_str(), this);
 }
 
+bool Request::isWorkerOperation() const
+{
+  // Receive operations go through ucp_tag_recv_nbx / ucp_am_recv_nbx and friends,
+  // which post requests against the UCP worker rather than a specific endpoint.
+  // ucp_ep_close_nbx(FORCE) does not cancel them, so closeBlocking() must do so
+  // explicitly.  Send / RMA operations are endpoint-bound and are cleaned up by
+  // FORCE close itself.
+  return std::visit(data::dispatch{
+                      [](const data::TagReceive&) { return true; },
+                      [](const data::TagReceiveWithHandle&) { return true; },
+                      [](const data::TagMultiReceive&) { return true; },
+                      [](const data::AmReceive&) { return true; },
+                      [](const data::StreamReceive&) { return true; },
+                      [](const auto&) { return false; },
+                    },
+                    _requestData);
+}
+
 void Request::cancel()
 {
   std::lock_guard<std::recursive_mutex> lock(_mutex);

python/distributed-ucxx/distributed_ucxx/ucxx.py

@@ -498,9 +498,9 @@ async def write(
                 for each_frame in send_frames:
                     await self.ep.send(each_frame)
             return sum(sizes)
-        except ucxx.exceptions.UCXError:
+        except BaseException as e:
             self.abort()
-            raise CommClosedError("While writing, the connection was closed")
+            raise CommClosedError("While writing, the connection was closed") from e
 
     @log_errors
     async def read(self, deserializers=("cuda", "dask", "pickle", "error")):
@@ -729,6 +729,8 @@ async def serve_forever(client_ep):
         self.ucxx_server = ucxx.create_listener(serve_forever, port=self._input_port)
 
     def stop(self):
+        if self.ucxx_server is not None:
+            self.ucxx_server.close()
         self.ucxx_server = None
         _deregister_dask_resource(self._resource_id)
 

python/distributed-ucxx/distributed_ucxx/utils_test.py

@@ -1,9 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: BSD-3-Clause
 
 from __future__ import annotations
 
 import asyncio
+import gc
 import logging
 import sys
 
@@ -84,6 +85,10 @@ def ucxx_loop(request):
 
     with check_thread_leak():
         yield loop
+
+        # Collect garbage to break any remaining reference cycles before reset.
+        gc.collect()
+
         if ignore_alive_references:
             try:
                 ucxx.reset()

python/ucxx/ucxx/_lib_async/application_context.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 import logging
@@ -311,7 +311,7 @@ def create_listener(
                 cb_args=(
                     loop,
                     callback_func,
-                    self,
+                    weakref.ref(self),
                     endpoint_error_handling,
                     connect_timeout,
                     listener_id,
@@ -321,6 +321,7 @@ def create_listener(
             ),
             listener_id,
             self._listener_active_clients,
+            ctx=self,
         )
         return ret