test_large_input.py

import numpy as np
import xgboost as xgb
import cupy as cp
import time
import pytest


# Test for integer overflow or out of memory exceptions
def test_large_input():
    available_bytes, _ = cp.cuda.runtime.memGetInfo()
    # 15 GB
    required_bytes = 1.5e+10
    if available_bytes < required_bytes:
        pytest.skip("Not enough memory on this device")
    n = 1000
    m = ((1 << 31) + n - 1) // n
    assert (np.log2(m * n) > 31)
    X = cp.ones((m, n), dtype=np.float32)
    y = cp.ones(m)
    dmat = xgb.DeviceQuantileDMatrix(X, y)
    xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)