Repository: ludwig-ai/ludwig
Pull request: Hyperopt steps per epoch not being computed correctly
Following up on [this issue](https://github.com/ludwig-ai/ludwig/issues/2144), when using windowing in the backend over a partitioned dataset (> 1 partition), we observed that the `steps_per_epoch` was being under-calculated. This results in each epoch only comprising of a fraction of the dataset.

This PR modifies the `RayDatasetShard` class to use size of the `RayDataset` to calculate the number of epochs instead of the size of the `DatasetPipeline` to calculate the number of epochs, bringing back the number of epochs to what we'd expect.

**Validation that this change works**:
- Dataset size (100MB) ~ 354000 rows in the training data
- Partitioning the dataset into 10 blocks
- Setting the loader's window size to 2048 bytes

**Before**:
- Calling `train_cli()` with the dataset resulted in:
    - 10 windows being created
    - Just two steps per epoch, and with a batch size of ~32K, that meant we would only train over 64K rows per epoch

**Now**
- Calling `train_cli()` with the dataset results in
    - 10 windows being created
    - 11 steps per epoch, and with a batch size of ~32K, that means ~354K rows per epoch which is correct

``` 
2022-06-21 19:36:36,768 INFO dataset.py:2643 -- Created DatasetPipeline with 10 windows: 1.75MiB min, 1.78MiB max, 1.77MiB mean
2022-06-21 19:36:36,802 INFO dataset.py:2643 -- Created DatasetPipeline with 10 windows: 0.52MiB min, 0.54MiB max, 0.53MiB mean
2022-06-21 19:36:36,834 INFO dataset.py:2643 -- Created DatasetPipeline with 10 windows: 0.27MiB min, 0.29MiB max, 0.28MiB mean
``` 

```
100 [00:00<?, ?it/s](BaseWorkerMixin pid=13772) Note: steps_per_checkpoint (was 2000) is now set to the number of steps per epoch: 11.
(BaseWorkerMixin pid=13772) 
(BaseWorkerMixin pid=13772) Training for 1100 step(s), approximately 100 epoch(s).
(BaseWorkerMixin pid=13772) Early stopping policy: 10 round(s) of evaluation, or 110 step(s), approximately 10 epoch(s).
```

Problem statement:
[Hyperopt] Steps per epoch not being computed correctly
- Total dataset size = 570M (Including test, train, val)
- Batch size = 32K
- Expected steps per epoch = 570M / 32K = ~17800
- Observed steps per epoch = 354 (From stdout at start of experiment training `training for 1000 epochs, and 354000 steps`)
- Not performing distributed training
- To reproduce, run hyperopt on the dataset with 8 parallel trials.

Hints from issue discussion:
I suspect this is because of the windowing. [Here](https://github.com/ludwig-ai/ludwig/blob/master/ludwig/data/dataset/ray.py#L98) we define the window over the dataset, and later [here](https://github.com/ludwig-ai/ludwig/blob/master/ludwig/data/dataset/ray.py#L175) we take the size over this result to determine the total length of the dataset. But if the length reported is the length of the window (not the length of the full dataset) then it makes sense it would significantly under-report.
@ShreyaR I see that you report the total dataset size (train + val + test), but wouldn't it be actually correct to only report the train dataset size since that's what's being used to compute the number of epochs based on batch size?

For e.g., I ran the 100MB dataset with and without hyperopt with the following config:

- 100MB dataset has ~ 354K rows in the train set
- Batch size: 32K
- Window size: 16K
- Expected steps per epoch: 354K / 32K ~ 11
- The output says the same thing: 11 steps per epoch, training for 1100 steps, approximately 100 epochs = 11 * 32K = ~ 354K

I feel like this functionality is working correctly, but maybe I'm still missing the crux of the issue
Quick update: 
1. This issue should show up when window bytes size is small enough to create multiple windows 
2. We're going to try and reproduce the issue again since I haven't been able to with what currently exists in master. Will continue investigating once we have a good way of reproducing the issue
@ShreyaR mentioned that > 1 windows are created only when the dataset has more than one partition, so we repartitioned the dataset into more than one partition and that causes the `steps_per_checkpoint` issue to show up (see the example below - it should be 11 steps per epoch, but it says 2 for the 100MB dataset)

```
(BaseWorkerMixin pid=25002) Note: steps_per_checkpoint (was 2000) is now set to the number of steps per epoch: 2.
(BaseWorkerMixin pid=25002) 
(BaseWorkerMixin pid=25002) Training for 200 step(s), approximately 100 epoch(s).
(BaseWorkerMixin pid=25002) Early stopping policy: 10 round(s) of evaluation, or 20 step(s), approximately 10 epoch(s).
(BaseWorkerMixin pid=25002) 
(BaseWorkerMixin pid=25002) Starting with step 0, epoch: 0
```

Commit message: Modify RayDatasetShard length to factor in windowing and dataset partitions

Patch to review:
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
index 47005b7b54f..a234082d1b6 100644
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -17,7 +17,6 @@
 import contextlib
 import copy
 import logging
-from distutils.version import LooseVersion
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
@@ -27,6 +26,7 @@
 import ray
 import torch
 import tqdm
+from packaging import version
 from ray import ObjectRef
 from ray.data.dataset_pipeline import DatasetPipeline
 from ray.data.extensions import TensorDtype
@@ -44,7 +44,8 @@
 from ludwig.utils.torch_utils import get_torch_device, initialize_pytorch
 from ludwig.utils.types import Series
 
-_ray112 = LooseVersion("1.12") <= LooseVersion(ray.__version__) < LooseVersion("1.13")
+_ray112 = version.parse("1.12") <= version.parse(ray.__version__) < version.parse("1.13")
+
 import ray.train as rt  # noqa: E402
 from ray.train.trainer import Trainer  # noqa: E402
 
@@ -90,7 +91,7 @@ def get_trainer_kwargs(use_gpu=None):
         num_workers = int(ray.cluster_resources().get("GPU", 0))
     else:
         # TODO: use placement groups or otherwise spread across nodes
-        node_resources = [node["Resources"] for node in ray.state.nodes()]
+        node_resources = [node["Resources"] for node in ray.nodes()]
         num_workers = len(node_resources)
 
     return dict(
@@ -149,6 +150,7 @@ def _get_df_engine(processor):
 def train_fn(
     executable_kwargs: Dict[str, Any] = None,
     model_ref: ObjectRef = None,  # noqa: F821
+    shards_metadata: Dict[str, Any] = None,
     training_set_metadata: Dict[str, Any] = None,
     features: Dict[str, Dict] = None,
     **kwargs,
@@ -162,6 +164,8 @@ def train_fn(
             rt.get_dataset_shard("train"),
             features,
             training_set_metadata,
+            shards_metadata.get("training_num_partitions"),
+            shards_metadata.get("training_window_status"),
         )
 
         try:
@@ -174,6 +178,8 @@ def train_fn(
                 val_shard,
                 features,
                 training_set_metadata,
+                shards_metadata.get("validation_num_partitions"),
+                shards_metadata.get("validation_window_status"),
             )
 
         try:
@@ -186,6 +192,8 @@ def train_fn(
                 test_shard,
                 features,
                 training_set_metadata,
+                shards_metadata.get("test_num_partitions"),
+                shards_metadata.get("test_window_status"),
             )
 
         model = ray.get(model_ref)
@@ -227,9 +235,7 @@ def tune_batch_size_fn(
 
         pipe = dataset.pipeline(shuffle=False, **data_loader_kwargs)
         train_shard = RayDatasetShard(
-            pipe,
-            features,
-            training_set_metadata,
+            pipe, features, training_set_metadata, dataset.num_partitions, dataset.window_status
         )
 
         device = get_torch_device()
@@ -260,9 +266,7 @@ def tune_learning_rate_fn(
 
         pipe = dataset.pipeline(shuffle=False, **data_loader_kwargs)
         train_shard = RayDatasetShard(
-            pipe,
-            features,
-            training_set_metadata,
+            pipe, features, training_set_metadata, dataset.num_partitions, dataset.window_status
         )
 
         device = get_torch_device()
@@ -340,6 +344,13 @@ def train(
         test_set: Optional[RayDataset] = None,
         **kwargs,
     ):
+
+        dataset = {"train": training_set.pipeline(**self.data_loader_kwargs)}
+        if validation_set is not None:
+            dataset["val"] = validation_set.pipeline(shuffle=False, **self.data_loader_kwargs)
+        if test_set is not None:
+            dataset["test"] = test_set.pipeline(shuffle=False, **self.data_loader_kwargs)
+
         executable_kwargs = self.executable_kwargs
 
         kwargs = {
@@ -348,16 +359,24 @@ def train(
             **kwargs,
         }
 
-        dataset = {"train": training_set.pipeline(**self.data_loader_kwargs)}
-        if validation_set is not None:
-            dataset["val"] = validation_set.pipeline(shuffle=False, **self.data_loader_kwargs)
-        if test_set is not None:
-            dataset["test"] = test_set.pipeline(shuffle=False, **self.data_loader_kwargs)
+        shards_metadata = {
+            "training_num_partitions": training_set.num_partitions,
+            "training_window_status": training_set.window_status,
+            "validation_num_partitions": validation_set.num_partitions if validation_set else 0,
+            "validation_window_status": validation_set.window_status if validation_set else False,
+            "test_num_partitions": validation_set.num_partitions if validation_set else 0,
+            "test_window_status": validation_set.window_status if validation_set else False,
+        }
 
         with self.create_runner() as runner:
             results, self._validation_field, self._validation_metric = runner.run(
                 lambda config: train_fn(**config),
-                config={"executable_kwargs": executable_kwargs, "model_ref": ray.put(self.model), **kwargs},
+                config={
+                    "executable_kwargs": executable_kwargs,
+                    "model_ref": ray.put(self.model),
+                    "shards_metadata": shards_metadata,
+                    **kwargs,
+                },
                 callbacks=[TqdmCallback()],
                 dataset=dataset,
             )[0]
@@ -463,6 +482,7 @@ def legacy_train_fn(
     train_shards: List[DatasetPipeline] = None,
     val_shards: List[DatasetPipeline] = None,
     test_shards: List[DatasetPipeline] = None,
+    shards_metadata: Dict[str, Any] = None,
     **kwargs,
 ):
     # Pin GPU before loading the model to prevent memory leaking onto other devices
@@ -473,6 +493,8 @@ def legacy_train_fn(
         train_shards[hvd.rank()],
         features,
         training_set_metadata,
+        shards_metadata.get("training_num_partitions"),
+        shards_metadata.get("training_window_status"),
     )
 
     val_shard = val_shards[hvd.rank()] if val_shards else None
@@ -481,6 +503,8 @@ def legacy_train_fn(
             val_shard,
             features,
             training_set_metadata,
+            shards_metadata.get("validation_num_partitions"),
+            shards_metadata.get("validation_window_status"),
         )
 
     test_shard = test_shards[hvd.rank()] if test_shards else None
@@ -489,6 +513,8 @@ def legacy_train_fn(
             test_shard,
             features,
             training_set_metadata,
+            shards_metadata.get("test_num_partitions"),
+            shards_metadata.get("test_window_status"),
         )
 
     results = trainer.train(train_shard, val_shard, test_shard, **kwargs)
@@ -516,6 +542,7 @@ def __init__(self, horovod_kwargs, executable_kwargs):
 
     def train(self, model, training_set, validation_set=None, test_set=None, **kwargs):
         workers = self.executor.driver.workers
+
         train_shards = training_set.pipeline().split(n=len(workers), locality_hints=workers, equal=True)
         val_shards = (
             validation_set.pipeline(shuffle=False).split(n=len(workers), locality_hints=workers)
@@ -526,6 +553,15 @@ def train(self, model, training_set, validation_set=None, test_set=None, **kwarg
             test_set.pipeline(shuffle=False).split(n=len(workers), locality_hints=workers) if test_set else None
         )
 
+        shards_metadata = {
+            "training_num_partitions": training_set.num_partitions,
+            "training_window_status": training_set.window_status,
+            "validation_num_partitions": validation_set.num_partitions if validation_set else 0,
+            "validation_window_status": training_set.window_status if validation_set else False,
+            "test_num_partitions": test_set.num_partitions if test_set else 0,
+            "test_window_status": test_set.window_status if test_set else False,
+        }
+
         results = self.executor.execute(
             lambda trainer: legacy_train_fn(
                 trainer,
@@ -535,6 +571,7 @@ def train(self, model, training_set, validation_set=None, test_set=None, **kwarg
                 train_shards,
                 val_shards,
                 test_shards,
+                shards_metadata,
                 **kwargs,
             )
         )
@@ -563,6 +600,7 @@ def eval_fn(
     model_ref: ObjectRef = None,  # noqa: F821
     training_set_metadata: Dict[str, Any] = None,
     features: Dict[str, Dict] = None,
+    shards_metadata: Dict[str, Any] = None,
     **kwargs,
 ):
     # Pin GPU before loading the model to prevent memory leaking onto other devices
@@ -574,6 +612,8 @@ def eval_fn(
             rt.get_dataset_shard("eval"),
             features,
             training_set_metadata,
+            shards_metadata.get("evaluation_num_partitions"),
+            shards_metadata.get("evaluation_window_status"),
         )
 
         model = ray.get(model_ref)
@@ -666,6 +706,10 @@ def batch_evaluation(self, dataset: RayDataset, collect_predictions: bool = Fals
                 **self.predictor_kwargs,
                 "collect_predictions": False,
             }
+            shards_metadata = {
+                "evaluation_num_partitions": dataset.num_partitions,
+                "evaluation_window_status": dataset.window_status,
+            }
             eval_stats, _ = runner.run(
                 lambda config: eval_fn(**config),
                 config={
@@ -673,6 +717,7 @@ def batch_evaluation(self, dataset: RayDataset, collect_predictions: bool = Fals
                     "model_ref": ray.put(self.model),
                     "training_set_metadata": dataset.training_set_metadata,
                     "features": dataset.features,
+                    "shards_metadata": shards_metadata,
                     **kwargs,
                 },
                 dataset=datasets,
diff --git a/ludwig/data/dataset/ray.py b/ludwig/data/dataset/ray.py
index 165e2ab6f66..7e751f7125f 100644
--- a/ludwig/data/dataset/ray.py
+++ b/ludwig/data/dataset/ray.py
@@ -17,13 +17,13 @@
 import math
 import queue
 import threading
-from distutils.version import LooseVersion
 from functools import lru_cache
 from typing import Any, Dict, Iterator, Optional, Union
 
 import numpy as np
 import pandas as pd
 import ray
+from packaging import version
 from pyarrow.fs import FSSpecHandler, PyFileSystem
 from ray.data import read_parquet
 from ray.data.dataset_pipeline import DatasetPipeline
@@ -38,8 +38,7 @@
 from ludwig.utils.misc_utils import get_proc_features
 from ludwig.utils.types import DataFrame
 
-_ray113 = LooseVersion(ray.__version__) == LooseVersion("1.13.0")
-
+_ray113 = version.parse(ray.__version__) == version.parse("1.13.0")
 
 _SCALAR_TYPES = {BINARY, CATEGORY, NUMBER}
 
@@ -65,6 +64,7 @@ def __init__(
         self.training_set_metadata = training_set_metadata
         self.data_hdf5_fp = training_set_metadata.get(DATA_TRAIN_HDF5_FP)
         self.data_parquet_fp = training_set_metadata.get(DATA_TRAIN_PARQUET_FP)
+        self.dataset_window_status = False
 
         # TODO ray 1.8: convert to Tensors before shuffle
         # def to_tensors(df: pd.DataFrame) -> pd.DataFrame:
@@ -96,6 +96,7 @@ def pipeline(
             pipe = self.ds.repeat()
         else:
             pipe = self.ds.window(bytes_per_window=window_size_bytes).repeat()
+            self.dataset_window_status = True
         if shuffle:
             pipe = pipe.random_shuffle_each_window()
         return pipe
@@ -117,6 +118,14 @@ def __len__(self):
     def size(self):
         return len(self)
 
+    @property
+    def num_partitions(self):
+        return self.ds.num_blocks()
+
+    @property
+    def window_status(self):
+        return self.dataset_window_status
+
     def to_df(self):
         return self.df_engine.from_ray_dataset(self.ds)
 
@@ -153,10 +162,15 @@ def __init__(
         dataset_shard: DatasetPipeline,
         features: Dict[str, Dict],
         training_set_metadata: Dict[str, Any],
+        num_dataset_partitions: int,
+        window_status: bool,
     ):
         self.dataset_shard = dataset_shard
         self.features = features
         self.training_set_metadata = training_set_metadata
+        self.num_dataset_partitions = num_dataset_partitions
+        self.window_status = window_status
+
         self.dataset_iter = dataset_shard.iter_datasets()
 
     @contextlib.contextmanager
@@ -172,7 +186,10 @@ def initialize_batcher(self, batch_size=128, should_shuffle=True, seed=0, ignore
     @lru_cache(1)
     def __len__(self):
         # TODO(travis): find way to avoid calling this, as it's expensive
-        return next(self.dataset_iter).count()
+        next_iteration_length = next(self.dataset_iter).count()
+        if self.num_dataset_partitions > 1 and self.window_status:
+            return next_iteration_length * self.num_dataset_partitions
+        return next_iteration_length
 
     @property
     def size(self):


Code review comments:
diff_hunk: @@ -172,7 +186,10 @@ def initialize_batcher(self, batch_size=128, should_shuffle=True, seed=0, ignore
     @lru_cache(1)
     def __len__(self):
         # TODO(travis): find way to avoid calling this, as it's expensive
-        return next(self.dataset_iter).count()
+        next_iteration_length = next(self.dataset_iter).count()
+        if self.num_dataset_partitions > 1 and self.window_status:
+            return next_iteration_length * self.num_dataset_partitions
path: ludwig/data/dataset/ray.py
text: @user1:
My concern with this solution is that this might not be accurate -- two partitions may not have the same number of examples.