[jvm-packages] Small cleanup for GPU QDM. (#11179)

Author: trivialfis

jvm-packages/create_jni.py

@@ -95,11 +95,17 @@ def native_build(cli_args: argparse.Namespace) -> None:
             CONFIG["USE_DLOPEN_NCCL"] = "OFF"
 
         args = ["-D{0}:BOOL={1}".format(k, v) for k, v in CONFIG.items()]
+        if sys.platform != "win32":
+            try:
+                subprocess.check_call(["ninja", "--version"])
+                args.append("-GNinja")
+            except FileNotFoundError:
+                pass
 
         # if enviorment set GPU_ARCH_FLAG
         gpu_arch_flag = os.getenv("GPU_ARCH_FLAG", None)
         if gpu_arch_flag is not None:
-            args.append("%s" % gpu_arch_flag)
+            args.append("-DCMAKE_CUDA_ARCHITECTURES=%s" % gpu_arch_flag)
 
         with cd(build_dir):
             lib_dir = os.path.join(os.pardir, "lib")

jvm-packages/xgboost4j-spark-gpu/src/main/java/ml/dmlc/xgboost4j/java/QuantileDMatrix.java

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2021-2024 by Contributors
+ Copyright (c) 2021-2025 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -15,7 +15,40 @@
  */
 package ml.dmlc.xgboost4j.java;
 
+import java.io.IOException;
 import java.util.Iterator;
+import java.util.Map;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonSerializer;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+
+class F64NaNSerializer extends JsonSerializer<Double> {
+  @Override
+  public void serialize(Double value, JsonGenerator gen,
+                        SerializerProvider serializers) throws IOException {
+    if (value.isNaN()) {
+      gen.writeRawValue("NaN"); // Write NaN without quotes
+    } else {
+      gen.writeNumber(value);
+    }
+  }
+}
+
+class F32NaNSerializer extends JsonSerializer<Float> {
+  @Override
+  public void serialize(Float value, JsonGenerator gen,
+                        SerializerProvider serializers) throws IOException {
+    if (value.isNaN()) {
+      gen.writeRawValue("NaN"); // Write NaN without quotes
+    } else {
+      gen.writeNumber(value);
+    }
+  }
+}
 
 /**
  * QuantileDMatrix will only be used to train
@@ -112,8 +145,23 @@ public void setGroup(int[] group) throws XGBoostError {
   }
 
   private String getConfig(float missing, int maxBin, int nthread) {
-    return String.format("{\"missing\":%f,\"max_bin\":%d,\"nthread\":%d}",
-                         missing, maxBin, nthread);
+    Map<String, Object> conf = new java.util.HashMap<>();
+    conf.put("missing", missing);
+    conf.put("max_bin", maxBin);
+    conf.put("nthread", nthread);
+    ObjectMapper mapper = new ObjectMapper();
+
+    // Handle NaN values. Jackson by default serializes NaN values into strings.
+    SimpleModule module = new SimpleModule();
+    module.addSerializer(Double.class, new F64NaNSerializer());
+    module.addSerializer(Float.class, new F32NaNSerializer());
+    mapper.registerModule(module);
+
+    try {
+      String config = mapper.writeValueAsString(conf);
+      return config;
+    } catch (JsonProcessingException e) {
+      throw new RuntimeException("Failed to serialize configuration", e);
+    }
   }
-
 }

jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/Communicator.java

@@ -3,8 +3,6 @@
 import java.io.Serializable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.util.LinkedList;
-import java.util.List;
 import java.util.Map;
 
 import com.fasterxml.jackson.core.JsonProcessingException;

jvm-packages/xgboost4j/src/native/jvm_utils.h

@@ -1,15 +1,31 @@
+/**
+ *  Copyright 2014-2025, XGBoost Contributors
+ */
 #ifndef JVM_UTILS_H_
 #define JVM_UTILS_H_
 
-#define JVM_CHECK_CALL(__expr)                                                 \
-  {                                                                            \
-    int __errcode = (__expr);                                                  \
-    if (__errcode != 0) {                                                      \
-      return __errcode;                                                        \
-    }                                                                          \
+#include <jni.h>
+
+#include "xgboost/logging.h"  // for Check
+
+#define JVM_CHECK_CALL(__expr) \
+  {                            \
+    int __errcode = (__expr);  \
+    if (__errcode != 0) {      \
+      return __errcode;        \
+    }                          \
   }
 
-JavaVM*& GlobalJvm();
-void setHandle(JNIEnv *jenv, jlongArray jhandle, void* handle);
+JavaVM *&GlobalJvm();
+void setHandle(JNIEnv *jenv, jlongArray jhandle, void *handle);
+
+template <typename T>
+T CheckJvmCall(T const &v, JNIEnv *jenv) {
+  if (!v) {
+    CHECK(jenv->ExceptionOccurred());
+    jenv->ExceptionDescribe();
+  }
+  return v;
+}
 
 #endif  // JVM_UTILS_H_