Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
aniketmaurya committed Sep 20, 2024
1 parent 25be47a commit 65b0d59
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions src/litserve/loggers.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import functools
import multiprocessing as mp
from abc import ABC, abstractmethod
from typing import List, Optional, Union, TYPE_CHECKING
Expand Down Expand Up @@ -95,22 +96,24 @@ def add_logger(self, logger: Logger):
if "mount" in logger._config:
self._mount(logger._config["mount"]["path"], logger._config["mount"]["app"])

def _process_logger_queue(self, loggers: List[Logger], queue):
@staticmethod
def _process_logger_queue(loggers: List[Logger], queue):
while True:
key, value = queue.get()
for logger in loggers:
logger.process(key, value)

@functools.cache # Run once per LitServer instance
def run(self):
ctx = mp.get_context("spawn")
queue = self._lit_server.logger_queue
# Disconnect the logger connector from the LitServer to avoid pickling issues
del self._lit_server
self._lit_server = None
if not self._loggers:
return

process = ctx.Process(
target=self._process_logger_queue,
target=_LoggerConnector._process_logger_queue,
args=(
self._loggers,
queue,
Expand Down

0 comments on commit 65b0d59

Please sign in to comment.