Model Registry Operations with XGBoost

Model Registry Operations with XGBoost#

Register XGBoost models to SageMaker Model Registry, create models from existing registry entries, and manage model approval workflows.

from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.serve.mode.function_pointers import Mode
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import image_uris
import boto3

sagemaker_session = Session()
role = get_execution_role()
region = sagemaker_session.boto_region_name

image_uri = image_uris.retrieve(
    framework="xgboost",
    region=region,
    version="1.0-1",
    py_version="py3",
    instance_type="ml.m5.xlarge",
)

s3 = boto3.client('s3')
bucket = sagemaker_session.default_bucket()

# this is pre-trained xgboost model
s3.upload_file('model/model.tar.gz', bucket, 'registry-testing/model.tar.gz')
s3_url = f"s3://{bucket}/registry-testing/model.tar.gz"

Build and register a model from model artifact#

model_builder = ModelBuilder(
    s3_model_data_url=s3_url,
    image_uri=image_uri,
    sagemaker_session=sagemaker_session,
    role_arn=role,
)

# Build the model
model_builder.build(model_name="my-model")

# Register the model
model_builder.register(
    model_package_group_name="my-model-package-group",
    content_types=["application/json"],
    response_types=["application/json"],
    inference_instances=["ml.m5.xlarge"],
    approval_status="Approved"
)

Register from existing SageMaker Model resource#

This assumes that you already have a SageMaker Model resource.

from sagemaker.core.resources import Model
from sagemaker.serve.model_builder import ModelBuilder

# your sagemaker model is called my-model
model = Model.get(model_name="my-model")

# Get image URI and S3 model data URI
image_uri = model.primary_container.image
s3_model_data_uri = model.primary_container.model_data_url

model_builder = ModelBuilder(
    image_uri=image_uri,
    s3_model_data_url=s3_model_data_uri,
    role_arn=role,
)

registered_model_package_arn = model_builder.register(
    model_package_group_name="my-model-group-from-existing-model",
    content_types=["application/json"],
    response_types=["application/json"]
)

Create a SageMaker model from specific registry#

Here we assume you already have a registered model version in model group, and we create a sagemaker model from this version.

# approve the version before creating model

# There is a gap that API response for a versioned model package doesn't include model_package_name
sagemaker_client = boto3.client('sagemaker', region_name=region)
sagemaker_client.update_model_package(
    ModelPackageArn=registered_model_package_arn,
    ModelApprovalStatus="Approved"
)

from sagemaker.core.resources import ModelPackage
from sagemaker.serve.model_builder import ModelBuilder

# Get the registered ModelPackage using the ARN as model_package_name
model_package = ModelPackage.get(model_package_name=registered_model_package_arn)

s3_model_data_uri = model_package.inference_specification.containers[0].model_data_url
image_uri = model_package.inference_specification.containers[0].image

model_builder = ModelBuilder(
    s3_model_data_url=s3_model_data_uri,
    image_uri=image_uri,
    role_arn=role,
    sagemaker_session=sagemaker_session
)

# Build the model
model = model_builder.build(model_name="my-model-from-registry")
Contents