如何捕获joblib loky子进程输出并导入Python logging？

问题背景 #

在使用scikit-learn时，我习惯开启verbose参数查看训练过程输出，同时希望日志统一包含时间戳、进程ID等元信息，因此使用Python的logging模块管理日志。虽然已有将sklearn输出重定向到logging的实现，但当使用joblib的loky多进程后端时，子进程的stdout无法被捕获，直接输出到控制台；而threading后端则可以正常工作。

原实现代码 #

import logging
import sys
import contextlib

class LogAdapter:
    def __init__(self, level, logger) -> None:
        if level == 'INFO':
            self.report = logger.info
        elif level == 'ERROR':
            self.report = logger.error

    def write(self, msg):
        stripped = msg.rstrip()
        if len(stripped) > 0:
            self.report(stripped)

    def flush(self):
        pass

@contextlib.contextmanager
def redirect_to_log(logger):
    originals = sys.stdout, sys.stderr
    sys.stdout = LogAdapter(level='INFO', logger=logger)
    sys.stderr = LogAdapter(level='ERROR', logger=logger)
    yield
    sys.stdout, sys.stderr = originals


def test_case():
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.utils import parallel_backend
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(process)d | %(asctime)s | %(name)14s | %(levelname)7s | %(message)s",
    )

    for backend_name in ['loky','threading']:
        logger.info(f"Testing backend {backend_name}")
        with parallel_backend(backend_name), redirect_to_log(logger):
            clf = RandomForestClassifier(2, verbose=4)
            X = [[0, 0], [1, 1]]
            Y = [0, 1]
            clf = clf.fit(X, Y)

if __name__ == "__main__":
    test_case()

原运行输出 #

19320 | 2022-11-30 17:49:16,938 |       __main__ |    INFO | Testing backend loky
19320 | 2022-11-30 17:49:16,951 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
building tree 1 of 2
building tree 2 of 2
19320 | 2022-11-30 17:49:18,923 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    1.9s remaining:    0.0s
19320 | 2022-11-30 17:49:18,923 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    1.9s finished
19320 | 2022-11-30 17:49:18,924 |       __main__ |    INFO | Testing backend threading
19320 | 2022-11-30 17:49:18,925 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Using backend ThreadingBackend with 8 concurrent workers.
19320 | 2022-11-30 17:49:18,932 |       __main__ |    INFO | building tree 1 of 2
19320 | 2022-11-30 17:49:18,932 |       __main__ |    INFO | building tree 2 of 2
19320 | 2022-11-30 17:49:18,934 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    0.0s remaining:    0.0s
19320 | 2022-11-30 17:49:18,934 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    0.0s finished

解决方案 #

Loky采用多进程模式，子进程不会继承主进程的sys.stdout/sys.stderr重定向，因此需要在每个子进程启动时单独初始化日志配置和输出重定向。可以通过parallel_backend的backend_params参数，传递子进程初始化函数来实现：

修改后的完整代码 #

import logging
import sys
import contextlib

class LogAdapter:
    def __init__(self, level, logger) -> None:
        if level == 'INFO':
            self.report = logger.info
        elif level == 'ERROR':
            self.report = logger.error

    def write(self, msg):
        stripped = msg.rstrip()
        if len(stripped) > 0:
            self.report(stripped)

    def flush(self):
        pass

@contextlib.contextmanager
def redirect_to_log(logger):
    originals = sys.stdout, sys.stderr
    sys.stdout = LogAdapter(level='INFO', logger=logger)
    sys.stderr = LogAdapter(level='ERROR', logger=logger)
    yield
    sys.stdout, sys.stderr = originals

def worker_init_logging(logger_name):
    """子进程初始化函数：配置日志并重定向输出"""
    logger = logging.getLogger(logger_name)
    # 避免重复添加handler
    if not logger.handlers:
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(process)d | %(asctime)s | %(name)14s | %(levelname)7s | %(message)s",
        )
    # 重定向子进程的stdout/stderr到logging
    sys.stdout = LogAdapter(level='INFO', logger=logger)
    sys.stderr = LogAdapter(level='ERROR', logger=logger)

def test_case():
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.utils import parallel_backend
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(process)d | %(asctime)s | %(name)14s | %(levelname)7s | %(message)s",
    )

    for backend_name in ['loky','threading']:
        logger.info(f"Testing backend {backend_name}")
        # 针对loky后端传递子进程初始化参数
        backend_params = {}
        if backend_name == 'loky':
            backend_params = {
                'worker_initializer': worker_init_logging,
                'worker_initargs': (__name__,)
            }
        
        with parallel_backend(backend_name, backend_params=backend_params), redirect_to_log(logger):
            clf = RandomForestClassifier(2, verbose=4)
            X = [[0, 0], [1, 1]]
            Y = [0, 1]
            clf = clf.fit(X, Y)

if __name__ == "__main__":
    test_case()

效果说明 #

修改后，loky子进程的输出会被正确捕获并通过logging模块输出，每条日志会显示子进程的PID，示例输出如下：

19320 | 2022-11-30 17:55:20,123 |       __main__ |    INFO | Testing backend loky
19320 | 2022-11-30 17:55:20,135 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
19321 | 2022-11-30 17:55:20,241 |       __main__ |    INFO | building tree 1 of 2
19322 | 2022-11-30 17:55:20,242 |       __main__ |    INFO | building tree 2 of 2
19320 | 2022-11-30 17:55:22,110 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    1.9s remaining:    0.0s
19320 | 2022-11-30 17:55:22,110 |       __main__ |   ERROR | [Parallel(n_jobs=-1)]: Done   2 out of   2 | elapsed:    1.9s finished
...

内容的提问来源于stack exchange，提问作者LudvigH