feature/pyodide package cache

src/py/mat3ra/notebooks_utils/packages.py

@@ -1,23 +1,37 @@
+"""Install notebook dependencies from config.yml.
+
+Routes to the appropriate backend: micropip (Pyodide) or pip (Python).
+In Pyodide, downloaded wheels are cached in IndexedDB so subsequent
+kernel restarts skip network downloads.
+"""
+
 from .ipython.packages.install import install_packages_python
 from .primitive.environment import is_pyodide_environment
 
 
-async def install_packages(notebook_name_pattern: str, config_file_path: str = "", verbose: bool = True):
-    """
-    Install the packages listed in config.yml for the given notebook name pattern.
+async def install_packages(
+    notebook_name_pattern: str,
+    config_file_path: str = "",
+    verbose: bool = True,
+    force: bool = False,
+    use_cache: bool = True,
+):
+    """Install the packages listed in config.yml for the given notebook name pattern.
 
     Usage in notebooks:
         from mat3ra.notebooks_utils.packages import install_packages
         await install_packages("my_notebook")
 
     Args:
-        notebook_name_pattern (str): Pattern matched against notebook names in config.yml.
-        config_file_path (str): Path to config.yml; empty string uses the JupyterLite default (/drive/config.yml).
-        verbose (bool): Whether to print install progress.
+        notebook_name_pattern: Pattern matched against notebook names in config.yml.
+        config_file_path: Path to config.yml; empty string uses the JupyterLite default.
+        verbose: Whether to print install progress.
+        force: Clear wheel cache and reinstall from network (Pyodide only).
+        use_cache: Use IndexedDB wheel cache for faster restarts (Pyodide only, default True).
     """
     if is_pyodide_environment():
         from .pyodide.packages.install import install_packages_pyodide
 
-        await install_packages_pyodide(notebook_name_pattern, verbose)
+        await install_packages_pyodide(notebook_name_pattern, verbose, force=force, use_cache=use_cache)
     else:
         install_packages_python(notebook_name_pattern, verbose)

src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py

@@ -8,8 +8,19 @@
     """
 (function() {
     const open = () => new Promise(resolve => {
-        const req = indexedDB.open("mat3ra", 1);
-        req.onupgradeneeded = () => req.result.createObjectStore("tokens");
+        const req = indexedDB.open("mat3ra", 2);
+        req.onupgradeneeded = (event) => {
+            const db = event.target.result;
+            if (!db.objectStoreNames.contains("tokens")) {
+                db.createObjectStore("tokens");
+            }
+            if (!db.objectStoreNames.contains("wheel_manifest")) {
+                db.createObjectStore("wheel_manifest");
+            }
+            if (!db.objectStoreNames.contains("wheel_data")) {
+                db.createObjectStore("wheel_data");
+            }
+        };
         req.onsuccess = () => resolve(req.result);
         req.onerror = () => resolve(null);
     });

src/py/mat3ra/notebooks_utils/pyodide/io.py

@@ -5,6 +5,11 @@
 
 from IPython.display import Javascript, display
 
+try:
+    from pyodide.http import pyfetch  # type: ignore
+except ImportError:
+    pyfetch = None  # type: ignore
+
 from ..core.io import set_data_python
 from ..primitive.logger import log
 
@@ -20,8 +25,6 @@ async def read_from_url_pyodide(url: str, as_bytes: bool = False) -> Union[str,
     Returns:
         str or bytes: The content.
     """
-    # `http` is a Pyodide module that will be installed in the Pyodide environment by default.
-    from pyodide.http import pyfetch  # type: ignore
 
     # Per https://developer.mozilla.org/en-US/docs/Web/API/Fetch_API/Using_Fetch
     response = await pyfetch(url)

src/py/mat3ra/notebooks_utils/pyodide/packages/install.py

@@ -1,11 +1,21 @@
+"""Pyodide package installer with IndexedDB wheel caching.
+
+On first install, packages are downloaded from PyPI and cached in IndexedDB.
+On subsequent kernel starts, cached wheels are written to emfs and installed
+locally — eliminating network downloads.
+"""
+
 import json
 import os
 import re
 import sys
+import time
 from typing import List, Tuple, Union
 
 from ...primitive.environment import ENVIRONMENT
 from ...primitive.logger import log
+from ..io import read_from_url_pyodide, write_to_file
+from . import wheel_cache
 
 try:
     import micropip  # type: ignore
@@ -14,8 +24,12 @@
 
 NODEPS_PREFIX = "nodeps:"
 URL_PREFIXES = ("http://", "https://", "emfs:/")
+LOCAL_URL_PREFIXES = ("emfs:/",)
+REMOTE_URL_PREFIXES = ("http://", "https://")
 VERSION_SPECIFIERS = ("==", ">=", "<=", "!=", "~=", ">", "<")
 
+WHEEL_CACHE_EMFS_DIRECTORY = "/tmp/wheel_cache"
+
 
 def get_config_yml_file_path(config_file_path: str) -> str:
     """
@@ -35,6 +49,10 @@ def get_config_yml_file_path(config_file_path: str) -> str:
 
 async def read_config_into_dict(config_file_path: str) -> dict:
     with open(get_config_yml_file_path(config_file_path), "r") as f:
+        # NOTE: PyYAML (yaml) is preinstalled at the JupyterLite build level as a required
+        # environment requirement. The import is placed inside this function (rather than
+        # at the top level of the module) to prevent any potential boot/load order issues
+        # during fresh kernel initialization.
         # import micropip  # type: ignore
         #
         # await micropip.install("pyyaml")
@@ -101,25 +119,37 @@ def should_install_packages(previous_hash: Union[str, None], requirements_hash:
     return previous_hash != requirements_hash
 
 
-def is_url_package(pkg: str) -> bool:
-    return pkg.startswith(URL_PREFIXES)
+def is_url_package(package_spec: str) -> bool:
+    return package_spec.startswith(URL_PREFIXES)
+
 
+def is_local_url_package(package_spec: str) -> bool:
+    return package_spec.startswith(LOCAL_URL_PREFIXES)
 
-def remove_nodeps_prefix(pkg: str) -> str:
-    return pkg.replace(NODEPS_PREFIX, "", 1) if pkg.startswith(NODEPS_PREFIX) else pkg
 
+def is_remote_url_package(package_spec: str) -> bool:
+    return package_spec.startswith(REMOTE_URL_PREFIXES)
 
-def package_has_version_specifier(pkg: str) -> bool:
-    spec = remove_nodeps_prefix(pkg)
+
+def remove_nodeps_prefix(package_spec: str) -> str:
+    return package_spec.replace(NODEPS_PREFIX, "", 1) if package_spec.startswith(NODEPS_PREFIX) else package_spec
+
+
+def package_has_version_specifier(package_spec: str) -> bool:
+    spec = remove_nodeps_prefix(package_spec)
     return any(op in spec for op in VERSION_SPECIFIERS)
 
 
-def should_reinstall_package(pkg: str, profile_changed: bool) -> bool:
-    return profile_changed and package_has_version_specifier(pkg) and not is_url_package(remove_nodeps_prefix(pkg))
+def should_reinstall_package(package_spec: str, profile_changed: bool) -> bool:
+    return (
+        profile_changed
+        and package_has_version_specifier(package_spec)
+        and not is_url_package(remove_nodeps_prefix(package_spec))
+    )
 
 
-def get_package_name(pkg: str) -> Union[str, None]:
-    spec = remove_nodeps_prefix(pkg)
+def get_package_name(package_spec: str) -> Union[str, None]:
+    spec = remove_nodeps_prefix(package_spec)
     match = re.match(r"^[A-Za-z0-9_.-]+", spec)
     return match.group(0) if match else None
 
@@ -128,15 +158,21 @@ def get_import_package_name(package_name: str) -> str:
     return package_name.replace("-", "_")
 
 
+def get_display_name(package_spec: str) -> str:
+    if "://" in package_spec:
+        return package_spec.split("/")[-1].split("-")[0]
+    return package_spec.split("==")[0]
+
+
 def clear_imported_package_modules(package_name: str):
     import_name = get_import_package_name(package_name)
     module_names = [name for name in sys.modules if name == import_name or name.startswith(f"{import_name}.")]
     for module_name in module_names:
         sys.modules.pop(module_name, None)
 
 
-async def uninstall_package_pyodide(pkg: str):
-    package_name = get_package_name(pkg)
+async def uninstall_package_pyodide(package_spec: str):
+    package_name = get_package_name(package_spec)
     if not package_name:
         return
     if not hasattr(micropip, "uninstall"):
@@ -148,56 +184,166 @@ async def uninstall_package_pyodide(pkg: str):
     clear_imported_package_modules(package_name)
 
 
-def get_install_spec_and_deps(pkg: str) -> Tuple[str, bool]:
-    if pkg.startswith(NODEPS_PREFIX):
-        return remove_nodeps_prefix(pkg), False
-    return pkg, not is_url_package(pkg)
+def get_install_spec_and_deps(package_spec: str) -> Tuple[str, bool]:
+    if package_spec.startswith(NODEPS_PREFIX):
+        return remove_nodeps_prefix(package_spec), False
+    return package_spec, not is_url_package(package_spec)
 
 
-async def install_package_pyodide(pkg: str, verbose: bool = True, reinstall: bool = False):
+def get_cache_key(package_spec: str) -> str:
     """
-    Install a package in a Pyodide environment.
+    Normalize a package spec into a cache key.
+    """
+    spec = remove_nodeps_prefix(package_spec)
+    return spec.lower().strip()
 
-    Args:
-        pkg (str): The name of the package to install. Can be prefixed with 'nodeps:' to skip dependencies.
-        verbose (bool): Whether to print the name of the installed package.
 
-    Examples:
-        await install_package_pyodide("numpy")  # installs with deps
-        await install_package_pyodide("nodeps:e3nn==0.4.4")  # installs without deps
+async def install_package_pyodide(
+    package_spec: str,
+    verbose: bool = True,
+    reinstall: bool = False,
+    use_cache: bool = True,
+):
+    """
+    Install a package in a Pyodide environment, using wheel cache when possible.
+
+    Args:
+        package_spec: Package name, version spec, or URL. Can be prefixed with 'nodeps:'.
+        verbose: Whether to log the installed package name.
+        reinstall: Whether to uninstall first.
+        use_cache: Whether to use the IndexedDB wheel cache.
     """
-    pkg, are_dependencies_installed = get_install_spec_and_deps(pkg)
+    raw_spec, are_dependencies_installed = get_install_spec_and_deps(package_spec)
+
     if reinstall:
-        await uninstall_package_pyodide(pkg)
+        await uninstall_package_pyodide(raw_spec)
+
+    installed_from_cache = False
+
+    if use_cache and is_remote_url_package(raw_spec):
+        installed_from_cache = await _install_from_cache_or_download(raw_spec, are_dependencies_installed)
+
+    if not installed_from_cache:
+        await micropip.install(raw_spec, deps=are_dependencies_installed)
 
-    await micropip.install(pkg, deps=are_dependencies_installed)
-    pkg_name = pkg.split("/")[-1].split("-")[0] if "://" in pkg else pkg.split("==")[0]
     if verbose:
-        log(f"Installed {pkg_name}", force_verbose=verbose)
+        log(f"Installed {get_display_name(raw_spec)}", force_verbose=verbose)
 
 
-async def install_packages_pyodide(notebook_name_pattern: str, verbose: bool = True):
+async def _install_from_cache_or_download(url: str, install_dependencies: bool) -> bool:
     """
-    Install the given packages, skipping if the list is unchanged since last run.
+    Try to install from cache; if miss, download, cache, and install. Returns True on success.
+    """
+    try:
+        # Cache key is the full URL, which ensures that if a new version is published
+        # to PyPI, the URL will change and the cache will automatically invalidate.
+        cache_key = url
+        filename = url.split("/")[-1].split("?")[0]
+
+        # Micropip requires a physical filepath or a web URL to perform installation;
+        # it cannot install directly from in-memory bytes. Therefore, we write the
+        # wheel bytes to Pyodide's local EMFS (Emscripten File System) temp directory first.
+        os.makedirs(WHEEL_CACHE_EMFS_DIRECTORY, exist_ok=True)
+        emfs_path = os.path.join(WHEEL_CACHE_EMFS_DIRECTORY, filename)
+
+        # 1. Attempt to hit the persistent IndexedDB cache.
+        cached_bytes = await wheel_cache.get_cached_wheel(cache_key)
+
+        if cached_bytes is not None:
+            # Cache Hit: Write in-memory bytes to EMFS and install locally (fast restart).
+            await write_to_file(emfs_path, cached_bytes)
+            await micropip.install(emfs_path, deps=install_dependencies)
+            return True
+
+        # 2. Cache Miss: Download from the remote URL, store in IndexedDB, and write to EMFS.
+        wheel_bytes = await read_from_url_pyodide(url, as_bytes=True)
+        if not isinstance(wheel_bytes, bytes):
+            raise TypeError("Expected bytes from read_from_url_pyodide")
+        await wheel_cache.put_cached_wheel(cache_key, wheel_bytes)
+
+        await write_to_file(emfs_path, wheel_bytes)
+        await micropip.install(emfs_path, deps=install_dependencies)
+
+        # 3. Update the cache manifest with the newly added package metadata.
+        manifest = await wheel_cache.read_manifest()
+        manifest[cache_key] = {
+            "filename": filename,
+            "size": len(wheel_bytes),
+            "cached_at": time.time(),
+        }
+        await wheel_cache.write_manifest(manifest)
+
+        return True
+    except Exception:
+        # If any part of the cache system fails (DB error, disk full), fail gracefully and
+        # let the caller fall back to micropip's default remote HTTP installer.
+        return False
+
+
+async def install_packages_pyodide(
+    notebook_name_pattern: str,
+    verbose: bool = True,
+    force: bool = False,
+    use_cache: bool = True,
+):
+    """
+    Install packages from config.yml, using IndexedDB wheel cache for persistence.
 
     Args:
-        notebook_name_pattern (str): Pattern matched against notebook names in config.yml.
-        verbose (bool): Whether to print the names of the installed packages.
+        notebook_name_pattern: Pattern matched against notebook names in config.yml.
+        verbose: Whether to log install progress.
+        force: Clear the wheel cache and reinstall everything from network.
+        use_cache: Whether to use the IndexedDB wheel cache (default True).
     """
+    if force and use_cache:
+        # A forced reinstall clears the whole wheel cache first to guarantee downloading
+        # the latest packages and dependencies from PyPI.
+        await wheel_cache.clear_cache()
+        if verbose:
+            log("Wheel cache cleared.", force_verbose=verbose)
+
+    # 1. Resolve and hash the list of packages specified for this notebook profile.
     packages = await get_package_list_from_config(get_config_yml_file_path(""), notebook_name_pattern)
     requirements_hash = str(hash(json.dumps(packages)))
+
+    # 2. Check if this profile's packages have already been installed in this session.
+    #    os.environ contains volatile session state which is lost on kernel restart.
     previous_hash = os.environ.get("requirements_hash")
+
+    # 3. If this is a fresh kernel session (os.environ is empty) and caching is enabled,
+    #    we read the persistent hash from IndexedDB. This prevents redundant reinstall
+    #    checks across kernel restarts.
+    if use_cache and previous_hash is None:
+        manifest = await wheel_cache.read_manifest()
+        previous_hash = manifest.get("_requirements_hash")
+
     profile_changed = previous_hash is not None and previous_hash != requirements_hash
-    if should_install_packages(previous_hash, requirements_hash):
-        for pkg in packages:
+
+    # 4. Trigger installation if either:
+    #    a) The profile is different or has never been run (previous_hash != requirements_hash)
+    #    b) Force install was explicitly requested (force=True)
+    if should_install_packages(previous_hash, requirements_hash) or force:
+        for package_spec in packages:
             await install_package_pyodide(
-                pkg,
+                package_spec,
                 verbose,
-                reinstall=should_reinstall_package(pkg, profile_changed),
+                reinstall=should_reinstall_package(package_spec, profile_changed),
+                use_cache=use_cache,
             )
+
         if verbose:
             log("Packages installed successfully.", force_verbose=verbose)
+
+        # 5. Persist the new requirements hash both in volatile memory (for this session)
+        #    and in persistent storage (for subsequent kernel restarts).
         os.environ["requirements_hash"] = requirements_hash
+
+        if use_cache:
+            manifest = await wheel_cache.read_manifest()
+            manifest["_requirements_hash"] = requirements_hash
+            await wheel_cache.write_manifest(manifest)
     else:
+        # If the package profile matches perfectly, we skip micropip entirely,
+        # ensuring instant start times for the user.
         if verbose:
             log("Packages are already installed.", force_verbose=verbose)