diff --git a/reme4/application.py b/reme4/application.py index fb610868..8a19eb49 100644 --- a/reme4/application.py +++ b/reme4/application.py @@ -2,6 +2,7 @@ import asyncio import heapq +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import AsyncGenerator, TypeVar @@ -168,6 +169,10 @@ class Application(BaseComponent): async def _start(self) -> None: """Start components in dependency order, then jobs (background last).""" + pool_size = self.config.thread_pool_max_workers + if pool_size > 0: + self.context.thread_pool = ThreadPoolExecutor(max_workers=pool_size) + self.logger.info(f"Thread pool created with max_workers={pool_size}") components = self._topological_order() jobs = list(self.context.jobs.values()) # Background jobs come last so they observe a fully wired system. @@ -194,6 +199,9 @@ class Application(BaseComponent): except Exception as e: self.logger.exception(f"Failed to close {c.component_type.value}:{c.name}: {e}") self._started_components.clear() + if self.context.thread_pool is not None: + self.context.thread_pool.shutdown(wait=True) + self.context.thread_pool = None # ----- Job execution ------------------------------------------------- diff --git a/reme4/components/application_context.py b/reme4/components/application_context.py index 0feac294..4b482b7b 100644 --- a/reme4/components/application_context.py +++ b/reme4/components/application_context.py @@ -1,5 +1,6 @@ """Application context: shared state container for components, jobs, and service.""" +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING from ..enumeration import ComponentEnum @@ -27,3 +28,4 @@ class ApplicationContext: self.service: "BaseService | None" = None self.components: dict[ComponentEnum, dict[str, "BaseComponent"]] = {} self.jobs: dict[str, "BaseJob"] = {} + self.thread_pool: ThreadPoolExecutor | None = None diff --git a/reme4/components/embedding/base_embedding_model.py b/reme4/components/embedding/base_embedding_model.py index a149cfcf..b11071a3 100644 --- a/reme4/components/embedding/base_embedding_model.py +++ b/reme4/components/embedding/base_embedding_model.py @@ -17,7 +17,7 @@ Miss = tuple[int, str, str] # (result_index, text, cache_key) class BaseEmbeddingModel(BaseComponent): - """Embedding model with LRU cache, disk persistence, and concurrent batching.""" + """Embedding model with LRU cache, disk persistence, and serial batching.""" component_type = ComponentEnum.EMBEDDING_MODEL @@ -31,7 +31,6 @@ class BaseEmbeddingModel(BaseComponent): max_batch_size: int = 10, max_input_length: int = 8192, max_cache_size: int = 10000, - max_concurrency: int = 2, enable_cache: bool = True, cache_version: str = "v1", max_retries: int = 3, @@ -46,7 +45,6 @@ class BaseEmbeddingModel(BaseComponent): self.max_batch_size = max_batch_size self.max_input_length = max_input_length self.max_cache_size = max_cache_size - self.max_concurrency = max_concurrency self.enable_cache = enable_cache self.cache_version = cache_version self.max_retries = max_retries @@ -90,7 +88,7 @@ class BaseEmbeddingModel(BaseComponent): return results[0] if results else None async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: - """Get embeddings for texts. Cache hits return immediately; misses run concurrently.""" + """Get embeddings for texts. Cache hits return immediately; misses run in serial batches.""" texts = [self._truncate(t) for t in input_text] results, misses = self._partition_by_cache(texts) if misses: @@ -129,17 +127,11 @@ class BaseEmbeddingModel(BaseComponent): return results, misses async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None: - """Compute miss embeddings in concurrent batches and write into results + cache.""" + """Compute miss embeddings in serial batches and write into results + cache.""" size = self.max_batch_size batches = [misses[i : i + size] for i in range(0, len(misses), size)] - sem = asyncio.Semaphore(self.max_concurrency) - - async def run(batch: list[Miss]) -> list[tuple[int, str, np.ndarray]]: - async with sem: - return await self._compute_batch(batch, **kwargs) - - for done in await asyncio.gather(*(run(b) for b in batches)): - for idx, key, emb in done: + for batch in batches: + for idx, key, emb in await self._compute_batch(batch, **kwargs): results[idx] = emb self._cache_put(key, emb) diff --git a/reme4/components/job/background_job.py b/reme4/components/job/background_job.py index c72de54f..fa623d78 100644 --- a/reme4/components/job/background_job.py +++ b/reme4/components/job/background_job.py @@ -3,6 +3,7 @@ import asyncio import contextlib import random +import threading import time from .base_job import BaseJob @@ -36,6 +37,7 @@ class BackgroundJob(BaseJob): close_timeout: float = 5.0, attempt_reset_after: float = 60.0, enable_serve: bool = False, + use_thread_pool: bool = False, **kwargs, ): super().__init__(enable_serve=enable_serve, **kwargs) @@ -44,13 +46,21 @@ class BackgroundJob(BaseJob): self.backoff_cap: float = backoff_cap self.close_timeout: float = close_timeout self.attempt_reset_after: float = attempt_reset_after - self._stop_event: asyncio.Event | None = None + self.use_thread_pool: bool = use_thread_pool + self._stop_event: asyncio.Event | threading.Event | None = None self._task: asyncio.Task | None = None async def _start(self) -> None: await super()._start() - self._stop_event = asyncio.Event() - self._task = asyncio.create_task(self._run_with_supervisor()) + if self.use_thread_pool and self.app_context.thread_pool: + self._stop_event = threading.Event() + self._task = asyncio.get_event_loop().run_in_executor( + self.app_context.thread_pool, + lambda: asyncio.run(self._run_with_supervisor()), + ) + else: + self._stop_event = asyncio.Event() + self._task = asyncio.create_task(self._run_with_supervisor()) async def _close(self) -> None: if self._stop_event is not None: @@ -82,10 +92,13 @@ class BackgroundJob(BaseJob): async def _wait_or_stop(self, delay: float) -> None: """Sleep up to delay, returning immediately when stop_event is set.""" assert self._stop_event is not None - try: - await asyncio.wait_for(self._stop_event.wait(), timeout=delay) - except asyncio.TimeoutError: - pass + if isinstance(self._stop_event, threading.Event): + self._stop_event.wait(timeout=delay) + else: + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=delay) + except asyncio.TimeoutError: + pass async def _run_with_supervisor(self) -> None: assert self._stop_event is not None diff --git a/reme4/pyproject.toml b/reme4/pyproject.toml new file mode 100644 index 00000000..abba8874 --- /dev/null +++ b/reme4/pyproject.toml @@ -0,0 +1,84 @@ +[project] +name = "reme4" +dynamic = ["version"] +description = "Remember Me, Refine Me." +authors = [ + { name = "EconML team of Alibaba Tongyi Lab", email = "jinli.yl@alibaba-inc.com" }, +] +license = "Apache-2.0" +requires-python = ">=3.11" +keywords = ["llm", "memory", "agent", "agentscope", "ai", "mcp", "reme"] + +classifiers = [ + "Development Status :: 4 - Beta", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Operating System :: OS Independent", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] + +dependencies = [ + "aiofiles>=24.1.0", + "fastapi>=0.135.1", + "fastmcp>=3.1.0", + "httpx>=0.28.1", + "loguru>=0.7.3", + "mistletoe>=1.5.1", + "numpy>=2.2.6", + "openai>=2.26.0", + "pydantic>=2.12.5", + "python-frontmatter>=1.1.0", + "pyyaml>=6.0.3", + "rich>=14.3.3", + "uvicorn>=0.41.0", + "watchfiles>=1.1.1", +] + +[project.optional-dependencies] +core = [ + "agentscope>=2.0.0", + "faiss-cpu>=1.13.2", + "jieba>=0.42.1", + "rjieba>=0.2.1", + "neo4j>=6.2.0", + "networkx>=3.4.2", +] +dev = [ + "pre-commit", + "pytest>=8.0", + "pytest-asyncio>=0.23", +] +full = [ + "reme4[core]", + "reme4[dev]", +] + +[project.urls] +Homepage = "https://github.com/agentscope-ai/ReMe" +Documentation = "https://reme.agentscope.io/" +Repository = "https://github.com/agentscope-ai/ReMe" + +[project.scripts] +reme4 = "reme4.reme:main" + +[tool.setuptools] +packages = { find = { where = [".."], include = ["reme4*"] } } +include-package-data = true + +[tool.setuptools.package-data] +"*" = ["py.typed", "**/*.yaml", "**/*.json"] + +[tool.setuptools.dynamic] +version = { attr = "reme4.__version__" } + +[build-system] +requires = ["setuptools>=45", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.pytest.ini_options] +asyncio_default_fixture_loop_scope = "function" +testpaths = ["../tests4"] +python_files = ["test_*.py"] +python_functions = ["test_*"] diff --git a/reme4/schema/application_config.py b/reme4/schema/application_config.py index 3df6ff1d..35924bb4 100644 --- a/reme4/schema/application_config.py +++ b/reme4/schema/application_config.py @@ -46,6 +46,10 @@ class ApplicationConfig(BaseModel): default_factory=dict, description="Job definitions keyed by job name", ) + thread_pool_max_workers: int = Field( + default=0, + description="Max worker threads in the shared thread pool; 0 to disable", + ) components: dict[ComponentEnum, dict[str, ComponentConfig]] = Field( default_factory=dict, description="Component registry keyed by type then name",