Amazon SageMakerの推論エンドポイントでA/Bテストのためのバリアント追加・削除を試してみる

# 一つ目のモデル
from sagemaker.inputs import TrainingInput

estimator = sagemaker.estimator.Estimator(
    sagemaker_session=sagemaker_session,
    base_job_name="xgboost-iris",
    image_uri=image, 
    role=role,
    instance_count=1, 
    instance_type='ml.m4.xlarge', 
    volume_size=30,
    output_path=s3_output_path,
    max_run=600)

estimator.set_hyperparameters(
    objective="multi:softmax",
    num_class=3,
    num_round=10,
    eval_metric="merror")

train_input = TrainingInput(s3_input_path, content_type="csv")

estimator.fit({'train': train_input})

# 二つ目のモデル（処理内容自体は一つ目と全く同じ）
from sagemaker.inputs import TrainingInput

estimator_2 = sagemaker.estimator.Estimator(
    sagemaker_session=sagemaker_session,
    base_job_name="xgboost-iris2",
    image_uri=image, 
    role=role,
    instance_count=1, 
    instance_type='ml.m4.xlarge', 
    volume_size=30,
    output_path=s3_output_path,
    max_run=600)

# num_roundだけ変える
estimator_2.set_hyperparameters(
    objective="multi:softmax",
    num_class=3,
    num_round=15,
    eval_metric="merror")

train_input = TrainingInput(s3_input_path, content_type="csv")

estimator_2.fit({'train': train_input})

from sagemaker.model import Model

model = Model(
    image_uri=image,
    model_data=estimator.model_data,
    name='iris-model-a',
    role=role)
model.create(
    instance_type="ml.m5.xlarge"
)

model2 = Model(
    image_uri=image,
    model_data=estimator_2.model_data,
    name='iris-model-b',
    role=role)
model2.create(
    instance_type="ml.m5.xlarge"
)

from sagemaker.session import production_variant

variant1 = production_variant(
    model_name=model.name,
    instance_type="ml.c5.4xlarge",
    initial_instance_count=1,
    variant_name="Variant1",
    initial_weight=1,
)
variant2 = production_variant(
    model_name=model2.name,
    instance_type="ml.c5.4xlarge",
    initial_instance_count=1,
    variant_name="Variant2",
    initial_weight=1,
)

endpoint_name = "DEMO-xgb-iris-pred"
print(f"EndpointName={endpoint_name}")

sagemaker_session.endpoint_from_production_variants(
    name=endpoint_name, production_variants=[variant1, variant2]
)

from sagemaker.predictor import Predictor
from sagemaker.serializers import CSVSerializer

predictor = Predictor(endpoint_name=endpoint_name, serializer=CSVSerializer())
response = predictor.predict(data=test_data_no_target.values, target_variant="Variant1")
print("Endpoint Response:", response.decode("utf-8"))

response = predictor.predict(data=test_data_no_target.values, target_variant="Variant2")
print("Endpoint Response:", response.decode("utf-8"))

sm_runtime = boto3.Session().client("sagemaker-runtime")
csv_text = '6.1 2.8 4.7 1.2'
response = sm_runtime.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType='text/csv',
    Body=csv_text
)
response['InvokedProductionVariant']
# 'Variant2'

import time
sm = boto_session.client("sagemaker")

sm.update_endpoint_weights_and_capacities(
    EndpointName=endpoint_name,
    DesiredWeightsAndCapacities=[
        {"DesiredWeight": 0, "VariantName": variant1["VariantName"]},
        {"DesiredWeight": 1, "VariantName": variant2["VariantName"]},
    ],
)
print("Waiting for update to complete")
while True:
    status = sm.describe_endpoint(EndpointName=endpoint_name)["EndpointStatus"]
    if status in ["InService", "Failed"]:
        print("Done")
        break
    print(".", end="", flush=True)
    time.sleep(1)

{
    variant["VariantName"]: variant["CurrentWeight"]
    for variant in sm.describe_endpoint(EndpointName=endpoint_name)["ProductionVariants"]
}

# Waiting for update to complete
# ....................................................................Done
# {'Variant1': 0.0, 'Variant2': 1.0}

sm.create_endpoint_config(
    EndpointConfigName = 'DEMO-xgb-iris-pred-after-ab-testing',
    ProductionVariants = [
        {
            "VariantName": "Variant2",
            "InstanceType": "ml.c5.4xlarge",
            "InitialInstanceCount": 1,
            "InitialVariantWeight": 1,
            "ModelName": model2.name
        }
    ]
)

sm.update_endpoint(
    EndpointName=endpoint_name,
    EndpointConfigName='DEMO-xgb-iris-pred-after-ab-testing'
)

response = sm.delete_endpoint(
    EndpointName=endpoint_name
)