示例#1
0
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()
示例#2
0
    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
            )
示例#3
0
    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
            )
示例#4
0
    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)