diff --git a/application/Dockerfile b/application/Dockerfile index 849bcc79..0eb037d3 100644 --- a/application/Dockerfile +++ b/application/Dockerfile @@ -8,7 +8,7 @@ RUN apt-get update && \ add-apt-repository ppa:deadsnakes/ppa && \ apt-get update && \ apt-get install -y --no-install-recommends gcc g++ wget unzip libc6-dev python3.12 python3.12-venv python3.12-dev && \ - rm -rf /var/lib/apt/lists/* + rm -rf /var/lib/apt/lists/* # Verify Python installation and setup symlink RUN if [ -f /usr/bin/python3.12 ]; then \ @@ -73,7 +73,7 @@ COPY --from=builder /models /app/models COPY . /app/application # Change the ownership of the /app directory to the appuser - + RUN mkdir -p /app/application/inputs/local RUN chown -R appuser:appuser /app @@ -82,6 +82,11 @@ ENV FLASK_APP=app.py \ FLASK_DEBUG=true \ PATH="/venv/bin:$PATH" +ENV MALLOC_ARENA_MAX=2 \ + OMP_NUM_THREADS=4 \ + MKL_NUM_THREADS=4 \ + OPENBLAS_NUM_THREADS=4 + # Expose the port the app runs on EXPOSE 7091 diff --git a/application/celery_init.py b/application/celery_init.py index 69fdd28a..6e400105 100644 --- a/application/celery_init.py +++ b/application/celery_init.py @@ -1,5 +1,8 @@ +import ctypes +import gc import inspect import logging +import sys import threading from celery import Celery @@ -98,6 +101,34 @@ def _unbind_task_log_context(task_id, **_): ) +def _trim_native_heap() -> None: + """Return freed glibc heap pages to the OS (Linux only; no-op elsewhere).""" + # docling/torch parsing makes large transient allocations; glibc keeps the + # freed pages in per-thread malloc arenas rather than returning them, so a + # long-lived worker child's RSS only ever climbs. malloc_trim hands them + # back. The symbol is glibc-only — absent in macOS libc. + if not sys.platform.startswith("linux"): + return + try: + ctypes.CDLL("libc.so.6").malloc_trim(0) + except (OSError, AttributeError): + pass + + +@task_postrun.connect +def _reclaim_memory_after_task(*args, **kwargs): + """Drop per-task allocations so the prefork child's RSS doesn't ratchet.""" + gc.collect() + torch = sys.modules.get("torch") + if torch is not None: + try: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + _trim_native_heap() + + @worker_ready.connect def _run_version_check(*args, **kwargs): """Kick off the anonymous version check on worker startup. diff --git a/application/celeryconfig.py b/application/celeryconfig.py index eb060bae..82d9ca0d 100644 --- a/application/celeryconfig.py +++ b/application/celeryconfig.py @@ -31,3 +31,10 @@ worker_prefetch_multiplier = settings.CELERY_WORKER_PREFETCH_MULTIPLIER broker_transport_options = {"visibility_timeout": settings.CELERY_VISIBILITY_TIMEOUT} result_expires = 86400 * 7 task_track_started = True + +# Recycle the prefork worker child to bound native-heap growth from +# docling/torch parsing. Left unset (Celery's unlimited default) when 0. +if settings.CELERY_WORKER_MAX_MEMORY_PER_CHILD > 0: + worker_max_memory_per_child = settings.CELERY_WORKER_MAX_MEMORY_PER_CHILD +if settings.CELERY_WORKER_MAX_TASKS_PER_CHILD > 0: + worker_max_tasks_per_child = settings.CELERY_WORKER_MAX_TASKS_PER_CHILD diff --git a/application/core/settings.py b/application/core/settings.py index dfb4cc66..e243bcbd 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -36,6 +36,11 @@ class Settings(BaseSettings): # and Dify defaults; long ingests can override via env. CELERY_WORKER_PREFETCH_MULTIPLIER: int = 1 CELERY_VISIBILITY_TIMEOUT: int = 3600 + # Recycle the prefork worker child once its resident size crosses this many + # kilobytes — backstops native-heap growth from docling/torch parsing. 0 disables. + CELERY_WORKER_MAX_MEMORY_PER_CHILD: int = 4194304 + # Recycle the child after this many tasks; 0 disables (memory cap is the primary knob). + CELERY_WORKER_MAX_TASKS_PER_CHILD: int = 0 # Only consulted when VECTOR_STORE=mongodb or when running scripts/db/backfill.py; user data lives in Postgres. MONGO_URI: Optional[str] = None # User-data Postgres DB.