diff --git a/openml-xgboost/pom.xml b/openml-xgboost/pom.xml new file mode 100644 index 00000000..e3395386 --- /dev/null +++ b/openml-xgboost/pom.xml @@ -0,0 +1,158 @@ + + + + + + com.feedzai + openml-java + 0.0.0-SNAPSHOT + + 4.0.0 + + openml-xgboost + OpenML XGBoost + Provider that imports, scores and trains XGBoost models using the native xgboost4j JVM package. + + + + 3.4.0 + + + + + com.feedzai + openml-api + provided + + + com.feedzai + openml-utils + provided + + + + + ml.dmlc + xgboost4j_2.12 + ${xgboost.version} + + + org.scala-lang + scala-compiler + + + com.esotericsoftware + kryo + + + + + + com.esotericsoftware + kryo + 5.6.2 + provided + + + + com.google.guava + guava + + + org.slf4j + slf4j-api + + + com.google.auto.service + auto-service + + + + + com.feedzai + openml-utils + test-jar + test + + + junit + junit + test + + + org.assertj + assertj-core + test + + + org.apache.commons + commons-csv + test + + + commons-io + commons-io + test + + + ch.qos.logback + logback-classic + test + + + + + + + + org.apache.maven.plugins + maven-jar-plugin + + + + test-jar + + + + + + + diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostAlgorithms.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostAlgorithms.java new file mode 100644 index 00000000..9b2954cf --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostAlgorithms.java @@ -0,0 +1,59 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.provider.descriptor.MLAlgorithmDescriptor; +import com.feedzai.openml.provider.descriptor.MachineLearningAlgorithmType; +import com.feedzai.openml.util.algorithm.MLAlgorithmEnum; + +import static com.feedzai.openml.util.algorithm.MLAlgorithmEnum.createDescriptor; + +/** + * Specifies the XGBoost algorithms that can be imported and trained through this provider. + */ +public enum XgboostAlgorithms implements MLAlgorithmEnum { + + /** + * XGBoost binary classifier. + */ + XGBOOST_BINARY_CLASSIFIER(createDescriptor( + "DMLC - XGBoost", + XgboostDescriptorUtil.PARAMS, + MachineLearningAlgorithmType.SUPERVISED_BINARY_CLASSIFICATION, + "https://xgboost.readthedocs.io/" + )); + + /** + * {@link MLAlgorithmDescriptor} for this algorithm. + */ + private final MLAlgorithmDescriptor descriptor; + + /** + * Constructor. + * + * @param descriptor {@link MLAlgorithmDescriptor} for this algorithm. + */ + XgboostAlgorithms(final MLAlgorithmDescriptor descriptor) { + this.descriptor = descriptor; + } + + @Override + public MLAlgorithmDescriptor getAlgorithmDescriptor() { + return this.descriptor; + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostClassificationModel.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostClassificationModel.java new file mode 100644 index 00000000..8791140b --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostClassificationModel.java @@ -0,0 +1,157 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.data.Instance; +import com.feedzai.openml.data.schema.DatasetSchema; +import com.feedzai.openml.model.ClassificationMLModel; +import com.feedzai.openml.provider.exception.ModelLoadingException; +import ml.dmlc.xgboost4j.java.Booster; +import ml.dmlc.xgboost4j.java.XGBoostError; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.file.Path; + +/** + * A classification model backed by a native XGBoost {@link Booster}, used for real-time single-instance + * scoring. + * + *

Scoring uses {@link Booster#inplace_predict(float[], int, int, float)} on a single-row feature + * vector, which avoids allocating a {@code DMatrix} per prediction. The native booster handle is not + * thread-safe, so predictions are serialized on a private lock (mirrors the H2O provider's approach). + */ +public class XgboostClassificationModel implements ClassificationMLModel { + + /** + * Logger for this class. + */ + private static final Logger logger = LoggerFactory.getLogger(XgboostClassificationModel.class); + + /** + * Value used to signal a missing feature to XGBoost. + */ + private static final float MISSING_VALUE = Float.NaN; + + /** + * The native XGBoost booster. + */ + private final Booster booster; + + /** + * The schema the model uses. + */ + private final DatasetSchema schema; + + /** + * The number of predictive features expected by the model. + */ + private final int numFeatures; + + /** + * Lock serializing access to the non-thread-safe native booster during prediction. + */ + private final Object predictLock = new Object(); + + /** + * Constructor. + * + * @param booster The trained/loaded native XGBoost booster. + * @param schema The {@link DatasetSchema} the model uses. + */ + XgboostClassificationModel(final Booster booster, final DatasetSchema schema) { + this.booster = booster; + this.schema = schema; + this.numFeatures = XgboostSchemaUtils.numFeatures(schema); + } + + @Override + public double[] getClassDistribution(final Instance instance) { + final float[] row = XgboostSchemaUtils.featureRow(instance, this.schema); + + final float[][] predictions; + try { + // The native booster handle is not thread-safe; serialize predictions. + synchronized (this.predictLock) { + predictions = this.booster.inplace_predict(row, 1, this.numFeatures, MISSING_VALUE); + } + } catch (final XGBoostError e) { + throw new RuntimeException("XGBoost failed to score the instance.", e); + } + + return toClassDistribution(predictions[0]); + } + + @Override + public int classify(final Instance instance) { + final double[] distribution = getClassDistribution(instance); + + int argMax = 0; + for (int i = 1; i < distribution.length; i++) { + if (distribution[i] > distribution[argMax]) { + argMax = i; + } + } + return argMax; + } + + @Override + public boolean save(final Path dir, final String name) { + try { + this.booster.saveModel(dir.resolve(XgboostModelCreator.MODEL_BINARY_RESOURCE_FILE_NAME).toString()); + return true; + } catch (final XGBoostError e) { + logger.error("Failed to save XGBoost model {} to {}.", name, dir, e); + return false; + } + } + + @Override + public DatasetSchema getSchema() { + return this.schema; + } + + @Override + public void close() { + this.booster.dispose(); + } + + /** + * Converts a raw XGBoost prediction row into a class distribution aligned with the schema's target + * classes. + * + *

For binary objectives XGBoost outputs a single value - the probability of the positive class - + * which is expanded to {@code [1 - p, p]}. For multi-class objectives ({@code multi:softprob}) the + * per-class probability vector is returned as-is. + * + * @param prediction The raw prediction row for a single instance. + * @return The class distribution. + */ + static double[] toClassDistribution(final float[] prediction) { + if (prediction.length == 1) { + final double positiveProbability = prediction[0]; + return new double[]{1.0 - positiveProbability, positiveProbability}; + } + + final double[] distribution = new double[prediction.length]; + for (int i = 0; i < prediction.length; i++) { + distribution[i] = prediction[i]; + } + return distribution; + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostDescriptorUtil.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostDescriptorUtil.java new file mode 100644 index 00000000..4c410c38 --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostDescriptorUtil.java @@ -0,0 +1,265 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.provider.descriptor.ModelParameter; +import com.feedzai.openml.provider.descriptor.fieldtype.ChoiceFieldType; +import com.feedzai.openml.provider.descriptor.fieldtype.NumericFieldType; +import com.google.common.collect.ImmutableSet; + +import java.util.Set; + +/** + * Organizes the Machine Learning hyper-parameters exposed for training XGBoost models. + * + *

The parameter names match the native XGBoost parameter names (see + * XGBoost Parameters) so they can + * be forwarded directly to the {@code xgboost4j} training API. + */ +final class XgboostDescriptorUtil { + + /** + * Alias to ease readability of mandatory parameters. + */ + private static final boolean MANDATORY = true; + + /** + * Alias to ease readability of non-mandatory parameters. + */ + private static final boolean NOT_MANDATORY = false; + + /** + * The learning task and objective. Restricted to {@code binary:logistic} - the only objective whose + * single probability output maps cleanly onto the binary {@code [1 - p, p]} class distribution this + * provider produces. + */ + static final String OBJECTIVE_PARAMETER_NAME = "objective"; + + /** + * The number of boosting rounds (trees). Passed as the {@code nrounds} argument of + * {@code XGBoost.train}, not as a booster parameter. + */ + static final String NUM_ROUND_PARAMETER_NAME = "num_round"; + + /** + * Random seed parameter name. + */ + static final String SEED_PARAMETER_NAME = "seed"; + + /** + * Number of parallel threads parameter name. + */ + static final String NTHREAD_PARAMETER_NAME = "nthread"; + + /** + * The set of parameters accepted when training an XGBoost model. + */ + static final Set PARAMS = ImmutableSet.of( + new ModelParameter( + OBJECTIVE_PARAMETER_NAME, + "Objective", + "The learning task and corresponding objective. Only 'binary:logistic' is supported:\n" + + "it outputs the probability of the positive class, which this provider maps\n" + + "directly onto a [1 - p, p] class distribution.", + MANDATORY, + new ChoiceFieldType( + ImmutableSet.of("binary:logistic"), + "binary:logistic" + ) + ), + new ModelParameter( + NUM_ROUND_PARAMETER_NAME, + "Number of boosting rounds", + "Number of boosting iterations (trees) to build.", + MANDATORY, + intRange(1, Integer.MAX_VALUE, 50) + ), + new ModelParameter( + "eta", + "Learning rate (eta)", + "Step size shrinkage used in updates to prevent over-fitting. Also named 'learning_rate'.", + NOT_MANDATORY, + doubleRange(0.0, 1.0, 0.3) + ), + new ModelParameter( + "max_depth", + "Maximum tree depth", + "Maximum depth of a tree. Increasing this value makes the model more complex and more\n" + + "likely to over-fit. 0 means no limit.", + NOT_MANDATORY, + intRange(0, Integer.MAX_VALUE, 6) + ), + new ModelParameter( + "min_child_weight", + "Minimum child weight", + "Minimum sum of instance weight (hessian) needed in a child. Larger values are more\n" + + "conservative.", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 1.0) + ), + new ModelParameter( + "max_delta_step", + "Maximum delta step", + "Maximum delta step allowed for each leaf output. 0 means no constraint. Can help make\n" + + "the update more conservative, notably with very imbalanced classes.", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 0.0) + ), + new ModelParameter( + "gamma", + "Minimum split loss (gamma)", + "Minimum loss reduction required to make a further partition on a leaf node.", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 0.0) + ), + new ModelParameter( + "subsample", + "Subsample ratio", + "Subsample ratio of the training instances. Setting it to 0.5 means XGBoost randomly\n" + + "samples half of the training data prior to growing trees.", + NOT_MANDATORY, + doubleRange(1E-6, 1.0, 1.0) + ), + new ModelParameter( + "colsample_bytree", + "Column subsample ratio by tree", + "Subsample ratio of columns when constructing each tree.", + NOT_MANDATORY, + doubleRange(1E-6, 1.0, 1.0) + ), + new ModelParameter( + "colsample_bylevel", + "Column subsample ratio by level", + "Subsample ratio of columns for each level. Subsampling occurs once for every new depth\n" + + "level reached in a tree, over the columns already chosen for the current tree.", + NOT_MANDATORY, + doubleRange(1E-6, 1.0, 1.0) + ), + new ModelParameter( + "lambda", + "L2 regularization (lambda)", + "L2 regularization term on weights. Increasing this value makes the model more conservative.", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 1.0) + ), + new ModelParameter( + "alpha", + "L1 regularization (alpha)", + "L1 regularization term on weights. Increasing this value makes the model more conservative.", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 0.0) + ), + new ModelParameter( + "scale_pos_weight", + "Scale positive weight", + "Controls the balance of positive and negative weights, useful for imbalanced classes.\n" + + "A typical value to consider is sum(negative instances) / sum(positive instances).", + NOT_MANDATORY, + doubleRange(0.0, Double.MAX_VALUE, 1.0) + ), + new ModelParameter( + "tree_method", + "Tree method", + "The tree construction algorithm used in XGBoost:\n" + + "'auto' (same as 'hist'), 'exact', 'approx' and 'hist'.", + NOT_MANDATORY, + new ChoiceFieldType( + ImmutableSet.of("auto", "exact", "approx", "hist"), + "auto" + ) + ), + new ModelParameter( + "grow_policy", + "Grow policy", + "Controls the way new nodes are added to the tree:\n" + + "'depthwise' splits at nodes closest to the root, 'lossguide' splits at nodes with\n" + + "the highest loss change.", + NOT_MANDATORY, + new ChoiceFieldType( + ImmutableSet.of("depthwise", "lossguide"), + "depthwise" + ) + ), + new ModelParameter( + "max_leaves", + "Maximum number of leaves", + "Maximum number of nodes to be added. 0 means no limit. Not used by the 'exact' tree\n" + + "method.", + NOT_MANDATORY, + intRange(0, Integer.MAX_VALUE, 0) + ), + new ModelParameter( + "max_bin", + "Maximum number of bins", + "Maximum number of discrete bins used to bucket continuous features. Only used by the\n" + + "'hist' and 'approx' tree methods. Increasing it improves optimality of splits at the\n" + + "cost of higher computation time.", + NOT_MANDATORY, + intRange(2, Integer.MAX_VALUE, 256) + ), + new ModelParameter( + SEED_PARAMETER_NAME, + "Seed", + "Random number seed used for reproducibility.", + NOT_MANDATORY, + intRange(0, Integer.MAX_VALUE, 0) + ), + new ModelParameter( + NTHREAD_PARAMETER_NAME, + "Number of threads", + "Number of parallel threads used to run XGBoost. 0 (the default) uses all available\n" + + "threads, matching the behavior of the previous H2O-backed implementation.", + NOT_MANDATORY, + intRange(0, Integer.MAX_VALUE, 0) + ) + ); + + /** + * This class is not meant to be instantiated. + */ + private XgboostDescriptorUtil() { + } + + /** + * Helper that builds a {@code DOUBLE} numeric range. + * + * @param minValue Minimum allowed value. + * @param maxValue Maximum allowed value. + * @param defaultValue Default value. + * @return The numeric field type. + */ + private static NumericFieldType doubleRange(final double minValue, + final double maxValue, + final double defaultValue) { + return NumericFieldType.range(minValue, maxValue, NumericFieldType.ParameterConfigType.DOUBLE, defaultValue); + } + + /** + * Helper that builds an {@code INT} numeric range. + * + * @param minValue Minimum allowed value. + * @param maxValue Maximum allowed value. + * @param defaultValue Default value. + * @return The numeric field type. + */ + private static NumericFieldType intRange(final int minValue, + final int maxValue, + final int defaultValue) { + return NumericFieldType.range(minValue, maxValue, NumericFieldType.ParameterConfigType.INT, defaultValue); + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelCreator.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelCreator.java new file mode 100644 index 00000000..6b879752 --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelCreator.java @@ -0,0 +1,369 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.data.Dataset; +import com.feedzai.openml.data.Instance; +import com.feedzai.openml.data.schema.CategoricalValueSchema; +import com.feedzai.openml.data.schema.DatasetSchema; +import com.feedzai.openml.data.schema.FieldSchema; +import com.feedzai.openml.data.schema.StringValueSchema; +import com.feedzai.openml.provider.descriptor.fieldtype.ParamValidationError; +import com.feedzai.openml.provider.exception.ModelLoadingException; +import com.feedzai.openml.provider.exception.ModelTrainingException; +import com.feedzai.openml.provider.model.MachineLearningModelTrainer; +import com.feedzai.openml.util.load.LoadModelUtils; +import com.feedzai.openml.util.load.LoadSchemaUtils; +import com.feedzai.openml.util.validate.ValidationUtils; +import com.google.common.collect.ImmutableList; +import ml.dmlc.xgboost4j.java.Booster; +import ml.dmlc.xgboost4j.java.DMatrix; +import ml.dmlc.xgboost4j.java.XGBoost; +import ml.dmlc.xgboost4j.java.XGBoostError; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Random; +import java.util.stream.Stream; + +/** + * Loads and trains XGBoost models through the native {@code xgboost4j} JVM package. + * + *

Training is fully in-process (no Spark), mirroring how the H2O provider trains inside an + * embedded in-JVM instance: the dataset is materialized into an in-memory {@link DMatrix}, + * {@link XGBoost#train} builds the booster, the booster is exported and then reloaded. Because + * {@code xgboost4j} is a thin JNI wrapper it runs on Java 8-25 and on ARM (Graviton / Apple Silicon). + */ +public class XgboostModelCreator implements MachineLearningModelTrainer { + + /** + * Logger for this class. + */ + private static final Logger logger = LoggerFactory.getLogger(XgboostModelCreator.class); + + /** + * Name of the model file written inside the model folder, using XGBoost's portable UBJSON format. + */ + public static final String MODEL_BINARY_RESOURCE_FILE_NAME = "XGBoost_model.ubj"; + + /** + * Prefix for the temporary directory holding a freshly trained model before it is reloaded. + */ + private static final String EXPORT_DIR_PREFIX = "fdz_xgboost_"; + + /** + * Value used to signal a missing feature to XGBoost. + */ + private static final float MISSING_VALUE = Float.NaN; + + /** + * Default number of boosting rounds used if the parameter is absent. Matches the default exposed in + * {@link XgboostDescriptorUtil} and the previous (DS Python package) implementation. + */ + private static final int DEFAULT_NUM_ROUND = 50; + + /** + * Error reported when the schema contains free-text (string) fields, which XGBoost cannot consume. + */ + static final String ERROR_MSG_SCHEMA_HAS_STRING_FIELDS = + "XGBoost only supports numeric and categorical fields; the schema contains string fields."; + + /** + * Error reported when the target field is not binary. This provider only offers binary classification. + */ + static final String ERROR_MSG_NON_BINARY_TARGET = + "XGBoost is offered here as a binary classifier; the target field must have exactly two classes."; + + @Override + public XgboostClassificationModel loadModel(final Path modelPath, final DatasetSchema schema) + throws ModelLoadingException { + + logger.info("Loading XGBoost model from [{}].", modelPath); + final String modelFilePath = resolveModelFile(modelPath).toAbsolutePath().toString(); + + try { + final Booster booster = XGBoost.loadModel(modelFilePath); + logger.info("XGBoost model loaded successfully."); + return new XgboostClassificationModel(booster, schema); + } catch (final XGBoostError e) { + throw new ModelLoadingException( + String.format("Failed to load the XGBoost model from [%s].", modelFilePath), e); + } + } + + /** + * Resolves the actual model file to load, supporting both layouts used across the codebase: + *

+ * + * @param modelPath The path provided to {@link #loadModel(Path, DatasetSchema)}. + * @return The path of the model file to load. + * @throws ModelLoadingException If the model file cannot be located within the model folder layout. + */ + static Path resolveModelFile(final Path modelPath) throws ModelLoadingException { + if (!Files.isDirectory(modelPath)) { + return modelPath; + } + if (Files.isDirectory(modelPath.resolve(LoadModelUtils.MODEL_FOLDER))) { + return LoadModelUtils.getModelFilePath(modelPath); + } + return modelPath.resolve(MODEL_BINARY_RESOURCE_FILE_NAME); + } + + @Override + public DatasetSchema loadSchema(final Path modelPath) throws ModelLoadingException { + return LoadSchemaUtils.datasetSchemaFromJson(modelPath); + } + + @Override + public List validateForLoad(final Path modelPath, + final DatasetSchema schema, + final Map params) { + final ImmutableList.Builder errorBuilder = ImmutableList.builder(); + + errorBuilder.addAll(ValidationUtils.baseLoadValidations(schema, params)); + errorBuilder.addAll(ValidationUtils.validateModelInDir(modelPath)); + ValidationUtils.validateCategoricalSchema(schema).ifPresent(errorBuilder::add); + if (schemaHasStringFields(schema)) { + errorBuilder.add(new ParamValidationError(ERROR_MSG_SCHEMA_HAS_STRING_FIELDS)); + } + + return errorBuilder.build(); + } + + @Override + public XgboostClassificationModel fit(final Dataset dataset, + final Random random, + final Map params) throws ModelTrainingException { + + final DatasetSchema schema = dataset.getSchema(); + + DMatrix trainMatrix = null; + Booster booster = null; + Path exportDir = null; + try { + trainMatrix = buildTrainMatrix(dataset); + + final Map boosterParams = toBoosterParams(params, random); + final int numRound = numRoundOf(params); + + booster = XGBoost.train(trainMatrix, boosterParams, numRound, new HashMap<>(), null, null); + + // Export the freshly trained booster to a temporary folder and reload it through loadModel so + // that fit() returns a model built by the exact same code path as model import. This guarantees + // a freshly-trained model scores identically to the same model once saved and reloaded, and + // keeps a single load/scoring path (the LightGBM provider follows the same approach). + exportDir = exportModel(booster); + return loadModel(exportDir, schema); + } catch (final XGBoostError | IOException | ModelLoadingException e) { + throw new ModelTrainingException("Failed to train the XGBoost model.", e); + } finally { + if (booster != null) { + booster.dispose(); + } + if (trainMatrix != null) { + trainMatrix.dispose(); + } + deleteExportDir(exportDir); + } + } + + @Override + public List validateForFit(final Path pathToPersist, + final DatasetSchema schema, + final Map params) { + final ImmutableList.Builder errorBuilder = ImmutableList.builder(); + + errorBuilder.addAll(ValidationUtils.validateModelPathToTrain(pathToPersist)); + errorBuilder.addAll(ValidationUtils.checkParams( + XgboostAlgorithms.XGBOOST_BINARY_CLASSIFIER.getAlgorithmDescriptor(), params)); + ValidationUtils.validateCategoricalSchema(schema).ifPresent(errorBuilder::add); + if (schemaHasStringFields(schema)) { + errorBuilder.add(new ParamValidationError(ERROR_MSG_SCHEMA_HAS_STRING_FIELDS)); + } + if (getNumTargetClasses(schema).orElse(-1) != 2) { + errorBuilder.add(new ParamValidationError(ERROR_MSG_NON_BINARY_TARGET)); + } + + return errorBuilder.build(); + } + + /** + * Checks whether the schema contains any free-text (string) field. XGBoost consumes only numeric + * input, and Pulse encodes categorical values as numeric indices, so numeric and categorical fields + * are supported but {@link StringValueSchema} fields are not. + * + * @param schema The schema to inspect. + * @return {@code true} if at least one field uses a {@link StringValueSchema}. + */ + private static boolean schemaHasStringFields(final DatasetSchema schema) { + return schema.getFieldSchemas().stream() + .anyMatch(field -> field.getValueSchema() instanceof StringValueSchema); + } + + /** + * Returns the number of classes of the schema's target field, mirroring the LightGBM provider. The + * result is empty when there is no target field or the target is not categorical, which the binary + * validation treats the same as a non-binary target. + * + * @param schema The schema to inspect. + * @return The number of target classes, or empty if there is no categorical target field. + */ + private static Optional getNumTargetClasses(final DatasetSchema schema) { + return schema.getTargetFieldSchema() + .map(FieldSchema::getValueSchema) + .filter(valueSchema -> valueSchema instanceof CategoricalValueSchema) + .map(valueSchema -> ((CategoricalValueSchema) valueSchema).getNominalValues().size()); + } + + /** + * Materializes the whole dataset into an in-memory dense {@link DMatrix} with its label column set. + * + * @param dataset The training dataset. + * @return The training {@link DMatrix}. + * @throws XGBoostError If the native matrix cannot be created. + * @throws ModelTrainingException If the dataset is empty. + */ + private static DMatrix buildTrainMatrix(final Dataset dataset) throws XGBoostError, ModelTrainingException { + final DatasetSchema schema = dataset.getSchema(); + final int numFeatures = XgboostSchemaUtils.numFeatures(schema); + // Supervised training requires a target; enforced by validateForFit via validateCategoricalSchema. + final int targetIndex = schema.getTargetIndex().orElseThrow( + () -> new IllegalStateException("Supervised training requires a schema with a target field.")); + + final List rows = new ArrayList<>(); + final List labels = new ArrayList<>(); + + final Iterator iterator = dataset.getInstances(); + while (iterator.hasNext()) { + final Instance instance = iterator.next(); + labels.add((float) instance.getValue(targetIndex)); + rows.add(XgboostSchemaUtils.featureRow(instance, schema)); + } + + final int numRows = rows.size(); + if (numRows == 0) { + throw new ModelTrainingException("Received an empty training dataset for XGBoost."); + } + + final float[] flatFeatures = new float[numRows * numFeatures]; + final float[] labelArray = new float[numRows]; + for (int row = 0; row < numRows; row++) { + System.arraycopy(rows.get(row), 0, flatFeatures, row * numFeatures, numFeatures); + labelArray[row] = labels.get(row); + } + + final DMatrix trainMatrix = new DMatrix(flatFeatures, numRows, numFeatures, MISSING_VALUE); + trainMatrix.setLabel(labelArray); + return trainMatrix; + } + + /** + * Translates the Pulse string parameters into the {@code Map} expected by + * {@code xgboost4j}. The {@value XgboostDescriptorUtil#NUM_ROUND_PARAMETER_NAME} entry is excluded + * because it is passed as the {@code nrounds} argument of {@link XGBoost#train}. A seed is derived + * from the supplied {@link Random} when not explicitly provided, for reproducibility. + * + * @param params The Pulse model parameters. + * @param random The source of randomness. + * @return The XGBoost booster parameters. + */ + static Map toBoosterParams(final Map params, final Random random) { + final Map boosterParams = new HashMap<>(); + + params.forEach((name, value) -> { + if (!XgboostDescriptorUtil.NUM_ROUND_PARAMETER_NAME.equals(name) && value != null && !value.isEmpty()) { + boosterParams.put(name, value); + } + }); + + boosterParams.putIfAbsent(XgboostDescriptorUtil.OBJECTIVE_PARAMETER_NAME, "binary:logistic"); + boosterParams.putIfAbsent(XgboostDescriptorUtil.SEED_PARAMETER_NAME, random.nextInt(Integer.MAX_VALUE)); + + return boosterParams; + } + + /** + * Reads the number of boosting rounds from the parameters, falling back to a default. + * + * @param params The Pulse model parameters. + * @return The number of boosting rounds. + */ + static int numRoundOf(final Map params) { + final String numRound = params.get(XgboostDescriptorUtil.NUM_ROUND_PARAMETER_NAME); + if (numRound == null || numRound.isEmpty()) { + return DEFAULT_NUM_ROUND; + } + return Integer.parseInt(numRound.trim()); + } + + /** + * Exports a trained booster following the Pulse model folder convention + * ({@code /model/}), so it can be reloaded through {@link #loadModel}. + * + * @param booster The trained booster. + * @return The export directory root. + * @throws IOException If the export directories/files cannot be created. + * @throws XGBoostError If the booster cannot be serialized. + */ + private static Path exportModel(final Booster booster) throws IOException, XGBoostError { + final Path exportDir = Files.createTempDirectory(EXPORT_DIR_PREFIX); + final Path modelDir = Files.createDirectory(exportDir.resolve(LoadModelUtils.MODEL_FOLDER)); + booster.saveModel(modelDir.resolve(MODEL_BINARY_RESOURCE_FILE_NAME).toString()); + return exportDir; + } + + /** + * Best-effort recursive deletion of the temporary export directory created by {@link #exportModel}. + * The directory only exists to hand the freshly trained model to {@link #loadModel}; once the model + * is loaded into memory the files are no longer needed. Failures are logged and swallowed, mirroring + * how the LightGBM provider cleans up its temporary model file. + * + * @param exportDir The temporary export directory, or {@code null} if it was never created. + */ + private static void deleteExportDir(final Path exportDir) { + if (exportDir == null) { + return; + } + try (Stream paths = Files.walk(exportDir)) { + paths.sorted(Comparator.reverseOrder()).forEach(path -> { + try { + Files.delete(path); + } catch (final IOException e) { + logger.warn("Could not delete temporary export path [{}]: {}", path, e.getMessage()); + } + }); + } catch (final IOException e) { + logger.warn("Could not clean up temporary export directory [{}]: {}", exportDir, e.getMessage()); + } + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelProvider.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelProvider.java new file mode 100644 index 00000000..bea1eaf8 --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostModelProvider.java @@ -0,0 +1,60 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.provider.MachineLearningProvider; +import com.feedzai.openml.provider.TrainingMachineLearningProvider; +import com.feedzai.openml.provider.descriptor.MLAlgorithmDescriptor; +import com.feedzai.openml.util.algorithm.MLAlgorithmEnum; +import com.google.auto.service.AutoService; + +import java.util.Optional; +import java.util.Set; + +/** + * Feedzai OpenML {@link MachineLearningProvider} for XGBoost, backed by the native {@code xgboost4j} JVM + * package. + * + *

The provider is discovered by Pulse through the standard Java {@link java.util.ServiceLoader} + * mechanism (via {@link AutoService}), so no changes to Pulse core are required to make it available - + * only adding this module to the runtime classpath. + */ +@AutoService(MachineLearningProvider.class) +public class XgboostModelProvider implements TrainingMachineLearningProvider { + + /** + * The reported name of this provider. + */ + public static final String PROVIDER_NAME = "XGBoost"; + + @Override + public String getName() { + return PROVIDER_NAME; + } + + @Override + public Set getAlgorithms() { + return MLAlgorithmEnum.getDescriptors(XgboostAlgorithms.values()); + } + + @Override + public Optional getModelCreator(final String algorithmName) { + return MLAlgorithmEnum.getByName(XgboostAlgorithms.values(), algorithmName) + .map(algorithm -> new XgboostModelCreator()); + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostSchemaUtils.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostSchemaUtils.java new file mode 100644 index 00000000..ad19a6da --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostSchemaUtils.java @@ -0,0 +1,76 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.data.Instance; +import com.feedzai.openml.data.schema.DatasetSchema; + +/** + * Shared helpers to turn Pulse {@link Instance}s into the flat {@code float[]} feature vectors XGBoost + * expects. + * + *

Why this is shared between scoring and training: XGBoost is purely numeric and positional - + * it has no notion of feature names or categorical domains. Therefore the exact same column ordering and + * encoding must be used when a model is trained and when it is later scored. Centralizing the feature + * vector construction here guarantees that parity: both {@link XgboostModelCreator} (training) and + * {@link XgboostClassificationModel} (scoring) build rows through this class. + * + *

Categorical fields arrive already encoded as {@code double} indices in the {@link Instance} (Pulse's + * standard encoding), so they are copied as-is - identical to the LightGBM provider's behavior. + */ +final class XgboostSchemaUtils { + + /** + * This class is not meant to be instantiated. + */ + private XgboostSchemaUtils() { + } + + /** + * The number of predictive (non-target) features described by the schema. + * + * @param schema The dataset schema. + * @return The number of predictive features. + */ + static int numFeatures(final DatasetSchema schema) { + return schema.getPredictiveFields().size(); + } + + /** + * Builds the feature vector for a single {@link Instance}, in schema field order, skipping the target + * field if one is present. + * + * @param instance The instance to convert. + * @param schema The dataset schema the instance conforms to. + * @return A dense {@code float[]} with one entry per predictive feature. + */ + static float[] featureRow(final Instance instance, final DatasetSchema schema) { + final int numFields = schema.getFieldSchemas().size(); + final int targetIndex = schema.getTargetIndex().orElse(-1); + final float[] row = new float[numFeatures(schema)]; + + int featureIdx = 0; + for (int fieldIdx = 0; fieldIdx < numFields; fieldIdx++) { + if (fieldIdx == targetIndex) { + continue; + } + row[featureIdx++] = (float) instance.getValue(fieldIdx); + } + return row; + } +} diff --git a/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/package-info.java b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/package-info.java new file mode 100644 index 00000000..f8a2da0f --- /dev/null +++ b/openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/package-info.java @@ -0,0 +1,21 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +/** + * OpenML XGBoost provider, backed by the native {@code xgboost4j} JVM package. + */ +package com.feedzai.openml.provider.xgboost; diff --git a/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostInternalsTest.java b/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostInternalsTest.java new file mode 100644 index 00000000..636cf8e6 --- /dev/null +++ b/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostInternalsTest.java @@ -0,0 +1,238 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.data.Dataset; +import com.feedzai.openml.data.schema.CategoricalValueSchema; +import com.feedzai.openml.data.schema.DatasetSchema; +import com.feedzai.openml.data.schema.FieldSchema; +import com.feedzai.openml.data.schema.NumericValueSchema; +import com.feedzai.openml.mocks.MockDataset; +import com.feedzai.openml.mocks.MockInstance; +import com.feedzai.openml.provider.descriptor.ModelParameter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.junit.Test; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.offset; + +/** + * Unit tests for the pure-logic helpers of the XGBoost provider that do not require the native library, + * so they run on every platform (including musl/Alpine where the native library is unavailable). + */ +public class XgboostInternalsTest { + + /** + * Schema with the target in the middle position (index 1) and two numeric features. + */ + private static final DatasetSchema SCHEMA_TARGET_IN_MIDDLE = new DatasetSchema(1, ImmutableList.of( + new FieldSchema("f0", 0, new NumericValueSchema(false)), + new FieldSchema("target", 1, new CategoricalValueSchema(false, ImmutableSet.of("0", "1"))), + new FieldSchema("f2", 2, new NumericValueSchema(false)) + )); + + // --- XgboostSchemaUtils --- + + /** + * The feature vector is built in field order, skipping the target column. + */ + @Test + public void featureRowSkipsTargetColumn() { + assertThat(XgboostSchemaUtils.numFeatures(SCHEMA_TARGET_IN_MIDDLE)).isEqualTo(2); + + final MockInstance instance = new MockInstance(ImmutableList.of(10.0, 0.0, 30.0)); + final float[] row = XgboostSchemaUtils.featureRow(instance, SCHEMA_TARGET_IN_MIDDLE); + + assertThat(row).containsExactly(10.0f, 30.0f); + } + + // --- XgboostModelCreator.numRoundOf --- + + /** + * The number of rounds falls back to the default when absent or empty, and is parsed otherwise. + */ + @Test + public void numRoundOfHandlesDefaultAndValue() { + assertThat(XgboostModelCreator.numRoundOf(ImmutableMap.of())).isEqualTo(50); + assertThat(XgboostModelCreator.numRoundOf(ImmutableMap.of("num_round", ""))).isEqualTo(50); + assertThat(XgboostModelCreator.numRoundOf(ImmutableMap.of("num_round", "120"))).isEqualTo(120); + } + + // --- XgboostModelCreator.toBoosterParams --- + + /** + * {@code num_round} and empty values are excluded, and objective/seed defaults are injected. + */ + @Test + public void toBoosterParamsExcludesNumRoundAndInjectsDefaults() { + final Map params = ImmutableMap.of( + "num_round", "10", + "eta", "0.3", + "max_depth", "" + ); + + final Map boosterParams = XgboostModelCreator.toBoosterParams(params, new Random(0)); + + assertThat(boosterParams).doesNotContainKey("num_round"); + assertThat(boosterParams).doesNotContainKey("max_depth"); + assertThat(boosterParams).containsEntry("eta", "0.3"); + assertThat(boosterParams).containsEntry("objective", "binary:logistic"); + assertThat(boosterParams).containsKey("seed"); + } + + /** + * Explicit objective and seed are preserved (not overridden by defaults). Note that + * {@code toBoosterParams} is a transport-only mapper; objective validity is enforced upstream by + * {@code validateForFit} through the algorithm descriptor, not here. + */ + @Test + public void toBoosterParamsPreservesExplicitValues() { + final Map params = ImmutableMap.of( + "objective", "binary:logistic", + "seed", "42" + ); + + final Map boosterParams = XgboostModelCreator.toBoosterParams(params, new Random(0)); + + assertThat(boosterParams).containsEntry("objective", "binary:logistic"); + assertThat(boosterParams).containsEntry("seed", "42"); + } + + // --- XgboostModelCreator.resolveModelFile --- + + /** + * A direct path to a file is returned unchanged. + * + * @throws Exception If file operations fail. + */ + @Test + public void resolveModelFileReturnsDirectFile() throws Exception { + final Path file = Files.createTempFile("xgb_model_", ".ubj"); + assertThat(XgboostModelCreator.resolveModelFile(file)).isEqualTo(file); + } + + /** + * A directory with a {@code model/} sub-folder resolves to the file within it (Pulse layout). + * + * @throws Exception If file operations fail. + */ + @Test + public void resolveModelFileResolvesPulseModelFolderLayout() throws Exception { + final Path root = Files.createTempDirectory("xgb_root_"); + final Path modelDir = Files.createDirectory(root.resolve("model")); + final Path modelFile = Files.createFile(modelDir.resolve(XgboostModelCreator.MODEL_BINARY_RESOURCE_FILE_NAME)); + + assertThat(XgboostModelCreator.resolveModelFile(root)).isEqualTo(modelFile); + } + + /** + * A directory without a {@code model/} sub-folder resolves to the model file at its root. + * + * @throws Exception If file operations fail. + */ + @Test + public void resolveModelFileResolvesRootLayout() throws Exception { + final Path root = Files.createTempDirectory("xgb_root_flat_"); + + assertThat(XgboostModelCreator.resolveModelFile(root)) + .isEqualTo(root.resolve(XgboostModelCreator.MODEL_BINARY_RESOURCE_FILE_NAME)); + } + + // --- XgboostClassificationModel.toClassDistribution --- + + /** + * A single-value (binary) prediction expands to {@code [1 - p, p]}. + */ + @Test + public void toClassDistributionExpandsBinaryPrediction() { + final double[] distribution = XgboostClassificationModel.toClassDistribution(new float[]{0.3f}); + + assertThat(distribution).hasSize(2); + assertThat(distribution[1]).isCloseTo(0.3f, offset(1e-6)); + assertThat(distribution[0]).isCloseTo(1.0 - 0.3f, offset(1e-6)); + } + + /** + * A multi-value (multi-class) prediction is returned as-is. + */ + @Test + public void toClassDistributionPassesThroughMulticlassPrediction() { + final double[] distribution = XgboostClassificationModel.toClassDistribution(new float[]{0.2f, 0.5f, 0.3f}); + + assertThat(distribution).hasSize(3); + assertThat(distribution[0]).isCloseTo(0.2f, offset(1e-6)); + assertThat(distribution[1]).isCloseTo(0.5f, offset(1e-6)); + assertThat(distribution[2]).isCloseTo(0.3f, offset(1e-6)); + } + + // --- misc no-native paths --- + + /** + * {@link XgboostClassificationModel#getSchema()} returns the schema it was constructed with. + */ + @Test + public void getSchemaReturnsProvidedSchema() { + final XgboostClassificationModel model = new XgboostClassificationModel(null, SCHEMA_TARGET_IN_MIDDLE); + assertThat(model.getSchema()).isSameAs(SCHEMA_TARGET_IN_MIDDLE); + } + + /** + * Fitting with a schema that has no target field fails fast (before any native call). + */ + @Test + public void fitWithoutTargetSchemaThrows() { + final DatasetSchema noTargetSchema = new DatasetSchema(ImmutableList.of( + new FieldSchema("f0", 0, new NumericValueSchema(false)), + new FieldSchema("f1", 1, new NumericValueSchema(false)) + )); + final Dataset dataset = new MockDataset(noTargetSchema, 5, new Random(0)); + + assertThatThrownBy(() -> new XgboostModelCreator().fit(dataset, new Random(0), ImmutableMap.of())) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("target"); + } + + // --- XgboostDescriptorUtil parameter surface --- + + /** + * The exposed parameter surface includes the hyper-parameters requested in review (added alongside + * the original set). + */ + @Test + public void descriptorExposesExpectedParameters() { + final Set paramNames = XgboostDescriptorUtil.PARAMS.stream() + .map(ModelParameter::getName) + .collect(Collectors.toSet()); + + assertThat(paramNames).contains( + "objective", "num_round", "eta", "max_depth", "min_child_weight", "gamma", + "subsample", "colsample_bytree", "colsample_bylevel", "max_delta_step", "lambda", + "alpha", "scale_pos_weight", "tree_method", "grow_policy", "max_leaves", "max_bin", + "seed", "nthread"); + } +} diff --git a/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostModelProviderTest.java b/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostModelProviderTest.java new file mode 100644 index 00000000..02a8b28a --- /dev/null +++ b/openml-xgboost/src/test/java/com/feedzai/openml/provider/xgboost/XgboostModelProviderTest.java @@ -0,0 +1,347 @@ +/* + * Copyright 2026 Feedzai + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package com.feedzai.openml.provider.xgboost; + +import com.feedzai.openml.data.Dataset; +import com.feedzai.openml.data.schema.CategoricalValueSchema; +import com.feedzai.openml.data.schema.DatasetSchema; +import com.feedzai.openml.data.schema.FieldSchema; +import com.feedzai.openml.data.schema.NumericValueSchema; +import com.feedzai.openml.data.schema.StringValueSchema; +import com.feedzai.openml.mocks.MockDataset; +import com.feedzai.openml.mocks.MockInstance; +import com.feedzai.openml.provider.descriptor.MLAlgorithmDescriptor; +import com.feedzai.openml.provider.descriptor.fieldtype.ParamValidationError; +import com.feedzai.openml.provider.exception.ModelLoadingException; +import com.feedzai.openml.provider.exception.ModelTrainingException; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import ml.dmlc.xgboost4j.java.DMatrix; +import org.junit.Assume; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.offset; + +/** + * Round-trip tests for the XGBoost provider: train (in-process, no Spark) -> export -> load -> score. + * + *

These tests exercise the native {@code xgboost4j} library, validating that the provider works + * end-to-end on the host architecture (including ARM / Apple Silicon). + */ +public class XgboostModelProviderTest { + + /** + * Binary target nominal values. + */ + private static final Set TARGET_VALUES = ImmutableSet.of("false", "true"); + + /** + * Number of predictive features in the test schema. + */ + private static final int NUM_FEATURES = 4; + + /** + * Schema used across the tests (4 numeric features + binary categorical target). + */ + private static DatasetSchema schema; + + /** + * Whether the native {@code xgboost4j} library can be loaded on this platform. Tests that train or + * score are skipped when it cannot (e.g. musl/Alpine, for which XGBoost ships no native library). + */ + private static boolean nativeAvailable; + + /** + * Sets up the shared schema and detects native library availability. + */ + @BeforeClass + public static void setUp() { + schema = MockDataset.generateDefaultSchema(TARGET_VALUES, NUM_FEATURES); + nativeAvailable = xgboostNativeAvailable(); + } + + /** + * Probes whether the native XGBoost library can be loaded on the current platform. + * + *

The published {@code xgboost4j} jar bundles glibc Linux (x86_64, aarch64), macOS (x86_64, + * Apple Silicon) and Windows natives, but no musl build - so on Alpine/musl the load fails. + * + * @return {@code true} if the native library initializes successfully. + */ + private static boolean xgboostNativeAvailable() { + try { + new DMatrix(new float[]{0f}, 1, 1, Float.NaN).dispose(); + return true; + } catch (final Throwable t) { + // UnsatisfiedLinkError / NoClassDefFoundError / XGBoostError: native unsupported here. + return false; + } + } + + /** + * Valid training parameters. + * + * @return The parameters map. + */ + private static Map trainParams() { + return ImmutableMap.of( + XgboostDescriptorUtil.OBJECTIVE_PARAMETER_NAME, "binary:logistic", + XgboostDescriptorUtil.NUM_ROUND_PARAMETER_NAME, "10", + "max_depth", "3", + "eta", "0.3", + XgboostDescriptorUtil.NTHREAD_PARAMETER_NAME, "1" + ); + } + + /** + * The provider exposes the XGBoost algorithm and resolves its creator by name. + */ + @Test + public void providerExposesXgboostAlgorithm() { + final XgboostModelProvider provider = new XgboostModelProvider(); + + assertThat(provider.getName()).isEqualTo("XGBoost"); + assertThat(provider.getAlgorithms()) + .extracting(MLAlgorithmDescriptor::getAlgorithmName) + .contains("DMLC - XGBoost"); + assertThat(provider.getModelCreator("DMLC - XGBoost")).isPresent(); + assertThat(provider.getModelCreator("Non Existing Algorithm")).isEmpty(); + } + + /** + * Valid fit parameters produce no validation errors. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForFitAcceptsValidParams() throws Exception { + final Path tmpDir = Files.createTempDirectory("xgb_fit_validation_"); + final List errors = + new XgboostModelCreator().validateForFit(tmpDir, schema, trainParams()); + + assertThat(errors).isEmpty(); + } + + /** + * Trains a model in-process, then scores an instance: the class distribution must be a valid + * probability distribution and {@code classify} must return the arg-max class. + * + * @throws Exception If training/scoring fails. + */ + @Test + public void trainsAndScoresInProcess() throws Exception { + Assume.assumeTrue("XGBoost native library unavailable on this platform (e.g. musl/Alpine).", nativeAvailable); + + final XgboostModelCreator creator = new XgboostModelCreator(); + final Dataset trainDataset = new MockDataset(schema, 200, new Random(0)); + + final XgboostClassificationModel model = creator.fit(trainDataset, new Random(0), trainParams()); + + final MockInstance instance = new MockInstance(schema, new Random(7)); + final double[] distribution = model.getClassDistribution(instance); + + assertThat(distribution).hasSize(TARGET_VALUES.size()); + assertThat(distribution[0] + distribution[1]).isCloseTo(1.0, offset(1e-6)); + assertThat(distribution[0]).isBetween(0.0, 1.0); + assertThat(distribution[1]).isBetween(0.0, 1.0); + + final int classIndex = model.classify(instance); + assertThat(classIndex).isBetween(0, 1); + assertThat(distribution[classIndex]).isGreaterThanOrEqualTo(distribution[1 - classIndex]); + + model.close(); + } + + /** + * A model saved to disk and reloaded produces identical scores (export -> load round-trip). + * + * @throws Exception If training/saving/loading fails. + */ + @Test + public void savedModelReloadsWithIdenticalScores() throws Exception { + Assume.assumeTrue("XGBoost native library unavailable on this platform (e.g. musl/Alpine).", nativeAvailable); + + final XgboostModelCreator creator = new XgboostModelCreator(); + final Dataset trainDataset = new MockDataset(schema, 200, new Random(1)); + + final XgboostClassificationModel trainedModel = creator.fit(trainDataset, new Random(1), trainParams()); + + final MockInstance instance = new MockInstance(schema, new Random(11)); + final double[] originalDistribution = trainedModel.getClassDistribution(instance); + + final Path saveDir = Files.createTempDirectory("xgb_save_"); + assertThat(trainedModel.save(saveDir, "reloaded")).isTrue(); + trainedModel.close(); + + final XgboostClassificationModel reloadedModel = creator.loadModel(saveDir, schema); + final double[] reloadedDistribution = reloadedModel.getClassDistribution(instance); + + assertThat(reloadedDistribution).containsExactly(originalDistribution, offset(1e-9)); + + reloadedModel.close(); + } + + /** + * Training on an empty dataset raises a {@link ModelTrainingException}. + */ + @Test + public void trainingOnEmptyDatasetThrows() { + final XgboostModelCreator creator = new XgboostModelCreator(); + final Dataset emptyDataset = new MockDataset(schema, 0, new Random(0)); + + assertThatThrownBy(() -> creator.fit(emptyDataset, new Random(0), trainParams())) + .isInstanceOf(ModelTrainingException.class) + .hasMessageContaining("empty"); + } + + /** + * {@code validateForLoad} runs its validations and reports an error when no model exists in the + * given directory. This path does not touch the native library. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForLoadReportsMissingModel() throws Exception { + final Path emptyDir = Files.createTempDirectory("xgb_no_model_"); + + final List errors = + new XgboostModelCreator().validateForLoad(emptyDir, schema, ImmutableMap.of()); + + assertThat(errors).isNotEmpty(); + } + + /** + * Loading an invalid (non-XGBoost) model file raises a {@link ModelLoadingException}. + * + * @throws Exception If file operations fail. + */ + @Test + public void loadModelThrowsOnInvalidModelFile() throws Exception { + Assume.assumeTrue("XGBoost native library unavailable on this platform (e.g. musl/Alpine).", nativeAvailable); + + final Path root = Files.createTempDirectory("xgb_bad_model_"); + final Path modelDir = Files.createDirectory(root.resolve("model")); + Files.write(modelDir.resolve(XgboostModelCreator.MODEL_BINARY_RESOURCE_FILE_NAME), + "this is not a valid xgboost model".getBytes(StandardCharsets.UTF_8)); + + assertThatThrownBy(() -> new XgboostModelCreator().loadModel(root, schema)) + .isInstanceOf(ModelLoadingException.class); + } + + /** + * A schema containing a free-text (string) feature field, which XGBoost cannot consume. + */ + private static DatasetSchema schemaWithStringField() { + return new DatasetSchema(1, ImmutableList.of( + new FieldSchema("stringFeature", 0, new StringValueSchema(false)), + new FieldSchema("target", 1, new CategoricalValueSchema(false, ImmutableSet.of("0", "1"))) + )); + } + + /** + * {@code validateForFit} rejects schemas with string fields. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForFitRejectsStringFields() throws Exception { + final Path tmpDir = Files.createTempDirectory("xgb_fit_string_"); + + final List errors = + new XgboostModelCreator().validateForFit(tmpDir, schemaWithStringField(), trainParams()); + + assertThat(errors).extracting(ParamValidationError::getMessage) + .contains(XgboostModelCreator.ERROR_MSG_SCHEMA_HAS_STRING_FIELDS); + } + + /** + * {@code validateForLoad} rejects schemas with string fields. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForLoadRejectsStringFields() throws Exception { + final Path emptyDir = Files.createTempDirectory("xgb_load_string_"); + + final List errors = + new XgboostModelCreator().validateForLoad(emptyDir, schemaWithStringField(), ImmutableMap.of()); + + assertThat(errors).extracting(ParamValidationError::getMessage) + .contains(XgboostModelCreator.ERROR_MSG_SCHEMA_HAS_STRING_FIELDS); + } + + /** + * A schema whose target field has more than two classes, which this binary provider cannot train. + */ + private static DatasetSchema schemaWithNonBinaryTarget() { + return new DatasetSchema(1, ImmutableList.of( + new FieldSchema("f0", 0, new NumericValueSchema(false)), + new FieldSchema("target", 1, new CategoricalValueSchema(false, ImmutableSet.of("a", "b", "c"))) + )); + } + + /** + * {@code validateForFit} rejects a schema whose target is not binary. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForFitRejectsNonBinaryTarget() throws Exception { + final Path tmpDir = Files.createTempDirectory("xgb_fit_nonbinary_"); + + final List errors = + new XgboostModelCreator().validateForFit(tmpDir, schemaWithNonBinaryTarget(), trainParams()); + + assertThat(errors).extracting(ParamValidationError::getMessage) + .contains(XgboostModelCreator.ERROR_MSG_NON_BINARY_TARGET); + } + + /** + * {@code validateForFit} rejects an objective outside the supported set (only {@code binary:logistic} + * is allowed), enforced through the algorithm descriptor's choice field. + * + * @throws Exception If the temporary directory cannot be created. + */ + @Test + public void validateForFitRejectsUnsupportedObjective() throws Exception { + final Path tmpDir = Files.createTempDirectory("xgb_fit_objective_"); + final Map params = ImmutableMap.of( + XgboostDescriptorUtil.OBJECTIVE_PARAMETER_NAME, "multi:softprob", + XgboostDescriptorUtil.NUM_ROUND_PARAMETER_NAME, "10" + ); + + final List errors = + new XgboostModelCreator().validateForFit(tmpDir, schema, params); + + assertThat(errors).extracting(ParamValidationError::getMessage) + .anySatisfy(message -> assertThat(message).contains(XgboostDescriptorUtil.OBJECTIVE_PARAMETER_NAME)); + } +} diff --git a/openml-xgboost/src/test/resources/logback-test.xml b/openml-xgboost/src/test/resources/logback-test.xml new file mode 100644 index 00000000..d3498ba5 --- /dev/null +++ b/openml-xgboost/src/test/resources/logback-test.xml @@ -0,0 +1,12 @@ + + + + + %d{HH:mm:ss.SSS} [%thread] %-5level %logger{0} - %msg%n + + + + + + + diff --git a/pom.xml b/pom.xml index a6a71beb..a193ffff 100644 --- a/pom.xml +++ b/pom.xml @@ -30,6 +30,7 @@ openml-h2o openml-java-utils openml-lightgbm + openml-xgboost Java OpenML Main