Driver should have slicing diabled in serial mode

tanjialiang · May 12, 2024 · 4c3610c · 4c3610c
1 parent 2c98308
commit 4c3610c
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 20 deletions.
diff --git a/velox/exec/Driver.cpp b/velox/exec/Driver.cpp
@@ -288,7 +288,7 @@ void Driver::init(
     std::vector<std::unique_ptr<Operator>> operators) {
   VELOX_CHECK_NULL(ctx_);
   ctx_ = std::move(ctx);
-  cpuSliceMs_ = ctx_->queryConfig().driverCpuTimeSliceLimitMs();
+  cpuSliceMs_ = task()->driverCpuTimeSliceLimitMs();
   VELOX_CHECK(operators_.empty());
   operators_ = std::move(operators);
   curOperatorId_ = operators_.size() - 1;

diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
@@ -412,6 +412,12 @@ void Task::removeSpillDirectoryIfExists() {
   }
 }
 
+uint64_t Task::driverCpuTimeSliceLimitMs() const {
+  return mode_ == Task::ExecutionMode::kSerial
+      ? 0
+      : queryCtx_->queryConfig().driverCpuTimeSliceLimitMs();
+}
+
 void Task::initTaskPool() {
   VELOX_CHECK_NULL(pool_);
   pool_ = queryCtx_->pool()->addAggregateChild(

diff --git a/velox/exec/Task.h b/velox/exec/Task.h
@@ -84,6 +84,10 @@ class Task : public std::enable_shared_from_this<Task> {
       ConsumerSupplier consumerSupplier,
       std::function<void(std::exception_ptr)> onError = nullptr);
 
+  /// Convenience function for shortening a Presto taskId. To be used
+  /// in debugging messages and listings.
+  static std::string shortId(const std::string& id);
+
   ~Task();
 
   /// Specify directory to which data will be spilled if spilling is enabled and
@@ -116,9 +120,9 @@ class Task : public std::enable_shared_from_this<Task> {
     return destination_;
   }
 
-  // Convenience function for shortening a Presto taskId. To be used
-  // in debugging messages and listings.
-  static std::string shortId(const std::string& id);
+  /// Configured cpu slice time limit for drivers. 0 (meaning slicing/yield
+  /// disabled) when task is under serial mode.
+  uint64_t driverCpuTimeSliceLimitMs() const;
 
   /// Returns QueryCtx specified in the constructor.
   const std::shared_ptr<core::QueryCtx>& queryCtx() const {

diff --git a/velox/exec/tests/DriverTest.cpp b/velox/exec/tests/DriverTest.cpp
@@ -1474,8 +1474,19 @@ DEBUG_ONLY_TEST_F(DriverTest, driverCpuTimeSlicingCheck) {
         makeRowVector({"c0"}, {makeFlatVector<int32_t>({1, 2, 3})}));
   }
 
-  for (const auto& hasCpuTimeSliceLimit : {false, true}) {
-    SCOPED_TRACE(fmt::format("hasCpuSliceLimit: {}", hasCpuTimeSliceLimit));
+  struct TestParam {
+    bool hasCpuTimeSliceLimit;
+    Task::ExecutionMode executionMode;
+  };
+  std::vector<TestParam> testParams{
+      {true, Task::ExecutionMode::kParallel},
+      {false, Task::ExecutionMode::kParallel},
+      {true, Task::ExecutionMode::kSerial},
+      {false, Task::ExecutionMode::kSerial}};
+
+  for (const auto& testParam : testParams) {
+    SCOPED_TRACE(
+        fmt::format("hasCpuSliceLimit: {}", testParam.hasCpuTimeSliceLimit));
     SCOPED_TESTVALUE_SET(
         "facebook::velox::exec::Values::getOutput",
         std::function<void(const exec::Values*)>([&](const exec::Values*
@@ -1485,7 +1496,7 @@ DEBUG_ONLY_TEST_F(DriverTest, driverCpuTimeSlicingCheck) {
           ASSERT_NE(
               values->testingOperatorCtx()->driver()->state().startExecTimeMs,
               0);
-          if (hasCpuTimeSliceLimit) {
+          if (testParam.hasCpuTimeSliceLimit) {
             std::this_thread::sleep_for(std::chrono::seconds(1)); // NOLINT
             ASSERT_GT(
                 values->testingOperatorCtx()->driver()->state().execTimeMs(),
@@ -1496,23 +1507,39 @@ DEBUG_ONLY_TEST_F(DriverTest, driverCpuTimeSlicingCheck) {
     auto fragment =
         PlanBuilder(planNodeIdGenerator).values(batches).planFragment();
     std::unordered_map<std::string, std::string> queryConfig;
-    if (hasCpuTimeSliceLimit) {
+    if (testParam.hasCpuTimeSliceLimit) {
       queryConfig.emplace(core::QueryConfig::kDriverCpuTimeSliceLimitMs, "500");
     }
     const uint64_t oldYieldCount = Driver::yieldCount();
-    auto task = Task::create(
-        "t0",
-        fragment,
-        0,
-        std::make_shared<core::QueryCtx>(
-            driverExecutor_.get(), std::move(queryConfig)),
-        Task::ExecutionMode::kParallel,
-        [](RowVectorPtr /*unused*/, ContinueFuture* /*unused*/) {
-          return exec::BlockingReason::kNotBlocked;
-        });
-    task->start(1, 1);
+
+    std::shared_ptr<Task> task;
+    if (testParam.executionMode == Task::ExecutionMode::kParallel) {
+      task = Task::create(
+          "t0",
+          fragment,
+          0,
+          std::make_shared<core::QueryCtx>(
+              driverExecutor_.get(), std::move(queryConfig)),
+          testParam.executionMode,
+          [](RowVectorPtr /*unused*/, ContinueFuture* /*unused*/) {
+            return exec::BlockingReason::kNotBlocked;
+          });
+      task->start(1, 1);
+    } else {
+      task = Task::create(
+          "t0",
+          fragment,
+          0,
+          std::make_shared<core::QueryCtx>(
+              driverExecutor_.get(), std::move(queryConfig)),
+          testParam.executionMode);
+      while (task->next() != nullptr) {
+      }
+    }
+
     ASSERT_TRUE(waitForTaskCompletion(task.get(), 600'000'000));
-    if (hasCpuTimeSliceLimit) {
+    if (testParam.hasCpuTimeSliceLimit &&
+        testParam.executionMode == Task::ExecutionMode::kParallel) {
       // NOTE: there is one additional yield for the empty output.
       ASSERT_GE(Driver::yieldCount(), oldYieldCount + numBatches + 1);
     } else {