From f7c641d980a1a40df5395cd9b12589834860879a Mon Sep 17 00:00:00 2001
From: James Lamb <jaylamb20@gmail.com>
Date: Sun, 2 Mar 2025 09:08:46 -0600
Subject: [PATCH] [ci] [dask] test `lightgbm.dask` on macOS (#6677)

* [ci] [Dask] enable Dask tests on macOS

* fix windows condition

* fix merge conflicts
---
 .vsts-ci.yml                           |  4 ++++
 docs/Parallel-Learning-Guide.rst       |  2 +-
 tests/python_package_test/test_dask.py | 11 ++++++-----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/.vsts-ci.yml b/.vsts-ci.yml
index 0f44f4d4365e..883b519f5f9f 100644
--- a/.vsts-ci.yml
+++ b/.vsts-ci.yml
@@ -254,6 +254,9 @@ jobs:
       - script: |
           git clean -d -f -x
         displayName: 'Clean source directory'
+      # LGBM_SKIP_DASK_TESTS=true is set below only because running the tests under emulation is so slow...
+      # in theory, 'lightgbm.dask' should work without issue on aarch64 Linux systems.
+      # That could probably be removed as part of https://github.com/microsoft/LightGBM/issues/6788
       - script: |
           cat > docker-script.sh <<EOF
           export CONDA=\$HOME/miniforge
@@ -273,6 +276,7 @@ jobs:
             --env BUILD_ARTIFACTSTAGINGDIRECTORY=$BUILD_ARTIFACTSTAGINGDIRECTORY \
             --env BUILD_DIRECTORY=$BUILD_DIRECTORY \
             --env COMPILER=$COMPILER \
+            --env LGBM_SKIP_DASK_TESTS=true \
             --env METHOD=$METHOD \
             --env OS_NAME=linux \
             --env PRODUCES_ARTIFACTS=$PRODUCES_ARTIFACTS \
diff --git a/docs/Parallel-Learning-Guide.rst b/docs/Parallel-Learning-Guide.rst
index a3d0a9d14f89..1ac8bd1ea782 100644
--- a/docs/Parallel-Learning-Guide.rst
+++ b/docs/Parallel-Learning-Guide.rst
@@ -64,7 +64,7 @@ LightGBM's Python-package supports distributed learning via `Dask`_. This integr
 
 .. warning::
 
-    Dask integration is only tested on Linux.
+    Dask integration is only tested on macOS and Linux.
 
 Dask Examples
 '''''''''''''
diff --git a/tests/python_package_test/test_dask.py b/tests/python_package_test/test_dask.py
index ad13734187d8..df9eada87a7a 100644
--- a/tests/python_package_test/test_dask.py
+++ b/tests/python_package_test/test_dask.py
@@ -5,7 +5,6 @@
 import socket
 from itertools import groupby
 from os import getenv
-from platform import machine
 from sys import platform
 from urllib.parse import urlparse
 
@@ -16,10 +15,8 @@
 
 from .utils import sklearn_multiclass_custom_objective
 
-if not platform.startswith("linux"):
-    pytest.skip("lightgbm.dask is currently supported in Linux environments", allow_module_level=True)
-if machine() != "x86_64":
-    pytest.skip("lightgbm.dask tests are currently skipped on some architectures like arm64", allow_module_level=True)
+if platform in {"cygwin", "win32"}:
+    pytest.skip("lightgbm.dask is not currently supported on Windows", allow_module_level=True)
 if not lgb.compat.DASK_INSTALLED:
     pytest.skip("Dask is not installed", allow_module_level=True)
 
@@ -58,6 +55,10 @@
     pytest.mark.skipif(getenv("TASK", "") == "mpi", reason="Fails to run with MPI interface"),
     pytest.mark.skipif(getenv("TASK", "") == "gpu", reason="Fails to run with GPU interface"),
     pytest.mark.skipif(getenv("TASK", "") == "cuda", reason="Fails to run with CUDA interface"),
+    pytest.mark.skipif(
+        getenv("LGBM_SKIP_DASK_TESTS", "") == "true",
+        reason="Skipping lightgbm.dask tests (found env variable LGBM_SKIP_DASK_TESTS=true)",
+    ),
 ]