fix: reclaim queued Pi parent capacity

This commit is contained in:
2026-08-23 13:53:28 +08:00
parent a10b98e484
commit 1727b75f30
3 changed files with 171 additions and 17 deletions

View File

@@ -351,11 +351,20 @@ export class PiWorkerPool {
while (true) {
if (signal?.aborted) throw new Error('Pi idle worker reclaim cancelled');
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
const record = [...this.workers.values()]
const idleRecord = [...this.workers.values()]
.filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
if (record) {
if (await this.evict(record)) return true;
if (idleRecord) {
if (await this.evict(idleRecord)) return true;
continue;
}
const queuedRecord = [...this.workers.values()]
.filter((candidate) => candidate.state === 'queued'
&& candidate.processLease !== null
&& !candidate.processStopFlight)
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
if (queuedRecord) {
if (await this.suspendQueuedWorker(queuedRecord)) return true;
continue;
}
await this.waitForReclaimableWorker(signal);
@@ -501,6 +510,7 @@ export class PiWorkerPool {
this.waitingRuns.push(pending);
pending.queuedAt = this.now();
record.state = 'queued';
this.notifyReclaimableWorker();
return { queuePosition: this.waitingRuns.length, accepted };
}
@@ -550,7 +560,7 @@ export class PiWorkerPool {
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
if (this.workers.get(conversationId) === record) this.workers.delete(conversationId);
await this.stopAndRelease(record);
await this.ensureStoppedAndReleased(record);
}
private async performShutdown(): Promise<void> {
@@ -569,7 +579,7 @@ export class PiWorkerPool {
record.unsubscribeInvalidation();
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
await this.stopAndRelease(record);
await this.ensureStoppedAndReleased(record);
}));
}
@@ -651,6 +661,10 @@ export class PiWorkerPool {
if (!current) throw new Error('Conversation worker is no longer available');
if (current !== record) return await this.ensureFresh(current);
if (record.rebuildFlight) return await record.rebuildFlight;
if (!record.processLease) {
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
return await record.rebuildFlight;
}
const action = this.revisions.beforePrompt(record.revisionWorkerId);
if (action.action !== 'rebuild-before-prompt') return record;
record.rebuildFlight = this.beginRebuild(record, action.revision);
@@ -817,7 +831,23 @@ export class PiWorkerPool {
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
this.workers.delete(conversationId);
await this.stopAndRelease(record);
await this.ensureStoppedAndReleased(record);
return true;
}
private async suspendQueuedWorker(record: WorkerRecord): Promise<boolean> {
const conversationId = record.conversation.conversationId;
if (this.workers.get(conversationId) !== record
|| record.state !== 'queued'
|| !record.processLease
|| record.processStopFlight) {
return false;
}
record.unsubscribeEvent();
record.unsubscribeInvalidation();
this.cancelGenerationResources(record);
record.processStopFlight = this.stopAndRelease(record);
await record.processStopFlight;
return true;
}
@@ -924,6 +954,13 @@ export class PiWorkerPool {
}
}
private async ensureStoppedAndReleased(record: WorkerRecord): Promise<void> {
if (!record.processStopFlight) {
record.processStopFlight = this.stopAndRelease(record);
}
await record.processStopFlight;
}
private waitForReclaimableWorker(signal?: AbortSignal): Promise<void> {
return new Promise((resolve, reject) => {
const cleanup = () => {