Add offline installation of plugin dependencies

.github/workflows/_build.yml

@@ -428,9 +428,19 @@ jobs:
           $env:USERPROFILE = Join-Path $env:RUNNER_TEMP "datalab-profile"
           New-Item -ItemType Directory -Force $env:USERPROFILE | Out-Null
           $env:DATALAB_PLUGIN_PROBE_MARKER = Join-Path $env:RUNNER_TEMP "plugin-probe.json"
-          python scripts/frozen_plugin_probe.py install
+          if (Test-Path -LiteralPath $env:DATALAB_PLUGIN_PROBE_MARKER) { Remove-Item -LiteralPath $env:DATALAB_PLUGIN_PROBE_MARKER }
+          $env:DATALAB_PLUGIN_PROBE_WHEELHOUSE = Join-Path $env:RUNNER_TEMP "plugin-probe-catalog\dependencies"
+          python scripts/frozen_plugin_probe.py prepare
+          if ($LASTEXITCODE -ne 0) { throw "Cannot prepare the dependency wheelhouse" }
+          $probe = (Resolve-Path scripts/frozen_plugin_probe.py).Path
+          $log = Join-Path $env:RUNNER_TEMP "plugin-install.log"
+          $code = "import runpy,sys; sys.stdout = sys.stderr = open(r'$log', 'w', encoding='utf-8'); sys.argv = [r'$probe', 'install']; runpy.run_path(r'$probe', run_name='__main__')"
+          $installer = Start-Process dist\DataLab\DataLab.exe -ArgumentList "-c", ('"' + $code + '"') -PassThru
+          if (-not $installer.WaitForExit(300000)) { Stop-Process -Id $installer.Id; throw "Plugin installation did not finish" }
+          if ($installer.ExitCode -ne 0) { Get-Content $log; throw "Offline plugin installation failed" }
           $process = Start-Process dist/DataLab/DataLab.exe -ArgumentList "--unattended", "--delay", "5000" -PassThru
           if (-not $process.WaitForExit(300000)) { $process.Kill(); throw "DataLab did not quit" }
+          if ($process.ExitCode -ne 0) { throw "Frozen DataLab exited with code $($process.ExitCode)" }
           python scripts/frozen_plugin_probe.py verify
 
       - name: Package executable as zip (internal, required by MSI workflow)

DataLab.spec

@@ -7,7 +7,9 @@ from PyInstaller.utils.hooks import collect_submodules, collect_data_files, copy
 all_hidden_imports = collect_submodules('datalab')
 # Plugins installed from a file may import any Sigima module
 all_hidden_imports += collect_submodules('sigima', filter=lambda name: '.tests' not in name)
+all_hidden_imports += collect_submodules('pip')
 datas = collect_data_files('datalab') + [('datalab\\plugins\\builtin', 'datalab\\plugins\\builtin')]
+datas += collect_data_files('pip') + copy_metadata('pip')
 # Macro templates are read with ``pkgutil.get_data`` and must exist on disk,
 # not only inside the PYZ archive (``collect_data_files`` skips ``.py`` files):
 datas += [('datalab\\gui\\macros_templates', 'datalab\\gui\\macros_templates')]
@@ -37,6 +39,8 @@ a = Analysis(
     runtime_hooks=[],
     excludes=[],
     noarchive=False,
+    # distlib's resource finder requires a standard file-system loader.
+    module_collection_mode={'pip': 'py'},
 )
 pyz = PYZ(a.pure)
 

datalab/__init__.py

@@ -24,7 +24,7 @@
     # this module is imported more than once, e.g. when running tests)
     pass
 
-__version__ = "1.3.0"
+__version__ = "1.4.0"
 __docurl__ = __homeurl__ = "https://datalab-platform.com/"
 __supporturl__ = "https://github.com/DataLab-Platform/DataLab/issues/new/choose"
 

datalab/gui/plugins/available.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-import shutil
+import os.path as osp
 import sys
 import tempfile
 from collections.abc import Callable, Sequence
@@ -32,42 +32,21 @@
     get_plugin_catalog_paths,
     normalize_plugin_paths,
 )
-from datalab.gui.plugins.install import InstalledPluginsWidget
+from datalab.gui.plugins.install import InstalledPluginsWidget, run_plugin_task
 from datalab.plugins.catalog import (
     CatalogPlugin,
     CatalogRelease,
     CatalogsContent,
+    download_dependency_wheels,
     download_release,
     load_catalogs,
 )
+from datalab.plugins.dependencies import DEPENDENCIES_DIRECTORY
 from datalab.plugins.store import WHEEL_KIND
 from datalab.widgets.expandabletext import apply_palette_color, apply_subdued_color
 
 CATALOG_PAGE_URL = "https://datalab-platform.com/plugins/"
 
-#: Running tasks, kept alive until their thread finishes
-_RUNNING_TASKS: set[_Task] = set()
-
-
-class _Task(QC.QThread):
-    """Run a function in a background thread."""
-
-    succeeded = QC.Signal(object)
-    failed = QC.Signal(str)
-
-    def __init__(self, function: Callable[[], object]) -> None:
-        super().__init__()
-        self._function = function
-
-    def run(self) -> None:  # pragma: no cover - thread entry point
-        """Call the function and emit its result or error message."""
-        try:
-            result = self._function()
-        except (OSError, ValueError) as exc:
-            self.failed.emit(str(exc))
-            return
-        self.succeeded.emit(result)
-
 
 def _is_newer(version: str, installed: str) -> bool:
     try:
@@ -199,8 +178,11 @@ def __init__(
         self._installer = installer
         self.plugins: list[CatalogPlugin] | None = None
         self.item_widgets: list[AvailablePluginItemWidget] = []
-        self._download: tuple[CatalogPlugin, str] | None = None
+        self._download: tuple[CatalogPlugin, tempfile.TemporaryDirectory] | None = None
         self._busy = False
+        self._installer.busy_changed.connect(
+            lambda busy: self._set_busy(_("Installing plugin...") if busy else None)
+        )
         self._directories = merge_catalog_directories(get_plugin_catalog_paths())
         #: Catalog URL and local folders of the last loading
         self._loaded_sources: tuple[str, tuple[str, ...]] | None = None
@@ -286,12 +268,7 @@ def run_task(
         on_failure: Callable[[str], None],
     ) -> None:
         """Run a function in a background thread and report its outcome."""
-        task = _Task(function)
-        _RUNNING_TASKS.add(task)
-        task.succeeded.connect(on_success)
-        task.failed.connect(on_failure)
-        task.finished.connect(lambda: _RUNNING_TASKS.discard(task))
-        task.start()
+        run_plugin_task(function, on_success, on_failure)
 
     def refresh(self) -> None:
         """Load the catalogs again."""
@@ -315,15 +292,24 @@ def install(self, plugin: CatalogPlugin) -> None:
         release = plugin.get_installable_release()
         if release is None or self._busy:
             return
-        directory = tempfile.mkdtemp(prefix="datalab-plugin-")
+        directory = tempfile.TemporaryDirectory(prefix="datalab-plugin-")
         self._download = (plugin, directory)
         self._set_busy(_("Downloading %s...") % release.filename)
         self.run_task(
-            lambda: download_release(release, directory),
+            lambda: self._download_release(release, directory.name),
             self._on_downloaded,
             self._on_download_failed,
         )
 
+    @staticmethod
+    def _download_release(release: CatalogRelease, directory: str) -> str:
+        path = download_release(release, directory)
+        if release.dependency_wheels:
+            download_dependency_wheels(
+                release, str(osp.join(directory, DEPENDENCIES_DIRECTORY))
+            )
+        return path
+
     def _set_busy(self, message: str | None) -> None:
         self._busy = message is not None
         self.refresh_button.setEnabled(not self._busy)
@@ -364,20 +350,26 @@ def _on_downloaded(self, path: str) -> None:
         self._download = None
         self._set_busy(None)
         self.status_label.setText("")
-        try:
-            if plugin.source:
-                origin = _("Local plugin catalog (%s)") % plugin.source
-            else:
-                origin = _("DataLab plugin catalog (%s)") % _tier_label(plugin)
-            self._installer.install_path(path, origin=origin)
-        finally:
-            shutil.rmtree(directory, ignore_errors=True)
-        self._populate()
+        if plugin.source:
+            origin = _("Local plugin catalog (%s)") % plugin.source
+        else:
+            origin = _("DataLab plugin catalog (%s)") % _tier_label(plugin)
+        wheelhouse = osp.join(directory.name, DEPENDENCIES_DIRECTORY)
+        if plugin.source and not osp.isdir(wheelhouse):
+            wheelhouse = osp.join(plugin.source, DEPENDENCIES_DIRECTORY)
+
+        def finished() -> None:
+            directory.cleanup()
+            self._populate()
+
+        self._installer.install_path(
+            path, origin=origin, wheelhouse=wheelhouse, on_finished=finished
+        )
 
     def _on_download_failed(self, message: str) -> None:
         _plugin, directory = self._download
         self._download = None
-        shutil.rmtree(directory, ignore_errors=True)
+        directory.cleanup()
         self._set_busy(None)
         self.status_label.setText("")
         QW.QMessageBox.warning(self, _("Install plugin"), message)