From c60f8dacf191c5255ef8fca462d6d3d571b86b07 Mon Sep 17 00:00:00 2001
From: pyranota <92104930+pyranota@users.noreply.github.com>
Date: Wed, 29 Jan 2025 21:13:47 +0300
Subject: [PATCH] feat(python): add `custom_wheels` directory to `PYTHONPATH`
 (#5169)

* fix(python): fix uv can't find ssl certificates

- Add `PY_NATIVE_CERT` flag, forces UV to use native tls
- Rename `PIP_INDEX_CERT` to `PY_INDEX_CERT`
- Rename `PIP_TRUSTED_HOST` to `PY_TRUSTED_HOST`

For backwards compatibility PIP* variables are still accessible

* feat(python): add `custom_wheels` directory to PYTHONPATH

Add global directory by path `<CACHE_DIR>/python_xyz/custom_wheels`

For example for scripts running python 3.11, in every execution
`<CACHE_DIR>/python_311/custom_wheels` will be accessible and all wheels placed there could be imported and used.

This is usefull for preinstalling wheels before runtime

* Make it work with Nsjail

* Rework and make custom_wheels optional

* Remove `create_dir_all` from imports

* Use sync version of metadata

* Rename `custom_wheels` to `global-site-packages`
---
 .../nsjail/run.python3.config.proto           |  7 ++++++
 .../windmill-worker/src/python_executor.rs    | 25 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/backend/windmill-worker/nsjail/run.python3.config.proto b/backend/windmill-worker/nsjail/run.python3.config.proto
index a9d61dc8b1..b49b9cfbfe 100644
--- a/backend/windmill-worker/nsjail/run.python3.config.proto
+++ b/backend/windmill-worker/nsjail/run.python3.config.proto
@@ -123,6 +123,13 @@ mount {
     is_bind: true
 }
 
+mount {
+    src: "{GLOBAL_SITE_PACKAGES}"
+    dst: "{GLOBAL_SITE_PACKAGES}"
+    is_bind: true
+   	mandatory: false
+}
+
 {SHARED_MOUNT}
 
 {SHARED_DEPENDENCIES}
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index 630cb0ced1..8c134166a0 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -902,6 +902,8 @@ pub async fn handle_python_job(
         tracing::debug!("Finished deps postinstall stage");
     }
 
+
+
     if no_uv {
         append_logs(
             &job.id,
@@ -1050,7 +1052,27 @@ except BaseException as e:
 
     let client = client.get_authed().await;
     let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?;
-    let additional_python_paths_folders = additional_python_paths.iter().join(":");
+
+    // Add /tmp/windmill/cache/python_xyz/global-site-packages to PYTHONPATH.
+    // Usefull if certain wheels needs to be preinstalled before execution.
+    let global_site_packages_path = py_version.to_cache_dir() + "/global-site-packages";
+    let additional_python_paths_folders = {
+        let mut paths= additional_python_paths.clone();
+        if std::fs::metadata(&global_site_packages_path).is_ok() {
+            // We want global_site_packages_path to be included in additonal_python_paths_folders, but
+            // we don't want it to be included in global_site_packages_path.
+            // The reason for this is that additional_python_paths_folders is used to fill PYTHONPATH env variable for jailed script
+            // When global_site_packages_path used to place mount point of wheels to the jail config.
+            // Since we handle mount of global_site_packages on our own, we don't want it to be mounted automatically.
+            // We do this because existence of every wheel in cache is mandatory and if it is not there and nsjail expects it, it is a bug.
+            // On the other side global_site_packages is purely optional.
+            // NOTE: This behaviour can be changed in future, so verification of wheels can be offloaded from nsjail to windmill
+            paths.insert(0, global_site_packages_path.clone());
+            //    ^^^^^^^^
+            // We also want this be priorotized, that's why we insert it to the beginning
+        }
+        paths.iter().join(":")
+    };
 
     #[cfg(windows)]
     let additional_python_paths_folders = additional_python_paths_folders.replace(":", ";");
@@ -1080,6 +1102,7 @@ mount {{
                 .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string())
                 .replace("{SHARED_MOUNT}", shared_mount)
                 .replace("{SHARED_DEPENDENCIES}", shared_deps.as_str())
+                .replace("{GLOBAL_SITE_PACKAGES}", &global_site_packages_path)
                 .replace("{MAIN}", format!("{dirs}/{last}").as_str())
                 .replace(
                     "{ADDITIONAL_PYTHON_PATHS}",