fix: improve event system for inference servers (#886)

Signed-off-by: axel7083 <[email protected]>

Author: axel7083

packages/backend/src/managers/inference/inferenceManager.spec.ts

@@ -538,3 +538,88 @@ describe('Request Create Inference Server', () => {
     });
   });
 });
+
+describe('containerRegistry events', () => {
+  test('container die event', async () => {
+    mockListContainers([
+      {
+        Id: 'dummyId',
+        engineId: 'dummyEngineId',
+        Labels: {
+          [LABEL_INFERENCE_SERVER]: '[]',
+        },
+      },
+    ]);
+    const disposableMock = vi.fn();
+    const deferred = new Promise<(status: string) => void>((resolve, reject) => {
+      vi.mocked(containerRegistryMock.subscribe).mockImplementation((containerId, listener) => {
+        if (containerId !== 'dummyId') reject(new Error('invalid container id'));
+        else resolve(listener);
+        return {
+          dispose: disposableMock,
+        };
+      });
+    });
+
+    const inferenceManager = await getInitializedInferenceManager();
+    const listener = await deferred;
+
+    const server = inferenceManager.get('dummyId');
+    expect(server.status).toBe('running');
+    expect(containerEngine.inspectContainer).toHaveBeenCalledOnce();
+
+    vi.mocked(containerEngine.inspectContainer).mockResolvedValue({
+      State: {
+        Status: 'stopped',
+        Health: undefined,
+      },
+    } as unknown as ContainerInspectInfo);
+
+    listener('die');
+
+    await vi.waitFor(() => {
+      expect(inferenceManager.get('dummyId').status).toBe('stopped');
+      expect(containerEngine.inspectContainer).toHaveBeenCalledTimes(2);
+    });
+
+    // we should not have disposed the subscriber, as the container is only stopped, not removed
+    expect(disposableMock).not.toHaveBeenCalled();
+  });
+
+  test('container remove event', async () => {
+    mockListContainers([
+      {
+        Id: 'dummyId',
+        engineId: 'dummyEngineId',
+        Labels: {
+          [LABEL_INFERENCE_SERVER]: '[]',
+        },
+      },
+    ]);
+    const disposableMock = vi.fn();
+    const deferred = new Promise<(status: string) => void>((resolve, reject) => {
+      vi.mocked(containerRegistryMock.subscribe).mockImplementation((containerId, listener) => {
+        if (containerId !== 'dummyId') reject(new Error('invalid container id'));
+        else resolve(listener);
+        return {
+          dispose: disposableMock,
+        };
+      });
+    });
+
+    const inferenceManager = await getInitializedInferenceManager();
+    const listener = await deferred;
+
+    const server = inferenceManager.get('dummyId');
+    expect(server.status).toBe('running');
+
+    listener('remove');
+
+    await vi.waitFor(() => {
+      expect(inferenceManager.get('dummyId')).toBeUndefined();
+    });
+
+    // we should have disposed the subscriber, as the container is removed
+    expect(disposableMock).toHaveBeenCalled();
+  });
+});

packages/backend/src/managers/inference/inferenceManager.ts

@@ -287,6 +287,10 @@ export class InferenceManager extends Publisher<InferenceServer[]> implements Di
     // Subscribe to container status update
     const disposable = this.containerRegistry.subscribe(containerId, (status: string) => {
       switch (status) {
+        case 'die':
+          this.updateServerStatus(engineId, containerId);
+          clearInterval(intervalId);
+          break;
         case 'remove':
           // Update the list of servers
           this.removeInferenceServer(containerId);