Reverts 1c6f26d

PiperOrigin-RevId: 646524880
openxla · Jun 27, 2024 · 17d2736 · 17d2736
1 parent da7a0b3
commit 17d2736
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/xla/tsl/framework/serving_device_selector.h b/xla/tsl/framework/serving_device_selector.h
@@ -152,6 +152,14 @@ class ServingDeviceSelector {
   virtual DeviceReservation ReserveDevice(
       absl::string_view program_fingerprint) = 0;
 
+  // Enqueues a program on the given device. Used only for load tracking
+  // purposes when the device selection feature is unused.
+  virtual void Enqueue(int32_t device_index, absl::string_view fingerprint) = 0;
+
+  // Marks the completion of a program on the given device. Used only for load
+  // tracking purposes when the device selection feature is unused.
+  virtual void Completed(int32_t device_index, bool had_error) = 0;
+
  protected:
   // A helper function for Enqueue. The EnqueueHelper does the following things.
   //  1. If there are programs in the scheduled_programs queue of the given

diff --git a/xla/tsl/framework/test_util/mock_serving_device_selector.h b/xla/tsl/framework/test_util/mock_serving_device_selector.h
@@ -29,6 +29,8 @@ class MockServingDeviceSelector : public tsl::ServingDeviceSelector {
  public:
   MOCK_METHOD(tsl::DeviceReservation, ReserveDevice, (absl::string_view),
               (override));
+  MOCK_METHOD(void, Enqueue, (int32_t, absl::string_view), (override));
+  MOCK_METHOD(void, Completed, (int32_t, bool), (override));
   MOCK_METHOD(void, FreeDeviceReservation, (const tsl::DeviceReservation&),
               (override));
 };