From 6286815a6d28facb29822ce4881b9eba2cfbd511 Mon Sep 17 00:00:00 2001
From: Travis Addair <tgaddair@gmail.com>
Date: Fri, 10 Jun 2022 07:46:15 -0700
Subject: [PATCH] Added ray 1.13.0 and nightly wheel tests to CI (#2128)

---
 .github/workflows/pytest.yml | 11 ++++++++---
 ludwig/backend/ray.py        |  2 +-
 ludwig/data/dataset/ray.py   | 10 ++++++----
 ludwig/hyperopt/sampling.py  |  6 +++---
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 608e45357bc..7ac4c110410 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -32,11 +32,11 @@ jobs:
           - python-version: 3.8
             pytorch-version: 1.11.0
             torchscript-version: 1.10.2
-            ray-version: 1.12.0
+            ray-version: 1.13.0
           - python-version: 3.9
             pytorch-version: nightly
             torchscript-version: 1.10.2
-            ray-version: 1.12.0
+            ray-version: nightly
     env:
       PYTORCH: ${{ matrix.pytorch-version }}
       MARKERS: ${{ matrix.test-markers }}
@@ -102,7 +102,12 @@ jobs:
           pip install protobuf==3.20.1 # https://github.com/databrickslabs/dbx/issues/257
 
           if [ "$MARKERS" == "distributed" ]; then
-            pip install ray==$RAY_VERSION
+            if [ "$RAY_VERSION" == "nightly" ]; then
+              # NOTE: hardcoded for python 3.9 on Linux
+              pip install https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl
+            else
+              pip install ray==$RAY_VERSION
+            fi
             ray_expected=$(python -c "import ray; print(ray.__version__)")
           fi
 
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
index b8d39840ace..47005b7b54f 100644
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -44,7 +44,7 @@
 from ludwig.utils.torch_utils import get_torch_device, initialize_pytorch
 from ludwig.utils.types import Series
 
-_ray112 = LooseVersion(ray.__version__) >= LooseVersion("1.12")
+_ray112 = LooseVersion("1.12") <= LooseVersion(ray.__version__) < LooseVersion("1.13")
 import ray.train as rt  # noqa: E402
 from ray.train.trainer import Trainer  # noqa: E402
 
diff --git a/ludwig/data/dataset/ray.py b/ludwig/data/dataset/ray.py
index 1a59c769ed2..165e2ab6f66 100644
--- a/ludwig/data/dataset/ray.py
+++ b/ludwig/data/dataset/ray.py
@@ -38,7 +38,7 @@
 from ludwig.utils.misc_utils import get_proc_features
 from ludwig.utils.types import DataFrame
 
-_ray112 = LooseVersion(ray.__version__) >= LooseVersion("1.12")
+_ray113 = LooseVersion(ray.__version__) == LooseVersion("1.13.0")
 
 
 _SCALAR_TYPES = {BINARY, CATEGORY, NUMBER}
@@ -83,10 +83,12 @@ def pipeline(
             window_size_bytes: If not None, windowing is enabled and this parameter specifies the window size in bytes
                     for the dataset.
         """
-        if not fully_executed and not _ray112:
-            raise ValueError(f"Cannot set fully_execute=False in ray {ray.__version__}")
+        if fully_executed:
+            if _ray113:
+                # Workaround for: https://github.com/ray-project/ray/issues/25643
+                # TODO(travis): remove after 1.13.1
+                self.ds = self.ds.map_batches(lambda x: x, batch_size=None)
 
-        if fully_executed and _ray112:
             # set instance state so calls to __len__ will also use the fully_executed version
             self.ds = self.ds.fully_executed()
 
diff --git a/ludwig/hyperopt/sampling.py b/ludwig/hyperopt/sampling.py
index 3c1baefde2f..cd795265c65 100644
--- a/ludwig/hyperopt/sampling.py
+++ b/ludwig/hyperopt/sampling.py
@@ -23,14 +23,14 @@
 try:
     from ray import tune
     from ray.tune.schedulers.resource_changing_scheduler import (
-        evenly_distribute_cpus_gpus,
+        DistributeResources,
         PlacementGroupFactory,
         ResourceChangingScheduler,
     )
 
     _HAS_RAY_TUNE = True
 except ImportError:
-    evenly_distribute_cpus_gpus = None
+    logging.exception("ImportError (sampling.py) failed to import ray tune")
     _HAS_RAY_TUNE = False
 
 
@@ -41,7 +41,7 @@ def ray_resource_allocation_function(
     scheduler: "ResourceChangingScheduler",
 ):
     """Determine resources to allocate to running trials."""
-    pgf = evenly_distribute_cpus_gpus(trial_runner, trial, result, scheduler)
+    pgf = DistributeResources(trial_runner, trial, result, scheduler)
     # restore original base trial resources
 
     # create bundles