Implementación de un modelo personalizado (SDK de AWS) - Amazon Bedrock

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Implementación de un modelo personalizado (SDK de AWS)

Para crear un modelo personalizado a partir de un modelo de Amazon Nova entrenado por SageMaker AI y almacenado en Amazon S3, se utiliza la operación de la API CreateCustomModel. Puede usar el código siguiente para crear un modelo personalizado con el SDK para Python (Boto3). El código crea un modelo personalizado y, a continuación, comprueba su estado hasta que el modelo esté ACTIVE y listo para usarse.

Para usar el código, actualice los siguientes parámetros. El ejemplo de código también incluye parámetros opcionales, como clientRequestToken para la idempotencia y modelTags para el etiquetado de recursos.

  • modelName: asigne al modelo un nombre único.

  • s3Uri: especifique la ruta al bucket de Amazon S3 administrado por Amazon que almacena los artefactos del modelo. SageMaker AI crea este bucket cuando ejecuta su primer trabajo de entrenamiento de SageMaker AI.

  • roleArn: especifique el Nombre de recurso de Amazon (ARN) de un rol de servicio de IAM que Amazon Bedrock asuma para realizar tareas en su nombre. Para obtener más información sobre la creación de este rol, consulte Creación de un rol de servicio para la importación de modelos preentrenados.

  • modelKmsKeyArn (opcional): especifique una clave de AWS KMS para cifrar el modelo en Amazon Bedrock. Si no proporciona una clave de AWS KMS, Amazon Bedrock utiliza una clave de AWS KMS administrada de AWS para cifrar el modelo. Para obtener información acerca del cifrado, consulte Cifrado de modelos personalizados importados.

Tras crear el modelo personalizado, el modelo aparece en la respuesta ListCustomModels con un customizationType de imported. Para realizar un seguimiento del estado del nuevo modelo, se utiliza la operación de la API GetCustomModel.

import boto3
import uuid
from botocore.exceptions import ClientError
import time

def create_custom_model(bedrock_client):
    """
    Creates a custom model in Amazon Bedrock from a SageMaker AI-trained Amazon Nova model stored in Amazon S3.
    Args:
        bedrock_client: The Amazon Bedrock client instance
    Returns:
        dict: Response from the CreateCustomModel API call
    """
    try:
        # Create a unique client request token for idempotency
        client_request_token = str(uuid.uuid4())

        # Define the model source configuration
        model_source_config = {
            's3DataSource': {
                's3Uri': 's3://amzn-s3-demo-bucket/folder/',
            }
        }

        # Create the custom model
        response = bedrock_client.create_custom_model(
            # Required parameters
            modelName='modelName',
            roleArn='serviceRoleArn',
            modelSourceConfig=model_source_config,

            # Optional parameters
            clientRequestToken=client_request_token,
            modelKmsKeyArn='keyArn',
            modelTags=[
                {
                    'key': 'Environment',
                    'value': 'Production'
                },
                {
                    'key': 'Project',
                    'value': 'AIInference'
                }
            ]
        )

        print(f"Custom model creation initiated. Model ARN: {response['modelArn']}")

        return response

    except ClientError as e:
        print(f"Error creating custom model: {e}")
        raise

def list_custom_models(bedrock_client):
    """
    Lists all custom models in Amazon Bedrock.

    Args:
        bedrock_client: An Amazon Bedrock client.

    Returns:
        dict: Response from the ListCustomModels API call

    """

    try:
        response = bedrock_client.list_custom_models()
        print(f"Total number of custom models: {len(response['modelSummaries'])}")

        for model in response['modelSummaries']:
            print("ARN: " + model['modelArn'])
            print("Name: " + model['modelName'])
            print("Status: " + model['modelStatus'])
            print("Customization type: " + model['customizationType'])
            print("------------------------------------------------------")

        return response

    except ClientError as e:
        print(f"Error listing custom models: {e}")
        raise

def check_model_status(bedrock_client, model_arn):
    """
    Checks the status of a custom model creation.

    Args:
        model_arn (str): The ARN of the custom model
        bedrock_client: An Amazon Bedrock client.

    Returns:
        dict: Response from the GetCustomModel API call

    """

    try:
        max_time = time.time() + 60 * 60  # 1 hour

        while time.time() < max_time:
            response = bedrock_client.get_custom_model(modelIdentifier=model_arn)
            status = response.get('modelStatus')
            print(f"Job status: {status}")
            if status == 'Failed':
                print(f"Failure reason: {response.get('failureMessage')}")
                break
            if status == 'Active':
                print("Model is ready for use.")
                break
            time.sleep(60)

    except ClientError as e:
        print(f"Error checking model status: {e}")
        raise


def main():
    bedrock_client = boto3.client(service_name='bedrock', region_name='REGION')
    
    # Create the custom model
    model_arn = create_custom_model(bedrock_client)["modelArn"]

    # Check the status of the model
    if model_arn:
        check_model_status(bedrock_client, model_arn)

    # View all custom models
    list_custom_models(bedrock_client)


if __name__ == "__main__":
    main()