[TEST][FLAKY] test_auto_scheduler_cost_model.py::test_xgb_model #6305
Description
As part of a recent PR #6302, I noticed a test that had nothing to do with the implementation is failing for unknown reasons.
That is tests/python/unittest/test_auto_scheduler_cost_model.py::test_xgb_model, and error is shown below:
________________________________ test_xgb_model ________________________________
def test_xgb_model():
task, dag, inputs, results = get_sample_records(50)
model = auto_scheduler.XGBModel(num_warmup_sample=-1)
model.update(inputs, results)
preds = model.predict(task, [x.state for x in inputs])
assert len(preds) == len(inputs)
costs = [np.mean([x.value for x in res.costs]) for res in results]
throughputs = np.min(costs) / costs
rmse = np.sqrt(np.mean([np.square(pred - label) for pred, label in zip(preds, throughputs)]))
> assert rmse <= 0.05
E assert 0.09984155049112678 <= 0.05
tests/python/unittest/test_auto_scheduler_cost_model.py:68: AssertionError