SageMaker V3 JumpStart E2E Training Example

SageMaker V3 JumpStart E2E Training Example#

This notebook demonstrates how to use SageMaker V3 to train a JumpStart model from scratch and deploy it for inference.

Prerequisites#

Note: Ensure you have sagemaker and ipywidgets installed in your environment. The ipywidgets package is required to monitor endpoint deployment progress in Jupyter notebooks.

# Import required libraries
import json
import uuid

from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.resources import EndpointConfig

Step 1: Configure JumpStart Model for Training#

We’ll train a HuggingFace Falcon model using JumpStart.

# Configuration
MODEL_ID = "huggingface-spc-bert-base-cased"
MODEL_NAME_PREFIX = "js-e2e-example-model"
ENDPOINT_NAME_PREFIX = "js-e2e-example-endpoint"

# Generate unique identifiers
unique_id = str(uuid.uuid4())[:8]
training_job_name = f"js-training-{unique_id}"
model_name = f"{MODEL_NAME_PREFIX}-{unique_id}"
endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}"

print(f"Training job name: {training_job_name}")
print(f"Model name: {model_name}")
print(f"Endpoint name: {endpoint_name}")

Step 2: Train the Model#

Use ModelTrainer to train the JumpStart model. The training job may take 30+ minutes to complete.

# Create JumpStart configuration
jumpstart_config = JumpStartConfig(model_id=MODEL_ID)

# Initialize ModelTrainer from JumpStart config
model_trainer = ModelTrainer.from_jumpstart_config(
    jumpstart_config=jumpstart_config, 
    base_job_name=training_job_name, 
    hyperparameters={"epochs": 1}
)

# Train the model
print("Starting model training...")
model_trainer.train()
print(f"Training completed: {training_job_name}")

Step 3: Build ModelBuilder from Trained Model#

Create a ModelBuilder from the training artifacts.

# Build ModelBuilder from trained model
model_builder = ModelBuilder(
    model=model_trainer,
    dependencies={"auto": False}
)

Step 4: Build the Model#

Build the model artifacts and prepare for deployment.

# Build the model
core_model = model_builder.build(model_name=model_name)
print(f"Model Successfully Created: {core_model.model_name}")

Step 5: Deploy the Trained Model#

Deploy the trained model to a SageMaker endpoint for real-time inference.

# Deploy the trained model to an endpoint
core_endpoint = model_builder.deploy(endpoint_name=endpoint_name)
print(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}")

Step 6: Test the Endpoint#

This endpoint performs text entailment classification, determining the logical relationship between pairs of sentences. The returned scores indicate the model’s confidence for different entailment categories (e.g., entailment, contradiction, neutral) - higher scores indicate stronger predictions for each relationship type.

This sends a test request to the deployed endpoint.

# Test with a sample query for the trained model
test_data = ["The weather is sunny today", "It is not raining"]

result = core_endpoint.invoke(
    body=json.dumps(test_data),
    content_type="application/list-text"
)

entailment_scores = json.loads(result.body.read().decode('utf-8'))
print(f"Result of invoking trained endpoint: {entailment_scores}")

Step 7: Clean Up Resources#

Clean up the created resources to avoid ongoing charges.

# Clean up resources
core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name)

# Delete in the correct order
core_model.delete()
core_endpoint.delete()
core_endpoint_config.delete()

print("Trained Model and Endpoint Successfully Deleted!")

Summary#

This notebook demonstrated:

  1. Training a JumpStart model using ModelTrainer

  2. Building a ModelBuilder from training artifacts

  3. Building the model for deployment

  4. Deploying to a SageMaker endpoint

  5. Making inference requests

  6. Cleaning up resources

The V3 ModelTrainer and ModelBuilder provide a seamless end-to-end workflow from training to deployment!

Contents