[DistDGL related] When and how to fetch features on the remote machine

tonyjie · February 4, 2022, 4:30am

I’m looking into DistDGL’s implementation on neighbor sampling. I found that it would first check if required nodes are on the local machine. If it is a local node, it would call sample_neighbor function on the local node. Otherwise, it would send a SampleRequest to the remote machines. Then they would be merged together to form a sampled graph. (If there’s any mistake of my understanding, please point out! Thanks. )

But my confusion is about how to fetch node/edge features on remote machines. I notice the SampleRequest would return the node/edge structure & id, but I don’t find the code for sending feature message. Would the server send feature as well with the ID? Or the feature is fetched in other way?

Rhett-Ying · February 7, 2022, 7:31am

Features are stored/fetched in/via DistTensor in DistDGL. please refer to below link for more details: dgl.distributed — DGL 0.8 documentation

tonyjie · February 16, 2022, 5:31am

I read the detailed description in DistTensor, but I’m still unclear on how one machine fetches the DistTensor which stores on the other machine. Could you please clarify a little more?

I know that features are stored in DistTensor data structures. It is sharded and stored in a cluster of machines, which means that each machine hold one part of feature matrix (for the nodes in this part). So when Machine 0 needs the feature DistTensor from machine 1, how does it fetch DistTensor? Is it still like RPC request, similar with what I have found for communications on graph structures, or there are some other methods working on that?

Thanks

Rhett-Ying · February 16, 2022, 5:59am

yes, remote data is fetched via RPC. the details of impl is a bit complicated. In short, local ids will be converted to global ids and send request to target machines which stores the target data, receive response. please refer to below code/call path for more details.

github.com

dmlc/dgl/blob/42f8c8f3357b4f8eca63da44c75ccf47774fa43a/python/dgl/distributed/dist_tensor.py#L175

      
        
            
            
def __del__(self):
                initialized = os.environ.get('DGL_DIST_MODE', 'standalone') == 'standalone' \
                        or is_initialized()
                if not self._persistent and self._owner and initialized:
                    self.kvstore.delete_data(self._name)
            
            
def __getitem__(self, idx):
                idx = utils.toindex(idx)
                idx = idx.tousertensor()
                return self.kvstore.pull(name=self._name, id_tensor=idx)
            
            
def __setitem__(self, idx, val):
                idx = utils.toindex(idx)
                idx = idx.tousertensor()
                # TODO(zhengda) how do we want to support broadcast (e.g., G.ndata['h'][idx] = 1).
                self.kvstore.push(name=self._name, id_tensor=idx, data_tensor=val)
            
            
def __len__(self):
                return self._shape[0]

github.com

dmlc/dgl/blob/42f8c8f3357b4f8eca63da44c75ccf47774fa43a/python/dgl/distributed/kvstore.py#L1284

      
        
                        # communication-local_push here
                        local_id = self._part_policy[name].to_local(partial_id)
                        local_data = partial_data
                    else: # push data to remote server
                        request = PushRequest(name, partial_id, partial_data)
                        rpc.send_request_to_machine(machine_idx, request)
                    start += count[idx]
                if local_id is not None: # local push
                    self._push_handlers[name](self._data_store, name, local_id, local_data)
            
            
def pull(self, name, id_tensor):
                """Pull message from KVServer.
            
            
    Parameters
                ----------
                name : str
                    data name
                id_tensor : tensor
                    a vector storing the ID list
            
            
    Returns

tonyjie · February 16, 2022, 5:00pm

Got it! The last question is that when do we fetch remote data via RPC? Does it happen when forming blocks, or when doing forward computation for one batch, or something else?
Thanks for clarification.

Rhett-Ying · February 17, 2022, 1:05am

when tensor data is accessed, namely __getitem__ is called such as g.ndata['h'][0:10] or dgl/train_dist.py at 39121dfdb80e36620fa62517da020f7a8c12e51a · dmlc/dgl · GitHub

system · March 19, 2022, 1:06am

This topic was automatically closed 30 days after the last reply. New replies are no longer allowed.