ray/release/xgboost_tests/workloads/train_moderate.py
Kai Fricke 18d535f290
[ci/release] Migrate LightGBM tests (#22952)
Note that LightGBM release tests were previously not enabled.
https://buildkite.com/ray-project/release-tests-branch/builds/113
https://buildkite.com/ray-project/release-tests-branch/builds/114

Co-authored-by: Antoni Baum <antoni.baum@protonmail.com>
2022-03-10 08:14:31 +00:00

49 lines
1.1 KiB
Python

"""Moderate cluster training
This training run will start 32 workers on 32 nodes (including head node).
Test owner: krfricke
Acceptance criteria: Should run through and report final results.
"""
import json
import os
import time
import ray
from xgboost_ray import RayParams
from ray.util.xgboost.release_test_util import train_ray
if __name__ == "__main__":
ray.init(address="auto")
ray_params = RayParams(
elastic_training=False,
max_actor_restarts=2,
num_actors=32,
cpus_per_actor=4,
gpus_per_actor=0,
)
start = time.time()
train_ray(
path="/data/classification.parquet",
num_workers=None,
num_boost_rounds=100,
num_files=128,
regression=False,
use_gpu=False,
ray_params=ray_params,
xgboost_params=None,
)
taken = time.time() - start
result = {
"time_taken": taken,
}
test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/train_moderate.json")
with open(test_output_json, "wt") as f:
json.dump(result, f)
print("PASSED.")