make virtual env creation much faster

We first create a base venv that is not to be deleted and can be re-used easily.
There we install all default packages.

Then, we create another, kernel-specific, venv without any additional packages.
Instead, we create a link that points to the packages installed in the base image.

The base image is intended to stay immutable. All newly installed packages end up
in the kernel-specific venv. While they are lost on restart, the packages inside the
base image are not.

Author: dasmy

gpt_code_ui/kernel_program/kernel_manager.py

@@ -10,6 +10,7 @@
 import threading
 import time
 import traceback
+import venv
 
 from time import sleep
 from jupyter_client import BlockingKernelClient
@@ -20,8 +21,6 @@
 import gpt_code_ui.kernel_program.utils as utils
 import gpt_code_ui.kernel_program.config as config
 
-USE_SEPARATE_VENV = False
-
 # Set up globals
 messaging = None
 logger = config.get_logger()
@@ -170,9 +169,58 @@ def flush_kernel_msgs(kc, tries=1, timeout=0.2):
         logger.debug(f"{e} [{type(e)}")
 
 
+def create_venv(venv_dir: pathlib.Path, install_default_packages: bool) -> pathlib.Path:
+    venv_bindir = venv_dir / 'bin'
+    venv_python_executable = venv_bindir / os.path.basename(sys.executable)
+
+    if not os.path.isdir(venv_dir):
+        # create virtual env inside venv_dir directory
+        venv.create(venv_dir, system_site_packages=True, with_pip=True, upgrade_deps=True)
+
+        if install_default_packages:
+            # install wheel because some packages do not like being installed without
+            subprocess.run([venv_python_executable, '-m', 'pip', 'install', 'wheel>=0.41,<1.0'])
+            # install all default packages into the venv
+            default_packages = [
+                "ipykernel>=6,<7",
+                "numpy>=1.24,<1.25",
+                "dateparser>=1.1,<1.2",
+                "pandas>=1.5,<1.6",
+                "geopandas>=0.13,<0.14",
+                "tabulate>=0.9.0<1.0",
+                "PyPDF2>=3.0,<3.1",
+                "pdfminer>=20191125,<20191200",
+                "pdfplumber>=0.9,<0.10",
+                "matplotlib>=3.7,<3.8",
+                "openpyxl>=3.1.2,<4",
+            ]
+            subprocess.run([venv_python_executable, '-m', 'pip', 'install'] + default_packages)
+
+    # get base env library path as we need this to refer to this form a derived venv
+    site_packages = subprocess.check_output([venv_python_executable, '-c', 'import sysconfig; print(sysconfig.get_paths()["purelib"])'])
+    site_packages = site_packages.decode('utf-8').split('\n')[0]
+
+    return pathlib.Path(site_packages)
+
+
+def create_derived_venv(base_venv: pathlib.Path, venv_dir: pathlib.Path):
+    site_packages_base = create_venv(base_venv, install_default_packages=True)
+    site_packages_derived = create_venv(venv_dir, install_default_packages=False)
+
+    # create a link from derived venv into the base venv, see https://stackoverflow.com/a/75545634
+    with open(site_packages_derived / '_base_packages.pth', 'w') as pth:
+        pth.write(f'{site_packages_base}\n')
+
+    venv_bindir = venv_dir / 'bin'
+    venv_python_executable = venv_bindir / os.path.basename(sys.executable)
+
+    return venv_bindir, venv_python_executable
+
+
 def start_kernel(id: str):
     cwd = pathlib.Path(os.getcwd())
     kernel_dir = cwd / f'kernel.{id}'
+    base_dir = cwd / 'kernel.base'
 
     # Cleanup potential leftovers
     shutil.rmtree(kernel_dir, ignore_errors=True)
@@ -182,33 +230,9 @@ def start_kernel(id: str):
     kernel_connection_file = kernel_dir / "kernel_connection_file.json"
     launch_kernel_script_path = pathlib.Path(__file__).parent.resolve() / "launch_kernel.py"
 
-    if not USE_SEPARATE_VENV:
-        kernel_python_executable = sys.executable  # TODO: here we could also pick up an already existing env by simply selecting the corresponding Python binary
-    else:
-        kernel_venv = kernel_dir / 'venv'
-        kernel_venv_bindir = kernel_venv / 'bin'
-        kernel_python_executable = kernel_venv_bindir / os.path.basename(sys.executable)
-        kernel_env['PATH'] = str(kernel_venv_bindir) + os.pathsep + kernel_env['PATH']
-
-        # create virtual env inside kernel_venv directory
-        subprocess.run([sys.executable, '-m', 'venv', kernel_venv, '--upgrade-deps', '--system-site-packages'])
-        # install wheel because some packages do not like being installed without
-        subprocess.run([kernel_python_executable, '-m', 'pip', 'install', 'wheel>=0.41,<1.0'])
-        # install all default packages into the venv
-        default_packages = [
-            "ipykernel>=6,<7",
-            "numpy>=1.24,<1.25",
-            "dateparser>=1.1,<1.2",
-            "pandas>=1.5,<1.6",
-            "geopandas>=0.13,<0.14",
-            "tabulate>=0.9.0<1.0",
-            "PyPDF2>=3.0,<3.1",
-            "pdfminer>=20191125,<20191200",
-            "pdfplumber>=0.9,<0.10",
-            "matplotlib>=3.7,<3.8",
-            "openpyxl>=3.1.2,<4",
-        ]
-        subprocess.run([kernel_python_executable, '-m', 'pip', 'install'] + default_packages)
+    kernel_venv_dir = kernel_dir / 'venv'
+    kernel_venv_bindir, kernel_python_executable = create_derived_venv(base_dir, kernel_venv_dir)
+    kernel_env['PATH'] = str(kernel_venv_bindir) + os.pathsep + kernel_env['PATH']
 
     # start the kernel using the virtual env python executable
     kernel_process = subprocess.Popen(