diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala
index a9e771ae7..7cc5f7658 100644
--- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala
+++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala
@@ -65,7 +65,8 @@ object XGBoost extends Serializable {
         if (trainingSamples.hasNext) {
           val cacheFileName: String = {
             if (useExternalMemory && trainingSamples.hasNext) {
-              s"$appName-dtrain_cache-${TaskContext.getPartitionId()}"
+              s"$appName-${TaskContext.get().stageId()}-" +
+                s"dtrain_cache-${TaskContext.getPartitionId()}"
             } else {
               null
             }
diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala
index 6dad15bc4..8df7ef389 100644
--- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala
+++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala
@@ -42,7 +42,7 @@ class XGBoostModel(_booster: Booster) extends Serializable {
         Rabit.init(rabitEnv.asJava)
         val cacheFileName = {
           if (useExternalCache) {
-            s"$appName-dtest_cache-${TaskContext.getPartitionId()}"
+            s"$appName-${TaskContext.get().stageId()}-dtest_cache-${TaskContext.getPartitionId()}"
           } else {
             null
           }