Skip to content

Commit 9f20bd6

Browse files
authored
Use actual cores count (#8576)
1 parent 23dcf10 commit 9f20bd6

File tree

2 files changed

+28
-6
lines changed

2 files changed

+28
-6
lines changed

ydb/core/kqp/runtime/kqp_compute_scheduler.cpp

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,8 @@ class TSchedulerEntity {
126126
std::atomic<i64> DelayedSumBatches = 0;
127127
std::atomic<i64> DelayedCount = 0;
128128

129+
double Share;
130+
129131
TMultithreadPublisher<TGroupMutableStats> MutableStats;
130132
};
131133

@@ -215,13 +217,16 @@ struct TComputeScheduler::TImpl {
215217

216218
TDuration MaxDelay = TDuration::Seconds(10);
217219

218-
void AssignWeights() { }
220+
void AssignWeights() {
221+
for (auto& record : Records) {
222+
record->MutableStats.Next()->Weight = SumCores * record->Share;
223+
}
224+
}
219225

220-
void CreateGroup(TString groupName, double maxShare, NMonotonic::TMonotonic now) {
226+
void CreateGroup(TString groupName, double maxShare) {
221227
PoolId[groupName] = Records.size();
222228
auto group = std::make_unique<TSchedulerEntity::TGroupRecord>();
223-
group->MutableStats.Next()->LastNowRecalc = now;
224-
group->MutableStats.Next()->Weight = maxShare;
229+
group->Share = maxShare;
225230
Records.push_back(std::move(group));
226231
}
227232
};
@@ -369,14 +374,18 @@ bool TComputeScheduler::Disable(TString group, TMonotonic now) {
369374
void TComputeScheduler::UpdateMaxShare(TString group, double share, TMonotonic now) {
370375
auto ptr = Impl->PoolId.FindPtr(group);
371376
if (!ptr) {
372-
Impl->CreateGroup(group, share, now);
377+
Impl->CreateGroup(group, share);
373378
} else {
374379
auto& record = Impl->Records[*ptr];
375-
record->MutableStats.Next()->Weight = share;
380+
record->Share = share;
376381
}
382+
Impl->AssignWeights();
377383
AdvanceTime(now);
378384
}
379385

386+
void TComputeScheduler::SetCapacity(ui64 cores) {
387+
Impl->SumCores = cores;
388+
}
380389

381390
::NMonitoring::TDynamicCounters::TCounterPtr TComputeScheduler::GetGroupUsageCounter(TString group) const {
382391
return Impl->Counters
@@ -418,6 +427,16 @@ class TSchedulerActor : public TActorBootstrapped<TSchedulerActor> {
418427
IEventHandle::FlagTrackDelivery);
419428

420429
Become(&TSchedulerActor::State);
430+
SetCapacity(SelfId().PoolID());
431+
}
432+
433+
void SetCapacity(ui32 pool) {
434+
NActors::TExecutorPoolStats poolStats;
435+
TVector<NActors::TExecutorThreadStats> threadsStats;
436+
TlsActivationContext->ActorSystem()->GetPoolStats(pool, poolStats, threadsStats);
437+
ui64 threads = Max<ui64>(poolStats.MaxThreadCount, 1);
438+
Opts.Counters->SchedulerCapacity->Set(threads);
439+
Opts.Scheduler->SetCapacity(threads);
421440
}
422441

423442
STATEFN(State) {
@@ -467,6 +486,7 @@ class TSchedulerActor : public TActorBootstrapped<TSchedulerActor> {
467486
}
468487

469488
void Handle(TEvents::TEvWakeup::TPtr&) {
489+
SetCapacity(SelfId().PoolID());
470490
Opts.Scheduler->AdvanceTime(TlsActivationContext->Monotonic());
471491
Schedule(Opts.AdvanceTimeInterval, new TEvents::TEvWakeup());
472492
}

ydb/core/kqp/runtime/kqp_compute_scheduler.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ class TComputeScheduler {
6262
void ReportCounters(TIntrusivePtr<TKqpCounters>);
6363

6464
void UpdateMaxShare(TString, double, TMonotonic now);
65+
66+
void SetCapacity(ui64 cores);
6567

6668
void SetMaxDeviation(TDuration);
6769
void SetForgetInterval(TDuration);

0 commit comments

Comments
 (0)