Skip to content

Commit

Permalink
fix: refactor pullApplication
Browse files Browse the repository at this point in the history 
Signed-off-by: lstocchi <[email protected]>
  • Loading branch information
@lstocchi
lstocchi committed Jan 24, 2024
1 parent 580f364 commit 333bd19
Show file tree
Hide file tree
Showing 3 changed files with 141 additions and 108 deletions.
2 changes: 1 addition & 1 deletion packages/backend/src/managers/applicationManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ describe('pullApplication', () => {
} as unknown as ModelsManager,
);

downloadModelMainSpy = vi.spyOn(manager, 'downloadModelMain');
downloadModelMainSpy = vi.spyOn(manager, 'doDownloadModelWrapper');
downloadModelMainSpy.mockResolvedValue('');
}

Expand Down
242 changes: 135 additions & 107 deletions packages/backend/src/managers/applicationManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import * as https from 'node:https';
import * as path from 'node:path';
import { containerEngine } from '@podman-desktop/api';
import type { RecipeStatusRegistry } from '../registries/RecipeStatusRegistry';
import type { AIConfig } from '../models/AIConfig';
import type { AIConfig, AIConfigFile, ContainerConfig } from '../models/AIConfig';
import { parseYaml } from '../models/AIConfig';
import type { Task } from '@shared/src/models/ITask';
import { RecipeStatusUtils } from '../utils/recipeStatusUtils';
Expand Down Expand Up @@ -53,111 +53,23 @@ export class ApplicationManager {

const localFolder = path.join(this.appUserDirectory, recipe.id);

// Adding checkout task
const checkoutTask: Task = {
id: 'checkout',
name: 'Checkout repository',
state: 'loading',
labels: {
git: 'checkout',
},
};
taskUtil.setTask(checkoutTask);

// We might already have the repository cloned
if (fs.existsSync(localFolder) && fs.statSync(localFolder).isDirectory()) {
// Update checkout state
checkoutTask.name = 'Checkout repository (cached).';
checkoutTask.state = 'success';
} else {
// Create folder
fs.mkdirSync(localFolder, { recursive: true });

// Clone the repository
console.log(`Cloning repository ${recipe.repository} in ${localFolder}.`);
await this.git.cloneRepository(recipe.repository, localFolder);

// Update checkout state
checkoutTask.state = 'success';
}
// Update task
taskUtil.setTask(checkoutTask);
// clone the recipe repository on the local folder
await this.doCheckout(recipe.repository, localFolder, taskUtil);

// Adding loading configuration task
const loadingConfiguration: Task = {
id: 'loading-config',
name: 'Loading configuration',
state: 'loading',
};
taskUtil.setTask(loadingConfiguration);
// load and parse the recipe configuration file
const aiConfigFile = this.getConfiguration(recipe.config, localFolder, taskUtil);

let configFile: string;
if (recipe.config !== undefined) {
configFile = path.join(localFolder, recipe.config);
} else {
configFile = path.join(localFolder, CONFIG_FILENAME);
}
// get model by downloading it or retrieving locally
await this.downloadModel(model, taskUtil);

if (!fs.existsSync(configFile)) {
loadingConfiguration.state = 'error';
taskUtil.setTask(loadingConfiguration);
throw new Error(`The file located at ${configFile} does not exist.`);
}
// filter the containers based on architecture, gpu accelerator and backend (that define which model supports)
const filteredContainers: ContainerConfig[] = this.filterContainers(aiConfigFile.aiConfig);

// If the user configured the config as a directory we check for "ai-studio.yaml" inside.
if (fs.statSync(configFile).isDirectory()) {
const tmpPath = path.join(configFile, CONFIG_FILENAME);
// If it has the ai-studio.yaml we use it.
if (fs.existsSync(tmpPath)) {
configFile = tmpPath;
}
}

// Parsing the configuration
console.log(`Reading configuration from ${configFile}.`);
const rawConfiguration = fs.readFileSync(configFile, 'utf-8');
let aiConfig: AIConfig;
try {
aiConfig = parseYaml(rawConfiguration, arch());
} catch (err) {
// Mask task as failed
loadingConfiguration.state = 'error';
taskUtil.setTask(loadingConfiguration);
throw new Error('Cannot load configuration file.');
}

// Mark as success.
loadingConfiguration.state = 'success';
taskUtil.setTask(loadingConfiguration);

// Filter the containers based on architecture
const filteredContainers = aiConfig.application.containers.filter(
container => container.arch === undefined || container.arch === arch(),
);

if (!this.modelsManager.isModelOnDisk(model.id)) {
// Download model
taskUtil.setTask({
id: model.id,
state: 'loading',
name: `Downloading model ${model.name}`,
labels: {
'model-pulling': model.id,
},
});

await this.downloadModelMain(model.id, model.url, taskUtil);
} else {
taskUtil.setTask({
id: model.id,
state: 'success',
name: `Model ${model.name} already present on disk`,
labels: {
'model-pulling': model.id,
},
});
}
// build all images, one per container (for a basic sample we should have 2 containers = sample app + model service)
await this.buildImages(filteredContainers, aiConfigFile.path, taskUtil);
}

async buildImages(filteredContainers: ContainerConfig[], configPath: string, taskUtil: RecipeStatusUtils) {
filteredContainers.forEach(container => {
taskUtil.setTask({
id: container.name,
Expand All @@ -167,10 +79,10 @@ export class ApplicationManager {
});

// Promise all the build images
return Promise.all(
await Promise.all(
filteredContainers.map(container => {
// We use the parent directory of our configFile as the rootdir, then we append the contextDir provided
const context = path.join(getParentDirectory(configFile), container.contextdir);
const context = path.join(getParentDirectory(configPath), container.contextdir);
console.log(`Application Manager using context ${context} for container ${container.name}`);

// Ensure the context provided exist otherwise throw an Error
Expand Down Expand Up @@ -211,7 +123,123 @@ export class ApplicationManager {
);
}

downloadModelMain(modelId: string, url: string, taskUtil: RecipeStatusUtils, destFileName?: string): Promise<string> {
filterContainers(aiConfig: AIConfig): ContainerConfig[] {
return aiConfig.application.containers.filter(
container => container.arch === undefined || container.arch === arch(),
);
}

async downloadModel(model: ModelInfo, taskUtil: RecipeStatusUtils) {
if (!this.modelsManager.isModelOnDisk(model.id)) {
// Download model
taskUtil.setTask({
id: model.id,
state: 'loading',
name: `Downloading model ${model.name}`,
labels: {
'model-pulling': model.id,
},
});

await this.doDownloadModelWrapper(model.id, model.url, taskUtil);
} else {
taskUtil.setTask({
id: model.id,
state: 'success',
name: `Model ${model.name} already present on disk`,
labels: {
'model-pulling': model.id,
},
});
}
}

getConfiguration(recipeConfig: string, localFolder: string,taskUtil: RecipeStatusUtils): AIConfigFile {
// Adding loading configuration task
const loadingConfiguration: Task = {
id: 'loading-config',
name: 'Loading configuration',
state: 'loading',
};
taskUtil.setTask(loadingConfiguration);

let configFile: string;
if (recipeConfig !== undefined) {
configFile = path.join(localFolder, recipeConfig);
} else {
configFile = path.join(localFolder, CONFIG_FILENAME);
}

if (!fs.existsSync(configFile)) {
loadingConfiguration.state = 'error';
taskUtil.setTask(loadingConfiguration);
throw new Error(`The file located at ${configFile} does not exist.`);
}

// If the user configured the config as a directory we check for "ai-studio.yaml" inside.
if (fs.statSync(configFile).isDirectory()) {
const tmpPath = path.join(configFile, CONFIG_FILENAME);
// If it has the ai-studio.yaml we use it.
if (fs.existsSync(tmpPath)) {
configFile = tmpPath;
}
}

// Parsing the configuration
console.log(`Reading configuration from ${configFile}.`);
const rawConfiguration = fs.readFileSync(configFile, 'utf-8');
let aiConfig: AIConfig;
try {
aiConfig = parseYaml(rawConfiguration, arch());
} catch (err) {
// Mask task as failed
loadingConfiguration.state = 'error';
taskUtil.setTask(loadingConfiguration);
throw new Error('Cannot load configuration file.');
}

// Mark as success.
loadingConfiguration.state = 'success';
taskUtil.setTask(loadingConfiguration);
return {
aiConfig,
path: configFile,
}
}

async doCheckout(repository: string, localFolder: string, taskUtil: RecipeStatusUtils) {
// Adding checkout task
const checkoutTask: Task = {
id: 'checkout',
name: 'Checkout repository',
state: 'loading',
labels: {
git: 'checkout',
},
};
taskUtil.setTask(checkoutTask);

// We might already have the repository cloned
if (fs.existsSync(localFolder) && fs.statSync(localFolder).isDirectory()) {
// Update checkout state
checkoutTask.name = 'Checkout repository (cached).';
checkoutTask.state = 'success';
} else {
// Create folder
fs.mkdirSync(localFolder, { recursive: true });

// Clone the repository
console.log(`Cloning repository ${repository} in ${localFolder}.`);
await this.git.cloneRepository(repository, localFolder);

// Update checkout state
checkoutTask.state = 'success';
}
// Update task
taskUtil.setTask(checkoutTask);
}

doDownloadModelWrapper(modelId: string, url: string, taskUtil: RecipeStatusUtils, destFileName?: string): Promise<string> {
return new Promise((resolve, reject) => {
const downloadCallback = (result: DownloadModelResult) => {
if (result.result) {
Expand All @@ -229,11 +257,11 @@ export class ApplicationManager {
return;
}

this.downloadModel(modelId, url, taskUtil, downloadCallback, destFileName);
this.doDownloadModel(modelId, url, taskUtil, downloadCallback, destFileName);
});
}

private downloadModel(
private doDownloadModel(
modelId: string,
url: string,
taskUtil: RecipeStatusUtils,
Expand All @@ -253,7 +281,7 @@ export class ApplicationManager {
let progress = 0;
https.get(url, resp => {
if (resp.headers.location) {
this.downloadModel(modelId, resp.headers.location, taskUtil, callback, destFileName);
this.doDownloadModel(modelId, resp.headers.location, taskUtil, callback, destFileName);
return;
} else {
if (totalFileSize === 0 && resp.headers['content-length']) {
Expand Down
5 changes: 5 additions & 0 deletions packages/backend/src/models/AIConfig.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,11 @@ export interface AIConfig {
};
}

export interface AIConfigFile {
aiConfig: AIConfig;
path: string;
}

export function isString(value: unknown): value is string {
return (!!value && typeof value === 'string') || value instanceof String;
}
Expand Down

0 comments on commit 333bd19

Please sign in to comment.