better:优化部分导入以更快的启动

This commit is contained in:
SengokuCola
2026-04-25 14:26:02 +08:00
parent dcc6748a76
commit 8168fe0d8a
5 changed files with 64 additions and 27 deletions

View File

@@ -4,7 +4,6 @@ import asyncio
import json
import pickle
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
@@ -19,7 +18,7 @@ from src.services.llm_service import LLMServiceClient
from ...paths import default_data_dir, resolve_repo_path
from ..embedding import create_embedding_api_adapter
from ..retrieval import RetrievalResult, SparseBM25Config, SparseBM25Index, TemporalQueryOptions
from ..retrieval import RetrievalResult, SparseBM25Config, SparseBM25Index
from ..storage import GraphStore, MetadataStore, QuantizationType, SparseMatrixFormat, VectorStore
from ..utils.aggregate_query_service import AggregateQueryService
from ..utils.episode_retrieval_service import EpisodeRetrievalService
@@ -634,23 +633,18 @@ class SDKMemoryKernel:
model_name=str(self._cfg("embedding.model_name", "auto") or "auto"),
retry_config=self._cfg("embedding.retry", {}) or {},
)
detected_dimension = int(await self.embedding_manager._detect_dimension())
self.embedding_dimension = detected_dimension
dimension_detection_task = asyncio.create_task(
asyncio.to_thread(lambda: asyncio.run(self.embedding_manager._detect_dimension()))
)
await asyncio.sleep(0)
stored_dimension = self._stored_vector_dimension()
if stored_dimension is not None and stored_dimension != detected_dimension:
raise RuntimeError(
self._vector_mismatch_error(
stored_dimension=stored_dimension,
detected_dimension=detected_dimension,
)
)
provisional_dimension = stored_dimension or self.embedding_dimension
matrix_format = str(self._cfg("graph.sparse_matrix_format", "csr") or "csr").strip().lower()
graph_format = SparseMatrixFormat.CSC if matrix_format == "csc" else SparseMatrixFormat.CSR
self.vector_store = VectorStore(
dimension=detected_dimension,
dimension=provisional_dimension,
quantization_type=QuantizationType.INT8,
data_dir=self.data_dir / "vectors",
)
@@ -658,9 +652,11 @@ class SDKMemoryKernel:
self.metadata_store = MetadataStore(data_dir=self.data_dir / "metadata")
self.metadata_store.connect()
if self.vector_store.has_data():
vector_store_loaded = False
if stored_dimension is not None and self.vector_store.has_data():
self.vector_store.load()
self.vector_store.warmup_index(force_train=True)
vector_store_loaded = True
if self.graph_store.has_data():
self.graph_store.load()
@@ -674,6 +670,33 @@ class SDKMemoryKernel:
if getattr(self.sparse_index.config, "enabled", False):
self.sparse_index.ensure_loaded()
try:
detected_dimension = int(await dimension_detection_task)
except Exception:
if not dimension_detection_task.done():
dimension_detection_task.cancel()
raise
self.embedding_dimension = detected_dimension
if stored_dimension is not None and stored_dimension != detected_dimension:
raise RuntimeError(
self._vector_mismatch_error(
stored_dimension=stored_dimension,
detected_dimension=detected_dimension,
)
)
if self.vector_store.dimension != detected_dimension:
self.vector_store = VectorStore(
dimension=detected_dimension,
quantization_type=QuantizationType.INT8,
data_dir=self.data_dir / "vectors",
)
if not vector_store_loaded and self.vector_store.has_data():
self.vector_store.load()
self.vector_store.warmup_index(force_train=True)
self.relation_write_service = RelationWriteService(
metadata_store=self.metadata_store,
graph_store=self.graph_store,

View File

@@ -279,9 +279,14 @@ class AMemorixHostService:
async with self._lock:
if self._kernel is None:
config = self._read_config()
self._kernel = SDKMemoryKernel(plugin_root=repo_root(), config=config)
await self._kernel.initialize()
set_runtime_kernel(self._kernel)
kernel = SDKMemoryKernel(plugin_root=repo_root(), config=config)
try:
await kernel.initialize()
except Exception:
kernel.close()
raise
self._kernel = kernel
set_runtime_kernel(kernel)
return self._kernel
def _read_config(self) -> Dict[str, Any]: