fix: reflect logger changes during life time of task (#155)

fal-ai · Sep 18, 2024 · a431995 · a431995
1 parent 47c4a1b
commit a431995
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/isolate/server/server.py b/src/isolate/server/server.py
@@ -211,7 +211,7 @@ def _run_task(self, task: RunTask) -> Iterator[definitions.PartialRunResult]:
                 StatusCode.INVALID_ARGUMENT,
             )
 
-        log_handler = LogHandler(messages, logger=task.logger)
+        log_handler = LogHandler(messages, task=task)
 
         run_settings = replace(
             self.default_settings,
@@ -480,10 +480,11 @@ def _proxy_to_queue(
 @dataclass
 class LogHandler:
     messages: Queue
-    logger: IsolateLogger = ENV_LOGGER
+    # Reference to the task so we can change the logger
+    task: RunTask
 
     def handle(self, log: Log) -> None:
-        self.logger.log(log.level, log.message, source=log.source)
+        self.task.logger.log(log.level, log.message, source=log.source)
         self._add_log_to_queue(log)
 
     def _add_log_to_queue(self, log: Log) -> None: