Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Create Sagemaker pipeline schedules if specified #3271

Open
wants to merge 28 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
5570c0c
Create Sagemaker pipeline schedules if specified
htahir1 Dec 20, 2024
6f7bf28
Add property to check if orchestrator is schedulable
htahir1 Dec 20, 2024
0212388
Add EventBridge rule for SageMaker pipeline execution
htahir1 Dec 20, 2024
2125312
Update IAM policy and trust relationship for EventBridge
htahir1 Dec 20, 2024
966f712
Refactor schedule metadata generation for Sagemaker orchestrator
htahir1 Dec 20, 2024
dd7110c
Add scheduling support for SageMaker orchestrator
htahir1 Dec 20, 2024
fcf934e
Remove trust relationship logic in Sagemaker orchestrator
htahir1 Dec 20, 2024
358b3e8
Handle unsupported schedule in custom orchestrator
htahir1 Dec 20, 2024
7ce08b6
Refactor yield statement to use 'yield from' syntax
htahir1 Dec 20, 2024
72bdae1
Ensure IAM permissions for scheduled SageMaker pipelines
htahir1 Dec 20, 2024
abf4610
Update authentication instructions for SageMaker orchestrator
htahir1 Dec 20, 2024
67705c2
Refactor Sagemaker orchestrator metadata handling
htahir1 Dec 20, 2024
d190f31
Add unit tests for SageMaker orchestrator metadata
htahir1 Dec 20, 2024
f1cabc7
Add exception handling for pipeline preparation errors
htahir1 Dec 20, 2024
726b47a
Add timezone information to first execution message
htahir1 Dec 20, 2024
80c8e8e
Add timezone support to AWS SageMaker orchestrator
htahir1 Dec 20, 2024
da8fd35
Update error handling in SagemakerOrchestrator
htahir1 Dec 20, 2024
ab0c06d
Update error handling messages for AWS in Sagemaker orchestrator
htahir1 Dec 20, 2024
3c12d45
Refactor error handling in SagemakerOrchestrator
htahir1 Dec 20, 2024
51a499a
Handle insufficient permissions creating EventBridge rules
htahir1 Dec 20, 2024
7d65a20
Update error message for EventBridge creation failure
htahir1 Dec 20, 2024
e2ecaf3
Remove logging in SagemakerOrchestrator class
htahir1 Dec 20, 2024
a5fd82b
Refactor orchestrator metadata computation logic
htahir1 Dec 20, 2024
d229460
Merge branch 'develop' into feature/add-sagemaker-schedule
htahir1 Dec 22, 2024
a005e5a
Update handling of scheduled pipeline updates in SageMaker.md
htahir1 Dec 22, 2024
645903a
Merge branch 'feature/add-sagemaker-schedule' of github.com:zenml-io/…
htahir1 Dec 22, 2024
8811f3f
Add optional IAM permissions for policy updates
htahir1 Dec 23, 2024
44fa6c5
Remove redundant code for getting SageMaker session
htahir1 Dec 23, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/book/component-guide/orchestrators/custom.md
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ To see a full end-to-end worked example of a custom orchestrator, [see here](htt

There are some additional optional features that your orchestrator can implement:

* **Running pipelines on a schedule**: if your orchestrator supports running pipelines on a schedule, make sure to handle `deployment.schedule` if it exists. If your orchestrator does not support schedules, you should either log a warning and or even raise an exception in case the user tries to schedule a pipeline.
* **Running pipelines on a schedule**: if your orchestrator supports running pipelines on a schedule, make sure to handle `deployment.schedule` if it exists. If your orchestrator ules, you should either log a warning and or even raise an exception in case the user tries to schedule a pipeline.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
* **Running pipelines on a schedule**: if your orchestrator supports running pipelines on a schedule, make sure to handle `deployment.schedule` if it exists. If your orchestrator ules, you should either log a warning and or even raise an exception in case the user tries to schedule a pipeline.
* **Running pipelines on a schedule**: if your orchestrator supports running pipelines on a schedule, make sure to handle `deployment.schedule` if it exists. If your orchestrator rules, you should either log a warning and or even raise an exception in case the user tries to schedule a pipeline.

Issues:

  • Style Guide - (Spelling-error) Did you really mean 'ules'?

Fix Explanation:

The word 'ules' appears to be a typographical error for 'rules'. Correcting this spelling error improves the clarity and accuracy of the documentation.

* **Specifying hardware resources**: If your orchestrator supports setting resources like CPUs, GPUs or memory for the pipeline or specific steps, make sure to handle the values defined in `step.config.resource_settings`. See the code sample below for additional helper methods to check whether any resources are required from your orchestrator.

### Code sample
Expand Down
100 changes: 95 additions & 5 deletions docs/book/component-guide/orchestrators/sagemaker.md
Original file line number Diff line number Diff line change
Expand Up @@ -153,10 +153,6 @@ Alternatively, for a more detailed view of log messages during SageMaker pipelin

![SageMaker CloudWatch Logs](../../.gitbook/assets/sagemaker-cloudwatch-logs.png)

### Run pipelines on a schedule

The ZenML Sagemaker orchestrator doesn't currently support running pipelines on a schedule. We maintain a public roadmap for ZenML, which you can find [here](https://zenml.io/roadmap). We welcome community contributions (see more [here](https://github.com/zenml-io/zenml/blob/main/CONTRIBUTING.md)) so if you want to enable scheduling for Sagemaker, please [do let us know](https://zenml.io/slack)!

### Configuration at pipeline or step level

When running your ZenML pipeline with the Sagemaker orchestrator, the configuration set when configuring the orchestrator as a ZenML component will be used by default. However, it is possible to provide additional configuration at the pipeline or step level. This allows you to run whole pipelines or individual steps with alternative configurations. For example, this allows you to run the training process with a heavier, GPU-enabled instance type, while running other steps with lighter instances.
Expand Down Expand Up @@ -339,4 +335,98 @@ This approach allows for more granular tagging, giving you flexibility in how yo

Note that if you wish to use this orchestrator to run steps on a GPU, you will need to follow [the instructions on this page](../../how-to/pipeline-development/training-with-gpus/README.md) to ensure that it works. It requires adding some extra settings customization and is essential to enable CUDA for the GPU to give its full acceleration.

<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>
### Scheduling Pipelines

The SageMaker orchestrator supports running pipelines on a schedule using AWS EventBridge. You can configure schedules in three ways:

* Using a cron expression
* Using a fixed interval
* Running once at a specific time

```python
from zenml import pipeline
from datetime import datetime, timedelta

# Using a cron expression (runs daily at 2 AM UTC)
@pipeline(schedule=Schedule(cron_expression="0 2 * * *"))
def my_scheduled_pipeline():
# Your pipeline steps here
pass

# Using an interval (runs every 2 hours)
@pipeline(schedule=Schedule(interval_second=timedelta(hours=2)))
def my_interval_pipeline():
# Your pipeline steps here
pass

# Running once at a specific time
@pipeline(schedule=Schedule(run_once_start_time=datetime(2024, 12, 31, 23, 59)))
def my_one_time_pipeline():
# Your pipeline steps here
pass
```

When you deploy a scheduled pipeline, ZenML will:
1. Create an EventBridge rule with the specified schedule
2. Configure the necessary IAM permissions
3. Set up the SageMaker pipeline as the target

#### Required IAM Permissions

When using scheduled pipelines, you need to ensure your IAM role has the correct permissions and trust relationships:

1. Trust Relationships
Your execution role needs to trust both SageMaker and EventBridge services. Add this trust relationship to your role:

```json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": {
"Service": [
"sagemaker.amazonaws.com",
"events.amazonaws.com"
]
},
"Action": "sts:AssumeRole"
}
]
}
```

2. Required IAM Policies
In addition to the basic SageMaker permissions, you'll need:

```json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": [
"events:PutRule",
"events:PutTargets",
"events:DeleteRule",
"events:RemoveTargets",
"events:DescribeRule",
"events:ListTargetsByRule"
],
"Resource": "arn:aws:events:*:*:rule/zenml-*"
},
{
"Effect": "Allow",
"Action": [
"iam:GetRole",
"iam:GetRolePolicy",
"iam:PutRolePolicy",
"iam:UpdateAssumeRolePolicy"
],
"Resource": "arn:aws:iam::*:role/*"
}
]
}
```

<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ Schedules don't work for all orchestrators. Here is a list of all supported orch
| [KubernetesOrchestrator](../../../component-guide/orchestrators/kubernetes.md) | ✅ |
| [LocalOrchestrator](../../../component-guide/orchestrators/local.md) | ⛔️ |
| [LocalDockerOrchestrator](../../../component-guide/orchestrators/local-docker.md) | ⛔️ |
| [SagemakerOrchestrator](../../../component-guide/orchestrators/sagemaker.md) | ⛔️ |
| [SagemakerOrchestrator](../../../component-guide/orchestrators/sagemaker.md) | |
| [SkypilotAWSOrchestrator](../../../component-guide/orchestrators/skypilot-vm.md) | ⛔️ |
| [SkypilotAzureOrchestrator](../../../component-guide/orchestrators/skypilot-vm.md) | ⛔️ |
| [SkypilotGCPOrchestrator](../../../component-guide/orchestrators/skypilot-vm.md) | ⛔️ |
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,15 @@ class SagemakerOrchestratorSettings(BaseSettings):
("processor_role", "execution_role"), ("processor_tags", "tags")
)

@property
def is_schedulable(self) -> bool:
"""Whether the orchestrator is schedulable or not.

Returns:
Whether the orchestrator is schedulable or not.
"""
return True

@model_validator(mode="before")
def validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
"""Check if model is configured correctly.
Expand Down
Loading
Loading