Move plasma retry logic into plasma store provider (#7328)

src/ray/common/ray_config_def.h

@@ -238,3 +238,10 @@ RAY_CONFIG(uint32_t, object_store_get_max_ids_to_print_in_warning, 20)
 /// Note: this only takes effect when gcs service is enabled.
 RAY_CONFIG(int64_t, gcs_service_connect_retries, 50)
 RAY_CONFIG(int64_t, gcs_service_connect_wait_milliseconds, 100)
+
+/// Maximum number of times to retry putting an object when the plasma store is full.
+/// Can be set to -1 to enable unlimited retries.
+RAY_CONFIG(int32_t, object_store_full_max_retries, 5)
+/// Duration to sleep after failing to put an object in plasma because it is full.
+/// This will be exponentially increased for each retry.
+RAY_CONFIG(uint32_t, object_store_full_initial_delay_ms, 1000)

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -55,28 +55,46 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr<Buffer> &meta
                                              const size_t data_size,
                                              const ObjectID &object_id,
                                              std::shared_ptr<Buffer> *data) {
-  auto plasma_id = object_id.ToPlasmaId();
-  std::shared_ptr<arrow::Buffer> arrow_buffer;
-  {
-    std::lock_guard<std::mutex> guard(store_client_mutex_);
-    arrow::Status status =
-        store_client_.Create(plasma_id, data_size, metadata ? metadata->Data() : nullptr,
-                             metadata ? metadata->Size() : 0, &arrow_buffer);
-    if (plasma::IsPlasmaObjectExists(status)) {
-      RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: "
-                       << object_id << ".";
-      return Status::OK();
+  int32_t retries = 0;
+  int32_t max_retries = RayConfig::instance().object_store_full_max_retries();
+  uint32_t delay = RayConfig::instance().object_store_full_initial_delay_ms();
+  Status status;
+  bool should_retry = true;
+  while (should_retry) {
+    should_retry = false;
+    arrow::Status plasma_status;
+    std::shared_ptr<arrow::Buffer> arrow_buffer;
+    {
+      std::lock_guard<std::mutex> guard(store_client_mutex_);
+      plasma_status = store_client_.Create(
+          object_id.ToPlasmaId(), data_size, metadata ? metadata->Data() : nullptr,
+          metadata ? metadata->Size() : 0, &arrow_buffer);
     }
-    if (plasma::IsPlasmaStoreFull(status)) {
+    if (plasma::IsPlasmaStoreFull(plasma_status)) {
       std::ostringstream message;
       message << "Failed to put object " << object_id << " in object store because it "
               << "is full. Object size is " << data_size << " bytes.";
-      return Status::ObjectStoreFull(message.str());
+      status = Status::ObjectStoreFull(message.str());
+      if (max_retries < 0 || retries < max_retries) {
+        RAY_LOG(ERROR) << message.str() << " Plasma store status:\n"
+                       << MemoryUsageString() << "\nWaiting " << delay
+                       << "ms for space to free up...";
+        usleep(1000 * delay);
+        delay *= 2;
+        retries += 1;
+        should_retry = true;
+      }
+    } else if (plasma::IsPlasmaObjectExists(plasma_status)) {
+      RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: "
+                       << object_id << ".";
+      status = Status::OK();
+    } else {
+      RAY_ARROW_RETURN_NOT_OK(plasma_status);
+      *data = std::make_shared<PlasmaBuffer>(PlasmaBuffer(arrow_buffer));
+      status = Status::OK();
     }
-    RAY_ARROW_RETURN_NOT_OK(status);
   }
-  *data = std::make_shared<PlasmaBuffer>(PlasmaBuffer(arrow_buffer));
-  return Status::OK();
+  return status;
 }
 
 Status CoreWorkerPlasmaStoreProvider::Seal(const ObjectID &object_id) {