relax local compute tasks

ydb-platform · May 3, 2024 · cb8cfb3 · cb8cfb3
1 parent eebd1ab
commit cb8cfb3
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 2 deletions.
diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -2408,7 +2408,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
 
         const bool singlePartitionOptAllowed = !HasOlapTable && !UnknownAffectedShardCount && !HasExternalSources && DatashardTxs.empty() && EvWriteTxs.empty();
         const bool useDataQueryPool = !(HasExternalSources && DatashardTxs.empty() && EvWriteTxs.empty());
-        const bool localComputeTasks = !((HasExternalSources || HasOlapTable || HasDatashardSourceScan) && DatashardTxs.empty());
+        const bool localComputeTasks = !DatashardTxs.empty();
 
         Planner = CreateKqpPlanner({
             .TasksGraph = TasksGraph,

diff --git a/ydb/core/kqp/executer_actor/kqp_planner.cpp b/ydb/core/kqp/executer_actor/kqp_planner.cpp
@@ -250,6 +250,10 @@ std::unique_ptr<IEventHandle> TKqpPlanner::AssignTasksToNodes() {
         return nullptr;
     }
 
+    if (ResourcesSnapshot.empty()) {
+        ResourcesSnapshot = std::move(GetKqpResourceManager()->GetClusterResources());
+    }
+
     if (ResourcesSnapshot.empty() || (ResourcesSnapshot.size() == 1 && ResourcesSnapshot[0].GetNodeId() == ExecuterId.NodeId())) {
         // try to run without memory overflow settings
         if (LocalRunMemoryEst <= localResources.Memory[NRm::EKqpMemoryPool::ScanQuery] &&
@@ -402,6 +406,8 @@ std::unique_ptr<IEventHandle> TKqpPlanner::PlanExecution() {
 
     nComputeTasks = ComputeTasks.size();
 
+    // explicit requirement to execute task on the same node because it has dependencies
+    // on datashard tx.
     if (LocalComputeTasks) {
         bool shareMailbox = (ComputeTasks.size() <= 1);
         for (ui64 taskId : ComputeTasks) {

diff --git a/ydb/core/kqp/executer_actor/kqp_planner.h b/ydb/core/kqp/executer_actor/kqp_planner.h
@@ -103,7 +103,7 @@ class TKqpPlanner {
     const bool WithSpilling;
     const TMaybe<NKikimrKqp::TRlPath> RlPath;
     THashSet<ui32> TrackingNodes;
-    const TVector<NKikimrKqp::TKqpNodeResources> ResourcesSnapshot;
+    TVector<NKikimrKqp::TKqpNodeResources> ResourcesSnapshot;
     NWilson::TSpan& ExecuterSpan;
     const NKikimrConfig::TTableServiceConfig::TExecuterRetriesConfig& ExecuterRetriesConfig;
     ui64 LocalRunMemoryEst = 0;

diff --git a/ydb/core/kqp/rm_service/kqp_rm_service.cpp b/ydb/core/kqp/rm_service/kqp_rm_service.cpp
@@ -599,6 +599,23 @@ class TKqpResourceManager : public IKqpResourceManager {
         FireResourcesPublishing();
     }
 
+    TVector<NKikimrKqp::TKqpNodeResources> GetClusterResources() const override {
+        TVector<NKikimrKqp::TKqpNodeResources> resources;
+        Y_ABORT_UNLESS(PublishResourcesByExchanger);
+
+        if (PublishResourcesByExchanger) {
+            std::shared_ptr<TVector<NKikimrKqp::TKqpNodeResources>> infos;
+            with_lock (ResourceSnapshotState->Lock) {
+                infos = ResourceSnapshotState->Snapshot;
+            }
+            if (infos != nullptr) {
+                resources = *infos;
+            }
+        }
+
+        return resources;
+    }
+
     void RequestClusterResourcesInfo(TOnResourcesSnapshotCallback&& callback) override {
         LOG_AS_D("Schedule Snapshot request");
         if (PublishResourcesByExchanger) {

diff --git a/ydb/core/kqp/rm_service/kqp_rm_service.h b/ydb/core/kqp/rm_service/kqp_rm_service.h
@@ -91,6 +91,7 @@ class IKqpResourceManager : private TNonCopyable {
 
     virtual void RequestClusterResourcesInfo(TOnResourcesSnapshotCallback&& callback) = 0;
 
+    virtual TVector<NKikimrKqp::TKqpNodeResources> GetClusterResources() const = 0;
     virtual TKqpLocalNodeResources GetLocalResources() const = 0;
     virtual NKikimrConfig::TTableServiceConfig::TResourceManager GetConfig() = 0;