[RFC] Design of Checkpoint Mechanism in XGBoost-Spark

## Checkpoint Cleanup

The current implementation in XGBoost-Spark does not clean up the checkpoint file after a successful training. As a result, the user may get a wrong training process without a careful setup. For instance, the following the process will lead to a wrongly-formed model in the second run.

* User starts training by setting checkpoint interval as 2 iterations and round as 10 iterations
* After a successful training, the user starts a new training job with a new group of parameters (except the changing of checkpoint directory)

The second training job will finish “too fast” for two reasons:

The checkpoint built at the 8th iteration in the first was left in the checkpoint directory
Without a different checkpoint dir setup, the second run will load the left checkpoint file and only run for 2 iterations

Therefore, we propose to always cleanup the checkpoint directory with a successful training.

this has been merged with https://github.com/dmlc/xgboost/commit/7b5cbcc8468448245a1c2ab698d07cb05be75e94

## Deterministic Partitioning

Based on the definition of gradient boosting, we should ensure that the input dataset to a booster is fixed for each iteration. However, this would be hard to achieve in the distributed training with Spark.

The current partitioning mechanism in XGBoost-Spark is to ensure that the input data can be dispatched to the correct number of partitions. The number of partitions is controlled by the user’s configuration on parameter numWorkers.

The design of deterministic partitioning is to guarantee that given a fixed input dataset the partitioning of the strategy is always the same. This goal is interpreted as follows:

* We rely on the user to ensure that the data feed to XGBoost-Spark is fixed for each checkpoint interval
* XGBoost-Spark guarantees that the partitioning of the given dataset is deterministic

The current partitioning strategy in XGBoost-Spark cannot achieve the second part of the goal. The current strategy is based on the repartition API in Spark RDD which is implemented to ensure an even distribution. The execution of the current partitioning is like the following:

* We have A partitions and would like to repartition to B partitions
* Each element in each partition of A is dispatched to B partitions but randomly start from one of B partitions.

Because of the random start, we cannot guarantee the fixed partitioning strategy for each checkpoint interval or each recovery from the failure

We propose the following mechanism to have a deterministic partitioning and achieve load balancing with the best effort

partitionId = math.abs(preventOverflow(row.getHashCode + row.getAs[Float](math.abs(row.getHashCode) % number_non_zero_features)) % numWorkers

## Avoid Multiple Jobs for Checkpointing

The current checkpoint is to collect the booster produced at the last iteration of each checkpoint internal to Driver and persist it in HDFS. The major issue with this approach is that it needs to re-perform the data preparation for training if the user didn’t choose to cache the training dataset.

The proposed change is to build the external memory checkpoint in XGBoost4J layer as well so that we can instruct XGBoost4J to save the checkpoint content to HDFS from XGBoost-Spark layer after moving forward with a specified number of iterations.

In the engineering perspective, XGBoost-Spark passes in three variables to XGBoost4J, buildExternalCache, interval and a OutputStream. then in XGBoost4J layer, we simply feed booster to OutputStream in the partition where buildExternalCache is true when it has moved interval iterations

## Prediction Cache Building

The key factor leading to the performance drop with checkpoint mechanism is that we lost the prediction cache after the first checkpoint was made so that we need to go through each tree to calculate the residual even we performed the same computation for evaluation in the last round.

There are two potential fixes to address this problem:

* Rebuild the whole prediction cache in the first iteration in each checkpoint interval: this is the simplest fix as we only need to refill the prediction cache at the beginning of each checkpoint interval. However, it also implicitly pushes the user to set checkpoint interval to a relatively large value to avoid performance issues.
* Checkpoint prediction cache and add API to load it: the current checkpoint only contains the booster itself. The proposed fix is to add prediction to the content to each checkpoint (not to change the booster format but wrap them in XGBoost-Spark layer). We also need to expose APIs to load the checkpointed cache. However, the major overhead of this fix is the request to expose new APIs which is likely only used by this XGBoost-Spark functionality. Additionally, the proposed API is to expose an internal concept of XGBoost, “Prediction Cache”, to end-users.

Based on the above analysis and the conflicts with `Avoid Multiple Jobs for Checkpointing` in the second approach, we adopt the first approach for pursuing the best practice of design.

No files, tests, or entry points are named. Start by reviewing the existing XGBoost-Spark checkpoint implementation and the linked commit, then determine which proposed areas remain unresolved: cleanup, deterministic partitioning, checkpoint job reduction, and prediction-cache rebuilding. Done would require an agreed scope and validation for the selected design.

Written by the indexing model from the issue text.