Skip to content

Commit

Permalink
fix: handle multi download model, uninstall script (#932)
Browse files Browse the repository at this point in the history
  • Loading branch information
marknguyen1302 authored Jul 29, 2024
1 parent a8eb3b4 commit 0954152
Show file tree
Hide file tree
Showing 11 changed files with 82 additions and 23 deletions.
3 changes: 2 additions & 1 deletion cortex-js/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@
"build:dev": "yarn build && run-script-os && npm link",
"build:dev:windows": "echo 'Windows build complete'",
"build:dev:macos": "chmod +x ./dist/src/command.js",
"build:dev:linux": "chmod +x ./dist/src/command.js"
"build:dev:linux": "chmod +x ./dist/src/command.js",
"preuninstall": "node ./uninstall.js"
},
"dependencies": {
"@cortexso/cortex.js": "^0.1.3",
Expand Down
7 changes: 4 additions & 3 deletions cortex-js/src/infrastructure/commanders/chat.command.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,9 @@ import { FileManagerService } from '../services/file-manager/file-manager.servic
import { isRemoteEngine } from '@/utils/normalize-model-id';
import { Cortex } from '@cortexso/cortex.js';
import { ChatClient } from './services/chat-client';
import { downloadModelProgress } from '@/utils/pull-model';
import { downloadProgress } from '@/utils/download-progress';
import { CortexClient } from './services/cortex.client';
import { DownloadType } from '@/domain/models/download.interface';

type ChatOptions = {
threadId?: string;
Expand Down Expand Up @@ -92,7 +93,7 @@ export class ChatCommand extends BaseCommand {
) {
console.log('Downloading engine...');
await this.cortex.engines.init(engine);
await downloadModelProgress(this.cortex);
await downloadProgress(this.cortex, undefined, DownloadType.Engine)
}

if (!message) options.attach = true;
Expand All @@ -107,7 +108,7 @@ export class ChatCommand extends BaseCommand {
);

const preset = await this.fileService.getPreset(options.preset);

return this.cortex.models.start(modelId, preset).then(() =>
this.chatClient.chat(
modelId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { BaseCommand } from '../base.command';
import { defaultInstallationOptions } from '@/utils/init';
import { Presets, SingleBar } from 'cli-progress';
import { CortexClient } from '../services/cortex.client';
import ora from 'ora';

@SubCommand({
name: '<name> init',
Expand Down Expand Up @@ -44,9 +45,11 @@ export class EnginesInitCommand extends BaseCommand {
const host = configs.cortexCppHost;
const port = configs.cortexCppPort;
// Should stop cortex before installing engine
const stopCortexSpinner = ora('Stopping cortex...').start();
if (await this.cortexUsecases.healthCheck(host, port)) {
await this.cortexUsecases.stopCortex();
}
stopCortexSpinner.succeed('Cortex stopped');
console.log(`Installing engine ${engine}...`);
await this.cortex.engines.init(engine, params);
const response = await this.cortex.events.downloadEvent();
Expand All @@ -68,6 +71,9 @@ export class EnginesInitCommand extends BaseCommand {
}
}
progressBar.stop();
const startCortexSpinner = ora('Starting cortex...').start();
await this.cortexUsecases.startCortex();
startCortexSpinner.succeed('Cortex started');
console.log('Engine installed successfully');
process.exit(0);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,9 @@ import { checkModelCompatibility } from '@/utils/model-check';
import { Engines } from '../types/engine.interface';
import { CortexUsecases } from '@/usecases/cortex/cortex.usecases';
import { BaseCommand } from '../base.command';
import { downloadModelProgress } from '@/utils/pull-model';
import { downloadProgress } from '@/utils/download-progress';
import { CortexClient } from '../services/cortex.client';
import { DownloadType } from '@/domain/models/download.interface';

@SubCommand({
name: 'pull',
Expand Down Expand Up @@ -58,7 +59,7 @@ export class ModelPullCommand extends BaseCommand {
exit(1);
});

await downloadModelProgress(this.cortex, modelId);
await downloadProgress(this.cortex, modelId);

const existingModel = await this.cortex.models.retrieve(modelId);
const engine = existingModel?.engine || Engines.llamaCPP;
Expand All @@ -70,7 +71,7 @@ export class ModelPullCommand extends BaseCommand {
console.log('\n');
console.log('Downloading engine...');
await this.cortex.engines.init(engine);
await downloadModelProgress(this.cortex);
await downloadProgress(this.cortex, undefined, DownloadType.Engine);
}
this.telemetryUsecases.sendEvent(
[
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ import { Engines } from '../types/engine.interface';
import { checkModelCompatibility } from '@/utils/model-check';
import { BaseCommand } from '../base.command';
import { isRemoteEngine } from '@/utils/normalize-model-id';
import { downloadModelProgress } from '@/utils/pull-model';
import { downloadProgress } from '@/utils/download-progress';
import { CortexClient } from '../services/cortex.client';
import { DownloadType } from '@/domain/models/download.interface';

type ModelStartOptions = {
attach: boolean;
Expand Down Expand Up @@ -73,7 +74,7 @@ export class ModelStartCommand extends BaseCommand {
) {
console.log('Downloading engine...');
await this.cortex.engines.init(engine);
await downloadModelProgress(this.cortex);
await downloadProgress(this.cortex, undefined, DownloadType.Engine);
}

// Attached - stdout logs
Expand Down
7 changes: 4 additions & 3 deletions cortex-js/src/infrastructure/commanders/run.command.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,9 @@ import { checkModelCompatibility } from '@/utils/model-check';
import { BaseCommand } from './base.command';
import { isRemoteEngine } from '@/utils/normalize-model-id';
import { ChatClient } from './services/chat-client';
import { downloadModelProgress } from '@/utils/pull-model';
import { downloadProgress } from '@/utils/download-progress';
import { CortexClient } from './services/cortex.client';
import { DownloadType } from '@/domain/models/download.interface';

type RunOptions = {
threadId?: string;
Expand Down Expand Up @@ -65,7 +66,7 @@ export class RunCommand extends BaseCommand {
checkingSpinner.fail(e.message ?? e);
exit(1);
});
await downloadModelProgress(this.cortex, modelId);
await downloadProgress(this.cortex, modelId);
}

// Second check if model is available
Expand All @@ -84,7 +85,7 @@ export class RunCommand extends BaseCommand {
) {
console.log('Downloading engine...');
await this.cortex.engines.init(engine);
await downloadModelProgress(this.cortex);
await downloadProgress(this.cortex, undefined, DownloadType.Engine);
}

const startingSpinner = ora('Loading model...').start();
Expand Down
14 changes: 10 additions & 4 deletions cortex-js/src/infrastructure/controllers/engines.controller.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ import {
Post,
Body,
Patch,
Res,
} from '@nestjs/common';
import { ApiOperation, ApiParam, ApiTags, ApiResponse } from '@nestjs/swagger';
import { Response } from 'express';
import { TransformInterceptor } from '../interceptors/transform.interceptor';
import { EnginesUsecases } from '@/usecases/engines/engines.usecase';
import { EngineDto, InitEngineDto } from '../dtos/engines/engines.dto';
Expand Down Expand Up @@ -78,11 +80,15 @@ export class EnginesController {
description: 'The unique identifier of the engine.',
})
@Post(':name(*)/init')
initialize(@Param('name') name: string, @Body() body: InitEngineDto | undefined) {
initialize(@Param('name') name: string, @Body() body: InitEngineDto | undefined, @Res() res: Response) {
try{
this.initUsescases.installEngine(body, 'latest', name, true);
return {
message: 'Engine initialization started successfully.',
};
res.json({
message: 'Engine initialization started successfully.',
})
} catch (error) {
res.status(400).send(error.message);
}
}

@HttpCode(200)
Expand Down
13 changes: 7 additions & 6 deletions cortex-js/src/usecases/engines/engines.usecase.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,11 @@ export class EnginesUsecases {
force: boolean = false,
): Promise<any> => {
// Use default option if not defined

if (!options && engine === Engines.llamaCPP) {
options = await defaultInstallationOptions();
}
const installPackages = [];
// Ship Llama.cpp engine by default
if (
!existsSync(
Expand All @@ -93,7 +95,7 @@ export class EnginesUsecases {
engine === Engines.llamaCPP &&
(options?.vulkan ||
(options?.runMode === 'GPU' && options?.gpuType !== 'Nvidia'));
await this.installAcceleratedEngine(version, engine, [
installPackages.push(this.installAcceleratedEngine(version, engine, [
process.platform === 'win32'
? '-windows'
: process.platform === 'darwin'
Expand All @@ -116,7 +118,7 @@ export class EnginesUsecases {
? '-arm64'
: '-amd64'
: '',
]);
]));
}

if (
Expand All @@ -126,12 +128,12 @@ export class EnginesUsecases {
options?.gpuType === 'Nvidia' &&
!options?.vulkan)
)
await this.installCudaToolkitDependency(
installPackages.push(this.installCudaToolkitDependency(
engine === Engines.tensorrtLLM
? MIN_CUDA_VERSION
: options?.cudaVersion,
);

));
await Promise.all(installPackages);
// Update states
await this.extensionRepository.findOne(engine).then((e) => {
if (e) e.status = EngineStatus.READY;
Expand Down Expand Up @@ -254,7 +256,6 @@ export class EnginesUsecases {
console.log(
`Could not find engine file for platform ${process.platform}`,
);
exit(1);
}

const engineDir = await this.fileManagerService.getCortexCppEnginePath();
Expand Down
7 changes: 7 additions & 0 deletions cortex-js/src/usecases/models/models.usecases.ts
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,14 @@ export class ModelsUsecases {
}

const modelFolder = join(modelsContainerDir, normalizeModelId(id));
const model = await this.getModelOrThrow(id);
const engine = (await this.extensionRepository.findOne(
model!.engine ?? Engines.llamaCPP,
)) as EngineExtension | undefined;

if (engine) {
await engine.unloadModel(id, model.engine || Engines.llamaCPP).catch(() => {}); // Silent fail
}
return this.modelRepository
.remove(id)
.then(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import { Presets, SingleBar } from "cli-progress";
import { Cortex } from "@cortexso/cortex.js";
import { exit, stdin, stdout } from 'node:process';
import { DownloadType } from "@/domain/models/download.interface";

export const downloadModelProgress = async (cortex: Cortex, downloadId?: string) => {
export const downloadProgress = async (cortex: Cortex, downloadId?: string, downloadType?: DownloadType) => {
const response = await cortex.events.downloadEvent();

const rl = require('readline').createInterface({
Expand All @@ -27,6 +28,7 @@ export const downloadModelProgress = async (cortex: Cortex, downloadId?: string)
for await (const stream of response) {
if (stream.length) {
const data = stream[0] as any;
if (downloadId && data.id !== downloadId || downloadType && data.type !== downloadType) continue;

if (data.status === 'downloaded') break;

Expand Down
32 changes: 32 additions & 0 deletions cortex-js/uninstall.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
const { promises } = require('node:fs');
const os = require('os');
const fs = require('fs');
const path = require('path');
const yaml = require('js-yaml');

const uninstall = async () => {
const cortexConfigPath = path.join(os.homedir(), '.cortexrc');
if (fs.existsSync(cortexConfigPath)) {
const content = await promises.readFile(cortexConfigPath, 'utf8');
const config = yaml.load(content);
if(!config) {
return;
}
const { dataFolderPath } = config;
const modelsFolderPath = path.join(dataFolderPath, 'models');
// remove all data in data folder path except models
const files = fs.readdirSync(dataFolderPath);
for (const file of files) {
const fileStat = fs.statSync(path.join(dataFolderPath, file));
if (file !== 'models') {
if (fileStat.isDirectory()) {
fs.rmSync(path.join(dataFolderPath, file), { recursive: true });
} else {
fs.unlinkSync(path.join(dataFolderPath, file));
}
}
}
}
};

uninstall();

0 comments on commit 0954152

Please sign in to comment.