From af998ee0a7f3cda50e2f91036b12f9524104da44 Mon Sep 17 00:00:00 2001 From: Timothy Carambat Date: Mon, 19 Jan 2026 15:08:07 -0800 Subject: [PATCH] Docker model runner download from UI (#4884) * Enable downloads of DMR models from UI * add utils + dev build * linting --- .github/workflows/dev-build.yaml | 2 +- .../DockerModelRunnerOptions/index.jsx | 44 ++++++- frontend/src/models/utils/dmrUtils.js | 77 ++++++++++++ server/endpoints/utils.js | 5 + .../endpoints/utils/dockerModelRunnerUtils.js | 110 ++++++++++++++++++ 5 files changed, 232 insertions(+), 6 deletions(-) create mode 100644 frontend/src/models/utils/dmrUtils.js create mode 100644 server/endpoints/utils/dockerModelRunnerUtils.js diff --git a/.github/workflows/dev-build.yaml b/.github/workflows/dev-build.yaml index 446a8ff4..2f28b4fd 100644 --- a/.github/workflows/dev-build.yaml +++ b/.github/workflows/dev-build.yaml @@ -6,7 +6,7 @@ concurrency: on: push: - branches: ['model-table-component'] # put your current branch to create a build. Core team only. + branches: ['docker-model-runner-download-from-ui'] # put your current branch to create a build. Core team only. paths-ignore: - '**.md' - 'cloud-deployments/*' diff --git a/frontend/src/components/LLMSelection/DockerModelRunnerOptions/index.jsx b/frontend/src/components/LLMSelection/DockerModelRunnerOptions/index.jsx index b209c24b..56320cfa 100644 --- a/frontend/src/components/LLMSelection/DockerModelRunnerOptions/index.jsx +++ b/frontend/src/components/LLMSelection/DockerModelRunnerOptions/index.jsx @@ -10,6 +10,8 @@ import { Link } from "react-router-dom"; import ModelTable from "@/components/lib/ModelTable"; import ModelTableLayout from "@/components/lib/ModelTable/layout"; import ModelTableLoadingSkeleton from "@/components/lib/ModelTable/loading"; +import DMRUtils from "@/models/utils/dmrUtils"; +import showToast from "@/utils/toast"; export default function DockerModelRunnerOptions({ settings }) { const { @@ -237,12 +239,44 @@ function DockerModelRunnerModelSelection({ setFilteredModels(Array.from(filteredModels.values())); }, [searchQuery]); - function downloadModel(modelId, _fileSize, progressCallback) { - const [name, tag] = modelId.split(":"); + async function downloadModel(modelId, fileSize, progressCallback) { + try { + if ( + !window.confirm( + `Are you sure you want to download this model? It is ${fileSize} in size and may take a while to download.` + ) + ) + return; + const { success, error } = await DMRUtils.downloadModel( + modelId, + basePath, + progressCallback + ); + if (!success) + throw new Error( + error || "An error occurred while downloading the model" + ); + progressCallback(100); + handleSetActiveModel(modelId); - // Open the model in the Docker Hub (via browser since they may not be installed locally) - window.open(`https://hub.docker.com/layers/${name}/${tag}`, "_blank"); - progressCallback(100); + const existingModels = [...customModels]; + const newModel = existingModels.find((model) => model.id === modelId); + if (newModel) { + newModel.downloaded = true; + setCustomModels(existingModels); + setFilteredModels(existingModels); + setSearchQuery(""); + } + } catch (e) { + console.error("Error downloading model:", e); + showToast( + e.message || "An error occurred while downloading the model", + "error", + { clear: true } + ); + } finally { + setLoading(false); + } } function groupModelsByAlias(models) { diff --git a/frontend/src/models/utils/dmrUtils.js b/frontend/src/models/utils/dmrUtils.js new file mode 100644 index 00000000..6103ef3b --- /dev/null +++ b/frontend/src/models/utils/dmrUtils.js @@ -0,0 +1,77 @@ +import { API_BASE } from "@/utils/constants"; +import { baseHeaders } from "@/utils/request"; +import { safeJsonParse } from "@/utils/request"; + +const DMRUtils = { + /** + * Download a DMR model. + * @param {string} modelId - The ID of the model to download. + * @param {(percentage: number) => void} progressCallback - The callback to receive the progress percentage. If the model is already downloaded, it will be called once with 100. + * @returns {Promise<{success: boolean, error: string|null}>} + */ + downloadModel: async function ( + modelId, + basePath = "", + progressCallback = () => {} + ) { + // eslint-disable-next-line no-async-promise-executor + return new Promise(async (resolve) => { + try { + const response = await fetch(`${API_BASE}/utils/dmr/download-model`, { + method: "POST", + headers: baseHeaders(), + body: JSON.stringify({ modelId, basePath }), + }); + + if (!response.ok) + throw new Error("Error downloading model: " + response.statusText); + const reader = response.body.getReader(); + let done = false; + + while (!done) { + const { value, done: readerDone } = await reader.read(); + if (readerDone) { + done = true; + resolve({ success: true }); + } else { + const chunk = new TextDecoder("utf-8").decode(value); + const lines = chunk.split("\n"); + for (const line of lines) { + if (line.startsWith("data:")) { + const data = safeJsonParse(line.slice(5)); + switch (data?.type) { + case "success": + done = true; + resolve({ success: true }); + break; + case "error": + done = true; + resolve({ + success: false, + error: data?.error || data?.message, + }); + break; + case "progress": + progressCallback(data?.percentage); + break; + default: + break; + } + } + } + } + } + } catch (error) { + console.error("Error downloading model:", error); + resolve({ + success: false, + error: + error?.message || "An error occurred while downloading the model", + }); + } + }); + }, + // Uninstall a DMR model is not supported via the API +}; + +export default DMRUtils; diff --git a/server/endpoints/utils.js b/server/endpoints/utils.js index 30d2e9b7..aa449709 100644 --- a/server/endpoints/utils.js +++ b/server/endpoints/utils.js @@ -21,6 +21,11 @@ function utilEndpoints(app) { response.sendStatus(500).end(); } }); + + const { + dockerModelRunnerUtilsEndpoints, + } = require("./utils/dockerModelRunnerUtils"); + dockerModelRunnerUtilsEndpoints(app); } function getGitVersion() { diff --git a/server/endpoints/utils/dockerModelRunnerUtils.js b/server/endpoints/utils/dockerModelRunnerUtils.js new file mode 100644 index 00000000..6e7f0edc --- /dev/null +++ b/server/endpoints/utils/dockerModelRunnerUtils.js @@ -0,0 +1,110 @@ +const { validatedRequest } = require("../../utils/middleware/validatedRequest"); +const { + flexUserRoleValid, + ROLES, +} = require("../../utils/middleware/multiUserProtected"); +const { reqBody } = require("../../utils/http"); +const { safeJsonParse } = require("../../utils/http"); + +/** + * Decode HTML entities from a string. + * The DMR response is encoded with HTML entities, so we need to decode them + * so we can parse the JSON and report the progress percentage. + * @param {string} str - The string to decode. + * @returns {string} The decoded string. + */ +function decodeHtmlEntities(str) { + return str + .replace(/"/g, '"') + .replace(/'/g, "'") + .replace(/</g, "<") + .replace(/>/g, ">") + .replace(/&/g, "&"); +} + +function dockerModelRunnerUtilsEndpoints(app) { + if (!app) return; + const { + parseDockerModelRunnerEndpoint, + } = require("../../utils/AiProviders/dockerModelRunner"); + + app.post( + "/utils/dmr/download-model", + [validatedRequest, flexUserRoleValid([ROLES.admin])], + async (request, response) => { + try { + const { modelId, basePath = "" } = reqBody(request); + const dmrUrl = new URL( + parseDockerModelRunnerEndpoint( + basePath ?? process.env.DOCKER_MODEL_RUNNER_BASE_PATH, + "dmr" + ) + ); + dmrUrl.pathname = "/models/create"; + response.writeHead(200, { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }); + + const dmrResponse = await fetch(dmrUrl.toString(), { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ from: String(modelId) }), + }); + if (!dmrResponse.ok) + throw new Error( + dmrResponse.statusText || + "An error occurred while downloading the model" + ); + const reader = dmrResponse.body.getReader(); + let done = false; + while (!done) { + const { value, done: readerDone } = await reader.read(); + if (readerDone) done = true; + + const chunk = new TextDecoder("utf-8").decode(value); + const lines = chunk.split("\n"); + for (const line of lines) { + if (!line.trim()) continue; + const decodedLine = decodeHtmlEntities(line); + const data = safeJsonParse(decodedLine); + if (!data) continue; + + if (data.type === "error") { + throw new Error( + data.message || "An error occurred while downloading the model" + ); + } else if (data.type === "success") { + response.write( + `data: ${JSON.stringify({ type: "success", percentage: 100, message: "Model downloaded successfully" })}\n\n` + ); + done = true; + } else if (data.type === "progress") { + const percentage = + data.total > 0 + ? Math.round((data.pulled / data.total) * 100) + : 0; + response.write( + `data: ${JSON.stringify({ type: "progress", percentage, message: data.message })}\n\n` + ); + } + } + } + } catch (e) { + console.error(e); + response.write( + `data: ${JSON.stringify({ type: "error", message: e.message })}\n\n` + ); + } finally { + response.end(); + } + } + ); +} + +module.exports = { + dockerModelRunnerUtilsEndpoints, +};