wangxuming/makelore
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix(pi): converge worker failures and thinking state

Browse Source
This commit is contained in:
brother7
2026-08-25 11:45:14 +08:00
parent 274187e3cf
commit 61817b161f
28 changed files with 2019 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
electron/coding-runtime/pi/worker-pool.ts
View File

@@ -10,7 +10,11 @@ import type {
PiRpcRequestOptions,
PiRpcResponse,
} from './rpc-client';
import type { PiWorkerStopResult } from './worker-process';
import type {
PiWorkerProofFailure,
PiWorkerStopReason,
PiWorkerStopResult,
} from './worker-process';
import {
PiManagedInputRevisionCoordinator,
type PiManagedInputRevision,
@@ -19,6 +23,7 @@ import {
createPiRuntimeTelemetryEvent,
type PiRuntimeTelemetryEvent,
} from './telemetry';
import { logger } from '../../utils/logger';
export interface PiConversationWorker {
readonly id: string;
@@ -30,7 +35,8 @@ export interface PiConversationWorker {
send(command: PiRpcCommand): Promise<void>;
subscribe(listener: (event: PiRpcEvent) => void): () => void;
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
stop(): Promise<PiWorkerStopResult>;
stop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult>;
injectFailureForProof?(failure: PiWorkerProofFailure): Promise<void>;
}
export interface PiWorkerSessionBinding {
@@ -100,9 +106,17 @@ export type PiWorkerPoolEvent =
type: 'worker.replaced';
conversationId: string;
generation: number;
reason: PiWorkerReplacementReason;
state: PiWorkerPoolState;
};
export type PiWorkerReplacementReason =
| 'stale_resource_rebuild'
| 'recover'
| 'model_reconfiguration'
| 'process_capacity_reopen'
| 'fork_replacement';
export interface PiWorkerPoolOptions {
openWorker(input: PiWorkerOpenInput): Promise<PiWorkerOpenResult>;
maxRunning?: number;
@@ -365,6 +379,22 @@ export class PiWorkerPool {
return record ? this.publicState(record) : null;
}
async injectFailureForProof(
conversationId: string,
failure: PiWorkerProofFailure,
): Promise<{ generation: number }> {
const record = this.workers.get(conversationId);
if (!record || record.state === 'crashed') {
throw new Error('Conversation worker is not available for proof failure injection');
}
if (!record.worker.injectFailureForProof) {
throw new Error('Conversation worker does not support proof failure injection');
}
const generation = record.generation;
await record.worker.injectFailureForProof(failure);
return { generation };
}
getDiagnostics(): CodingRuntimeDiagnostics {
const stage = (
state: PiWorkerPoolState['state'],
@@ -386,6 +416,22 @@ export class PiWorkerPool {
};
}
getResilienceProofDiagnostics(): {
processBudget: { active: number; waiting: number };
runs: { active: number; waiting: number };
} {
return {
processBudget: {
active: this.processBudget.activeCount,
waiting: this.processBudget.waitingCount,
},
runs: {
active: this.activeRuns.size,
waiting: this.waitingRuns.length,
},
};
}
async reclaimIdleWorker(signal?: AbortSignal): Promise<boolean> {
while (true) {
if (signal?.aborted) throw new Error('Pi idle worker reclaim cancelled');
@@ -505,7 +551,11 @@ export class PiWorkerPool {
record.reconfigureAfterSettled = true;
return null;
}
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
record.rebuildFlight = this.beginRebuild(
record,
this.revisions.current,
'model_reconfiguration',
);
return this.publicState(await record.rebuildFlight);
}
@@ -581,7 +631,7 @@ export class PiWorkerPool {
break;
}
}
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
record.rebuildFlight = this.beginRebuild(record, this.revisions.current, 'recover');
return this.publicState(await record.rebuildFlight);
}
@@ -599,7 +649,7 @@ export class PiWorkerPool {
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
if (this.workers.get(conversationId) === record) this.workers.delete(conversationId);
await this.ensureStoppedAndReleased(record);
await this.ensureStoppedAndReleased(record, 'dispose');
}
private async performShutdown(): Promise<void> {
@@ -618,7 +668,7 @@ export class PiWorkerPool {
record.unsubscribeInvalidation();
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
await this.ensureStoppedAndReleased(record);
await this.ensureStoppedAndReleased(record, 'app_shutdown');
}));
}
@@ -710,6 +760,9 @@ export class PiWorkerPool {
revisionAction.action === 'rebuild-after-settled'
? revisionAction.revision
: this.revisions.current,
revisionAction.action === 'rebuild-after-settled'
? 'stale_resource_rebuild'
: 'model_reconfiguration',
);
void record.rebuildFlight.catch(() => undefined);
}
@@ -731,20 +784,31 @@ export class PiWorkerPool {
if (record.rebuildFlight) return await record.rebuildFlight;
}
if (!record.processLease) {
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
record.rebuildFlight = this.beginRebuild(
record,
this.revisions.current,
'process_capacity_reopen',
);
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);
record.rebuildFlight = this.beginRebuild(record, action.revision, 'stale_resource_rebuild');
return await record.rebuildFlight;
}
private beginRebuild(
record: WorkerRecord,
revision: PiManagedInputRevision,
reason: PiWorkerReplacementReason,
): Promise<WorkerRecord> {
const flight = this.rebuild(record, revision).finally(() => {
logger.info('[PiWorkerLifecycle]', {
event: 'worker.replacement_started',
conversationId: record.conversation.conversationId,
generation: record.generation,
reason,
});
const flight = this.rebuild(record, revision, reason).finally(() => {
this.rebuildFlights.delete(flight);
});
this.rebuildFlights.add(flight);
@@ -754,6 +818,7 @@ export class PiWorkerPool {
private async rebuild(
record: WorkerRecord,
revision: PiManagedInputRevision,
reason: PiWorkerReplacementReason,
): Promise<WorkerRecord> {
const conversationId = record.conversation.conversationId;
record.state = 'spawning';
@@ -763,7 +828,7 @@ export class PiWorkerPool {
this.revisions.removeWorker(record.revisionWorkerId);
try {
if (record.processStopFlight) await record.processStopFlight;
else await record.worker.stop();
else await record.worker.stop(reason);
let lease = record.processLease;
if (!lease) lease = await this.processBudget.acquire(this.shutdownController.signal);
if (this.shuttingDown) {
@@ -781,7 +846,7 @@ export class PiWorkerPool {
existingSession: structuredClone(record.session),
});
if (this.shuttingDown) {
await opened.worker.stop().catch(() => undefined);
await opened.worker.stop('app_shutdown').catch(() => undefined);
lease.release();
if (record.processLease === lease) record.processLease = null;
throw new Error('Pi worker pool is shutting down');
@@ -793,7 +858,7 @@ export class PiWorkerPool {
}
if (opened.session.piSessionId !== record.session.piSessionId
|| opened.session.sessionKey !== record.session.sessionKey) {
await opened.worker.stop().catch(() => undefined);
await opened.worker.stop('session_binding_mismatch').catch(() => undefined);
lease.release();
if (record.processLease === lease) record.processLease = null;
throw new Error('Reopened Pi worker returned a different session binding');
@@ -817,8 +882,15 @@ export class PiWorkerPool {
type: 'worker.replaced',
conversationId,
generation,
reason,
state: this.publicState(replacement),
});
logger.info('[PiWorkerLifecycle]', {
event: reason === 'recover' ? 'worker.recovered' : 'worker.replacement_ready',
conversationId,
generation,
reason,
});
await this.trimIdleWorkers();
return replacement;
} catch (error) {
@@ -903,7 +975,7 @@ export class PiWorkerPool {
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
this.workers.delete(conversationId);
await this.ensureStoppedAndReleased(record);
await this.ensureStoppedAndReleased(record, 'idle_eviction');
return true;
}
@@ -918,7 +990,7 @@ export class PiWorkerPool {
record.unsubscribeEvent();
record.unsubscribeInvalidation();
this.cancelGenerationResources(record);
record.processStopFlight = this.stopAndRelease(record);
record.processStopFlight = this.stopAndRelease(record, 'queued_suspension');
await record.processStopFlight;
return true;
}
@@ -944,7 +1016,12 @@ export class PiWorkerPool {
this.revisions.removeWorker(record.revisionWorkerId);
this.cancelGenerationResources(record);
if (!record.processStopFlight) {
record.processStopFlight = this.stopAndRelease(record).catch(() => undefined);
record.processStopFlight = this.stopAndRelease(
record,
error.code === 'PI_RPC_PROTOCOL_ERROR'
? 'protocol_invalidation'
: 'unexpected_exit_cleanup',
).catch(() => undefined);
}
const active = this.activeRuns.get(conversationId);
@@ -959,6 +1036,12 @@ export class PiWorkerPool {
pending.reject(error);
}
this.launchWaitingRuns();
logger.warn('[PiWorkerLifecycle]', {
event: 'worker.crashed',
conversationId,
generation: record.generation,
code: error.code,
});
this.emit({
type: 'worker.crashed',
conversationId,
@@ -1006,7 +1089,7 @@ export class PiWorkerPool {
try {
const opened = await this.openWorker(input);
if (this.shuttingDown) {
await opened.worker.stop().catch(() => undefined);
await opened.worker.stop('app_shutdown').catch(() => undefined);
lease.release();
throw new Error('Pi worker pool is shutting down');
}
@@ -1017,18 +1100,21 @@ export class PiWorkerPool {
}
}
private async stopAndRelease(record: WorkerRecord): Promise<void> {
private async stopAndRelease(record: WorkerRecord, reason: PiWorkerStopReason): Promise<void> {
try {
await record.worker.stop();
await record.worker.stop(reason);
} finally {
record.processLease?.release();
record.processLease = null;
}
}
private async ensureStoppedAndReleased(record: WorkerRecord): Promise<void> {
private async ensureStoppedAndReleased(
record: WorkerRecord,
reason: PiWorkerStopReason,
): Promise<void> {
if (!record.processStopFlight) {
record.processStopFlight = this.stopAndRelease(record);
record.processStopFlight = this.stopAndRelease(record, reason);
}
await record.processStopFlight;
}