From d3f2dbe64ff3cfd1106c57c83bf00ae652a16a3a Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Fri, 26 Jan 2024 02:09:38 +0800
Subject: [PATCH] [dask] Add seed to demos. (#10009)

---
 demo/dask/cpu_training.py | 5 +++--
 demo/dask/gpu_training.py | 7 +++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/demo/dask/cpu_training.py b/demo/dask/cpu_training.py
index 0f3316741..00453740f 100644
--- a/demo/dask/cpu_training.py
+++ b/demo/dask/cpu_training.py
@@ -14,8 +14,9 @@ def main(client):
     # generate some random data for demonstration
     m = 100000
     n = 100
-    X = da.random.random(size=(m, n), chunks=100)
-    y = da.random.random(size=(m,), chunks=100)
+    rng = da.random.default_rng(1)
+    X = rng.normal(size=(m, n))
+    y = X.sum(axis=1)
 
     # DaskDMatrix acts like normal DMatrix, works as a proxy for local
     # DMatrix scatter around workers.
diff --git a/demo/dask/gpu_training.py b/demo/dask/gpu_training.py
index fd5b35bf3..6096b87fc 100644
--- a/demo/dask/gpu_training.py
+++ b/demo/dask/gpu_training.py
@@ -2,6 +2,7 @@
 Example of training with Dask on GPU
 ====================================
 """
+import cupy as cp
 import dask_cudf
 from dask import array as da
 from dask import dataframe as dd
@@ -72,10 +73,12 @@ if __name__ == "__main__":
     with LocalCUDACluster(n_workers=2, threads_per_worker=4) as cluster:
         with Client(cluster) as client:
             # generate some random data for demonstration
+            rng = da.random.default_rng(1)
+
             m = 100000
             n = 100
-            X = da.random.random(size=(m, n), chunks=10000)
-            y = da.random.random(size=(m,), chunks=10000)
+            X = rng.normal(size=(m, n))
+            y = X.sum(axis=1)
 
             print("Using DaskQuantileDMatrix")
             from_ddqdm = using_quantile_device_dmatrix(client, X, y)