Function _graph_debug

python/oneflow/test/graph/test_graph_debug.py:29–57

(test_case, v_level=0, ranks=None, max_py_stack_depth=2)

def _graph_debug(test_case, v_level=0, ranks=None, max_py_stack_depth=2):
    class DebugGraph(flow.nn.Graph):
        def __init__(self):
            super().__init__()
            self.m = flow.nn.Linear(3, 3)

        def build(self, x):
            return x

    d_g = DebugGraph()
    d_g.debug(v_level, ranks=ranks, max_py_stack_depth=max_py_stack_depth)

    if ranks is None:
        rank_list = [0]
    elif isinstance(ranks, int):
        rank_list = [ranks]
    elif isinstance(ranks, list):
        rank_list = ranks

    if (
        -1 in rank_list or rank in rank_list
    ) and v_level >= 0:  # v_level == -1 means debug mode is closed
        test_case.assertTrue(d_g._debug)
        test_case.assertTrue(d_g.m.to(GraphModule)._debug)
        print(f"ranks {ranks} rank {rank} debug is opened.")
    else:
        test_case.assertTrue(not d_g._debug)
        test_case.assertTrue(not d_g.m.to(GraphModule)._debug)
        print(f"ranks {ranks} rank {rank} debug is closed.")


@unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")