sagemaker-train

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)