TensorFlow
AWS
Cloud Computing
Machine Learning
Clustering

How to run TensorFlow on an AWS cluster?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Running TensorFlow on an AWS cluster usually means distributing training across multiple EC2 instances, each acting as a worker in the same job. The core ingredients are compute instances, shared access to data, network connectivity between workers, and a TF_CONFIG definition so TensorFlow knows the cluster layout. Once those pieces are in place, tf.distribute.MultiWorkerMirroredStrategy is the usual starting point for synchronous training.

Provision The Cluster Nodes

At a minimum, each node needs:

  • the same Python and TensorFlow environment
  • network access to the other workers
  • access to the training data, often through S3
  • enough CPU or GPU resources for the workload

A practical EC2 setup usually starts with one AMI or Docker image that you reuse on every worker so the software stack stays identical.

If your data lives in S3, make sure the instances have an IAM role or credentials that allow reads from the training bucket.

Use TF_CONFIG To Describe The Cluster

TensorFlow distributed training relies on the TF_CONFIG environment variable to tell each process what role it has.

For a two-worker example, worker 0 might use:

bash
1export TF_CONFIG='{
2  "cluster": {
3    "worker": ["10.0.1.10:12345", "10.0.1.11:12345"]
4  },
5  "task": {"type": "worker", "index": 0}
6}'

Worker 1 uses the same cluster block but a different task index:

bash
1export TF_CONFIG='{
2  "cluster": {
3    "worker": ["10.0.1.10:12345", "10.0.1.11:12345"]
4  },
5  "task": {"type": "worker", "index": 1}
6}'

Security groups must allow the chosen port between the worker machines.

A Minimal Multi-Worker Training Script

Once the environment is ready, TensorFlow code can use MultiWorkerMirroredStrategy.

python
1import tensorflow as tf
2
3strategy = tf.distribute.MultiWorkerMirroredStrategy()
4
5with strategy.scope():
6    model = tf.keras.Sequential([
7        tf.keras.layers.Dense(64, activation="relu"),
8        tf.keras.layers.Dense(10, activation="softmax"),
9    ])
10    model.compile(
11        optimizer="adam",
12        loss="sparse_categorical_crossentropy",
13        metrics=["accuracy"],
14    )
15
16dataset = tf.data.Dataset.from_tensor_slices((
17    tf.random.normal((1000, 20)),
18    tf.random.uniform((1000,), maxval=10, dtype=tf.int32),
19)).batch(32)
20
21model.fit(dataset, epochs=3)

Each worker runs the same script. TensorFlow coordinates the replicas based on TF_CONFIG.

Put Data Somewhere All Workers Can Reach

Small experiments can load local files, but a real cluster usually needs shared storage. S3 is a common choice because every instance can read from the same bucket.

python
1import tensorflow as tf
2
3files = tf.io.gfile.glob("s3://my-training-bucket/data/*.tfrecord")
4dataset = tf.data.TFRecordDataset(files)

In practice, many teams stage TFRecord files or other preprocessed data in S3 and let each worker read its shard through the input pipeline.

Consider Containers And Managed Alternatives

If you want easier environment consistency, package the training job in Docker and run the same container on every EC2 instance. That reduces drift between nodes and makes redeployment easier.

If the real goal is simply distributed TensorFlow on AWS, also ask whether SageMaker or EKS would remove operational work you do not want to manage yourself. A self-managed EC2 cluster gives control, but it also means you own setup, scaling, health checks, and cleanup.

Common Pitfalls

The most common mistake is misconfiguring TF_CONFIG. If the cluster addresses or task indices differ between workers, the job will hang or fail to start correctly.

Another issue is forgetting network rules. Workers must be able to reach each other on the TensorFlow communication port.

Environment mismatch is also common. If one instance has a different TensorFlow version or CUDA stack, debugging becomes painful very quickly.

Finally, do not underestimate data throughput. Distributed training often shifts the bottleneck from compute to data loading, so S3 access patterns and input pipeline performance matter.

Summary

  • A TensorFlow AWS cluster usually means multiple EC2 workers running the same training script.
  • Use identical environments on every node and give them shared access to data.
  • Define cluster roles with TF_CONFIG and use MultiWorkerMirroredStrategy for synchronous training.
  • Open the required network ports between workers.
  • Consider Docker, EKS, or SageMaker if you want easier environment and cluster management.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.