koriweb
6d06311d60
VRAM 부족으로 12b 등 다른 모델 로드 실패하던 문제 강화. - lifecycleManager.doSwitch: 추적 중인 currentModel 만이 아니라 listLoaded() 기반으로 *로드된 모든 LLM* 을 타깃 전 언로드(VRAM 회수). draft 모델·임베딩 모델은 보호. listLoaded 실패 시 기존 동작(tracked unload)으로 폴백. - extension.ts: defaultModel 설정 변경(설정 패널/settings.json 포함) 시 lifecycle.onModelSelected 호출 → 설정 패널 전환도 unload→load 발동. - 테스트 FakeLMStudioClient 가 실제 로드 상태를 추적하도록 갱신. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
331 lines
12 KiB
TypeScript
331 lines
12 KiB
TypeScript
/**
|
|
* Unit tests for ModelLifecycleManager.
|
|
*
|
|
* Strategy: inject mock ILMStudioClient and a simple in-memory IActivityTracker.
|
|
* No real LM Studio or SDK is touched — the manager file does not import the
|
|
* SDK directly (only types via `import type`).
|
|
*/
|
|
|
|
import {
|
|
ModelLifecycleManager,
|
|
LifecycleConfig,
|
|
LifecycleManagerDeps,
|
|
} from '../src/lmstudio/lifecycleManager';
|
|
import type { ILMStudioClient } from '../src/lmstudio/client';
|
|
import type { IActivityTracker } from '../src/lmstudio/activityTracker';
|
|
|
|
class FakeActivityTracker implements IActivityTracker {
|
|
private listeners: Array<(_: void) => any> = [];
|
|
public readonly onActivity = ((listener: (_: void) => any) => {
|
|
this.listeners.push(listener);
|
|
return { dispose: () => { this.listeners = this.listeners.filter(l => l !== listener); } };
|
|
}) as any;
|
|
bump(): void {
|
|
for (const l of this.listeners.slice()) l();
|
|
}
|
|
}
|
|
|
|
class FakeLMStudioClient implements ILMStudioClient {
|
|
public loadCalls: string[] = [];
|
|
public unloadCalls: string[] = [];
|
|
public listLoadedCalls = 0;
|
|
public failNextLoad: Error | null = null;
|
|
public failNextUnload: Error | null = null;
|
|
public loadDelayMs = 0;
|
|
public lastLoadSignal: AbortSignal | undefined;
|
|
/** 실제 로드 상태 추적 — listLoaded()가 이를 반영해야 lifecycle 의 '전체 언로드'를 검증할 수 있다. */
|
|
public loaded = new Set<string>();
|
|
|
|
setBaseUrl(_: string): void { /* noop */ }
|
|
|
|
async load(modelKey: string, signal?: AbortSignal): Promise<void> {
|
|
this.loadCalls.push(modelKey);
|
|
this.lastLoadSignal = signal;
|
|
if (this.loadDelayMs > 0) {
|
|
await new Promise<void>((resolve, reject) => {
|
|
const t = setTimeout(resolve, this.loadDelayMs);
|
|
if (signal) {
|
|
const onAbort = () => { clearTimeout(t); reject(new Error('aborted')); };
|
|
if (signal.aborted) onAbort();
|
|
else signal.addEventListener('abort', onAbort);
|
|
}
|
|
});
|
|
}
|
|
if (this.failNextLoad) {
|
|
const err = this.failNextLoad;
|
|
this.failNextLoad = null;
|
|
throw err;
|
|
}
|
|
this.loaded.add(modelKey); // 성공 시에만 로드 상태로
|
|
}
|
|
|
|
async unload(modelKey: string): Promise<void> {
|
|
this.unloadCalls.push(modelKey);
|
|
if (this.failNextUnload) {
|
|
const err = this.failNextUnload;
|
|
this.failNextUnload = null;
|
|
throw err; // 실패 시 로드 상태 유지
|
|
}
|
|
this.loaded.delete(modelKey);
|
|
}
|
|
|
|
async listLoaded(): Promise<string[]> {
|
|
this.listLoadedCalls++;
|
|
return [...this.loaded];
|
|
}
|
|
|
|
async isReachable(): Promise<boolean> {
|
|
return true;
|
|
}
|
|
|
|
async listLoadedCached(): Promise<string[]> {
|
|
return [...this.loaded];
|
|
}
|
|
|
|
async listDownloaded(): Promise<string[]> {
|
|
return [];
|
|
}
|
|
|
|
async listDownloadedCached(): Promise<string[]> {
|
|
return [];
|
|
}
|
|
|
|
async getModelHandle(_modelKey: string): Promise<any> {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
function makeManager(overrides: Partial<LifecycleConfig> = {}, depOverrides: Partial<LifecycleManagerDeps> = {}) {
|
|
const client = new FakeLMStudioClient();
|
|
const activity = new FakeActivityTracker();
|
|
const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true, ...overrides };
|
|
const errors: string[] = [];
|
|
const manager = new ModelLifecycleManager({
|
|
client,
|
|
activity,
|
|
getConfig: () => config,
|
|
notifyError: (m) => errors.push(m),
|
|
switchDebounceMs: 0,
|
|
initialEngine: 'lmstudio',
|
|
...depOverrides,
|
|
});
|
|
return { manager, client, activity, config, errors };
|
|
}
|
|
|
|
const flush = () => new Promise<void>((r) => setImmediate(r));
|
|
|
|
describe('ModelLifecycleManager', () => {
|
|
beforeEach(() => {
|
|
jest.useFakeTimers({ doNotFake: ['setImmediate'] });
|
|
});
|
|
|
|
afterEach(async () => {
|
|
jest.useRealTimers();
|
|
});
|
|
|
|
test('loads model on selection and arms idle timer', async () => {
|
|
const { manager, client } = makeManager();
|
|
manager.onModelSelected('llama-3.2-3b');
|
|
await flush();
|
|
expect(client.loadCalls).toEqual(['llama-3.2-3b']);
|
|
expect(manager._getState()).toBe('loaded');
|
|
expect(manager._getCurrentModel()).toBe('llama-3.2-3b');
|
|
expect(manager._hasIdleTimer()).toBe(true);
|
|
});
|
|
|
|
test('idle timer triggers unload after timeout', async () => {
|
|
const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(manager._getState()).toBe('loaded');
|
|
|
|
jest.advanceTimersByTime(1000);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
expect(manager._getState()).toBe('idle');
|
|
expect(manager._getCurrentModel()).toBe(null);
|
|
});
|
|
|
|
test('idleTimeoutMs <= 0 disables auto-eject', async () => {
|
|
const { manager, client } = makeManager({ idleTimeoutMs: 0 });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(manager._getState()).toBe('loaded');
|
|
expect(manager._hasIdleTimer()).toBe(false);
|
|
|
|
jest.advanceTimersByTime(1_000_000);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual([]);
|
|
expect(manager._getState()).toBe('loaded');
|
|
});
|
|
|
|
test('activity bump resets idle timer', async () => {
|
|
const { manager, client, activity } = makeManager({ idleTimeoutMs: 1000 });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
|
|
jest.advanceTimersByTime(900);
|
|
activity.bump();
|
|
jest.advanceTimersByTime(900); // total 1800ms but timer reset at 900
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual([]);
|
|
expect(manager._getState()).toBe('loaded');
|
|
|
|
jest.advanceTimersByTime(200); // 200 + 900 since reset = ~1100 since reset
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
});
|
|
|
|
test('streamStart pauses idle timer; streamEnd resumes', async () => {
|
|
const { manager, client } = makeManager({ idleTimeoutMs: 500 });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(manager._hasIdleTimer()).toBe(true);
|
|
|
|
manager.onStreamStart();
|
|
expect(manager._getState()).toBe('streaming');
|
|
expect(manager._hasIdleTimer()).toBe(false);
|
|
|
|
jest.advanceTimersByTime(10000);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual([]); // never ejected during stream
|
|
|
|
manager.onStreamEnd();
|
|
expect(manager._getState()).toBe('loaded');
|
|
expect(manager._hasIdleTimer()).toBe(true);
|
|
|
|
jest.advanceTimersByTime(500);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
});
|
|
|
|
test('model switch unloads previous and loads next', async () => {
|
|
const { manager, client } = makeManager();
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
|
|
manager.onModelSelected('m2');
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
expect(client.loadCalls).toEqual(['m1', 'm2']);
|
|
expect(manager._getCurrentModel()).toBe('m2');
|
|
expect(manager._getState()).toBe('loaded');
|
|
});
|
|
|
|
test('switch during streaming defers until streamEnd', async () => {
|
|
const { manager, client } = makeManager();
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
manager.onStreamStart();
|
|
|
|
manager.onModelSelected('m2');
|
|
await flush();
|
|
// No switch yet — still in streaming with m1
|
|
expect(client.unloadCalls).toEqual([]);
|
|
expect(client.loadCalls).toEqual(['m1']);
|
|
|
|
manager.onStreamEnd();
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
expect(client.loadCalls).toEqual(['m1', 'm2']);
|
|
expect(manager._getCurrentModel()).toBe('m2');
|
|
});
|
|
|
|
test('load failure surfaces via notifyError and resets to idle', async () => {
|
|
const { manager, client, errors } = makeManager();
|
|
client.failNextLoad = new Error('LM Studio is not running');
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(manager._getState()).toBe('idle');
|
|
expect(manager._getCurrentModel()).toBe(null);
|
|
expect(errors.length).toBe(1);
|
|
expect(errors[0]).toContain('LM Studio is not running');
|
|
});
|
|
|
|
test('engine change to non-lmstudio clears state without unloading', async () => {
|
|
const { manager, client } = makeManager();
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(manager._getState()).toBe('loaded');
|
|
|
|
manager.setEngine('ollama');
|
|
expect(manager._getState()).toBe('idle');
|
|
expect(manager._getCurrentModel()).toBe(null);
|
|
expect(manager._hasIdleTimer()).toBe(false);
|
|
expect(client.unloadCalls).toEqual([]); // explicitly does not call unload
|
|
});
|
|
|
|
test('selection while engine is non-lmstudio is a no-op', async () => {
|
|
const { manager, client } = makeManager({}, { initialEngine: 'ollama' });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(client.loadCalls).toEqual([]);
|
|
expect(manager._getState()).toBe('idle');
|
|
});
|
|
|
|
test('autoLoadOnSelect=false skips loading', async () => {
|
|
const { manager, client } = makeManager({ autoLoadOnSelect: false });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(client.loadCalls).toEqual([]);
|
|
expect(manager._getState()).toBe('idle');
|
|
});
|
|
|
|
test('rapid switch debounce: only the last selection wins', async () => {
|
|
// Re-create manager with non-zero debounce
|
|
const client = new FakeLMStudioClient();
|
|
const activity = new FakeActivityTracker();
|
|
const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true };
|
|
const manager = new ModelLifecycleManager({
|
|
client,
|
|
activity,
|
|
getConfig: () => config,
|
|
switchDebounceMs: 300,
|
|
initialEngine: 'lmstudio',
|
|
});
|
|
manager.onModelSelected('m1');
|
|
manager.onModelSelected('m2');
|
|
manager.onModelSelected('m3');
|
|
// Before debounce expires no load fires
|
|
expect(client.loadCalls).toEqual([]);
|
|
jest.advanceTimersByTime(300);
|
|
await flush();
|
|
expect(client.loadCalls).toEqual(['m3']);
|
|
});
|
|
|
|
test('disposeAndUnload ejects loaded model', async () => {
|
|
const { manager, client } = makeManager();
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
await manager.disposeAndUnload(2000);
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
expect(manager._getState()).toBe('idle');
|
|
});
|
|
|
|
test('disposeAndUnload while idle is a no-op', async () => {
|
|
const { manager, client } = makeManager();
|
|
await manager.disposeAndUnload(500);
|
|
expect(client.unloadCalls).toEqual([]);
|
|
});
|
|
|
|
test('selecting same model does not re-load but refreshes timer', async () => {
|
|
const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(client.loadCalls.length).toBe(1);
|
|
|
|
jest.advanceTimersByTime(800);
|
|
manager.onModelSelected('m1');
|
|
await flush();
|
|
expect(client.loadCalls.length).toBe(1); // no reload
|
|
// Timer was reset; should survive past original 1000ms
|
|
jest.advanceTimersByTime(700);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual([]);
|
|
|
|
jest.advanceTimersByTime(400);
|
|
await flush();
|
|
expect(client.unloadCalls).toEqual(['m1']);
|
|
});
|
|
});
|