[DataLoader] User-defined Dataloader Problem

[jalencato]

🐛 Bug

GraphStorm is trying to use own-defined sampler to create one dataloader. The code example should look like:

        sampler = MultiLayerNeighborSamplerForReconstruct(sampler,
                    dataset, construct_feat_ntype, construct_feat_fanout)
        loader = dgl.dataloading.DistNodeDataLoader(g, target_idx, sampler,
            batch_size=batch_size, shuffle=train_task)

The error log is like:

[rank3]: Traceback (most recent call last):
[rank3]:   File "/graphstorm/python/graphstorm/run/gsgnn_np/gsgnn_np.py", line 209, in <module>
[rank3]:     main(gs_args)
[rank3]:   File "/graphstorm/python/graphstorm/run/gsgnn_np/gsgnn_np.py", line 118, in main
[rank3]:     dataloader = GSgnnNodeDataLoader(train_data, train_idxs,
[rank3]:   File "/graphstorm/python/graphstorm/dataloading/dataloading.py", line 1645, in __init__
[rank3]:     self.dataloader = self._prepare_dataloader(dataset,
[rank3]:   File "/graphstorm/python/graphstorm/dataloading/dataloading.py", line 1673, in _prepare_dataloader
[rank3]:     loader = dgl.dataloading.DistNodeDataLoader(g, target_idx, sampler,
[rank3]:   File "/opt/gs-venv/lib/python3.10/site-packages/dgl/dataloading/dist_dataloader.py", line 641, in __init__
[rank3]:     self.collator = NodeCollator(g, nids, graph_sampler, **collator_kwargs)
[rank3]:   File "/opt/gs-venv/lib/python3.10/site-packages/dgl/dataloading/dist_dataloader.py", line 206, in __init__
[rank3]:     Collator.add_edge_attribute_to_graph(self.g, self.graph_sampler.prob)
[rank3]: AttributeError: 'MultiLayerNeighborSamplerForReconstruct' object has no attribute 'prob'

Expected behavior

The whole process works good with dgl 1.1.3, but failed for dgl >= 2.3.0. (Perhaps 2.0+ version). The line which causes problem is here: https://github.com/dmlc/dgl/blob/d650422402aa770fbf7ec05e7230405c45a3bdfa/python/dgl/dataloading/dist_dataloader.py#L205.

Environment

• DGL Version (e.g., 1.0): 2.3.0 + GPU
• Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): Torch 2.3.0, CUDA 12.1
• OS (e.g., Linux): Amazon Linux
• How you installed DGL (conda, pip, source): pip
• Build command you used (if compiling from source): NA
• Python version: 3.10.0
• CUDA/cuDNN version (if applicable): 12.1
• GPU models and configuration (e.g. V100): AWS Instance G5.48xlarge
• Any other relevant information:

Additional context

[rudongyu]

Thanks for reporting the bug. As a temporary workaround, you may add an empty attribute to the customized sampler as self.prob=None.