Repository: horovod/horovod

Issue:
Pytorch_lightning's TorchEstimator doesn't pass transformation_fn correctly
**Environment:**
1. Framework: Pytorch Lightning
2. Framework version: 1.5.0
3. Horovod version: 0.23.0
4. MPI version:
5. CUDA version:
6. NCCL version:
7. Python version:
8. Spark / PySpark version: 3.2.0
9. Ray version:
10. OS and version:
11. GCC version:
12. CMake version:

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?

**Bug report:**
Please describe erroneous behavior you're observing and steps to reproduce it.

Pytorch_lightning's TorchEstimator accepts `transformation_fn` as a parameter but it's not actually passing it through petaStormDataModule so that the transformation_fn used in the training/validation is always the default value of `torch.as_tensor`. From the source code I found here (https://github.com/horovod/horovod/blob/464c82ef5307986531909861e01ee1c5b235640a/horovod/spark/lightning/datamodule.py#:~:text=transform_spec%20%3D%20TransformSpec,%3DFalse)) it assigned a value to `transform_spec` but only used it to distinguish the reader type and didn't pass it as a parameter to the reader.

A reproduce of the problem would be assigning a transformation_fn to horovod.spark.lightning.TorchEstimator and it doesn't use my transformation_fn to transform the row; the same transformation_fn works for horovod.spark.torch.TorchEstimator.




Discussion:
@chongxiaoc Hey could you take a look at this issue cuz it seems you're the owner of that datamodule file 😃
@serena-ruan 

Do you have a user case and example scripts for `transformation_fn`?
When we introduced `datamodule`, we found that we didn't use `transformation_fn` (from Spark Estimator users we know), so it was taken out for simplicity.

Also, feel free to contribute a PR if you find out where to insert `transformation_fn`.
I guess we can just pass it in here in file horovod/spark/lightning/datamodule.py:
<img width="675" alt="image" src="https://user-images.githubusercontent.com/82044803/167557563-5fd0d6c6-dfdc-4f9c-b705-c65272e6b9db.png">

This should pass through the param `transform_spec` into the Reader, so that when it runs PyDictReaderWorker to generate the reader it can call the transform_spec internally and changed the data schema. I don't have the dev environment setup ready so could you check if the above change works? Or could you guide me about the developing environment for raising a PR in this repo? 

Besides, when you say `we didn't use transformation_fn from Spark Estimator users` do you mean they just don't use this parameter at all? Because from my side I'm seeing the use case for example I want to pass in a dataframe for training but the dataframe just includes the paths to my dataset images, and I want to pass in a transformation_fn to read the image and transform it to tensors then work with my model.
@serena-ruan Our test cases are most recommendation system models which don't need `transform_spec`.
It is kind of supervised to me Spark estimator is used outside of Uber, which is a good sign though.
Ah I see, if we make the transform_spec work then it must be more useful for TorchEstimator usage in general!

> @serena-ruan Our test cases are most recommendation system models which don't need `transform_spec`. It is kind of supervised to me Spark estimator is used outside of Uber, which is a good sign though.



Resolving patch:
diff --git a/horovod/spark/lightning/datamodule.py b/horovod/spark/lightning/datamodule.py
--- a/horovod/spark/lightning/datamodule.py
+++ b/horovod/spark/lightning/datamodule.py
@@ -13,7 +13,7 @@ class PetastormDataModule(pl.LightningDataModule):
     def __init__(self, train_dir: str, val_dir: str, num_train_epochs: int=1, has_val: bool=True,
                  train_batch_size: int=32, val_batch_size: int=32, shuffle_size: int=1000,
                  num_reader_epochs=None, reader_pool_type: str="process",
-                 reader_worker_count: int=2, transform_spec=None, inmemory_cache_all=False,
+                 reader_worker_count: int=2, transformation=None, inmemory_cache_all=False,
                  cur_shard: int=0, shard_count: int=1, schema_fields=None, storage_options=None,
                  steps_per_epoch_train: int=1, steps_per_epoch_val: int=1, verbose=True,
                  debug_data_loader: bool=False, train_async_data_loader_queue_size: int=None,
@@ -29,7 +29,7 @@ def __init__(self, train_dir: str, val_dir: str, num_train_epochs: int=1, has_va
         self.num_reader_epochs = num_reader_epochs
         self.reader_pool_type = reader_pool_type
         self.reader_worker_count = reader_worker_count
-        self.transform_spec = transform_spec
+        self.transformation = transformation
         self.inmemory_cache_all = inmemory_cache_all
         self.cur_shard = cur_shard
         self.shard_count = shard_count
@@ -49,13 +49,15 @@ def __init__(self, train_dir: str, val_dir: str, num_train_epochs: int=1, has_va
     def setup(self, stage=None):
         # Assign train/val datasets for use in dataloaders
         if stage == 'fit' or stage is None:
-            transform_spec = TransformSpec(self.transform_spec) if self.transform_spec else None
+            transform_spec = TransformSpec(self.transformation) if self.transformation else None
             # In general, make_batch_reader is faster than make_reader for reading the dataset.
             # However, we found out that make_reader performs data transformations much faster than
             # make_batch_reader with parallel worker processes. Therefore, the default reader
             # we choose is make_batch_reader unless there are data transformations.
+            reader_factory_kwargs = dict()
             if transform_spec:
                 reader_factory = make_reader
+                reader_factory_kwargs['pyarrow_serialize'] = True
             else:
                 reader_factory = make_batch_reader
 
@@ -64,15 +66,19 @@ def setup(self, stage=None):
                                                hdfs_driver=PETASTORM_HDFS_DRIVER,
                                                schema_fields=self.schema_fields,
                                                storage_options=self.storage_options,
+                                               transform_spec=transform_spec,
                                                # Don't shuffle row groups without shuffling.
-                                               shuffle_row_groups=True if self.shuffle_size > 0 else False)
+                                               shuffle_row_groups=True if self.shuffle_size > 0 else False,
+                                               **reader_factory_kwargs)
             if self.has_val:
                 self.val_reader = reader_factory(self.val_dir, num_epochs=self.num_reader_epochs,
                                                  cur_shard=self.cur_shard, shard_count=self.shard_count,
                                                  hdfs_driver=PETASTORM_HDFS_DRIVER,
                                                  schema_fields=self.schema_fields,
                                                  storage_options=self.storage_options,
-                                                 shuffle_row_groups=False)
+                                                 transform_spec=transform_spec,
+                                                 shuffle_row_groups=False,
+                                                 **reader_factory_kwargs)
 
     def teardown(self, stage=None):
         if stage == "fit" or stage is None:
diff --git a/horovod/spark/lightning/remote.py b/horovod/spark/lightning/remote.py
--- a/horovod/spark/lightning/remote.py
+++ b/horovod/spark/lightning/remote.py
@@ -266,7 +266,7 @@ def on_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
                 'num_reader_epochs': loader_num_epochs,
                 'reader_pool_type': reader_pool_type,
                 'reader_worker_count': train_reader_worker_count,
-                'transform_spec': transformation,
+                'transformation': transformation,
                 'inmemory_cache_all': inmemory_cache_all,
                 'cur_shard': hvd.rank(),
                 'shard_count': hvd.size(),
diff --git a/horovod/spark/torch/remote.py b/horovod/spark/torch/remote.py
--- a/horovod/spark/torch/remote.py
+++ b/horovod/spark/torch/remote.py
@@ -264,7 +264,7 @@ def save_checkpoint():
                                 transform_spec=transform_spec,
                                 storage_options=storage_options,
                                 # Don't shuffle row groups without shuffling.
-                                shuffle_row_groups=True if shuffle_buffer_size > 0 else False
+                                shuffle_row_groups=True if shuffle_buffer_size > 0 else False,
                                 **reader_factory_kwargs) as train_reader:
                 with reader_factory(remote_store.val_data_path,
                                     num_epochs=None,


Test patch: