Introduction

In SageMaker training jobs, dependencies must be available inside the remote training container, not just on your local machine. Teams often hit import errors because a package exists in a notebook kernel but is missing when the job starts on managed infrastructure. The reliable approach is to package dependencies with your training code, choose the installation method based on complexity, and keep versions pinned for reproducibility. SageMaker supports several patterns: requirements.txt in script mode, custom Docker images, and packaging local modules through source_dir. This guide explains when to use each option and how to avoid brittle training setups.

Core Sections

Use requirements.txt for standard Python dependencies

For built-in framework estimators (like TensorFlow or PyTorch), script mode with requirements.txt is the fastest path.

Project structure:

train/
  train.py
  requirements.txt

requirements.txt example:

pandas==2.2.3
scikit-learn==1.5.2
xgboost==2.1.2

Estimator setup:

1from sagemaker.pytorch import PyTorch
2
3estimator = PyTorch(
4    entry_point="train.py",
5    source_dir="train",
6    framework_version="2.1.0",
7    py_version="py310",
8    role=role,
9    instance_count=1,
10    instance_type="ml.m5.xlarge",
11)
12
13estimator.fit({"train": "s3://my-bucket/data/train"})

SageMaker installs packages at job startup.

Package your own modules with source_dir

If you have internal helper code, include it in source_dir and import it normally.

1train/
2  train.py
3  requirements.txt
4  mylib/
5    __init__.py
6    features.py

In train.py:

from mylib.features import build_features

This avoids publishing private packages externally just to run training.

Use a custom image for system level dependencies

If you need OS packages, CUDA specific tools, or strict runtime control, build and push a custom ECR image.

1FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04
2RUN apt-get update && apt-get install -y libgomp1 && rm -rf /var/lib/apt/lists/*
3COPY requirements.txt /opt/ml/code/requirements.txt
4RUN pip install --no-cache-dir -r /opt/ml/code/requirements.txt

Then run:

1from sagemaker.estimator import Estimator
2
3estimator = Estimator(
4    image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-train:latest",
5    role=role,
6    instance_count=1,
7    instance_type="ml.g4dn.xlarge",
8    entry_point="train.py",
9    source_dir="train"
10)

Custom images reduce startup surprises and make runs more reproducible.

Operational practices that prevent failures

Pin versions for all critical packages. Keep startup logs visible in CloudWatch and fail fast on missing imports at the top of train.py. For large dependency sets, prefer prebuilt images so jobs do not spend minutes installing packages every run. If your organization has private package indexes, configure credentials securely through environment variables or AWS Secrets Manager rather than hardcoding tokens.

Common Pitfalls

• Installing packages in the notebook kernel and assuming the same environment exists inside the remote training container.
• Leaving dependency versions unpinned, causing silent behavior changes between training runs.
• Forgetting to include local modules in source_dir, which leads to ModuleNotFoundError at runtime.
• Relying on runtime pip install for heavy dependencies, increasing startup time and failure risk.
• Mixing incompatible framework, Python, and CUDA versions when building custom images.

Production Readiness Check

Before closing the task, run a short validation loop on representative inputs and one intentional failure case. Confirm that your code path behaves correctly for normal data, empty data, and malformed data. Capture at least one measurable signal such as runtime, memory use, or error rate, then compare it to your baseline so regressions are visible. Keep this check lightweight so it can run in local development and CI without slowing feedback too much. A simple checklist plus one executable smoke test prevents most regressions after refactors and library upgrades.

11. Run happy-path example
22. Run edge-case example
33. Run failure-path example
44. Capture one performance or reliability metric
55. Verify output format and error handling

Summary

Installing modules in SageMaker training jobs is mostly an environment packaging problem. Use requirements.txt for typical Python dependencies, source_dir for local code, and custom images when you need OS-level control or strict reproducibility. Pin versions, inspect CloudWatch logs, and test the container entry point locally when possible. Once dependency management is treated as part of training infrastructure, SageMaker jobs become predictable and much easier to debug.