Skip to content

Add Pythonic SageMaker training and processing - #3459

Draft
rohitrsh wants to merge 2 commits into
flyteorg:masterfrom
rohitrsh:feat/sagemaker-pythonic-jobs
Draft

Add Pythonic SageMaker training and processing#3459
rohitrsh wants to merge 2 commits into
flyteorg:masterfrom
rohitrsh:feat/sagemaker-pythonic-jobs

Conversation

@rohitrsh

@rohitrsh rohitrsh commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Stacked PR 2/2 — do not merge out of order.

This branch currently includes the commit from #3456. After #3456 merges, this branch will be rebased onto master so the final review contains only the Pythonic SageMaker change.

Summary

  • Adds Pythonic SageMaker Training and Processing tasks that execute ordinary Flyte @task functions inside SageMaker-managed containers.
  • Introduces shared entrypoint, environment, job naming, output, error, retry, and cancellation handling for Pythonic SageMaker jobs.
  • Requires an explicit container image and region, limits Pythonic jobs to one instance, validates SageMaker entrypoint limits, and rejects incompatible Processing network isolation.
  • Documents image requirements, IAM behavior, model artifacts, inputs/outputs, and the distinction between declarative and Pythonic modes.
  • Treats RETRYABLE_FAILED as terminal for local connector execution and adds regression coverage.

Stack order:

  1. Add managed SageMaker job connectors #3456 — Managed SageMaker job connectors
  2. Pythonic SageMaker Training and Processing (this PR)

Test plan

  • Focused Pythonic and managed Training/Processing connector suite — 74 passed
  • Core connector utility suite — 15 passed
  • SageMaker plugin lint and formatting checks
  • Plugin packaging, serialization, and import smoke tests
  • DCO, confidentiality, attribution, and internal-reference audits

Merge only after #3456. Before merge, rebase this branch onto upstream master, rerun all local and hosted verification gates, and mark the PR ready for review.

Expand the existing SageMaker plugin with async lifecycle support for training, processing, batch transform, hyperparameter tuning, and Inference Recommender jobs. Include typed task wrappers, stable projected outputs, idempotent retries, public documentation, and comprehensive unit tests without requiring Propeller changes.

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>
Allow typed Flyte task functions to run directly inside single-node SageMaker Training and Processing jobs through rendered container entrypoints. Preserve Flyte outputs and error semantics, validate SageMaker runtime constraints, and document the ECR, IAM, S3, networking, and secret-management requirements.

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant