Function test_local_sampling_heterograph

tests/distributed/test_distributed_sampling.py:1766–1858

(num_parts, use_graphbolt, prob_or_mask)

@pytest.mark.parametrize("use_graphbolt", [False])
@pytest.mark.parametrize("prob_or_mask", ["prob", "mask"])
def test_local_sampling_heterograph(num_parts, use_graphbolt, prob_or_mask):
    reset_envs()
    os.environ["DGL_DIST_MODE"] = "distributed"
    with tempfile.TemporaryDirectory() as test_dir:
        g = create_random_hetero()
        for c_etype in g.canonical_etypes:
            prob = torch.rand(g.num_edges(c_etype))
            mask = prob > 0.2
            prob[torch.randperm(len(prob))[: int(len(prob) * 0.5)]] = 0.0
            g.edges[c_etype].data["prob"] = prob
            g.edges[c_etype].data["mask"] = mask
        graph_name = "test_local_sampling"

        _, orig_eids = partition_graph(
            g,
            graph_name,
            num_parts,
            test_dir,
            num_hops=1,
            part_method="metis",
            return_mapping=True,
            use_graphbolt=use_graphbolt,
            store_eids=True,
            store_inner_node=True,
            store_inner_edge=True,
        )

        part_config = os.path.join(test_dir, f"{graph_name}.json")
        for part_id in range(num_parts):
            local_g, _, edge_feats, gpb, _, _, _ = load_partition(
                part_config,
                part_id,
                load_feats=True,
                use_graphbolt=use_graphbolt,
            )
            inner_global_nids = [
                gpb.map_to_homo_nid(gpb.partid2nids(part_id, ntype), ntype)
                for ntype in gpb.ntypes
            ]
            inner_global_nids = torch.cat(inner_global_nids)
            inner_global_eids = {
                c_etype: gpb.partid2eids(part_id, c_etype)
                for c_etype in gpb.canonical_etypes
            }
            inner_node_data = (
                local_g.node_attributes["inner_node"]
                if use_graphbolt
                else local_g.ndata["inner_node"]
            )
            inner_edge_data = (
                local_g.edge_attributes["inner_edge"]
                if use_graphbolt
                else local_g.edata["inner_edge"]
            )
            assert len(inner_global_nids) == inner_node_data.sum()
            num_inner_global_eids = sum(
                [len(eids) for eids in inner_global_eids.values()]
            )