Repository: horovod/horovod

Issue:
TF estimators and horovod.spark
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet)
TensorFlow
2. Framework version:
1.13.1
3. Horovod version:
0.16
4. MPI version:
3.0.0
5. CUDA version:
N/A
6. NCCL version:
N/A
7. Python version:
3.5
8. OS and version:
Amazon Linux AMI release 2018.03

**Your question:**
Is there any guidance on how to use horovod.spark with TF estimators? I've been able to train models using TF-Keras and horovod.spark in the style of the keras_spark_rossman.py example script, but I'm curious about how horovod.spark would work with Estimator's?

Moreover, I think I'm just confused about how horovod.spark actually works - is it okay to instantiate an Estimator instance inside the input_fn for a horovod.spark.run() call, or will that mean each worker has a different starting point for the model?


Discussion:
@bluesummers1129, you can instantiate `model_fn` outside of Horovod's `train_fn`, like it's done here - https://github.com/horovod/horovod/blob/master/examples/tensorflow_mnist_estimator.py#L32.  Horovod in Spark transmits both `train_fn` and all of it dependent functions, so it's fine to define functions outside of `train_fn`.
When I try to modify a TF workflow that works on a single machine to do this I run into issues with how feature_columns are passed to the model_fn. I'm getting a strange error in the _normalize_feature_columns() call when in horovod.spark.run mode:

```
[1,1]<stderr>:    features, labels, model_fn_lib.ModeKeys.TRAIN, self.config)
[1,1]<stderr>:  File "/usr/local/lib/python3.5/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1112, in _call_model_fn
[1,1]<stderr>:    model_fn_results = self._model_fn(features=features, **kwargs)
[1,1]<stderr>:  File "/mnt/Optimitron/optimitron_spark_tf.py", line 616, in tf_model_fn
[1,1]<stderr>:  File "/usr/local/lib/python3.5/site-packages/tensorflow/python/feature_column/feature_column.py", line 302, in input_layer
[1,1]<stderr>:    cols_to_output_tensors=cols_to_output_tensors)
[1,1]<stderr>:  File "/usr/local/lib/python3.5/site-packages/tensorflow/python/feature_column/feature_column.py", line 181, in _internal_input_layer
[1,1]<stderr>:    feature_columns = _normalize_feature_columns(feature_columns)
[1,1]<stderr>:  File "/usr/local/lib/python3.5/site-packages/tensorflow/python/feature_column/feature_column.py", line 2263, in _normalize_feature_columns
[1,1]<stderr>:    'Given (type {}): {}.'.format(type(column), column))
[1,1]<stderr>:ValueError: Items of feature_columns must be a _FeatureColumn. Given (type <class 'collections.IndicatorColumn'>): IndicatorColumn(categorical_column=VocabularyListCategoricalColumn(...
```

I generate feature columns outside of my horovod.spark call using:
```
def get_tf_features(indexers, whitelist_features=None):
    """

    Args:
        indexers (dict):
        whitelist_features (list[str]): optional list of features we want to include

    Returns:
        list[tf.feature_column._FeatureColumn:

    """
    feature_columns = []
    if whitelist_features:
        selected_features = set(indexers.keys()).intersection(set(whitelist_features))
    else:
        selected_features = set(indexers.keys())

    for feature in selected_features:
        feature_type = FEATURE_INDEXERS.get(feature, OneHotIndexer).FEATURE_TYPE
        if feature_type in (FeatureType["ONE_HOT"], FeatureType["MULTI_HOT"]):
            col = tf.feature_column.indicator_column(
                tf.feature_column.categorical_column_with_vocabulary_list(
                    key=feature,
                    vocabulary_list=[k for k in indexers[feature]],
                    default_value=0
                )
            )
        elif feature_type == FeatureType["MULTI_HOT_SCORED"]:
            col = tf.feature_column.weighted_categorical_column(
                tf.feature_column.categorical_column_with_vocabulary_list(
                    key=feature + INDEX_SUFFIX,
                    vocabulary_list=[k for k in indexers[feature]],
                    default_value=0),
                feature + SCORE_SUFFIX,
                )
        else:
            raise Exception("whoops")

        feature_columns.append(col)

    return feature_columns

feature_columns = get_tf_columns(indexers)
```
so I'm sure that my columns are correct, and then define the model_fn as
```
def tf_model_fn(features, labels, mode, params):
    """

    Args:
        features: This is the first item returned from the input_fn passed to train, evaluate, and predict.
            This should be a single tf.Tensor or dict of same.
        labels: This is the second item returned from the input_fn passed to train, evaluate, and predict.
            This should be a single tf.Tensor or dict of same (for multi-head models).
            If mode is tf.estimator.ModeKeys.PREDICT, labels=None will be passed.
            If the model_fn's signature does not accept mode, the model_fn must still be able to handle labels=None.
        mode (tf.estimator.ModeKeys): Optional. Specifies if this training, evaluation or prediction. See tf.estimator.ModeKeys.
        params (dict): optional dict of hyperparameters, received from Estimator instantiation

    Returns:
        tf.estimator.EstimatorSpec:

    """
    import horovod.tensorflow as hvd

    # Build the dense model
    net = tf.feature_column.input_layer(features, list(params['feature_columns']))
    for units in params['hidden_units']:
        net = tf.layers.dense(net, units=units, activation=tf.nn.relu)
    logits = tf.layers.dense(net, 1, activation=None)

    # Generate predictions (for PREDICT and EVAL mode)
    probabilities = tf.nn.sigmoid(logits)
    predictions = {
        'probabilities': probabilities,
        'logits': logits,
    }
    # Compute log-loss (for sigmoid activation)
    # y_train.reshape((y_train.shape[0], 1))
    labels_reshaped = tf.reshape(labels, [-1, 1])
    loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=labels_reshaped, logits=logits)
    # loss = tf.losses.log_loss(labels, probabilities)

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = params.get('optimizer', hvd.DistributedOptimizer(tf.train.FtrlOptimizer(learning_rate=params["learning_rate"])))

        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step()
        )
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    elif mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss)
    elif mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    else:
        raise Exception("incorrect mode!")
```

and run it as
```
def input_fn_tf_estimator(model_fn, feature_columns, feature_types, data_path, train_rows, valid_rows, batch_size=BATCH_SIZE):
    """An input function is a function that returns a tf.data.Dataset object which outputs the following two-element tuple:

    Args:
        model_fn: model function for invoking an Estimator
        feature_columns (list[tf.feature_column._FeatureColumn]):
        feature_types (dict[str, str]):
        data_path (str):
        train_rows (int):
        valid_rows (int):
        batch_size (int):

    Returns:

    """

    # # Make sure pyarrow is referenced before anything else to avoid segfault due to conflict
    # # with TensorFlow libraries.  Use `pa` package reference to ensure it's loaded before
    # # functions like `deserialize_model` which are implemented at the top level.
    # # See https://jira.apache.org/jira/browse/ARROW-3346
    # pa
    import tensorflow as tf

    import atexit
    import horovod.tensorflow as hvd
    import os
    from petastorm import make_batch_reader
    from petastorm.tf_utils import make_petastorm_dataset
    import tempfile
    import shutil

    # Horovod: initialize Horovod inside the trainer.
    hvd.init()

    # Horovod: save checkpoints only on worker 0 to prevent other workers from
    # corrupting them.
    model_dir = './tf_model' if hvd.rank() == 0 else None

    # # Horovod: pin GPU to be used to process local rank (one GPU per process), if GPUs are available.
    config = tf.ConfigProto()

    training_path = u'{}/{}={}/'.format(data_path, DATASET_COLUMN, TRAINING_SET_LABEL)
    valid_path = u'{}/{}={}/'.format(data_path, DATASET_COLUMN, VALID_SET_LABEL)
    print("Training: {}".format(training_path))
    print("Valid: {}".format(valid_path))
    print("steps per epoch: {}".format(int(train_rows / batch_size / hvd.size())))
    print("validation steps: {}".format(int(valid_rows / batch_size / hvd.size())))

    def train_input_fn(train_data, feature_types, batch_size):
        reader = make_batch_reader(train_data, num_epochs=None, cur_shard=hvd.rank(), shard_count=hvd.size(),
                                   hdfs_driver=PETASTORM_HDFS_DRIVER, schema_fields=list(feature_types.keys()) + ["label"])
        ds = make_petastorm_dataset(reader) \
            .apply(tf.data.experimental.unbatch()) \
            .batch(batch_size)

        def row_mapper(row):
            out_dict = dict()
            for feature in feature_types:
                feature_type = feature_types[feature]
                if feature_type == "ONE_HOT":
                    out_dict[feature] = getattr(row, feature)
                elif feature_type == "MULTI_HOT":
                    indices_and_scores_as_strings = tf.string_split(getattr(row, feature), ",")  # Still strings!
                    out_dict[feature] = indices_and_scores_as_strings
                elif feature_type == "MULTI_HOT_SCORED":
                    raise NotImplementedError("No implementation for multi_hot_scored yet!")
                else:
                    raise Exception()

            label = getattr(row, "label")
            return out_dict, label

        return ds.map(row_mapper)

    # Create model
    classifier = tf.estimator.Estimator(
        model_fn=model_fn,
        model_dir=model_dir,
        config=tf.estimator.RunConfig(session_config=config),
        params={
            "hidden_units": [10, ],
            "feature_columns": feature_columns,
            "learning_rate": 0.1
        }
    )
    bcast_hook = hvd.BroadcastGlobalVariablesHook(0)

    # Train it
    classifier.train(
        input_fn=lambda: train_input_fn(training_path, feature_types, batch_size),
        steps=int(train_rows / batch_size / hvd.size()),
        hooks=[bcast_hook, ]
    )


horovod.spark.run(
                    input_fn_tf_estimator,
                    args=(tf_model_fn, feature_columns, feature_types, data_path, num_train_rows, num_valid_rows),
                    kwargs={
                        "batch_size": batch_size
                    },
                    num_proc=obj.num_workers,
                    verbose=2
                )[0]
```
I can get this to work in a regular tf.Session() as opposed to a hvd.init() and spark session, so it seems to be related to how Spark is interacting with horovod?
@bluesummers1129, I wonder if `cloudpickle` fails to properly serialize TensorFlow's FeatureColumns?  Is it possible to call the function that generates feature columns inside your model function?
I tried this as well and just sent the indexers dict to the input_fn_tf_estimator(), but this had the same issue. I also tried moving the tf_model_fn definition itself into the input_fn, and this also didn't work. Just for completeness I also refactored it as a simple horovodrun script without spark, and it still didn't work
Small update - I rewrote the script from scratch as a pure horovorun execution, and had it import horovod.tensorflow and pyspark at the top level rather than in the individual functions, and it works as a call to horovodrun. It still doesn't work as a horovod.spark run, though, even if I move the model_fn definition into the function being called by horovod.spark.run(). It really does seem like this is explicitly something to do with how horovod is handling executions within a Spark context, but I'm not sure how to further debug this
Assigned to @tgaddair to take a look.
Hey @bluesummers1129, this appears to be a bug with PySpark.  I've filed a bug with them [here](https://issues.apache.org/jira/browse/SPARK-27810).

In the meantime, you should be able to workaround this issue by doing the following at the top of your script:

```
import collections
collections.namedtuple.__hijack = 1
``` 

This will prevent pyspark from overwriting the collections.namedtuple class and breaking cloudpickle.
I tried putting this at the top of my script before all the other imports but am still seeing the same issue. Do I also have to import this inside the input_fn itself?
Possible that we need to do the same on the backend of horovod when deserializing the function.  I'll take a look and get back to you.
Any update on this yet? I'm making things work through horovodrun rather than horovod.spark, but just want to know if this is fixed/will get fixed sometime soon so I can revert to the horovod.spark implementation
Hey @bluesummers1129, thanks for following up on this.  Did some more digging and created #1180 to fix this.  Feel free to try the branch and let me know if it works!
I tried to build that branch but it seems to have an updated third-party package (gloo) that's throwing an error when I try to git clone --recursive it:

error: Server does not allow request for unadvertised object d58881cfd20dfe8c94ad3b6d88efbcd14d3f2ebe
Fetched in submodule path 'third_party/gloo', but it did not contain d58881cfd20dfe8c94ad3b6d88efbcd14d3f2ebe. Direct fetching of that commit failed.

Weirdly, I can see that commit in the gloo package itself, I guess I don't have permission to that commit?

Resolving patch:
diff --git a/horovod/spark/__init__.py b/horovod/spark/__init__.py
--- a/horovod/spark/__init__.py
+++ b/horovod/spark/__init__.py
@@ -13,6 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
+# Workaround for https://issues.apache.org/jira/browse/SPARK-22674
+# This fix also requires the user to make this same change at the top of their
+# training script before importing pyspark (on serialization).
+import collections
+collections.namedtuple.__hijack = 1
+
 import os
 import pyspark
 from six.moves import queue


Test patch: