要約
- setup.pyでmultiprocessモジュールを含むunittestを実行すると無限ループになる
- どうやらmultiprocessモジュールの仕様らしい
- 回避策として「__main__」を一時的に書き換える
前提
pythonでマルチプロセスで関数を動作させるのは「multiprocess」モジュールを使用します。
使い方やサンプルコードは公式が十分に分かりやすいです。
https://docs.python.org/ja/3/library/multiprocessing.html
さて、次のコードを見てください
from multiprocessing import Pool
def f(x):
return x*x
def main():
with Pool(5) as p:
print(p.map(f, [1, 2, 3]))
main()
公式のサンプルコードから「if __name__ == '__main__':」を省きました。
これを実行するとえらいことになります。
(無限ループに入ってプロセスキルが必要になります)
公式にも記載があるので、これは仕様のようです
https://docs.python.org/ja/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods
ひとまず
multiprocessを使用したコードでは「if __name__ == '__main__':」を省いてはいけない
ということになります。
本題
さて、マルチプロセスを使ったモジュールを作成して、unittestを実行するとします。
from multiprocessing import Pool
def f(x):
return x*x
def main():
with Pool(5) as p:
print(p.map(f, [1, 2, 3]))
if __name__ == '__main__':
pass
"""Tests for `multiprocess_test` package."""
import unittest
from mp_test import mp_test
class TestMultiprocess_test(unittest.TestCase):
"""Tests for `multiprocess_test` package."""
def setUp(self):
"""Set up test fixtures, if any."""
def tearDown(self):
"""Tear down test fixtures, if any."""
def test_000_something(self):
mp_test.main()
if __name__ == "__main__":
unittest.main()
ではテストしてみましょう。
$python test_mp_test.py
これは実行できます。
では、このテスト、setup.pyを経由して実行するとどうなるでしょうか。
mp_test
|-mp_test.py
tests
|-test_mp_test.py
setup.py
#!/usr/bin/env python
"""The setup script."""
from setuptools import setup, find_packages
setup(
author="rr28",
author_email='rr28_yosizumi@hotmail.com',
python_requires='>=3.5',
description="multiprocess test.",
entry_points={
'console_scripts': [
'mp_test=mp_test.mp_test:main',
],
},
name='mp_test',
packages=find_packages(include=['mp_test', 'mp_test.*']),
test_suite='tests',
version='0.1.0',
)
実行
$python setup.py test
そうです。「if name == 'main':」を省いていないにも関わらず、無限ループになります。
実行結果
======================================================================
ERROR: test_000_something (tests.test_mp_test.TestMp_test)
Test something.
----------------------------------------------------------------------
Traceback (most recent call last):
~省略~
RuntimeError:
An attempt has been made to start a new process before the
current process has finished its bootstrapping phase.
This probably means that you are not using fork to start your
child processes and you have forgotten to use the proper idiom
in the main module:
if __name__ == '__main__':
freeze_support()
...
The "freeze_support()" line can be omitted if the program
is not going to be frozen to produce an executable.
----------------------------------------------------------------------
Ran 1 test in 0.007s
FAILED (errors=1)
Test failed: <unittest.runner.TextTestResult run=1 errors=1 failures=0>
test_000_something (tests.test_mp_test.TestMp_test)
Test something. ... ERROR
以下ループ
・
・
・
回避策
setup.pyのテスト実行プロセスがmultiprocessの仕様にあっていないのではとは思うのですが、
行きついた回避策は次のようなものでした
"""Tests for `multiprocess_test` package."""
import unittest
import sys
from mp_test import mp_test
class TestMp_test(unittest.TestCase):
"""Tests for `mp_test` package."""
def setUp(self):
"""Set up test fixtures, if any."""
def tearDown(self):
"""Tear down test fixtures, if any."""
def test_000_something(self):
old_main = sys.modules["__main__"]
old_main_file = sys.modules["__main__"].__file__
sys.modules["__main__"] = sys.modules["mp_test.mp_test"]
sys.modules["__main__"].__file__ = sys.modules["mp_test.mp_test"].__file__
mp_test.main()
sys.modules["__main__"] = old_main
sys.modules["__main__"].__file__ = old_main_file
if __name__ == "__main__":
unittest.main()
実行中の"__main__"を書き換えます。
こうすることで、setup.pyからのunittestも実行可能です。
また、複数のテストケースでマルチプロセスによる処理を実行するのであれば、
専用のテストクラスにしてsetUpとtearDownに記載してもよいでしょう。
import unittest
import sys
from mp_test import mp_test
class TestMp_test(unittest.TestCase):
"""Tests for `mp_test` package."""
def setUp(self):
"""Set up test fixtures, if any."""
self._old_main = sys.modules["__main__"]
self._old_main_file = sys.modules["__main__"].__file__
sys.modules["__main__"] = sys.modules["mp_test.mp_test"]
sys.modules["__main__"].__file__ = sys.modules["mp_test.mp_test"].__file__
def tearDown(self):
"""Tear down test fixtures, if any."""
sys.modules["__main__"] = self._old_main
sys.modules["__main__"].__file__ = self._old_main_file
def test_000_something(self):
"""Test something."""
mp_test.main()
if __name__=="__main__":
unittest.main()
おわりに
toxでテスト環境を構築しているときに気が付いて調査しました。
VSCや単体でのUnittestでは気が付かないので、開発終盤で「わぁぁぁぁぁっぁぁ」となることもあるのではないでしょうか。
ただ、もっとスマートな回避策がありそうなのですが、知っている人がいれば教えてもらえるとありがたいです。