fix: propagate custom reranker registry to spawned workers (#799)

TextCrossEncoder(custom).rerank_pairs(..., parallel=N) with input >= batch_size
spawned worker processes whose fresh interpreter had an empty registry, so the
custom model could not be resolved and the worker raised "Model not supported".
The embedding side already handled this; the reranker did not.

Mirror the embedding fix: OnnxCrossEncoderModel gains _extra_worker_params() and
the parallel pool path now does params.update(self._extra_worker_params());
CustomTextCrossEncoder gains _export_registry/_import_registry (idempotent) +
_extra_worker_params (carries custom_registry) + a CustomTextCrossEncoderWorker
that re-imports the registry before constructing the model in the worker.

Co-authored-by: n24q02m <n24q02m@outlook.com>

Author: n24q02m

qwen3_embed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -1,5 +1,8 @@
+from typing import Any
+
 from qwen3_embed.common.model_description import BaseModelDescription
 from qwen3_embed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from qwen3_embed.rerank.cross_encoder.onnx_text_model import TextRerankerWorker
 
 
 class CustomTextCrossEncoder(OnnxTextCrossEncoder):
@@ -16,3 +19,48 @@ def add_model(
     ) -> None:
         cls._clear_model_cache()
         cls.SUPPORTED_MODELS.append(model_description)
+
+    @classmethod
+    def _export_registry(cls) -> list[BaseModelDescription]:
+        # Snapshot the runtime registry so it can be pickled into spawned
+        # worker processes, which start with a fresh interpreter and an empty
+        # SUPPORTED_MODELS. BaseModelDescription is a frozen dataclass (picklable).
+        return list(cls.SUPPORTED_MODELS)
+
+    @classmethod
+    def _import_registry(cls, payload: list[BaseModelDescription]) -> None:
+        # Re-register custom models in a worker process, idempotently (same id
+        # imported twice must not create a duplicate entry).
+        existing = {m.model.lower() for m in cls.SUPPORTED_MODELS}
+        for desc in payload:
+            if desc.model.lower() not in existing:
+                cls.SUPPORTED_MODELS.append(desc)
+                existing.add(desc.model.lower())
+
+    @classmethod
+    def _get_worker_class(cls) -> type["CustomTextCrossEncoderWorker"]:
+        return CustomTextCrossEncoderWorker
+
+    def _extra_worker_params(self) -> dict[str, Any]:
+        # Propagate the runtime registry so spawned workers (fresh interpreters
+        # with an empty SUPPORTED_MODELS) can resolve + re-register this custom
+        # model instead of raising "Model ... not supported".
+        return {"custom_registry": self._export_registry()}
+
+
+class CustomTextCrossEncoderWorker(TextRerankerWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> CustomTextCrossEncoder:
+        registry = kwargs.pop("custom_registry", None)
+        if registry is not None:
+            CustomTextCrossEncoder._import_registry(registry)
+        return CustomTextCrossEncoder(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

qwen3_embed/rerank/cross_encoder/onnx_text_model.py

@@ -24,6 +24,14 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
     def _get_worker_class(cls) -> type["TextRerankerWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
+    def _extra_worker_params(self) -> dict[str, Any]:
+        """Extra kwargs injected into each spawned worker's init.
+
+        Subclasses override this to carry state (e.g. a runtime-registered custom
+        model registry) into worker processes that start with a fresh interpreter.
+        """
+        return {}
+
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
         assert self.tokenizer is not None
         return self.tokenizer.encode_batch(pairs)
@@ -110,6 +118,11 @@ def _rerank_pairs(
             if extra_session_options is not None:
                 params.update(extra_session_options)
 
+            # Carry per-instance worker state (e.g. a runtime custom-model
+            # registry) into the spawned workers, which start with a fresh
+            # interpreter + empty registry. Mirrors the embedding pool path.
+            params.update(self._extra_worker_params())
+
             pool = ParallelWorkerPool(
                 worker=self._get_worker_class(),
                 config=PoolConfig(

tests/test_custom_reranker_registry.py

@@ -0,0 +1,97 @@
+"""F5: custom reranker registry must propagate to spawned multiprocessing workers.
+
+The embedding side already does this (CustomTextEmbedding._export_registry /
+_import_registry / _extra_worker_params + the parallel pool params.update). The
+reranker side did not, so TextCrossEncoder(custom).rerank_pairs(..., parallel=N)
+with input >= batch_size raised "Model ... not supported" in the spawned worker
+(fresh interpreter with an empty registry). These tests guard the mirror fix.
+"""
+
+import sys
+
+import pytest
+
+from qwen3_embed.common.model_description import BaseModelDescription, ModelSource
+from qwen3_embed.rerank.cross_encoder.custom_text_cross_encoder import (
+    CustomTextCrossEncoder,
+    CustomTextCrossEncoderWorker,
+)
+
+
+def _clear() -> None:
+    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()
+
+
+def test_custom_reranker_registry_survives_serialization():
+    _clear()
+    CustomTextCrossEncoder.add_model(
+        BaseModelDescription(model="Org/My-Reranker", sources=ModelSource(hf="Org/My-Reranker"))
+    )
+    payload = CustomTextCrossEncoder._export_registry()
+    _clear()
+    assert CustomTextCrossEncoder.SUPPORTED_MODELS == []  # fresh-worker simulation
+    CustomTextCrossEncoder._import_registry(payload)
+    try:
+        models = [m.model for m in CustomTextCrossEncoder._list_supported_models()]
+        assert "Org/My-Reranker" in models
+    finally:
+        _clear()
+
+
+def test_custom_reranker_import_registry_is_idempotent():
+    _clear()
+    desc = BaseModelDescription(model="Org/Dup", sources=ModelSource(hf="Org/Dup"))
+    CustomTextCrossEncoder.add_model(desc)
+    payload = CustomTextCrossEncoder._export_registry()
+    CustomTextCrossEncoder._import_registry(payload)  # re-import same payload
+    try:
+        ids = [m.model for m in CustomTextCrossEncoder._list_supported_models()]
+        assert ids.count("Org/Dup") == 1  # no duplicate entry
+    finally:
+        _clear()
+
+
+def test_custom_reranker_extra_worker_params_carries_registry():
+    _clear()
+    CustomTextCrossEncoder.add_model(
+        BaseModelDescription(model="Org/R", sources=ModelSource(hf="Org/R"))
+    )
+    try:
+        enc = CustomTextCrossEncoder.__new__(CustomTextCrossEncoder)  # skip ONNX load
+        params = enc._extra_worker_params()
+        assert "custom_registry" in params
+        assert any(d.model == "Org/R" for d in params["custom_registry"])
+    finally:
+        _clear()
+
+
+def test_custom_reranker_uses_custom_worker_class():
+    # The base OnnxTextCrossEncoder worker would construct a plain (non-custom)
+    # cross-encoder + never import the registry. The custom subclass must use a
+    # worker that re-registers the custom model in the spawned process.
+    assert CustomTextCrossEncoder._get_worker_class() is CustomTextCrossEncoderWorker
+
+
+@pytest.mark.integration
+@pytest.mark.skipif(sys.platform == "win32", reason="multiprocessing spawn deadlock on Windows")
+def test_custom_reranker_parallel_resolves_in_workers():
+    """Register a custom reranker, rerank_pairs with parallel=2 + batch_size=1
+    (forces the worker-pool path), assert no 'not supported' error + scores."""
+    _clear()
+    model_id = "Org/Custom-Reranker"
+    CustomTextCrossEncoder.add_model(
+        BaseModelDescription(
+            model=model_id,
+            sources=ModelSource(hf="n24q02m/Qwen3-Reranker-0.6B-ONNX"),
+            model_file="onnx/model_quantized.onnx",
+        )
+    )
+    try:
+        from qwen3_embed.rerank.cross_encoder.text_cross_encoder import TextCrossEncoder
+
+        ce = TextCrossEncoder(model_name=model_id)
+        pairs = [("q", "d1"), ("q", "d2"), ("q", "d3")]
+        scores = list(ce.rerank_pairs(pairs, batch_size=1, parallel=2))
+        assert len(scores) == 3
+    finally:
+        _clear()