diff --git a/packages/leann-core/src/leann/chat.py b/packages/leann-core/src/leann/chat.py index a36383f5..2e7fca1f 100644 --- a/packages/leann-core/src/leann/chat.py +++ b/packages/leann-core/src/leann/chat.py @@ -8,6 +8,7 @@ import logging import os from abc import ABC, abstractmethod +from collections.abc import Callable from typing import Any, Optional, cast from .settings import ( @@ -27,6 +28,36 @@ logger = logging.getLogger(__name__) +def _run_with_optional_posix_alarm( + operation: Callable[[], Any], timeout_seconds: int, timeout_message: str +) -> Any: + """Run an operation with a POSIX alarm where the platform supports it.""" + try: + import signal + except ImportError: + return operation() + + if not hasattr(signal, "SIGALRM") or not hasattr(signal, "alarm"): + return operation() + + def timeout_handler(signum, frame): + raise TimeoutError(timeout_message) + + try: + old_handler = signal.signal(signal.SIGALRM, timeout_handler) + except ValueError: + # signal.signal only works from the main thread. Loading without an + # alarm is preferable to failing before model loading starts. + return operation() + + signal.alarm(timeout_seconds) + try: + return operation() + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + def check_ollama_models(host: str) -> list[str]: """Check available Ollama models and return a list""" try: @@ -606,20 +637,12 @@ def __init__( self.device = "cpu" logger.info("No GPU detected. Using CPU.") - # Load tokenizer and model with timeout protection + # Load tokenizer and model with timeout protection when POSIX alarms are available. try: - import signal - - def timeout_handler(signum, frame): - raise TimeoutError("Model download/loading timed out") - - # Set timeout for model loading (60 seconds) - old_handler = signal.signal(signal.SIGALRM, timeout_handler) - signal.alarm(60) - try: + def load_model_assets(): logger.info(f"Loading tokenizer for {model_name}...") - self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( model_name, trust_remote_code=self.trust_remote_code ) @@ -634,16 +657,20 @@ def timeout_handler(signum, frame): # Auto mode: let HuggingFace distribute across available GPUs device_map = "auto" - self.model = AutoModelForCausalLM.from_pretrained( + model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float16 if self.device != "cpu" else torch.float32, device_map=device_map, trust_remote_code=self.trust_remote_code, ) logger.info(f"Successfully loaded {model_name}") - finally: - signal.alarm(0) # Cancel the alarm - signal.signal(signal.SIGALRM, old_handler) # Restore old handler + return tokenizer, model + + self.tokenizer, self.model = _run_with_optional_posix_alarm( + load_model_assets, + timeout_seconds=60, + timeout_message="Model download/loading timed out", + ) except TimeoutError: logger.error(f"Model loading timed out for {model_name}") diff --git a/packages/leann-core/src/leann/embedding_compute.py b/packages/leann-core/src/leann/embedding_compute.py index 7e51d4fb..aac11a56 100644 --- a/packages/leann-core/src/leann/embedding_compute.py +++ b/packages/leann-core/src/leann/embedding_compute.py @@ -278,7 +278,7 @@ def _query_lmstudio_context_limit(model_name: str, base_url: str) -> Optional[in # Append to existing NODE_PATH if present existing_node_path = env.get("NODE_PATH", "") env["NODE_PATH"] = ( - f"{global_modules}:{existing_node_path}" + os.pathsep.join([global_modules, existing_node_path]) if existing_node_path else global_modules ) @@ -495,7 +495,8 @@ def compute_embeddings_sentence_transformers( # TODO: Haven't tested this yet torch.set_num_threads(min(8, os.cpu_count() or 4)) try: - torch.backends.mkldnn.enabled = True + mkldnn_backend = cast(Any, torch.backends.mkldnn) + mkldnn_backend.enabled = True except AttributeError: pass diff --git a/tests/test_windows_native_polish.py b/tests/test_windows_native_polish.py new file mode 100644 index 00000000..53a91be2 --- /dev/null +++ b/tests/test_windows_native_polish.py @@ -0,0 +1,54 @@ +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +from leann import embedding_compute +from leann.chat import _run_with_optional_posix_alarm +from leann.embedding_compute import _query_lmstudio_context_limit + + +def test_hf_loader_timeout_helper_runs_without_sigalrm(monkeypatch): + monkeypatch.setitem(sys.modules, "signal", SimpleNamespace()) + + assert _run_with_optional_posix_alarm(lambda: "loaded", 60, "timeout") == "loaded" + + +def test_hf_loader_timeout_helper_runs_when_signal_registration_is_unavailable(monkeypatch): + fake_signal = SimpleNamespace( + SIGALRM=14, + alarm=Mock(side_effect=AssertionError("alarm should not be called")), + signal=Mock(side_effect=ValueError("signal only works in main thread")), + ) + monkeypatch.setitem(sys.modules, "signal", fake_signal) + + assert _run_with_optional_posix_alarm(lambda: "loaded", 60, "timeout") == "loaded" + + +def test_lmstudio_node_path_uses_platform_separator(monkeypatch): + captured_node_path = None + + def mock_run(cmd, **kwargs): + nonlocal captured_node_path + if cmd == ["npm", "root", "-g"]: + result = Mock() + result.returncode = 0 + result.stdout = r"C:\npm\node_modules" + "\n" + result.stderr = "" + return result + + assert cmd[0] == "node" + captured_node_path = kwargs["env"]["NODE_PATH"] + result = Mock() + result.returncode = 0 + result.stdout = '{"contextLength": 8192, "identifier": "custom-model"}' + result.stderr = "" + return result + + monkeypatch.setattr(embedding_compute.subprocess, "run", mock_run) + monkeypatch.setattr(embedding_compute.os, "pathsep", ";") + monkeypatch.setenv("NODE_PATH", r"C:\existing;D:\more") + + limit = _query_lmstudio_context_limit(model_name="custom-model", base_url="ws://localhost:1234") + + assert limit == 8192 + assert captured_node_path == r"C:\npm\node_modules;C:\existing;D:\more" diff --git a/uv.lock b/uv.lock index a32f0382..cd7fa7e9 100644 --- a/uv.lock +++ b/uv.lock @@ -2005,6 +2005,30 @@ requires-dist = [ ] provides-extras = ["query-server"] +[[package]] +name = "leann-backend-flashlib-ivf" +version = "0.3.6" +source = { editable = "packages/leann-backend-flashlib-ivf" } +dependencies = [ + { name = "flashlib" }, + { name = "leann-core" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "torch" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashlib" }, + { name = "leann-backend-hnsw", marker = "extra == 'query-server'", specifier = ">=0.3.6" }, + { name = "leann-core", specifier = ">=0.3.6" }, + { name = "msgpack", marker = "extra == 'query-server'", specifier = ">=1.0.0" }, + { name = "numpy", specifier = ">=1.20.0" }, + { name = "pyzmq", marker = "extra == 'query-server'", specifier = ">=23.0.0" }, + { name = "torch" }, +] +provides-extras = ["query-server"] + [[package]] name = "leann-backend-hnsw" version = "0.3.7" @@ -2158,6 +2182,9 @@ documents = [ flashlib = [ { name = "leann-backend-flashlib" }, ] +flashlib-ivf = [ + { name = "leann-backend-flashlib-ivf" }, +] [package.dev-dependencies] dev = [ @@ -2190,6 +2217,7 @@ requires-dist = [ { name = "ipykernel", specifier = "==6.29.5" }, { name = "leann-backend-diskann", marker = "extra == 'diskann'", editable = "packages/leann-backend-diskann" }, { name = "leann-backend-flashlib", marker = "extra == 'flashlib'", editable = "packages/leann-backend-flashlib" }, + { name = "leann-backend-flashlib-ivf", marker = "extra == 'flashlib-ivf'", editable = "packages/leann-backend-flashlib-ivf" }, { name = "leann-backend-hnsw", editable = "packages/leann-backend-hnsw" }, { name = "leann-core", editable = "packages/leann-core" }, { name = "llama-index", specifier = ">=0.12.44" }, @@ -2230,7 +2258,7 @@ requires-dist = [ { name = "tree-sitter-typescript", specifier = ">=0.20.0" }, { name = "typer", specifier = ">=0.12.3" }, ] -provides-extras = ["diskann", "flashlib", "documents"] +provides-extras = ["diskann", "flashlib", "flashlib-ivf", "documents"] [package.metadata.requires-dev] dev = [