4.1.3. Containers for ML and AI Workloads
💡 First Principle: SageMaker AI runs every training job and every endpoint inside a container, and Amazon Bedrock AgentCore Runtime runs agents inside containers too. The engineering choice is how much of that image you own. The less you own, the less you patch; the more you own, the more you can customise. The exam tests whether you pick the lowest-ownership option that still meets the requirement.
| Option | You own | Use when |
|---|---|---|
| Prebuilt AWS Deep Learning Container + Script Mode | Only your train.py / inference.py and a requirements.txt | Supported framework (PyTorch, TensorFlow, scikit-learn, XGBoost, Hugging Face) with pip-installable extras |
| Large model inference (LMI) container | Configuration (model ID, tensor parallel degree, backend such as vLLM) | Serving LLMs on SageMaker AI without writing a model server |
| Extended prebuilt image | A Dockerfile that starts FROM the AWS image and adds system libraries | You need OS-level or compiled dependencies but still want the AWS framework stack |
| Fully custom image | Everything, including the web server | Unsupported language or runtime |
Script Mode remains the default for custom model code: AWS maintains and patches the framework image and you supply scripts. If a question mentions "custom training logic with PyTorch", Script Mode with the PyTorch container is the answer.
Extending a prebuilt image is the middle path. When a dependency cannot be installed with pip (a compiled C++ library, a system codec), write a Dockerfile based on the AWS framework image, add the dependency, push the result to Amazon ECR and reference that image URI in the model definition. You keep the AWS-tested framework layer and own only the delta.
Serving contract for any custom inference image. SageMaker AI calls GET /ping for health checks and POST /invocations for predictions on port 8080; an endpoint whose container does not answer /ping in time never reaches InService. AgentCore Runtime uses the same /ping and /invocations pattern on port 8080 for HTTP agents, but its images must be built for ARM64.
Maintaining images. Tag images with immutable version identifiers rather than latest so a training job or endpoint can be reproduced, enable Amazon ECR image scanning (backed by Amazon Inspector for enhanced scanning) so vulnerabilities fail the pipeline, and rebuild extended images when AWS publishes a patched base image.
⚠️ Exam Trap: A fully custom image is the last resort, not the "most flexible, therefore best" answer. If the framework is supported and only a library is missing, extend the prebuilt image. If the library is pip-installable, a requirements.txt with Script Mode is enough and involves no Docker work at all.
Reflection Question: A team's PyTorch inference code needs a compiled audio-processing library that is not on PyPI. Which container option adds the library with the least ongoing maintenance, and what must the resulting image expose to SageMaker AI?