dmlc / dmlc/xgboost

[spark] xgboost.spark supports saving model that is very large

Open
#8,818 14 comments 0 reactions 0 assignees View on GitHub
feature-request
Dominant language
C++
Stars
28.8k
Forks
8.9k
Avg merge
1d 12h
Merged PRs (30d)
54

Description

Currently, ` xgboost.spark ` saves model by:

```
booster = xgb_model.get_booster().save_raw("json").decode("utf-8")
_get_spark_session().sparkContext.parallelize([booster], 1).saveAsTextFile(
model_save_path
)
```

when `booster` is a string larger than 2GB, spark will raise error like:
```
java.lang.OutOfMemoryError: Requested array size exceeds VM limit
at java.lang.StringCoding.encode(StringCoding.java:350)
at java.lang.String.getBytes(String.java:941)
at org.apache.spark.unsafe.types.UTF8String.fromString(UTF8String.java:163)
at org.apache.spark.sql.execution.python.EvaluatePython$$anonfun$$nestedInanonfun$makeFromJava$11$1.applyOrElse(EvaluatePython.scala:149)
at org.apache.spark.sql.execution.python.EvaluatePython$.nullSafeConvert(EvaluatePython.scala:213)
at org.apache.spark.sql.execution.python.EvaluatePython$.$anonfun$makeFromJava$11(EvaluatePython.scala:148)
at org.apache.spark.sql.execution.python.EvaluatePython$$$Lambda$8923/940497638.apply(Unknown Source)
at org.apache.spark.sql.execution.python.EvaluatePython$$anonfun$$nestedInanonfun$makeFromJava$16$1.applyOrElse(EvaluatePython.scala:195)
at org.apache.spark.sql.execution.python.EvaluatePython$.nullSafeConvert(EvaluatePython.scala:213)
at org.apache.spark.sql.execution.python.EvaluatePython$.$anonfun$makeFromJava$16(EvaluatePython.scala:182)
at org.apache.spark.sql.execution.python.EvaluatePython$$$Lambda$8925/1092357589.apply(Unknown Source)
at org.apache.spark.sql.SparkSession.$anonfun$applySchemaToPythonRDD$2(SparkSession.scala:960)
at org.apache.spark.sql.SparkSession$$Lambda$8929/213457688.apply(Unknown Source)
```

We need to consider split the model string and save it.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start from the xgboost.spark save path shown in the issue, where save_raw("json") is parallelized and written with saveAsTextFile. Reproduce or trace the >2GB failure, then determine a storage approach that avoids a single oversized string and verify large-model save and load behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, spark
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.