[Draft] Testing inference/remap optimizations

python/graphstorm/gconstruct/id_map.py

@@ -12,23 +12,19 @@
     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     See the License for the specific language governing permissions and
     limitations under the License.
-
-    Generate example graph data using built-in datasets for node classification,
-    node regression, edge classification and edge regression.
 """
-import os
 import logging
+import os
+import time
+
+import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
-import numpy as np
 
 from graphstorm.data.constants import (
-    GSP_MAPPING_INPUT_ID,
-    GSP_MAPPING_OUTPUT_ID,
     MAPPING_INPUT_ID,
     MAPPING_OUTPUT_ID,
 )
-from .file_io import read_data_parquet
 from .utils import ExtMemArrayWrapper
 
 GIB_BYTES = 1024**3
@@ -85,23 +81,43 @@ class IdReverseMap:
             Id mapping file prefix
     """
     def __init__(self, id_map_prefix):
+        load_data_start = time.perf_counter()
         assert os.path.exists(id_map_prefix), \
             f"{id_map_prefix} does not exist."
-        try:
-            data = read_data_parquet(id_map_prefix, [MAPPING_INPUT_ID, MAPPING_OUTPUT_ID])
-        except AssertionError:
-            # To maintain backwards compatibility with GraphStorm v0.2.1
-            data = read_data_parquet(id_map_prefix, [GSP_MAPPING_INPUT_ID, GSP_MAPPING_OUTPUT_ID])
-            data[MAPPING_OUTPUT_ID] = data[GSP_MAPPING_OUTPUT_ID]
-            data[MAPPING_INPUT_ID] = data[GSP_MAPPING_INPUT_ID]
-            data.pop(GSP_MAPPING_INPUT_ID)
-            data.pop(GSP_MAPPING_OUTPUT_ID)
-
-        sort_idx = np.argsort(data[MAPPING_OUTPUT_ID])
-        self._ids = data[MAPPING_INPUT_ID][sort_idx]
+
+        mapping_files = os.listdir(id_map_prefix)
+
+        col_names = None
+        for filename in mapping_files:
+            if filename.endswith(".parquet"):
+                col_names = pq.read_metadata(os.path.join(id_map_prefix, filename)).schema.names
+                break
+        assert col_names is not None, \
+            f"No parquet file found in id map directory {id_map_prefix}"
+
+        if "node_str_id" in col_names:
+            map_schema = pa.schema(
+                    [("node_str_id", pa.large_string()),
+                    ("node_int_id", pa.int64())]
+                )
+        else:
+            map_schema = pa.schema([("orig", pa.large_string()), ("new", pa.int64())])
+
+        mapping_table = pq.read_table(
+            id_map_prefix,
+            memory_map=True,
+            schema=map_schema)  # type: pa.Table
+        if "node_str_id" in mapping_table.column_names:
+            mapping_table = mapping_table.rename_columns(["orig", "new"])
+
+            logging.info("Time to load id data: %f", time.perf_counter() - load_data_start)
+
+        sort_ids_start = time.perf_counter()
+        self._ordered_raw_ids = mapping_table.sort_by("new")['orig'].to_numpy()
+        logging.info("Time to sort id data: %f", time.perf_counter() - sort_ids_start)
 
     def __len__(self):
-        return len(self._ids)
+        return len(self._ordered_raw_ids)
 
     def map_range(self, start, end):
         """ Map a range of GraphStorm IDs to the raw IDs.
@@ -117,24 +133,25 @@ def map_range(self, start, end):
         -------
         tensor: A numpy array of raw IDs.
         """
-        return self._ids[start:end]
+        return self._ordered_raw_ids[start:end]
 
-    def map_id(self, ids):
-        """ Map the GraphStorm IDs to the raw IDs.
+    def map_id(self, dgl_ids: np.ndarray) -> np.ndarray:
+        """Maps GraphStorm IDs to raw IDs.
 
         Parameters
         ----------
-        ids : numpy array
-            The input IDs
+        ids : np.ndarray
+            The input Graphstorm/DGL IDs
 
         Returns
         -------
-        tensor: A numpy array of raw IDs.
+        np.ndarray
+            A numpy array of raw IDs.
         """
-        if len(ids) == 0:
+        if len(dgl_ids) == 0:
             return np.array([], dtype=np.str)
 
-        return self._ids[ids]
+        return self._ordered_raw_ids[dgl_ids]
 
 class IdMap:
     """ Map an ID to a new ID.

python/graphstorm/gconstruct/remap_result.py

@@ -26,11 +26,14 @@
 import sys
 import math
 from functools import partial
+from typing import Callable, Dict
 
+from joblib import Parallel, delayed
+import numpy as np
 import pandas as pd
 import torch as th
+
 from ..model.utils import pad_file_index
-from .file_io import write_data_parquet
 from .id_map import IdReverseMap
 from ..utils import get_log_level
 from .utils import multiprocessing_exec_no_return as multiprocessing_remap
@@ -58,14 +61,15 @@
                          GS_REMAP_DST_NID_COL,
                          GS_REMAP_EMBED_COL]
 
+
 # Id_maps is a global variable.
 # When using multi-processing to do id remap,
 # we do not want to pass id_maps to each worker process
 # through argument which uses Python pickle to copy
 # data. By making id_maps as a global variable, we
 # can rely on Linux copy-on-write to provide a zero-copy
 # id_maps to each worker process.
-id_maps = {}
+id_maps = {} # type: Dict[str, IdReverseMap]
 
 def write_data_parquet_file(data, file_prefix, col_name_map=None):
     """ Write data into disk using parquet format.
@@ -89,7 +93,7 @@ def write_data_parquet_file(data, file_prefix, col_name_map=None):
         data = updated_data
 
     output_fname = f"{file_prefix}.parquet"
-    write_data_parquet(data, output_fname)
+    pd.DataFrame.from_dict(data).to_parquet(output_fname)
 
 def write_data_csv_file(data, file_prefix, delimiter=",", col_name_map=None):
     """ Write data into disk using csv format.
@@ -148,40 +152,74 @@ def write_data_csv_file(data, file_prefix, delimiter=",", col_name_map=None):
     data_frame = pd.DataFrame(csv_data)
     data_frame.to_csv(output_fname, index=False, sep=delimiter)
 
+
 def worker_remap_node_data(data_file_path, nid_path, ntype, data_col_key,
-    output_fname_prefix, chunk_size, output_func):
+    output_fname_prefix, chunk_size, output_func: Callable[[Dict, str], None]):
     """ Do one node prediction remapping task
 
         Parameters
         ----------
         data_file_path: str
             The path to the node data.
         nid_path: str
-            The path to the file storing node ids
+            The path to the file storing DGL/GraphStorm node ids
         ntype: str
             Node type.
         data_col_key: str
             Column key of the node data
         output_fname_prefix: str
             Output file name prefix.
         chunk_size: int
-            Max number of raws per output file.
-        output_func: func
-            Function used to write data to disk.
+            Max number of node ids per output file.
+        output_func: Callable[[Dict, str], None]
+            Function used to write data dictionary to disk.
+            The function must accept two arguments, its first argument a dict from
+            column name(s) to an array-like, the second argument must be
+            filepath string.
     """
+    # rank = get_rank()
     node_data = th.load(data_file_path).numpy()
-    nids = th.load(nid_path).numpy()
-    nid_map = id_maps[ntype]
+    dgl_ids = th.load(nid_path).numpy()
+    # nid_map = id_maps[ntype] # type: IdReverseMap
     num_chunks = math.ceil(len(node_data) / chunk_size)
 
-    for i in range(num_chunks):
-        start = i * chunk_size
-        end = (i + 1) * chunk_size if i + 1 < num_chunks else len(node_data)
-        data = node_data[start:end]
-        nid = nid_map.map_id(nids[start:end])
-        data = {data_col_key: data,
-                GS_REMAP_NID_COL: nid}
-        output_func(data, f"{output_fname_prefix}_{pad_file_index(i)}")
+    data_chunks = np.array_split(node_data, num_chunks)
+    dgl_ids_chunks = np.array_split(dgl_ids, num_chunks)
+
+    def thread_remap_node_data(i, node_data_chunk, dgl_ids_chunk):
+        # chunk_start = time.perf_counter()
+        nid_map = id_maps[ntype] # type: IdReverseMap
+
+        output_func(
+            {
+                data_col_key: node_data_chunk.tolist(),
+                GS_REMAP_NID_COL: nid_map.map_id(dgl_ids_chunk).tolist()
+            },
+            f"{output_fname_prefix}_{pad_file_index(i)}"
+        )
+        # logging.info("Rank %d: Finished remapping chunk %d/%d in %.2f seconds.",
+        #                 rank, i+1, num_chunks, time.perf_counter() - chunk_start)
+
+    Parallel(n_jobs=int(os.cpu_count()/4), backend="threading", verbose=1)(
+        delayed(thread_remap_node_data)(
+            i, data_chunks[i], dgl_ids_chunks[i]) for i in range(num_chunks)
+    )
+
+    # for i in range(num_chunks):
+    #     chunk_start = time.perf_counter()
+    #     start = i * chunk_size
+    #     end = (i + 1) * chunk_size if i + 1 < num_chunks else len(node_data)
+
+    #     output_func(
+    #         {
+    #             data_col_key: node_data[start:end].tolist(),
+    #             GS_REMAP_NID_COL: nid_map.map_id(dgl_ids[start:end]).tolist()
+    #         },
+    #         f"{output_fname_prefix}_{pad_file_index(i)}"
+    #     )
+    #     logging.info("Rank %d: Finished remapping chunk %d/%d in %.2f seconds.",
+    #                  rank, i, num_chunks, time.perf_counter() - chunk_start)
+
 
 def worker_remap_edge_pred(pred_file_path, src_nid_path,
     dst_nid_path, src_type, dst_type,
@@ -286,7 +324,7 @@ def _remove_inputs(with_shared_fs, files_to_remove,
 def remap_node_emb(emb_ntypes, node_emb_dir,
                    output_dir, out_chunk_size,
                    num_proc, rank, world_size,
-                   with_shared_fs, output_func):
+                   with_shared_fs, output_func,):
     """ Remap node embeddings.
 
         The function will iterate all the node types that
@@ -297,7 +335,7 @@ def remap_node_emb(emb_ntypes, node_emb_dir,
 
         Example
         --------
-        # embedddings:
+        # embeddings:
         #   ntype0:
         #     embed_nids-00000.pt
         #     embed_nids-00001.pt
@@ -317,7 +355,7 @@ def remap_node_emb(emb_ntypes, node_emb_dir,
 
         Example
         --------
-        # embedddings:
+        # embeddings:
         #   ntype0:
         #     embed-00000_00000.parquet
         #     embed-00000_00001.parquet
@@ -344,7 +382,7 @@ def remap_node_emb(emb_ntypes, node_emb_dir,
         world_size: int
             The total number of processes in the cluster.
         with_shared_fs: bool
-            Whether shared file system is avaliable.
+            Whether shared file system is available.
         output_func: func
             Function used to write data to disk.
 
@@ -353,11 +391,13 @@ def remap_node_emb(emb_ntypes, node_emb_dir,
         list of str
             The list of files to be removed.
     """
+    embedding_remap_start = time.perf_counter()
     task_list = []
     files_to_remove = []
     for ntype in emb_ntypes:
         input_emb_dir = os.path.join(node_emb_dir, ntype)
         out_embdir = os.path.join(output_dir, ntype)
+        os.makedirs(out_embdir, exist_ok=True)
         ntype_emb_files = os.listdir(input_emb_dir)
         # please note nid_files can be empty.
         nid_files = [fname for fname in ntype_emb_files \
@@ -404,7 +444,10 @@ def remap_node_emb(emb_ntypes, node_emb_dir,
                 "output_func": output_func,
             })
 
+    logging.info("Rank %d: Length of task list: %d", rank, len(task_list))
     multiprocessing_remap(task_list, num_proc, worker_remap_node_data)
+    logging.info("Rank %d: Remap node embeddings done. Time elapsed: %f",
+                 rank, time.perf_counter() - embedding_remap_start)
     return files_to_remove
 
 def remap_node_pred(pred_ntypes, pred_dir,
@@ -715,6 +758,7 @@ def main(args, gs_config_args):
     rank = args.rank
     world_size = args.world_size
     with_shared_fs = args.with_shared_fs
+    remap_program_start = time.perf_counter()
 
     if args.yaml_config_file is not None:
         # Case 1: remap_result is called right after the
@@ -949,6 +993,8 @@ def main(args, gs_config_args):
                         "Embeddings will remain in PyTorch format.")
         sys.exit(0)
 
+    logging.info("Retrieving id_maps from %s", id_mapping_path)
+    id_map_start = time.perf_counter()
     for ntype in set(ntypes):
         mapping_prefix = os.path.join(id_mapping_path, ntype)
         logging.debug("loading mapping file %s",
@@ -962,8 +1008,11 @@ def main(args, gs_config_args):
                  "Embeddings will remain in PyTorch format."),
                 mapping_prefix)
             sys.exit(0)
+    logging.info(
+        "Rank %d: Retrieving id_maps took %f seconds", rank, time.perf_counter() - id_map_start)
 
     num_proc = args.num_processes if args.num_processes > 0 else 1
+    logging.info("Number of processes %d", num_proc)
     col_name_map = None
     if args.column_names is not None:
         col_name_map = {}
@@ -1073,6 +1122,9 @@ def main(args, gs_config_args):
         _remove_inputs(with_shared_fs, files_to_remove, rank, world_size,
                        node_emb_dir if node_emb_dir is not None else predict_dir)
 
+    logging.info("Rank %d: Remapping program finished in %.4f seconds",
+                 rank, time.perf_counter() - remap_program_start)
+
 def add_distributed_remap_args(parser):
     """ Distributed remapping only
 

python/graphstorm/model/utils.py

@@ -18,6 +18,7 @@
 import os
 import math
 import json
+import pickle
 import shutil
 import logging
 
@@ -884,15 +885,17 @@ def save_pytorch_embeddings(emb_path, embeddings, rank, world_size,
         for name, emb in embeddings.items():
             os.makedirs(os.path.join(emb_path, name), exist_ok=True)
             th.save(emb, os.path.join(os.path.join(emb_path, name),
-                                      f'embed-{pad_file_index(rank)}.pt'))
+                                      f'embed-{pad_file_index(rank)}.pt'),
+                                      pickle_protocol=pickle.HIGHEST_PROTOCOL,)
             emb_info["emb_name"].append(name)
             emb_info["emb_dim"][name] = emb.shape[1]
     else:
         os.makedirs(os.path.join(emb_path, NTYPE), exist_ok=True)
         # There is no ntype for the embedding
         # use NTYPE
         th.save(embeddings, os.path.join(os.path.join(emb_path, NTYPE),
-                                         f'embed-{pad_file_index(rank)}.pt'))
+                                         f'embed-{pad_file_index(rank)}.pt'),
+                                         pickle_protocol=pickle.HIGHEST_PROTOCOL,)
         emb_info["emb_name"] = NTYPE
         emb_info["emb_dim"][NTYPE] = embeddings.shape[1]
 
@@ -999,9 +1002,11 @@ def save_shuffled_node_embeddings(shuffled_embs, save_embed_path, save_embed_for
         assert len(nids) == len(embs), \
             f"The embeding length {len(embs)} does not match the node id length {len(nids)}"
         th.save(embs, os.path.join(os.path.join(save_embed_path, ntype),
-                                  f'embed-{pad_file_index(rank)}.pt'))
+                                  f'embed-{pad_file_index(rank)}.pt'),
+                                  pickle_protocol=pickle.HIGHEST_PROTOCOL)
         th.save(nids, os.path.join(os.path.join(save_embed_path, ntype),
-                                  f'embed_nids-{pad_file_index(rank)}.pt'))
+                                  f'embed_nids-{pad_file_index(rank)}.pt'),
+                                  pickle_protocol=pickle.HIGHEST_PROTOCOL)
         emb_info["emb_name"].append(ntype)
 
         if len(embs.shape) == 1: # Embedding is a 1D tensor