feat(application): add thread pool support and background job threading capabilities (#270)

* feat(application): add thread pool support and background job threading capabilities

- Integrate ThreadPoolExecutor for shared thread pool management in application context
- Add thread_pool_max_workers configuration option with default value of 0 (disabled)
- Implement thread pool creation and shutdown in application lifecycle methods
- Add use_thread_pool option to background job configuration for thread-based execution
- Support both asyncio event loop and threading event for background job stop mechanism
- Implement _run_in_thread method for running supervisors in dedicated threads
- Modify embedding models to use serial batching instead of concurrent batching
- Remove max_concurrency parameter from base embedding model component
- Update embedding model documentation to reflect serial batching implementation
- Add pyproject.toml with project metadata, dependencies, and build configuration

* refactor(job): simplify background job thread execution

- Replace separate _run_in_thread method with direct lambda execution
- Remove unnecessary ensure_future wrapper for thread pool execution
- Simplify asyncio event loop usage in thread pool mode
- Maintain same background job functionality with cleaner implementation
- Remove redundant method definition and streamline execution flow

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 -------------------------------------------------
 

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

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)
 

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

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_*"]

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",