Thread-Safe Predictions with XGBoost

When making predictions with an XGBoost model in a parallel processing environment, it’s crucial to ensure that the predictions are thread-safe.

By default, XGBoost’s predict() method is thread-safe.

This means it can be called concurrently from multiple threads, for example:

import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from concurrent.futures import ThreadPoolExecutor

# Generate synthetic data
X, y = make_classification(n_samples=10000000, n_features=5, n_informative=3, n_redundant=1, random_state=42)
# Split the dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)
# Configure the XGBoost model
model = XGBClassifier()
# Train the XGBoost model
model.fit(X_train, y_train)

# Make predictions in parallel with thread safety
def predict(data):
    return model.predict(data)

# start 4 threads for concurrent predictions
with ThreadPoolExecutor(4) as executor:
    results = list(executor.map(predict, np.array_split(X_test, 4)))

predictions = np.concatenate(results)
print(f'Total predictions: {len(predictions)}')

The same model may also be used for inference from multiple child processes.

Note, this example makes a pickled copy of the model and data which is transmitted to each child processes.

For example:

import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from concurrent.futures import ProcessPoolExecutor

# Generate synthetic data
X, y = make_classification(n_samples=10000, n_features=5, n_informative=3, n_redundant=1, random_state=42)
# Split the dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)
# Configure the XGBoost model
model = XGBClassifier()
# Train the XGBoost model
model.fit(X_train, y_train)

# Make predictions in parallel with thread safety
def predict(data):
    return model.predict(data)

if __name__ == '__main__':
    # start 4 threads for concurrent predictions
    with ProcessPoolExecutor(4) as executor:
        results = list(executor.map(predict, np.array_split(X_test, 4)))

    predictions = np.concatenate(results)
    print(f'Total predictions: {len(predictions)}')

Only the predict() function is thread-safe. Other functions on the model may not be thread safe.

See Also