fix/desktop install builtin packages (#572)

* Fix the worker to install the requirements when the snapshot file is also specified

* Update bin/dump_artifacts to vendor the necessary builtin wheels and generates requirements.txt to call micropip.install() at runtime

* Fix reading requirements.txt

* Fix dump_artifacts.ts to fetch the Pyodide files instead of reading the predownloaded local ones

packages/desktop/bin/dump_artifacts.ts

@@ -6,12 +6,20 @@ import path from "path";
 import fsPromises from "fs/promises";
 import fsExtra from "fs-extra";
 import fetch from "node-fetch";
-import { loadPyodide, PyodideInterface } from "pyodide";
+import {
+  loadPyodide,
+  type PyodideInterface,
+  version as pyodideVersion,
+} from "pyodide";
 import { parseRequirementsTxt } from "@stlite/common";
 
 // @ts-ignore
 global.fetch = fetch; // The global `fetch()` is necessary for micropip.install() to load the remote packages.
 
+function makePyodideUrl(filename: string): string {
+  return `https://cdn.jsdelivr.net/pyodide/v${pyodideVersion}/full/${filename}`;
+}
+
 interface CopyBuildDirectoryOptions {
   keepOld: boolean;
   copyTo: string;
@@ -54,6 +62,50 @@ async function copyBuildDirectory(options: CopyBuildDirectoryOptions) {
   await fsExtra.copy(sourceDir, options.copyTo);
 }
 
+interface InspectUsedBuiltinPackagesOptions {
+  requirements: string[];
+  useLocalKernelWheels: boolean;
+}
+/**
+ * Get the list of the built-in packages used by the given requirements.
+ * These package files (`pyodide/*.whl`) will be vendored in the app executable
+ * and loaded at runtime to avoid problems such as https://github.com/whitphx/stlite/issues/558
+ */
+async function inspectUsedBuiltinPackages(
+  options: InspectUsedBuiltinPackagesOptions
+): Promise<string[]> {
+  if (options.requirements.length === 0) {
+    return [];
+  }
+
+  const pyodide = await loadPyodide();
+
+  await pyodide.loadPackage("micropip");
+  const micropip = pyodide.pyimport("micropip");
+
+  micropip.add_mock_package("streamlit", "1.21.0");
+
+  await micropip.install(options.requirements);
+  return Object.entries(pyodide.loadedPackages)
+    .filter(([, channel]) => channel === "default channel")
+    .map(([name]) => name);
+}
+
+async function loadPyodideBuiltinPackageData(): Promise<
+  Record<
+    string,
+    { name: string; version: string; file_name: string; depends: string[] }
+  >
+> {
+  const url = makePyodideUrl("repodata.json");
+
+  console.log(`Load the Pyodide repodata.json from ${url}`);
+  const res = await fetch(url);
+  const resJson = await res.json();
+
+  return resJson.packages;
+}
+
 async function installLocalWheel(pyodide: PyodideInterface, localPath: string) {
   console.log(`Install the local wheel ${localPath}`);
 
@@ -70,6 +122,7 @@ async function installLocalWheel(pyodide: PyodideInterface, localPath: string) {
 interface CreateSitePackagesSnapshotOptions {
   useLocalKernelWheels: boolean;
   requirements: string[];
+  usedBuiltinPackages: string[];
   saveTo: string;
 }
 async function createSitePackagesSnapshot(
@@ -119,10 +172,29 @@ async function createSitePackagesSnapshot(
     await micropip.install.callKwargs(wheelUrls, { keep_going: true });
   }
 
+  const micropip = pyodide.pyimport("micropip");
+
+  const pyodideBuiltinPackageMap = await loadPyodideBuiltinPackageData();
+
+  if (options.usedBuiltinPackages.length > 0) {
+    console.log(
+      "Mocking builtin packages so that they will not be included in the site-packages snapshot because these will be installed from the vendored wheel files at runtime..."
+    );
+    options.usedBuiltinPackages.forEach((pkg) => {
+      const packageInfo = pyodideBuiltinPackageMap[pkg];
+      if (packageInfo == null) {
+        throw new Error(`Package ${pkg} is not found in the lock file.`);
+      }
+
+      console.log(`Mock ${packageInfo.name} ${packageInfo.version}`);
+      micropip.add_mock_package(packageInfo.name, packageInfo.version);
+    });
+  }
+
   console.log(
     `Install the requirements ${JSON.stringify(options.requirements)}`
   );
-  const micropip = pyodide.pyimport("micropip");
+
   await micropip.install.callKwargs(options.requirements, { keep_going: true });
 
   console.log("Archive the site-packages director(y|ies)");
@@ -167,6 +239,16 @@ async function readRequirements(
   return parseRequirementsTxt(requirementsTxtData);
 }
 
+async function writeRequirements(
+  requirementsTxtPath: string,
+  requirements: string[]
+): Promise<void> {
+  const requirementsTxtData = requirements.join("\n");
+  await fsPromises.writeFile(requirementsTxtPath, requirementsTxtData, {
+    encoding: "utf-8",
+  });
+}
+
 // Original: kernel/src/requirements.ts
 // TODO: Be DRY
 function verifyRequirements(requirements: string[]) {
@@ -188,6 +270,42 @@ function verifyRequirements(requirements: string[]) {
   });
 }
 
+interface DownloadPyodideBuiltinPackageWheelsOptions {
+  packages: string[];
+  destDir: string;
+}
+async function downloadPyodideBuiltinPackageWheels(
+  options: DownloadPyodideBuiltinPackageWheelsOptions
+) {
+  const pyodideBuiltinPackages = await loadPyodideBuiltinPackageData();
+  const usedBuiltInPackages = options.packages.map(
+    (pkgName) => pyodideBuiltinPackages[pkgName]
+  );
+  const usedBuiltinPackageUrls = usedBuiltInPackages.map((pkg) =>
+    makePyodideUrl(pkg.file_name)
+  );
+
+  console.log("Downloading the used built-in packages...");
+  await Promise.all(
+    usedBuiltinPackageUrls.map(async (pkgUrl) => {
+      const dstPath = path.resolve(
+        options.destDir,
+        "./pyodide",
+        path.basename(pkgUrl)
+      );
+      console.log(`Download ${pkgUrl} to ${dstPath}`);
+      const res = await fetch(pkgUrl);
+      if (!res.ok) {
+        throw new Error(
+          `Failed to download ${pkgUrl}: ${res.status} ${res.statusText}`
+        );
+      }
+      const buf = await res.arrayBuffer();
+      await fsPromises.writeFile(dstPath, Buffer.from(buf));
+    })
+  );
+}
+
 yargs(hideBin(process.argv))
   .command(
     "* <appHomeDirSource> [packages..]",
@@ -246,14 +364,34 @@ yargs(hideBin(process.argv))
     }
     verifyRequirements(requirements);
 
+    const usedBuiltinPackages = await inspectUsedBuiltinPackages({
+      requirements: requirements,
+      useLocalKernelWheels: args.localKernelWheels,
+    });
+    console.log("The built-in packages loaded for the given requirements:");
+    console.log(usedBuiltinPackages);
+
     await copyBuildDirectory({ copyTo: destDir, keepOld: args.keepOldBuild });
     await createSitePackagesSnapshot({
       useLocalKernelWheels: args.localKernelWheels,
       requirements: requirements,
+      usedBuiltinPackages,
       saveTo: path.resolve(destDir, "./site-packages-snapshot.tar.gz"), // This path will be loaded in the `readSitePackagesSnapshot` handler in electron/main.ts.
     });
+    // The `requirements.txt` file will be needed to call `micropip.install()` at runtime.
+    // The built-in packages will be vendored in the build artifact as wheel files
+    // and `micropip.install()` will install them at runtime,
+    // while the packages downloaded from PyPI will have been included in the site-packages snapshot.
+    await writeRequirements(
+      path.resolve(destDir, "./requirements.txt"), // This path will be loaded in the `readRequirements` handler in electron/main.ts.
+      requirements
+    );
     await copyStreamlitAppDirectory({
       sourceDir: args.appHomeDirSource,
       copyTo: path.resolve(destDir, "./streamlit_app"), // This path will be loaded in the `readStreamlitAppDirectory` handler in electron/main.ts.
     });
+    await downloadPyodideBuiltinPackageWheels({
+      packages: usedBuiltinPackages,
+      destDir,
+    });
   });

packages/desktop/electron/main.ts

@@ -49,6 +49,20 @@ const createWindow = () => {
     );
     return fsPromises.readFile(archiveFilePath);
   });
+  ipcMain.handle("readRequirements", async (ev): Promise<string[]> => {
+    if (!isValidIpcSender(ev.senderFrame)) {
+      throw new Error("Invalid IPC sender");
+    }
+
+    const requirementsTxtPath = path.resolve(__dirname, "../requirements.txt");
+    const requirementsTxtData = await fsPromises.readFile(requirementsTxtPath, {
+      encoding: "utf-8",
+    });
+    return requirementsTxtData
+      .split("\n")
+      .map((r) => r.trim())
+      .filter((r) => r.length > 0); // Assuming that the input `requirements.txt` file is generated by `dump-artifacts.js` and can be parsed with this simple logic.
+  });
   ipcMain.handle(
     "readStreamlitAppDirectory",
     async (ev): Promise<Record<string, Buffer>> => {

packages/desktop/electron/preload.ts

@@ -3,6 +3,7 @@ import { contextBridge, ipcRenderer } from "electron";
 contextBridge.exposeInMainWorld("archives", {
   readSitePackagesSnapshot: () =>
     ipcRenderer.invoke("readSitePackagesSnapshot"),
+  readRequirements: () => ipcRenderer.invoke("readRequirements"),
   readStreamlitAppDirectory: () =>
     ipcRenderer.invoke("readStreamlitAppDirectory"),
 });

packages/desktop/package.json

@@ -25,6 +25,7 @@
     "start": "concurrently \"cross-env BROWSER=none yarn start:web\" \"wait-on http://localhost:3000 && yarn start:electron\" \"wait-on http://localhost:3000 && tsc -p electron && cross-env NODE_ENV=\"development\" electron .\"",
     "build:app": "yarn build:web && yarn build:electron && yarn build:pyodide",
     "build": "yarn build:app && yarn build:bin",
+    "dump:dev": "ts-node ./bin/dump_artifacts.ts",
     "dump": "dump-stlite-desktop-artifacts",
     "serve": "cross-env NODE_ENV=production electron .",
     "pack": "electron-builder --dir",

packages/desktop/src/App.tsx

@@ -24,38 +24,41 @@ function App() {
 
     Promise.all([
       window.archives.readSitePackagesSnapshot(),
+      window.archives.readRequirements(),
       window.archives.readStreamlitAppDirectory(),
-    ]).then(([sitePackagesSnapshotFileBin, streamlitAppFiles]) => {
-      if (unmounted) {
-        return;
-      }
+    ]).then(
+      ([sitePackagesSnapshotFileBin, requirements, streamlitAppFiles]) => {
+        if (unmounted) {
+          return;
+        }
 
-      const files: StliteKernelOptions["files"] = {};
-      Object.keys(streamlitAppFiles).forEach((path) => {
-        const data = streamlitAppFiles[path];
-        files[path] = {
-          data,
-        };
-      });
+        const files: StliteKernelOptions["files"] = {};
+        Object.keys(streamlitAppFiles).forEach((path) => {
+          const data = streamlitAppFiles[path];
+          files[path] = {
+            data,
+          };
+        });
 
-      const mountedSitePackagesSnapshotFilePath =
-        "/tmp/site-packages-snapshot.tar.gz";
-      kernel = new StliteKernel({
-        entrypoint: "streamlit_app.py",
-        files: {
-          ...files,
-          [mountedSitePackagesSnapshotFilePath]: {
-            data: sitePackagesSnapshotFileBin,
+        const mountedSitePackagesSnapshotFilePath =
+          "/tmp/site-packages-snapshot.tar.gz";
+        kernel = new StliteKernel({
+          entrypoint: "streamlit_app.py",
+          files: {
+            ...files,
+            [mountedSitePackagesSnapshotFilePath]: {
+              data: sitePackagesSnapshotFileBin,
+            },
           },
-        },
-        archives: [],
-        requirements: [],
-        mountedSitePackagesSnapshotFilePath,
-        pyodideEntrypointUrl,
-        ...makeToastKernelCallbacks(),
-      });
-      setKernel(kernel);
-    });
+          archives: [],
+          requirements,
+          mountedSitePackagesSnapshotFilePath,
+          pyodideEntrypointUrl,
+          ...makeToastKernelCallbacks(),
+        });
+        setKernel(kernel);
+      }
+    );
 
     return () => {
       unmounted = true;

packages/desktop/src/global.d.ts

@@ -2,6 +2,7 @@ export declare global {
   interface Window {
     archives: {
       readSitePackagesSnapshot: () => Promise<Uint8Array>;
+      readRequirements: () => Promise<string[]>;
       readStreamlitAppDirectory: () => Promise<Record<string, Buffer>>;
     };
   }

packages/kernel/src/worker.ts

@@ -147,14 +147,18 @@ async function loadPyodideAndPackages() {
     console.debug("Mock pyarrow");
     mockPyArrow(pyodide);
     console.debug("Mocked pyarrow");
+  } else {
+    throw new Error(`Neither snapshot nor wheel files are provided.`);
+  }
 
+  if (requirements.length > 0) {
     postProgressMessage("Installing the requirements.");
     console.debug("Installing the requirements:", requirements);
     verifyRequirements(requirements); // Blocks the not allowed wheel URL schemes.
+    await pyodide.loadPackage("micropip");
+    const micropip = pyodide.pyimport("micropip");
     await micropip.install.callKwargs(requirements, { keep_going: true });
     console.debug("Installed the requirements:", requirements);
-  } else {
-    throw new Error(`Neither snapshot nor wheel files are provided.`);
   }
 
   // The following code is necessary to avoid errors like  `NameError: name '_imp' is not defined`