Revert "[core] Correct bytes in flight when objects <5mb (#54349)" (#56387)

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

src/ray/common/ray_config_def.h

@@ -334,7 +334,9 @@ RAY_CONFIG(uint64_t, object_manager_default_chunk_size, 5 * 1024 * 1024)
 
 /// The maximum number of outbound bytes to allow to be outstanding. This avoids
 /// excessive memory usage during object broadcast to many receivers.
-RAY_CONFIG(int64_t, object_manager_max_bytes_in_flight, (int64_t)2 * 1024 * 1024 * 1024)
+RAY_CONFIG(uint64_t,
+           object_manager_max_bytes_in_flight,
+           ((uint64_t)2) * 1024 * 1024 * 1024)
 
 /// Maximum number of ids in one batch to send to GCS to delete keys.
 RAY_CONFIG(uint32_t, maximum_gcs_deletion_batch_size, 1000)

src/ray/object_manager/chunk_object_reader.h

@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <algorithm>
 #include <memory>
 #include <string>
 
@@ -42,10 +41,6 @@ class ChunkObjectReader {
 
   const IObjectReader &GetObject() const { return *object_; }
 
-  uint64_t ChunkSize() const {
-    return std::min(chunk_size_, object_->GetDataSize() + object_->GetMetadataSize());
-  }
-
  private:
   const std::shared_ptr<IObjectReader> object_;
   const uint64_t chunk_size_;

src/ray/object_manager/object_manager.cc

@@ -116,7 +116,10 @@ ObjectManager::ObjectManager(
       get_spilled_object_url_(std::move(get_spilled_object_url)),
       pull_retry_timer_(*main_service_,
                         boost::posix_time::milliseconds(config.timer_freq_ms)),
-      push_manager_(std::make_unique<PushManager>(config_.max_bytes_in_flight)) {
+      push_manager_(std::make_unique<PushManager>(/* max_chunks_in_flight= */ std::max(
+          static_cast<int64_t>(1L),
+          static_cast<int64_t>(config_.max_bytes_in_flight /
+                               config_.object_chunk_size)))) {
   RAY_CHECK_GT(config_.rpc_service_threads_number, 0);
 
   pull_retry_timer_.async_wait([this](const boost::system::error_code &e) { Tick(e); });
@@ -488,13 +491,8 @@ void ObjectManager::PushObjectInternal(const ObjectID &object_id,
       << ", total data size: " << chunk_reader->GetObject().GetObjectSize();
 
   auto push_id = UniqueID::FromRandom();
-  uint64_t push_max_chunk_size = chunk_reader->ChunkSize();
   push_manager_->StartPush(
-      node_id,
-      object_id,
-      chunk_reader->GetNumChunks(),
-      push_max_chunk_size,
-      [=](int64_t chunk_id) {
+      node_id, object_id, chunk_reader->GetNumChunks(), [=](int64_t chunk_id) {
         rpc_service_.post(
             [=]() {
               // Post to the multithreaded RPC event loop so that data is copied
@@ -505,14 +503,11 @@ void ObjectManager::PushObjectInternal(const ObjectID &object_id,
                   node_id,
                   chunk_id,
                   rpc_client,
-                  [this, push_max_chunk_size](const Status &status) {
+                  [=](const Status &status) {
                     // Post back to the main event loop because the
                     // PushManager is not thread-safe.
-                    this->main_service_->post(
-                        [this, push_max_chunk_size]() {
-                          this->push_manager_->OnChunkComplete(push_max_chunk_size);
-                        },
-                        "ObjectManager.Push");
+                    main_service_->post([this]() { push_manager_->OnChunkComplete(); },
+                                        "ObjectManager.Push");
                   },
                   chunk_reader,
                   from_disk);

src/ray/object_manager/object_manager.h

@@ -53,7 +53,7 @@ struct ObjectManagerConfig {
   /// Object chunk size, in bytes
   uint64_t object_chunk_size;
   /// Max object push bytes in flight.
-  int64_t max_bytes_in_flight;
+  uint64_t max_bytes_in_flight;
   /// The store socket name.
   std::string store_socket_name;
   /// The time in milliseconds to wait until a Push request

src/ray/object_manager/push_manager.cc

@@ -24,7 +24,6 @@ namespace ray {
 void PushManager::StartPush(const NodeID &dest_id,
                             const ObjectID &obj_id,
                             int64_t num_chunks,
-                            int64_t max_chunk_size,
                             std::function<void(int64_t)> send_chunk_fn) {
   auto push_id = std::make_pair(dest_id, obj_id);
   RAY_CHECK(num_chunks > 0);
@@ -38,7 +37,6 @@ void PushManager::StartPush(const NodeID &dest_id,
         dest_id,
         obj_id,
         num_chunks,
-        max_chunk_size,
         std::move(send_chunk_fn));
   } else {
     RAY_LOG(DEBUG) << "Duplicate push request " << push_id.first << ", " << push_id.second
@@ -49,8 +47,8 @@ void PushManager::StartPush(const NodeID &dest_id,
   ScheduleRemainingPushes();
 }
 
-void PushManager::OnChunkComplete(int64_t push_max_chunk_size) {
-  bytes_in_flight_ -= push_max_chunk_size;
+void PushManager::OnChunkComplete() {
+  chunks_in_flight_ -= 1;
   chunks_remaining_ -= 1;
   ScheduleRemainingPushes();
 }
@@ -64,17 +62,17 @@ void PushManager::ScheduleRemainingPushes() {
 
   // Loop over all active pushes for approximate round-robin prioritization.
   bool keep_looping = true;
-  while (bytes_in_flight_ < max_bytes_in_flight_ && keep_looping) {
+  while (chunks_in_flight_ < max_chunks_in_flight_ && keep_looping) {
     // Loop over each active push and try to send another chunk.
-    // If we could push out a chunk and haven't reached the max_bytes_in_flight_ limit,
+    // If we could push out a chunk and haven't reached the chunks_in_flight_ limit,
     // we'll loop again to try to send more chunks.
     keep_looping = false;
     auto iter = push_requests_with_chunks_to_send_.begin();
     while (iter != push_requests_with_chunks_to_send_.end() &&
-           bytes_in_flight_ < max_bytes_in_flight_) {
+           chunks_in_flight_ < max_chunks_in_flight_) {
       auto &push_state = *iter;
       push_state.SendOneChunk();
-      bytes_in_flight_ += push_state.max_chunk_size_;
+      chunks_in_flight_ += 1;
       if (push_state.num_chunks_to_send_ == 0) {
         auto push_state_map_iter = push_state_map_.find(push_state.node_id_);
         RAY_CHECK(push_state_map_iter != push_state_map_.end());
@@ -109,16 +107,18 @@ void PushManager::HandleNodeRemoved(const NodeID &node_id) {
 
 void PushManager::RecordMetrics() const {
   ray::stats::STATS_push_manager_num_pushes_remaining.Record(
-      push_requests_with_chunks_to_send_.size());
-  ray::stats::STATS_push_manager_chunks.Record(chunks_remaining_, "Remaining");
+      NumPushRequestsWithChunksToSend());
+  ray::stats::STATS_push_manager_chunks.Record(NumChunksInFlight(), "InFlight");
+  ray::stats::STATS_push_manager_chunks.Record(NumChunksRemaining(), "Remaining");
 }
 
 std::string PushManager::DebugString() const {
   std::stringstream result;
   result << "PushManager:";
-  result << "\n- num pushes remaining: " << push_requests_with_chunks_to_send_.size();
-  result << "\n- num chunks remaining: " << chunks_remaining_;
-  result << "\n- max bytes allowed: " << max_bytes_in_flight_;
+  result << "\n- num pushes remaining: " << NumPushRequestsWithChunksToSend();
+  result << "\n- num chunks in flight: " << NumChunksInFlight();
+  result << "\n- num chunks remaining: " << NumChunksRemaining();
+  result << "\n- max chunks allowed: " << max_chunks_in_flight_;
   return result.str();
 }
 

src/ray/object_manager/push_manager.h

@@ -28,10 +28,12 @@ class PushManager {
  public:
   /// Create a push manager.
   ///
-  /// \param max_bytes_in_flight Max number of bytes allowed to be in flight
+  /// \param max_chunks_in_flight Max number of chunks allowed to be in flight
   ///                             from this PushManager (this raylet).
-  explicit PushManager(int64_t max_bytes_in_flight)
-      : max_bytes_in_flight_(max_bytes_in_flight){};
+  explicit PushManager(int64_t max_chunks_in_flight)
+      : max_chunks_in_flight_(max_chunks_in_flight) {
+    RAY_CHECK_GT(max_chunks_in_flight_, 0);
+  };
 
   /// Start pushing an object subject to max chunks in flight limit.
   ///
@@ -40,39 +42,40 @@ class PushManager {
   /// \param dest_id The node to send to.
   /// \param obj_id The object to send.
   /// \param num_chunks The total number of chunks to send.
-  /// \param max_chunk_size See comment for max_chunk_size_ in PushState.
   /// \param send_chunk_fn This function will be called with args 0...{num_chunks-1}.
   ///                      The caller promises to call PushManager::OnChunkComplete()
   ///                      once a call to send_chunk_fn finishes.
   void StartPush(const NodeID &dest_id,
                  const ObjectID &obj_id,
                  int64_t num_chunks,
-                 int64_t max_chunk_size,
                  std::function<void(int64_t)> send_chunk_fn);
 
   /// Called every time a chunk completes to trigger additional sends.
   /// TODO(ekl) maybe we should cancel the entire push on error.
-  void OnChunkComplete(int64_t push_max_chunk_size);
+  void OnChunkComplete();
 
   /// Cancel all pushes that have not yet been sent to the removed node.
   void HandleNodeRemoved(const NodeID &node_id);
 
-  void RecordMetrics() const;
-
-  int64_t BytesInFlight() const { return bytes_in_flight_; }
-
-  int64_t ChunksRemaining() const { return chunks_remaining_; }
+  /// Return the number of chunks currently in flight. For metrics and testing.
+  int64_t NumChunksInFlight() const { return chunks_in_flight_; };
 
-  int64_t PushesInFlight() const { return push_state_map_.size(); }
+  /// Return the number of chunks remaining. For metrics and testing.
+  int64_t NumChunksRemaining() const { return chunks_remaining_; }
 
-  int64_t PushRequestsRemaining() const {
+  /// Return the number of push requests with remaining chunks. For metrics and testing.
+  int64_t NumPushRequestsWithChunksToSend() const {
     return push_requests_with_chunks_to_send_.size();
-  }
+  };
+
+  /// Record the internal metrics.
+  void RecordMetrics() const;
 
   std::string DebugString() const;
 
  private:
   FRIEND_TEST(TestPushManager, TestPushState);
+  FRIEND_TEST(TestPushManager, TestNodeRemoved);
 
   /// Tracks the state of an active object push to another node.
   struct PushState {
@@ -81,13 +84,8 @@ class PushManager {
 
     /// total number of chunks of this object.
     int64_t num_chunks_;
-    /// the max size of a chunk for this object in bytes, used to count bytes_in_flight_
-    /// and assure it stays under max_bytes_in_flight_. This means we can overcount for
-    /// the last chunk but we're accepting that to keep the code simpler.
-    int64_t max_chunk_size_;
     /// The function to send chunks with.
     std::function<void(int64_t)> chunk_send_fn_;
-
     /// The index of the next chunk to send.
     int64_t next_chunk_id_ = 0;
     /// The number of chunks remaining to send.
@@ -96,12 +94,10 @@ class PushManager {
     PushState(NodeID node_id,
               ObjectID object_id,
               int64_t num_chunks,
-              int64_t max_chunk_size,
               std::function<void(int64_t)> chunk_send_fn)
         : node_id_(node_id),
           object_id_(object_id),
           num_chunks_(num_chunks),
-          max_chunk_size_(max_chunk_size),
           chunk_send_fn_(std::move(chunk_send_fn)),
           num_chunks_to_send_(num_chunks) {}
 
@@ -126,11 +122,11 @@ class PushManager {
   /// Called on completion events to trigger additional pushes.
   void ScheduleRemainingPushes();
 
-  /// Max number of bytes in flight allowed.
-  const int64_t max_bytes_in_flight_;
+  /// Max number of chunks in flight allowed.
+  const int64_t max_chunks_in_flight_;
 
-  /// Running count of bytes in flight
-  int64_t bytes_in_flight_ = 0;
+  /// Running count of chunks in flight, used to limit progress of in_flight_pushes_.
+  int64_t chunks_in_flight_ = 0;
 
   /// Remaining count of chunks to push to other nodes.
   int64_t chunks_remaining_ = 0;