Limit concurrent analyses with asyncio.Semaphore
Unbounded asyncio.create_task calls could exhaust CPU and memory on modest hardware. Add a configurable semaphore (MAX_CONCURRENT_ANALYSES, default 3) so excess jobs queue instead of running all at once.
This commit is contained in:
parent
7bd7d0f793
commit
c2e91262c6
4 changed files with 74 additions and 4 deletions
|
|
@ -13,6 +13,7 @@ class Settings:
|
|||
docling_serve_url: str = "http://localhost:5001"
|
||||
docling_serve_api_key: str | None = None
|
||||
conversion_timeout: int = 600
|
||||
max_concurrent_analyses: int = 3
|
||||
upload_dir: str = "./uploads"
|
||||
db_path: str = "./data/docling_studio.db"
|
||||
cors_origins: list[str] = field(
|
||||
|
|
@ -29,6 +30,7 @@ class Settings:
|
|||
docling_serve_url=os.environ.get("DOCLING_SERVE_URL", "http://localhost:5001"),
|
||||
docling_serve_api_key=os.environ.get("DOCLING_SERVE_API_KEY"),
|
||||
conversion_timeout=int(os.environ.get("CONVERSION_TIMEOUT", "600")),
|
||||
max_concurrent_analyses=int(os.environ.get("MAX_CONCURRENT_ANALYSES", "3")),
|
||||
upload_dir=os.environ.get("UPLOAD_DIR", "./uploads"),
|
||||
db_path=os.environ.get("DB_PATH", "./data/docling_studio.db"),
|
||||
cors_origins=[o.strip() for o in cors_raw.split(",")],
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ def _build_analysis_service() -> AnalysisService:
|
|||
converter=converter,
|
||||
chunker=chunker,
|
||||
conversion_timeout=settings.conversion_timeout,
|
||||
max_concurrent=settings.max_concurrent_analyses,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -34,6 +34,10 @@ def _chunk_to_dict(c: ChunkResult) -> dict:
|
|||
}
|
||||
|
||||
|
||||
# Maximum number of concurrent analysis jobs to prevent resource exhaustion.
|
||||
_DEFAULT_MAX_CONCURRENT = 3
|
||||
|
||||
|
||||
class AnalysisService:
|
||||
"""Orchestrates document analysis using an injected converter."""
|
||||
|
||||
|
|
@ -42,10 +46,12 @@ class AnalysisService:
|
|||
converter: DocumentConverter,
|
||||
chunker: DocumentChunker | None = None,
|
||||
conversion_timeout: int = 600,
|
||||
max_concurrent: int = _DEFAULT_MAX_CONCURRENT,
|
||||
):
|
||||
self._converter = converter
|
||||
self._chunker = chunker
|
||||
self._conversion_timeout = conversion_timeout
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
async def create(
|
||||
self,
|
||||
|
|
@ -113,7 +119,25 @@ class AnalysisService:
|
|||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> None:
|
||||
"""Background task: run conversion and optionally chunk."""
|
||||
"""Background task: run conversion and optionally chunk.
|
||||
|
||||
Acquires the concurrency semaphore to limit parallel conversions
|
||||
and prevent CPU/memory exhaustion on modest hardware.
|
||||
"""
|
||||
async with self._semaphore:
|
||||
await self._run_analysis_inner(
|
||||
job_id, file_path, filename, pipeline_options, chunking_options
|
||||
)
|
||||
|
||||
async def _run_analysis_inner(
|
||||
self,
|
||||
job_id: str,
|
||||
file_path: str,
|
||||
filename: str,
|
||||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> None:
|
||||
"""Inner analysis logic — called under the concurrency semaphore."""
|
||||
try:
|
||||
job = await analysis_repo.find_by_id(job_id)
|
||||
if not job:
|
||||
|
|
|
|||
|
|
@ -1,13 +1,13 @@
|
|||
"""Tests for AnalysisService — focus on _on_task_done callback (bug #1 fix)."""
|
||||
"""Tests for AnalysisService — callbacks, concurrency, and orchestration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from services.analysis_service import _on_task_done
|
||||
from services.analysis_service import AnalysisService, _on_task_done
|
||||
|
||||
|
||||
class TestOnTaskDone:
|
||||
|
|
@ -69,3 +69,46 @@ class TestOnTaskDone:
|
|||
await asyncio.sleep(0)
|
||||
|
||||
mock_mark.assert_not_called()
|
||||
|
||||
|
||||
class TestAnalysisServiceConcurrency:
|
||||
"""Verify that the semaphore limits concurrent analysis jobs."""
|
||||
|
||||
def test_semaphore_initialized_with_max_concurrent(self):
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter, max_concurrent=5)
|
||||
assert service._semaphore._value == 5
|
||||
|
||||
def test_default_max_concurrent(self):
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter)
|
||||
assert service._semaphore._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semaphore_limits_parallel_jobs(self):
|
||||
"""Only max_concurrent jobs should run in parallel; others must wait."""
|
||||
call_order: list[str] = []
|
||||
blocker = asyncio.Event()
|
||||
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter, max_concurrent=1)
|
||||
|
||||
async def fake_inner(self, *args, **kwargs):
|
||||
call_order.append("start")
|
||||
await blocker.wait()
|
||||
call_order.append("end")
|
||||
|
||||
with patch.object(AnalysisService, "_run_analysis_inner", fake_inner):
|
||||
t1 = asyncio.create_task(service._run_analysis("j1", "/f", "f.pdf"))
|
||||
t2 = asyncio.create_task(service._run_analysis("j2", "/f", "f.pdf"))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
# With max_concurrent=1, only one task should have started
|
||||
assert call_order.count("start") == 1
|
||||
|
||||
blocker.set()
|
||||
await asyncio.gather(t1, t2)
|
||||
|
||||
# Both should have completed
|
||||
assert call_order.count("start") == 2
|
||||
assert call_order.count("end") == 2
|
||||
|
|
|
|||
Loading…
Reference in a new issue