def test_wrap_bad(self):
none = ""
stdout_log = os.path.join(self.test_dir, "stdout.log")
stderr_log = os.path.join(self.test_dir, "stderr.log")
redirs = [
(none, none),
(none, stderr_log),
(stdout_log, none),
(stdout_log, stderr_log),
]
for stdout_redir, stderr_redir in redirs:
queue = multiprocessing.SimpleQueue()
worker_finished_event_mock = mock.Mock()
_wrap(
local_rank=0,
fn=echo1,
args={0: ("hello",)},
envs={0: {"RANK": "0"}},
stdout_redirects={0: stdout_redir},
stderr_redirects={0: stderr_redir},
ret_vals={0: queue},
queue_finished_reading_event=worker_finished_event_mock,
)
self.assertEqual("hello_0", queue.get())
if stdout_redir:
self.assert_in_file(["hello stdout from 0"], stdout_log)
if stderr_redir:
self.assert_in_file(["hello stderr from 0"], stderr_log)
worker_finished_event_mock.wait.assert_called_once()
def test_wrap(self):
none = ""
stdout_log = os.path.join(self.test_dir, "stdout.log")
stderr_log = os.path.join(self.test_dir, "stderr.log")
redirs = [
(none, none),
(none, stderr_log),
(stdout_log, none),
(stdout_log, stderr_log),
]
for stdout_redir, stderr_redir in redirs:
queue = multiprocessing.SimpleQueue()
_wrap(
local_rank=0,
fn=echo1,
args={0: ("hello",)},
envs={0: {"RANK": "0"}},
stdout_redirects={0: stdout_redir},
stderr_redirects={0: stderr_redir},
ret_vals={0: queue},
)
self.assertEqual("hello_0", queue.get())
if stdout_redir:
self.assert_in_file(["hello stdout from 0"], stdout_log)
if stderr_redir:
self.assert_in_file(["hello stderr from 0"], stderr_log)
