# 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
import json
from datetime import datetime
from os import environ
import boto3
from airflow.decorators import task, task_group
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
from airflow.operators.empty import EmptyOperator
from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
from airflow.providers.amazon.aws.operators.bedrock import (
BedrockCreateProvisionedModelThroughputOperator,
BedrockCustomizeModelOperator,
BedrockInvokeModelOperator,
)
from airflow.providers.amazon.aws.operators.s3 import (
S3CreateBucketOperator,
S3CreateObjectOperator,
S3DeleteBucketOperator,
)
from airflow.providers.amazon.aws.sensors.bedrock import (
BedrockCustomizeModelCompletedSensor,
BedrockProvisionModelThroughputCompletedSensor,
)
from airflow.utils.edgemodifier import Label
from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
# Externally fetched variables:
[docs]ROLE_ARN_KEY = "ROLE_ARN"
[docs]sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
[docs]DAG_ID = "example_bedrock"
# Creating a custom model takes nearly two hours. If SKIP_LONG_TASKS
# is True then these tasks will be skipped. This way we can still have
# the code snippets for docs, and we can manually run the full tests.
[docs]SKIP_LONG_TASKS = environ.get("SKIP_LONG_SYSTEM_TEST_TASKS", default=True)
[docs]LLAMA_SHORT_MODEL_ID = "meta.llama2-13b-chat-v1"
[docs]TITAN_MODEL_ID = "amazon.titan-text-express-v1:0:8k"
[docs]TITAN_SHORT_MODEL_ID = TITAN_MODEL_ID.split(":")[0]
[docs]PROMPT = "What color is an orange?"
[docs]TRAIN_DATA = {"prompt": "what is AWS", "completion": "it's Amazon Web Services"}
[docs]HYPERPARAMETERS = {
"epochCount": "1",
"batchSize": "1",
"learningRate": ".0005",
"learningRateWarmupSteps": "0",
}
@task_group
[docs]def customize_model_workflow():
# [START howto_operator_customize_model]
customize_model = BedrockCustomizeModelOperator(
task_id="customize_model",
job_name=custom_model_job_name,
custom_model_name=custom_model_name,
role_arn=test_context[ROLE_ARN_KEY],
base_model_id=f"{model_arn_prefix}{TITAN_SHORT_MODEL_ID}",
hyperparameters=HYPERPARAMETERS,
training_data_uri=training_data_uri,
output_data_uri=f"s3://{bucket_name}/myOutputData",
)
# [END howto_operator_customize_model]
# [START howto_sensor_customize_model]
await_custom_model_job = BedrockCustomizeModelCompletedSensor(
task_id="await_custom_model_job",
job_name=custom_model_job_name,
)
# [END howto_sensor_customize_model]
@task
def delete_custom_model():
BedrockHook().conn.delete_custom_model(modelIdentifier=custom_model_name)
@task.branch
def run_or_skip():
return end_workflow.task_id if SKIP_LONG_TASKS else customize_model.task_id
run_or_skip = run_or_skip()
end_workflow = EmptyOperator(task_id="end_workflow", trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
chain(run_or_skip, Label("Long-running tasks skipped"), end_workflow)
chain(run_or_skip, customize_model, await_custom_model_job, delete_custom_model(), end_workflow)
@task
[docs]def delete_provision_throughput(provisioned_model_id: str):
BedrockHook().conn.delete_provisioned_model_throughput(provisionedModelId=provisioned_model_id)
with DAG(
dag_id=DAG_ID,
schedule="@once",
start_date=datetime(2021, 1, 1),
tags=["example"],
catchup=False,
) as dag:
[docs] 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"
training_data_uri = f"s3://{bucket_name}/{input_data_s3_key}"
custom_model_name = f"CustomModel{env_id}"
custom_model_job_name = f"CustomizeModelJob{env_id}"
provisioned_model_name = f"ProvisionedModel{env_id}"
model_arn_prefix = f"arn:aws:bedrock:{boto3.session.Session().region_name}::foundation-model/"
create_bucket = S3CreateBucketOperator(
task_id="create_bucket",
bucket_name=bucket_name,
)
upload_training_data = S3CreateObjectOperator(
task_id="upload_data",
s3_bucket=bucket_name,
s3_key=input_data_s3_key,
data=json.dumps(TRAIN_DATA),
)
# [START howto_operator_invoke_llama_model]
invoke_llama_model = BedrockInvokeModelOperator(
task_id="invoke_llama",
model_id=LLAMA_SHORT_MODEL_ID,
input_data={"prompt": PROMPT},
)
# [END howto_operator_invoke_llama_model]
# [START howto_operator_invoke_titan_model]
invoke_titan_model = BedrockInvokeModelOperator(
task_id="invoke_titan",
model_id=TITAN_SHORT_MODEL_ID,
input_data={"inputText": PROMPT},
)
# [END howto_operator_invoke_titan_model]
# [START howto_operator_provision_throughput]
provision_throughput = BedrockCreateProvisionedModelThroughputOperator(
task_id="provision_throughput",
model_units=1,
provisioned_model_name=provisioned_model_name,
model_id=f"{model_arn_prefix}{TITAN_MODEL_ID}",
)
# [END howto_operator_provision_throughput]
provision_throughput.wait_for_completion = False
# [START howto_sensor_provision_throughput]
await_provision_throughput = BedrockProvisionModelThroughputCompletedSensor(
task_id="await_provision_throughput",
model_id=provision_throughput.output,
)
# [END howto_sensor_provision_throughput]
delete_bucket = S3DeleteBucketOperator(
task_id="delete_bucket",
trigger_rule=TriggerRule.ALL_DONE,
bucket_name=bucket_name,
force_delete=True,
)
chain(
# TEST SETUP
test_context,
create_bucket,
upload_training_data,
# TEST BODY
[invoke_llama_model, invoke_titan_model],
customize_model_workflow(),
provision_throughput,
await_provision_throughput,
# TEST TEARDOWN
delete_provision_throughput(provision_throughput.output),
delete_bucket,
)
from tests.system.utils.watcher import watcher
# This test needs watcher in order to properly mark success/failure
# when "tearDown" task with trigger rule is part of the DAG
list(dag.tasks) >> watcher()
from tests.system.utils import get_test_run # noqa: E402
# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
[docs]test_run = get_test_run(dag)