[tf.data service] Disable special handling for local tasks during coo…

…rdinated reads. PiperOrigin-RevId: 397380918 Change-Id: I89fe76de4292624621b7055a6a7c4eed744eeb92
Marbck · Sep 17, 2021 · 6ae7c19 · 6ae7c19
1 parent 7217151
commit 6ae7c19
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/data/experimental/data_service_dataset_op.cc b/tensorflow/core/kernels/data/experimental/data_service_dataset_op.cc
@@ -709,6 +709,9 @@ class DataServiceDatasetOp::Dataset : public DatasetBase {
     }
 
     void UpdateLocalTasks() TF_LOCKS_EXCLUDED(mu_) {
+      if (StrictRoundRobin()) {
+        return;
+      }
       std::vector<std::shared_ptr<Task>> tasks;
       absl::flat_hash_map<std::string, std::shared_ptr<Task>>
           previous_local_tasks, local_tasks;