Re: [PR] Amazon Bedrock - Model Throughput Provisioning [airflow]

2024-04-11 Thread via GitHub


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]

2024-04-10 Thread via GitHub


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]

2024-04-10 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-09 Thread via GitHub


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]

2024-04-08 Thread via GitHub


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]

2024-04-08 Thread via GitHub


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]

2024-04-08 Thread via GitHub


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