Repository: ludwig-ai/ludwig

Issue:
Parquet dataset does not have a `self.features` attribute, which is needed in the IterableBatcher
**Describe the bug**

The `IterableBatcher` will loop through all the features of a dataset (see[ this line of code](https://github.com/ludwig-ai/ludwig/blob/master/ludwig/data/batcher/iterable.py#L40)) object via an attribute `self.features` of a `Dataset` object

When the dataset is a parquet dataset (a new feature merged to master in [PR#1090](https://github.com/ludwig-ai/ludwig/commit/01d114ba76667dc52d1673d9be0809ee1e58087c)) (note that parquet is a required format when running ludwig on Ray), the `IteratableBatcher` cannot find this attribute `self.features`.

By going through the code of `parquet.py` I wasn't able to fnd a `self.features` attributes. There is only a self.reshape_features, which does not seem to be compatible with `IterableBatcher`.


**To Reproduce**
See my code on petfinder dataset: https://github.com/zhisbug/ludwig-petfinder/blob/master/petfinder-adoption-prediction/train.py#L11


**Expected behavior**
I guess a fix is needed to make IteratableBatcher compatible with ParquetDataset?


**Environment (please complete the following information):**
 - OS: Ubuntu 18
- Python version: 3.6
- Ludwig version: master

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -28,7 +28,7 @@
 from ludwig.data.dataframe.dask import DaskEngine
 from ludwig.models.predictor import BasePredictor, RemotePredictor
 from ludwig.models.trainer import BaseTrainer, RemoteTrainer
-from ludwig.utils.tf_utils import initialize_tensorflow
+from ludwig.utils.tf_utils import initialize_tensorflow, save_weights_to_buffer, load_weights_from_buffer
 
 
 logger = logging.getLogger(__name__)
@@ -73,12 +73,14 @@ def get_total_resources(bucket):
 
 class RayRemoteModel:
     def __init__(self, model):
+        buf = save_weights_to_buffer(model)
         self.cls, self.args, state = list(model.__reduce__())
-        self.state = ray.put(state)
+        self.state = ray.put(buf)
 
     def load(self):
         obj = self.cls(*self.args)
-        obj.__setstate__(ray.get(self.state))
+        buf = ray.get(self.state)
+        load_weights_from_buffer(obj, buf)
         return obj
 
 
@@ -90,13 +92,13 @@ def train(self, *args, **kwargs):
         results = super().train(*args, **kwargs)
         if results is not None:
             model, *stats = results
-            results = (model.get_weights(), *stats)
+            results = (save_weights_to_buffer(model), *stats)
         return results
 
     def train_online(self, *args, **kwargs):
         results = super().train_online(*args, **kwargs)
         if results is not None:
-            results = results.get_weights()
+            results = save_weights_to_buffer(results)
         return results
 
 
@@ -114,7 +116,7 @@ def train(self, model, *args, **kwargs):
         )
 
         weights, *stats = results[0]
-        model.set_weights(weights)
+        load_weights_from_buffer(model, weights)
         return (model, *stats)
 
     def train_online(self, model, *args, **kwargs):
@@ -124,7 +126,7 @@ def train_online(self, model, *args, **kwargs):
         )
 
         weights = results[0]
-        model.set_weights(weights)
+        load_weights_from_buffer(model, weights)
         return model
 
     @property
diff --git a/ludwig/data/dataset/parquet.py b/ludwig/data/dataset/parquet.py
--- a/ludwig/data/dataset/parquet.py
+++ b/ludwig/data/dataset/parquet.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import math
+
 import tensorflow as tf
 
 from petastorm import make_batch_reader
@@ -27,6 +29,7 @@
 class ParquetDataset(Dataset):
     def __init__(self, url, features, training_set_metadata):
         self.url = url
+        self.features = [feature[PROC_COLUMN] for feature in features]
         self.training_set_metadata = training_set_metadata
 
         with make_batch_reader(self.url) as reader:
@@ -76,7 +79,7 @@ def initialize_batcher(self,
             dataset = dataset.shuffle(buffer_size)
         dataset = dataset.batch(batch_size)
 
-        steps_per_epoch = int(local_samples / batch_size)
+        steps_per_epoch = math.ceil(local_samples / batch_size)
 
         batcher = IterableBatcher(self,
                                   dataset,
diff --git a/ludwig/utils/tf_utils.py b/ludwig/utils/tf_utils.py
--- a/ludwig/utils/tf_utils.py
+++ b/ludwig/utils/tf_utils.py
@@ -14,10 +14,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import io
+import os
+import shutil
+import tempfile
 import multiprocessing
 import warnings
+import zipfile
 
 import tensorflow as tf
+from ludwig.globals import MODEL_WEIGHTS_FILE_NAME
 
 _TF_INIT_PARAMS = None
 
@@ -139,3 +145,21 @@ def get_available_gpus_cuda_string():
     if len(gpus) == 0:
         return None
     return ','.join(gpus)
+
+
+def save_weights_to_buffer(model):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        weights_path = os.path.join(tmpdir, MODEL_WEIGHTS_FILE_NAME)
+        model.save_weights(weights_path)
+        with tempfile.TemporaryDirectory() as zipdir:
+            shutil.make_archive(os.path.join(zipdir, MODEL_WEIGHTS_FILE_NAME), 'zip', tmpdir)
+            with open(os.path.join(zipdir, f'{MODEL_WEIGHTS_FILE_NAME}.zip'), 'rb') as f:
+                return f.read()
+
+
+def load_weights_from_buffer(model, buf):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with zipfile.ZipFile(io.BytesIO(buf)) as zip_ref:
+            zip_ref.extractall(tmpdir)
+        weights_path = os.path.join(tmpdir, MODEL_WEIGHTS_FILE_NAME)
+        model.load_weights(weights_path)


Test patch:
diff --git a/tests/integration_tests/test_hyperopt.py b/tests/integration_tests/test_hyperopt.py
--- a/tests/integration_tests/test_hyperopt.py
+++ b/tests/integration_tests/test_hyperopt.py
@@ -88,6 +88,10 @@
 def test_hyperopt_executor(sampler, executor, csv_filename,
                            validate_output_feature=False,
                            validation_metric=None):
+    if executor['type'] == 'fiber' and sampler['type'] == 'grid':
+        # This test is very slow and doesn't give us additional converage
+        pytest.skip('Skipping Fiber grid search')
+
     input_features = [
         text_feature(name="utterance", cell_type="lstm", reduce_output="sum"),
         category_feature(vocab_size=2, reduce_input="sum")]
diff --git a/tests/integration_tests/test_ray.py b/tests/integration_tests/test_ray.py
--- a/tests/integration_tests/test_ray.py
+++ b/tests/integration_tests/test_ray.py
@@ -38,8 +38,8 @@
 
 
 @pytest.fixture
-def ray_start_4_cpus():
-    res = ray.init(num_cpus=4)
+def ray_start_2_cpus():
+    res = ray.init(num_cpus=2)
     try:
         yield res
     finally:
@@ -50,7 +50,7 @@ def run_api_experiment(config, data_parquet):
     # Sanity check that we get 4 slots over 1 host
     kwargs = get_horovod_kwargs()
     assert kwargs.get('num_hosts') == 1
-    assert kwargs.get('num_slots') == 4
+    assert kwargs.get('num_slots') == 2
 
     # Train on Parquet
     dask_backend = RayBackend()
@@ -70,7 +70,7 @@ def run_test_parquet(
         'input_features': input_features,
         'output_features': output_features,
         'combiner': {'type': 'concat', 'fc_size': 14},
-        'training': {'epochs': 2}
+        'training': {'epochs': 2, 'batch_size': 8}
     }
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -86,7 +86,7 @@ def run_test_parquet(
 
 
 @pytest.mark.distributed
-def test_ray_tabular(ray_start_4_cpus):
+def test_ray_tabular(ray_start_2_cpus):
     input_features = [
         sequence_feature(reduce_output='sum'),
         numerical_feature(normalization='zscore'),