fix: reclaim queued Pi parent capacity
This commit is contained in:
@@ -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 = () => {
|
||||
|
||||
Reference in New Issue
Block a user