Ensure host has GPU when trying to execute locally with GPU image

sagemaker/pipeline/create_sm_pipeline.py

@@ -199,6 +199,7 @@ def _create_pipeline_parameters(self, args: PipelineArgs):
         self.partition_algorithm_param = self._create_string_parameter(
             "PartitionAlgorithm", args.partition_config.partition_algorithm
         )
+        # TODO: Probably should not be a parameter
         self.graph_name_param = self._create_string_parameter(
             "GraphName", args.task_config.graph_name
         )
@@ -314,7 +315,7 @@ def _create_gconstruct_step(self, args: PipelineArgs) -> ProcessingStep:
         gc_proc_output = ProcessingOutput(
             source=gc_local_output_path,
             destination=gconstruct_s3_output,
-            output_name=f"{self.graph_name_param}-gconstruct",
+            output_name=f"{self.args.task_config.graph_name}-gconstruct",
         )
 
         gconstruct_arguments = [

sagemaker/pipeline/execute_sm_pipeline.py

@@ -18,6 +18,7 @@
 
 import argparse
 import os
+import subprocess
 import sys
 import warnings
 
@@ -140,6 +141,15 @@ def main():
     deploy_time_hash = pipeline_deploy_args.get_hash_hex()
 
     if args.local_execution:
+        # Ensure GPU is available if trying to execute with GPU locally
+        if not pipeline_deploy_args.instance_config.train_on_cpu:
+            try:
+                subprocess.check_output('nvidia-smi')
+            except Exception:
+                raise RuntimeError(
+                    'Need host with NVidia GPU to run training on GPU! '
+                    "Try re-deploying the pipeline with --train-on-cpu set."
+                    )
         # Use local pipeline and session
         local_session = LocalPipelineSession()
         pipeline_generator = GraphStormPipelineGenerator(