SageMaker Training (Deprecated)

1.7.1 · deprecated · verified Sat Apr 11

This library, `sagemaker-train`, is deprecated and no longer maintained. It previously provided functionalities for defining and running training jobs on Amazon SageMaker. Users are strongly advised to migrate to the main `sagemaker` Python SDK (also known as `sagemaker-python-sdk`) for all SageMaker interactions, including training, which offers comprehensive and actively developed features. The last stable version of `sagemaker-train` is 1.7.1.

Warnings

Install

Imports

Quickstart

This quickstart demonstrates how to initiate a training job using the recommended `sagemaker` Python SDK. It creates a simple training script and configures a PyTorch Estimator to run it on SageMaker. Ensure you replace the placeholder IAM role ARN with your actual SageMaker execution role and adjust instance types/framework versions as needed.

import sagemaker
from sagemaker.pytorch import PyTorch
import os

# --- THIS IS A REPLACEMENT EXAMPLE FOR THE DEPRECATED `sagemaker-train` LIBRARY ---
# It demonstrates how to train a model using the recommended `sagemaker` SDK.

# 1. Configure SageMaker session and IAM role
sagemaker_session = sagemaker.Session()
# IMPORTANT: Replace with your actual AWS IAM role ARN for SageMaker execution.
# This role grants SageMaker permissions to access resources like S3 and ECR.
# Example: "arn:aws:iam::123456789012:role/SageMakerExecutionRole"
role = os.environ.get('SAGEMAKER_ROLE_ARN', 'arn:aws:iam::123456789012:role/SageMakerExecutionRole')
if '123456789012' in role:
    print("WARNING: Placeholder IAM role ARN detected. Please replace 'SAGEMAKER_ROLE_ARN' "
          "with your actual SageMaker execution role ARN to run this code on AWS.")

# 2. Create a dummy local training script for demonstration purposes.
# In a real scenario, this would be your actual training script (e.g., train.py).
script_path = 'my_training_script.py'
with open(script_path, 'w') as f:
    f.write("""
import argparse, os, logging
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=1)
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO)
    logging.info(f"Starting dummy training for {args.epochs} epochs.")
    # SageMaker automatically sets SM_MODEL_DIR for model artifacts
    model_output_path = os.environ.get('SM_MODEL_DIR')
    if model_output_path:
        with open(os.path.join(model_output_path, 'model.txt'), 'w') as mf:
            mf.write('dummy_model_content')
    logging.info("Dummy training finished.")
""")

# 3. Define the Estimator for your training job.
# This example uses a PyTorch Estimator. Other framework estimators (TensorFlow, SKLearn, etc.)
# are also available in the sagemaker SDK.
pytorch_estimator = PyTorch(
    entry_point=script_path,
    role=role,
    instance_count=1,
    instance_type='ml.m5.large', # Choose an appropriate SageMaker instance type
    framework_version='1.13.1',  # Specify your desired PyTorch version
    py_version='py39',          # Specify your desired Python version
    hyperparameters={'epochs': 2},
    sagemaker_session=sagemaker_session
)

print(f"Estimator configured for script: {script_path}")
print("To start a training job on SageMaker, uncomment the line below:")
# pytorch_estimator.fit() 

# Optional: Clean up the dummy script
# os.remove(script_path)

view raw JSON →