Fix Windows shutdown deadlock and leak with the Ruy backend

python/tests/test_translator.py

@@ -4,6 +4,9 @@
 import logging
 import os
 import shutil
+import subprocess
+import sys
+import textwrap
 
 import numpy as np
 import pytest
@@ -824,3 +827,53 @@ def test_logging():
     with wurlitzer.pipes() as (_, err):
         _get_transliterator()
     assert not err.read()
+
+
+# Child process for test_shutdown_does_not_deadlock: build an int8 CPU Translator,
+# run one large translation batch to spin up the Ruy thread pool, then destroy it.
+# The deadlock (if present) happens during that destruction, so the interesting part
+# is whether the process returns at all.
+#
+# The batch is deliberately large: the hang only happens once Ruy has actually
+# spawned its internal thread pool, which it does only when the GEMM is big enough.
+# A tiny batch runs single-threaded, has no threads to join, and shuts down cleanly
+# even on an unpatched build -- so it would not catch the regression.
+_SHUTDOWN_CHILD = textwrap.dedent(
+    """
+    import sys
+    import ctranslate2
+
+    translator = ctranslate2.Translator(
+        sys.argv[1], device="cpu", compute_type="int8",
+        inter_threads=2, intra_threads=4,
+    )
+    source = [["آ", "ت", "ز", "م", "و", "ن"]] * 512
+    translator.translate_batch(source)
+    del translator
+    """
+)
+
+
+@pytest.mark.skipif(
+    sys.platform != "win32",
+    reason="Shutdown deadlock is Windows-specific (loader lock on thread exit)",
+)
+def test_shutdown_does_not_deadlock():
+    # Regression test for the Windows shutdown deadlock in the Ruy backend: the
+    # thread_local ruy::Context destructor joined ruy's internal thread pool while
+    # the owning worker thread was exiting (holding the loader lock), hanging
+    # process shutdown indefinitely. See ctranslate2-rs#64.
+    #
+    # Run the whole lifecycle in a subprocess so a hang fails via `timeout` instead
+    # of blocking the test runner. int8 selects the Ruy backend; the large batch (in
+    # the child) ensures Ruy spawns the thread pool whose join deadlocks when unpatched.
+    model_path = _get_model_path()
+    try:
+        subprocess.run(
+            [sys.executable, "-c", _SHUTDOWN_CHILD, model_path],
+            timeout=30,
+            check=True,
+            capture_output=True,
+        )
+    except subprocess.TimeoutExpired:
+        pytest.fail("Translator did not shut down within 30s (deadlock regression)")

src/cpu/backend.cc

@@ -106,10 +106,36 @@ namespace ctranslate2 {
     }
 
 #ifdef CT2_WITH_RUY
+#  ifdef _WIN32
+    // Windows only. The ruy::Context destructor joins ruy's internal thread pool;
+    // if that join runs while the owning thread is exiting (under the loader lock),
+    // it deadlocks ThreadPool shutdown (jkawamoto/ctranslate2-rs#64). So instead of
+    // a `thread_local` object destroyed at thread exit, we heap-allocate it and
+    // destroy it explicitly via clear_ruy_context() from ReplicaWorker::finalize(),
+    // which runs on the worker thread in a normal context (not thread exit), so the
+    // join completes and no memory (incl. ruy's prepacked cache) leaks.
+    static thread_local ruy::Context* ruy_context = nullptr;
+
+    ruy::Context *get_ruy_context() {
+      if (!ruy_context)
+        ruy_context = new ruy::Context();
+      return ruy_context;
+    }
+
+    void clear_ruy_context() {
+      delete ruy_context;
+      ruy_context = nullptr;
+    }
+#  else
+    // Other platforms keep the plain thread_local RAII: the join at thread exit is
+    // harmless here, and this avoids introducing a manual-cleanup path (and the leak
+    // that would follow if get_ruy_context() were ever called on a thread that never
+    // reaches finalize()).
     ruy::Context *get_ruy_context() {
       static thread_local ruy::Context context;
       return &context;
     }
+#  endif
 #endif
   }
 }

src/cpu/backend.h

@@ -28,6 +28,12 @@ namespace ctranslate2 {
     bool pack_gemm_weights(ComputeType compute_type);
 #ifdef CT2_WITH_RUY
     ruy::Context *get_ruy_context();
+#  ifdef _WIN32
+    // Windows only. Destroy the calling thread's ruy::Context (joins ruy's thread
+    // pool). Must be called from a normal execution context, not during thread exit
+    // — see backend.cc. Other platforms rely on thread_local RAII instead.
+    void clear_ruy_context();
+#  endif
 #endif
 
   }

src/devices.cc

@@ -4,6 +4,9 @@
 #  include "cuda/utils.h"
 #  include "cuda/random.h"
 #endif
+#ifdef CT2_WITH_RUY
+#  include "cpu/backend.h"
+#endif
 #ifdef CT2_WITH_TENSOR_PARALLEL
 #  include <unistd.h>
 #endif
@@ -125,9 +128,17 @@ namespace ctranslate2 {
       if (device == Device::CUDA) {
           cuda::free_curand_states();
       }
-#else
-      (void)device;
 #endif
+#if defined(CT2_WITH_RUY) && defined(_WIN32)
+      if (device == Device::CPU) {
+          // Windows only. Release this worker thread's ruy::Context here (a normal
+          // execution context) rather than at thread exit, where joining ruy's
+          // internal threads deadlocks ThreadPool shutdown (ctranslate2-rs#64).
+          // Other platforms rely on thread_local RAII — see cpu/backend.cc.
+          cpu::clear_ruy_context();
+      }
+#endif
+      (void)device;
   }
 
   // Initialize the static member variable

tools/ruy_shutdown_repro/CMakeLists.txt

@@ -0,0 +1,24 @@
+cmake_minimum_required(VERSION 3.15)
+project(ct2_ruy_shutdown_repro CXX)
+
+# Point this at a CTranslate2 source checkout (submodules initialized).
+#   cmake -DCT2_DIR=/path/to/CTranslate2 ...
+if(NOT DEFINED CT2_DIR)
+  message(FATAL_ERROR "Pass -DCT2_DIR=<path to CTranslate2 source checkout>")
+endif()
+
+# Ruy-only CPU build: this is the backend that exhibits the shutdown deadlock.
+set(WITH_RUY  ON  CACHE BOOL   "" FORCE)
+set(WITH_MKL  OFF CACHE BOOL   "" FORCE)
+set(WITH_DNNL OFF CACHE BOOL   "" FORCE)
+set(WITH_CUDA OFF CACHE BOOL   "" FORCE)
+set(BUILD_CLI OFF CACHE BOOL   "" FORCE)
+set(OPENMP_RUNTIME "COMP" CACHE STRING "" FORCE)
+
+# Build CTranslate2 as part of this project so the `ctranslate2` target carries all
+# transitive static deps (ruy, cpu_features, ...) automatically.
+add_subdirectory(${CT2_DIR} ctranslate2_build)
+
+add_executable(repro repro.cpp)
+target_link_libraries(repro PRIVATE ctranslate2)
+target_compile_features(repro PRIVATE cxx_std_17)

tools/ruy_shutdown_repro/README.md

@@ -0,0 +1,53 @@
+# Windows Ruy shutdown deadlock — minimal repro
+
+Self-contained repro for the shutdown hang fixed by
+[OpenNMT/CTranslate2#2076](https://github.com/OpenNMT/CTranslate2/pull/2076)
+(reported downstream as jkawamoto/ctranslate2-rs#64).
+
+`repro.cpp` builds a CPU int8 `Translator` (Ruy backend) with worker threads, runs one
+large translation batch, then destroys the `Translator`. On an unpatched build the
+per-thread `ruy::Context` destructor joins Ruy's thread pool from a worker thread that
+is exiting under the Windows loader lock, and that join deadlocks.
+
+## The one thing that matters: batch size
+
+The batch must be large enough that Ruy actually spawns its internal thread pool. The
+destructor only deadlocks when there are Ruy worker threads to join; a tiny batch runs
+single-threaded, has nothing to join, and shuts down cleanly even unpatched. That is
+easy to trip over when writing a test. This repro uses a 512-sentence batch.
+
+It is *not* specific to a CRT model or link mode. Measured on Windows 11, MSVC 14.44,
+x64, unpatched CTranslate2 `0d8bcd36`, Ruy int8, `intra_threads=4`, `inter_threads=2`,
+512-sentence batch:
+
+| Build of CTranslate2 + this repro | Result |
+|---|---|
+| static lib + static CRT (`/MT`)   | **hangs on shutdown** |
+| shared lib + dynamic CRT (`/MD`)  | **hangs on shutdown** |
+
+The `/MD` shared build is the configuration of the official wheels, so they are affected
+too. Rebuilding either with #2076 applied, both exit cleanly (`SURVIVED`).
+
+## Build & run
+
+Needs CMake, MSVC, and a CTranslate2 source checkout with submodules initialized
+(`git submodule update --init --recursive`). No CUDA / MKL / oneDNN required. `CT2_DIR`
+points at a CTranslate2 source tree; from here in the tree that is the repo root,
+`../..`. Run these from this directory (`tools/ruy_shutdown_repro`):
+
+```sh
+cmake -G "Visual Studio 17 2022" -A x64 \
+  -DCMAKE_POLICY_DEFAULT_CMP0091=NEW \
+  -DCT2_DIR=../.. \
+  -DBUILD_SHARED_LIBS=OFF -DCMAKE_MSVC_RUNTIME_LIBRARY=MultiThreaded \
+  -S . -B build
+cmake --build build --config Release --target repro
+build/Release/repro.exe ../../tests/data/models/v2/aren-transliteration
+```
+
+For the shared `/MD` variant use `-DBUILD_SHARED_LIBS=ON
+-DCMAKE_MSVC_RUNTIME_LIBRARY=MultiThreadedDLL`, and put the built `ctranslate2.dll`
+(under `build/ctranslate2_build/Release`) on `PATH` before running.
+
+Unpatched, the run prints `destroying Translator ...` and then hangs (kill it). With
+#2076 applied it prints `SURVIVED: clean shutdown, no deadlock`.

tools/ruy_shutdown_repro/repro.cpp

@@ -0,0 +1,50 @@
+// Minimal repro for the Windows Ruy shutdown deadlock (ctranslate2-rs#64 / #2076).
+//
+// It creates a CPU int8 Translator (Ruy backend) with worker threads, runs one
+// translation, then destroys the Translator. On an unpatched build the per-thread
+// ruy::Context destructor joins Ruy's internal thread pool from a worker thread
+// that is exiting under the Windows loader lock, and that join deadlocks.
+//
+// The point of this repro: it only hangs when CTranslate2 is built with the STATIC
+// CRT (/MT) and linked statically. Built with the dynamic CRT (/MD) as a shared
+// library -- the configuration of the official wheels -- the same code shuts down
+// cleanly. See README.md.
+#include <ctranslate2/translator.h>
+#include <ctranslate2/models/model.h>
+
+#include <iostream>
+
+int main(int argc, char** argv) {
+  if (argc < 2) {
+    std::cerr << "usage: repro <model_dir>\n";
+    return 2;
+  }
+
+  ctranslate2::ReplicaPoolConfig pool_config;
+  pool_config.num_threads_per_replica = 4;   // intra_threads: give Ruy a real thread pool
+
+  ctranslate2::models::ModelLoader model_loader(argv[1]);
+  model_loader.device = ctranslate2::Device::CPU;
+  model_loader.compute_type = ctranslate2::ComputeType::INT8;
+  model_loader.num_replicas_per_device = 2;  // inter_threads: use worker threads
+
+  // A large batch so the int8 GEMM is big enough that Ruy actually spawns its
+  // internal thread pool. That is a precondition for the hang: the destructor only
+  // deadlocks if there are Ruy worker threads to join. A tiny batch runs
+  // single-threaded and shuts down cleanly even unpatched.
+  const std::vector<std::string> sentence = {"آ", "ت", "ز", "م", "و", "ن"};
+  const std::vector<std::vector<std::string>> source(512, sentence);
+
+  std::cout << "translating a batch of " << source.size() << "..." << std::endl;
+  {
+    ctranslate2::Translator translator(model_loader, pool_config);
+    const auto results = translator.translate_batch(source);
+    std::cout << "output[0]:";
+    for (const auto& token : results[0].hypotheses[0])
+      std::cout << ' ' << token;
+    std::cout << "\ndestroying Translator (Ruy thread-pool join happens here)..."
+              << std::endl;
+  }  // <-- unpatched + static /MT: the join deadlocks here and the process hangs.
+  std::cout << "SURVIVED: clean shutdown, no deadlock" << std::endl;
+  return 0;
+}