Concurrently Issue All `dist_feature.async_get()`Requests in Base Sampler `_collate_fn()`

[jchmura-sc]

TLDR
In our samplers, we should consider firing off all async_get() requests to all distributed feature stores in parallel, and awaiting them all at once when building the SampleMessage. Currently, we await them serially, potentially limiting concurrency.

Description
In the base sampler _collate_fn, we issue async_get() calls on each distributed feature store to retrieve the tensors required for constructing a SampleMessage. There is a synchronous version (self.use_all2all=True) which we don't use so is not relevant here.

For each feature store (e.g. node features, edge features, node labels, ...) we essentially do something like:

# In base_sampler._collate_fn()
for feat in self.distributed_features: # e.g. node_feats, edge_feats, ...
    fut_dict = {}
    for k_type, k in output.items(): # looping over node types / edge types
         fut_dict[k_type] = feat.async_get(k, k_type)
    for k_type, fut in fut_dict.items():
        result[...] = await wrap_torch_tensor(fut)

I suspect that we might be able to get slightly better feature fetch throughput from more concurrency by issuing all async_get() requests in parallel, across feature stores. Conceptually:

# In base_sampler._collate_fn()
fut_dict = {}
for feat in self.distributed_features: # e.g. node_feats, edge_feats, ...
    for k_type, k in output.items(): # looping over node types / edge types
         fut_dict[(feat_type, k_type)] = feat.async_get(k, k_type)

for (feat_type, k_type), fut in fut_dict.items():
    result[...] = await wrap_torch_tensor(fut)

Scope
Change is minimal. Testing on a workload with multiple large feature stores (i.e. edge features) would help. Quantization work might need this to hide overhead of managing a new quantized node feature store