Tune XGBoost on 1,000 CPUs

In this example we:

• Download a 1.4GB Kaggle flight-delay dataset.

• Boot 13 machines with 80 CPUs each.

• Train 36 XGBoost models in parallel and pick the best one.

Dataset: 2022 flight delays

The dataset is Combined_Flights_2022.csv, a commercial flight-delay CSV from Kaggle.

The question is simple: given the route, airline, schedule, airport ids, and date fields, can we predict whether a flight arrives at least 15 minutes late?

This is the kind of job where I don't want to build a training platform. I want to try a bunch of XGBoost settings, look at the AUCs, and move on.

Step 1: Upload the CSV

Download the CSV here:

Flight Status Predictionwww.kaggle.com

Then upload it to the Burla filesystem.

(upload section is fast-forwarded)

Files uploaded here appear at ./shared inside every worker container. Files your code writes there show up back in the dashboard, so you can download results later.

Step 2: Boot some VMs

For this run we boot 13 VMs, each with 80 CPUs.

(Booting is fast-forwarded, this cluster actually took 1.5 min. to boot up)

The workers use the python:3.12 image. That image does not already have XGBoost, pandas, or sklearn installed, which is fine. Burla detects the local packages the function needs and installs them inside the worker containers when the job starts.

Step 3: Write the model function

Each function call trains one model with one set of hyperparameters. It loads the CSV from ./shared, cleans a few columns, trains XGBoost with all 80 CPUs, and returns the AUC.

import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

columns_to_train_on = [
    "Airline", "Operating_Airline", "Marketing_Airline_Network",
    "IATA_Code_Marketing_Airline", "IATA_Code_Operating_Airline",
    "Origin", "Dest", "OriginState", "DestState", "OriginWac", "DestWac",
    "OriginAirportID", "DestAirportID", "OriginCityMarketID", "DestCityMarketID",
    "Distance", "DistanceGroup", "CRSElapsedTime",
    "DayOfWeek", "Month", "Quarter", "DayofMonth", "Year",
    "DepTimeBlk", "ArrTimeBlk", "CRSDepTime", "CRSArrTime",
]

def train_model(params):
    print("Loading dataset ...")
    df = pd.read_csv("./shared/Combined_Flights_2022.csv")
    df = df.dropna(subset=["ArrDel15"])

    y = df["ArrDel15"].astype(int)
    X = pd.get_dummies(df[columns_to_train_on], dummy_na=False)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    data_size_gb = (X_train.memory_usage(deep=True).sum() + y_train.memory_usage(deep=True)) / (1024 ** 3)

    print(f"Training on {format(len(df), ',')} rows ({data_size_gb:.1f}GB) using 80 CPUs ...")
    model = xgb.XGBClassifier(tree_method="hist", n_jobs=80, eval_metric="auc", **params)
    model.fit(X_train, y_train)

    auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    print(f"Done training, [AUC={auc:.4f}] Hyperparameters: {params}")
    return {"auc": auc, "params": params}

Step 4: Test one model

Before launching the whole grid, run one model. This is the smoke test. It tells us the CSV path, packages, memory, and function shape are all sane.

from burla import remote_parallel_map

parameter_grid = [{"n_estimators": 300, "max_depth": 4, "eta": 0.1}]
results = remote_parallel_map(train_model, parameter_grid, func_cpu=80)

print(f"Model with params {parameter_grid[0]} yielded an AUC of: {results[0]}!")

We pass func_cpu=80 because the XGBoost call uses n_jobs=80. There is no magic here. The function asks for the same hardware the training code is going to use.

Step 5: Run the grid

Now we send 36 parameter sets to the cluster.

parameter_grid = [
    {"n_estimators": n, "max_depth": d, "eta": e}
    for n in [300, 600, 900]
    for d in [4, 6, 8, 10]
    for e in [0.05, 0.1, 0.15]
]

results = remote_parallel_map(train_model, parameter_grid, func_cpu=80)

best = max(results, key=lambda r: r["auc"])
print("Best:", best)

Because we have 13 machines, the first 13 models start immediately and the rest wait in the queue. Burla can queue millions of inputs, so 36 models is nothing.

You can watch logs from the Jobs tab while the models train.

What's the point?

The point is not that XGBoost is hard to parallelize. It isn't.

The annoying part is usually the stuff around it: picking machines, getting the CSV onto them, installing packages, running many trainings at once, and collecting the scores without turning the notebook into a platform project.

If someone asked me to try a parameter grid on this dataset, this is the version I would actually want to run. Same CSV, same model code, no rewrite into a different framework just because my laptop is too small.