example sagemaker pipeline

dagster_with_sagemaker.py

import boto3
import time


# Define an asset that uses the custom resource
@asset
def sagemaker_pipeline() -> str:
  # note: the functions to connect and use Sagemaker would likely be put into a Dagster Resource
  # this is for illustration purposes for how you might integrate a system like Sagemaker into Dagster as simply as possible
  # Step 1: Set up the SageMaker session and client
  sagemaker_client = boto3.client('sagemaker')

  # Step 2: Define the pipeline name
  pipeline_name = 'your-pipeline-name'

  # Step 3: Start the pipeline execution
  response = sagemaker_client.start_pipeline_execution(
      PipelineName=pipeline_name
  )

  # Get the execution ARN
  execution_arn = response['PipelineExecutionArn']
  print(f'Started pipeline execution: {execution_arn}')

  # Step 4: Monitor the pipeline execution status
  while True:
      response = sagemaker_client.describe_pipeline_execution(
          PipelineExecutionArn=execution_arn
      )
      status = response['PipelineExecutionStatus']
      print(f'Pipeline execution status: {status}')

      if status in ['Succeeded', 'Failed', 'Stopped']:
          break

      time.sleep(30)  # Wait for 30 seconds before checking the status again

  print(f'Pipeline execution completed with status: {status}')


# Define the Dagster definitions with the custom resource
defs = Definitions(
    assets=[greet_asset],
    resources={
        "my_resource": MyCustomResource(greeting="Hello"),
    },
)