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

Amazon Bedrock - Model Customization Jobs #38693

Merged
merged 12 commits into from
Apr 8, 2024

Conversation

ferruzzi
Copy link
Contributor

@ferruzzi ferruzzi commented Apr 3, 2024

Adds support for Model Customization Jobs. Includes changes to the hook(s), new Operator, Sensor, Trigger, Waiter, unit tests for the above, updates to the doc page, and additions to the system test.

Manually tested in Breeze using (wait_for_completion=False with a Sensor), wait_for_completion=True, and deferrable=True. Don't mind the first two, bugs were squashed , as seen in the rest of the runs :P
image

@ferruzzi
Copy link
Contributor Author

ferruzzi commented Apr 3, 2024

I also didn't add a BedrockDeleteCustomModelOperator because it is literally just a hook call: BedrockHook().conn.delete_custom_model(modelIdentifier=custom_model_arn). It's instantaneous and irreversible, so no sensors or waiters or triggers or anything like that, just a bulky wrapper around a hook call. Should I add it for completion sake?

airflow/providers/amazon/aws/hooks/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/sensors/bedrock.py Outdated Show resolved Hide resolved
tests/providers/amazon/aws/operators/test_bedrock.py Outdated Show resolved Hide resolved
tests/providers/amazon/aws/operators/test_bedrock.py Outdated Show resolved Hide resolved
@ferruzzi ferruzzi changed the title Amazon Bedrock -Model Customization Jobs Amazon Bedrock - Model Customization Jobs Apr 5, 2024
Copy link
Contributor

@syedahsn syedahsn left a comment

Choose a reason for hiding this comment

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

LGTM. Nice work, I love all the optimizations we've made with supporting deferrable operators (AwsBaseWaiterTrigger, custom waiters, consistent structures etc.) Really simplifies the whole process 😄

@ferruzzi
Copy link
Contributor Author

ferruzzi commented Apr 5, 2024

@o-nikolas @Taragolis @syedahsn - I refactored how we handle a name conflict. Can you please have another look and re-approve (or comment) when you get time?

@o-nikolas
Copy link
Contributor

@o-nikolas @Taragolis @syedahsn - I refactored how we handle a name conflict. Can you please have another look and re-approve (or comment) when you get time?

The code looks functionally correct to me. Although the while True loop worries me a bit, it's very easy for a small bug to get introduced into something like that and then suddenly be stuck in an infinite loop.

@ferruzzi
Copy link
Contributor Author

ferruzzi commented Apr 5, 2024

@o-nikolas Reworked the while True so it won't repeat by default

@ferruzzi ferruzzi added use public runners Makes sure that Public runners are used even if commiters creates the PR (useful for testing) use self-hosted runners When set on a PR run self-hosted runners will be used by default and removed use public runners Makes sure that Public runners are used even if commiters creates the PR (useful for testing) use self-hosted runners When set on a PR run self-hosted runners will be used by default labels Apr 5, 2024
@ferruzzi ferruzzi merged commit 7ed31d5 into apache:main Apr 8, 2024
40 checks passed
@ferruzzi ferruzzi deleted the ferruzzi/bedrock/m2-customize-model branch April 8, 2024 20:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants