[jvm-packages] Unbounded native memory growth per DMatrix used with Booster.predict() on Linux (2.1.4 – 3.4.0)

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.lang.management.ManagementFactory;
import java.util.HashMap;
import java.util.Map;

import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import ml.dmlc.xgboost4j.java.XGBoost;

/**
 * Minimal reproducer for a native memory leak in xgboost4j on linux/x86_64:
 * RSS grows linearly (~0.5-0.9 KB per cycle) as long as the JVM keeps doing
 * "new DMatrix -> Booster.predict -> DMatrix.dispose" cycles, and never plateaus.
 * Construct+dispose without predict is flat; a single reused DMatrix is flat.
 */
public class LeakRepro {

    public static void main(String[] args) throws Exception {
        final int rows = 1440, cols = 14;
        float[] data = new float[rows * cols];
        float[] labels = new float[rows];
        java.util.Random rnd = new java.util.Random(42);
        for (int i = 0; i < data.length; i++) data[i] = rnd.nextFloat() * 100f;
        for (int i = 0; i < rows; i++) labels[i] = rnd.nextFloat() * 100f;

        // Train a small booster (quantile objective, as used by our production models).
        Map<String, Object> params = new HashMap<>();
        params.put("objective", "reg:quantileerror");
        params.put("quantile_alpha", 0.5);
        params.put("max_depth", 4);
        params.put("nthread", 2);
        DMatrix train = new DMatrix(data, rows, cols);
        train.setLabel(labels);
        Booster booster = XGBoost.train(train, params, 20, new HashMap<String, DMatrix>(), null, null);
        train.dispose();

        // Warm up: native lib load, JIT, OpenMP thread pool.
        for (int i = 0; i < 1000; i++) {
            DMatrix dm = new DMatrix(data, rows, cols);
            booster.predict(dm, false, 0);
            dm.dispose();
        }

        // Measured loop: new DMatrix + predict + dispose.
        for (int i = 1; i <= 100_000; i++) {
            DMatrix dm = new DMatrix(data, rows, cols);
            try {
                booster.predict(dm, false, 0);
            } finally {
                dm.dispose();
            }
            if (i % 10_000 == 0) {
                System.gc(); // keep heap noise down; native leak is GC-independent
                Thread.sleep(50);
                System.out.println(i + "," + rssKb());
            }
        }
        booster.dispose();
    }

    /** Process RSS in KB (ps works on linux and macOS). */
    static long rssKb() throws Exception {
        String pid = ManagementFactory.getRuntimeMXBean().getName().split("@")[0];
        Process p = Runtime.getRuntime().exec(new String[]{"/bin/sh", "-c", "ps -o rss= -p " + pid});
        BufferedReader r = new BufferedReader(new InputStreamReader(p.getInputStream()));
        String line = r.readLine();
        p.waitFor();
        return Long.parseLong(line.trim());
    }
}

docker run --rm --platform linux/amd64 -v "$PWD":/repro -w /repro --memory 1g eclipse-temurin:8-jdk sh -c '
  apt-get update -qq && apt-get install -y -qq libgomp1
  M=https://repo1.maven.org/maven2
  curl -sLO $M/ml/dmlc/xgboost4j_2.12/3.4.0/xgboost4j_2.12-3.4.0.jar
  curl -sLO $M/com/esotericsoftware/kryo/5.6.2/kryo-5.6.2.jar
  curl -sLO $M/org/objenesis/objenesis/3.4/objenesis-3.4.jar
  curl -sLO $M/com/esotericsoftware/minlog/1.3.1/minlog-1.3.1.jar
  curl -sLO $M/commons-logging/commons-logging/1.2/commons-logging-1.2.jar
  javac -cp xgboost4j_2.12-3.4.0.jar:kryo-5.6.2.jar LeakRepro.java
  java -Xmx256m -cp .:xgboost4j_2.12-3.4.0.jar:kryo-5.6.2.jar:objenesis-3.4.jar:minlog-1.3.1.jar:commons-logging-1.2.jar LeakRepro
'