Skip to content
This repository was archived by the owner on Jul 4, 2025. It is now read-only.

Commit ca5fb29

Browse files
marknguyen1302marknguyen1302
authored
feat: download model checksum (#986)
1 parent 4ac07e2 commit ca5fb29

File tree

13 files changed

+144
-34
lines changed

13 files changed

+144
-34
lines changed

‎cortex-js/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
"preuninstall": "node ./uninstall.js"
3939
},
4040
"dependencies": {
41-
"@cortexso/cortex.js": "^0.1.5",
41+
"@cortexso/cortex.js": "^0.1.7",
4242
"@nestjs/axios": "^3.0.2",
4343
"@nestjs/common": "^10.0.0",
4444
"@nestjs/config": "^3.2.2",

‎cortex-js/src/domain/models/huggingface.interface.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ export interface HuggingFaceRepoSibling {
1010
downloadUrl?: string;
1111
fileSize?: number;
1212
quantization?: Quantization;
13+
lfs?: { oid?: string };
1314
}
1415
export interface HuggingFaceRepoData {
1516
id: string;

‎cortex-js/src/infrastructure/commanders/base.command.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,4 +34,7 @@ export abstract class BaseCommand extends CommandRunner {
3434
}
3535
await this.runCommand(passedParam, options);
3636
}
37+
protected setCortex(cortex: CortexClient) {
38+
this.cortex = cortex;
39+
}
3740
}

‎cortex-js/src/infrastructure/commanders/chat.command.ts‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { Cortex } from '@cortexso/cortex.js';
1818
import { ChatClient } from './services/chat-client';
1919
import { downloadProgress } from '@/utils/download-progress';
2020
import { DownloadType } from '@/domain/models/download.interface';
21-
import { CortexClient } from './services/cortex.client';
21+
import { checkRequiredVersion } from '@/utils/model-check';
2222

2323
type ChatOptions = {
2424
threadId?: string;
@@ -96,8 +96,20 @@ export class ChatCommand extends BaseCommand {
9696
await downloadProgress(this.cortex, undefined, DownloadType.Engine);
9797
}
9898

99+
const { version: engineVersion } =
100+
await this.cortex.engines.retrieve(engine);
101+
if (
102+
existingModel.engine_version &&
103+
!checkRequiredVersion(existingModel.engine_version, engineVersion)
104+
) {
105+
console.log(
106+
`Model engine version ${existingModel.engine_version} is not compatible with engine version ${engineVersion}`,
107+
);
108+
process.exit(1);
109+
}
110+
99111
if (!message) options.attach = true;
100-
this.telemetryUsecases.sendEvent(
112+
void this.telemetryUsecases.sendEvent(
101113
[
102114
{
103115
name: EventName.CHAT,

‎cortex-js/src/infrastructure/commanders/engines.command.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ export class EnginesCommand extends BaseCommand {
6161
) {
6262
const commandInstance = this.moduleRef.get(commandClass, { strict: false });
6363
if (commandInstance) {
64+
commandInstance.setCortex(this.cortex);
6465
await commandInstance.runCommand(params, options);
6566
} else {
6667
console.error('Command not found.');

‎cortex-js/src/infrastructure/commanders/models.command.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ export class ModelsCommand extends BaseCommand {
6767
) {
6868
const commandInstance = this.moduleRef.get(commandClass, { strict: false });
6969
if (commandInstance) {
70+
commandInstance.setCortex(this.cortex);
7071
await commandInstance.runCommand(params, options);
7172
} else {
7273
console.error('Command not found.');

‎cortex-js/src/infrastructure/commanders/run.command.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import ora from 'ora';
55
import { existsSync } from 'fs';
66
import { join } from 'path';
77
import { Engines } from './types/engine.interface';
8-
import { checkModelCompatibility } from '@/utils/model-check';
8+
import { checkModelCompatibility, checkRequiredVersion } from '@/utils/model-check';
99
import { BaseCommand } from './base.command';
1010
import { isRemoteEngine } from '@/utils/normalize-model-id';
1111
import { ChatClient } from './services/chat-client';
@@ -98,6 +98,11 @@ export class RunCommand extends BaseCommand {
9898
await this.cortex.engines.init(engine);
9999
await downloadProgress(this.cortex, undefined, DownloadType.Engine);
100100
}
101+
const { version: engineVersion } = await this.cortex.engines.retrieve(engine);
102+
if(existingModel.engine_version && !checkRequiredVersion(existingModel.engine_version, engineVersion)) {
103+
console.log(`Model engine version ${existingModel.engine_version} is not compatible with engine version ${engineVersion}`);
104+
process.exit(1);
105+
}
101106

102107
const startingSpinner = ora('Loading model...').start();
103108

‎cortex-js/src/infrastructure/services/download-manager/download-manager.service.ts‎

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { Presets, SingleBar } from 'cli-progress';
1111
import { createWriteStream, unlinkSync } from 'node:fs';
1212
import { basename } from 'node:path';
1313
import { firstValueFrom } from 'rxjs';
14+
import crypto from 'crypto';
1415

1516
@Injectable()
1617
export class DownloadManagerService {
@@ -51,7 +52,13 @@ export class DownloadManagerService {
5152
downloadId: string,
5253
title: string,
5354
downloadType: DownloadType,
54-
urlToDestination: Record<string, string>,
55+
urlToDestination: Record<
56+
string,
57+
{
58+
destination: string;
59+
checksum?: string;
60+
}
61+
>,
5562
finishedCallback?: () => Promise<void>,
5663
inSequence: boolean = true,
5764
) {
@@ -65,7 +72,7 @@ export class DownloadManagerService {
6572

6673
const downloadItems: DownloadItem[] = Object.keys(urlToDestination).map(
6774
(url) => {
68-
const destination = urlToDestination[url];
75+
const { destination, checksum } = urlToDestination[url];
6976
const downloadItem: DownloadItem = {
7077
id: destination,
7178
time: {
@@ -78,6 +85,7 @@ export class DownloadManagerService {
7885
},
7986
progress: 0,
8087
status: DownloadStatus.Downloading,
88+
checksum,
8189
};
8290

8391
return downloadItem;
@@ -119,15 +127,15 @@ export class DownloadManagerService {
119127
if (!inSequence) {
120128
Promise.all(
121129
Object.keys(urlToDestination).map((url) => {
122-
const destination = urlToDestination[url];
123-
return this.downloadFile(downloadId, url, destination);
130+
const { destination, checksum } = urlToDestination[url];
131+
return this.downloadFile(downloadId, url, destination, checksum);
124132
}),
125133
).then(callBack);
126134
} else {
127135
// Download model file in sequence
128136
for (const url of Object.keys(urlToDestination)) {
129-
const destination = urlToDestination[url];
130-
await this.downloadFile(downloadId, url, destination);
137+
const { destination, checksum } = urlToDestination[url];
138+
await this.downloadFile(downloadId, url, destination, checksum);
131139
}
132140
return callBack();
133141
}
@@ -137,7 +145,14 @@ export class DownloadManagerService {
137145
downloadId: string,
138146
url: string,
139147
destination: string,
148+
checksum?: string,
140149
) {
150+
console.log('Downloading', {
151+
downloadId,
152+
url,
153+
destination,
154+
checksum,
155+
});
141156
const controller = new AbortController();
142157
// adding to abort controllers
143158
this.abortControllers[downloadId][destination] = controller;
@@ -155,6 +170,7 @@ export class DownloadManagerService {
155170
}
156171

157172
const writer = createWriteStream(destination);
173+
const hash = crypto.createHash('sha256');
158174
const totalBytes = Number(response.headers['content-length']);
159175

160176
// update download state
@@ -214,8 +230,23 @@ export class DownloadManagerService {
214230
const downloadItem = currentDownloadState?.children.find(
215231
(downloadItem) => downloadItem.id === destination,
216232
);
233+
const isFileBroken = checksum && checksum === hash.digest('hex');
217234
if (downloadItem) {
218-
downloadItem.status = DownloadStatus.Downloaded;
235+
downloadItem.status = isFileBroken
236+
? DownloadStatus.Error
237+
: DownloadStatus.Downloaded;
238+
if (isFileBroken) {
239+
downloadItem.error = 'Checksum is not matched';
240+
this.handleError(
241+
new Error('Checksum is not matched'),
242+
downloadId,
243+
destination,
244+
);
245+
}
246+
}
247+
if (isFileBroken) {
248+
currentDownloadState.status = DownloadStatus.Error;
249+
currentDownloadState.error = 'Checksum is not matched';
219250
}
220251
this.eventEmitter.emit('download.event', this.allDownloadStates);
221252
} finally {
@@ -234,6 +265,7 @@ export class DownloadManagerService {
234265
});
235266

236267
response.data.on('data', (chunk: any) => {
268+
hash.update(chunk);
237269
resetTimeout();
238270
transferredBytes += chunk.length;
239271

‎cortex-js/src/infrastructure/services/file-manager/file-manager.service.ts‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ const writeAsync = promisify(write);
2828

2929
@Injectable()
3030
export class FileManagerService {
31-
private configFile = '.cortexrc';
3231
private cortexDirectoryName = 'cortex';
3332
private modelFolderName = 'models';
3433
private presetFolderName = 'presets';
@@ -44,7 +43,10 @@ export class FileManagerService {
4443
*/
4544
async getConfig(dataFolderPath?: string): Promise<Config & object> {
4645
const homeDir = os.homedir();
47-
const configPath = join(homeDir, this.configFile);
46+
const configPath = join(
47+
homeDir,
48+
this.getConfigFileName(this.configProfile),
49+
);
4850
const config = this.defaultConfig();
4951
const dataFolderPathUsed = dataFolderPath || config.dataFolderPath;
5052
if (!existsSync(configPath) || !existsSync(dataFolderPathUsed)) {
@@ -55,8 +57,7 @@ export class FileManagerService {
5557

5658
try {
5759
const content = await promises.readFile(configPath, 'utf8');
58-
const configs = yaml.load(content) as Record<string, Config>;
59-
const config = configs?.[this.configProfile] ?? {};
60+
const config = yaml.load(content) as Config & object;
6061
return {
6162
...this.defaultConfig(),
6263
...config,
@@ -72,17 +73,20 @@ export class FileManagerService {
7273

7374
async writeConfigFile(config: Config & object): Promise<void> {
7475
const homeDir = os.homedir();
75-
const configPath = join(homeDir, this.configFile);
76+
const configPath = join(
77+
homeDir,
78+
this.getConfigFileName(this.configProfile),
79+
);
7680

7781
// write config to file as yaml
7882
if (!existsSync(configPath)) {
7983
await promises.writeFile(configPath, '', 'utf8');
8084
}
8185
const content = await promises.readFile(configPath, 'utf8');
82-
const currentConfig = yaml.load(content) as Record<string, Config>;
86+
const currentConfig = yaml.load(content) as Config & object;
8387
const configString = yaml.dump({
8488
...currentConfig,
85-
[this.configProfile]: config,
89+
...config,
8690
});
8791
await promises.writeFile(configPath, configString, 'utf8');
8892
}
@@ -345,12 +349,17 @@ export class FileManagerService {
345349
*/
346350
getServerConfig(): { host: string; port: number } {
347351
const homeDir = os.homedir();
348-
const configPath = join(homeDir, this.configFile);
352+
const configPath = join(
353+
homeDir,
354+
this.getConfigFileName(this.configProfile),
355+
);
349356
let config = this.defaultConfig();
350357
try {
351358
const content = readFileSync(configPath, 'utf8');
352-
const configs = (yaml.load(content) as Record<string, Config>) ?? {};
353-
config = configs?.[this.configProfile] ?? config;
359+
const currentConfig = (yaml.load(content) as Config & object) ?? {};
360+
if (currentConfig) {
361+
config = currentConfig;
362+
}
354363
} catch {}
355364
return {
356365
host: config.apiServerHost ?? '127.0.0.1',
@@ -366,15 +375,22 @@ export class FileManagerService {
366375
}
367376
public profileConfigExists(profile: string): boolean {
368377
const homeDir = os.homedir();
369-
const configPath = join(homeDir, this.configFile);
378+
const configPath = join(homeDir, this.getConfigFileName(profile));
370379
try {
371380
const content = readFileSync(configPath, 'utf8');
372-
const configs = (yaml.load(content) as Record<string, Config>) ?? {};
373-
return !!configs[profile];
381+
const config = yaml.load(content) as Config & object;
382+
return !!config;
374383
} catch {
375384
return false;
376385
}
377386
}
387+
388+
private getConfigFileName(configProfile: string): string {
389+
if (configProfile === 'default') {
390+
return '.cortexrc';
391+
}
392+
return `.${configProfile}rc`;
393+
}
378394
}
379395

380396
export const fileManagerService = new FileManagerService();

‎cortex-js/src/usecases/engines/engines.usecase.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ export class EnginesUsecases {
199199
url,
200200
'Cuda Toolkit Dependencies',
201201
DownloadType.Engine,
202-
{ [url]: destination },
202+
{ [url]: { destination } },
203203
async () => {
204204
try {
205205
await decompress(
@@ -274,7 +274,7 @@ export class EnginesUsecases {
274274
toDownloadAsset.browser_download_url,
275275
engine,
276276
DownloadType.Engine,
277-
{ [toDownloadAsset.browser_download_url]: destination },
277+
{ [toDownloadAsset.browser_download_url]: { destination } },
278278
// On completed - post processing
279279
async () => {
280280
try {

0 commit comments

Comments
 (0)