Elastic failure handling

.github/workflows/test.yaml

@@ -5,6 +5,7 @@ on: [push, pull_request]
 jobs:
   test_lint:
     runs-on: ubuntu-latest
+    timeout-minutes: 3
 
     steps:
     - uses: actions/checkout@v2
@@ -22,8 +23,8 @@ jobs:
         ./format.sh --all
 
   test_linux_ray_master:
-
     runs-on: ubuntu-latest
+    timeout-minutes: 10
 
     steps:
     - uses: actions/checkout@v2
@@ -51,15 +52,16 @@ jobs:
         echo "running smoke test on benchmark_cpu_gpu.py" && python release/benchmark_cpu_gpu.py 2 10 20 --smoke-test
         popd
         pushd examples/
-        echo "running simple.py" && python simple.py
+        ray stop || true
+        echo "running simple.py" && python simple.py --smoke-test
         echo "running simple_predict.py" && python simple_predict.py
-        echo "running simple_tune.py" && python simple_tune.py
+        echo "running simple_tune.py" && python simple_tune.py --smoke-test
         echo "running train_on_test_data.py" && python train_on_test_data.py --smoke-test
       # for f in *.py; do echo "running $f" && python "$f" || exit 1 ; done
 
   test_linux_ray_release:
-
     runs-on: ubuntu-latest
+    timeout-minutes: 10
 
     steps:
     - uses: actions/checkout@v2
@@ -87,8 +89,9 @@ jobs:
         echo "running smoke test on benchmark_cpu_gpu.py" && python release/benchmark_cpu_gpu.py 2 10 20 --smoke-test
         popd
         pushd examples/
-        echo "running simple.py" && python simple.py
+        ray stop || true
+        echo "running simple.py" && python simple.py --smoke-test
         echo "running simple_predict.py" && python simple_predict.py
-        echo "running simple_tune.py" && python simple_tune.py
+        echo "running simple_tune.py" && python simple_tune.py --smoke-test
         echo "running train_on_test_data.py" && python train_on_test_data.py --smoke-test
       # for f in *.py; do echo "running $f" && python "$f" || exit 1 ; done

examples/simple.py

@@ -3,6 +3,8 @@
 from sklearn import datasets
 from sklearn.model_selection import train_test_split
 
+import ray
+
 from xgboost_ray import RayDMatrix, train, RayParams
 
 
@@ -60,12 +62,16 @@ def main(cpus_per_actor, num_actors):
     parser.add_argument(
         "--num-actors",
         type=int,
-        default=1,
+        default=4,
         help="Sets number of xgboost workers to use.")
+    parser.add_argument(
+        "--smoke-test", action="store_true", default=False, help="gpu")
 
     args, _ = parser.parse_known_args()
 
-    import ray
-    ray.init(address=args.address)
+    if args.smoke_test:
+        ray.init(num_cpus=args.num_actors)
+    else:
+        ray.init(address=args.address)
 
     main(args.cpus_per_actor, args.num_actors)

examples/simple_tune.py

@@ -5,6 +5,7 @@
 from sklearn.model_selection import train_test_split
 import xgboost as xgb
 
+import ray
 from ray import tune
 
 from xgboost_ray import train, RayDMatrix, RayParams
@@ -98,9 +99,14 @@ def main(cpus_per_actor, num_actors, num_samples):
         type=int,
         default=4,
         help="Number of samples to use for Tune.")
+    parser.add_argument(
+        "--smoke-test", action="store_true", default=False, help="gpu")
 
     args, _ = parser.parse_known_args()
 
-    import ray
-    ray.init(address=args.address)
+    if args.smoke_test:
+        ray.init(num_cpus=args.num_actors * args.num_samples)
+    else:
+        ray.init(address=args.address)
+
     main(args.cpus_per_actor, args.num_actors, args.num_samples)