Skip to content

Commit 20892ee

Browse files
committed
Add executor affinity for numpy/torch thread safety
Workers are now assigned a fixed executor_id at creation. All calls from the same worker go to the same executor thread, preventing thread state corruption in libraries like numpy and PyTorch that have thread-local state. This fixes segfaults when using sentence-transformers or other ML libraries that depend on consistent thread affinity.
1 parent 3a34236 commit 20892ee

File tree

3 files changed

+18
-2
lines changed

3 files changed

+18
-2
lines changed

c_src/py_exec.c

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -791,8 +791,16 @@ static int executor_enqueue(py_request_t *req) {
791791

792792
case PY_MODE_MULTI_EXECUTOR:
793793
if (atomic_load(&g_multi_executor_initialized)) {
794-
/* Route to multi-executor pool */
795-
int exec_id = select_executor();
794+
/* Route to multi-executor pool.
795+
* Use worker's assigned executor for thread affinity if available.
796+
* This ensures libraries like numpy/torch that have thread-local
797+
* state always run on the same thread for a given worker. */
798+
int exec_id;
799+
if (req->worker != NULL && req->worker->executor_id >= 0) {
800+
exec_id = req->worker->executor_id % g_num_executors;
801+
} else {
802+
exec_id = select_executor();
803+
}
796804
multi_executor_enqueue(exec_id, req);
797805
return 0;
798806
}

c_src/py_nif.c

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1418,6 +1418,11 @@ static ERL_NIF_TERM nif_worker_new(ErlNifEnv *env, int argc, const ERL_NIF_TERM
14181418
worker->has_callback_handler = false;
14191419
worker->callback_env = NULL;
14201420

1421+
/* Assign executor affinity for thread-safe library support (numpy, torch).
1422+
* Each worker gets a fixed executor to ensure all calls from the same
1423+
* worker go to the same thread, preventing thread state corruption. */
1424+
worker->executor_id = select_executor();
1425+
14211426
PyGILState_Release(gstate);
14221427

14231428
ERL_NIF_TERM result = enif_make_resource(env, worker);

c_src/py_nif.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,9 @@ typedef struct {
390390

391391
/** @brief Environment for building callback messages */
392392
ErlNifEnv *callback_env;
393+
394+
/** @brief Assigned executor ID for thread affinity (-1 = round-robin) */
395+
int executor_id;
393396
} py_worker_t;
394397

395398
/* async_pending_t and py_async_worker_t removed - async workers replaced by event loop model */

0 commit comments

Comments
 (0)