def METHOD_NAME(self):
    """Verify that metric data is complete during fit when using
    ParameterServerStrategy
    """
    strategy = tf.distribute.ParameterServerStrategy(
        self.cluster_resolver,
        variable_partitioner=None,
    )
    class BatchCount(metrics_module.Sum):
        def __init__(self, name="batch_count", dtype=tf.int64):
            super().__init__(name=name, dtype=dtype)
        def update_state(self, y_true, y_pred, sample_weight=None):
            return super().update_state(1, sample_weight)
    # Build and compile model within strategy scope.
    with strategy.scope():
        inputs = layers_module.Input((1,))
        outputs = layers_module.Dense(1)(inputs)
        model = training_module.Model(inputs, outputs)
        model.compile(
            loss="mse", metrics=[BatchCount()], steps_per_execution=2
        )
    BATCH_SIZE = 10
    x, y = np.ones((400, 1)), np.ones((400, 1))
    val_x, val_y = np.ones((100, 1)), np.ones((100, 1))
    train_dataset = tf.data.Dataset.from_tensor_slices((x, y))
    train_dataset = train_dataset.batch(BATCH_SIZE)
    val_dataset = tf.data.Dataset.from_tensor_slices((val_x, val_y))
    val_dataset = val_dataset.batch(BATCH_SIZE)
    train_batch_count = x.shape[0] // BATCH_SIZE
    val_batch_count = val_x.shape[0] // BATCH_SIZE
    # Verify that Model fit doesn't drop any batches
    hist = model.fit(
        train_dataset,
        steps_per_epoch=train_batch_count,
        validation_data=val_dataset,
        validation_steps=val_batch_count,
        epochs=5,
    )
    # Verify that min and max value of batch count metric is accurate
    self.assertEqual(max(hist.history["batch_count"]), train_batch_count)
    self.assertEqual(min(hist.history["batch_count"]), train_batch_count)
    self.assertEqual(max(hist.history["val_batch_count"]), val_batch_count)
    self.assertEqual(min(hist.history["val_batch_count"]), val_batch_count)