You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ferruzzi (via GitHub)" <gi...@apache.org> on 2024/04/03 01:25:23 UTC
[PR] Amazon Bedrock -Model Customization Jobs [airflow]
ferruzzi opened a new pull request, #38693:
URL: https://github.com/apache/airflow/pull/38693
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
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550638175
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
> We don't test the waiters for other services
TBH we test it, if we skip some of them then we should add it, some issue which happen in the past because we have lack of the simple tests https://github.com/apache/airflow/pull/33656
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552529331
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -37,30 +81,84 @@
) as dag:
test_context = sys_test_context_task()
env_id = test_context["ENV_ID"]
+ bucket_name = f"{env_id}-bedrock"
+ input_data_s3_key = f"{env_id}/train.jsonl"
Review Comment:
Good catch, but the filename came straight out of the Bedrock docs page. I hadn't heard of it either and thought it was odd as well, but Google said it is a known file format and it works in manual testing. So I'm reasonably sure it's not a typo.
```
The key difference [between json and jsonl] is in how they handle multiple
JSON objects. Regular JSON files are typically a single, self-contained
structure, while JSON Lines use a line-by-line format, allowing for easier
streaming and processing of individual objects.
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552138597
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
And I guess Batch has custom waiters into the two places 🤣
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi merged PR #38693:
URL: https://github.com/apache/airflow/pull/38693
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556213468
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +98,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param ensure_unique_job_name: If set to true, operator will check whether a model customization
+ job already exists for the name in the config and append the current timestamp if there is a
+ name conflict. (Default: True)
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "ensure_unique_job_name",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ ensure_unique_job_name: bool = True,
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.ensure_unique_job_name = ensure_unique_job_name
+ self.customization_job_kwargs = customization_job_kwargs or {}
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def execute(self, context: Context) -> dict:
+ response = {}
+ retry = True
+ while retry:
+ # If there is a name conflict and ensure_unique_job_name is True, append the current timestamp
+ # to the name and retry until there is no name conflict.
+ # - Break the loop when the API call returns success.
+ # - If the API returns an exception other than a name conflict, raise that exception.
+ # - If the API returns a name conflict and ensure_unique_job_name is false, raise that exception.
+ try:
+ # Ensure the loop is executed at least once, and not repeat unless explicitly set to do so.
+ retry = False
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
+ )
+ except ClientError as error:
+ if error.response["Error"]["Message"] != "The provided job name is currently in use.":
+ raise error
+ if not self.ensure_unique_job_name:
+ raise error
+ retry = True
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ if response["ResponseMetadata"]["HTTPStatusCode"] != 201:
+ raise AirflowException(f"Bedrock model customization job creation failed: {response}")
Review Comment:
and yes, according to [their docs](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelCustomizationJob.html), a successful request will always return a 201:
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556206876
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +98,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param ensure_unique_job_name: If set to true, operator will check whether a model customization
+ job already exists for the name in the config and append the current timestamp if there is a
+ name conflict. (Default: True)
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "ensure_unique_job_name",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ ensure_unique_job_name: bool = True,
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.ensure_unique_job_name = ensure_unique_job_name
+ self.customization_job_kwargs = customization_job_kwargs or {}
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def execute(self, context: Context) -> dict:
+ response = {}
+ retry = True
+ while retry:
+ # If there is a name conflict and ensure_unique_job_name is True, append the current timestamp
+ # to the name and retry until there is no name conflict.
+ # - Break the loop when the API call returns success.
+ # - If the API returns an exception other than a name conflict, raise that exception.
+ # - If the API returns a name conflict and ensure_unique_job_name is false, raise that exception.
+ try:
+ # Ensure the loop is executed at least once, and not repeat unless explicitly set to do so.
+ retry = False
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
+ )
+ except ClientError as error:
+ if error.response["Error"]["Message"] != "The provided job name is currently in use.":
+ raise error
+ if not self.ensure_unique_job_name:
+ raise error
+ retry = True
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ if response["ResponseMetadata"]["HTTPStatusCode"] != 201:
+ raise AirflowException(f"Bedrock model customization job creation failed: {response}")
Review Comment:
I am only handling the error case where the job name already exists so I can retry it, so I think this is still needed for other error cases:
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554297267
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
PR to move them all into their own files is here: https://github.com/apache/airflow/pull/38785
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550254242
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,150 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "check_if_job_exists",
+ "action_if_job_exists",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ check_if_job_exists: bool = True,
+ action_if_job_exists: str = "timestamp",
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.check_if_job_exists = check_if_job_exists
+ self.customization_job_kwargs = customization_job_kwargs or {}
+ if action_if_job_exists in {"timestamp", "fail"}:
+ self.action_if_job_exists = action_if_job_exists
+ else:
+ raise AirflowException(
+ f"Argument action_if_job_exists accepts only 'timestamp', and 'fail'. \
+ Provided value: '{action_if_job_exists}."
+ )
Review Comment:
done
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1553998433
##########
airflow/providers/amazon/aws/triggers/bedrock.py:
##########
@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when a Bedrock model customization job is complete.
+
+ :param job_name: The name of the Bedrock model customization job.
+ :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+ :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ *,
+ job_name: str,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ aws_conn_id: str | None = None,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"job_name": job_name},
Review Comment:
Looks right to me. I'll leave this open for now in case anyone else chimes 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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554216885
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
Alright, I decided to tweak this. Let's leave this comment open as a failsafe against merging for 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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552583500
##########
tests/providers/amazon/aws/waiters/test_bedrock.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+
+
+class TestBedrockCustomWaiters:
+ def test_service_waiters(self):
+ assert "model_customization_job_complete" in BedrockHook().list_waiters()
+
+
+class TestBedrockCustomWaitersBase:
Review Comment:
I jumped the gun on that a little. I'm already implementing more waiters and they'll be sharing that same fixture (which is used below [on L45](https://github.com/apache/airflow/pull/38693/files#diff-9d844f38bae728c36b62862b4197ffc0d54b238cf85d7841b40d7aa8571acf6cR45) )to mock the boto client.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552557529
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
I've actually been second-guessing this one. A parameter with only two possible values sounds a lot like a bool to me, in hindsight. I can't think of a good succinct name for it as a bool though.
To maybe shake some ideas loose, if it stays as a string, I'll add this to the description:
```
If "timestamp" is used and the job name already exists, the current timestamp
will be appended to the name in order to make it unique.
```
But perhaps I should replace `action_if_job_exists: str` (and all the associated input validation for it) with `append_timestamp_on_name_conflict: bool`?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552580948
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -18,42 +18,153 @@
from __future__ import annotations
import json
-from typing import Generator
+from typing import TYPE_CHECKING, Generator
from unittest import mock
import pytest
from moto import mock_aws
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
-
-MODEL_ID = "meta.llama2-13b-chat-v1"
-PROMPT = "A very important question."
-GENERATED_RESPONSE = "An important answer."
-MOCK_RESPONSE = json.dumps(
- {
- "generation": GENERATED_RESPONSE,
- "prompt_token_count": len(PROMPT),
- "generation_token_count": len(GENERATED_RESPONSE),
- "stop_reason": "stop",
- }
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
)
-
-@pytest.fixture
-def runtime_hook() -> Generator[BedrockRuntimeHook, None, None]:
- with mock_aws():
- yield BedrockRuntimeHook(aws_conn_id="aws_default")
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
class TestBedrockInvokeModelOperator:
- @mock.patch.object(BedrockRuntimeHook, "conn")
- def test_invoke_model_prompt_good_combinations(self, mock_conn):
- mock_conn.invoke_model.return_value["body"].read.return_value = MOCK_RESPONSE
+ def setup_method(self):
+ self.model_id = "meta.llama2-13b-chat-v1"
+ self.prompt = "A very important question."
+ self.generated_response = "An important answer."
Review Comment:
(Here and above) Yeah, it felt odd having them just hanging there outside of any methods. I'll move them to class-level fields.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552137707
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
I guess it is started into the custom_waiters, and after that it added into the separate one. So many thing need to be improved 🙄
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552560048
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "check_if_job_exists",
+ "action_if_job_exists",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ check_if_job_exists: bool = True,
+ action_if_job_exists: str = "timestamp",
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.check_if_job_exists = check_if_job_exists
+ self.customization_job_kwargs = customization_job_kwargs or {}
+ self.action_if_job_exists = action_if_job_exists.lower()
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def _validate_action_if_job_exists(self):
+ if self.action_if_job_exists not in self.valid_action_if_job_exists:
+ raise AirflowException(
+ f"Invalid value for argument action_if_job_exists {self.action_if_job_exists}; "
+ f"must be one of: {self.valid_action_if_job_exists}."
+ )
+
+ def execute(self, context: Context) -> dict:
+ self._validate_action_if_job_exists()
+
+ if self.check_if_job_exists and self.hook.job_name_exists(self.job_name):
+ if self.action_if_job_exists == "fail":
+ raise AirflowException(f"A Bedrock job with name {self.job_name} already exists.")
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
Review Comment:
Sorry, I don't follow. Is the suggestion to build something like a "job_parameters" dict in the init and just pass/unpack that here?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556206876
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +98,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param ensure_unique_job_name: If set to true, operator will check whether a model customization
+ job already exists for the name in the config and append the current timestamp if there is a
+ name conflict. (Default: True)
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "ensure_unique_job_name",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ ensure_unique_job_name: bool = True,
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.ensure_unique_job_name = ensure_unique_job_name
+ self.customization_job_kwargs = customization_job_kwargs or {}
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def execute(self, context: Context) -> dict:
+ response = {}
+ retry = True
+ while retry:
+ # If there is a name conflict and ensure_unique_job_name is True, append the current timestamp
+ # to the name and retry until there is no name conflict.
+ # - Break the loop when the API call returns success.
+ # - If the API returns an exception other than a name conflict, raise that exception.
+ # - If the API returns a name conflict and ensure_unique_job_name is false, raise that exception.
+ try:
+ # Ensure the loop is executed at least once, and not repeat unless explicitly set to do so.
+ retry = False
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
+ )
+ except ClientError as error:
+ if error.response["Error"]["Message"] != "The provided job name is currently in use.":
+ raise error
+ if not self.ensure_unique_job_name:
+ raise error
+ retry = True
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ if response["ResponseMetadata"]["HTTPStatusCode"] != 201:
+ raise AirflowException(f"Bedrock model customization job creation failed: {response}")
Review Comment:
I am only catching "job already exists" errors above so I can retry them, so i think this is still needed for other error cases:
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552191237
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
Yeah, the list just keeps growing. I don't know how you find time to do as much as you do for cleaning up old code. I think I want to split the test_custom_waiters monlith into per-service files for now, then we can worry about doing some kind of baseTest class for them later. That will at least make it obvious that they are expected. But I'll do that in another PR.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552572554
##########
tests/providers/amazon/aws/hooks/test_bedrock.py:
##########
@@ -16,7 +16,70 @@
# under the License.
from __future__ import annotations
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
+from unittest import mock
+
+import pytest
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+
+JOB_NAME = "testJobName"
+EXPECTED_STATUS = "InProgress"
+
+
+@pytest.fixture
+def mock_conn():
+ with mock.patch.object(BedrockHook, "conn") as _conn:
+ _conn.get_model_customization_job.return_value = {"jobName": JOB_NAME, "status": EXPECTED_STATUS}
+ yield _conn
+
+
+class TestBedrockHook:
+ def setup_method(self):
+ self.hook = BedrockHook()
+
+ self.validation_exception_error = ClientError(
+ error_response={"Error": {"Code": "ValidationException", "Message": ""}},
+ operation_name="GetModelCustomizationJob",
+ )
+
+ self.unexpected_exception = ClientError(
+ error_response={"Error": {"Code": "ExpiredTokenException", "Message": ""}},
+ operation_name="GetModelCustomizationJob",
+ )
Review Comment:
Moved to a class-level field.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550257495
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
I'm, not sure about this one. We don't test the other services and it's not really testing anything. In the Neptune one, we're mocking the return and not even making any asserts, just making sure they don't crash. I don't see any benefit 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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556195811
##########
airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -19,6 +19,37 @@
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+class BedrockHook(AwsBaseHook):
+ """
+ Interact with Amazon Bedrock.
+
+ Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.Client>`.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ client_type = "bedrock"
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = self.client_type
+ super().__init__(*args, **kwargs)
+
+ def _get_job_by_name(self, job_name: str):
+ return self.conn.get_model_customization_job(jobIdentifier=job_name)
+
+ def get_customize_model_job_state(self, job_name: str) -> str:
+ state = self._get_job_by_name(job_name)["status"]
+ self.log.info("Job '%s' state: %s", job_name, state)
+ return state
+
+ def get_job_arn(self, job_name: str) -> str:
+ return self._get_job_by_name(job_name)["jobArn"]
Review Comment:
Yeah, they felt handy and made the Sensor unit tests a little easier to mock, but I'll drop them.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554213131
##########
tests/providers/amazon/aws/waiters/test_bedrock.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+
+
+class TestBedrockCustomWaiters:
+ def test_service_waiters(self):
+ assert "model_customization_job_complete" in BedrockHook().list_waiters()
+
+
+class TestBedrockCustomWaitersBase:
Review Comment:
Yeah, the autouse made sense after your previous message. My last one was just referencing that you have a base class only used once, but as suspected and you've confirmed, more uses are coming. So it's all good :)
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550257495
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
I'm not sure about this one. We don't test the waiters for other services and it's not really testing anything. In the Neptune one that was just added, we're mocking the return and not even making any asserts, just making sure they don't crash. I don't see any benefit 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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on PR #38693:
URL: https://github.com/apache/airflow/pull/38693#issuecomment-2040697189
@o-nikolas Reworked the `while True` so it won't repeat by default
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on PR #38693:
URL: https://github.com/apache/airflow/pull/38693#issuecomment-2033530120
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?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1549147317
##########
airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -16,9 +16,54 @@
# under the License.
from __future__ import annotations
+from botocore.exceptions import ClientError
+
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+class BedrockHook(AwsBaseHook):
+ """
+ Interact with Amazon Bedrock.
+
+ Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.Client>`.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ client_type = "bedrock"
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = self.client_type
+ super().__init__(*args, **kwargs)
+
+ def _get_job_by_name(self, job_name: str):
+ return self.conn.get_model_customization_job(jobIdentifier=job_name)
+
+ def get_customize_model_job_state(self, job_name) -> str:
+ state = self._get_job_by_name(job_name)["status"]
+ self.log.info("Job '%s' state: %s", job_name, state)
+ return state
+
+ def job_name_exists(self, job_name: str) -> bool:
+ try:
+ self._get_job_by_name(job_name)
+ self.log.info("Verified that job name '%s' does exist.", job_name)
+ return True
+ except ClientError as e:
+ if e.response["Error"]["Code"] == "ValidationException":
+ self.log.info("Job name '%s' does not exist.", job_name)
+ return False
+ else:
+ raise e
Review Comment:
```suggestion
raise
```
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
Worthwhile to tests waiter separately, time to time we have invalid definitions in it:
https://github.com/apache/airflow/tree/main/tests/providers/amazon/aws/waiters
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,150 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "check_if_job_exists",
+ "action_if_job_exists",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ check_if_job_exists: bool = True,
+ action_if_job_exists: str = "timestamp",
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.check_if_job_exists = check_if_job_exists
+ self.customization_job_kwargs = customization_job_kwargs or {}
+ if action_if_job_exists in {"timestamp", "fail"}:
+ self.action_if_job_exists = action_if_job_exists
+ else:
+ raise AirflowException(
+ f"Argument action_if_job_exists accepts only 'timestamp', and 'fail'. \
+ Provided value: '{action_if_job_exists}."
+ )
Review Comment:
`action_if_job_exists` is a templated field, so better move this validation into place after templates are rendered, e.g. into the execute method.
Pros: It does not raise random error on Jinja template or XComArg (or simmilar)
Cons: No fast fail
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550640629
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
You know what, I'll admit that the fact that some are combined into [test_custom_waiters.py](https://github.com/apache/airflow/blob/main/tests/providers/amazon/aws/waiters/test_custom_waiters.py) and others were in their own files messed me up.
I thought we only had tests for Neptune and Databrew because the others are "hiding".
I still don't really see the point of those ones, but I'll add them in their own file, and we can split the combined file out into one-file-per-service like everywhere else so it's more obvious in the future.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550689743
##########
airflow/providers/amazon/aws/waiters/bedrock.json:
##########
Review Comment:
Added some. I think they about cover it. Let me know what you think.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552578537
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
I'm happy either way, I don't want to be too prescriptive. But if you were going to change it to a bool, I much prefer your approach of `ensure_unique_name` (or something like that).
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552581272
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -16,17 +16,61 @@
# under the License.
from __future__ import annotations
+import json
from datetime import datetime
+from botocore.exceptions import ClientError
+
+from airflow.decorators import task
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
-sys_test_context_task = SystemTestContextBuilder().build()
+# Externally fetched variables:
+ROLE_ARN_KEY = "ROLE_ARN"
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
DAG_ID = "example_bedrock"
+
+# Creating a custom model takes nearly two hours. If SKIP_LONG_TASKS is True then set
Review Comment:
That's very sad. Okay, this is a nice workaround if that's the case :+1:
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552442572
##########
docs/apache-airflow-providers-amazon/operators/bedrock.rst:
##########
@@ -65,6 +65,44 @@ To invoke an Amazon Titan model you would use:
For details on the different formats, see `Inference parameters for foundation models <https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html>`__
+.. _howto/operator:BedrockCustomizeModelOperator:
+
+Customize an existing Amazon Bedrock Model
+==========================================
+
+To create a fine-tuning job to customize a base model, you can use
+:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockCustomizeModelOperator`.
+
+Model-customization jobs are asynchronous and the completion time depends on the base model
+and the training/validation data size. To monitor the state of the job, you can use the
+"model_customization_job_complete" Waiter, the
+:class:`~airflow.providers.amazon.aws.sensors.bedrock.BedrockCustomizeModelCompletedSensor` Sensor,
+or the :class:`~airflow.providers.amazon.aws.triggersBedrockCustomizeModelCompletedTrigger` Trigger.
Review Comment:
missed a period in the class path, surprised the build-docs test didn't catch that. Will fix.
##########
docs/apache-airflow-providers-amazon/operators/bedrock.rst:
##########
@@ -65,6 +65,44 @@ To invoke an Amazon Titan model you would use:
For details on the different formats, see `Inference parameters for foundation models <https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html>`__
+.. _howto/operator:BedrockCustomizeModelOperator:
+
+Customize an existing Amazon Bedrock Model
+==========================================
+
+To create a fine-tuning job to customize a base model, you can use
+:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockCustomizeModelOperator`.
+
+Model-customization jobs are asynchronous and the completion time depends on the base model
+and the training/validation data size. To monitor the state of the job, you can use the
+"model_customization_job_complete" Waiter, the
+:class:`~airflow.providers.amazon.aws.sensors.bedrock.BedrockCustomizeModelCompletedSensor` Sensor,
+or the :class:`~airflow.providers.amazon.aws.triggersBedrockCustomizeModelCompletedTrigger` Trigger.
+
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_bedrock.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_customize_model]
+ :end-before: [END howto_operator_customize_model]
+
+
+Sensors
+-------
+
+.. _howto/sensor:BedrockCustomizeModelCompletedSensor:
+
+Wait for an Amazon Bedrock customize model job
+==============================================
+
+To wait on the state of an AWS CloudFormation stack creation until it reaches a terminal state you can use
Review Comment:
copypasta :( Will 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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554074878
##########
tests/providers/amazon/aws/waiters/test_bedrock.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+
+
+class TestBedrockCustomWaiters:
+ def test_service_waiters(self):
+ assert "model_customization_job_complete" in BedrockHook().list_waiters()
+
+
+class TestBedrockCustomWaitersBase:
Review Comment:
I think maybe the autouse is throwing you off. It creates `self.client` which is then inherited/used/modified in the mock_get_job in TestModelCustomizationJobCompleteWaiter.
In addition, it will also be used by TestProvisionThroughputJobCompleteWaiter which I will have a PR up for either later today or Monday. I know we generally don't create the helper until it is used a second time, but I got ahead of myself on this one, but I assure you it will get reused :P
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554032096
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
Other than this dangling thread, it looks good to me, approving :+1:
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552532744
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -16,17 +16,61 @@
# under the License.
from __future__ import annotations
+import json
from datetime import datetime
+from botocore.exceptions import ClientError
+
+from airflow.decorators import task
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
-sys_test_context_task = SystemTestContextBuilder().build()
+# Externally fetched variables:
+ROLE_ARN_KEY = "ROLE_ARN"
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
DAG_ID = "example_bedrock"
+
+# Creating a custom model takes nearly two hours. If SKIP_LONG_TASKS is True then set
Review Comment:
Negative. That is the smallest possible job I can come up with (feeding the smallest model one datapoint to learn), there is no cancel command, and it can not be deleted until it is in a terminal state, whether that is a successful build or a failed job.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552529331
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -37,30 +81,84 @@
) as dag:
test_context = sys_test_context_task()
env_id = test_context["ENV_ID"]
+ bucket_name = f"{env_id}-bedrock"
+ input_data_s3_key = f"{env_id}/train.jsonl"
Review Comment:
Good catch, but the filename came straight out of the Bedrock docs page. I hadn't hard of it either and thought it was odd as well, but Google said it is a known file format and it works in manual testing. So I'm reasonably sure it's not a typo.
```
The key difference [between json and jsonl] is in how they handle multiple JSON objects. Regular JSON files are typically a single, self-contained structure, while JSON Lines use a line-by-line format, allowing for easier streaming and processing of individual objects.
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552586854
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -16,17 +16,61 @@
# under the License.
from __future__ import annotations
+import json
from datetime import datetime
+from botocore.exceptions import ClientError
+
+from airflow.decorators import task
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
-sys_test_context_task = SystemTestContextBuilder().build()
+# Externally fetched variables:
+ROLE_ARN_KEY = "ROLE_ARN"
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
DAG_ID = "example_bedrock"
+
+# Creating a custom model takes nearly two hours. If SKIP_LONG_TASKS is True then set
Review Comment:
I've been putting thought into using a BranchingOperator to just skip them entirely if the SKIP_LONG flag is true, rather than assign them an improbable trigger_rule. Then maybe we can figure out some trick to get our dashboard CI to run the long version from time to time by setting an envvar or something.... I don't know, needs some thinking, but I think this is a decent compromise for the time being.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1550181713
##########
airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -16,9 +16,54 @@
# under the License.
from __future__ import annotations
+from botocore.exceptions import ClientError
+
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+class BedrockHook(AwsBaseHook):
+ """
+ Interact with Amazon Bedrock.
+
+ Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.Client>`.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ client_type = "bedrock"
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = self.client_type
+ super().__init__(*args, **kwargs)
+
+ def _get_job_by_name(self, job_name: str):
+ return self.conn.get_model_customization_job(jobIdentifier=job_name)
+
+ def get_customize_model_job_state(self, job_name) -> str:
+ state = self._get_job_by_name(job_name)["status"]
+ self.log.info("Job '%s' state: %s", job_name, state)
+ return state
+
+ def job_name_exists(self, job_name: str) -> bool:
+ try:
+ self._get_job_by_name(job_name)
+ self.log.info("Verified that job name '%s' does exist.", job_name)
+ return True
+ except ClientError as e:
+ if e.response["Error"]["Code"] == "ValidationException":
+ self.log.info("Job name '%s' does not exist.", job_name)
+ return False
+ else:
+ raise e
Review Comment:
fair.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552557529
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
UI've actually been second-guessing this one. A parameter with only two possible values sounds a lot like a bool to me, in hindsight. I can't think of a good succinct name for it as a bool though.
To maybe shake some ideas loose, if it stays as a string, I'll add this to the description:
```
If "timestamp" is used and a the job name already exists, the current timestamp
will be appended tot he name in order to make it unique.
```
But perhaps I should replace `action_if_job_exists: str` (and all the associated input validation for it) with `append_timestamp_on_name_conflict: bool`?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on PR #38693:
URL: https://github.com/apache/airflow/pull/38693#issuecomment-2033487899
Just realized I didn't finish making the sensor deferrable, I'll add the execute_complete() to that tomorrow.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552580790
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "check_if_job_exists",
+ "action_if_job_exists",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ check_if_job_exists: bool = True,
+ action_if_job_exists: str = "timestamp",
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.check_if_job_exists = check_if_job_exists
+ self.customization_job_kwargs = customization_job_kwargs or {}
+ self.action_if_job_exists = action_if_job_exists.lower()
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def _validate_action_if_job_exists(self):
+ if self.action_if_job_exists not in self.valid_action_if_job_exists:
+ raise AirflowException(
+ f"Invalid value for argument action_if_job_exists {self.action_if_job_exists}; "
+ f"must be one of: {self.valid_action_if_job_exists}."
+ )
+
+ def execute(self, context: Context) -> dict:
+ self._validate_action_if_job_exists()
+
+ if self.check_if_job_exists and self.hook.job_name_exists(self.job_name):
+ if self.action_if_job_exists == "fail":
+ raise AirflowException(f"A Bedrock job with name {self.job_name} already exists.")
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
Review Comment:
I think this is fine, it was really a comment on the boto API that it just takes all the customization kwargs as top level args. I expected some struct that held all the customization values.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552558849
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
Or maybe `ensure_unique_name: bool`? It returns the job name anyway, do the users need to care HOW it will be uniquified? I think this is a good one.... I'll think it over a bit more.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on PR #38693:
URL: https://github.com/apache/airflow/pull/38693#issuecomment-2040573795
@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?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552493564
##########
airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -16,9 +16,53 @@
# under the License.
from __future__ import annotations
+from botocore.exceptions import ClientError
+
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+class BedrockHook(AwsBaseHook):
+ """
+ Interact with Amazon Bedrock.
+
+ Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.Client>`.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ client_type = "bedrock"
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = self.client_type
+ super().__init__(*args, **kwargs)
+
+ def _get_job_by_name(self, job_name: str):
+ return self.conn.get_model_customization_job(jobIdentifier=job_name)
+
+ def get_customize_model_job_state(self, job_name) -> str:
Review Comment:
Missing type annotation on job_name
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -18,42 +18,153 @@
from __future__ import annotations
import json
-from typing import Generator
+from typing import TYPE_CHECKING, Generator
from unittest import mock
import pytest
from moto import mock_aws
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
-
-MODEL_ID = "meta.llama2-13b-chat-v1"
-PROMPT = "A very important question."
-GENERATED_RESPONSE = "An important answer."
-MOCK_RESPONSE = json.dumps(
- {
- "generation": GENERATED_RESPONSE,
- "prompt_token_count": len(PROMPT),
- "generation_token_count": len(GENERATED_RESPONSE),
- "stop_reason": "stop",
- }
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
)
-
-@pytest.fixture
-def runtime_hook() -> Generator[BedrockRuntimeHook, None, None]:
- with mock_aws():
- yield BedrockRuntimeHook(aws_conn_id="aws_default")
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
class TestBedrockInvokeModelOperator:
- @mock.patch.object(BedrockRuntimeHook, "conn")
- def test_invoke_model_prompt_good_combinations(self, mock_conn):
- mock_conn.invoke_model.return_value["body"].read.return_value = MOCK_RESPONSE
+ def setup_method(self):
+ self.model_id = "meta.llama2-13b-chat-v1"
+ self.prompt = "A very important question."
+ self.generated_response = "An important answer."
Review Comment:
This can just be done once?
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "check_if_job_exists",
+ "action_if_job_exists",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ check_if_job_exists: bool = True,
+ action_if_job_exists: str = "timestamp",
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.check_if_job_exists = check_if_job_exists
+ self.customization_job_kwargs = customization_job_kwargs or {}
+ self.action_if_job_exists = action_if_job_exists.lower()
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def _validate_action_if_job_exists(self):
+ if self.action_if_job_exists not in self.valid_action_if_job_exists:
+ raise AirflowException(
+ f"Invalid value for argument action_if_job_exists {self.action_if_job_exists}; "
+ f"must be one of: {self.valid_action_if_job_exists}."
+ )
+
+ def execute(self, context: Context) -> dict:
+ self._validate_action_if_job_exists()
+
+ if self.check_if_job_exists and self.hook.job_name_exists(self.job_name):
+ if self.action_if_job_exists == "fail":
+ raise AirflowException(f"A Bedrock job with name {self.job_name} already exists.")
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
Review Comment:
Interesting that those are all top-level and don't live within some structure.
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
Review Comment:
Fail is pretty self explanatory, but can you expand what `timestamp` will do? I had to read the code to figure it out.
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger
+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 BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ """
+ Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor`
+
+
+ :param job_name: The name of the Bedrock model customization job.
+
+ :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore
Review Comment:
```suggestion
:param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
```
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +96,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param check_if_job_exists: If set to true, operator will check whether a model customization
+ job already exists for the name in the config. (Default: True)
+ :param action_if_job_exists: Behavior if the job name already exists. Options are "timestamp" (default),
+ and "fail"
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
Review Comment:
Mention the defaults for these too?
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger
+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 BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ """
+ Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor`
+
+
+ :param job_name: The name of the Bedrock model customization job.
+
+ :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore
+ module to be installed.
+ (default: False, but can be overridden in config file by setting default_deferrable to True)
+ :param max_retries: Number of times before returning the current state, defaults to None
+ :param poke_interval: Polling period in seconds to check for the status of the job.
Review Comment:
Mention defaults?
##########
airflow/providers/amazon/aws/triggers/bedrock.py:
##########
@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when a Bedrock model customization job is complete.
+
+ :param job_name: The name of the Bedrock model customization job.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
Review Comment:
Mention defaults?
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -18,42 +18,153 @@
from __future__ import annotations
import json
-from typing import Generator
+from typing import TYPE_CHECKING, Generator
from unittest import mock
import pytest
from moto import mock_aws
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
-
-MODEL_ID = "meta.llama2-13b-chat-v1"
-PROMPT = "A very important question."
-GENERATED_RESPONSE = "An important answer."
-MOCK_RESPONSE = json.dumps(
- {
- "generation": GENERATED_RESPONSE,
- "prompt_token_count": len(PROMPT),
- "generation_token_count": len(GENERATED_RESPONSE),
- "stop_reason": "stop",
- }
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
)
-
-@pytest.fixture
-def runtime_hook() -> Generator[BedrockRuntimeHook, None, None]:
- with mock_aws():
- yield BedrockRuntimeHook(aws_conn_id="aws_default")
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
class TestBedrockInvokeModelOperator:
- @mock.patch.object(BedrockRuntimeHook, "conn")
- def test_invoke_model_prompt_good_combinations(self, mock_conn):
- mock_conn.invoke_model.return_value["body"].read.return_value = MOCK_RESPONSE
+ def setup_method(self):
+ self.model_id = "meta.llama2-13b-chat-v1"
+ self.prompt = "A very important question."
+ self.generated_response = "An important answer."
+
+ @pytest.fixture
+ def mock_runtime_conn(self) -> Generator[BaseAwsConnection, None, None]:
+ with mock.patch.object(BedrockRuntimeHook, "conn") as _conn:
+ _conn.invoke_model.return_value["body"].read.return_value = json.dumps(
+ {
+ "generation": self.generated_response,
+ "prompt_token_count": len(self.prompt),
+ "generation_token_count": len(self.generated_response),
+ "stop_reason": "stop",
+ }
+ )
+ yield _conn
+
+ @pytest.fixture
+ def runtime_hook(self) -> Generator[BedrockRuntimeHook, None, None]:
+ with mock_aws():
+ yield BedrockRuntimeHook(aws_conn_id="aws_default")
+
+ def test_invoke_model_prompt_good_combinations(self, mock_runtime_conn):
operator = BedrockInvokeModelOperator(
- task_id="test_task", model_id=MODEL_ID, input_data={"input_data": {"prompt": PROMPT}}
+ task_id="test_task", model_id=self.model_id, input_data={"input_data": {"prompt": self.prompt}}
)
response = operator.execute({})
- assert response["generation"] == GENERATED_RESPONSE
+ assert response["generation"] == self.generated_response
+
+
+class TestBedrockCustomizeModelOperator:
+ @pytest.fixture
+ def mock_conn(self) -> Generator[BaseAwsConnection, None, None]:
+ with mock.patch.object(BedrockHook, "conn") as _conn:
+ _conn.create_model_customization_job.return_value = {
+ "ResponseMetadata": {"HTTPStatusCode": 201},
+ "jobArn": self.custom_job_arn,
+ }
+ _conn.get_model_customization_job.return_value = {
+ "jobName": self.customize_model_job_name,
+ "status": "InProgress",
+ }
+ yield _conn
+
+ @pytest.fixture
+ def bedrock_hook(self) -> Generator[BedrockHook, None, None]:
+ with mock_aws():
+ hook = BedrockHook(aws_conn_id="aws_default")
+ yield hook
+
+ def setup_method(self):
+ self.custom_job_arn = "valid_arn"
+ self.customize_model_job_name = "testModelJob"
+
+ self.operator = BedrockCustomizeModelOperator(
+ task_id="test_task",
+ job_name=self.customize_model_job_name,
+ custom_model_name="testModelName",
+ role_arn="valid_arn",
+ base_model_id="base_model_id",
+ hyperparameters={
+ "epochCount": "1",
+ "batchSize": "1",
+ "learningRate": ".0005",
+ "learningRateWarmupSteps": "0",
+ },
+ training_data_uri="s3://uri",
+ output_data_uri="s3://uri/output",
+ )
+ self.operator.defer = mock.MagicMock()
+
+ @pytest.mark.parametrize(
+ "wait_for_completion, deferrable",
+ [
+ pytest.param(False, False, id="no_wait"),
+ pytest.param(True, False, id="wait"),
+ pytest.param(False, True, id="defer"),
+ ],
+ )
+ @mock.patch.object(BedrockHook, "get_waiter")
+ def test_customize_model_wait_combinations(
+ self, _, wait_for_completion, deferrable, mock_conn, bedrock_hook
+ ):
+ self.operator.wait_for_completion = wait_for_completion
+ self.operator.deferrable = deferrable
+
+ response = self.operator.execute({})
+
+ assert response == self.custom_job_arn
+ assert bedrock_hook.get_waiter.call_count == wait_for_completion
+ assert self.operator.defer.call_count == deferrable
+
+ @pytest.mark.parametrize(
+ "action_if_job_exists, succeeds",
+ [
+ pytest.param("timestamp", True, id="timestamp"),
+ pytest.param("fail", True, id="fail"),
+ pytest.param("call me maybe", False, id="invalid"),
Review Comment:
:laughing:
##########
tests/providers/amazon/aws/waiters/test_bedrock.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+
+
+class TestBedrockCustomWaiters:
+ def test_service_waiters(self):
+ assert "model_customization_job_complete" in BedrockHook().list_waiters()
+
+
+class TestBedrockCustomWaitersBase:
Review Comment:
There are no test methods in this test class? Just an unused fixture? Or am I missing something?
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -37,30 +81,84 @@
) as dag:
test_context = sys_test_context_task()
env_id = test_context["ENV_ID"]
+ bucket_name = f"{env_id}-bedrock"
+ input_data_s3_key = f"{env_id}/train.jsonl"
Review Comment:
`json1` is a weird file name, is that a typo?
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -16,17 +16,61 @@
# under the License.
from __future__ import annotations
+import json
from datetime import datetime
+from botocore.exceptions import ClientError
+
+from airflow.decorators import task
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
-sys_test_context_task = SystemTestContextBuilder().build()
+# Externally fetched variables:
+ROLE_ARN_KEY = "ROLE_ARN"
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
DAG_ID = "example_bedrock"
+
+# Creating a custom model takes nearly two hours. If SKIP_LONG_TASKS is True then set
Review Comment:
This is a shame, is there no way to speed this up, or even cancel it once we ensure it has started correctly?
##########
tests/providers/amazon/aws/hooks/test_bedrock.py:
##########
@@ -16,7 +16,70 @@
# under the License.
from __future__ import annotations
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
+from unittest import mock
+
+import pytest
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+
+JOB_NAME = "testJobName"
+EXPECTED_STATUS = "InProgress"
+
+
+@pytest.fixture
+def mock_conn():
+ with mock.patch.object(BedrockHook, "conn") as _conn:
+ _conn.get_model_customization_job.return_value = {"jobName": JOB_NAME, "status": EXPECTED_STATUS}
+ yield _conn
+
+
+class TestBedrockHook:
+ def setup_method(self):
+ self.hook = BedrockHook()
+
+ self.validation_exception_error = ClientError(
+ error_response={"Error": {"Code": "ValidationException", "Message": ""}},
+ operation_name="GetModelCustomizationJob",
+ )
+
+ self.unexpected_exception = ClientError(
+ error_response={"Error": {"Code": "ExpiredTokenException", "Message": ""}},
+ operation_name="GetModelCustomizationJob",
+ )
Review Comment:
Why are you creating these before each test? They seem pretty static to me.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock -Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552529331
##########
tests/system/providers/amazon/aws/example_bedrock.py:
##########
@@ -37,30 +81,84 @@
) as dag:
test_context = sys_test_context_task()
env_id = test_context["ENV_ID"]
+ bucket_name = f"{env_id}-bedrock"
+ input_data_s3_key = f"{env_id}/train.jsonl"
Review Comment:
Good catch, but the filename came straight out of the Bedrock docs page. I hadn't hard of it either and thought it was odd as well, but Google said it is a known file format and it works in manual testing. So I'm reasonably sure it's not a typo.
```
The key difference [between json and jsonl] is in how they handle multiple
JSON objects. Regular JSON files are typically a single, self-contained
structure, while JSON Lines use a line-by-line format, allowing for easier
streaming and processing of individual objects.
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1555532965
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -91,3 +98,155 @@ def execute(self, context: Context) -> dict[str, str | int]:
self.log.info("Bedrock %s prompt: %s", self.model_id, self.input_data)
self.log.info("Bedrock model response: %s", response_body)
return response_body
+
+
+class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create a fine-tuning job to customize a base model.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockCustomizeModelOperator`
+
+ :param job_name: A unique name for the fine-tuning job.
+ :param custom_model_name: A name for the custom model being created.
+ :param role_arn: The Amazon Resource Name (ARN) of an IAM role that Amazon Bedrock can assume
+ to perform tasks on your behalf.
+ :param base_model_id: Name of the base model.
+ :param training_data_uri: The S3 URI where the training data is stored.
+ :param output_data_uri: The S3 URI where the output data is stored.
+ :param hyperparameters: Parameters related to tuning the model.
+ :param ensure_unique_job_name: If set to true, operator will check whether a model customization
+ job already exists for the name in the config and append the current timestamp if there is a
+ name conflict. (Default: True)
+ :param customization_job_kwargs: Any optional parameters to pass to the API.
+
+ :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
+ :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_name",
+ "custom_model_name",
+ "role_arn",
+ "base_model_id",
+ "hyperparameters",
+ "ensure_unique_job_name",
+ "customization_job_kwargs",
+ )
+
+ def __init__(
+ self,
+ job_name: str,
+ custom_model_name: str,
+ role_arn: str,
+ base_model_id: str,
+ training_data_uri: str,
+ output_data_uri: str,
+ hyperparameters: dict[str, str],
+ ensure_unique_job_name: bool = True,
+ customization_job_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ self.job_name = job_name
+ self.custom_model_name = custom_model_name
+ self.role_arn = role_arn
+ self.base_model_id = base_model_id
+ self.training_data_config = {"s3Uri": training_data_uri}
+ self.output_data_config = {"s3Uri": output_data_uri}
+ self.hyperparameters = hyperparameters
+ self.ensure_unique_job_name = ensure_unique_job_name
+ self.customization_job_kwargs = customization_job_kwargs or {}
+
+ self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+ event = validate_execute_complete_event(event)
+
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
+ return self.hook.get_job_arn(event["job_name"])
+
+ def execute(self, context: Context) -> dict:
+ response = {}
+ retry = True
+ while retry:
+ # If there is a name conflict and ensure_unique_job_name is True, append the current timestamp
+ # to the name and retry until there is no name conflict.
+ # - Break the loop when the API call returns success.
+ # - If the API returns an exception other than a name conflict, raise that exception.
+ # - If the API returns a name conflict and ensure_unique_job_name is false, raise that exception.
+ try:
+ # Ensure the loop is executed at least once, and not repeat unless explicitly set to do so.
+ retry = False
+ self.log.info("Creating Bedrock model customization job '%s'.", self.job_name)
+
+ response = self.hook.conn.create_model_customization_job(
+ jobName=self.job_name,
+ customModelName=self.custom_model_name,
+ roleArn=self.role_arn,
+ baseModelIdentifier=self.base_model_id,
+ trainingDataConfig=self.training_data_config,
+ outputDataConfig=self.output_data_config,
+ hyperParameters=self.hyperparameters,
+ **self.customization_job_kwargs,
+ )
+ except ClientError as error:
+ if error.response["Error"]["Message"] != "The provided job name is currently in use.":
+ raise error
+ if not self.ensure_unique_job_name:
+ raise error
+ retry = True
+ self.job_name = f"{self.job_name}-{int(utcnow().timestamp())}"
+ self.log.info("Changed job name to '%s' to avoid collision.", self.job_name)
+
+ if response["ResponseMetadata"]["HTTPStatusCode"] != 201:
+ raise AirflowException(f"Bedrock model customization job creation failed: {response}")
Review Comment:
I've just wondering is it always return [201](https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/201) status code? And is this check required if we other errors didn't raised?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1553990772
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger
+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 BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ """
+ Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor`
+
+
+ :param job_name: The name of the Bedrock model customization job.
+
+ :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore
Review Comment:
HAH! That was actually copypasta fro SQS and I didn't catch it. I'll fix this one now and fix SQS in another PR
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1553827975
##########
airflow/providers/amazon/aws/triggers/bedrock.py:
##########
@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when a Bedrock model customization job is complete.
+
+ :param job_name: The name of the Bedrock model customization job.
+ :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+ :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ *,
+ job_name: str,
+ waiter_delay: int = 120,
+ waiter_max_attempts: int = 75,
+ aws_conn_id: str | None = None,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"job_name": job_name},
Review Comment:
It may be just something on my screen, but this looks like its indented incorrectly? Not a big deal either way.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1554031431
##########
tests/providers/amazon/aws/waiters/test_bedrock.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.sensors.bedrock import BedrockCustomizeModelCompletedSensor
+
+
+class TestBedrockCustomWaiters:
+ def test_service_waiters(self):
+ assert "model_customization_job_complete" in BedrockHook().list_waiters()
+
+
+class TestBedrockCustomWaitersBase:
Review Comment:
Ahh I see, so this base class will be used by future waiter test classes?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #38693:
URL: https://github.com/apache/airflow/pull/38693#issuecomment-2040644595
> @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.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556199352
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger
+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 BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ """
+ Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor`
+
+
+ :param job_name: The name of the Bedrock model customization job.
+
+ :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+ module to be installed.
+ (default: False, but can be overridden in config file by setting default_deferrable to True)
+ :param max_retries: Number of times before returning the current state. (default: 75)
+ :param poke_interval: Polling period in seconds to check for the status of the job. (default: 120)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ INTERMEDIATE_STATES = ("InProgress",)
+ FAILURE_STATES = ("Failed", "Stopping", "Stopped")
+ SUCCESS_STATES = ("Completed",)
+ FAILURE_MESSAGE = "Bedrock model customization job sensor failed."
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields("job_name")
+ ui_color = "#66c3ff"
+
+ def __init__(
+ self,
+ *,
+ job_name: str,
+ max_retries: int = 75,
+ poke_interval: int = 120,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.job_name = job_name
+ self.poke_interval = poke_interval
+ self.max_retries = max_retries
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> Any:
+ if self.deferrable:
+ self.defer(
+ trigger=BedrockCustomizeModelCompletedTrigger(
+ job_name=self.job_name,
+ waiter_delay=int(self.poke_interval),
+ waiter_max_attempts=self.max_retries,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="poke",
+ )
+ else:
+ super().execute(context=context)
+
+ def poke(self, context: Context) -> bool:
+ state = self.hook.get_customize_model_job_state(self.job_name)
+
+ if state in self.FAILURE_STATES:
+ # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
+ if self.soft_fail:
+ raise AirflowSkipException(self.FAILURE_MESSAGE)
+ raise AirflowException(self.FAILURE_MESSAGE)
+
+ if state in self.INTERMEDIATE_STATES:
+ return False
+ return True
Review Comment:
HAH! Yes, of course.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Model Customization Jobs [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1556043966
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.bedrock import BedrockCustomizeModelCompletedTrigger
+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 BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ """
+ Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor`
+
+
+ :param job_name: The name of the Bedrock model customization job.
+
+ :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+ module to be installed.
+ (default: False, but can be overridden in config file by setting default_deferrable to True)
+ :param max_retries: Number of times before returning the current state. (default: 75)
+ :param poke_interval: Polling period in seconds to check for the status of the job. (default: 120)
+ :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).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ INTERMEDIATE_STATES = ("InProgress",)
+ FAILURE_STATES = ("Failed", "Stopping", "Stopped")
+ SUCCESS_STATES = ("Completed",)
+ FAILURE_MESSAGE = "Bedrock model customization job sensor failed."
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields("job_name")
+ ui_color = "#66c3ff"
+
+ def __init__(
+ self,
+ *,
+ job_name: str,
+ max_retries: int = 75,
+ poke_interval: int = 120,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.job_name = job_name
+ self.poke_interval = poke_interval
+ self.max_retries = max_retries
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> Any:
+ if self.deferrable:
+ self.defer(
+ trigger=BedrockCustomizeModelCompletedTrigger(
+ job_name=self.job_name,
+ waiter_delay=int(self.poke_interval),
+ waiter_max_attempts=self.max_retries,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="poke",
+ )
+ else:
+ super().execute(context=context)
+
+ def poke(self, context: Context) -> bool:
+ state = self.hook.get_customize_model_job_state(self.job_name)
+
+ if state in self.FAILURE_STATES:
+ # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
+ if self.soft_fail:
+ raise AirflowSkipException(self.FAILURE_MESSAGE)
+ raise AirflowException(self.FAILURE_MESSAGE)
+
+ if state in self.INTERMEDIATE_STATES:
+ return False
+ return True
Review Comment:
```suggestion
return not state in self.INTERMEDIATE_STATES
```
##########
airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -19,6 +19,37 @@
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+class BedrockHook(AwsBaseHook):
+ """
+ Interact with Amazon Bedrock.
+
+ Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.Client>`.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ client_type = "bedrock"
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = self.client_type
+ super().__init__(*args, **kwargs)
+
+ def _get_job_by_name(self, job_name: str):
+ return self.conn.get_model_customization_job(jobIdentifier=job_name)
+
+ def get_customize_model_job_state(self, job_name: str) -> str:
+ state = self._get_job_by_name(job_name)["status"]
+ self.log.info("Job '%s' state: %s", job_name, state)
+ return state
+
+ def get_job_arn(self, job_name: str) -> str:
+ return self._get_job_by_name(job_name)["jobArn"]
Review Comment:
Sorry I am late to the party but are these methods really needed? They seem to jsut call the boto3 api with no additional work
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org