Great question. Short answer: there isn’t a built-in SageMaker Pipelines integration in Prefect 3.x today, but the recommended approach is straightforward—use

prefect-aws

for auth and boto3 to start the pipeline, then poll

DescribePipelineExecution

until it reaches a terminal state from within your flow. Key points - No native SageMaker Pipeline block/tasks in Prefect 3.x. - Use `prefect-aws`’s

AwsCredentials

for auth and boto3 for API calls. - Start with

StartPipelineExecution

, then poll

DescribePipelineExecution

until status is Succeeded/Failed/Stopped. - There is no boto3 waiter for pipeline executions, so polling is the reliable pattern. - Prefer running your flows on AWS (e.g., ECS worker) so IAM roles can be used instead of long‑lived keys. Example: start and monitor a SageMaker Pipeline from a Prefect 3 flow

from prefect import flow, task, get_run_logger
from prefect_aws import AwsCredentials
import time

@task
def start_sagemaker_pipeline(pipeline_name: str, parameters: dict, aws: AwsCredentials) -> str:
    """
    Start a SageMaker pipeline, return the execution ARN.
    """
    client = aws.get_boto3_session().client("sagemaker")
    params = [{"Name": k, "Value": str(v)} for k, v in (parameters or {}).items()]
    resp = client.start_pipeline_execution(
        PipelineName=pipeline_name,
        PipelineParameters=params
    )
    return resp["PipelineExecutionArn"]

@task
def wait_for_pipeline(execution_arn: str, aws: AwsCredentials, poll_seconds: int = 30) -> dict:
    """
    Poll the SageMaker pipeline execution until it completes.
    """
    logger = get_run_logger()
    client = aws.get_boto3_session().client("sagemaker")
    terminal = {"Succeeded", "Failed", "Stopped"}

    while True:
        desc = client.describe_pipeline_execution(PipelineExecutionArn=execution_arn)
        status = desc["PipelineExecutionStatus"]
        <http://logger.info|logger.info>(f"SageMaker pipeline status: {status}")

        if status in terminal:
            if status != "Succeeded":
                reason = desc.get("PipelineExecutionFailureReason", "Unknown")
                raise RuntimeError(f"Pipeline ended with status {status}: {reason}")
            return desc

        time.sleep(poll_seconds)

@flow
def run_sagemaker_pipeline(
    pipeline_name: str,
    parameters: dict | None = None,
    aws_block_name: str = "my-aws-credentials",
    poll_seconds: int = 30
):
    aws = AwsCredentials.load(aws_block_name)  # or AwsCredentials() to rely on IAM role/env
    arn = start_sagemaker_pipeline(pipeline_name, parameters or {}, aws)
    result = wait_for_pipeline(arn, aws, poll_seconds=poll_seconds)
    return {"execution_arn": arn, "status": result["PipelineExecutionStatus"]}

Authentication patterns - Best: run Prefect workers in AWS (ECS/EKS/EC2) and rely on IAM roles; your

AwsCredentials()

can be empty and boto3 will use the role. - Alternatively: create an

AwsCredentials

block with keys and region and load it in your flow. - Required IAM permissions typically include:

sagemaker:StartPipelineExecution

and

sagemaker:DescribePipelineExecution

on the pipeline resources. Helpful links - Prefect AWS collection: prefect-aws docs - prefect-aws repo: Prefect AWS on GitHub - Boto3 StartPipelineExecution: start_pipeline_execution - Boto3 DescribePipelineExecution: describe_pipeline_execution Optional: reusable abstraction

- If you want a “click and reuse” primitive for teams, you can wrap the start/poll logic into a small helper class or even a custom Prefect Block that exposes

trigger()

and

wait_for_completion()

. The flow-level tasks above are usually sufficient and simpler to maintain. Would you prefer: - fire-and-forget (trigger and exit), or - synchronous monitoring (start and wait), and where are your flows running today (local, ECS, EKS)? I can tailor the example (credentials, retries, notifications) to your setup.

Totally fine to build the SageMaker pipeline inside your Prefect flow. The general pattern is: - Construct the pipeline (and optionally upsert it) inside a Prefect task - Start an execution from the flow or another task - Poll

DescribePipelineExecution

until it finishes - Use

prefect-aws

for credentials and IAM role–based auth if possible Two important notes for your snippet: - If you deploy this flow, avoid passing numpy arrays as parameters—Prefect deployment parameters are JSON-serialized. Prefer storing inputs in S3 and passing URIs, or convert to plain lists. - Whether you’re using SageMaker “remote functions”/`@remote` or the Pipelines SDK classes, the run/wait pattern is the same. Here’s a minimal, flow-first approach that creates the pipeline in a task, starts it, and waits: ``` from prefect import flow, task, get_run_logger from prefect_aws import AwsCredentials import time # If you’re using SageMaker Pipelines: from sagemaker.workflow.pipeline import Pipeline from sagemaker.workflow.pipeline_context import PipelineSession # If you’re using remote functions, import the right decorator for your SDK version: # from sagemaker.remote_function import remote @task def build_pipeline(aws: AwsCredentials, pipeline_name: str) -> Pipeline: # Create a SageMaker PipelineSession that uses Prefect’s AWS creds sm_session = PipelineSession(boto_session=aws.get_boto3_session()) # Define your steps here (examples only; wire up your real steps) # If you're using remote functions: # @remote(image_uri="public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu", # instance_type="ml.m4.4xlarge", # sagemaker_session=sm_session) # def matrix_multiply(a, b): # import numpy as np # return np.matmul(a, b) # # @remote(image_uri="public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu", # instance_type="ml.m4.4xlarge", # sagemaker_session=sm_session) # def divide_matrix(a, b): # import numpy as np # return np.matmul(a, b) # Or, if you’re using ProcessingStep/TrainingStep, construct those and pass to Pipeline pipe = Pipeline( name=pipeline_name, steps=[/* your steps here */], sagemaker_session=sm_session, ) # Ensure the pipeline exists/updated (you’ll need a role with pipeline perms) # pipe.upsert(role_arn="arnawsiam:<acct>role/SageMakerPipelineRole") return pipe @task def start_and_wait(pipe: Pipeline, aws: AwsCredentials, parameters: dict | None = None, poll_seconds: int = 30) -> str: logger = get_run_logger() client = aws.get_boto3_session().client("sagemaker") # Prefer idempotency if retried: use a stable token (e.g., flow_run_id) # from prefect import runtime # token = str(runtime.flow_run.id) # start_resp = client.start_pipeline_execution( # PipelineName=pipe.name, # PipelineParameters=[{"Name": k, "Value": str(v)} for k, v in (parameters or {}).items()], # ClientRequestToken=token, # ) # If using SDK’s Pipeline.start(): start_resp = pipe.start(parameters=parameters or {}) exec_arn = start_resp["PipelineExecutionArn"] logger.info(f"Started SageMaker pipeline execution: {exec_arn}") terminal = {"Succeeded", "Failed", "Stopped"} while True: desc = client.describe_pipeline_execution(PipelineExecutionArn=exec_arn) status = desc["PipelineExecutionStatus"] logger.info(f"Pipeline status: {status}") if status in terminal: if status != "Succeeded": raise RuntimeError(f"Pipeline {status}: {desc.get('PipelineExecutionFailureReason')}") return exec_arn time.sleep(poll_seconds) @flow def main_flow( pipeline_name: str = "my-prefect-built-pipeline", parameters: dict | None = None, aws_block_name: str = "my-aws-credentials", ):

aws = AwsCredentials.load(aws_block_name) # or AwsCredentials() to rely on IAM role pipe = build_pipeline(aws, pipeline_name) exec_arn = start_and_wait(pipe, aws, parameters=parameters) return exec_arn ``` Best practices - Put “build/upsert” in its own task (or its own provisioning flow) and “run execution + wait” in another task. This prevents recreating/upserting the pipeline on every run. - Use IAM roles on your worker (ECS/EKS/EC2) rather than long-lived keys.

AwsCredentials()

with no args will let boto3 pick up the role automatically. - Use an idempotency token (e.g.,

prefect.runtime.flow_run.id

) when starting executions so retries don’t create duplicates. - For large binary inputs (numpy arrays), write them to S3 and pass S3 URIs/parameters to the pipeline; fetch results from S3 in downstream tasks. AWS/SageMaker references - SageMaker Pipelines SDK reference: Pipeline class - Remote functions overview: SageMaker remote functions - Boto3 start/describe execution: start_pipeline_execution, describe_pipeline_execution - Prefect AWS integration: prefect-aws docs Quick clarification to tailor this precisely: - Are you using SageMaker “remote functions” decorators or the Pipelines SDK steps (ProcessingStep/TrainingStep/etc.)? - Where will this flow run (local, ECS, EKS)? I can adjust the auth + packaging guidance accordingly.