diff --git a/src/py/mat3ra/notebooks_utils/packages.py b/src/py/mat3ra/notebooks_utils/packages.py index 12396cd9..6b64b634 100644 --- a/src/py/mat3ra/notebooks_utils/packages.py +++ b/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) diff --git a/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py b/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py index 7e5414a5..3cb86559 100644 --- a/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py +++ b/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); }); diff --git a/src/py/mat3ra/notebooks_utils/pyodide/io.py b/src/py/mat3ra/notebooks_utils/pyodide/io.py index 9dbabfeb..aaf40f63 100644 --- a/src/py/mat3ra/notebooks_utils/pyodide/io.py +++ b/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) diff --git a/src/py/mat3ra/notebooks_utils/pyodide/packages/install.py b/src/py/mat3ra/notebooks_utils/pyodide/packages/install.py index 5a8e10ca..85c69185 100644 --- a/src/py/mat3ra/notebooks_utils/pyodide/packages/install.py +++ b/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,6 +158,12 @@ 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}.")] @@ -135,8 +171,8 @@ def clear_imported_package_modules(package_name: str): 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) diff --git a/src/py/mat3ra/notebooks_utils/pyodide/packages/wheel_cache.py b/src/py/mat3ra/notebooks_utils/pyodide/packages/wheel_cache.py new file mode 100644 index 00000000..c8aa64c2 --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/pyodide/packages/wheel_cache.py @@ -0,0 +1,151 @@ +"""IndexedDB-based wheel cache for Pyodide package installations. + +Caches downloaded wheel files in IndexedDB so they persist across kernel +restarts. On subsequent installs, wheels are loaded from cache and written +to emfs for local installation (no network download needed). +""" + +import json + +try: + from pyodide.code import run_js # type: ignore + from pyodide.ffi import to_js # type: ignore + + _HAS_PYODIDE = True +except ImportError: + run_js = None # type: ignore + to_js = None # type: ignore + _HAS_PYODIDE = False + + +_cache = ( + run_js( + """ +(function() { + const DB_NAME = "mat3ra"; + const DB_VERSION = 2; + const MANIFEST_STORE = "wheel_manifest"; + const WHEEL_STORE = "wheel_data"; + + const open = () => new Promise((resolve, reject) => { + const req = indexedDB.open(DB_NAME, DB_VERSION); + req.onupgradeneeded = (event) => { + const db = event.target.result; + if (!db.objectStoreNames.contains("tokens")) { + db.createObjectStore("tokens"); + } + if (!db.objectStoreNames.contains(MANIFEST_STORE)) { + db.createObjectStore(MANIFEST_STORE); + } + if (!db.objectStoreNames.contains(WHEEL_STORE)) { + db.createObjectStore(WHEEL_STORE); + } + }; + req.onsuccess = () => resolve(req.result); + req.onerror = () => reject(req.error); + }); + + return { + getManifest: async () => { + const db = await open(); + return new Promise(resolve => { + const tx = db.transaction(MANIFEST_STORE); + const req = tx.objectStore(MANIFEST_STORE).get("manifest"); + req.onsuccess = () => { db.close(); resolve(req.result || null); }; + req.onerror = () => { db.close(); resolve(null); }; + }); + }, + + setManifest: async (data) => { + const db = await open(); + const tx = db.transaction(MANIFEST_STORE, "readwrite"); + tx.objectStore(MANIFEST_STORE).put(data, "manifest"); + return new Promise(resolve => { + tx.oncomplete = () => { db.close(); resolve(); }; + }); + }, + + getWheel: async (key) => { + const db = await open(); + return new Promise(resolve => { + const tx = db.transaction(WHEEL_STORE); + const req = tx.objectStore(WHEEL_STORE).get(key); + req.onsuccess = () => { db.close(); resolve(req.result || null); }; + req.onerror = () => { db.close(); resolve(null); }; + }); + }, + + putWheel: async (key, data) => { + const db = await open(); + const tx = db.transaction(WHEEL_STORE, "readwrite"); + tx.objectStore(WHEEL_STORE).put(data, key); + return new Promise(resolve => { + tx.oncomplete = () => { db.close(); resolve(); }; + }); + }, + + clear: async () => { + const db = await open(); + const tx = db.transaction([MANIFEST_STORE, WHEEL_STORE], "readwrite"); + tx.objectStore(MANIFEST_STORE).clear(); + tx.objectStore(WHEEL_STORE).clear(); + return new Promise(resolve => { + tx.oncomplete = () => { db.close(); resolve(); }; + }); + } + }; +})() +""" + ) + if _HAS_PYODIDE + else None +) + + +async def read_manifest() -> dict: + """ + Read the cache manifest from IndexedDB. + + The manifest maps cache keys (wheel URLs) to metadata (filename, size, cached_at). + It also stores the '_requirements_hash' of the currently installed packages profile, + which survives kernel restarts and lets us skip running micropip if unchanged. + """ + result = await _cache.getManifest() + return json.loads(str(result)) if result else {} + + +async def write_manifest(manifest: dict) -> None: + """ + Write the cache manifest back to IndexedDB as a JSON string. + """ + await _cache.setManifest(json.dumps(manifest)) + + +async def get_cached_wheel(cache_key: str): + """ + Retrieve cached wheel bytes from IndexedDB, or None if not cached. + + Returns standard Python bytes by converting the JS Uint8Array proxy object + using .to_bytes() for high-performance memory translation. + """ + result = await _cache.getWheel(cache_key) + if result is None: + return None + return result.to_bytes() + + +async def put_cached_wheel(cache_key: str, wheel_bytes: bytes) -> None: + """ + Store wheel bytes in IndexedDB. + + Converts Python bytes to a JS Uint8Array using to_js() to ensure + zero-copy buffer transfer and avoid slow byte-by-byte duplication. + """ + await _cache.putWheel(cache_key, to_js(wheel_bytes)) + + +async def clear_cache() -> None: + """ + Wipe all cached wheels and the manifest from IndexedDB. + """ + await _cache.clear()