From 698c30f7b8fd36ba6beeba02f6d95b0d11efde73 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 20:50:37 +0800 Subject: [PATCH 1/4] fix: configure logging once per process --- src/memos/log.py | 23 ++++++++++++++++++++++- tests/test_log.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/memos/log.py b/src/memos/log.py index c0bb5bf31..ede5e8280 100644 --- a/src/memos/log.py +++ b/src/memos/log.py @@ -27,6 +27,8 @@ load_dotenv() selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING +_LOGGING_CONFIG_LOCK = threading.RLock() +_LOGGING_CONFIGURED_PID: int | None = None def _setup_logfile() -> Path: @@ -224,12 +226,31 @@ def close(self): } +def configure_logging(force: bool = False) -> None: + """Configure process-local logging once. + + Re-running dictConfig replaces and closes existing handlers. Guarding it avoids races with + background threads that may be emitting log records while other modules import loggers. + """ + global _LOGGING_CONFIGURED_PID + + current_pid = os.getpid() + if not force and current_pid == _LOGGING_CONFIGURED_PID: + return + + with _LOGGING_CONFIG_LOCK: + current_pid = os.getpid() + if force or current_pid != _LOGGING_CONFIGURED_PID: + dictConfig(LOGGING_CONFIG) + _LOGGING_CONFIGURED_PID = current_pid + + def get_logger(name: str | None = None) -> logging.Logger: """returns the project logger, scoped to a child name if provided Args: name: will define a child logger """ - dictConfig(LOGGING_CONFIG) + configure_logging() parent_logger = logging.getLogger("") if name: diff --git a/tests/test_log.py b/tests/test_log.py index fbd8791ee..e799ff1c4 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -28,3 +28,34 @@ def test_get_logger_returns_logger(): assert any(isinstance(h, logging.StreamHandler) for h in logger.parent.handlers) or any( isinstance(h, logging.FileHandler) for h in logger.parent.handlers ) + + +def test_get_logger_configures_logging_once_per_process(monkeypatch): + calls = [] + + monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) + monkeypatch.setattr(log.os, "getpid", lambda: 123) + monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) + + log.get_logger("first") + log.get_logger("second") + + assert len(calls) == 1 + + +def test_get_logger_reconfigures_after_process_fork(monkeypatch): + calls = [] + pid = 123 + + def getpid(): + return pid + + monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) + monkeypatch.setattr(log.os, "getpid", getpid) + monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) + + log.get_logger("parent") + pid = 456 + log.get_logger("child") + + assert len(calls) == 2 From 2017cd5ecb67c5a960d0dfedfe877614c69e15b0 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 20:50:53 +0800 Subject: [PATCH 2/4] fix: close scheduler resources on API shutdown --- src/memos/api/lifecycle.py | 26 ++++++++++++++++++++++++++ src/memos/api/server_api.py | 10 ++++++++-- src/memos/api/server_api_ext.py | 12 +++++++++--- tests/api/test_lifecycle.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 74 insertions(+), 5 deletions(-) create mode 100644 src/memos/api/lifecycle.py create mode 100644 tests/api/test_lifecycle.py diff --git a/src/memos/api/lifecycle.py b/src/memos/api/lifecycle.py new file mode 100644 index 000000000..e05a2f777 --- /dev/null +++ b/src/memos/api/lifecycle.py @@ -0,0 +1,26 @@ +from collections.abc import Mapping +from typing import Any + +from memos.log import get_logger + + +logger = get_logger(__name__) + + +def shutdown_components(components: Mapping[str, Any] | None) -> None: + """Release long-lived API components before the logging system shuts down.""" + if not components: + return + + mem_scheduler = components.get("mem_scheduler") + if mem_scheduler is None: + return + + for method_name in ("stop", "rabbitmq_close"): + method = getattr(mem_scheduler, method_name, None) + if not callable(method): + continue + try: + method() + except Exception: + logger.exception("Failed to run mem_scheduler.%s during API shutdown", method_name) diff --git a/src/memos/api/server_api.py b/src/memos/api/server_api.py index a9afe554c..b05e8ac7c 100644 --- a/src/memos/api/server_api.py +++ b/src/memos/api/server_api.py @@ -7,8 +7,9 @@ from starlette.staticfiles import StaticFiles from memos.api.exceptions import APIExceptionHandler +from memos.api.lifecycle import shutdown_components from memos.api.middleware.request_context import RequestContextMiddleware -from memos.api.routers.server_router import router as server_router +from memos.api.routers import server_router as server_router_module from memos.plugins.manager import plugin_manager @@ -35,7 +36,7 @@ app.add_middleware(RequestContextMiddleware, source="server_api") # Include routers -app.include_router(server_router) +app.include_router(server_router_module.router) @app.get("/health") @@ -48,6 +49,11 @@ def health_check(): } +@app.on_event("shutdown") +def shutdown_server_components() -> None: + shutdown_components(server_router_module.components) + + # Request validation failed app.exception_handler(RequestValidationError)(APIExceptionHandler.validation_error_handler) # Invalid business code parameters diff --git a/src/memos/api/server_api_ext.py b/src/memos/api/server_api_ext.py index 8c457e362..e28505332 100644 --- a/src/memos/api/server_api_ext.py +++ b/src/memos/api/server_api_ext.py @@ -26,11 +26,12 @@ from starlette.responses import Response # Import Krolik extensions +from memos.api.lifecycle import shutdown_components from memos.api.middleware.rate_limit import RateLimitMiddleware -from memos.api.routers.admin_router import router as admin_router # Import base routers from MemOS -from memos.api.routers.server_router import router as server_router +from memos.api.routers import server_router as server_router_module +from memos.api.routers.admin_router import router as admin_router # Try to import exception handlers (may vary between MemOS versions) @@ -94,7 +95,7 @@ async def dispatch(self, request: Request, call_next) -> Response: logger.info("Rate limiting enabled") # Include routers -app.include_router(server_router) +app.include_router(server_router_module.router) app.include_router(admin_router) # Exception handlers @@ -118,6 +119,11 @@ async def health_check(): } +@app.on_event("shutdown") +def shutdown_server_components() -> None: + shutdown_components(server_router_module.components) + + if __name__ == "__main__": import uvicorn diff --git a/tests/api/test_lifecycle.py b/tests/api/test_lifecycle.py new file mode 100644 index 000000000..1daa7a22e --- /dev/null +++ b/tests/api/test_lifecycle.py @@ -0,0 +1,31 @@ +from memos.api.lifecycle import shutdown_components + + +class SchedulerStub: + def __init__(self, fail_stop: bool = False): + self.fail_stop = fail_stop + self.calls = [] + + def stop(self): + self.calls.append("stop") + if self.fail_stop: + raise RuntimeError("stop failed") + + def rabbitmq_close(self): + self.calls.append("rabbitmq_close") + + +def test_shutdown_components_stops_scheduler_before_rabbitmq_close(): + scheduler = SchedulerStub() + + shutdown_components({"mem_scheduler": scheduler}) + + assert scheduler.calls == ["stop", "rabbitmq_close"] + + +def test_shutdown_components_still_closes_rabbitmq_when_stop_fails(): + scheduler = SchedulerStub(fail_stop=True) + + shutdown_components({"mem_scheduler": scheduler}) + + assert scheduler.calls == ["stop", "rabbitmq_close"] From 8aefda7b571f8d4991a5c7b6b5513e5c6daca91d Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 21:09:16 +0800 Subject: [PATCH 3/4] fix: configure logging once per process Avoid re-running dictConfig on every get_logger call, since repeated configuration can close and replace active handlers while background threads are emitting records. Keep logging configuration process-local and reconfigure only after a PID change. --- src/memos/log.py | 10 +++++----- tests/test_log.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/memos/log.py b/src/memos/log.py index ede5e8280..c18bd2118 100644 --- a/src/memos/log.py +++ b/src/memos/log.py @@ -226,6 +226,10 @@ def close(self): } +def _get_current_pid() -> int: + return os.getpid() + + def configure_logging(force: bool = False) -> None: """Configure process-local logging once. @@ -234,12 +238,8 @@ def configure_logging(force: bool = False) -> None: """ global _LOGGING_CONFIGURED_PID - current_pid = os.getpid() - if not force and current_pid == _LOGGING_CONFIGURED_PID: - return - with _LOGGING_CONFIG_LOCK: - current_pid = os.getpid() + current_pid = _get_current_pid() if force or current_pid != _LOGGING_CONFIGURED_PID: dictConfig(LOGGING_CONFIG) _LOGGING_CONFIGURED_PID = current_pid diff --git a/tests/test_log.py b/tests/test_log.py index e799ff1c4..5387adc34 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -34,7 +34,7 @@ def test_get_logger_configures_logging_once_per_process(monkeypatch): calls = [] monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) - monkeypatch.setattr(log.os, "getpid", lambda: 123) + monkeypatch.setattr(log, "_get_current_pid", lambda: 123) monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) log.get_logger("first") @@ -51,7 +51,7 @@ def getpid(): return pid monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) - monkeypatch.setattr(log.os, "getpid", getpid) + monkeypatch.setattr(log, "_get_current_pid", getpid) monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) log.get_logger("parent") From 9c7dd2fdbf132a75e9be0eea73b096e41d1ee385 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 21:09:22 +0800 Subject: [PATCH 4/4] fix: close scheduler resources via API lifespan Use FastAPI lifespan cleanup to stop scheduler and RabbitMQ resources before logging handlers are torn down during worker or pod shutdown. --- src/memos/api/server_api.py | 16 +++++++++++----- src/memos/api/server_api_ext.py | 15 ++++++++++----- tests/api/test_lifecycle.py | 7 +++++++ 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/src/memos/api/server_api.py b/src/memos/api/server_api.py index b05e8ac7c..1c6d93b9f 100644 --- a/src/memos/api/server_api.py +++ b/src/memos/api/server_api.py @@ -1,6 +1,9 @@ import logging import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + from dotenv import load_dotenv from fastapi import FastAPI, HTTPException from fastapi.exceptions import RequestValidationError @@ -26,10 +29,18 @@ os.getenv("MEMSCHEDULER_REDIS_STREAM_KEY_PREFIX"), ) + +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + yield + shutdown_components(server_router_module.components) + + app = FastAPI( title="MemOS Server REST APIs", description="A REST API for managing multiple users with MemOS Server.", version="1.0.1", + lifespan=lifespan, ) app.mount("/download", StaticFiles(directory=os.getenv("FILE_LOCAL_PATH")), name="static_mapping") @@ -49,11 +60,6 @@ def health_check(): } -@app.on_event("shutdown") -def shutdown_server_components() -> None: - shutdown_components(server_router_module.components) - - # Request validation failed app.exception_handler(RequestValidationError)(APIExceptionHandler.validation_error_handler) # Invalid business code parameters diff --git a/src/memos/api/server_api_ext.py b/src/memos/api/server_api_ext.py index e28505332..7b5aafc46 100644 --- a/src/memos/api/server_api_ext.py +++ b/src/memos/api/server_api_ext.py @@ -18,6 +18,9 @@ import logging import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware @@ -59,11 +62,18 @@ async def dispatch(self, request: Request, call_next) -> Response: return response +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + yield + shutdown_components(server_router_module.components) + + # Create FastAPI app app = FastAPI( title="MemOS Server REST APIs (Krolik Extended)", description="MemOS API with authentication, rate limiting, and admin endpoints.", version="2.0.3-krolik", + lifespan=lifespan, ) # CORS configuration @@ -119,11 +129,6 @@ async def health_check(): } -@app.on_event("shutdown") -def shutdown_server_components() -> None: - shutdown_components(server_router_module.components) - - if __name__ == "__main__": import uvicorn diff --git a/tests/api/test_lifecycle.py b/tests/api/test_lifecycle.py index 1daa7a22e..b2f7f7064 100644 --- a/tests/api/test_lifecycle.py +++ b/tests/api/test_lifecycle.py @@ -29,3 +29,10 @@ def test_shutdown_components_still_closes_rabbitmq_when_stop_fails(): shutdown_components({"mem_scheduler": scheduler}) assert scheduler.calls == ["stop", "rabbitmq_close"] + + +def test_shutdown_components_skips_missing_methods(): + class MinimalScheduler: + pass + + shutdown_components({"mem_scheduler": MinimalScheduler()})