TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: logs

Browse Source
This commit is contained in:
Andy Lee
2025-07-21 22:45:24 -07:00
parent f7af6805fa
commit 573313f0b6
6 changed files with 87 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
packages/leann-core/src/leann/api.py
View File

@@ -113,7 +113,7 @@ class PassageManager:
for source in passage_sources:
assert source["type"] == "jsonl", "only jsonl is supported"
passage_file = source["path"]
index_file = source["index_path"]
index_file = source["index_path"] # .idx file
if not Path(index_file).exists():
raise FileNotFoundError(f"Passage index file not found: {index_file}")
with open(index_file, "rb") as f:

13
packages/leann-core/src/leann/embedding_compute.py
View File

@@ -6,7 +6,7 @@ Preserves all optimization parameters to ensure performance
import numpy as np
import torch
from typing import List, Dict, Any, Optional
from typing import List, Dict, Any
import logging
logger = logging.getLogger(__name__)
@@ -16,7 +16,10 @@ _model_cache: Dict[str, Any] = {}
def compute_embeddings(
texts: List[str], model_name: str, mode: str = "sentence-transformers",is_build: bool = False
texts: List[str],
model_name: str,
mode: str = "sentence-transformers",
is_build: bool = False,
) -> np.ndarray:
"""
Unified embedding computation entry point
@@ -30,7 +33,9 @@ def compute_embeddings(
Normalized embeddings array, shape: (len(texts), embedding_dim)
"""
if mode == "sentence-transformers":
return compute_embeddings_sentence_transformers(texts, model_name, is_build=is_build)
return compute_embeddings_sentence_transformers(
texts, model_name, is_build=is_build
)
elif mode == "openai":
return compute_embeddings_openai(texts, model_name)
elif mode == "mlx":
@@ -65,7 +70,7 @@ def compute_embeddings_sentence_transformers(
# Create cache key
cache_key = f"sentence_transformers_{model_name}_{device}_{use_fp16}"
# Check if model is already cached
if cache_key in _model_cache:
print(f"INFO: Using cached model: {model_name}")

108
packages/leann-core/src/leann/embedding_server_manager.py
View File

@@ -1,14 +1,22 @@
import threading
import time
import atexit
import socket
import subprocess
import sys
import os
import logging
from pathlib import Path
from typing import Optional
import select
import psutil
# Set up logging based on environment variable
LOG_LEVEL = os.getenv("LEANN_LOG_LEVEL", "WARNING").upper()
logging.basicConfig(
level=getattr(logging, LOG_LEVEL, logging.INFO),
format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
def _check_port(port: int) -> bool:
"""Check if a port is in use"""
@@ -36,11 +44,11 @@ def _check_process_matches_config(
cmdline, port, expected_model, expected_passages_file
)
print(f"DEBUG: No process found listening on port {port}")
logger.debug(f"No process found listening on port {port}")
return False
except Exception as e:
print(f"WARNING: Could not check process on port {port}: {e}")
logger.warning(f"Could not check process on port {port}: {e}")
return False
@@ -61,7 +69,7 @@ def _check_cmdline_matches_config(
) -> bool:
"""Check if command line matches our expected configuration."""
cmdline_str = " ".join(cmdline)
print(f"DEBUG: Found process on port {port}: {cmdline_str}")
logger.debug(f"Found process on port {port}: {cmdline_str}")
# Check if it's our embedding server
is_embedding_server = any(
@@ -74,7 +82,7 @@ def _check_cmdline_matches_config(
)
if not is_embedding_server:
print(f"DEBUG: Process on port {port} is not our embedding server")
logger.debug(f"Process on port {port} is not our embedding server")
return False
# Check model name
@@ -84,8 +92,8 @@ def _check_cmdline_matches_config(
passages_matches = _check_passages_in_cmdline(cmdline, expected_passages_file)
result = model_matches and passages_matches
print(
f"DEBUG: model_matches: {model_matches}, passages_matches: {passages_matches}, overall: {result}"
logger.debug(
f"model_matches: {model_matches}, passages_matches: {passages_matches}, overall: {result}"
)
return result
@@ -132,10 +140,10 @@ def _find_compatible_port_or_next_available(
# Port is in use, check if it's compatible
if _check_process_matches_config(port, model_name, passages_file):
print(f"✅ Found compatible server on port {port}")
logger.info(f"Found compatible server on port {port}")
return port, True
else:
print(f"⚠️ Port {port} has incompatible server, trying next port...")
logger.info(f"Port {port} has incompatible server, trying next port...")
raise RuntimeError(
f"Could not find compatible or available port in range {start_port}-{start_port + max_attempts}"
@@ -194,17 +202,17 @@ class EmbeddingServerManager:
port, model_name, passages_file
)
except RuntimeError as e:
print(f"❌ {e}")
logger.error(str(e))
return False, port
if is_compatible:
print(f"✅ Using existing compatible server on port {actual_port}")
logger.info(f"Using existing compatible server on port {actual_port}")
self.server_port = actual_port
self.server_process = None # We don't own this process
return True, actual_port
if actual_port != port:
print(f"⚠️ Using port {actual_port} instead of {port}")
logger.info(f"Using port {actual_port} instead of {port}")
# Start new server
return self._start_new_server(actual_port, model_name, embedding_mode, **kwargs)
@@ -221,19 +229,21 @@ class EmbeddingServerManager:
return False
if _check_process_matches_config(self.server_port, model_name, passages_file):
print(
f"✅ Existing server process (PID {self.server_process.pid}) is compatible"
logger.info(
f"Existing server process (PID {self.server_process.pid}) is compatible"
)
return True
print("⚠️ Existing server process is incompatible. Should start a new server.")
logger.info(
"Existing server process is incompatible. Should start a new server."
)
return False
def _start_new_server(
self, port: int, model_name: str, embedding_mode: str, **kwargs
) -> tuple[bool, int]:
"""Start a new embedding server on the given port."""
print(f"INFO: Starting embedding server on port {port}...")
logger.info(f"Starting embedding server on port {port}...")
command = self._build_server_command(port, model_name, embedding_mode, **kwargs)
@@ -241,7 +251,7 @@ class EmbeddingServerManager:
self._launch_server_process(command, port)
return self._wait_for_server_ready(port)
except Exception as e:
print(f"❌ ERROR: Failed to start embedding server: {e}")
logger.error(f"Failed to start embedding server: {e}")
return False, port
def _build_server_command(
@@ -268,20 +278,18 @@ class EmbeddingServerManager:
def _launch_server_process(self, command: list, port: int) -> None:
"""Launch the server process."""
project_root = Path(__file__).parent.parent.parent.parent.parent
print(f"INFO: Command: {' '.join(command)}")
logger.info(f"Command: {' '.join(command)}")
# Let server output go directly to console
# The server will respect LEANN_LOG_LEVEL environment variable
self.server_process = subprocess.Popen(
command,
cwd=project_root,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
bufsize=1,
universal_newlines=True,
stdout=None, # Direct to console
stderr=None, # Direct to console
)
self.server_port = port
print(f"INFO: Server process started with PID: {self.server_process.pid}")
logger.info(f"Server process started with PID: {self.server_process.pid}")
# Register atexit callback only when we actually start a process
if not self._atexit_registered:
@@ -294,49 +302,19 @@ class EmbeddingServerManager:
max_wait, wait_interval = 120, 0.5
for _ in range(int(max_wait / wait_interval)):
if _check_port(port):
print("✅ Embedding server is ready!")
threading.Thread(target=self._log_monitor, daemon=True).start()
logger.info("Embedding server is ready!")
return True, port
if self.server_process.poll() is not None:
print("❌ ERROR: Server terminated during startup.")
self._print_recent_output()
if self.server_process and self.server_process.poll() is not None:
logger.error("Server terminated during startup.")
return False, port
time.sleep(wait_interval)
print(f"❌ ERROR: Server failed to start within {max_wait} seconds.")
logger.error(f"Server failed to start within {max_wait} seconds.")
self.stop_server()
return False, port
def _print_recent_output(self):
"""Print any recent output from the server process."""
if not self.server_process or not self.server_process.stdout:
return
try:
if select.select([self.server_process.stdout], [], [], 0)[0]:
output = self.server_process.stdout.read()
if output:
print(f"[{self.backend_module_name} OUTPUT]: {output}")
except Exception as e:
print(f"Error reading server output: {e}")
def _log_monitor(self):
"""Monitors and prints the server's stdout and stderr."""
if not self.server_process:
return
try:
if self.server_process.stdout:
while True:
line = self.server_process.stdout.readline()
if not line:
break
print(
f"[{self.backend_module_name} LOG]: {line.strip()}", flush=True
)
except Exception as e:
print(f"Log monitor error: {e}")
def stop_server(self):
"""Stops the embedding server process if it's running."""
if not self.server_process:
@@ -347,17 +325,17 @@ class EmbeddingServerManager:
self.server_process = None
return
print(
f"INFO: Terminating server process (PID: {self.server_process.pid}) for backend {self.backend_module_name}..."
logger.info(
f"Terminating server process (PID: {self.server_process.pid}) for backend {self.backend_module_name}..."
)
self.server_process.terminate()
try:
self.server_process.wait(timeout=5)
print(f"INFO: Server process {self.server_process.pid} terminated.")
logger.info(f"Server process {self.server_process.pid} terminated.")
except subprocess.TimeoutExpired:
print(
f"WARNING: Server process {self.server_process.pid} did not terminate gracefully, killing it."
logger.warning(
f"Server process {self.server_process.pid} did not terminate gracefully, killing it."
)
self.server_process.kill()