Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi merged PR #38850: URL: https://github.com/apache/airflow/pull/38850 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on PR #38850: URL: https://github.com/apache/airflow/pull/38850#issuecomment-2048444921 Gah, static check is just a merge issue, imported the sensor twice. Easy fix. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on PR #38850: URL: https://github.com/apache/airflow/pull/38850#issuecomment-2048127209 Dependency merged; I'll get this rebased on top of that one and merge it this afternoon. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558544946 ## airflow/providers/amazon/aws/operators/bedrock.py: ## @@ -25,7 +25,10 @@ from airflow.exceptions import AirflowException Review Comment: I had to refactor the existing portion of the Bedrock system test, so this PR will need to be rebased on top of that and adjusted accordingly. Do not merge until https://github.com/apache/airflow/pull/38887 is merged and rebased in. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558543614 ## tests/system/providers/amazon/aws/example_bedrock.py: ## @@ -127,6 +137,7 @@ def delete_custom_model(model_name: str): output_data_uri=f"s3://{bucket_name}/myOutputData", ) # [END howto_operator_customize_model] +customize_model.wait_for_completion = False Review Comment: That fix should be [here](https://github.com/apache/airflow/pull/38887) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558148127 ## tests/system/providers/amazon/aws/example_bedrock.py: ## @@ -127,6 +137,7 @@ def delete_custom_model(model_name: str): output_data_uri=f"s3://{bucket_name}/myOutputData", ) # [END howto_operator_customize_model] +customize_model.wait_for_completion = False Review Comment: I don't think so. I am going to change the system test to use a branching operator like we discussed, but this was something that caught my eye while I was in there. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
o-nikolas commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558137372 ## tests/system/providers/amazon/aws/example_bedrock.py: ## @@ -127,6 +137,7 @@ def delete_custom_model(model_name: str): output_data_uri=f"s3://{bucket_name}/myOutputData", ) # [END howto_operator_customize_model] +customize_model.wait_for_completion = False Review Comment: Will this fix the stability issues we're seeing in the system test right now? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558135437 ## airflow/providers/amazon/aws/sensors/bedrock.py: ## @@ -17,36 +17,89 @@ # under the License. from __future__ import annotations +import abc from typing import TYPE_CHECKING, Any, Sequence from airflow.configuration import conf +from airflow.exceptions import AirflowException, AirflowSkipException +from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor -from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger +from airflow.providers.amazon.aws.triggers.bedrock import ( +BedrockCustomizeModelCompletedTrigger, +BedrockProvisionModelThroughputCompletedTrigger, +) from airflow.providers.amazon.aws.utils.mixins import aws_template_fields if TYPE_CHECKING: from airflow.utils.context import Context -from airflow.exceptions import AirflowException, AirflowSkipException -from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook +class BaseBedrockSensor(AwsBaseSensor[BedrockHook]): +""" +General sensor behavior for Amazon Bedrock. + +Subclasses must implement following methods: +- ``get_state()`` -class BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]): +Subclasses must set the following fields: +- ``INTERMEDIATE_STATES`` +- ``FAILURE_STATES` +- ``SUCCESS_STATES`` +- ``FAILURE_MESSAGE`` + +:param aws_conn_id: The Airflow connection used for AWS credentials. +If this is None or empty then the default boto3 behaviour is used. If +running Airflow in a distributed manner and aws_conn_id is None or +empty, then default boto3 configuration would be used (and must be +maintained on each worker node). +""" + +INTERMEDIATE_STATES: tuple[str, ...] = () +FAILURE_STATES: tuple[str, ...] = () +SUCCESS_STATES: tuple[str, ...] = () +FAILURE_MESSAGE = "" + +aws_hook_class = BedrockHook +ui_color = "#66c3ff" + +def __init__( +self, +deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), Review Comment: Looks like somehow I put the aws_conn_id docstring there instead. I'll fix it. Good catch, thanks. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1558108644 ## airflow/providers/amazon/aws/operators/bedrock.py: ## @@ -25,7 +25,10 @@ from airflow.exceptions import AirflowException Review Comment: dependency cleared. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
vincbeck commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1557596522 ## airflow/providers/amazon/aws/sensors/bedrock.py: ## @@ -17,36 +17,89 @@ # under the License. from __future__ import annotations +import abc from typing import TYPE_CHECKING, Any, Sequence from airflow.configuration import conf +from airflow.exceptions import AirflowException, AirflowSkipException +from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor -from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger +from airflow.providers.amazon.aws.triggers.bedrock import ( +BedrockCustomizeModelCompletedTrigger, +BedrockProvisionModelThroughputCompletedTrigger, +) from airflow.providers.amazon.aws.utils.mixins import aws_template_fields if TYPE_CHECKING: from airflow.utils.context import Context -from airflow.exceptions import AirflowException, AirflowSkipException -from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook +class BaseBedrockSensor(AwsBaseSensor[BedrockHook]): +""" +General sensor behavior for Amazon Bedrock. + +Subclasses must implement following methods: +- ``get_state()`` -class BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]): +Subclasses must set the following fields: +- ``INTERMEDIATE_STATES`` +- ``FAILURE_STATES` +- ``SUCCESS_STATES`` +- ``FAILURE_MESSAGE`` + +:param aws_conn_id: The Airflow connection used for AWS credentials. +If this is None or empty then the default boto3 behaviour is used. If +running Airflow in a distributed manner and aws_conn_id is None or +empty, then default boto3 configuration would be used (and must be +maintained on each worker node). +""" + +INTERMEDIATE_STATES: tuple[str, ...] = () +FAILURE_STATES: tuple[str, ...] = () +SUCCESS_STATES: tuple[str, ...] = () +FAILURE_MESSAGE = "" + +aws_hook_class = BedrockHook +ui_color = "#66c3ff" + +def __init__( +self, +deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), Review Comment: nit: missing from docstring -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1556503642 ## airflow/providers/amazon/aws/operators/bedrock.py: ## @@ -25,7 +25,10 @@ from airflow.exceptions import AirflowException Review Comment: This comment is just a safety gate to make sure this PR doesn't get merged out of order; anyone can feel free to resolve it once https://github.com/apache/airflow/pull/38849 is merged. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi commented on code in PR #38850: URL: https://github.com/apache/airflow/pull/38850#discussion_r1556503642 ## airflow/providers/amazon/aws/operators/bedrock.py: ## @@ -25,7 +25,10 @@ from airflow.exceptions import AirflowException Review Comment: This comment is just a safety gate to make sure this one doesn't get merged out of order; anyone can feel free to resolve it once https://github.com/apache/airflow/pull/38849 is merged. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
[PR] Amazon Bedrock - Model Throughput Provisioning [airflow]
ferruzzi opened a new pull request, #38850: URL: https://github.com/apache/airflow/pull/38850 Adds support for provisioning model throughput on Amazon Bedrock including: operators, sensors, waiters, and triggerer, as well as unit tests and system tests for all of the above, and doc updates. Manually tested in Breeze with the Trigger and with the Sensor: ![image](https://github.com/apache/airflow/assets/1920178/97573595-935a-43e9-a0ad-be223d6d53b9) Sits on top of a change in https://github.com/apache/airflow/pull/38849, please make sure that is merged before this one. --- **^ Add meaningful description above** Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed. In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x). In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org