diff --git a/python/graphstorm/gconstruct/id_map.py b/python/graphstorm/gconstruct/id_map.py
index 015d997023..c5183bdd42 100644
--- a/python/graphstorm/gconstruct/id_map.py
+++ b/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.
diff --git a/python/graphstorm/gconstruct/remap_result.py b/python/graphstorm/gconstruct/remap_result.py
index 823e18659d..1d95f10fd8 100644
--- a/python/graphstorm/gconstruct/remap_result.py
+++ b/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,6 +61,7 @@
                          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
@@ -65,7 +69,7 @@
 # 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,8 +152,9 @@ 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
@@ -157,7 +162,7 @@ def worker_remap_node_data(data_file_path, nid_path, ntype, data_col_key,
         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
@@ -165,23 +170,56 @@ def worker_remap_node_data(data_file_path, nid_path, ntype, data_col_key,
         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
 
diff --git a/python/graphstorm/model/utils.py b/python/graphstorm/model/utils.py
index 7b3dbbdea7..937e5d2b2f 100644
--- a/python/graphstorm/model/utils.py
+++ b/python/graphstorm/model/utils.py
@@ -18,6 +18,7 @@
 import os
 import math
 import json
+import pickle
 import shutil
 import logging
 
@@ -884,7 +885,8 @@ 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:
@@ -892,7 +894,8 @@ def save_pytorch_embeddings(emb_path, embeddings, rank, world_size,
         # 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
diff --git a/python/graphstorm/run/launch.py b/python/graphstorm/run/launch.py
index 9c93388995..1e07b770fc 100644
--- a/python/graphstorm/run/launch.py
+++ b/python/graphstorm/run/launch.py
@@ -20,6 +20,7 @@
     python3 -m graphstorm.run.launch <Launch args> YOUR_SCRIPT.py <Train/Infer args>
 """
 import argparse
+import copy
 import json
 import logging
 import multiprocessing
@@ -31,7 +32,6 @@
 import sys
 import time
 import tempfile
-import copy
 
 from functools import partial
 from threading import Thread
@@ -500,10 +500,10 @@ def wrap_dist_remap_command(
         rank: int,
         world_size: int,
         with_shared_fs: bool,
-        num_trainers: int,
         output_chunk_size: int = 100000,
-        preserve_input: bool = False) -> str:
-    """ Wrap distributed remap command
+        preserve_input: bool = False,
+        num_trainers: int = 1) -> str:
+    """ Wrap distributed remap command.
 
         Parameters
         ----------
@@ -515,8 +515,6 @@ def wrap_dist_remap_command(
             Total number of workers
         with_shared_fs:
             Whether all files are stored in a shared fs.
-        num_trainers:
-            Number of trainers on each machine.
         output_chunk_size:
             Number of rows per output file.
         preserve_input:
@@ -537,6 +535,14 @@ def wrap_dist_remap_command(
     new_udf_command[0] = "graphstorm.gconstruct.remap_result"
 
     # Add remap related arguments
+    udf_command.extend([
+        "--rank", str(rank),
+        "--world-size", str(world_size),
+        "--with-shared-fs", "True" if with_shared_fs else "False",
+        "--num-processes", str(num_trainers),
+        "--output-chunk-size", str(output_chunk_size),
+        "--preserve-input", "True" if preserve_input else "False"]
+    )
     new_udf_command += ["--rank", str(rank)]
     new_udf_command += ["--world-size", str(world_size)]
     new_udf_command += ["--with-shared-fs", "True" if with_shared_fs else "False"]
@@ -544,7 +550,7 @@ def wrap_dist_remap_command(
     new_udf_command += ["--output-chunk-size", str(output_chunk_size)]
     new_udf_command += ["--preserve-input", "True" if preserve_input else "False"]
 
-    # transforms the udf_command from:
+    # transforms the new_udf_command from:
     #     path/to/dist_trainer.py arg0 arg1
     # to:
     #     python -m torch.distributed.launch [DIST TORCH ARGS] path/to/dist_trainer.py arg0 arg1
@@ -863,13 +869,13 @@ def submit_remap_jobs(args, udf_command, hosts, run_local):
 
     for node_id, host in enumerate(hosts):
         ip, _ = host
-        remap_dist_command = wrap_dist_remap_command(udf_command,
+        remap_dist_command = wrap_dist_remap_command(copy.deepcopy(udf_command),
                                                      node_id,
                                                      len(hosts),
                                                      args.with_shared_fs,
-                                                     args.num_trainers,
                                                      args.output_chunk_size,
-                                                     args.preserve_input)
+                                                     args.preserve_input,
+                                                     args.num_trainers)
 
         cmd = wrap_cmd_with_local_envvars(remap_dist_command, env_vars)
         cmd = (
@@ -895,11 +901,11 @@ def submit_remap_jobs(args, udf_command, hosts, run_local):
     conn1, conn2 = multiprocessing.Pipe()
     if run_local:
         func = partial(get_all_local_pids, udf_command)
-        process = multiprocessing.Process(target=local_cleanup_proc, args=(func, conn1))
+        cleanup_process = multiprocessing.Process(target=local_cleanup_proc, args=(func, conn1))
     else:
         func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command)
-        process = multiprocessing.Process(target=remote_cleanup_proc, args=(func, conn1))
-    process.start()
+        cleanup_process = multiprocessing.Process(target=remote_cleanup_proc, args=(func, conn1))
+    cleanup_process.start()
 
     def signal_handler(sig, frame): # pylint: disable=unused-argument
         logging.info("Stop launcher")
@@ -920,7 +926,7 @@ def signal_handler(sig, frame): # pylint: disable=unused-argument
 
     # The remap processes complete. We should tell the cleanup process to exit.
     conn2.send("exit")
-    process.join()
+    cleanup_process.join()
     if err != 0:
         logging.error("Remapping task failed")
         sys.exit(-1)
@@ -1131,10 +1137,12 @@ def signal_handler(sig, frame): # pylint: disable=unused-argument
         logging.error("Task failed")
         sys.exit(-1)
 
-    logging.info("Start doing node id remapping")
     if args.do_nid_remap:
+        remap_start = time.perf_counter()
+        logging.info("Start doing node id remapping")
         submit_remap_jobs(args, udf_command, hosts, run_local)
-    logging.info("Finish doing node id remapping")
+        logging.info("Finish doing node id remapping")
+        logging.info("Time to run node id remapping: %f", time.perf_counter() - remap_start)
 
 def get_argument_parser():
     """ Arguments listed here are those used by the launch script to launch
diff --git a/python/graphstorm/sagemaker/sagemaker_infer.py b/python/graphstorm/sagemaker/sagemaker_infer.py
index b97a27a953..e97cecfab0 100644
--- a/python/graphstorm/sagemaker/sagemaker_infer.py
+++ b/python/graphstorm/sagemaker/sagemaker_infer.py
@@ -18,18 +18,20 @@
     We have to put all code in one file.
 """
 # Install additional requirements
-import os
+import json
 import logging
+import os
+import queue
 import socket
-import time
-import json
 import subprocess
-from threading import Thread, Event
 import sys
-import queue
+import time
+from threading import Thread, Event
 
 import boto3
+import botocore
 import sagemaker
+
 from ..config import (BUILTIN_TASK_NODE_CLASSIFICATION,
                       BUILTIN_TASK_NODE_REGRESSION,
                       BUILTIN_TASK_EDGE_CLASSIFICATION,
@@ -44,16 +46,15 @@
                     barrier,
                     terminate_workers,
                     wait_for_exit,
-                    upload_data_to_s3,
                     update_gs_params,
                     download_model,
-                    upload_embs,
-                    remove_embs)
+                    remove_embs,
+                    upload_directory_parallel)
 
-def launch_infer_task(task_type, num_gpus, graph_config,
+def launch_infer_task(task_type, num_trainers, graph_config,
     load_model_path, save_emb_path, ip_list,
     yaml_path, extra_args, state_q, custom_script,
-    output_chunk_size=100000):
+    output_chunk_size=10**6):
     """ Launch SageMaker training task
 
     Parameters
@@ -62,8 +63,8 @@ def launch_infer_task(task_type, num_gpus, graph_config,
         Task type. It can be node classification/regression,
         edge classification/regression, link prediction, etc.
         Refer to graphstorm.config.config.SUPPORTED_TASKS for more details.
-    num_gpus: int
-        Number of gpus per instance
+    num_trainers: int
+        Number of trainers to use per instance
     graph_config: str
         Where does the graph partition config reside.
     load_model_path: str
@@ -82,7 +83,7 @@ def launch_infer_task(task_type, num_gpus, graph_config,
         Custom inference script provided by a customer to run customer inference logic.
     output_chunk_size: int
         Number of rows per chunked prediction result or node embedding file.
-        Default: 100000
+        Default: 10^6
 
     Return
     ------
@@ -108,7 +109,7 @@ def launch_infer_task(task_type, num_gpus, graph_config,
         raise RuntimeError(f"Unsupported task type {task_type}")
 
     launch_cmd = ["python3", "-u",  "-m", cmd,
-        "--num-trainers", f"{num_gpus if int(num_gpus) > 0 else 1}",
+        "--num-trainers", f"{num_trainers if int(num_trainers) > 0 else 1}",
         "--num-servers", "1",
         "--num-samplers", "0",
         "--part-config", f"{graph_config}",
@@ -164,27 +165,18 @@ def run_infer(args, unknownargs):
             customer training logic. Can be None.
         data_path: str
             Local working path.
-        num_gpus: int
-            Number of gpus.
+        num_trainers: int
+            Number of trainer processes to use during inference.
         sm_dist_env: json str
             SageMaker distributed env.
         region: str
             AWS Region.
     """
-    num_gpus = args.num_gpus
-    data_path = args.data_path
-    model_path = '/tmp/gsgnn_model'
-    output_path = '/tmp/infer_output'
-    os.makedirs(model_path, exist_ok=True)
-    os.makedirs(output_path, exist_ok=True)
-
-    # start the ssh server
-    subprocess.run(["service", "ssh", "start"], check=True)
-
-    logging.info("Known args %s", args)
-    logging.info("Unknown args %s", unknownargs)
-
-    train_env = json.loads(args.sm_dist_env)
+    try:
+        with open("/opt/ml/config/resourceconfig.json", "r", encoding="utf-8") as f:
+            train_env = json.load(f)
+    except FileNotFoundError:
+        train_env = json.loads(os.environ['SM_TRAINING_ENV'])
     hosts = train_env['hosts']
     current_host = train_env['current_host']
     world_size = len(hosts)
@@ -197,6 +189,19 @@ def run_infer(args, unknownargs):
         format=f'{current_host}: %(asctime)s - %(levelname)s - %(message)s',
         force=True)
 
+    num_trainers = args.num_trainers
+    data_path = args.data_path
+    model_path = '/tmp/gsgnn_model'
+    output_path = '/tmp/infer_output'
+    os.makedirs(model_path, exist_ok=True)
+    os.makedirs(output_path, exist_ok=True)
+
+    # start the ssh server
+    subprocess.run(["service", "ssh", "start"], check=True)
+
+    logging.info("Known args %s", args)
+    logging.info("Unknown args %s", unknownargs)
+
     try:
         for host in hosts:
             logging.info("The %s IP is %s", host, {socket.gethostbyname(host)})
@@ -228,6 +233,7 @@ def run_infer(args, unknownargs):
         logging.info("Connected")
 
     # write ip list info into disk
+    os.makedirs(data_path, exist_ok=True)
     ip_list_path = os.path.join(data_path, 'ip_list.txt')
     with open(ip_list_path, 'w', encoding='utf-8') as f:
         for host in hosts:
@@ -251,17 +257,26 @@ def run_infer(args, unknownargs):
         update_gs_params(gs_params, "--save-prediction-path", os.path.join(output_path, "predict"))
 
     ### Download Partitioned graph data
+    s3_client = boto3.client(
+        "s3",
+        config=botocore.config.Config(max_pool_connections=150),
+        region_name=args.region
+    )
     boto_session = boto3.session.Session(region_name=args.region)
     sagemaker_session = sagemaker.session.Session(boto_session=boto_session)
+    download_start = time.perf_counter()
     yaml_path = download_yaml_config(infer_yaml_s3,
         data_path, sagemaker_session)
     graph_config_path = download_graph(graph_data_s3, graph_name,
-        host_rank, world_size, data_path, sagemaker_session, args.raw_node_mappings_s3)
+        host_rank, world_size, data_path, sagemaker_session, args.raw_node_mappings_s3,
+        s3_client)
 
     # Download Saved model
     download_model(model_artifact_s3, model_path, sagemaker_session)
     logging.info("Successfully downloaded the model into %s.\n The model files are: %s.",
                  model_path, os.listdir(model_path))
+    logging.info("Rank %d: Time to download all data: %f",
+                 host_rank, time.perf_counter() - download_start)
 
     err_code = 0
     if host_rank == 0:
@@ -278,7 +293,7 @@ def run_infer(args, unknownargs):
             # launch distributed training here
             state_q = queue.Queue()
             train_task = launch_infer_task(task_type,
-                                           num_gpus,
+                                           num_trainers,
                                            graph_config_path,
                                            model_path,
                                            emb_path,
@@ -300,8 +315,11 @@ def run_infer(args, unknownargs):
         terminate_workers(client_list, world_size)
         logging.info("Master End")
         if err_code != -1:
-            upload_embs(output_emb_s3, emb_path, sagemaker_session)
+            upload_emb_start = time.perf_counter()
+            upload_directory_parallel(emb_path, output_emb_s3, s3_client)
             # clean embs, so SageMaker does not need to upload embs again
+            logging.info("Rank %d: Time to upload data: %f",
+                     host_rank, time.perf_counter() - upload_emb_start)
             remove_embs(emb_path)
     else:
         barrier(sock)
@@ -309,12 +327,13 @@ def run_infer(args, unknownargs):
         # Block util training finished
         # Listen to end command
         wait_for_exit(sock)
-        upload_embs(output_emb_s3, emb_path, sagemaker_session)
+        upload_emb_start = time.perf_counter()
+        upload_directory_parallel(emb_path, output_emb_s3, s3_client)
         # clean embs, so SageMaker does not need to upload embs again
+        logging.info("Rank %d: Time to upload data: %f",
+                     host_rank, time.perf_counter() - upload_emb_start)
         remove_embs(emb_path)
-        logging.info("Worker End")
 
-    sock.close()
     if err_code != 0:
         # Report an error
         logging.error("Task failed")
@@ -323,6 +342,8 @@ def run_infer(args, unknownargs):
     if args.output_prediction_s3 is not None:
         # remove tailing /
         output_prediction_s3 = args.output_prediction_s3.rstrip('/')
-        upload_data_to_s3(output_prediction_s3,
-                          os.path.join(output_path, "predict"),
-                          sagemaker_session)
+        upload_directory_parallel(
+            os.path.join(output_path, "predict"),
+            output_prediction_s3,
+            s3_client,
+        )
diff --git a/python/graphstorm/sagemaker/utils.py b/python/graphstorm/sagemaker/utils.py
index 71d51fbe5e..36cd53fb20 100644
--- a/python/graphstorm/sagemaker/utils.py
+++ b/python/graphstorm/sagemaker/utils.py
@@ -15,24 +15,31 @@
 
     sagemaker script utilities
 """
-
 import hashlib
+import math
 import logging
 import os
 import shutil
 import socket
 import subprocess
 import time
-from typing import Optional
+import warnings
+from typing import List, Tuple, Optional
 from urllib.parse import urlparse
 
 import boto3
+import botocore
 from botocore.errorfactory import ClientError
+from joblib import delayed, Parallel
 from sagemaker.s3 import S3Downloader
 from sagemaker.s3 import S3Uploader
 
+from graphstorm import get_rank
+
 PORT_MIN = 10000  # Avoid privileged ports
 PORT_MAX = 65535  # Maximum TCP port number
+DOWNLOAD_THREADS = min(16, os.cpu_count() or 16)
+UPLOAD_THREADS = min(16, os.cpu_count() or 16)
 
 def run(launch_cmd, state_q, env=None):
     """ Running cmd using shell
@@ -216,7 +223,8 @@ def download_model(model_artifact_s3, model_path, sagemaker_session):
 
 def download_graph(graph_data_s3, graph_name, part_id, world_size,
                    local_path, sagemaker_session,
-                   raw_node_mapping_prefix_s3=None):
+                   raw_node_mapping_prefix_s3=None,
+                   s3_client=None):
     """ download graph data
 
     Parameters
@@ -243,6 +251,7 @@ def download_graph(graph_data_s3, graph_name, part_id, world_size,
     """
     # Download partitioned graph data.
     # Each training instance only download 1 partition.
+    rank = get_rank()
     graph_part = f"part{part_id}"
 
     graph_path = os.path.join(local_path, graph_name)
@@ -252,21 +261,12 @@ def download_graph(graph_data_s3, graph_name, part_id, world_size,
 
     graph_data_s3 = graph_data_s3[:-1] if graph_data_s3.endswith('/') else graph_data_s3
 
-    # By default we assume the node mappings exist
-    # under the same path as the rest of the graph data
-    if not raw_node_mapping_prefix_s3:
-        raw_node_mapping_prefix_s3 = f"{graph_data_s3}/raw_id_mappings"
-    else:
-        raw_node_mapping_prefix_s3 = (
-            raw_node_mapping_prefix_s3[:-1] if raw_node_mapping_prefix_s3.endswith('/')
-            else raw_node_mapping_prefix_s3)
-
     # We split on '/' to get the bucket, as it's always the third split element in an S3 URI
     s3_input_bucket = graph_data_s3.split("/")[2]
     # Similarly, by having maxsplit=3 we get the S3 key value as the fourth element
     s3_input_key = graph_data_s3.split("/", maxsplit=3)[3]
 
-    s3_client = boto3.client('s3')
+    s3_client = boto3.client('s3') if s3_client is None else s3_client
     graph_config = None
     for config_name  in [f"{graph_name}.json", "metadata.json"]:
         try:
@@ -289,19 +289,69 @@ def download_graph(graph_data_s3, graph_name, part_id, world_size,
     assert graph_config, \
         (f"Could not find a graph config file named {graph_name}.json or metadata.json "
          f"under {graph_data_s3}")
+    graph_part_start = time.perf_counter()
+    # Download partition metadata file
     S3Downloader.download(os.path.join(graph_data_s3, graph_config),
             graph_path, sagemaker_session=sagemaker_session)
-    try:
-        logging.info("Download graph from %s to %s",
-                     os.path.join(os.path.join(graph_data_s3, graph_part), ""),
-                     graph_part_path)
-        # add tailing / to s3:/xxxx/partN
-        S3Downloader.download(os.path.join(os.path.join(graph_data_s3, graph_part), ""),
-            graph_part_path, sagemaker_session=sagemaker_session)
-    except Exception as err: # pylint: disable=broad-except
-        logging.error("Can not download graph_data from %s, %s.",
-                      graph_data_s3, str(err))
-        raise RuntimeError(f"Can not download graph_data from {graph_data_s3}, {err}.")
+
+    def s3_get_meta_data(client, bucket, key):
+        meta_data = client.head_object(
+            Bucket=bucket,
+            Key=key
+        )
+        return meta_data
+
+    def convert_size(size_bytes):
+        if size_bytes == 0:
+            return "0B"
+        size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
+        i = int(math.floor(math.log(size_bytes, 1024)))
+        denom = math.pow(1024, i)
+        s = round(size_bytes / denom, 2)
+        return "%s %s" % (s, size_name[i])
+
+
+    def get_cunks(size_bytes, desired_sections):
+        return size_bytes / desired_sections
+
+    def download_large_file(client, bucket, key, local_filepath, parallel_threads):
+        start = time.time()
+        s3_md = s3_get_meta_data(client, bucket, key)
+        chunk = get_cunks(s3_md["ContentLength"], parallel_threads)
+        logging.debug("Making %s parallel s3 calls with a chunk size of %s each...",
+            parallel_threads, convert_size(chunk)
+        )
+        client.download_file(
+            Bucket=bucket,
+            Filename=local_filepath,
+            Key=key,
+            Config=boto3.s3.transfer.TransferConfig(
+                max_concurrency=parallel_threads
+            )
+        )
+        end = time.time() - start
+        logging.debug("Finished downloading %s in %s seconds", key, end)
+
+
+    graph_part_s3_prefix = os.path.join(os.path.join(graph_data_s3, graph_part), "")
+    s3_graph_part_files = S3Downloader.list(
+        graph_part_s3_prefix,
+        sagemaker_session=sagemaker_session)
+
+    # Download graph structure, features and DGL mapping files
+    for s3_graph_part_file in s3_graph_part_files:
+        graph_part_key = s3_graph_part_file.split("/", maxsplit=3)[3]
+        local_part_path = os.path.join(graph_part_path, os.path.basename(graph_part_key))
+        download_large_file(
+            s3_client,
+            s3_input_bucket,
+            graph_part_key,
+            local_part_path,
+            min(DOWNLOAD_THREADS, os.cpu_count())
+        )
+
+    logging.info("Rank %d: Time to download graph part %s: %.2f seconds",
+                 rank, graph_part, time.perf_counter() - graph_part_start)
 
     node_id_mapping = "node_mapping.pt"
     # Try to download node id mapping file if any
@@ -318,28 +368,51 @@ def download_graph(graph_data_s3, graph_name, part_id, world_size,
                         "the node id mapping file created by gconstruct or gsprocessing.")
 
     if part_id == 0:
+        # The leader needs to download the DGL intermediate mapping files
+        lead_mapping_start = time.perf_counter()
         # It is possible that id mappings are generated by
         # dgl tools/distpartitioning/convert_partition.py
         for i in range(1, world_size):
             local_graph_part = f"part{i}"
-            graph_part_path = os.path.join(graph_path, local_graph_part)
-            os.makedirs(graph_part_path, exist_ok=True)
+            local_graph_part_path = os.path.join(graph_path, local_graph_part)
+            os.makedirs(local_graph_part_path, exist_ok=True)
 
             # Try to download node id mapping file if any
-            s3_path = os.path.join(graph_data_s3, local_graph_part, "orig_nids.dgl")
+            filename = "orig_nids.dgl"
+            s3_path = os.path.join(graph_data_s3, local_graph_part, filename)
             try:
-                logging.info("Try to download %s to %s", s3_path, graph_part_path)
-                S3Downloader.download(s3_path,
-                    graph_part_path, sagemaker_session=sagemaker_session)
+                logging.debug("Try to download %s to %s", s3_path, local_graph_part_path)
+                dgl_mapping_key = s3_path.split("/", maxsplit=3)[3]
+                download_large_file(
+                    s3_client,
+                    s3_input_bucket,
+                    dgl_mapping_key,
+                    os.path.join(local_graph_part_path, filename),
+                    min(DOWNLOAD_THREADS, os.cpu_count())
+                )
             except Exception: # pylint: disable=broad-except
-                logging.info("node id mapping file %s does not exist", s3_path)
+                logging.info("Could not download DGL node id mapping file %s", s3_path)
+        logging.info("Time to download DGL node ID mappings on leader: %f seconds",
+                     time.perf_counter() - lead_mapping_start)
 
     # Try to get GraphStorm ID to Original ID remapping files if any
     # The S3 path can be empty, which means no Raw ID mapping is needed.
     # For exampling during SageMaker training.
-    id_map_files = S3Downloader.list(
+    raw_id_mappings_start = time.perf_counter()
+
+    # By default we assume the node mappings exist
+    # under the same path as the rest of the graph data
+    if not raw_node_mapping_prefix_s3:
+        raw_node_mapping_prefix_s3 = f"{graph_data_s3}/raw_id_mappings"
+    else:
+        raw_node_mapping_prefix_s3 = (
+            raw_node_mapping_prefix_s3[:-1] if raw_node_mapping_prefix_s3.endswith('/')
+            else raw_node_mapping_prefix_s3)
+
+    # If no mappings exist this list will be empty
+    s3_id_map_files = S3Downloader.list(
         raw_node_mapping_prefix_s3, sagemaker_session=sagemaker_session)
-    for mapping_file in id_map_files:
+    for mapping_file in s3_id_map_files:
         # The expected layout for GConstruct mapping files on S3 is:
         # raw_id_mappings/node_type/part-xxxxx.parquet
         ntype = mapping_file.split("/")[-2]
@@ -348,14 +421,24 @@ def download_graph(graph_data_s3, graph_name, part_id, world_size,
             # Then we have raw_id_mappings/node_type/parquet/part-xxxxx.parquet
             ntype = mapping_file.split("/")[-3]
         os.makedirs(os.path.join(graph_path, "raw_id_mappings", ntype), exist_ok=True)
-        try:
-            S3Downloader.download(
-                mapping_file,
-                os.path.join(graph_path, "raw_id_mappings", ntype),
-                sagemaker_session=sagemaker_session)
-        except Exception: # pylint: disable=broad-except
-            logging.warning("Could not download node id remap file %s",
-                            mapping_file)
+
+    def download_raw_mapping_file(s3_mapping_file):
+        ntype = s3_mapping_file.split("/")[-2]
+        # This is the case where the output was generated by GSProcessing
+        if ntype == "parquet":
+            # Then we have raw_id_mappings/node_type/parquet/part-xxxxx.parquet
+            ntype = s3_mapping_file.split("/")[-3]
+        mapping_key = s3_mapping_file.split("/", maxsplit=3)[3]
+        filename = os.path.basename(mapping_key)
+        local_dl_path = os.path.join(graph_path, "raw_id_mappings", ntype, filename)
+        s3_client.download_file(
+            s3_input_bucket, mapping_key, local_dl_path)
+
+    # We expect the raw id mapping files to be many small files, so we download in parallel
+    Parallel(n_jobs=min(DOWNLOAD_THREADS, os.cpu_count()), prefer="threads")(
+        delayed(download_raw_mapping_file)(mapping_file) for mapping_file in s3_id_map_files)
+    logging.info("Rank %d: Time to download %d raw id mapping files: %f seconds",
+                 rank, len(s3_id_map_files), time.perf_counter() - raw_id_mappings_start)
 
     logging.info("Finished downloading graph data from %s", graph_data_s3)
     return os.path.join(graph_path, graph_config)
@@ -402,23 +485,44 @@ def upload_model_artifacts(model_s3_path, model_path, sagemaker_session):
     # Other ranks will only upload learnable embeddings owned by themselves.
     return upload_data_to_s3(model_s3_path, model_path, sagemaker_session)
 
-def upload_embs(emb_s3_path, emb_path, sagemaker_session):
-    """ Upload generated node embeddings into S3
-
-    As embeddding table is huge and each trainer/inferrer only
-    stores part of the embedding, we need to upload them
-    into S3.
+def upload_directory_parallel(local_prefix: str, s3_prefix: str, s3_client=None):
+    """Upload all files under a local prefix to an S3 prefix
 
     Parameters
     ----------
-    emb_s3_path: str
-        S3 uri to upload node embeddings
-    emb_path: str
-        Local embedding path
-    sagemaker_session: sagemaker.session.Session
-        sagemaker_session to run download
+    local_prefix : str
+        Local directory prefix
+    s3_prefix : str
+        S3 prefix under which files will be uploaded
+    s3_client : boto3.client, optional
+        S3 boto client, by default None
     """
-    return upload_data_to_s3(emb_s3_path, emb_path, sagemaker_session)
+    if not s3_client:
+        s3_client = boto3.client(
+            "s3",
+            config=botocore.config.Config(max_pool_connections=150),
+            region_name=os.environ.get("AWS_REGION", None)
+        )
+    rank = get_rank()
+
+    local_src_s3_dst_tuples = get_upload_tuples(local_prefix, s3_prefix, include_filename=True)
+
+    logging.info("Rank %d: Uploading %d embeddings files to %s",
+                rank, len(local_src_s3_dst_tuples), s3_prefix)
+
+    def upload_file(local_path: str, s3_uri: str):
+        bucket = s3_uri.split("/")[2]
+        key = s3_uri.split("/", maxsplit=3)[3]
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            s3_client.upload_file(local_path, bucket, key)
+
+    verbosity = 10 if rank == 0 else 0
+    # We know we'll get many 'WARNING - Connection pool is full' warnings here so we suppress them
+    Parallel(n_jobs=min(UPLOAD_THREADS, os.cpu_count()), prefer="threads", verbose=verbosity)(
+            delayed(upload_file)(local_path, s3_path)
+                for (local_path, s3_path) in local_src_s3_dst_tuples
+        )
 
 def update_gs_params(gs_params, param_name, param_value):
     """ Update the graphstorm parameter `param_name` with a new
@@ -462,6 +566,52 @@ def remove_embs(emb_path):
     """
     remove_data(emb_path)
 
+# From
+# https://github.com/aws/sagemaker-python-sdk/blob/fb16a269daf4db6a717ef26c1a6bf7631c0c8d2d/src/sagemaker/session.py#L390-L406
+def get_upload_tuples(
+        local_path: str,
+        key_prefix: str,
+        include_filename: bool = False
+    ) -> List[Tuple[str]]:
+    """Walks a directory to create a list of (local_src, s3_dst) paths for upload.
+
+    Parameters
+    ----------
+    local_path : str
+        A local path, can be a directory or single file.
+    key_prefix : str
+        An S3 key prefix under we want all local files uploaded.
+    include_filename: bool (default: False)
+        When True, will include the filename in the returned S3 URIs, otherwise
+        will just return the prefix
+    Returns
+    -------
+    List[Tuple[str]]
+        A list of (local_src_path, s3_dist_path) tuples, one for each file
+        under the input local_path.
+    """
+    # Generate a tuple for each file that we want to upload of the form (local_path, s3_key).
+    files = []
+    if os.path.isdir(local_path):
+        for dirpath, _, filenames in os.walk(local_path):
+            for name in filenames:
+                file_path = os.path.join(dirpath, name)
+                if local_path == dirpath:
+                    s3_relative_prefix = ""
+                else:
+                    s3_relative_prefix = os.path.relpath(dirpath, start=local_path) + "/"
+                if include_filename:
+                    s3_key = "{}/{}{}".format(key_prefix, s3_relative_prefix, name)
+                else:
+                    s3_key = "{}/{}".format(key_prefix, s3_relative_prefix)
+                files.append((file_path, s3_key))
+    else:
+        _, name = os.path.split(local_path)
+        s3_key = "{}/{}".format(key_prefix, name)
+        files.append((local_path, s3_key))
+
+    return files
+
 def is_port_available(port):
     """Check if a port is available."""
     try:
diff --git a/sagemaker/launch/common_parser.py b/sagemaker/launch/common_parser.py
index 55440003ae..5cdb1c9021 100644
--- a/sagemaker/launch/common_parser.py
+++ b/sagemaker/launch/common_parser.py
@@ -2,7 +2,7 @@
 Common parsers for all launcher scripts.
 """
 import argparse
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 from ast import literal_eval
 
 import boto3
@@ -103,9 +103,10 @@ def get_common_parser() -> argparse.ArgumentParser:
 
     return parser
 
-def parse_estimator_kwargs(arg_string: str) -> Dict[str, Any]:
-    """
-    Parses Estimator/Processor arguments for SageMaker tasks. See
+def parse_estimator_kwargs(arg_string: str, sm_job_type: Optional[str] = None) -> Dict[str, Any]:
+    """Parses Estimator/Processor arguments for SageMaker tasks.
+
+    See
     https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html
     for the full list of arguments for train/infer/partition tasks.
     For GConstruct see
@@ -113,9 +114,24 @@ def parse_estimator_kwargs(arg_string: str) -> Dict[str, Any]:
     Argument values are evaluated as Python
     literals using ast.literal_eval.
 
-    :param arg_string: String of arguments in the form of
-        <key>=<val> separated by spaces.
-    :return: Dictionary of parsed arguments.
+    Parameters
+    ----------
+    arg_string : str
+        String of arguments in the form of '<key>=<val>' separated by spaces.
+    sm_job_type : Optional[str], optional
+        Intended SageMaker job type, can be None, 'training', or 'processing',
+        by default None
+
+    Returns
+    -------
+    Dict[str, Any]
+        Dictionary of parsed arguments for use in a sagemaker.training.Estimator
+        or ScriptProcessor
+
+    Raises
+    ------
+    ValueError
+        If `sm_job_type` is not one of 'training', 'processing', or None
     """
     if arg_string is None:
         return {}
@@ -124,6 +140,25 @@ def parse_estimator_kwargs(arg_string: str) -> Dict[str, Any]:
         k, v =param.split("=")
         typed_args_dict[k] = literal_eval(v)
 
+    # Convert processing job args to training and vice versa
+    if sm_job_type:
+        if sm_job_type == "training":
+            if "volume_size_in_gb" in typed_args_dict:
+                typed_args_dict["volume_size"] = typed_args_dict["volume_size_in_gb"]
+                del typed_args_dict["volume_size_in_gb"]
+            if "max_runtime_in_seconds" in typed_args_dict:
+                typed_args_dict["max_run"] = typed_args_dict["max_runtime_in_seconds"]
+                del typed_args_dict["max_runtime_in_seconds"]
+        elif sm_job_type == "processing":
+            if "volume_size" in typed_args_dict:
+                typed_args_dict["volume_size_in_gb"] = typed_args_dict["volume_size"]
+                del typed_args_dict["volume_size"]
+            if "max_run" in typed_args_dict:
+                typed_args_dict["max_runtime_in_seconds"] = typed_args_dict["max_run"]
+                del typed_args_dict["max_run"]
+        else:
+            raise ValueError(f"Unknown SageMaker job type: {sm_job_type}")
+
     return typed_args_dict
 
 def parse_unknown_gs_args(unknown_args: List[str]) -> Dict[str, str]:
diff --git a/sagemaker/launch/launch_infer.py b/sagemaker/launch/launch_infer.py
index fbc68a833f..642b986bee 100644
--- a/sagemaker/launch/launch_infer.py
+++ b/sagemaker/launch/launch_infer.py
@@ -16,8 +16,10 @@
     Launch SageMaker inference task
 """
 import os
+from time import strftime, gmtime
 
 from sagemaker.pytorch.estimator import PyTorch
+from sagemaker.processing import ScriptProcessor
 
 
 from common_parser import (
@@ -30,7 +32,103 @@
 
 INSTANCE_TYPE = "ml.g4dn.12xlarge"
 
-def run_job(input_args, image, unknownargs):
+def run_processing_job(input_args, image, unknownargs):
+    timestamp = strftime("%Y-%m-%d-%H-%M-%S", gmtime())
+    # SageMaker base job name
+    sm_task_name = input_args.task_name if input_args.task_name else timestamp
+    role = input_args.role # SageMaker ARN role
+    instance_type = input_args.instance_type # SageMaker instance type
+    instance_count = input_args.instance_count # Number of infernece instances
+    region = input_args.region # AWS region
+    entry_point = input_args.entry_point # GraphStorm inference entry_point
+    task_type = input_args.task_type # Inference task type
+    graph_name = input_args.graph_name # Inference graph name
+    graph_data_s3 = input_args.graph_data_s3 # S3 location storing partitioned graph data
+    infer_yaml_s3 = input_args.yaml_s3 # S3 location storing the yaml file
+    output_emb_s3_path = input_args.output_emb_s3 # S3 location to save node embeddings
+    output_predict_s3_path = input_args.output_prediction_s3 # S3 location to save prediction results
+    model_artifact_s3 = input_args.model_artifact_s3 # S3 location of saved model artifacts
+    output_chunk_size = input_args.output_chunk_size # Number of rows per chunked prediction result or node embedding file.
+    log_level = input_args.log_level # SageMaker runner logging level
+
+    boto_session = boto3.session.Session(region_name=region)
+    sagemaker_client = boto_session.client(service_name="sagemaker", region_name=region)
+    # need to skip s3://
+    assert model_artifact_s3.startswith('s3://'), \
+        "Saved model artifact should be stored in S3"
+    sagemaker_session = sagemaker.session.Session(boto_session=boto_session,
+        sagemaker_client=sagemaker_client)
+
+    arguments = [
+        "--graph-data-s3", graph_data_s3,
+        "--graph-name", graph_name,
+        "--infer-yaml-s3", infer_yaml_s3,
+        "--log-level", log_level,
+        "--model-artifact-s3", model_artifact_s3,
+        "--num-trainers", input_args.num_trainers,
+        "--output-chunk-size", output_chunk_size,
+        "--output-emb-s3", output_emb_s3_path,
+        "--task-type", task_type,
+    ]
+
+    # In Link Prediction, no prediction outputs
+    if task_type not in ["link_prediction", "compute_emb"]:
+        arguments.extend(["--output-prediction-s3", output_predict_s3_path])
+    # If no raw mapping files are provided, remapping is skipped
+    if input_args.raw_node_mappings_s3 is not None:
+        arguments.extend(["--raw-node-mappings-s3", input_args.raw_node_mappings_s3])
+    # TODO: Handle unknown args
+    # We must handle cases like
+    # --target-etype query,clicks,asin query,search,asin
+    # --feat-name ntype0:feat0 ntype1:feat1
+    # --column-names nid,~id emb,embedding
+    # unknow_idx = 0
+    # while unknow_idx < len(unknownargs):
+    #     print(unknownargs[unknow_idx])
+    #     assert unknownargs[unknow_idx].startswith("--")
+    #     sub_params = []
+    #     for i in range(unknow_idx+1, len(unknownargs)+1):
+    #         # end of loop or stand with --
+    #         if i == len(unknownargs) or \
+    #             unknownargs[i].startswith("--"):
+    #             break
+    #         sub_params.append(unknownargs[i])
+    #     arguments[unknownargs[unknow_idx]] = ' '.join(sub_params)
+    #     unknow_idx = i
+
+    print(f"Parameters {arguments}")
+    print(f"GraphStorm Parameters {unknownargs}")
+
+    arguments = [str(arg) for arg in arguments]
+
+    estimator_kwargs = parse_estimator_kwargs(input_args.sm_estimator_parameters, sm_job_type="processing")
+
+    if input_args.sm_estimator_parameters:
+        print(f"SageMaker Estimator parameters: '{estimator_kwargs}'")
+
+
+    script_processor = ScriptProcessor(
+        image_uri=image,
+        role=role,
+        instance_count=instance_count,
+        instance_type=instance_type,
+        command=["python3"],
+        base_job_name=f"gs-infer-{sm_task_name}",
+        sagemaker_session=sagemaker_session,
+        tags=[{"Key":"GraphStorm", "Value":"beta"},
+              {"Key":"GraphStorm_Task", "Value":"Inference"}],
+        **estimator_kwargs
+    )
+
+    script_processor.run(
+        code=entry_point,
+        arguments=arguments,
+        inputs=[],
+        outputs=[],
+        wait=not input_args.async_execution
+    )
+
+def run_training_job(input_args, image, unknownargs):
     """ Run job using SageMaker estimator.PyTorch
 
         We use SageMaker training task to run offline inference.
@@ -74,7 +172,9 @@ def run_job(input_args, image, unknownargs):
         "graph-data-s3": graph_data_s3,
         "graph-name": graph_name,
         "infer-yaml-s3": infer_yaml_s3,
+        "log-level": log_level,
         "model-artifact-s3": model_artifact_s3,
+        "num-trainers": input_args.num_trainers,
         "output-chunk-size": output_chunk_size,
         "output-emb-s3": output_emb_s3_path,
         "task-type": task_type,
@@ -95,7 +195,7 @@ def run_job(input_args, image, unknownargs):
     if input_args.sm_estimator_parameters:
         print(f"SageMaker Estimator parameters: '{input_args.sm_estimator_parameters}'")
 
-    estimator_kwargs = parse_estimator_kwargs(input_args.sm_estimator_parameters)
+    estimator_kwargs = parse_estimator_kwargs(input_args.sm_estimator_parameters, sm_job_type="training")
 
     est = PyTorch(
         entry_point=os.path.basename(entry_point),
@@ -151,7 +251,7 @@ def get_inference_parser():
              "(Only works with node classification/regression " \
              "and edge classification/regression tasks)",
         default=None)
-    parser.add_argument("--output-chunk-size", type=int, default=100000,
+    parser.add_argument("--output-chunk-size", type=int, default=10**6,
         help="Number of rows per chunked prediction result or node embedding file.")
     inference_args.add_argument("--model-sub-path", type=str, default=None,
         help="Relative path to the trained model under <model_artifact_s3>."
@@ -159,6 +259,13 @@ def get_inference_parser():
              "<model_artifact_s3>, this argument is used to choose one.")
     inference_args.add_argument('--log-level', default='INFO',
         type=str, choices=['DEBUG', 'INFO', 'WARNING', 'CRITICAL', 'FATAL'])
+    inference_args.add_argument("--sagemaker-job-type", type=str, default="training",
+        help=("Choose the type of SageMaker job to run, "
+              "options are 'training' or 'processing'. Default: 'training'"),
+        choices=["training", "processing"]
+        )
+    inference_args.add_argument('--num-trainers', default=1,
+        type=int, help="Number of trainers to use during inference")
 
     return parser
 
@@ -173,4 +280,7 @@ def get_inference_parser():
         args.instance_type = INSTANCE_TYPE
 
 
-    run_job(args, infer_image, unknownargs)
+    if args.sagemaker_job_type == "training":
+        run_training_job(args, infer_image, unknownargs)
+    else:
+        run_processing_job(args, infer_image, unknownargs)
diff --git a/sagemaker/run/infer_entry.py b/sagemaker/run/infer_entry.py
index 50564f9bd2..902a489573 100644
--- a/sagemaker/run/infer_entry.py
+++ b/sagemaker/run/infer_entry.py
@@ -16,6 +16,7 @@
     SageMaker inference entry point
 """
 import argparse
+import json
 import os
 import subprocess
 
@@ -25,12 +26,18 @@
 def parse_inference_args():
     """  Add arguments for model inference
     """
+    try:
+        with open("/opt/ml/config/resourceconfig.json", "r", encoding="utf-8") as f:
+            train_env = json.load(f)
+    except FileNotFoundError:
+        train_env = json.loads(os.environ['SM_TRAINING_ENV'])
+
     parser = argparse.ArgumentParser(description='gs sagemaker inference pipeline')
 
     parser.add_argument("--task-type", type=str,
         help=f"task type, builtin task type includes: {SUPPORTED_TASKS}")
 
-    # disrributed training
+    # distributed training
     parser.add_argument("--graph-name", type=str, help="Graph name")
     parser.add_argument("--graph-data-s3", type=str,
         help="S3 location of input training graph",
@@ -59,12 +66,21 @@ def parse_inference_args():
         help="Number of rows per chunked prediction result or node embedding file.")
     parser.add_argument('--log-level', default='INFO',
         type=str, choices=['DEBUG', 'INFO', 'WARNING', 'CRITICAL', 'FATAL'])
+    parser.add_argument('--num-trainers', type=str, default='1')
 
+    # TODO: Remove dynamic args
+    if "SM_CHANNEL_TRAIN" in os.environ:
+        data_path = os.environ['SM_CHANNEL_TRAIN']
+    else:
+        data_path = "/opt/ml/input/data/training"
+    if "SM_MASTER_ADDR" in os.environ:
+        master_addr = os.environ['SM_MASTER_ADDR']
+    else:
+        master_addr = train_env["hosts"][0]
     # following arguments are required to launch a distributed GraphStorm training task
-    parser.add_argument('--data-path', type=str, default=os.environ['SM_CHANNEL_TRAIN'])
-    parser.add_argument('--num-gpus', type=str, default=os.environ['SM_NUM_GPUS'])
-    parser.add_argument('--sm-dist-env', type=str, default=os.environ['SM_TRAINING_ENV'])
-    parser.add_argument('--master-addr', type=str, default=os.environ['MASTER_ADDR'])
+    parser.add_argument('--data-path', type=str, default=data_path)
+    parser.add_argument('--sm-dist-env', type=str, default=train_env)
+    parser.add_argument('--master-addr', type=str, default=master_addr)
     parser.add_argument('--region', type=str, default=os.environ['AWS_REGION'])
 
     # Add your args if any