Skip to content

Commit 028fde5

Browse files
bhanutejagkBhanu Teja Goshikonda
bhanutejagk
and
Bhanu Teja Goshikonda
authored
docs: add tensorflow framework documentation for 2.21 release (#6399)
* docs: add tensorflow framework documentation for 2.21 release Adds a docs/tensorflow/ tree mirroring the docs/pytorch/ pattern so the aws.github.io site reflects the TF 2.21 SageMaker training DLC released 2026-07-10. Includes: - Overview page with images table, "What's Included" list, CUDA forward-compatibility note, and build info - SageMaker AI deployment guide with SM SDK v2 and v3 examples covering TensorFlow estimator and TensorFlowProcessor / FrameworkProcessor - Changelog with the initial 2.21 entry (AL2023, py312, CUDA 12.9.1, cuDNN 9.5.1.17, NCCL 2.30.7, EFA 1.49.0, OpenMPI 4.1.8) TF 2.21 is a SageMaker-only release (no EC2 variant), so the deployment/ tree omits ec2.md. Nav entry added to docs/.nav.yml. * docs(tensorflow): drop pytorch cross-reference in what's included The what's-included preamble compared the tensorflow DLC to the pytorch DLC, which is misleading now that a tensorflow inference DLC is planned. Each framework's docs should stand on its own. * docs(tensorflow): drop inference-DLC redirect note The training doc pointed inference users to vllm/vllm-omni/ray, which will be misleading once the tensorflow inference DLC ships. * docs(tensorflow): apply flowmark formatter Collapses pipe-aligned tables to compact form and re-wraps prose at 150 columns to satisfy the flowmark pre-commit hook. * docs: add TensorFlow card to Getting Started grid Mirrors the PyTorch training card. Requested by @sirutBuasai on PR review. * docs: club PyTorch and TensorFlow into a single training card Merges the separate framework cards into one "Train Models" card covering PyTorch and TensorFlow, with links to both framework guides. --------- Co-authored-by: Bhanu Teja Goshikonda <bhanugk@amazon.com>
1 parent d818339 commit 028fde5

6 files changed

Lines changed: 212 additions & 2 deletions

File tree

‎docs/.nav.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ nav:
88
- TEI: tei
99
- Ray: ray
1010
- PyTorch: pytorch
11+
- TensorFlow: tensorflow
1112
- Base: base
1213
- Blog Posts: tutorials
1314
- Resources:

‎docs/index.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,9 +72,9 @@ LLM serving is just one example. DLCs cover a range of AI/ML workloads — explo
7272

7373
---
7474

75-
Run distributed training with PyTorch on GPU or CPU, with EFA, NCCL, flash-attn, and DeepSpeed pre-installed.
75+
Run distributed training with PyTorch or TensorFlow on GPU or CPU, with EFA and NCCL pre-installed for multi-node workloads.
7676

77-
[PyTorch Guide](pytorch/index.md)
77+
[PyTorch Guide](pytorch/index.md) · [TensorFlow Guide](tensorflow/index.md)
7878

7979
- **Build Your Own Image**
8080

‎docs/tensorflow/.nav.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
nav:
2+
- Overview: index.md
3+
- Deployment:
4+
- Amazon SageMaker AI: deployment/sagemaker.md
5+
- Changelog: changelog/index.md

‎docs/tensorflow/changelog/index.md‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Changelog
2+
3+
Changelog for the Amazon Linux 2023-based TensorFlow SageMaker training images (`2.21-gpu-py312-cu129-amzn2023-sagemaker` and
4+
`2.21-cpu-py312-amzn2023-sagemaker`).
5+
6+
* * *
7+
8+
## TensorFlow 2.21 — 2026-07-10
9+
10+
**Tags:** `2.21-gpu-py312-cu129-amzn2023-sagemaker` · `2.21-cpu-py312-amzn2023-sagemaker` · `2.21.0-gpu-py312-cu129-amzn2023-sagemaker` ·
11+
`2.21.0-cpu-py312-amzn2023-sagemaker`
12+
13+
**Bundled versions:** TensorFlow 2.21.0 · Python 3.12 · CUDA 12.9.1 · cuDNN 9.5.1.17 · NCCL 2.30.7 · EFA 1.49.0 · OpenMPI 4.1.8
14+
15+
### Highlights
16+
17+
- Initial release of the TensorFlow DLC on Amazon Linux 2023 and Python 3.12
18+
- TensorFlow 2.21.0 for SageMaker AI training (CPU and GPU variants)
19+
- CUDA 12.9.1 with cuDNN 9.5.1.17 and NCCL 2.30.7 on the GPU variant
20+
- EFA 1.49.0 with OpenMPI 4.1.8 for multi-node training on EFA-capable instances
21+
- SageMaker toolkits pre-installed: `sagemaker-tensorflow-training`, `sagemaker-training`
22+
- MLflow 3.9+, `smclarify`, `sagemaker>=3.4.0`, and SageMaker Studio integrations bundled

‎docs/tensorflow/deployment/sagemaker.md‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
# Amazon SageMaker AI Deployment
2+
3+
Use the TensorFlow DLC for training jobs launched via the SageMaker Python SDK or `boto3`. The images bundle the `sagemaker-tensorflow-training`
4+
toolkit, OpenMPI for multi-node coordination, and SageMaker-specific Python libraries (`mlflow`, `smclarify`, `pandas`, `seaborn`, `s3fs`, etc.).
5+
6+
The TensorFlow 2.21 DLC is a **SageMaker-only** release — there is no EC2 or EKS variant of this image. Point your estimators, processors, and
7+
training jobs at the SageMaker image URIs shown below.
8+
9+
## SageMaker Python SDK v2
10+
11+
Pass the DLC image URI via `image_uri=` rather than `framework_version=`:
12+
13+
```python
14+
from sagemaker.tensorflow import TensorFlow
15+
16+
estimator = TensorFlow(
17+
image_uri="public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-gpu-py312-cu129-amzn2023-sagemaker",
18+
role="arn:aws:iam::<account_id>:role/<role_name>",
19+
entry_point="train.py",
20+
source_dir="src",
21+
instance_type="ml.p5.48xlarge",
22+
instance_count=2,
23+
distribution={"mpi": {"enabled": True, "processes_per_host": 8}},
24+
)
25+
26+
estimator.fit({"training": "s3://<bucket>/datasets/train/"})
27+
```
28+
29+
The toolkit invokes your `entry_point` script over MPI when `mpi.enabled` is set and as a single-process launcher otherwise.
30+
31+
## SageMaker Python SDK v3
32+
33+
```python
34+
from sagemaker.core.resources import TrainingJob
35+
from sagemaker.core.shapes import (
36+
AlgorithmSpecification,
37+
Channel,
38+
DataSource,
39+
InputDataConfig,
40+
OutputDataConfig,
41+
ResourceConfig,
42+
S3DataSource,
43+
StoppingCondition,
44+
)
45+
46+
job = TrainingJob.create(
47+
training_job_name="tensorflow-train",
48+
role_arn="arn:aws:iam::<account_id>:role/<role_name>",
49+
algorithm_specification=AlgorithmSpecification(
50+
training_image="public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-gpu-py312-cu129-amzn2023-sagemaker",
51+
training_input_mode="File",
52+
),
53+
resource_config=ResourceConfig(
54+
instance_type="ml.p5.48xlarge",
55+
instance_count=2,
56+
volume_size_in_gb=200,
57+
),
58+
input_data_config=[
59+
InputDataConfig(
60+
channel_name="training",
61+
data_source=DataSource(
62+
s3_data_source=S3DataSource(
63+
s3_data_type="S3Prefix",
64+
s3_uri="s3://<bucket>/datasets/train/",
65+
),
66+
),
67+
),
68+
],
69+
output_data_config=OutputDataConfig(s3_output_path="s3://<bucket>/output/"),
70+
stopping_condition=StoppingCondition(max_runtime_in_seconds=3600),
71+
)
72+
job.wait_for_status("Completed")
73+
```
74+
75+
## Processing Jobs
76+
77+
### `TensorFlowProcessor` (SDK v2)
78+
79+
```python
80+
from sagemaker.tensorflow.processing import TensorFlowProcessor
81+
82+
processor = TensorFlowProcessor(
83+
image_uri="public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-cpu-py312-amzn2023-sagemaker",
84+
framework_version="2.21", # workaround: value ignored when image_uri is set
85+
role="arn:aws:iam::<account_id>:role/<role_name>",
86+
instance_type="ml.m5.xlarge",
87+
instance_count=1,
88+
)
89+
```
90+
91+
### `FrameworkProcessor` (SDK v3)
92+
93+
```python
94+
from sagemaker.workflow.steps import ProcessingStep
95+
from sagemaker.processing import FrameworkProcessor
96+
97+
processor = FrameworkProcessor(
98+
image_uri="public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-cpu-py312-amzn2023-sagemaker",
99+
role="arn:aws:iam::<account_id>:role/<role_name>",
100+
instance_type="ml.m5.xlarge",
101+
instance_count=1,
102+
)
103+
```
104+
105+
## Multi-Node Training over EFA
106+
107+
SageMaker provisions EFA on the instance types that support it (e.g., `ml.p5.48xlarge`, `ml.p4d.24xlarge`). The GPU image ships with EFA 1.49.0 and
108+
OpenMPI 4.1.8 pre-installed — no extra plumbing in your training script.
109+
110+
Multi-node peer discovery is handled by the `sagemaker_tensorflow_container` training toolkit, which parses `/opt/ml/input/config/resourceconfig.json`
111+
at container start to enumerate hosts, current host rank, and network interface. Peers are addressed by the SageMaker-assigned hostnames (e.g.,
112+
`algo-1`, `algo-2`) that resolve inside the training cluster — your script does not need to know about this.
113+
114+
## Container Layout
115+
116+
| Path | Purpose |
117+
| --- | --- |
118+
| `/opt/ml/input/data/<channel_name>/` | Training data SageMaker mounts from S3 (one subdir per channel) |
119+
| `/opt/ml/model/` | Write your final model artifacts here — SageMaker uploads to `OutputDataConfig.s3_output_path` |
120+
| `/opt/ml/output/` | Auxiliary outputs (logs, checkpoints) |
121+
| `/opt/ml/code/` | Your training source dir (populated from the SDK's `source_dir`) |
122+
| `/opt/venv/` | Python venv with TensorFlow + DLC libraries |
123+
124+
## Notes
125+
126+
- The image sets `SAGEMAKER_TRAINING_MODULE=sagemaker_tensorflow_container.training:main` — this is the entry-point the SageMaker training toolkit
127+
invokes at container start, which in turn launches your `entry_point` script.
128+
- For a baseline driver/AMI compatible with these CUDA 12.9 images, request the latest SageMaker training AMI when launching jobs.

‎docs/tensorflow/index.md‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# ML Training using TensorFlow DLC
2+
3+
Production-ready Docker images for TensorFlow training workloads on {{ aws }}. Available in CPU and GPU variants, built on Amazon Linux 2023 with
4+
ongoing security patching.
5+
6+
These images bundle TensorFlow with the libraries needed for **distributed training on {{ sagemaker }}** — EFA for low-latency networking on
7+
EFA-capable instances, OpenMPI for multi-node coordination, and the SageMaker training toolkits for entry-point invocation and channel wiring.
8+
9+
## Images
10+
11+
| Platform | Variant | Image |
12+
| --- | --- | --- |
13+
| {{ sagemaker }} | GPU | `public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-gpu-py312-cu129-amzn2023-sagemaker` |
14+
| {{ sagemaker }} | CPU | `public.ecr.aws/deep-learning-containers/tensorflow-training:2.21-cpu-py312-amzn2023-sagemaker` |
15+
16+
All images are also available on the [ECR Public Gallery](https://gallery.ecr.aws/deep-learning-containers/tensorflow-training). For private ECR URIs,
17+
see [Image Access](../get_started/index.md).
18+
19+
The `2.21.0-*` fully-qualified tags (e.g. `2.21.0-gpu-py312-cu129-amzn2023-sagemaker`) are also published and pin to the same image digest.
20+
21+
## What's Included
22+
23+
The GPU image includes:
24+
25+
- **TensorFlow 2.21.0** (`tensorflow==2.21.0`)
26+
- **CUDA 12.9.1** with **cuDNN 9.5.1.17** (`nvidia-cudnn-cu12`) and **NCCL 2.30.7** (`nvidia-nccl-cu12`) for multi-GPU collectives
27+
- **[EFA](https://aws.amazon.com/hpc/efa/) 1.49.0** with **OpenMPI 4.1.8** for low-latency multi-node communication on EFA-capable instances
28+
- **[MPI for Python](https://mpi4py.readthedocs.io/) (`mpi4py`)** for multi-process Python coordination
29+
- **[SageMaker training toolkits](https://github.com/aws/sagemaker-tensorflow-training-toolkit)** — `sagemaker-tensorflow-training`,
30+
`sagemaker-training`
31+
- **[MLflow](https://mlflow.org/) 3.9+**, **`sagemaker>=3.4.0`**, **`smclarify`**, and **`sagemaker-experiments` 0.1.45**
32+
- **Data & ML tooling** — `tensorflow-io` 0.37, `tensorflow-datasets`, `pandas`, `scikit-learn`, `scipy`, `numpy` (2.1+), `Pillow`, `h5py`,
33+
`opencv-python`, `numba`, `plotly`, `seaborn`, `shap`, `bokeh`, `imageio`, `cloudpickle`
34+
- **AWS tooling** — `boto3`, `botocore`, `awscli` (<2), `s3fs`
35+
- **SageMaker Studio integration** — `sagemaker-studio-analytics-extension`, `sparkmagic` 0.22.0, `sagemaker-studio-sparkmagic-lib`
36+
- **Python 3.12** in a venv at `/opt/venv` (`PATH` already set)
37+
38+
The CPU variant includes the same TensorFlow ecosystem and SageMaker toolkits. EFA, CUDA, cuDNN, and NCCL are not present in the CPU image.
39+
40+
## CUDA Forward Compatibility
41+
42+
The GPU image entrypoint detects host NVIDIA driver versions older than the bundled `cuda-compat` layer and automatically prepends
43+
`/usr/local/cuda/compat` to `LD_LIBRARY_PATH`. No flag or env var needed — the check runs on every container start.
44+
45+
## How We Build
46+
47+
These images are curated builds tracking the [TensorFlow](https://www.tensorflow.org/) project:
48+
49+
- **Built from upstream TensorFlow wheels** published to PyPI
50+
- **Reproducible** — pinned via `pyproject.toml` + `uv.lock` for every image variant
51+
- **Security-patched** — continuously maintained with security patches from {{ aws }} on an Amazon Linux 2023 base
52+
53+
TensorFlow 2.21 is the first TensorFlow DLC on Amazon Linux 2023 and Python 3.12. Prior TensorFlow DLCs shipped on Ubuntu with earlier Python
54+
versions; this release moves the base OS and interpreter forward alongside the framework bump.

0 commit comments

Comments
 (0)