diff --git a/sagemaker/pipeline/create_sm_pipeline.py b/sagemaker/pipeline/create_sm_pipeline.py index d29c68b80..dbdf4dda9 100644 --- a/sagemaker/pipeline/create_sm_pipeline.py +++ b/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 = [ diff --git a/sagemaker/pipeline/execute_sm_pipeline.py b/sagemaker/pipeline/execute_sm_pipeline.py index 32dd31298..0e022aca3 100644 --- a/sagemaker/pipeline/execute_sm_pipeline.py +++ b/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(