Registry / aws / sagemaker-training

sagemaker-training

library5.1.1pypypi✓ verified 85d ago

The `sagemaker-training` library provides the core toolkit that runs inside Amazon SageMaker training containers. It handles downloading input data, parsing hyperparameters, executing user training scripts, and uploading model artifacts. It's currently at version 5.1.1 and has a relatively active release cadence, with minor versions released every few weeks/months and major versions less frequently.

aws ai-ml

pip install sagemaker-training

github ↗package ↗homepage ↗

INSTALL

IMPORT

SIG · SAGEMAKER-TRAINING

sagemaker-training

awspythonv5.1.1

Install

—

Import

—

Disk

—

Pass rate

0/ 10

Env Coverage0 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v? · pip install

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.920 runs

build_error

glibc

py 3.10–3.920 runs

build_error

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

environment

✓ from sagemaker_training import environment

✗ import sagemaker_training

Most common usage is to import the `environment` module to access training environment details.

get_environment

✓ from sagemaker_training.environment import get_environment

✗ from sagemaker_training import get_environment

The `get_environment` function is a member of the `environment` module.

get_hyperparameters

✓ from sagemaker_training.environment import get_hyperparameters

✗ sagemaker_training.get_hyperparameters()

Hyperparameters are accessed via the `environment` module.

This quickstart demonstrates a typical SageMaker training script entry point. It uses `sagemaker_training.environment` to retrieve hyperparameters and input/output paths, which are crucial for running user code within a SageMaker training container. The script should be placed at the root of your training code archive.

from sagemaker_training import environment import os def train(): # Get SageMaker training environment details env = environment.get_environment() # Access hyperparameters hyperparameters = env.hyperparameters learning_rate = hyperparameters.get('learning_rate', 0.01) # Access input data paths train_data_path = os.path.join(env.channel_input_dirs['training'], 'data.csv') # Access model output path model_dir = env.model_dir print(f"Learning Rate: {learning_rate}") print(f"Training data path: {train_data_path}") print(f"Model output directory: {model_dir}") # Your training logic here # Example: Save a dummy model artifact with open(os.path.join(model_dir, 'model.txt'), 'w') as f: f.write('My trained model output') if __name__ == '__main__': train()

Debug

Known issues

breakingVersion 5.0.0 updated the required `protobuf` dependency to version 5.28.1. This can cause significant conflicts and runtime errors (`AttributeError`, `TypeError`) if your custom training environment or other dependencies rely on an older (v3 or v4) `protobuf` version.

fix

Ensure all your dependencies, including those brought by SageMaker's base images, are compatible with `protobuf>=5.0.0`. You may need to pin specific versions of conflicting libraries or use a different base image if conflicts persist.

affects: >=5.0.0

gotchaThe `sagemaker-training` library is designed to run *inside* the SageMaker training container. Trying to run scripts locally that heavily rely on `sagemaker_training.environment` calls without mocking or setting up the corresponding environment variables will result in errors (e.g., `KeyError` for missing environment variables or incorrect paths).

fix

When testing locally, either mock the `sagemaker_training.environment` calls or ensure you set up dummy SageMaker environment variables (e.g., `SM_MODEL_DIR`, `SM_INPUT_DATA_CONFIG`, `SM_HYPERPARAMETERS`) for your test environment.

affects: All versions

gotchaMismatched `boto3` versions between the SageMaker Training Toolkit and your custom code or base image can lead to issues with S3 interactions (e.g., downloading data, uploading model artifacts) or credential handling.

fix

Regularly update `boto3` to the latest compatible version. If you encounter S3-related issues, verify that `boto3` versions are consistent across your environment and the toolkit, and consider explicitly pinning a compatible `boto3` version in your `requirements.txt`.

affects: All versions

gotchaThe training toolkit expects your user script to be at `/opt/ml/code/your_script.py` (or similar) within the container. Custom entrypoints or Dockerfiles that deviate from this structure without proper configuration can lead to `FileNotFoundError` or the script not being executed.

fix

Ensure your SageMaker estimator configuration or Dockerfile ENTRYPOINT/CMD correctly points to your training script. The `source_dir` argument in SageMaker SDK estimators handles this automatically for Python scripts.

affects: All versions

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'sagemaker_training'

The `sagemaker-training` library is not installed in the environment where the script is being executed (e.g., local machine or custom Docker container without proper installation).

fix

Run `pip install sagemaker-training` in your environment. If using a custom Dockerfile, add `RUN pip install sagemaker-training`.

AttributeError: module 'google.protobuf.descriptor' has no attribute '_HAS_OPTIONAL_FIELD_ACCESSORS'

This error typically occurs when `protobuf` v5+ is installed, but another library (or an older version of `sagemaker-training` itself) expects an API from `protobuf` v3 or v4. This was a common issue after `sagemaker-training` v5.0.0's `protobuf` upgrade.

fix

Verify all `protobuf` installations. Ensure all libraries in your environment are compatible with `protobuf>=5.0.0`. You may need to upgrade other dependencies or explicitly pin `protobuf` to a specific v5+ version (e.g., `protobuf>=5.28.1`).

KeyError: 'SM_HP_YOUR_HYPERPARAMETER'

Attempting to access a hyperparameter using `env.hyperparameters['YOUR_HYPERPARAMETER']` when that hyperparameter was not provided to the SageMaker training job.

fix

Always use `.get()` with a default value when accessing hyperparameters (e.g., `hyperparameters.get('your_hyperparameter', default_value)`). Double-check that the hyperparameter name passed to the SageMaker estimator matches the key used in your script.

FileNotFoundError: [Errno 2] No such file or directory: '/opt/ml/input/data/training/my_file.csv'

The script is trying to access an input file or directory that does not exist at the specified path within the SageMaker container, or the input data channel was not correctly configured.

fix

Ensure your SageMaker estimator's `inputs` argument correctly maps your S3 data to the expected channel name (e.g., 'training'). Verify the file exists in your S3 bucket and that your script uses the correct path derived from `env.channel_input_dirs['channel_name']`.

Upgrade

Version history

5.1.1latest on PyPI · released Sep 22, 2025

Audit

Dependencies

protobufrequiredRequired for internal data serialization; version updates (e.g., v5.0.0) can cause conflicts with other libraries.

boto3optionalUsed for interacting with AWS services like S3. While not a direct dependency of the core `sagemaker-training` package in all scenarios, version mismatches can cause runtime issues for functionality like S3 data transfer.

Agent activity

20 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/aws/sagemaker-training-toolkit ↗

packagehttps://pypi.org/project/sagemaker-training/ ↗

homepagehttps://aws.amazon.com/sagemaker/ ↗