Docker model runner download from UI (#4884)

* Enable downloads of DMR models from UI

* add utils + dev build

* linting

.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/*'

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) {

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;

server/endpoints/utils.js

@@ -21,6 +21,11 @@ function utilEndpoints(app) {
       response.sendStatus(500).end();
     }
   });
+
+  const {
+    dockerModelRunnerUtilsEndpoints,
+  } = require("./utils/dockerModelRunnerUtils");
+  dockerModelRunnerUtilsEndpoints(app);
 }
 
 function getGitVersion() {

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(/&lt;/g, "<")
+    .replace(/&gt;/g, ">")
+    .replace(/&amp;/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,
+};