Skip to content

Commit

Permalink
refactor: some cleanup application manager
Browse files Browse the repository at this point in the history 
Signed-off-by: axel7083 <[email protected]>
  • Loading branch information
@axel7083
axel7083 committed May 21, 2024
1 parent 6744a04 commit b130cbf
Show file tree
Hide file tree
Showing 7 changed files with 139 additions and 189 deletions.
100 changes: 4 additions & 96 deletions packages/backend/src/managers/applicationManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,6 @@ const mocks = vi.hoisted(() => {
getContainerConnectionsMock: vi.fn(),
pullImageMock: vi.fn(),
stopContainerMock: vi.fn(),
getFreePortMock: vi.fn(),
containerRegistrySubscribeMock: vi.fn(),
onPodStartMock: vi.fn(),
onPodStopMock: vi.fn(),
Expand Down Expand Up @@ -655,40 +654,6 @@ describe('filterContainers', () => {
});
});

describe('getRandomName', () => {
test('return base name plus random string', () => {
const manager = new ApplicationManager(
'/home/user/aistudio',
{} as unknown as GitManager,
taskRegistry,
{} as Webview,
{} as PodmanConnection,
{} as CatalogManager,
{} as unknown as ModelsManager,
telemetryLogger,
localRepositoryRegistry,
);
const randomName = manager.getRandomName('base');
expect(randomName).not.equal('base');
expect(randomName.length).toBeGreaterThan(4);
});
test('return random string when base is empty', () => {
const manager = new ApplicationManager(
'/home/user/aistudio',
{} as unknown as GitManager,
taskRegistry,
{} as Webview,
{} as PodmanConnection,
{} as CatalogManager,
{} as unknown as ModelsManager,
telemetryLogger,
localRepositoryRegistry,
);
const randomName = manager.getRandomName('');
expect(randomName.length).toBeGreaterThan(0);
});
});

describe('buildImages', () => {
const recipe = {
id: 'recipe1',
Expand Down Expand Up @@ -807,7 +772,6 @@ describe('createPod', async () => {
});
test('call createPod with sample app exposed port', async () => {
const images = [imageInfo1, imageInfo2];
vi.spyOn(manager, 'getRandomName').mockReturnValue('name');
vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9000');
vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9001');
vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9002');
Expand All @@ -817,7 +781,7 @@ describe('createPod', async () => {
});
await manager.createPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images);
expect(mocks.createPodMock).toBeCalledWith({
name: 'name',
name: expect.anything(),
portmappings: [
{
container_port: 8080,
Expand Down Expand Up @@ -1021,14 +985,13 @@ describe('createAndAddContainersToPod', () => {
id: 'container-1',
});
vi.spyOn(podman, 'isQEMUMachine').mockResolvedValue(false);
vi.spyOn(manager, 'getRandomName').mockReturnValue('name');
await manager.createAndAddContainersToPod(pod, [imageInfo1, imageInfo2], modelInfo, 'path');
expect(mocks.createContainerMock).toHaveBeenNthCalledWith(1, 'engine', {
Image: 'id',
Detach: true,
Env: ['MODEL_ENDPOINT=http://localhost:8085'],
start: false,
name: 'name',
name: expect.anything(),
pod: 'id',
HealthCheck: {
Interval: 5000000000,
Expand All @@ -1042,7 +1005,7 @@ describe('createAndAddContainersToPod', () => {
Detach: true,
Env: ['MODEL_PATH=/path', ...extraEnvs],
start: false,
name: 'name',
name: expect.anything(),
pod: 'id',
HostConfig: {
Mounts: [
Expand Down Expand Up @@ -1354,62 +1317,6 @@ describe('pod detection', async () => {
expect(mocks.stopPodMock).toHaveBeenCalledWith('engine-1', 'pod-1');
expect(mocks.removePodMock).toHaveBeenCalledWith('engine-1', 'pod-1');
});
});

describe('getImageTag', () => {
const manager = new ApplicationManager(
'/path/to/user/dir',
{} as GitManager,
taskRegistry,
{
postMessage: mocks.postMessageMock,
} as unknown as Webview,
{
onPodStart: mocks.onPodStartMock,
onPodStop: mocks.onPodStopMock,
onPodRemove: mocks.onPodRemoveMock,
startupSubscribe: mocks.startupSubscribeMock,
onMachineStop: mocks.onMachineStopMock,
} as unknown as PodmanConnection,
{
getRecipeById: vi.fn().mockReturnValue({ name: 'MyRecipe' } as Recipe),
} as unknown as CatalogManager,
{} as ModelsManager,
{} as TelemetryLogger,
localRepositoryRegistry,
);
test('return recipe-container tag if container image prop is not defined', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
} as ContainerConfig;
const imageTag = manager.getImageTag(recipe, container);
expect(imageTag).equals('recipe1-name:latest');
});
test('return container image prop is defined', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
image: 'quay.io/repo/image:v1',
} as ContainerConfig;
const imageTag = manager.getImageTag(recipe, container);
expect(imageTag).equals('quay.io/repo/image:v1');
});
test('append latest tag to container image prop if it has no tag', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
image: 'quay.io/repo/image',
} as ContainerConfig;
const imageTag = manager.getImageTag(recipe, container);
expect(imageTag).equals('quay.io/repo/image:latest');
});

test('adoptRunningApplications should check pods health', async () => {
mocks.listPodsMock.mockResolvedValue([
Expand Down Expand Up @@ -1474,3 +1381,4 @@ describe('getImageTag', () => {
expect(state[0].health).toEqual('healthy');
});
});

75 changes: 15 additions & 60 deletions packages/backend/src/managers/applicationManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import type { Task } from '@shared/src/models/ITask';
import { getParentDirectory } from '../utils/pathUtils';
import type { ModelInfo } from '@shared/src/models/IModelInfo';
import type { ModelsManager } from './modelsManager';
import { getPortsInfo } from '../utils/ports';
import { getPortsFromLabel, getPortsInfo } from '../utils/ports';
import { goarch } from '../utils/arch';
import { getDurationSecondsSince, timeout } from '../utils/utils';
import type { LocalRepositoryRegistry } from '../registries/LocalRepositoryRegistry';
Expand All @@ -50,6 +50,9 @@ import { Publisher } from '../utils/Publisher';
import { isQEMUMachine } from '../utils/podman';
import { SECOND } from '../utils/inferenceUtils';
import { getModelPropertiesForEnvironment } from '../utils/modelsUtils';
import { getPodHealth } from '../utils/podsUtils';
import { getImageTag } from '../utils/imagesUtils';
import { getRandomName } from '../utils/randomUtils';

export const LABEL_MODEL_ID = 'ai-lab-model-id';
export const LABEL_MODEL_PORTS = 'ai-lab-model-ports';
Expand Down Expand Up @@ -318,7 +321,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
};
}

const podifiedName = this.getRandomName(`${image.appName}-podified`);
const podifiedName = getRandomName(`${image.appName}-podified`);
await containerEngine.createContainer(podInfo.engineId, {
Image: image.id,
name: podifiedName,
Expand Down Expand Up @@ -384,7 +387,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
labels[LABEL_APP_PORTS] = appPorts.join(',');
}
const pod = await containerEngine.createPod({
name: this.getRandomName(`pod-${sampleAppImageInfo.appName}`),
name: getRandomName(`pod-${sampleAppImageInfo.appName}`),
portmappings: portmappings,
labels,
});
Expand All @@ -395,10 +398,6 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
};
}

getRandomName(base: string): string {
return `${base ?? ''}-${new Date().getTime()}`;
}

async buildImages(
recipe: Recipe,
containers: ContainerConfig[],
Expand Down Expand Up @@ -432,7 +431,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
throw new Error('Context configured does not exist.');
}

const imageTag = this.getImageTag(recipe, container);
const imageTag = getImageTag(recipe, container);
const buildOptions: BuildImageOptions = {
containerFile: container.containerfile,
tag: imageTag,
Expand Down Expand Up @@ -479,7 +478,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
await Promise.all(
containers.map(async container => {
const task = containerTasks[container.name];
const imageTag = this.getImageTag(recipe, container);
const imageTag = getImageTag(recipe, container);

const image = images.find(im => {
return im.RepoTags?.some(tag => tag.endsWith(imageTag));
Expand All @@ -506,14 +505,6 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
return imageInfoList;
}

getImageTag(recipe: Recipe, container: ContainerConfig) {
let tag = container.image ?? `${recipe.id}-${container.name}`;
if (!tag.includes(':')) {
tag += ':latest';
}
return tag;
}

getConfigAndFilterContainers(
recipeBaseDir: string | undefined,
localFolder: string,
Expand Down Expand Up @@ -672,7 +663,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
);
}

adoptPod(pod: PodInfo) {
private adoptPod(pod: PodInfo) {
if (!pod.Labels) {
return;
}
Expand All @@ -681,8 +672,8 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
if (!recipeId || !modelId) {
return;
}
const appPorts = this.getPortsFromLabel(pod.Labels, LABEL_APP_PORTS);
const modelPorts = this.getPortsFromLabel(pod.Labels, LABEL_MODEL_PORTS);
const appPorts = getPortsFromLabel(pod.Labels, LABEL_APP_PORTS);
const modelPorts = getPortsFromLabel(pod.Labels, LABEL_MODEL_PORTS);
if (this.#applications.has({ recipeId, modelId })) {
return;
}
Expand All @@ -697,7 +688,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
this.updateApplicationState(recipeId, modelId, state);
}

forgetPod(pod: PodInfo) {
private forgetPod(pod: PodInfo) {
if (!pod.Labels) {
return;
}
Expand All @@ -723,7 +714,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
}
}

forgetPodById(podId: string) {
private forgetPodById(podId: string) {
const app = Array.from(this.#applications.values()).find(p => p.pod.Id === podId);
if (!app) {
return;
Expand Down Expand Up @@ -753,7 +744,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
}
}

async checkPodsHealth() {
private async checkPodsHealth() {
const pods = await containerEngine
.listPods()
.then(pods => pods.filter(pod => LABEL_RECIPE_ID in pod.Labels && LABEL_MODEL_ID in pod.Labels));
Expand All @@ -770,7 +761,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
containerEngine.inspectContainer(pod.engineId, container.Id).then(data => data.State.Health?.Status),
),
);
const podHealth = this.getPodHealth(containerStates);
const podHealth = getPodHealth(containerStates);
const state = this.#applications.get({ recipeId, modelId });
if (state.health !== podHealth) {
state.health = podHealth;
Expand All @@ -788,20 +779,6 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
}
}

getPodHealth(infos: (string | undefined)[]): PodHealth {
const checked = infos.filter(info => !!info && info !== 'none' && info !== '');
if (!checked.length) {
return 'none';
}
if (infos.some(info => info === 'unhealthy')) {
return 'unhealthy';
}
if (infos.some(info => info === 'starting')) {
return 'starting';
}
return 'healthy';
}

updateApplicationState(recipeId: string, modelId: string, state: ApplicationState): void {
this.#applications.set({ recipeId, modelId }, state);
this.notify();
Expand Down Expand Up @@ -890,29 +867,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
);
}

getPortsFromLabel(labels: { [key: string]: string }, key: string): number[] {
if (!(key in labels)) {
return [];
}
const value = labels[key];
const portsStr = value.split(',');
const result: number[] = [];
for (const portStr of portsStr) {
const port = parseInt(portStr, 10);
if (isNaN(port)) {
// malformed label, just ignore it
return [];
}
result.push(port);
}
return result;
}

dispose(): void {
this.cleanDisposables();
}

private cleanDisposables(): void {
this.#disposables.forEach(disposable => disposable.dispose());
}
}
37 changes: 37 additions & 0 deletions packages/backend/src/utils/imagesUtils.spec.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import { expect, test } from 'vitest';
import type { Recipe } from '@shared/src/models/IRecipe';
import type { ContainerConfig } from '../models/AIConfig';
import { getImageTag } from './imagesUtils';

test('return recipe-container tag if container image prop is not defined', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
} as ContainerConfig;
const imageTag = getImageTag(recipe, container);
expect(imageTag).equals('recipe1-name:latest');
});
test('return container image prop is defined', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
image: 'quay.io/repo/image:v1',
} as ContainerConfig;
const imageTag = getImageTag(recipe, container);
expect(imageTag).equals('quay.io/repo/image:v1');
});
test('append latest tag to container image prop if it has no tag', () => {
const recipe = {
id: 'recipe1',
} as Recipe;
const container = {
name: 'name',
image: 'quay.io/repo/image',
} as ContainerConfig;
const imageTag = getImageTag(recipe, container);
expect(imageTag).equals('quay.io/repo/image:latest');
});
Loading

0 comments on commit b130cbf

Please sign in to comment.