Skip to content

Commit

Permalink
fix: listen to correct host port when checking model service (#168)
Browse files Browse the repository at this point in the history 
Signed-off-by: lstocchi <[email protected]>
  • Loading branch information
@lstocchi
lstocchi authored Jan 29, 2024
1 parent 6ba8fbd commit c9c97bb
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 42 deletions.
58 changes: 50 additions & 8 deletions packages/backend/src/managers/applicationManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ const mocks = vi.hoisted(() => {
startContainerMock: vi.fn(),
startPod: vi.fn(),
deleteContainerMock: vi.fn(),
inspectContainerMock: vi.fn(),
};
});
vi.mock('../models/AIConfig', () => ({
Expand All @@ -43,6 +44,7 @@ vi.mock('@podman-desktop/api', () => ({
startContainer: mocks.startContainerMock,
startPod: mocks.startPod,
deleteContainer: mocks.deleteContainerMock,
inspectContainer: mocks.inspectContainerMock,
},
}));
let setTaskMock: MockInstance;
Expand All @@ -69,6 +71,7 @@ describe('pullApplication', () => {
[modelId: string, url: string, taskUtil: RecipeStatusUtils, destFileName?: string],
Promise<string>
>;
vi.spyOn(utils, 'timeout').mockResolvedValue();
function mockForPullApplication(options: mockForPullApplicationOptions) {
vi.spyOn(os, 'homedir').mockReturnValue('/home/user');
vi.spyOn(fs, 'mkdirSync').mockReturnValue(undefined);
Expand Down Expand Up @@ -109,6 +112,11 @@ describe('pullApplication', () => {
],
},
});
mocks.inspectContainerMock.mockResolvedValue({
State: {
Running: true,
},
});
mocks.builImageMock.mockResolvedValue(undefined);
mocks.listImagesMock.mockResolvedValue([
{
Expand Down Expand Up @@ -170,6 +178,11 @@ describe('pullApplication', () => {
registry: '',
url: '',
};
mocks.inspectContainerMock.mockResolvedValue({
State: {
Running: true,
},
});
await manager.pullApplication(recipe, model);
if (process.platform === 'win32') {
expect(cloneRepositoryMock).toHaveBeenNthCalledWith(1, 'repo', '\\home\\user\\aistudio\\recipe1');
Expand Down Expand Up @@ -664,6 +677,10 @@ describe('createPod', async () => {
const images = [imageInfo1, imageInfo2];
vi.spyOn(manager, 'getRandomName').mockReturnValue('name');
vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValue('9000');
mocks.createPodMock.mockResolvedValue({
Id: 'podId',
engineId: 'engineId',
});
await manager.createPod(images);
expect(mocks.createPodMock).toBeCalledWith({
name: 'name',
Expand Down Expand Up @@ -720,6 +737,7 @@ describe('createApplicationPod', () => {
const pod: PodInfo = {
engineId: 'engine',
Id: 'id',
portmappings: [],
};
vi.spyOn(manager, 'createPod').mockResolvedValue(pod);
const createAndAddContainersToPodMock = vi
Expand All @@ -737,6 +755,7 @@ describe('createApplicationPod', () => {
const pod: PodInfo = {
engineId: 'engine',
Id: 'id',
portmappings: [],
};
vi.spyOn(manager, 'createPod').mockResolvedValue(pod);
vi.spyOn(manager, 'createAndAddContainersToPod').mockRejectedValue('error');
Expand Down Expand Up @@ -795,10 +814,11 @@ describe('doDownloadModelWrapper', () => {
});
});

describe('restartContainerWhenEndpointIsUp', () => {
describe('restartContainerWhenModelServiceIsUp', () => {
const containerAttachedInfo: ContainerAttachedInfo = {
name: 'name',
endPoint: 'endpoint',
modelService: false,
ports: ['9000'],
};
const manager = new ApplicationManager(
'/home/user/aistudio',
Expand All @@ -808,7 +828,7 @@ describe('restartContainerWhenEndpointIsUp', () => {
);
test('restart container if endpoint is alive', async () => {
vi.spyOn(utils, 'isEndpointAlive').mockResolvedValue(true);
await manager.restartContainerWhenEndpointIsUp('engine', containerAttachedInfo);
await manager.restartContainerWhenModelServiceIsUp('engine', 'endpoint', containerAttachedInfo);
expect(mocks.startContainerMock).toBeCalledWith('engine', 'name');
});
});
Expand All @@ -826,22 +846,43 @@ describe('runApplication', () => {
containers: [
{
name: 'first',
endPoint: 'url',
modelService: false,
ports: ['8080'],
},
{
name: 'second',
modelService: true,
ports: ['9000'],
},
],
portmappings: [
{
container_port: 9000,
host_port: 9001,
host_ip: '',
protocol: '',
range: -1,
},
],
};
test('check startPod is called and also restartContainerWhenEndpointIsUp for sample app', async () => {
const restartContainerWhenEndpointIsUpMock = vi
.spyOn(manager, 'restartContainerWhenEndpointIsUp')
.mockImplementation((_engineId: string, _container: ContainerAttachedInfo) => Promise.resolve());
.spyOn(manager, 'restartContainerWhenModelServiceIsUp')
.mockImplementation((_engineId: string, _modelEndpoint: string, _container: ContainerAttachedInfo) =>
Promise.resolve(),
);
vi.spyOn(utils, 'timeout').mockResolvedValue();
mocks.inspectContainerMock.mockResolvedValue({
State: {
Running: false,
},
});
await manager.runApplication(pod, taskUtils);
expect(mocks.startPod).toBeCalledWith(pod.engineId, pod.Id);
expect(restartContainerWhenEndpointIsUpMock).toBeCalledWith(pod.engineId, {
expect(restartContainerWhenEndpointIsUpMock).toBeCalledWith(pod.engineId, 'http://localhost:9001', {
name: 'first',
endPoint: 'url',
modelService: false,
ports: ['8080'],
});
});
});
Expand All @@ -856,6 +897,7 @@ describe('createAndAddContainersToPod', () => {
const pod: PodInfo = {
engineId: 'engine',
Id: 'id',
portmappings: [],
};
const imageInfo1: ImageInfo = {
id: 'id',
Expand Down
95 changes: 61 additions & 34 deletions packages/backend/src/managers/applicationManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,13 +55,15 @@ interface AIContainers {

export interface ContainerAttachedInfo {
name: string;
endPoint?: string;
modelService: boolean;
ports: string[];
}

export interface PodInfo {
engineId: string;
Id: string;
containers?: ContainerAttachedInfo[];
portmappings: PodCreatePortOptions[];
}

export interface ImageInfo {
Expand Down Expand Up @@ -120,16 +122,30 @@ export class ApplicationManager {

// most probably the sample app will fail at starting as it tries to connect to the model_service which is still loading the model
// so we check if the endpoint is ready before to restart the sample app
await Promise.all(
podInfo.containers?.map(async container => {
if (!container.endPoint) {
return;
// wait 5 secs to give the time to the sample app to connect to the model service
await timeout(5000);
// check if sample app container actually started
const sampleApp = podInfo.containers?.find(container => !container.modelService);
if (sampleApp) {
const sampleAppContainerInspectInfo = await containerEngine.inspectContainer(podInfo.engineId, sampleApp.name);
if (!sampleAppContainerInspectInfo.State.Running) {
const modelServiceContainer = podInfo.containers?.find(container => container.modelService);
if (modelServiceContainer) {
const modelServicePortMapping = podInfo.portmappings.find(
port => port.container_port === Number(modelServiceContainer.ports[0]),
);
if (modelServicePortMapping) {
await this.restartContainerWhenModelServiceIsUp(
podInfo.engineId,
`http://localhost:${modelServicePortMapping.host_port}`,
sampleApp,
).catch((e: unknown) => {
console.error(String(e));
});
}
}
return this.restartContainerWhenEndpointIsUp(podInfo.engineId, container).catch((e: unknown) => {
console.error(String(e));
});
}),
);
}
}

taskUtil.setTask({
id: `running-${podInfo.Id}`,
Expand All @@ -138,23 +154,29 @@ export class ApplicationManager {
});
}

async restartContainerWhenEndpointIsUp(engineId: string, container: ContainerAttachedInfo): Promise<void> {
const alive = await isEndpointAlive(container.endPoint);
async restartContainerWhenModelServiceIsUp(
engineId: string,
modelServiceEndpoint: string,
container: ContainerAttachedInfo,
): Promise<void> {
const alive = await isEndpointAlive(modelServiceEndpoint);
if (alive) {
await containerEngine.startContainer(engineId, container.name);
return;
}
await timeout(5000);
await this.restartContainerWhenEndpointIsUp(engineId, container).catch((error: unknown) => {
console.error('Error monitoring endpoint', error);
});
await this.restartContainerWhenModelServiceIsUp(engineId, modelServiceEndpoint, container).catch(
(error: unknown) => {
console.error('Error monitoring endpoint', error);
},
);
}

async createApplicationPod(images: ImageInfo[], modelPath: string, taskUtil: RecipeStatusUtils): Promise<PodInfo> {
// create empty pod
let pod: PodInfo;
let podInfo: PodInfo;
try {
pod = await this.createPod(images);
podInfo = await this.createPod(images);
} catch (e) {
console.error('error when creating pod');
taskUtil.setTask({
Expand All @@ -166,36 +188,36 @@ export class ApplicationManager {
}

taskUtil.setTask({
id: pod.Id,
id: podInfo.Id,
state: 'loading',
name: `Creating application`,
});

let attachedContainers: ContainerAttachedInfo[];
try {
attachedContainers = await this.createAndAddContainersToPod(pod, images, modelPath);
attachedContainers = await this.createAndAddContainersToPod(podInfo, images, modelPath);
} catch (e) {
console.error(`error when creating pod ${pod.Id}`);
console.error(`error when creating pod ${podInfo.Id}`);
taskUtil.setTask({
id: pod.Id,
id: podInfo.Id,
state: 'error',
name: 'Creating application',
});
throw e;
}

taskUtil.setTask({
id: pod.Id,
id: podInfo.Id,
state: 'success',
name: `Creating application`,
});

pod.containers = attachedContainers;
return pod;
podInfo.containers = attachedContainers;
return podInfo;
}

async createAndAddContainersToPod(
pod: PodInfo,
podInfo: PodInfo,
images: ImageInfo[],
modelPath: string,
): Promise<ContainerAttachedInfo[]> {
Expand All @@ -204,7 +226,6 @@ export class ApplicationManager {
images.map(async image => {
let hostConfig: unknown;
let envs: string[] = [];
let endPoint: string;
// if it's a model service we mount the model as a volume
if (image.modelService) {
const modelName = path.basename(modelPath);
Expand All @@ -226,11 +247,11 @@ export class ApplicationManager {
// TODO: remove static port
const modelService = images.find(image => image.modelService);
if (modelService && modelService.ports.length > 0) {
endPoint = `http://localhost:${modelService.ports[0]}`;
const endPoint = `http://localhost:${modelService.ports[0]}`;
envs = [`MODEL_ENDPOINT=${endPoint}`];
}
}
const createdContainer = await containerEngine.createContainer(pod.engineId, {
const createdContainer = await containerEngine.createContainer(podInfo.engineId, {
Image: image.id,
Detach: true,
HostConfig: hostConfig,
Expand All @@ -244,17 +265,18 @@ export class ApplicationManager {
await containerEngine.replicatePodmanContainer(
{
id: createdContainer.id,
engineId: pod.engineId,
engineId: podInfo.engineId,
},
{ engineId: pod.engineId },
{ pod: pod.Id, name: podifiedName },
{ engineId: podInfo.engineId },
{ pod: podInfo.Id, name: podifiedName },
);
containers.push({
name: podifiedName,
endPoint,
modelService: image.modelService,
ports: image.ports,
});
// remove the external container
await containerEngine.deleteContainer(pod.engineId, createdContainer.id);
await containerEngine.deleteContainer(podInfo.engineId, createdContainer.id);
} else {
throw new Error(`failed at creating container for image ${image.id}`);
}
Expand Down Expand Up @@ -288,10 +310,15 @@ export class ApplicationManager {
}

// create new pod
return await containerEngine.createPod({
const pod = await containerEngine.createPod({
name: this.getRandomName(`pod-${sampleAppImageInfo.appName}`),
portmappings: portmappings,
});
return {
Id: pod.Id,
engineId: pod.engineId,
portmappings: portmappings,
};
}

getRandomName(base: string): string {
Expand Down

0 comments on commit c9c97bb

Please sign in to comment.