fix: eliminate Linux Docker crash and restore MLX on native macOS

Dockerfile

@@ -36,6 +36,7 @@ RUN uv venv && source .venv/bin/activate && uv pip install -e .
 ENV PATH="/app/.venv/bin:$PATH" \
     PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
+    RERANKER_BACKEND_TYPE=pytorch \
     RERANKER_HOST=0.0.0.0 \
     RERANKER_PORT=8010
 

README.md

@@ -10,8 +10,8 @@ This project provides a FastAPI-based web service that implements a reranking AP
 
 *   **Jina Compatible API**: Implements `/v1/rerank` endpoint structure
 *   **Local Hosting**: Run reranker model entirely on your own infrastructure
-*   **Multiple Backends**: Supports both PyTorch and MLX backends for optimal performance
-*   **Apple Silicon Optimization**: MLX backend provides optimized performance for M1/M2/M3 chips
+*   **Multiple Backends**: Supports PyTorch everywhere and MLX on native Apple Silicon
+*   **Apple Silicon Optimization**: MLX backend is automatically installed on macOS arm64 (M1/M2/M3)
 *   **MLX Fallback Reranker**: Automatically wraps MLX-converted Hugging Face models that do not ship a `rerank.py` helper
 *   **Sentence Transformers**: Uses powerful `sentence-transformers` library for PyTorch backend
 *   **Configurable Model**: Easily switch between different reranker models and backends
@@ -29,10 +29,11 @@ This project provides a FastAPI-based web service that implements a reranking AP
 **PyTorch Backend:**
 *   PyTorch 2.5+ (automatically installed)
 *   CUDA/MPS support for GPU acceleration (optional)
+*   Recommended backend for Linux and Docker deployments
 
 **MLX Backend (Apple Silicon only):**
 *   Apple Silicon (M1/M2/M3) Mac
-*   MLX and MLX-LM libraries (automatically installed)
+*   MLX and MLX-LM are installed automatically on native macOS arm64
 *   Optimized for memory efficiency and performance on Apple chips
 
 ## Installation
@@ -50,6 +51,10 @@ source .venv/bin/activate
 uv pip install -e ".[dev]"
 ```
 
+On native macOS Apple Silicon, MLX dependencies are installed automatically.
+
+On Linux, including Docker, the default backend is `pytorch`. MLX is not supported there because Docker containers run a Linux userspace/kernel interface even on Apple hosts, and MLX requires Apple-native runtime libraries.
+
 ## Usage
 
 The CLI supports both modern subcommands and legacy arguments for backward compatibility.
@@ -74,7 +79,16 @@ cli --backend <backend_type> --model <model> --host <host> --port <port>
 ### Available Backends
 
 *   `pytorch`: PyTorch-based reranker (default, cross-platform)
-*   `mlx`: MLX-based reranker (Apple Silicon optimized)
+*   `mlx`: MLX-based reranker (native macOS Apple Silicon only)
+
+### Backend Support By Platform
+
+| Platform | Default backend | MLX support |
+| --- | --- | --- |
+| Linux / Docker | `pytorch` | Not supported |
+| macOS Intel | `pytorch` | Not supported |
+| macOS Apple Silicon (native) | `pytorch` | Auto-installed and supported |
+| Docker Desktop on Apple Silicon | `pytorch` | Not supported (still Linux container runtime) |
 
 ### Command Options
 
@@ -99,6 +113,48 @@ cli serve --backend pytorch --model jinaai/jina-reranker-v2-base-multilingual
 cli serve --backend mlx --model jinaai/jina-reranker-v3-mlx
 ```
 
+If you explicitly select `--backend mlx` on Linux or non-Apple hardware, startup fails fast with an actionable error telling you to switch back to `pytorch`.
+
+### Backend Selection Via Environment
+
+```bash
+# Default for Docker/Linux
+export RERANKER_BACKEND_TYPE=pytorch
+
+# Native macOS Apple Silicon (MLX auto-installed)
+export RERANKER_BACKEND_TYPE=mlx
+```
+
+Command-line flags still override the defaults:
+
+```bash
+cli serve --backend pytorch
+cli serve --backend mlx
+```
+
+### Docker
+
+Build and run the container with the default PyTorch backend:
+
+```bash
+docker build -t local-reranker .
+docker run --rm -p 8010:8010 \
+    -e RERANKER_BACKEND_TYPE=pytorch \
+    local-reranker
+```
+
+Minimal `docker compose` example:
+
+```yaml
+services:
+    local-reranker:
+        build: .
+        ports:
+            - "8010:8010"
+        environment:
+            RERANKER_BACKEND_TYPE: pytorch
+```
+
 **Development Mode:**
 
 ```bash
@@ -211,8 +267,8 @@ uv run ruff check && uv run mypy src/
 # Ensure you're on Apple Silicon
 uname -m  # Should show arm64
 
-# Install MLX dependencies
-uv add mlx mlx-lm safetensors
+# Reinstall project dependencies (MLX auto-installs on native macOS arm64)
+uv pip install -e ".[dev]"
 ```
 
 **Model download fails:**
@@ -271,8 +327,8 @@ top -o mem | grep python
 The application uses pydantic-settings for configuration management. You can set the following environment variables to override defaults:
 
 ```bash
-# Force MLX backend
-export RERANKER_RERANKER_TYPE=mlx
+# Force backend selection
+export RERANKER_BACKEND_TYPE=mlx
 
 # Custom model name
 export RERANKER_MODEL_NAME=custom-mlx-model
@@ -290,6 +346,9 @@ export RERANKER_RELOAD=true
 
 **Note**: Using the CLI command line options is recommended over environment variables for clarity.
 
+On Linux containers, keep `RERANKER_BACKEND_TYPE=pytorch`. Selecting `mlx` there will fail because Docker/Linux cannot provide the Apple MLX shared libraries required by `mlx.core`.
+On Docker Desktop running on Apple Silicon, containers are still Linux runtime environments, so MLX is still unsupported there.
+
 ## Project Structure
 
 ```

pyproject.toml

@@ -29,12 +29,12 @@ dependencies = [
     "pydantic-settings>=2.6.0",
     "sentence-transformers>=5.0.0",
     "torch>=2.5.0",
-    "transformers>=4.57.0",
+    "transformers>=4.57.0,<5",
     "tokenizers>=0.21.0",
     "accelerate>=1.2.0",
     "einops>=0.8.0",
-    "mlx>=0.30.0",
-    "mlx-lm>=0.28.0",
+    "mlx>=0.30.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
+    "mlx-lm>=0.28.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
     "safetensors>=0.5.0",
     "psutil>=7.0.0",
 ] 

src/local_reranker/__init__.py

@@ -6,7 +6,21 @@
 
 __version__ = "0.0.1"  # Placeholder version
 
-from .cli import main
-from .jina_mlx_reranker import JinaMLXReranker
-
 __all__ = ["main", "JinaMLXReranker"]
+
+
+def __getattr__(name: str):
+    """Lazily expose package entrypoints and optional MLX symbols."""
+    if name == "main":
+        from .cli import main
+
+        return main
+
+    if name == "JinaMLXReranker":
+        # Keep MLX optional so importing the package on Linux does not try to
+        # load Apple-only shared libraries during process startup.
+        from .jina_mlx_reranker import JinaMLXReranker
+
+        return JinaMLXReranker
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/local_reranker/api.py

@@ -11,8 +11,7 @@
 from .models import RerankRequest, RerankResponse
 from .reranker import Reranker as RerankerProtocol
 from .reranker_pytorch import Reranker as PyTorchReranker
-from .reranker_mlx import Reranker as MLXReranker
-from .config import Settings, get_effective_model_name
+from .config import Settings, get_effective_model_name, validate_backend_type
 
 # --- Logging Setup ---
 logger = logging.getLogger(__name__)
@@ -34,12 +33,16 @@ async def lifespan(app: FastAPI):
         logger.info(f"Loading reranker type: {settings.backend_type}")
         logger.info(f"Loading model: {model_name}")
 
+        validate_backend_type(settings.backend_type)
+
         # Initialize reranker based on type
         if settings.backend_type == "pytorch":
             reranker_instance = PyTorchReranker(
                 model_name=model_name, disable_batching=settings.disable_batching
             )
         elif settings.backend_type == "mlx":
+            from .reranker_mlx import Reranker as MLXReranker
+
             reranker_instance = MLXReranker(
                 model_name=model_name, disable_batching=settings.disable_batching
             )

src/local_reranker/cli.py

@@ -7,7 +7,12 @@
 import sys
 
 import uvicorn
-from .config import Settings, get_available_backends, get_effective_model_name
+from .config import (
+    Settings,
+    get_available_backends,
+    get_effective_model_name,
+    validate_backend_type,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -188,6 +193,11 @@ def main() -> None:
             disable_batching=args.disable_batching,
         )
 
+        try:
+            validate_backend_type(settings.backend_type)
+        except ValueError as exc:
+            parser.error(str(exc))
+
         logger.info(
             f"Starting Local Reranker API server on {settings.host}:{settings.port}"
         )

src/local_reranker/config.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 """Configuration management for the local reranker service."""
 
+import platform
 from typing import Dict, Optional
+
 from pydantic_settings import BaseSettings
 
 
@@ -56,3 +58,40 @@ def get_available_backends() -> Dict[str, str]:
         # "onnx": "ONNX-based reranker for CPU optimization",
         # "tensorflow": "TensorFlow-based reranker",
     }
+
+
+def get_mlx_platform_error() -> Optional[str]:
+    """Return an actionable MLX platform error, if any."""
+    system = platform.system()
+    machine = platform.machine().lower()
+
+    if system != "Darwin":
+        return (
+            "MLX is only supported on macOS Apple Silicon. "
+            f"Detected {system} on {machine or 'unknown'}, and Linux containers "
+            "cannot load the Apple MLX runtime libraries."
+        )
+
+    if machine not in {"arm64", "aarch64"}:
+        return (
+            "MLX requires Apple Silicon arm64 hardware. "
+            f"Detected macOS on {machine or 'unknown'}."
+        )
+
+    return None
+
+
+def validate_backend_type(backend_type: str) -> None:
+    """Validate that the requested backend can run on this platform."""
+    if backend_type != "mlx":
+        return
+
+    platform_error = get_mlx_platform_error()
+    if platform_error is None:
+        return
+
+    raise ValueError(
+        f"{platform_error} Use '--backend pytorch' or set "
+        "RERANKER_BACKEND_TYPE=pytorch. Install the optional MLX extras only "
+        "on macOS Apple Silicon with `pip install 'local-reranker[mlx]'`."
+    )

src/local_reranker/reranker_mlx.py

@@ -8,7 +8,7 @@
 from .models import RerankRequest, RerankResult, RerankDocument
 from .batch_manager import BatchManager
 from .batch_processor import BatchProcessor, DocumentTextExtractor
-from .jina_mlx_reranker import JinaMLXReranker
+from .config import get_mlx_platform_error
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +36,14 @@ def __init__(
             ImportError: If MLX dependencies are not installed.
             RuntimeError: If model loading fails.
         """
+        platform_error = get_mlx_platform_error()
+        if platform_error is not None:
+            raise RuntimeError(
+                f"{platform_error} Use '--backend pytorch' or set "
+                "RERANKER_BACKEND_TYPE=pytorch. Install `local-reranker[mlx]` "
+                "only on macOS Apple Silicon."
+            )
+
         self.model_name = model_name
         self.device = device
 
@@ -48,9 +56,12 @@ def __init__(
             self.model = self._load_mlx_reranker(model_path)
             logger.info(f"Successfully loaded MLX reranker with batching: {model_name}")
 
-        except ImportError as e:
+        except (ImportError, OSError) as e:
             raise ImportError(
-                "MLX dependencies not found. Install with: pip install mlx mlx-lm safetensors"
+                "MLX backend requested, but the optional MLX dependencies could "
+                "not be imported. This is expected on Linux containers because "
+                "Apple's MLX runtime is unavailable there. Use '--backend pytorch' "
+                "or install `local-reranker[mlx]` on macOS Apple Silicon."
             ) from e
         except Exception as e:
             raise RuntimeError(f"Failed to load MLX model '{model_name}': {e}") from e
@@ -85,6 +96,16 @@ def _load_mlx_reranker(self, model_path: str):
         Raises:
             RuntimeError: If model or projector loading fails.
         """
+        try:
+            from .jina_mlx_reranker import JinaMLXReranker
+        except (ImportError, OSError) as exc:
+            raise ImportError(
+                "MLX backend requested, but the MLX runtime could not be loaded. "
+                "On Docker/Linux this usually means the platform does not support "
+                "MLX shared libraries such as libmlx.so. Use '--backend pytorch' "
+                "or install `local-reranker[mlx]` on macOS Apple Silicon."
+            ) from exc
+
         projector_path = f"{model_path}/projector.safetensors"
 
         return JinaMLXReranker(

src/local_reranker/reranker_pytorch.py

@@ -2,6 +2,7 @@
 """PyTorch implementation of the reranker protocol."""
 
 import logging
+import warnings
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -12,6 +13,9 @@
 from .batch_manager import BatchManager
 from .batch_processor import BatchProcessor, DocumentTextExtractor
 
+# Suppress torch_dtype deprecation warning from Jina model files
+warnings.filterwarnings("ignore", message=".*torch_dtype.*is deprecated.*Use `dtype` instead")
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL_NAME = "jinaai/jina-reranker-v2-base-multilingual"

tests/test_cli.py

@@ -213,6 +213,26 @@ def test_main_invalid_port(self, mock_run_server):
         assert exc_info.value.code == 2  # argparse error code
         mock_run_server.assert_not_called()
 
+    @patch("local_reranker.cli.validate_backend_type")
+    @patch("local_reranker.cli.run_server")
+    @patch("sys.argv", ["local-reranker", "--backend", "mlx"])
+    def test_main_unsupported_mlx_backend(
+        self, mock_run_server, mock_validate_backend_type
+    ):
+        """Test main function with unsupported MLX backend."""
+        mock_validate_backend_type.side_effect = ValueError("MLX unavailable")
+
+        with pytest.raises(SystemExit) as exc_info:
+            main()
+
+        assert exc_info.value.code == 2
+        mock_run_server.assert_not_called()
+
+
+def test_package_import_keeps_mlx_lazy():
+    """Importing the package should not eagerly import MLX modules."""
+    assert "local_reranker.jina_mlx_reranker" not in sys.modules
+
 
 class TestArgumentParser:
     """Test cases for argument parsing logic."""