Open-Source Wikis

/

TensorFlow

/

Systems

/

Distributed runtime

tensorflow/tensorflow

Distributed runtime

Multi-host execution machinery: gRPC master/worker, rendezvous across hosts, collective communication. Lives in tensorflow/core/distributed_runtime/ and tensorflow/core/nccl/.

Purpose

  • Run a TF graph across multiple hosts and devices.
  • Provide gRPC-based master and worker services.
  • Stream tensors across hosts via RemoteRendezvous.
  • Implement collective ops (all-reduce, all-gather, broadcast) using ring algorithms or NCCL.
  • Underpin tf.distribute.MultiWorkerMirroredStrategy, ParameterServerStrategy, and remote eager.

Directory layout

tensorflow/core/distributed_runtime/
├── master.{h,cc}                  # Master service: orchestrates a graph across workers
├── master_session.{h,cc}          # Per-Session graph partitioning across workers
├── worker.{h,cc}                  # Worker service: runs a partition locally
├── worker_session.{h,cc}
├── remote_device.{h,cc}           # Device handle that lives on another worker
├── rendezvous_mgr_interface.{h,cc} # Inter-worker rendezvous abstraction
├── eager/                         # Remote eager execution
├── coordination/                  # tf.distribute coordination service
├── rpc/                           # gRPC service implementations
└── ...

tensorflow/core/nccl/               # NCCL-backed collective communication

Cluster model

A TF cluster is described by a ClusterDef proto: a list of jobs (worker, ps, chief, etc.), each with a list of task addresses. The master is whichever process holds the Session; workers run partitions on their local devices; PS (parameter servers) are workers that hold variables only.

Key abstractions

Type File Description
Master / MasterService tensorflow/core/distributed_runtime/master.h gRPC service for Session-style RPCs.
MasterSession tensorflow/core/distributed_runtime/master_session.h One Session, partitioned across workers.
Worker / WorkerService tensorflow/core/distributed_runtime/worker.h gRPC service that runs op partitions.
RemoteDevice tensorflow/core/distributed_runtime/remote_device.h Stub for a device on another worker.
RemoteRendezvous tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc Cross-worker tensor exchange.
CoordinationServiceAgent tensorflow/core/distributed_runtime/coordination/ Heartbeats, barrier, leader election for tf.distribute.
EagerService tensorflow/core/distributed_runtime/eager/ RPCs for remote eager.
CollectiveExecutor tensorflow/core/common_runtime/collective_executor_mgr.h Coordinates collective ops across replicas.

How a distributed graph runs

graph LR
    Client[User process: Session.run]
    Master[Master service]
    W1[Worker 1: GPU 0,1]
    W2[Worker 2: GPU 0,1]
    PS[Parameter server]

    Client -->|gRPC| Master
    Master -->|partition + assign devices| W1
    Master -->|partition + assign devices| W2
    Master -->|partition + assign devices| PS
    W1 <-->|RemoteRendezvous send/recv| W2
    W1 <-->|RemoteRendezvous read/update var| PS
  1. Client opens a Session whose target is the master's gRPC address.
  2. Master receives the GraphDef, places ops onto devices across the cluster (Placer + co-location), partitions per worker, ships partitions to each worker via RegisterGraph.
  3. On Run, the master tells each worker to run its piece. Workers run their own Executors locally and pull cross-worker tensors through RemoteRendezvous (gRPC-streamed).
  4. Variables on parameter servers are read/updated via Read/Apply-style ops; PSs are just workers with variable-holding kernels.

Collectives

Three collective implementations co-exist:

  • NCCL (tensorflow/core/nccl/) — preferred for GPU rings, especially intra-host NVLink. Used by MirroredStrategy and MultiWorkerMirroredStrategy when GPU is present.
  • Ring/Tree CPU collectivestensorflow/core/common_runtime/collective_* — fallback when NCCL isn't available.
  • TPU collectives — implemented through TPU runtime in tensorflow/core/tpu/.

The CollectiveExecutor in tensorflow/core/common_runtime/collective_executor_mgr.cc selects an algorithm and orchestrates rendezvous so each replica's CollectiveReduce op finds its peers.

Coordination service

tensorflow/core/distributed_runtime/coordination/ is a relatively recent addition (2021+) that provides barrier, heartbeat, and key-value store primitives for tf.distribute. It lets workers detect peer failure and synchronise startup. The gRPC service definition is in coordination_service.proto; client and server implementations are in the same directory.

Remote eager

EagerService (tensorflow/core/distributed_runtime/eager/) lets one process drive eager ops on another. The user calls tf.config.experimental_connect_to_cluster(...) and from that point EagerContext::Execute may dispatch to a remote worker via RemoteEagerExecute. This is the basis for multi-host eager and TPU eager.

gRPC details

  • Service definitions: tensorflow/core/protobuf/master_service.proto, worker_service.proto, eager_service.proto, coordination_service.proto.
  • Service implementations: tensorflow/core/distributed_runtime/rpc/.
  • A Server (tensorflow/core/distributed_runtime/server_lib.h) wraps the gRPC server and is what tf.distribute.Server constructs in Python.

Integration points

  • tf.distribute (tensorflow/python/distribute/) builds on this layer for multi-worker training. See features/distribution-strategy.
  • DTensor uses it for cross-host communication.
  • TPU uses a similar pattern with TPU-specific RPCs.

Entry points for modification

  • New gRPC service or RPC — define in tensorflow/core/protobuf/, implement under tensorflow/core/distributed_runtime/rpc/.
  • New collective algorithm — tensorflow/core/common_runtime/collective_* plus tensorflow/core/nccl/ for GPU.
  • Coordination behaviour — tensorflow/core/distributed_runtime/coordination/coordination_service.cc.

Built by Factory AutoWiki from public repository content. It is a generated preview for codebase exploration, not source-maintained documentation.

Distributed runtime – TensorFlow wiki | Factory