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-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
