perf(e2e): adopt a pre-baked toolkit venv instead of provisioning per suite

.github/workflows/e2e.yml

@@ -18,9 +18,17 @@ concurrency:
 
 jobs:
   e2e:
-    name: E2E (ExTester)
+    name: E2E (${{ matrix.group }})
     runs-on: ubuntu-latest
     timeout-minutes: 45
+    strategy:
+      # One shard per directory in test/e2e/suite/, grouped by what the suites cover rather than by
+      # measured time, so a new suite has an obvious home and the split does not skew as suites are
+      # added. Adding a directory means adding it here; verify-coverage below fails the run if you
+      # forget, since nothing would otherwise run it.
+      fail-fast: false
+      matrix:
+        group: [kernel, files, workspace]
     env:
       # Keep ExTester's downloads (test VS Code, ChromeDriver, settings, screenshots) inside the
       # workspace so the artifact-upload paths are predictable. Both this and .test-extensions are
@@ -39,6 +47,7 @@ jobs:
           node-version-file: '.nvmrc'
 
       - name: Setup Python # interpreter the Deepnote environment is created from
+        id: python
         uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version: '3.12'
@@ -72,29 +81,110 @@ jobs:
         run: npm run setup:e2e:deps
 
       - name: Cache pip wheel downloads
-        # Provisioning the Deepnote environment pip-installs the toolkit dependency tree into a
-        # fresh venv on first kernel connect — the bulk of the E2E runtime. Caching pip's wheel
-        # cache makes that warm on later runs (the installs use the cache; nothing passes
-        # --no-cache-dir). The key busts when the toolkit version / install set changes; the
-        # restore-keys prefix keeps the cache warm across unrelated changes, since pip's cache is
-        # additive.
+        # Only the deletion suite still pip-installs at run time (it creates a managed environment
+        # on purpose), so this is now a fallback rather than the main lever — see the baked venv
+        # below. The key carries github.run_id because actions/cache only writes on a miss: with a
+        # bare content hash the entry froze at its first save and the unpinned specs
+        # (ipykernel, python-lsp-server, deepnote-cli, all installed with --upgrade) were
+        # re-downloaded every run and never written back. restore-keys still resolves to the newest
+        # matching entry, so each run starts warm and saves a refreshed copy.
         uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
           path: ~/.cache/pip
-          key: pip-${{ runner.os }}-py312-${{ hashFiles('src/kernels/deepnote/types.ts', 'src/kernels/deepnote/deepnoteToolkitInstaller.node.ts') }}
+          key: pip-${{ runner.os }}-py${{ steps.python.outputs.python-version }}-${{ hashFiles('src/kernels/deepnote/types.ts', 'src/kernels/deepnote/deepnoteToolkitInstaller.node.ts') }}-${{ github.run_id }}
           restore-keys: |
-            pip-${{ runner.os }}-py312-
+            pip-${{ runner.os }}-py${{ steps.python.outputs.python-version }}-${{ hashFiles('src/kernels/deepnote/types.ts', 'src/kernels/deepnote/deepnoteToolkitInstaller.node.ts') }}-
+            pip-${{ runner.os }}-py${{ steps.python.outputs.python-version }}-
+
+      - name: Cache the baked Deepnote toolkit venv
+        # Every suite that connects a kernel adopts this venv instead of letting the extension build
+        # one, which is where most of the old E2E runtime went. Keyed on the resolved Python version
+        # as well as the install set: the venv records its base interpreter in pyvenv.cfg, so a
+        # patch bump to the hosted Python would leave a restored venv pointing at a path that no
+        # longer exists. prepareE2eVenv.js also self-heals — it discards and rebuilds any venv that
+        # cannot import deepnote_toolkit.
+        uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: .venv-e2e
+          key: e2e-venv-${{ runner.os }}-py${{ steps.python.outputs.python-version }}-${{ hashFiles('src/kernels/deepnote/types.ts', 'src/kernels/deepnote/deepnoteToolkitInstaller.node.ts') }}
+
+      - name: Bake the Deepnote toolkit venv
+        # No-op when the cache restored a usable venv; otherwise creates it and installs the same
+        # set deepnoteToolkitInstaller would have.
+        run: npm run setup:e2e:venv
 
       - name: Run E2E
         # VS Code launches with --no-sandbox (no AppArmor sysctl needed). Runs once; Mocha's retries:1
         # and rootHooks.ts (dismiss toasts between tests) handle flakiness in the single shared instance.
-        run: xvfb-run --auto-servernum --server-args='-screen 0 1920x1080x24' npm run test:e2e:prebuilt
+        run: |
+          xvfb-run --auto-servernum --server-args='-screen 0 1920x1080x24' \
+            npx extest run-tests "./out/e2e/suite/${{ matrix.group }}/*.e2e.test.js" \
+              -c max -o ./test/e2e/settings.generated.json -e .test-extensions -m ./test/e2e/.mocharc.js
+
+      # Written even when the tests fail, so a shard that ran is never mistaken for one that was
+      # never scheduled. verify-coverage collects these.
+      - name: Record the shard that ran
+        if: always()
+        run: mkdir -p shard-ran && touch "shard-ran/${{ matrix.group }}"
+
+      - name: Upload the shard record
+        if: always()
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: e2e-ran-${{ matrix.group }}
+          path: shard-ran/
 
       - name: Upload failure screenshots
         if: failure()
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
         with:
-          name: e2e-screenshots
+          name: e2e-screenshots-${{ matrix.group }}
           path: ${{ env.TEST_RESOURCES }}/screenshots/**/*.png
           if-no-files-found: ignore
           retention-days: 14
+
+  verify-coverage:
+    name: Verify shard coverage
+    needs: e2e
+    # Runs even when a shard failed: a directory nobody runs is a separate problem from a failing
+    # test, and it is the one that hides — nothing runs it and the build goes green sooner. Not on
+    # cancellation though, or superseding a run (cancel-in-progress) would report every directory as
+    # unrun.
+    if: ${{ !cancelled() }}
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+
+      - name: Collect what each shard reported
+        uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7.0.0
+        with:
+          merge-multiple: true
+          path: shard-ran
+          pattern: e2e-ran-*
+
+      - name: Every suite directory must have been run
+        shell: bash
+        run: |
+          set -uo pipefail
+          echo "Shards that reported in: $(ls shard-ran 2>/dev/null | tr '\n' ' ')"
+
+          status=0
+          for dir in test/e2e/suite/*/; do
+            group="$(basename "${dir}")"
+            if [ ! -e "shard-ran/${group}" ]; then
+              echo "::error::test/e2e/suite/${group}/ was not run by any shard — add it to the matrix in this workflow."
+              status=1
+            fi
+          done
+
+          # A suite outside a group directory is never globbed by any shard either.
+          if compgen -G 'test/e2e/suite/*.e2e.test.ts' > /dev/null; then
+            echo "::error::Suites sit directly in test/e2e/suite/; move them into a group directory:"
+            ls test/e2e/suite/*.e2e.test.ts
+            status=1
+          fi
+
+          exit "${status}"

.gitignore

@@ -79,3 +79,6 @@ tsconfig.tsbuildinfo
 # ExTester (vscode-extension-tester) E2E artifacts
 test-resources
 .test-extensions
+
+# Generated by build/e2e/prepareE2eVenv.js (absolute interpreter path resolved at run time).
+test/e2e/settings.generated.json

.vscodeignore

@@ -85,6 +85,7 @@ obj/**
 logs/**
 out/**
 .test-extensions/**
+.venv-e2e/**
 test-resources/**
 precommit.hook
 pythonFiles/.env

build/e2e/prepareE2eVenv.js

@@ -0,0 +1,98 @@
+// Bakes the Python venv the E2E suite adopts instead of letting the extension build one per test
+// run. `deepnoteEnvironmentManager.createEnvironment` adopts any interpreter that already lives in a
+// venv (`getVenvPathIfInVenv`), and `ensureVenvAndToolkit` returns early once `import
+// deepnote_toolkit` succeeds — so a pre-installed venv skips venv creation and the whole pip install.
+//
+// Also emits the settings file ExTester feeds VS Code, because the interpreter must be named by
+// absolute path and that path is only known at run time.
+
+const { execFileSync } = require('child_process');
+const fs = require('fs');
+const path = require('path');
+
+const repoRoot = path.resolve(__dirname, '..', '..');
+// Matches the `**/.venv*/` rule already in .gitignore.
+const venvDir = path.join(repoRoot, '.venv-e2e');
+const baseSettingsPath = path.join(repoRoot, 'test', 'e2e', 'settings.json');
+const generatedSettingsPath = path.join(repoRoot, 'test', 'e2e', 'settings.generated.json');
+
+/** Reads DEEPNOTE_TOOLKIT_VERSION from source so this script cannot drift from the extension. */
+function toolkitVersion() {
+    const types = fs.readFileSync(path.join(repoRoot, 'src', 'kernels', 'deepnote', 'types.ts'), 'utf8');
+    const match = types.match(/DEEPNOTE_TOOLKIT_VERSION\s*=\s*'([^']+)'/);
+    if (!match) {
+        throw new Error('Could not read DEEPNOTE_TOOLKIT_VERSION from src/kernels/deepnote/types.ts');
+    }
+
+    return match[1];
+}
+
+function venvPython() {
+    return process.platform === 'win32'
+        ? path.join(venvDir, 'Scripts', 'python.exe')
+        : path.join(venvDir, 'bin', 'python');
+}
+
+function run(command, args) {
+    execFileSync(command, args, { stdio: 'inherit' });
+}
+
+/** True when the venv exists and already imports the toolkit — the same check the extension makes. */
+function toolkitAlreadyInstalled() {
+    if (!fs.existsSync(venvPython())) {
+        return false;
+    }
+    try {
+        execFileSync(venvPython(), ['-c', 'import deepnote_toolkit'], { stdio: 'ignore' });
+
+        return true;
+    } catch {
+        return false;
+    }
+}
+
+function bakeVenv() {
+    if (toolkitAlreadyInstalled()) {
+        console.log(`[e2e-venv] reusing ${venvDir}`);
+
+        return;
+    }
+
+    // A half-built venv (interrupted run, partial cache) would make pip fail confusingly.
+    if (fs.existsSync(venvDir)) {
+        console.log(`[e2e-venv] discarding incomplete venv at ${venvDir}`);
+        fs.rmSync(venvDir, { recursive: true, force: true });
+    }
+
+    console.log(`[e2e-venv] creating ${venvDir}`);
+    run(process.env.PYTHON ?? 'python3', ['-m', 'venv', venvDir]);
+
+    // Mirrors deepnoteToolkitInstaller.installVenvAndToolkit so the baked venv satisfies the same
+    // checks the extension would have made after building it itself.
+    run(venvPython(), ['-m', 'pip', 'install', '--upgrade', 'pip']);
+    run(venvPython(), [
+        '-m',
+        'pip',
+        'install',
+        '--upgrade',
+        `deepnote-toolkit[server]==${toolkitVersion()}`,
+        'ipykernel',
+        'python-lsp-server[all]',
+        'deepnote-cli'
+    ]);
+}
+
+function writeSettings() {
+    const settings = JSON.parse(fs.readFileSync(baseSettingsPath, 'utf8'));
+
+    // venvPath makes the Python extension discover the baked venv so it reaches the interpreter
+    // quick pick; defaultInterpreterPath makes it the pick the suite lands on by default.
+    settings['python.venvPath'] = repoRoot;
+    settings['python.defaultInterpreterPath'] = venvPython();
+
+    fs.writeFileSync(generatedSettingsPath, `${JSON.stringify(settings, null, 4)}\n`);
+    console.log(`[e2e-venv] wrote ${path.relative(repoRoot, generatedSettingsPath)}`);
+}
+
+bakeVenv();
+writeSettings();

package.json

@@ -2672,9 +2672,10 @@
         "compile-e2e-watch": "tsc -p ./test/e2e/tsconfig.json --watch",
         "setup:e2e:vscode": "extest get-vscode -c max && extest get-chromedriver -c max",
         "setup:e2e:deps": "extest install-from-marketplace ms-python.python -e .test-extensions",
-        "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps",
-        "test:e2e": "extest setup-and-run \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js -i",
-        "test:e2e:prebuilt": "extest run-tests \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js",
+        "setup:e2e:venv": "node ./build/e2e/prepareE2eVenv.js",
+        "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps && npm run setup:e2e:venv",
+        "test:e2e": "extest setup-and-run \"./out/e2e/suite/**/*.e2e.test.js\" -c max -o ./test/e2e/settings.generated.json -e .test-extensions -m ./test/e2e/.mocharc.js -i",
+        "test:e2e:prebuilt": "extest run-tests \"./out/e2e/suite/**/*.e2e.test.js\" -c max -o ./test/e2e/settings.generated.json -e .test-extensions -m ./test/e2e/.mocharc.js",
         "test:unittests": "mocha --config ./build/.mocha.unittests.js.json ./out/**/*.unit.test.js",
         "test": "npm run test:unittests",
         "typecheck": "tsc -p ./ --noEmit",

test/e2e/helpers/constants.ts

@@ -47,3 +47,19 @@ export const FOLDER_OK_RETRY_DELAY = 400;
 // Selectors that only exist inside the notebook output iframe (`#active-frame`),
 // so reading them cannot accidentally match the cell's source in the editor.
 export const OUTPUT_SELECTOR = '.output_container, .output, .rendered-output';
+
+// The environment every suite shares. Creating a Deepnote environment is cheap (metadata only), but
+// the first kernel connect provisions a venv + the Deepnote toolkit, so a suite that invents its own
+// name pays that once more. Import this instead of writing a name; `E2E Delete Env` in
+// environment.e2e.test.ts is the one deliberate exception, since that suite deletes what it creates.
+export const SHARED_ENV_NAME = 'E2E Hello Env';
+
+// Venv baked by build/e2e/prepareE2eVenv.js, adopted by every suite that connects a kernel. Kept in
+// sync with that script by convention; the suite falls back to whatever the quick pick offers (and
+// warns) when it is absent, so a missed setup step is slow rather than red.
+export const PREBAKED_VENV_DIR_NAME = '.venv-e2e';
+
+// How long to wait for the interpreter quick pick to actually narrow to the baked venv. The
+// filter is near-instant when the venv is discoverable, so this only elapses when it is missing —
+// keep it well under QUICK_PICK_TIMEOUT so a forgotten setup step does not stall every suite.
+export const PREBAKED_VENV_FILTER_TIMEOUT = 10_000;