fix(logging): rotate exo's log file by size, not only when exo starts

src/exo/shared/logging.py

@@ -1,14 +1,21 @@
 import logging
 import sys
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from pathlib import Path
+from typing import TYPE_CHECKING, TextIO
 
 import zstandard
 from hypercorn import Config
 from hypercorn.logging import Logger as HypercornLogger
 from loguru import logger
 
+if TYPE_CHECKING:
+    from loguru import Message
+
 _MAX_LOG_ARCHIVES = 5
+# A node that runs for weeks would otherwise keep one ever-growing log file: at default
+# verbosity each request logs a few lines, ~200 MB a day under steady load
+_MAX_LOG_BYTES = 50 * 1024 * 1024
 
 
 def _zstd_compress(filepath: str) -> None:
@@ -26,6 +33,18 @@ def _once_then_never() -> Iterator[bool]:
         yield False
 
 
+def _rotate_at_start_and_by_size(
+    max_bytes: int,
+) -> "Callable[[Message, TextIO], bool]":
+    """Start a new log file when exo starts, and whenever the current one passes max_bytes."""
+    at_start = _once_then_never()
+
+    def should_rotate(message: "Message", file: TextIO) -> bool:
+        return next(at_start) or file.tell() + len(message) > max_bytes
+
+    return should_rotate
+
+
 class InterceptLogger(HypercornLogger):
     def __init__(self, config: Config):
         super().__init__(config)
@@ -73,14 +92,13 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
             enqueue=True,
         )
     if log_file:
-        rotate_once = _once_then_never()
         logger.add(
             log_file,
             format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
             level="DEBUG" if verbosity > 0 else "INFO",
             colorize=False,
             enqueue=True,
-            rotation=lambda _, __: next(rotate_once),
+            rotation=_rotate_at_start_and_by_size(_MAX_LOG_BYTES),
             retention=_MAX_LOG_ARCHIVES,
             compression=_zstd_compress,
         )

src/exo/shared/tests/test_log_rotation.py

@@ -0,0 +1,45 @@
+"""exo's log file is rotated when exo starts and whenever it grows past a size, so a node that runs
+for weeks keeps a bounded amount of log on disk."""
+
+import sys
+from pathlib import Path
+
+import pytest
+from loguru import logger
+
+from exo.shared import logging as exo_logging
+
+
+@pytest.fixture
+def small_logs(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(exo_logging, "_MAX_LOG_BYTES", 20_000)
+    log_file = tmp_path / "exo.log"
+    exo_logging.logger_setup(log_file, verbosity=0)
+    yield log_file
+    logger.remove()
+    logger.add(sys.stderr)
+
+
+def test_the_log_is_rotated_once_it_passes_the_size_limit(small_logs: Path):
+    for i in range(2_000):
+        logger.info(f"request {i}: " + "x" * 100)
+    exo_logging.logger_cleanup()
+
+    archives = list(small_logs.parent.glob("exo.*.log.zst"))
+    assert len(archives) > 1, "the log was only rotated when exo started"
+    # Old logs are kept compressed, up to the archive limit
+    assert len(archives) <= exo_logging._MAX_LOG_ARCHIVES  # pyright: ignore[reportPrivateUsage]
+    # The current log file stays near the limit
+    assert small_logs.stat().st_size <= 20_000 + 1_000
+    # Nothing written last is lost
+    assert "request 1999" in small_logs.read_text()
+
+
+def test_a_small_log_is_only_rotated_when_exo_starts(small_logs: Path):
+    for i in range(10):
+        logger.info(f"request {i}")
+    exo_logging.logger_cleanup()
+
+    # The one archive is the previous run's log, set aside when this run started
+    assert len(list(small_logs.parent.glob("exo.*.log.zst"))) == 1
+    assert "request 9" in small_logs.read_text()