Invalidate the cache.

Author: trivialfis

include/xgboost/cache.h

@@ -78,7 +78,8 @@ class DMatrixCache {
       auto p_fmat = queue_.front();
       auto it = container_.find(p_fmat);
       CHECK(it != container_.cend());
-      if (it->second.ref.expired()) {
+      // Re-new the cache if this has never been read.
+      if (it->second.ref.expired() || !it->second.ref.lock()->Info().HasBeenRead()) {
         expired.push_back(it->first);
       } else {
         remained.push(it->first);

include/xgboost/data.h

@@ -218,8 +218,9 @@ class MetaInfo {
    * @brief Setter for categories.
    */
   void Cats(std::shared_ptr<CatContainer> cats);
-
-  void SetReadFlag() { this->has_been_read_ = true; }
+  // Flag to indicate whether one needs to refresh the DMatrix cache.
+  void SetReadFlag(bool has_been_read) { this->has_been_read_ = has_been_read; }
+  [[nodiscard]] bool HasBeenRead() const { return this->has_been_read_; }
 
  private:
   void SetInfoFromHost(Context const* ctx, StringView key, Json arr);
@@ -743,7 +744,7 @@ class DMatrix {
 
 template <>
 inline BatchSet<SparsePage> DMatrix::GetBatches() {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetRowBatches();
 }
 
@@ -764,37 +765,37 @@ inline bool DMatrix::PageExists<SparsePage>() const {
 
 template <>
 inline BatchSet<SparsePage> DMatrix::GetBatches(Context const*) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetRowBatches();
 }
 
 template <>
 inline BatchSet<CSCPage> DMatrix::GetBatches(Context const* ctx) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetColumnBatches(ctx);
 }
 
 template <>
 inline BatchSet<SortedCSCPage> DMatrix::GetBatches(Context const* ctx) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetSortedColumnBatches(ctx);
 }
 
 template <>
 inline BatchSet<EllpackPage> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetEllpackBatches(ctx, param);
 }
 
 template <>
 inline BatchSet<GHistIndexMatrix> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetGradientIndex(ctx, param);
 }
 
 template <>
 inline BatchSet<ExtSparsePage> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
-  this->Info().SetReadFlag();
+  this->Info().SetReadFlag(true);
   return GetExtBatches(ctx, param);
 }
 }  // namespace xgboost

src/data/data.cc

@@ -513,8 +513,8 @@ void CopyTensorInfoImpl(Context const* ctx, Json arr_interface, linalg::Tensor<T
 }  // namespace
 
 void MetaInfo::SetInfo(Context const& ctx, StringView key, StringView interface_str) {
-  CHECK(!this->has_been_read_)
-      << "Don't modify the DMatrix it's used. This violates some caches in XGBoost.";
+  this->SetReadFlag(false);
+
   Json j_interface = Json::Load(interface_str);
   bool is_cuda{false};
   if (IsA<Array>(j_interface)) {