[dask] fix teardown issues in Dask tests (fixes #3829) (#3869)

* [dask] reduce teardown erros in Dask tests

* azure

* show logs

* try again

* more

* submodules

* try a bunch of sdist tests

* empty commit

* empty commit

* empty commit

* use sh-ubuntu

* 10 sdist tasks

* stuff

* empty commit

* empty commit

* empty commit

* empty commit

* try bdist

* empty commit

* empty commit

* empty commit

* empty commit

* py37

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* python 3.8

* empty commit

* empty commit

* empty commit

* empty commit

* empty commit

* cuda config

* comment out cuda again

* setting timeout

* put client close in the right place

* uncomment CI, make timeout 60

tests/python_package_test/test_dask.py

@@ -27,6 +27,10 @@ from sklearn.utils import check_random_state
 from .utils import make_ranking
 
 
+# time, in seconds, to wait for the Dask client to close. Used to avoid teardown errors
+# see https://distributed.dask.org/en/latest/api.html#distributed.Client.close
+CLIENT_CLOSE_TIMEOUT = 60
+
 data_output = ['array', 'scipy_csr_matrix', 'dataframe']
 data_centers = [[[-4, -4], [4, 4]], [[-4, -4], [4, 4], [-4, 4]]]
 group_sizes = [5, 5, 5, 10, 10, 10, 20, 20, 20, 50, 50]
@@ -172,7 +176,7 @@ def test_classifier(output, centers, client, listen_port):
     assert_eq(p1_local, p2)
     assert_eq(y, p1_local)
 
-    client.close()
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
 
 
 @pytest.mark.parametrize('output', data_output)
@@ -227,6 +231,8 @@ def test_classifier_pred_contrib(output, centers, client, listen_port):
             base_value_col = num_features * (i + 1) + i
             assert len(np.unique(preds_with_contrib[:, base_value_col]) == 1)
 
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
+
 
 def test_training_does_not_fail_on_port_conflicts(client):
     _, _, _, dX, dy, dw = _create_data('classification', output='array')
@@ -249,7 +255,7 @@ def test_training_does_not_fail_on_port_conflicts(client):
             )
             assert dask_classifier.booster_
 
-    client.close()
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
 
 
 @pytest.mark.parametrize('output', data_output)
@@ -292,7 +298,7 @@ def test_regressor(output, client, listen_port):
     assert_eq(y, p2, rtol=1., atol=50.)
     assert_eq(p1, p1_local)
 
-    client.close()
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
 
 
 @pytest.mark.parametrize('output', data_output)
@@ -328,6 +334,8 @@ def test_regressor_pred_contrib(output, client, listen_port):
     assert preds_with_contrib.shape[1] == num_features + 1
     assert preds_with_contrib.shape == local_preds_with_contrib.shape
 
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
+
 
 @pytest.mark.parametrize('output', data_output)
 @pytest.mark.parametrize('alpha', [.1, .5, .9])
@@ -362,7 +370,7 @@ def test_regressor_quantile(output, client, listen_port, alpha):
     np.testing.assert_allclose(q1, alpha, atol=0.2)
     np.testing.assert_allclose(q2, alpha, atol=0.2)
 
-    client.close()
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
 
 
 @pytest.mark.parametrize('output', ['array', 'dataframe'])
@@ -404,7 +412,7 @@ def test_ranker(output, client, listen_port, group):
     assert spearmanr(rnkvec_dask, rnkvec_local).correlation > 0.75
     assert_eq(rnkvec_dask, rnkvec_dask_local)
 
-    client.close()
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
 
 
 def test_find_open_port_works():
@@ -445,6 +453,8 @@ def test_warns_and_continues_on_unrecognized_tree_learner(client):
 
     assert dask_regressor.fitted_
 
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
+
 
 def test_warns_but_makes_no_changes_for_feature_or_voting_tree_learner(client):
     X = da.random.random((1e3, 10))
@@ -463,6 +473,8 @@ def test_warns_but_makes_no_changes_for_feature_or_voting_tree_learner(client):
         assert dask_regressor.fitted_
         assert dask_regressor.get_params()['tree_learner'] == tree_learner
 
+    client.close(timeout=CLIENT_CLOSE_TIMEOUT)
+
 
 @gen_cluster(client=True, timeout=None)
 def test_errors(c, s, a, b):