fix(inference): fix concurrent request Rust borrow checker panic

src/ogx/providers/inline/inference/transformers/transformers.py

@@ -39,6 +39,7 @@
 RERANKER_MODELS: dict[str, tuple] = {}
 
 RERANKER_MODELS_LOCK: asyncio.Lock = asyncio.Lock()
+TOKENIZER_LOCK: threading.Lock = threading.Lock()
 
 DEFAULT_RERANKER_INSTRUCTION = "Given the search query, retrieve relevant passages that answer the query"
 
@@ -191,27 +192,30 @@ def compute_reranked_scores(
         """
         # Reranker configuration
         max_length = 8192
-        token_true_id = reranker_tokenizer.convert_tokens_to_ids("yes")
-        token_false_id = reranker_tokenizer.convert_tokens_to_ids("no")
-
-        prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
-        suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
-        prefix_tokens = reranker_tokenizer.encode(prefix, add_special_tokens=False)
-        suffix_tokens = reranker_tokenizer.encode(suffix, add_special_tokens=False)
-
-        # Tokenize pairs
-        inputs = reranker_tokenizer(
-            pairs,
-            padding=False,
-            truncation="longest_first",
-            return_attention_mask=False,
-            max_length=max_length - len(prefix_tokens) - len(suffix_tokens),
-        )
-
-        for i, tokens in enumerate(inputs["input_ids"]):
-            inputs["input_ids"][i] = prefix_tokens + tokens + suffix_tokens
-
-        inputs = reranker_tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_length)
+        # We lock everything that touches reranker_tokenizer because it modifies
+        # its internal Rust state during these operations.
+        with TOKENIZER_LOCK:
+            token_true_id = reranker_tokenizer.convert_tokens_to_ids("yes")
+            token_false_id = reranker_tokenizer.convert_tokens_to_ids("no")
+
+            prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
+            suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+            prefix_tokens = reranker_tokenizer.encode(prefix, add_special_tokens=False)
+            suffix_tokens = reranker_tokenizer.encode(suffix, add_special_tokens=False)
+
+            # Tokenize pairs
+            inputs = reranker_tokenizer(
+                pairs,
+                padding=False,
+                truncation="longest_first",
+                return_attention_mask=False,
+                max_length=max_length - len(prefix_tokens) - len(suffix_tokens),
+            )
+
+            for i, tokens in enumerate(inputs["input_ids"]):
+                inputs["input_ids"][i] = prefix_tokens + tokens + suffix_tokens
+
+            inputs = reranker_tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_length)
 
         for key in inputs:
             inputs[key] = inputs[key].to(reranker_model.device)

tests/integration/inference/test_rerank.py

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import concurrent.futures
 
 import pytest
 from ogx_client import BadRequestError as OGXBadRequestError
@@ -104,6 +105,34 @@ def test_rerank_text(client_with_models, rerank_model_id, query, items, inferenc
     _validate_rerank_response(response, items)
 
 
+def test_rerank_text_parallel(client_with_models, rerank_model_id, inference_provider_type):
+    skip_if_provider_doesnt_support_rerank(inference_provider_type)
+    query_items = [
+        (DUMMY_STRING, [DUMMY_STRING, DUMMY_STRING2]),
+        (DUMMY_TEXT, [DUMMY_TEXT, DUMMY_TEXT2]),
+        (DUMMY_STRING, [DUMMY_STRING2, DUMMY_TEXT]),
+        (DUMMY_TEXT, [DUMMY_STRING, DUMMY_TEXT2]),
+    ]
+    future_to_items = {}
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=len(query_items)) as executor:
+        for query, items in query_items:
+            future = executor.submit(
+                client_with_models.alpha.inference.rerank, model=rerank_model_id, query=query, items=items
+            )
+            future_to_items[future] = items
+
+        for future in concurrent.futures.as_completed(future_to_items):
+            original_items = future_to_items[future]
+
+            # .result() will raise the 500 Error as RuntimeError here if the Rust borrower bug is present
+            response = future.result()
+
+            assert isinstance(response, list)
+            assert len(response) <= len(original_items)
+            _validate_rerank_response(response, original_items)
+
+
 @pytest.mark.parametrize(
     "query,items",
     [