[BUG] threads and streams issues

[davidkoski]

☑️ I understand it is strictly prohibited to use AI to write issues.

Investigation done against v0.32.2. See also #4480

This is observed in mlx-swift with cuda builds. 4480 fixed an issue where global command encoders were doing work on process shutdown after the cuda runtime terminated. This was sufficient on the multi-gpu, many core linux machines I was testing against, but when I run it under CI I found that sometimes it wasn't enough.

After tracking it down I think we have the following related issues:

• CudaHandle destructor calls throwing code that uses the cuda runtime
  • same as CommandEncoder
  • i don't think thread local resources can be safely leaked in the same way as the global resources
  • perhaps we should not allow any thrown behavior from destructors?
• I suspect this is reproducible from python if using unjoined streams that still exist at process exit
• determination of main thread may not be sufficient (but see next item, perhaps moot)
• functions like clear_streams() have special "main thread" handling that should be triggerable from any thread
• clear_streams() can only see global command encoders, not any of the thread local resources
• although mlx-swift only uses global streams (thread unsafe) there are a couple places that implicitly (and perhaps accidentally) create thread local streams

First, here is the crash I was observing:

 10 [ra] 0x0000614c5f960415 mlx::core::check_cuda_error(char const*, cudaError) + 352 in example1
 11 [ra] 0x0000614c5f4c435e mlx::core::CudaHandle<CUevent_st*, &cudaEventDestroy>::reset() + 59 in example1
 12 [ra] 0x0000614c5f4c36ee mlx::core::CudaHandle<CUevent_st*, &cudaEventDestroy>::~CudaHandle() + 41 in example1
 13 [ra] 0x0000614c5f4c2bbe mlx::core::cu::CudaEventHandle::~CudaEventHandle() + 27 in example1
 14 [ra] 0x0000614c5f4bea05 mlx::core::cu::CudaEvent::~CudaEvent() + 96 in example1
 15 [ra] 0x0000614c5f9a91da mlx::core::cu::Worker::~Worker() + 85 in example1
 16 [ra] 0x0000614c5f4b0255 void std::destroy_at<mlx::core::cu::Worker>(mlx::core::cu::Worker*) + 27 in example1
 17 [ra] 0x0000614c5f4b0236 void std::_Destroy<mlx::core::cu::Worker>(mlx::core::cu::Worker*) + 27 in example1
 18 [ra] 0x0000614c5f4b010a std::_Sp_counted_ptr_inplace<mlx::core::cu::Worker, std::allocator<void>, (__gnu_cxx::_Lock_policy)2>::_M_dispose() + 67 in example1
 19 [ra] 0x0000614c5e85edff std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release_last_use() + 38 in example1
 20 [ra] 0x0000614c5e85c950 std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release_last_use_cold() + 27 in example1
 21 [ra] 0x0000614c5e85bd0c std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release() + 303 in example1
 22 [ra] 0x0000614c5e85cb9b std::__shared_count<(__gnu_cxx::_Lock_policy)2>::~__shared_count() + 42 in example1
 23 [ra] 0x0000614c5f4a2ece std::__shared_ptr<mlx::core::cu::Worker, (__gnu_cxx::_Lock_policy)2>::~__shared_ptr() + 31 in example1
 24 [ra] 0x0000614c5f4a2eee std::shared_ptr<mlx::core::cu::Worker>::~shared_ptr() + 27 in example1
 25 [ra] 0x0000614c5f9aa570 std::_Head_base<1ul, std::shared_ptr<mlx::core::cu::Worker>, false>::~_Head_base() + 27 in example1
 26 [ra] 0x0000614c5f9aa590 std::_Tuple_impl<1ul, std::shared_ptr<mlx::core::cu::Worker> >::~_Tuple_impl() + 27 in example1
 27 [ra] 0x0000614c5f9aa5b0 std::_Tuple_impl<0ul, void (mlx::core::cu::Worker::*)(), std::shared_ptr<mlx::core::cu::Worker> >::~_Tuple_impl() + 27 in example1
 28 [ra] 0x0000614c5f9aa5d0 std::tuple<void (mlx::core::cu::Worker::*)(), std::shared_ptr<mlx::core::cu::Worker> >::~tuple() + 27 in example1
 29 [ra] 0x0000614c5f9aa5f0 std::thread::_Invoker<std::tuple<void (mlx::core::cu::Worker::*)(), std::shared_ptr<mlx::core::cu::Worker> > >::~_Invoker() + 27 in example1
 30 [ra] 0x0000614c5f9ad044 std::thread::_State_impl<std::thread::_Invoker<std::tuple<void (mlx::core::cu::Worker::*)(), std::shared_ptr<mlx::core::cu::Worker> > > >::~_State_impl() + 49 in example1
 31 [ra] 0x0000614c5f9ad070 std::thread::_State_impl<std::thread::_Invoker<std::tuple<void (mlx::core::cu::Worker::*)(), std::shared_ptr<mlx::core::cu::Worker> > > >::~_State_impl() + 27 in example1

The Worker owns a CudaEvenHandler, which calls into the cuda runtime and throws because it has already shut down. Same issue as 4480 but appears to be coming from a thread local source (which is curious because mlx-swift only uses the global stream resources).

It turns out eval() and compile() both implicitly use thread local streams. In python this is fine -- it just picks up the stream the thread was already using. In mlx-swift we can't use thread local resources because Tasks span threads (via async or continuations). Likely other async supporting languages would see the same issue.

array eval_impl(std::vector<array> outputs, bool async) {
  std::deque<array> tape;

  // Make an effort to choose a good output stream
  Stream stream = default_stream(default_device());
  for (auto& o : outputs) {
    if (o.status() == array::Status::unscheduled && o.has_primitive()) {
      stream = o.primitive().stream();
      break;
    }
  }


class CompileCache {

  std::tuple<CacheEntry&, std::shared_ptr<std::vector<CacheEntry>>> find(
      std::uintptr_t fun_id,
      const std::vector<array>& inputs,
      bool shapeless,
      const std::vector<uint64_t>& constants) {
...
    // Loop over entries and check:
    // - Default stream and device match the entry's default stream
    // - Inputs match i.e. shapes and types must be equal.
    auto stream = default_stream(default_device());
    for (CacheEntry& entry : entries) {
      // Check that the default stream and device match
      if (entry.stream != stream) {
        continue;
      }
      if (entry.shapeless != shapeless) {
        continue;
      }

      // Check the inputs match and return if so
      if (has_same_shape_and_dtype(inputs, entry.inputs) &&
          constants == entry.constants) {
        return {entry, std::move(entries_ptr)};
      }
    }

Previously (the latest mlx-swift release) there were no thread local streams, so this always got the default global stream. Now this produces thread local streams. This would be relatively harmless if they didn't crash on exit.

The next issue is how the main thread is determined:

bool is_main_thread() {
  static auto main_thread_id = std::this_thread::get_id();
  return main_thread_id == std::this_thread::get_id();
}

this will capture the thread id of the first caller. In mlx-swift this might any thread. Not only do we not know which thread, there is no way to get back to it. It might even exit and be permanently gone.

I am not sure mlx-swift can make use of it, but:

void clear_streams() {
  metal::get_command_encoders().clear();
  if (is_main_thread()) {
    metal::get_global_command_encoders().clear();
  }
}

and probably other code behaves differently if running on the main thread. I think this needs to take a parameter (perhaps a tri-state optional bool if we want the default behavior) so that callers that can't participate in the main thread can still use the call.

This might be useful in swift if 1) we could cause it to clear the global command encoders and 2) if there were not any secret thread local streams being created.

Suggestions

I think these changes might be useful:

• do not create thread local streams in eval_impl and CompileCache

  • are there more? can we make it impossible to accidentally add this back in?

• make Metal and Cuda resources never throw in their destructors

• ideally the notion of main thread would be a parameter (even if just an override)

[aleroot]

@davidkoski I tried to go down this rabbit hole even if I do not have any cuda card and added #4514, just to try to see if we can unblock the release on mlx-swift. Give it a try if you got some cuda hardware. Thanks.