Repository: ludwig-ai/ludwig
Pull request: Make Horovod an optional dependency when using Ray
Closes #3161.

Falls back to DDP if Horovod is not installed. We may choose to make DDP the default in v0.8.

Problem statement:
auto_train shouldn't require horovod
**Describe the bug**
Using the autoML feature in auto_train seems to require horovod, but it should be possible to leverage autoML on a single machine without requiring horovod (which is very difficult to install).  If this is indeed by design, it should be called out as a requirement in the auto_train documentation.

the error is:
ModuleNotFoundError: Horovod isn't installed. To install Horovod with PyTorch support, run 'pip install 'horovod[pytorch]''. To install Horovod with TensorFlow support, run 'pip install 'horovod[tensorflow]''.

**To Reproduce**
Sample code:

```
from ludwig.automl import auto_train
from ludwig.datasets import mushroom_edibility
from ludwig.utils.dataset_utils import get_repeatable_train_val_test_split

mushroom_df = mushroom_edibility.load()
mushroom_edibility_df = get_repeatable_train_val_test_split(mushroom_df, 'class', random_seed=42)

auto_train_results = auto_train(
    dataset=mushroom_edibility_df,
    target='class',
    time_limit_s=120,
    tune_for_memory=False,
    user_config={'preprocessing': {'split': {'column': 'split', 'type': 'fixed'}}},
)

pprint.pprint(auto_train_results)
```

**Expected behavior**
auto_train should run successfully on a machine without Horovod installed.

**Environment (please complete the following information):**

- OS: Linux
- Python version: 3.8
- Ludwig version: 0.7 installed via pip, along with dask, hummingbird-ml, and lightgbm

thanks

Hints from issue discussion:
Hey @skunkwerk, thanks for the feature request! Yes, this should be pretty doable. We should be able to have a fix in the next day or so, and include it in v0.7.1.

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/ludwig/backend/_ray210_compat.py b/ludwig/backend/_ray210_compat.py
index 9b48447bb94..a05c64f3e20 100644
--- a/ludwig/backend/_ray210_compat.py
+++ b/ludwig/backend/_ray210_compat.py
@@ -5,9 +5,6 @@
 
 import ray
 from ray.air.config import RunConfig
-from ray.air.result import Result
-from ray.train.base_trainer import TrainingFailedError
-from ray.train.horovod import HorovodTrainer
 from ray.tune.execution.trial_runner import _ResumeConfig
 from ray.tune.impl.tuner_internal import TunerInternal
 from ray.tune.trainable import Trainable
@@ -143,36 +140,3 @@ def _expected_utilization(self, cpus_per_trial, cpus_total):
             )
         )
         return (actual_concurrency * cpus_per_trial) / (cpus_total + 0.001)
-
-
-class HorovodTrainerRay210(HorovodTrainer):
-    """HACK(geoffrey): This is a temporary fix to support Ray 2.1.0.
-
-    Specifically, this Trainer ensures that TunerRay210 is called by the class.
-    For more details, see TunerRay210.
-    """
-
-    def fit(self) -> Result:
-        """Runs training.
-
-        Returns:
-            A Result object containing the training result.
-
-        Raises:
-            TrainingFailedError: If any failures during the execution of
-            ``self.as_trainable()``.
-        """
-        from ray.tune.error import TuneError
-
-        trainable = self.as_trainable()
-
-        tuner = TunerRay210(trainable=trainable, run_config=self.run_config)
-        result_grid = tuner.fit()
-        assert len(result_grid) == 1
-        try:
-            result = result_grid[0]
-            if result.error:
-                raise result.error
-        except TuneError as e:
-            raise TrainingFailedError from e
-        return result
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
index 371a5908493..7942e5b21d8 100644
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -32,12 +32,11 @@
 from ray.air import session
 from ray.air.checkpoint import Checkpoint
 from ray.air.config import DatasetConfig, RunConfig, ScalingConfig
-from ray.train.horovod import HorovodTrainer
 from ray.train.torch import TorchCheckpoint
 from ray.util.dask import ray_dask_get
 from ray.util.placement_group import placement_group, remove_placement_group
 
-from ludwig.distributed import get_current_dist_strategy, get_dist_strategy
+from ludwig.distributed import get_current_dist_strategy, get_default_strategy_name, get_dist_strategy
 from ludwig.utils.batch_size_tuner import BatchSizeEvaluator
 
 if TYPE_CHECKING:
@@ -64,16 +63,9 @@
 from ludwig.utils.torch_utils import get_torch_device, initialize_pytorch
 from ludwig.utils.types import DataFrame, Series
 
-_ray220 = version.parse(ray.__version__) >= version.parse("2.2.0")
 _ray230 = version.parse(ray.__version__) >= version.parse("2.3.0")
 
 
-if not _ray220:
-    from ludwig.backend._ray210_compat import HorovodTrainerRay210
-else:
-    HorovodTrainerRay210 = None
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -98,7 +90,7 @@ def get_trainer_kwargs(**kwargs) -> TrainerConfigDict:
     else:
         num_workers = _num_nodes()
 
-    strategy = kwargs.pop("strategy", "horovod")
+    strategy = kwargs.pop("strategy", get_default_strategy_name())
     backend = get_dist_strategy(strategy).get_ray_trainer_backend(**kwargs)
 
     # Remove params used by strategy but not the trainer here
@@ -106,6 +98,7 @@ def get_trainer_kwargs(**kwargs) -> TrainerConfigDict:
 
     defaults = dict(
         backend=backend,
+        strategy=strategy,
         num_workers=num_workers,
         use_gpu=use_gpu,
         resources_per_worker={
@@ -309,6 +302,7 @@ class RayAirRunner:
     def __init__(self, trainer_kwargs: Dict[str, Any]) -> None:
         trainer_kwargs = copy.copy(trainer_kwargs)
         self.backend_config = trainer_kwargs.pop("backend", None)
+        self.strategy = trainer_kwargs.pop("strategy", get_default_strategy_name())
 
         if "max_retries" in trainer_kwargs:
             logger.warning("`max_retries` is no longer supported as a trainer argument in Ray backend. Ignoring it.")
@@ -369,15 +363,15 @@ def run(
         stream_window_size: Dict[str, Union[None, float]],
         callbacks: List[Any] = [],
     ) -> Tuple[Dict, TorchCheckpoint]:
-        trainer_cls = HorovodTrainerRay210 if not _ray220 else HorovodTrainer
+        trainer_cls, kwargs = get_dist_strategy(self.strategy).get_trainer_cls(self.backend_config)
         trainer = trainer_cls(
             train_loop_per_worker=train_loop_per_worker,
             train_loop_config=config,
-            horovod_config=self.backend_config,
             datasets=dataset,
             scaling_config=self.scaling_config,
             dataset_config=self._get_dataset_configs(dataset, stream_window_size, data_loader_kwargs),
             run_config=RunConfig(callbacks=callbacks, verbose=0),
+            **kwargs,
         )
         return trainer.fit()
 
diff --git a/ludwig/distributed/__init__.py b/ludwig/distributed/__init__.py
index 6644c23a08d..dfc79c4ee88 100644
--- a/ludwig/distributed/__init__.py
+++ b/ludwig/distributed/__init__.py
@@ -46,3 +46,15 @@ def get_current_dist_strategy(allow_local=True) -> Type[DistributedStrategy]:
 
 def get_dist_strategy(name: str) -> Type[DistributedStrategy]:
     return STRATEGIES[name]()
+
+
+def get_default_strategy_name() -> str:
+    try:
+        # Use horovod by default for now if it's available
+        load_horovod()
+        return "horovod"
+    except ImportError:
+        pass
+
+    # Fallback to DDP if not
+    return "ddp"
diff --git a/ludwig/distributed/_ray_210_compat.py b/ludwig/distributed/_ray_210_compat.py
new file mode 100644
index 00000000000..59dd2962b5f
--- /dev/null
+++ b/ludwig/distributed/_ray_210_compat.py
@@ -0,0 +1,38 @@
+from ray.air.result import Result
+from ray.train.base_trainer import TrainingFailedError
+from ray.train.horovod import HorovodTrainer
+
+from ludwig.backend._ray210_compat import TunerRay210
+
+
+class HorovodTrainerRay210(HorovodTrainer):
+    """HACK(geoffrey): This is a temporary fix to support Ray 2.1.0.
+
+    Specifically, this Trainer ensures that TunerRay210 is called by the class.
+    For more details, see TunerRay210.
+    """
+
+    def fit(self) -> Result:
+        """Runs training.
+
+        Returns:
+            A Result object containing the training result.
+
+        Raises:
+            TrainingFailedError: If any failures during the execution of
+            ``self.as_trainable()``.
+        """
+        from ray.tune.error import TuneError
+
+        trainable = self.as_trainable()
+
+        tuner = TunerRay210(trainable=trainable, run_config=self.run_config)
+        result_grid = tuner.fit()
+        assert len(result_grid) == 1
+        try:
+            result = result_grid[0]
+            if result.error:
+                raise result.error
+        except TuneError as e:
+            raise TrainingFailedError from e
+        return result
diff --git a/ludwig/distributed/base.py b/ludwig/distributed/base.py
index 3bcd4676e69..08d534439f8 100644
--- a/ludwig/distributed/base.py
+++ b/ludwig/distributed/base.py
@@ -1,8 +1,10 @@
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional, Tuple, Type
 
 import torch
+from ray.train.backend import BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
 from torch import nn
 from torch.optim import Optimizer
 
@@ -87,6 +89,11 @@ def gather_all_tensors_fn(cls) -> Optional[Callable]:
     def get_ray_trainer_backend(cls, **kwargs) -> Optional[Any]:
         pass
 
+    @classmethod
+    @abstractmethod
+    def get_trainer_cls(cls, backend_config: BackendConfig) -> Tuple[Type[DataParallelTrainer], Dict[str, Any]]:
+        pass
+
     @abstractmethod
     def shutdown(self):
         pass
@@ -162,5 +169,9 @@ def gather_all_tensors_fn(cls) -> Optional[Callable]:
     def get_ray_trainer_backend(cls, **kwargs) -> Optional[Any]:
         return None
 
+    @classmethod
+    def get_trainer_cls(cls, backend_config: BackendConfig) -> Tuple[Type[DataParallelTrainer], Dict[str, Any]]:
+        raise ValueError("Cannot construct a trainer from a local strategy.")
+
     def shutdown(self):
         pass
diff --git a/ludwig/distributed/ddp.py b/ludwig/distributed/ddp.py
index 29c5810e430..0cbaab99482 100644
--- a/ludwig/distributed/ddp.py
+++ b/ludwig/distributed/ddp.py
@@ -1,10 +1,13 @@
 import contextlib
 import logging
 import socket
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple, Type
 
 import torch
 import torch.distributed as dist
+from ray.train.backend import BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.torch import TorchTrainer
 from torch import nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
@@ -84,6 +87,10 @@ def get_ray_trainer_backend(cls, **kwargs) -> Optional[Any]:
 
         return TorchConfig()
 
+    @classmethod
+    def get_trainer_cls(cls, backend_config: BackendConfig) -> Tuple[Type[DataParallelTrainer], Dict[str, Any]]:
+        return TorchTrainer, dict(torch_config=backend_config)
+
     def shutdown(self):
         # TODO(travis): currently Ray handles this for us, but is subject to hangs if one of the workers raises an
         # exception and the other makes a collective op. We should figure out a way to make this safe to call
diff --git a/ludwig/distributed/horovod.py b/ludwig/distributed/horovod.py
index e4c5eb0cd12..9fbbf3a9dec 100644
--- a/ludwig/distributed/horovod.py
+++ b/ludwig/distributed/horovod.py
@@ -1,16 +1,23 @@
 import contextlib
 import logging
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 import horovod.torch as hvd
+import ray
 import torch
 from horovod.torch.optimizer import _DistributedOptimizer
+from packaging import version
+from ray.train.backend import BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.horovod import HorovodTrainer
 from torch import nn
 from torch.optim import Optimizer
 
 from ludwig.distributed.base import DistributedStrategy
 from ludwig.utils.horovod_utils import gather_all_tensors, is_distributed_available
 
+_ray220 = version.parse(ray.__version__) >= version.parse("2.2.0")
+
 
 class HorovodStrategy(DistributedStrategy):
     def __init__(self):
@@ -81,5 +88,14 @@ def get_ray_trainer_backend(cls, nics: Optional[List[str]] = None, **kwargs) ->
             nics = set(nics)
         return HorovodConfig(nics=nics)
 
+    @classmethod
+    def get_trainer_cls(cls, backend_config: BackendConfig) -> Tuple[Type[DataParallelTrainer], Dict[str, Any]]:
+        if not _ray220:
+            from ludwig.distributed._ray_210_compat import HorovodTrainerRay210
+
+            return HorovodTrainerRay210, dict(horovod_config=backend_config)
+
+        return HorovodTrainer, dict(horovod_config=backend_config)
+
     def shutdown(self):
         hvd.shutdown()


Code review comments:
diff_hunk: @@ -1,8 +1,10 @@
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional, Tuple, Type
 
 import torch
+from ray.train.backend import BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
path: ludwig/distributed/base.py
text: @user1:
Probably need to gate these imports?

@author:
Good catch!