From a91fb61d394da92ef851a631e20c72e9617a9f7f Mon Sep 17 00:00:00 2001
From: Pang Wu <pang.wu@samsara.com>
Date: Sun, 8 Dec 2024 16:03:32 -0800
Subject: [PATCH] test against multiple ray version

---
 .github/workflows/ray_nightly_test.yml | 11 ++++-------
 .github/workflows/raydp.yml            |  5 +++--
 python/raydp/tf/estimator.py           |  2 +-
 3 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/ray_nightly_test.yml b/.github/workflows/ray_nightly_test.yml
index e49c8c69..465605c4 100644
--- a/.github/workflows/ray_nightly_test.yml
+++ b/.github/workflows/ray_nightly_test.yml
@@ -31,7 +31,7 @@ jobs:
     strategy:
       matrix:
         os: [ ubuntu-latest ]
-        python-version: [3.8, 3.9, 3.10.14]
+        python-version: [3.9, 3.10.14]
         spark-version: [3.2.4, 3.3.2, 3.4.0, 3.5.0]
 
     runs-on: ${{ matrix.os }}
@@ -82,15 +82,12 @@ jobs:
             pip install torch
           fi
           case $PYTHON_VERSION in
-            3.7)
-              pip install "ray[train] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl"
-              ;;
-            3.8)
-              pip install "ray[train] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp38-cp38-manylinux2014_x86_64.whl"
-              ;;
             3.9)
               pip install "ray[train] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl"
               ;;
+            3.10.14)
+              pip install "ray[train] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp310-cp310-manylinux2014_x86_64.whl"
+              ;;
           esac
           pip install pyarrow==6.0.1 pytest koalas tensorflow==2.13.1 tabulate grpcio-tools wget
           pip install "xgboost_ray[default]<=0.1.13"
diff --git a/.github/workflows/raydp.yml b/.github/workflows/raydp.yml
index d4b8f6a2..a1b10916 100644
--- a/.github/workflows/raydp.yml
+++ b/.github/workflows/raydp.yml
@@ -33,8 +33,9 @@ jobs:
     strategy:
       matrix:
         os: [ ubuntu-latest ]
-        python-version: [3.8, 3.9, 3.10.14]
+        python-version: [3.9, 3.10.14]
         spark-version: [3.2.4, 3.3.2, 3.4.0, 3.5.0]
+        ray-version: [2.34.0, 2.40.0]
 
     runs-on: ${{ matrix.os }}
 
@@ -82,7 +83,7 @@ jobs:
           else
             pip install torch
           fi
-          pip install pyarrow==6.0.1 ray[train] pytest tensorflow==2.13.1 tabulate grpcio-tools wget
+          pip install pyarrow==6.0.1 "ray[train]==${{ matrix.ray-version }}" pytest tensorflow==2.13.1 tabulate grpcio-tools wget
           pip install "xgboost_ray[default]<=0.1.13"
           pip install "xgboost<=2.0.3"
           pip install torchmetrics
diff --git a/python/raydp/tf/estimator.py b/python/raydp/tf/estimator.py
index 88bbb1eb..9fa21f1d 100644
--- a/python/raydp/tf/estimator.py
+++ b/python/raydp/tf/estimator.py
@@ -229,7 +229,7 @@ def fit(self,
                     # pylint: disable=E1123,E1120
                     if version.parse(ray.__version__) >= version.parse("2.39.0"):
                         preprocessor = Concatenator(
-                            columns=self._feature_columns - label_cols,
+                            columns=[col for col in self._feature_columns if col not in label_cols],
                             output_column_name="features",
                         )
                     else: