Skip to content

GitLab

Commit cc331215 authored by Timothy J. Baek's avatar Timothy J. Baek
Browse files

refac: model download

parent 4809d363
Hide whitespace changes
Inline Side-by-side
Showing with 119 additions and 129 deletions
src/lib/components/chat/Settings/Models.svelte
View file @ cc331215
...@@ -13,7 +13,7 @@ ...@@ -13,7 +13,7 @@
uploadModel uploadModel
} from '$lib/apis/ollama'; } from '$lib/apis/ollama';
import { WEBUI_API_BASE_URL, WEBUI_BASE_URL } from '$lib/constants'; import { WEBUI_API_BASE_URL, WEBUI_BASE_URL } from '$lib/constants';
import { WEBUI_NAME, models, user } from '$lib/stores'; import { WEBUI_NAME, models, MODEL_DOWNLOAD_POOL, user } from '$lib/stores';
import { splitStream } from '$lib/utils'; import { splitStream } from '$lib/utils';
import { onMount, getContext } from 'svelte'; import { onMount, getContext } from 'svelte';
import { addLiteLLMModel, deleteLiteLLMModel, getLiteLLMModelInfo } from '$lib/apis/litellm'; import { addLiteLLMModel, deleteLiteLLMModel, getLiteLLMModelInfo } from '$lib/apis/litellm';
...@@ -50,12 +50,6 @@ ...@@ -50,12 +50,6 @@
let showExperimentalOllama = false; let showExperimentalOllama = false;
let ollamaVersion = ''; let ollamaVersion = '';
const MAX_PARALLEL_DOWNLOADS = 3; const MAX_PARALLEL_DOWNLOADS = 3;
const modelDownloadQueue = queue(
(task: { modelName: string }, cb) =>
pullModelHandlerProcessor({ modelName: task.modelName, callback: cb }),
MAX_PARALLEL_DOWNLOADS
);
let modelDownloadStatus: Record<string, any> = {};
let modelTransferring = false; let modelTransferring = false;
let modelTag = ''; let modelTag = '';
...@@ -140,7 +134,8 @@ ...@@ -140,7 +134,8 @@
const pullModelHandler = async () => { const pullModelHandler = async () => {
const sanitizedModelTag = modelTag.trim().replace(/^ollama\s+(run|pull)\s+/, ''); const sanitizedModelTag = modelTag.trim().replace(/^ollama\s+(run|pull)\s+/, '');
if (modelDownloadStatus[sanitizedModelTag]) { console.log($MODEL_DOWNLOAD_POOL);
if ($MODEL_DOWNLOAD_POOL[sanitizedModelTag]) {
toast.error( toast.error(
$i18n.t(`Model '{{modelTag}}' is already in queue for downloading.`, { $i18n.t(`Model '{{modelTag}}' is already in queue for downloading.`, {
modelTag: sanitizedModelTag modelTag: sanitizedModelTag
...@@ -148,40 +143,117 @@ ...@@ -148,40 +143,117 @@
); );
return; return;
} }
if (Object.keys(modelDownloadStatus).length === 3) { if (Object.keys($MODEL_DOWNLOAD_POOL).length === MAX_PARALLEL_DOWNLOADS) {
toast.error( toast.error(
$i18n.t('Maximum of 3 models can be downloaded simultaneously. Please try again later.') $i18n.t('Maximum of 3 models can be downloaded simultaneously. Please try again later.')
); );
return; return;
} }
modelTransferring = true; const res = await pullModel(localStorage.token, sanitizedModelTag, '0').catch((error) => {
toast.error(error);
return null;
});
if (res) {
const reader = res.body
.pipeThrough(new TextDecoderStream())
.pipeThrough(splitStream('\n'))
.getReader();
modelDownloadQueue.push( while (true) {
{ modelName: sanitizedModelTag }, try {
async (data: { modelName: string; success: boolean; error?: Error }) => { const { value, done } = await reader.read();
const { modelName } = data; if (done) break;
// Remove the downloaded model
delete modelDownloadStatus[modelName];
modelDownloadStatus = { ...modelDownloadStatus }; let lines = value.split('\n');
for (const line of lines) {
if (line !== '') {
let data = JSON.parse(line);
console.log(data);
if (data.error) {
throw data.error;
}
if (data.detail) {
throw data.detail;
}
if (data.id) {
MODEL_DOWNLOAD_POOL.set({
...$MODEL_DOWNLOAD_POOL,
[sanitizedModelTag]: {
...$MODEL_DOWNLOAD_POOL[sanitizedModelTag],
requestId: data.id,
reader,
done: false
}
});
console.log(data);
}
if (data.status) {
if (data.digest) {
let downloadProgress = 0;
if (data.completed) {
downloadProgress = Math.round((data.completed / data.total) * 1000) / 10;
} else {
downloadProgress = 100;
}
if (!data.success) { MODEL_DOWNLOAD_POOL.set({
toast.error(data.error); ...$MODEL_DOWNLOAD_POOL,
} else { [sanitizedModelTag]: {
toast.success( ...$MODEL_DOWNLOAD_POOL[sanitizedModelTag],
$i18n.t(`Model '{{modelName}}' has been successfully downloaded.`, { modelName }) pullProgress: downloadProgress,
); digest: data.digest
}
});
} else {
toast.success(data.status);
const notification = new Notification($WEBUI_NAME, { MODEL_DOWNLOAD_POOL.set({
body: $i18n.t(`Model '{{modelName}}' has been successfully downloaded.`, { modelName }), ...$MODEL_DOWNLOAD_POOL,
icon: `${WEBUI_BASE_URL}/static/favicon.png` [sanitizedModelTag]: {
}); ...$MODEL_DOWNLOAD_POOL[sanitizedModelTag],
done: data.status === 'success'
}
});
}
}
}
}
} catch (error) {
console.log(error);
if (typeof error !== 'string') {
error = error.message;
}
models.set(await getModels()); toast.error(error);
// opts.callback({ success: false, error, modelName: opts.modelName });
} }
} }
);
console.log($MODEL_DOWNLOAD_POOL[sanitizedModelTag]);
if ($MODEL_DOWNLOAD_POOL[sanitizedModelTag].done) {
toast.success(
$i18n.t(`Model '{{modelName}}' has been successfully downloaded.`, {
modelName: sanitizedModelTag
})
);
models.set(await getModels(localStorage.token));
} else {
toast.error('Download canceled');
}
delete $MODEL_DOWNLOAD_POOL[sanitizedModelTag];
MODEL_DOWNLOAD_POOL.set({
...$MODEL_DOWNLOAD_POOL
});
}
modelTag = ''; modelTag = '';
modelTransferring = false; modelTransferring = false;
...@@ -352,88 +424,18 @@ ...@@ -352,88 +424,18 @@
models.set(await getModels()); models.set(await getModels());
}; };
const pullModelHandlerProcessor = async (opts: { modelName: string; callback: Function }) => { const cancelModelPullHandler = async (model: string) => {
const res = await pullModel(localStorage.token, opts.modelName, selectedOllamaUrlIdx).catch( const { reader, requestId } = $MODEL_DOWNLOAD_POOL[model];
(error) => { if (reader) {
opts.callback({ success: false, error, modelName: opts.modelName }); await reader.cancel();
return null;
}
);
if (res) {
const reader = res.body
.pipeThrough(new TextDecoderStream())
.pipeThrough(splitStream('\n'))
.getReader();
while (true) {
try {
const { value, done } = await reader.read();
if (done) break;
let lines = value.split('\n');
for (const line of lines) {
if (line !== '') {
let data = JSON.parse(line);
console.log(data);
if (data.error) {
throw data.error;
}
if (data.detail) {
throw data.detail;
}
if (data.id) {
modelDownloadStatus[opts.modelName] = {
...modelDownloadStatus[opts.modelName],
requestId: data.id,
reader,
done: false
};
console.log(data);
}
if (data.status) {
if (data.digest) {
let downloadProgress = 0;
if (data.completed) {
downloadProgress = Math.round((data.completed / data.total) * 1000) / 10;
} else {
downloadProgress = 100;
}
modelDownloadStatus[opts.modelName] = {
...modelDownloadStatus[opts.modelName],
pullProgress: downloadProgress,
digest: data.digest
};
} else {
toast.success(data.status);
modelDownloadStatus[opts.modelName] = {
...modelDownloadStatus[opts.modelName],
done: data.status === 'success'
};
}
}
}
}
} catch (error) {
console.log(error);
if (typeof error !== 'string') {
error = error.message;
}
opts.callback({ success: false, error, modelName: opts.modelName });
}
}
console.log(modelDownloadStatus[opts.modelName]);
if (modelDownloadStatus[opts.modelName].done) { await cancelOllamaRequest(localStorage.token, requestId);
opts.callback({ success: true, modelName: opts.modelName }); delete $MODEL_DOWNLOAD_POOL[model];
} else { MODEL_DOWNLOAD_POOL.set({
opts.callback({ success: false, error: 'Download canceled', modelName: opts.modelName }); ...$MODEL_DOWNLOAD_POOL
} });
await deleteModel(localStorage.token, model);
toast.success(`${model} download has been canceled`);
} }
}; };
...@@ -503,18 +505,6 @@ ...@@ -503,18 +505,6 @@
ollamaVersion = await getOllamaVersion(localStorage.token).catch((error) => false); ollamaVersion = await getOllamaVersion(localStorage.token).catch((error) => false);
liteLLMModelInfo = await getLiteLLMModelInfo(localStorage.token); liteLLMModelInfo = await getLiteLLMModelInfo(localStorage.token);
}); });
const cancelModelPullHandler = async (model: string) => {
const { reader, requestId } = modelDownloadStatus[model];
if (reader) {
await reader.cancel();
await cancelOllamaRequest(localStorage.token, requestId);
delete modelDownloadStatus[model];
await deleteModel(localStorage.token, model);
toast.success(`${model} download has been canceled`);
}
};
</script> </script>
<div class="flex flex-col h-full justify-between text-sm"> <div class="flex flex-col h-full justify-between text-sm">
...@@ -643,9 +633,9 @@ ...@@ -643,9 +633,9 @@
> >
</div> </div>
{#if Object.keys(modelDownloadStatus).length > 0} {#if Object.keys($MODEL_DOWNLOAD_POOL).length > 0}
{#each Object.keys(modelDownloadStatus) as model} {#each Object.keys($MODEL_DOWNLOAD_POOL) as model}
{#if 'pullProgress' in modelDownloadStatus[model]} {#if 'pullProgress' in $MODEL_DOWNLOAD_POOL[model]}
<div class="flex flex-col"> <div class="flex flex-col">
<div class="font-medium mb-1">{model}</div> <div class="font-medium mb-1">{model}</div>
<div class=""> <div class="">
...@@ -655,10 +645,10 @@ ...@@ -655,10 +645,10 @@
class="dark:bg-gray-600 bg-gray-500 text-xs font-medium text-gray-100 text-center p-0.5 leading-none rounded-full" class="dark:bg-gray-600 bg-gray-500 text-xs font-medium text-gray-100 text-center p-0.5 leading-none rounded-full"
style="width: {Math.max( style="width: {Math.max(
15, 15,
modelDownloadStatus[model].pullProgress ?? 0 $MODEL_DOWNLOAD_POOL[model].pullProgress ?? 0
)}%" )}%"
> >
{modelDownloadStatus[model].pullProgress ?? 0}% {$MODEL_DOWNLOAD_POOL[model].pullProgress ?? 0}%
</div> </div>
</div> </div>
...@@ -689,9 +679,9 @@ ...@@ -689,9 +679,9 @@
</button> </button>
</Tooltip> </Tooltip>
</div> </div>
{#if 'digest' in modelDownloadStatus[model]} {#if 'digest' in $MODEL_DOWNLOAD_POOL[model]}
<div class="mt-1 text-xs dark:text-gray-500" style="font-size: 0.5rem;"> <div class="mt-1 text-xs dark:text-gray-500" style="font-size: 0.5rem;">
{modelDownloadStatus[model].digest} {$MODEL_DOWNLOAD_POOL[model].digest}
</div> </div>
{/if} {/if}
</div> </div>
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment