# Buggy code
def distributed_train(self):
    """
    Distributed training.
    """
    if self._codec_type == "tf_example":
        codec = TFExampleCodec(self._all_columns)
    elif self._codec_type == "bytes":
        codec = BytesCodec(self._all_columns)
    else:
        raise ValueError("invalid codec_type: " + self._codec_type)
    while True:
        task = self.get_task()
        if not task.shard_file_name:
            # No more task
            break
        batch_size = task.minibatch_size
        err_msg = ""
        try:
            with recordio.File(
                task.shard_file_name, "r", decoder=codec.decode
            ) as rdio_r:
                reader = rdio_r.get_reader(task.start, task.end)
                min_model_version = task.model_version
                while True:
                    record_buf = list(itertools.islice(reader, 0, batch_size))
                    if not record_buf:
                        break

                    for _ in range(self._max_retrain_num):
                        # TODO: optimize the logic to avoid unnecessary get_model call.
                        self.get_model(max(self._model_version, min_model_version))

                        batch_input_data, batch_label = self._input_fn(record_buf)

                        with tf.GradientTape() as tape:
                            inputs = []
                            for f_col in self._feature_columns:
                                inputs.append(batch_input_data[f_col.key])
                            if len(inputs) == 1:
                                inputs = inputs[0]
                            outputs = self._model.call(inputs, training=True)
                            loss = self._loss(outputs, batch_label.flatten())

                            # TODO:  Add regularization loss if any,
                            #        which should be divided by the number of contributing workers.
                        grads = tape.gradient(loss, self._model.trainable_variables)
                        print("Loss is ", loss.numpy())

                        accepted, min_model_version = self.report_gradient(grads)
                        if accepted:
                            break
                    else:
                        # Worker got stuck, fail the task.
                        # TODO: stop the worker if it fails to make any progress for some time.
                        raise RuntimeError("Worker got stuck")

        except Exception as ex:
            err_msg = str(ex)
            traceback.print_exc()
        self.report_task_result(task.task_id, err_msg)


# Traceback (AssertionError)
python -m unittest elasticdl/worker/*_test.py
/usr/local/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
from ._conv import register_converters as _register_converters
2019-05-21 13:57:07.262725: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
WARNING:tensorflow:From /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:642: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.
must be str, not NoneType
FLoss is  3.090042
Loss is  1.4608976
Loss is  0.913306
Loss is  0.5969497
Loss is  0.66515267
Loss is  0.3935135
Loss is  0.37774342
Loss is  0.289928
.
======================================================================
FAIL: test_distributed_train (elasticdl.worker.worker_test.WorkerTest)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/Users/l.zou/git/elasticdl/elasticdl/worker/worker_test.py", line 96, in test_distributed_train
self.assertTrue(res)
AssertionError: False is not true

----------------------------------------------------------------------
Ran 2 tests in 0.165s

# Fixed code
def distributed_train(self):
    """
    Distributed training.
    """
    while True:
        task = self.get_task()
        if not task.shard_file_name:
            # No more task
            break
        batch_size = task.minibatch_size
        err_msg = ""
        try:
            with recordio.File(
                task.shard_file_name, "r", decoder=self._codec.decode
            ) as rdio_r:
                reader = rdio_r.get_reader(task.start, task.end)
                min_model_version = task.model_version
                while True:
                    record_buf = list(itertools.islice(reader, 0, batch_size))
                    if not record_buf:
                        break

                    for _ in range(self._max_retrain_num):
                        # TODO: optimize the logic to avoid unnecessary get_model call.
                        self.get_model(max(self._model_version, min_model_version))

                        batch_input_data, batch_label = self._input_fn(record_buf)

                        with tf.GradientTape() as tape:
                            inputs = []
                            for f_col in self._feature_columns:
                                inputs.append(batch_input_data[f_col.key])
                            if len(inputs) == 1:
                                inputs = inputs[0]
                            outputs = self._model.call(inputs, training=True)
                            loss = self._loss(outputs, batch_label.flatten())

                            # TODO:  Add regularization loss if any,
                            #        which should be divided by the number of contributing workers.
                        grads = tape.gradient(loss, self._model.trainable_variables)
                        print("Loss is ", loss.numpy())

                        accepted, min_model_version = self.report_gradient(grads)
                        if accepted:
                            break
                    else:
                        # Worker got stuck, fail the task.
                        # TODO: stop the worker if it fails to make any progress for some time.
                        raise RuntimeError("Worker got stuck")

        except Exception as ex:
            err_msg = str(ex)
            traceback.print_exc()
        self.report_task_result(task.task_id, err_msg)