Выполнение распределенного оценщика тензорного потока не запускает оценку или экспорт

Я тестирую распределенное обучение с помощью Оценщиков тензорного потока. В моем примере я использую простую синусоидальную функцию с настраиваемым оценщиком, используя _1 _. После обучения и оценки я хочу экспортировать модель, чтобы она была готова для обработки тензорного потока. Однако оценка и экспорт запускаются только при выполнении оценщика нераспределенным способом.

Модель и оценщики определены следующим образом:

def my_model(features, labels, mode):
    # define simple dense network
    net = tf.layers.dense(features['x'], units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)
    net = tf.layers.dense(net, units=8, activation=tf.nn.tanh)

    # output layer
    predictions = tf.layers.dense(net, units=1, activation=tf.nn.tanh)

    if mode == tf.estimator.ModeKeys.PREDICT:
        # define output message for tensorflow serving
        export_outputs = {'predict_output': tf.estimator.export.PredictOutput({"predictions": predictions})}

        return tf.estimator.EstimatorSpec(mode=mode, predictions={'predictions': predictions}, export_outputs=export_outputs)
    elif mode == tf.estimator.ModeKeys.EVAL:
        # for evaluation simply use mean squared error
        loss = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
        metrics = {'mse': tf.metrics.mean_squared_error(labels, predictions)}

        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
    elif mode == tf.estimator.ModeKeys.TRAIN:
        # train on mse with Adagrad optimizer
        loss = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
        optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
    else:
        raise ValueError("unhandled mode: %s" % str(mode))


def main(_):
    # prepare training data
    default_batch_size = 50
    examples = [{'x': x, 'y': math.sin(x)} for x in [random.random()*2*math.pi for _ in range(10000)]]

    estimator = tf.estimator.Estimator(model_fn=my_model,
                                       config=tf.estimator.RunConfig(model_dir='sin_model',
                                                                     save_summary_steps=100))

    # function converting examples to dataset
    def dataset_fn():
        # returns a dataset serving batched (feature_map, label)-pairs
        # e.g. ({'x': [1.0, 0.3, 1.1...]}, [0.84, 0.29, 0.89...])
        return tf.data.Dataset.from_generator(
            lambda: iter(examples),
            output_types={"x": tf.float32, "y": tf.float32},
            output_shapes={"x": [], "y": []}) \
            .map(lambda x: ({'x': [x['x']]}, [x['y']])) \
            .repeat() \
            .batch(default_batch_size)

    # function to export model to be used for serving
    feature_spec = {'x': tf.FixedLenFeature([1], tf.float32)}
    def serving_input_fn():
        serialized_tf_example = tf.placeholder(dtype=tf.string, shape=[default_batch_size])

        receiver_tensors = {'examples': serialized_tf_example}
        features = tf.parse_example(serialized_tf_example, feature_spec)
        return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

    # train, evaluate and export
    train_spec = tf.estimator.TrainSpec(input_fn=dataset_fn, max_steps=1000)
    eval_spec = tf.estimator.EvalSpec(input_fn=dataset_fn,
                                      steps=100,
                                      exporters=[tf.estimator.FinalExporter('sin', serving_input_fn)])

    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

if __name__ == '__main__':
    tf.app.run(main)

При выполнении этого кода в одном процессе я получаю выходную папку, содержащую контрольные точки модели, данные оценки и экспорт модели.

$ ls sin_model/
checkpoint                                  model.ckpt-0.index
eval                                        model.ckpt-0.meta
events.out.tfevents.1532426226.simon        model.ckpt-1000.data-00000-of-00001
export                                      model.ckpt-1000.index
graph.pbtxt                                 model.ckpt-1000.meta
model.ckpt-0.data-00000-of-00001

Однако при распространении процесса обучения (в этой тестовой настройке только на локальной машине) папки eval и export отсутствуют.

Я запускаю отдельные узлы, используя следующую конфигурацию кластера:

{"cluster": {
    "ps": ["localhost:2222"],
    "chief": ["localhost:2223"], 
    "worker": ["localhost:2224"]
}

Запуск сервера ps выглядит следующим образом

$ TF_CONFIG='{"cluster": {"chief": ["localhost:2223"], "worker": ["localhost:2224"], "ps": ["localhost:2222"]}, "task": {"type": "ps", "index": 0}}' CUDA_VISIBLE_DEVICES= python custom_estimator.py
2018-07-24 12:09:04.913967: E tensorflow/stream_executor/cuda/cuda_driver.cc:397] failed call to cuInit: CUDA_ERROR_NO_DEVICE
2018-07-24 12:09:04.914008: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:132] retrieving CUDA diagnostic information for host: simon
2018-07-24 12:09:04.914013: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:139] hostname: simon
2018-07-24 12:09:04.914035: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:163] libcuda reported version is: 384.130.0
2018-07-24 12:09:04.914059: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:167] kernel reported version is: 384.130.0
2018-07-24 12:09:04.914079: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:249] kernel version seems to match DSO: 384.130.0
2018-07-24 12:09:04.914961: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job chief -> {0 -> localhost:2223}
2018-07-24 12:09:04.914971: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job ps -> {0 -> localhost:2222}
2018-07-24 12:09:04.914976: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job worker -> {0 -> localhost:2224}
2018-07-24 12:09:04.915658: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Started server with target: grpc://localhost:2222

(Я добавил CUDA_VISIBLE_DEVICES= в командную строку, чтобы рабочий и руководитель не выделяли память графического процессора. Это вызывает ошибку failed call to cuInit: CUDA_ERROR_NO_DEVICE, которая, однако, не является критической)

Затем начальник запускается следующим образом

$ TF_CONFIG='{"cluster": {"chief": ["localhost:2223"], "worker": ["localhost:2224"], "ps": ["localhost:2222"]}, "task": {"type": "chief", "index": 0}}' CUDA_VISIBLE_DEVICES= python custom_estimator.py
2018-07-24 12:09:10.532171: E tensorflow/stream_executor/cuda/cuda_driver.cc:397] failed call to cuInit: CUDA_ERROR_NO_DEVICE
2018-07-24 12:09:10.532234: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:132] retrieving CUDA diagnostic information for host: simon
2018-07-24 12:09:10.532241: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:139] hostname: simon
2018-07-24 12:09:10.532298: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:163] libcuda reported version is: 384.130.0
2018-07-24 12:09:10.532353: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:167] kernel reported version is: 384.130.0
2018-07-24 12:09:10.532359: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:249] kernel version seems to match DSO: 384.130.0
2018-07-24 12:09:10.533195: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job chief -> {0 -> localhost:2223}
2018-07-24 12:09:10.533207: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job ps -> {0 -> localhost:2222}
2018-07-24 12:09:10.533211: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job worker -> {0 -> localhost:2224}
2018-07-24 12:09:10.533835: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Started server with target: grpc://localhost:2223
2018-07-24 12:09:14.038636: I tensorflow/core/distributed_runtime/master_session.cc:1165] Start master session 71a2748ad69725ae with config: allow_soft_placement: true graph_options { rewrite_options { meta_optimizer_iterations: ONE } }

Затем рабочий запускается следующим образом:

$ TF_CONFIG='{"cluster": {"chief": ["localhost:2223"], "worker": ["localhost:2224"], "ps": ["localhost:2222"]}, "task": {"type": "worker", "index": 0}}' CUDA_VISIBLE_DEVICES= python custom_estimator.py
2018-07-24 12:09:13.172260: E tensorflow/stream_executor/cuda/cuda_driver.cc:397] failed call to cuInit: CUDA_ERROR_NO_DEVICE
2018-07-24 12:09:13.172320: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:132] retrieving CUDA diagnostic information for host: simon
2018-07-24 12:09:13.172327: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:139] hostname: simon
2018-07-24 12:09:13.172362: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:163] libcuda reported version is: 384.130.0
2018-07-24 12:09:13.172399: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:167] kernel reported version is: 384.130.0
2018-07-24 12:09:13.172405: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:249] kernel version seems to match DSO: 384.130.0
2018-07-24 12:09:13.173230: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job chief -> {0 -> localhost:2223}
2018-07-24 12:09:13.173242: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job ps -> {0 -> localhost:2222}
2018-07-24 12:09:13.173247: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job worker -> {0 -> localhost:2224}
2018-07-24 12:09:13.173783: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:369] Started server with target: grpc://localhost:2224
2018-07-24 12:09:18.774264: I tensorflow/core/distributed_runtime/master_session.cc:1165] Start master session 1d13ac84816fdc80 with config: allow_soft_placement: true graph_options { rewrite_options { meta_optimizer_iterations: ONE } }

Через короткое время главный процесс останавливается, и папка sin_model существует с контрольной точкой модели, но без экспорта или оценки:

$ ls sin_model/
checkpoint                                  model.ckpt-0.meta
events.out.tfevents.1532426950.simon        model.ckpt-1001.data-00000-of-00001
graph.pbtxt                                 model.ckpt-1001.index
model.ckpt-0.data-00000-of-00001            model.ckpt-1001.meta
model.ckpt-0.index

Требуется ли какая-либо дополнительная настройка для оценки или экспорта в распределенной настройке?

Я работаю с python 3.5 и tensorflow 1.8