[Dist] format dtypes when loading graph in server (#4228)

* [Dist] format dtypes when loading graph in server

* add test

* refine

* add comments

python/dgl/distributed/dist_graph.py

@@ -336,6 +336,17 @@ class DistGraphServer(KVServer):
             self.client_g, _, _, self.gpb, graph_name, \
                     ntypes, etypes = load_partition(part_config, self.part_id, load_feats=False)
             print('load ' + graph_name)
+            # formatting dtype
+            # TODO(Rui) Formatting forcely is not a perfect solution.
+            #   We'd better store all dtypes when mapping to shared memory
+            #   and map back with original dtypes.
+            for k, dtype in FIELD_DICT.items():
+                if k in self.client_g.ndata:
+                    self.client_g.ndata[k] = F.astype(
+                        self.client_g.ndata[k], dtype)
+                if k in self.client_g.edata:
+                    self.client_g.edata[k] = F.astype(
+                        self.client_g.edata[k], dtype)
             # Create the graph formats specified the users.
             self.client_g = self.client_g.formats(graph_format)
             self.client_g.create_formats_()

tests/distributed/test_dist_graph_store.py

@@ -35,6 +35,15 @@ def run_server(graph_name, server_id, server_count, num_clients, shared_mem, kee
                         disable_shared_mem=not shared_mem,
                         graph_format=['csc', 'coo'], keep_alive=keep_alive)
     print('start server', server_id)
+    # verify dtype of underlying graph
+    cg = g.client_g
+    for k, dtype in dgl.distributed.dist_graph.FIELD_DICT.items():
+        if k in cg.ndata:
+            assert F.dtype(
+                cg.ndata[k]) == dtype, "Data type of {} in ndata should be {}.".format(k, dtype)
+        if k in cg.edata:
+            assert F.dtype(
+                cg.edata[k]) == dtype, "Data type of {} in edata should be {}.".format(k, dtype)
     g.start()
 
 def emb_init(shape, dtype):