Update

awslabs · Nov 22, 2023 · 2ae9c90 · 2ae9c90
1 parent 358f5bb
commit 2ae9c90
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/python/graphstorm/sagemaker/sagemaker_infer.py b/python/graphstorm/sagemaker/sagemaker_infer.py
@@ -144,7 +144,7 @@ def run_infer(args, unknownargs):
             The name of the graph.
         graph_data_s3: str
             S3 location of input training graph.
-        node_mappings_s3: str
+        raw_node_mappings_s3: str
             S3 location of node id mapping files.
         infer_yaml_s3: str
             S3 location of inference yaml file.
@@ -245,7 +245,7 @@ def run_infer(args, unknownargs):
     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.node_mappings_s3)
+        host_rank, world_size, data_path, sagemaker_session, args.raw_node_mappings_s3)
 
     # Download Saved model
     download_model(model_artifact_s3, model_path, sagemaker_session)

diff --git a/sagemaker/launch/launch_infer.py b/sagemaker/launch/launch_infer.py
@@ -73,7 +73,7 @@ def run_job(input_args, image, unknownargs):
         "graph-name": graph_name,
         "infer-yaml-s3": infer_yaml_s3,
         "model-artifact-s3": model_artifact_s3,
-        "node-mappings-s3": input_args.node_mappings_s3,
+        "raw-node-mappings-s3": input_args.raw_node_mappings_s3,
         "output-chunk-size": output_chunk_size,
         "output-emb-s3": output_emb_s3_path,
         "task-type": task_type,
@@ -148,7 +148,7 @@ def get_inference_parser():
     inference_args.add_argument("--model-artifact-s3", type=str,
         help="S3 location to load the saved model artifacts from",
         required=True)
-    inference_args.add_argument("--node-id-mappings-s3", type=str,
+    inference_args.add_argument("--raw-node-mappings-s3", type=str,
         help="S3 location to load the node id mappings from",
         required=False)
     inference_args.add_argument("--output-emb-s3", type=str,

diff --git a/sagemaker/run/infer_entry.py b/sagemaker/run/infer_entry.py
@@ -47,7 +47,7 @@ def parse_inference_args():
         default=None)
     parser.add_argument("--model-artifact-s3", type=str,
         help="S3 bucket to load the saved model artifacts")
-    parser.add_argument("--node-mappings-s3", type=str, required=False,
+    parser.add_argument("--raw-node-mappings-s3", type=str, required=False,
         default=None, help="S3 location where the node mappings exist.")
     parser.add_argument("--custom-script", type=str, default=None,
         help="Custom training script provided by a customer to run customer training logic. \