Skip to content

Amazon Bedrock - Model Customization Jobs - #38693

Merged
ferruzzi merged 12 commits into
apache:mainfrom
aws-mwaa:ferruzzi/bedrock/m2-customize-model
Apr 8, 2024
Merged

Amazon Bedrock - Model Customization Jobs#38693
ferruzzi merged 12 commits into
apache:mainfrom
aws-mwaa:ferruzzi/bedrock/m2-customize-model

Conversation

@ferruzzi

@ferruzziferruzzi commented Apr 3, 2024

Copy link
Copy Markdown
Contributor

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
Copy Markdown
ContributorAuthor

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?

Comment threadairflow/providers/amazon/aws/hooks/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/operators/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/waiters/bedrock.json
Comment threaddocs/apache-airflow-providers-amazon/operators/bedrock.rst Outdated
Comment threaddocs/apache-airflow-providers-amazon/operators/bedrock.rst Outdated
Comment threadairflow/providers/amazon/aws/hooks/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/operators/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/operators/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/operators/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/sensors/bedrock.py Outdated
Comment threadtests/providers/amazon/aws/operators/test_bedrock.py Outdated
Comment threadtests/providers/amazon/aws/operators/test_bedrock.py Outdated
Comment threadtests/providers/amazon/aws/waiters/test_bedrock.py
Comment threadtests/system/providers/amazon/aws/example_bedrock.py
Comment threadtests/system/providers/amazon/aws/example_bedrock.py
@ferruzziferruzzi changed the title Amazon Bedrock -Model Customization JobsAmazon Bedrock - Model Customization JobsApr 5, 2024
Comment threadairflow/providers/amazon/aws/triggers/bedrock.py

@syedahsnsyedahsn left a comment

Copy link
Copy Markdown
Contributor

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 😄

Comment threadairflow/providers/amazon/aws/operators/bedrock.py Outdated
@ferruzzi

Copy link
Copy Markdown
ContributorAuthor

@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
Copy Markdown
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
Copy Markdown
ContributorAuthor

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

@ferruzziferruzzi 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
Comment threadairflow/providers/amazon/aws/operators/bedrock.py
Comment threadairflow/providers/amazon/aws/hooks/bedrock.py Outdated
Comment threadairflow/providers/amazon/aws/sensors/bedrock.py Outdated
@ferruzzi
ferruzzi merged commit 7ed31d5 into apache:mainApr 8, 2024
@ferruzzi
ferruzzi deleted the ferruzzi/bedrock/m2-customize-model branch April 8, 2024 20:22
Sign up for freeto 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

@ferruzzi@o-nikolas@Taragolis@vincbeck@syedahsn