def _wrap(
local_rank: int,
fn: Callable,
args: Dict[int, Tuple],
envs: Dict[int, Dict[str, str]],
stdout_redirects: Dict[
int, str], # redirect file for stdout (to console if None)
stderr_redirects: Dict[
int, str], # redirect file for stderr (to console if None)
ret_vals: Dict[int, mp.SimpleQueue],
queue_finished_reading_event: synchronize.Event,
) -> None:
# get the per-rank params up front so we fail fast if no mapping is found
args_ = args[local_rank]
env_ = envs[local_rank]
ret_val_ = ret_vals[local_rank]
stdout_rd = stdout_redirects[local_rank]
stderr_rd = stderr_redirects[local_rank]
stdout_cm = redirect_stdout(stdout_rd) if stdout_rd else _nullcontext()
stderr_cm = redirect_stderr(stderr_rd) if stderr_rd else _nullcontext()
for k, v in env_.items():
os.environ[k] = v
with stdout_cm, stderr_cm:
ret = record(fn)(*args_)
ret_val_.put(ret)
queue_finished_reading_event.wait()
def test_redirect_stdout(self):
stdout_log = os.path.join(self.test_dir, "stdout.log")
# printing to stdout before redirect should go to console not stdout.log
print("foo first from python")
libc.printf(b"foo first from c\n")
os.system("echo foo first from cmd")
with redirect_stdout(stdout_log):
print("foo from python")
libc.printf(b"foo from c\n")
os.system("echo foo from cmd")
# make sure stdout is restored
print("foo again from python")
libc.printf(b"foo again from c\n")
os.system("echo foo again from cmd")
with open(stdout_log, "r") as f:
# since we print from python, c, cmd -> the stream is not ordered
# do a set comparison
lines = set(f.readlines())
self.assertEqual(
{"foo from python\n", "foo from c\n", "foo from cmd\n"}, lines
)
def test_redirect_both(self):
stdout_log = os.path.join(self.test_dir, "stdout.log")
stderr_log = os.path.join(self.test_dir, "stderr.log")
print("first stdout from python")
libc.printf(b"first stdout from c\n")
print("first stderr from python", file=sys.stderr)
libc.fprintf(c_stderr, b"first stderr from c\n")
with redirect_stdout(stdout_log), redirect_stderr(stderr_log):
print("redir stdout from python")
print("redir stderr from python", file=sys.stderr)
libc.printf(b"redir stdout from c\n")
libc.fprintf(c_stderr, b"redir stderr from c\n")
print("again stdout from python")
libc.fprintf(c_stderr, b"again stderr from c\n")
with open(stdout_log, "r") as f:
lines = set(f.readlines())
self.assertEqual(
{"redir stdout from python\n", "redir stdout from c\n"}, lines
)
with open(stderr_log, "r") as f:
lines = set(f.readlines())
self.assertEqual(
{"redir stderr from python\n", "redir stderr from c\n"}, lines
)
def _redirect_large_buffer(self, print_fn, num_lines=500_000):
stdout_log = os.path.join(self.test_dir, "stdout.log")
with redirect_stdout(stdout_log):
for i in range(num_lines):
print_fn(i)
with open(stdout_log) as fp:
actual = {int(line.split(":")[1]) for line in fp.readlines()}
expected = set(range(num_lines))
self.assertSetEqual(expected, actual)