炼丹师的工程修养之四： TensorFlow的分布式训练和K8S

def main():
  # 从环境变量中读取TF_CONFIG，这个JSON字符串描述了整个集群的配置情况
  tf_config = json.loads(os.environ.get('TF_CONFIG') or '{}')

  cluster_config = tf_config.get('cluster', {})
  # 参数服务器的地址
  ps_hosts = cluster_config.get('ps')
  # Worker的地址
  worker_hosts = cluster_config.get('worker')

  # 使用集群配置，分别在各个节点上启动Server，这个Server负责与其他所有节点的通信
  server = tf.train.Server(
      cluster, job_name=job_name, task_index=task_index)
  # 如果当前节点是PS节点，join等待。PS节点的代码执行到此就没有了
  if job_name == "ps":
    server.join()

  # 以下是Worker节点执行的代码
  # 默认情况下，第一个Worker节点是chief节点，chief除了要计算外，还要负责计算图的初始化和分配，
  # 模型文件保存、恢复等管理类的工作
  is_chief = (task_index == 0)
  # 各个节点根据配置信息，生成设备描述，CPU和GPU的设备描述不一样，分别处理
  if num_gpus > 0:
    gpu = (task_index % num_gpus)
    worker_device = "/job:worker/task:%d/gpu:%d" % (task_index, gpu)
  elif num_gpus == 0:
    cpu = 0
    worker_device = "/job:worker/task:%d/cpu:%d" % (task_index, cpu)

  # 使用tf.device这个Context Manager指定接下来所有的操作都在这个设备上面
  with tf.device(
      tf.train.replica_device_setter(
          worker_device=worker_device,
          ps_device="/job:ps/cpu:0",
          cluster=cluster)):

    # 构建模型，省略
    ...
    cross_entropy = ...
    # 创建优化器
    opt = tf.train.AdamOptimizer(learning_rate)

    # 如果是同步更新模式，使用TF的SyncReplicasOptimizer包装原来的优化器
    # 该优化器有两个重要的参数，replicas_to_aggregate和total_num_replicas，这两个参数有三种可能：
    # 1）如果replicas_to_aggregate < total_num_replicas，表示PS只需要接收到replicas_to_aggregate多个batch的梯度就可以更新，
    # 而剩下的部分可以丢弃，实际上起到备份的作用，当某个batch失败了，任务还可以继续进行
    # 2）如果replicas_to_aggregate > total_num_replicas，表示PS要接收到多于total_num_replicas数量的梯度才会更新，这可以看做是“能者多劳”模式，
    # 即如果集群中有某些Worker的性能更高，那么可以允许其比其他Worker多计算一些batch，充分利用其性能
    # 3）如果replicas_to_aggregate = total_num_replicas，其实就是“公平模式”，要求必须接收到规定数量的梯度才会更新
    if sync_replicas:
      opt = tf.train.SyncReplicasOptimizer(
          opt,
          replicas_to_aggregate=replicas_to_aggregate,
          total_num_replicas=num_workers,
          name="mnist_sync_replicas")
    # 最终的训练OP
    train_step = opt.minimize(cross_entropy, global_step=global_step)
    # 同步模式下，需要显式的初始化几个OP
    if sync_replicas:
      local_init_op = opt.local_step_init_op
      if is_chief:
        local_init_op = opt.chief_init_op

      ready_for_local_init_op = opt.ready_for_local_init_op

      # chief节点需要借助一个QueueRunner来完成同步操作
      chief_queue_runner = opt.get_chief_queue_runner()
      sync_init_op = opt.get_init_tokens_op()

    init_op = tf.global_variables_initializer()

    # 借助Supervisor类帮助进行训练初始化和部分管理工作，如果是同步模式，则需要多传入两个参数
    if sync_replicas:
      sv = tf.train.Supervisor(
          is_chief=is_chief,
          logdir=train_dir,
          init_op=init_op,
          local_init_op=local_init_op,
          ready_for_local_init_op=ready_for_local_init_op,
          recovery_wait_secs=1,
          global_step=global_step)
    else:
      sv = tf.train.Supervisor(
          is_chief=is_chief,
          logdir=train_dir,
          init_op=init_op,
          recovery_wait_secs=1,
          global_step=global_step)
    # 创建Session的配置对象
    sess_config = tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=False,
        device_filters=["/job:ps",
                        "/job:worker/task:%d" % task_index])
    # 如果是chief节点，手动初始化sync_init_op并启动QueueRunner
    if sync_replicas and is_chief:
      # Chief worker will start the chief queue runner and call the init op.
      sess.run(sync_init_op)
      sv.start_queue_runners(sess, [chief_queue_runner])

    # 开始训练的迭代，与单机训练无异
    while True:
      # Training feed
      batch_xs, batch_ys = mnist.train.next_batch(batch_size)
      train_feed = {x: batch_xs, y_: batch_ys}

      _, step = sess.run([train_step, global_step], feed_dict=train_feed)

{
  "cluster": {
    "ps": [
      "host1:2222"
    ],
    "worker": [
      "host2:3333",
      "host3:3333",
      "host3:4444"
    ]
  },
  "task": {
    "index": 0,
    "type": "ps"
  }
}

# 数据准备和输入，略过
def input_fn():
  ...

def main(args):

  # 构建计算图模型
  model = tf.keras.Sequential()
  model.add(tf.keras.layers.Dense(16, activation='relu', input_shape=(10,)))
  model.add(tf.keras.layers.Dense(1, activation='sigmoid'))

  # 创建优化器并compile模型
  optimizer = tf.train.GradientDescentOptimizer(0.2)
  model.compile(loss='binary_crossentropy', optimizer=optimizer)
  model.summary()
  tf.keras.backend.set_learning_phase(True)

  # 创建estimator的运行配置，在experimental_distribute参数里可以分别为训练和评估指定不同的策略
  # 这里我们指定在训练时用MultiWorkerMirroredStrategy策略，也可以使用其他的策略，比如ParameterServerStrategy等
  config = tf.estimator.RunConfig(
      experimental_distribute=tf.contrib.distribute.DistributeConfig(
          train_distribute=tf.contrib.distribute.MultiWorkerMirroredStrategy(
              num_gpus_per_worker=0),
          eval_distribute=tf.contrib.distribute.MirroredStrategy(
              num_gpus_per_worker=0)))
  # 构建estimator对象
  keras_estimator = tf.keras.estimator.model_to_estimator(
      keras_model=model, config=config, model_dir=model_dir)

  # 使用estimator启动分布式训练和评估
  tf.estimator.train_and_evaluate(
      keras_estimator,
      train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
      eval_spec=tf.estimator.EvalSpec(input_fn=input_fn))

# 数据准备 
def make_datasets_unbatched():
  ...
train_datasets = make_datasets_unbatched().batch(BATCH_SIZE)

#使用Keras API构建计算图模型，指定loss函数和优化器
def build_and_compile_cnn_model():
  model = tf.keras.Sequential([
      ...
  ])
  model.compile(
      loss=tf.keras.losses.sparse_categorical_crossentropy,
      optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
      metrics=['accuracy'])
  return model

# 通过Context Manager指定使用MultiWorkerMirroredStrategy策略，所有在这个Context下的操作，默认使用了该策略
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
with strategy.scope():
  # 这时才构建数据准备相关的Tensor和Op
  train_datasets = make_datasets_unbatched().batch(GLOBAL_BATCH_SIZE)
  # 这时才构建模型训练相关的Tensor和Op
  multi_worker_model = build_and_compile_cnn_model()
# 启动模型训练
multi_worker_model.fit(x=train_datasets, epochs=3, steps_per_epoch=5)

$ TF_CONFIG='{"cluster":{"worker":["p20677v.hulk.shbt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:4444"]},"task":{"index":0,"type":"worker"}}' python keras_with_strategy.py

2019-12-17 16:53:55.420119: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:250] Initialize GrpcChannelCache for job worker -> {0 -> localhost:3333, 1 -> k3711v.add.bjyt.qihoo.net:3333, 2 -> k3711v.add.bjyt.qihoo.net:4444}
2019-12-17 16:53:55.425259: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:365] Started server with target: grpc://localhost:3333
2019-12-17 16:53:55.425350: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Server already started (target: grpc://localhost:3333)

$ TF_CONFIG='{"cluster":{"worker":["p20677v.hulk.shbt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:4444"]},"task":{"index":1,"type":"worker"}}' python keras_with_strategy.py

2019-12-17 16:55:49.359229: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:250] Initialize GrpcChannelCache for job worker -> {0 -> p20677v.hulk.shbt.qihoo.net:3333, 1 -> localhost:3333, 2 -> k3711v.add.bjyt.qihoo.net:4444}
2019-12-17 16:55:49.362287: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:365] Started server with target: grpc://localhost:3333
2019-12-17 16:55:49.362364: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Server already started (target: grpc://localhost:3333)

$ TF_CONFIG='{"cluster":{"worker":["p20677v.hulk.shbt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:3333","k3711v.add.bjyt.qihoo.net:4444"]},"task":{"index":2,"type":"worker"}}' python keras_with_strategy.py

2019-12-17 16:58:32.681871: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:250] Initialize GrpcChannelCache for job worker -> {0 -> p20677v.hulk.shbt.qihoo.net:3333, 1 -> k3711v.add.bjyt.qihoo.net:3333, 2 -> localhost:4444}
2019-12-17 16:58:32.684099: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:365] Started server with target: grpc://localhost:4444
2019-12-17 16:58:32.684137: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Server already started (target: grpc://localhost:4444)

 1303/20000 [>.............................] - ETA: 9:33 - loss: 2.0348 - acc: 0.4191

FROM tensorflow/tensorflow:1.15.0
RUN pip install tensorflow_datasets -i https://pypi.douban.com/simple
ADD . /root
WORKDIR /root
ENTRYPOINT ["python", "/root/keras_with_strategy.py"]

apiVersion: batch/v1
kind: Job
metadata:
  name: worker-0
spec:
  template:
    spec:
      containers:
      - name: worker-0
        image: registry.cn-hangzhou.aliyuncs.com/chenzhen-docker-repo/keras_with_stragety:v3
        env:
        - name: TF_CONFIG
          value: '{"cluster":{"worker":["????:3333","????:3333","????:3333"]},"task":{"index":0,"type":"worker"}}'
        ports:
        - containerPort: 3333
      restartPolicy: Never
  backoffLimit: 4

kind: Service
apiVersion: v1
metadata:
  name: keras-with-strategy-worker-0
spec:
  type: LoadBalancer
  selector:
    job-name: worker-0
  ports:
  - port: 3333

...
        - name: TF_CONFIG
          value: '{"cluster":{"worker":["keras-with-strategy-worker-0:3333","keras-with-strategy-worker-1:3333","keras-with-strategy-worker-2:3333"]},"task":{"index":0,"type":"worker"}}'
...

$ kubectl create -f worker_0.yaml
service/keras-with-strategy-worker-0 created
job.batch/worker-0 created

$ kubectl create -f worker_1.yaml
service/keras-with-strategy-worker-1 created
job.batch/worker-1 created

$ kubectl create -f worker_2.yaml
service/keras-with-strategy-worker-2 created
job.batch/worker-2 created

$ kubectl get service
NAME                           TYPE           CLUSTER-IP       EXTERNAL-IP   PORT(S)          AGE
keras-with-strategy-worker-0   LoadBalancer   10.100.27.222    <pending>     3333:31342/TCP   16s
keras-with-strategy-worker-1   LoadBalancer   10.107.229.30    <pending>     3333:32135/TCP   10s
keras-with-strategy-worker-2   LoadBalancer   10.106.108.169   <pending>     3333:30083/TCP   7s
kubernetes                     ClusterIP      10.96.0.1        <none>        443/TCP          35d

$ kubectl get job
NAME       COMPLETIONS   DURATION   AGE
worker-0   0/1           3m         3m
worker-1   0/1           2m54s      2m54s
worker-2   0/1           2m51s      2m51s

$ kubectl get pods
NAME             READY   STATUS    RESTARTS   AGE
worker-0-2mgx4   1/1     Running   0          3m19s
worker-1-jvqff   1/1     Running   0          3m13s
worker-2-5xrkq   1/1     Running   0          3m10s

$ kubectl logs -f worker-0-2mgx4
...
Train on 20000 steps
Epoch 1/5
 1324/20000 [>.............................] - ETA: 12:26 - loss: 1.9863 - acc: 0.4678