Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: handle multi download model, uninstall script #932

Merged
merged 3 commits into from
Jul 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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": {
"@huggingface/gguf": "^0.1.5",
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 @@ -125,12 +127,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 @@ -253,7 +255,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();
Loading