Enable venv to inherit site packages from base python environment

frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java

@@ -221,7 +221,6 @@ private void setupModelVenv(Model model)
         commandParts.add("-m");
         commandParts.add("venv");
         commandParts.add("--clear");
-        commandParts.add("--system-site-packages");
         commandParts.add(venvPath.toString());
 
         ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
@@ -273,6 +272,57 @@ private void setupModelVenv(Model model)
             throw new ModelException(
                     "Virtual environment creation failed for model " + model.getModelName());
         }
+
+        // Inherit site-packages directories from the current environment torchserve is running in
+        // to the newly created virtual environment
+        commandParts.clear();
+        commandParts.add(configManager.getPythonExecutable());
+        commandParts.add(
+                Paths.get(
+                                configManager.getModelServerHome(),
+                                "ts",
+                                "utils",
+                                "inherit_site_packages.py")
+                        .toAbsolutePath()
+                        .toString());
+        commandParts.add(venvPath.toString());
+
+        processBuilder = new ProcessBuilder(commandParts);
+        processBuilder.directory(venvPath.getParentFile());
+        environment = processBuilder.environment();
+        envp =
+                EnvironmentUtils.getEnvString(
+                        configManager.getModelServerHome(),
+                        model.getModelDir().getAbsolutePath(),
+                        null);
+        for (String envVar : envp) {
+            String[] parts = envVar.split("=", 2);
+            if (parts.length == 2) {
+                environment.put(parts[0], parts[1]);
+            }
+        }
+        processBuilder.redirectErrorStream(true);
+        process = processBuilder.start();
+        exitCode = process.waitFor();
+        brdr = new BufferedReader(new InputStreamReader(process.getInputStream()));
+        outputString.setLength(0);
+        while ((line = brdr.readLine()) != null) {
+            outputString.append(line + "\n");
+        }
+
+        if (exitCode == 0) {
+            logger.debug(
+                    "Inherited site-packages directories to venv {}:\n{}",
+                    venvPath.toString(),
+                    outputString.toString());
+        } else {
+            logger.error(
+                    "Failed to inherit site-packages directories to venv {}:\n{}",
+                    venvPath.toString(),
+                    outputString.toString());
+            throw new ModelException(
+                    "Failed to inherit site-packages directories to venv " + venvPath.toString());
+        }
     }
 
     private void setupModelDependencies(Model model)
@@ -360,10 +410,7 @@ private void setupModelDependencies(Model model)
         }
 
         if (exitCode == 0) {
-            logger.info(
-                    "Installed custom pip packages for model {}:\n{}",
-                    model.getModelName(),
-                    outputString.toString());
+            logger.info("Installed custom pip packages for model {}", model.getModelName());
         } else {
             logger.error(
                     "Custom pip package installation failed for model {}:\n{}",

test/pytest/test_model_custom_dependencies.py

@@ -141,9 +141,7 @@ def register_model_and_make_inference_request(expect_model_load_failure=False):
 
 
 def test_install_dependencies_to_target_directory_with_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=True, use_venv=False)
@@ -162,14 +160,11 @@ def test_install_dependencies_to_target_directory_with_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=False)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_install_dependencies_to_target_directory_without_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=False, use_venv=False)
@@ -188,14 +183,11 @@ def test_install_dependencies_to_target_directory_without_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_disable_install_dependencies_to_target_directory_with_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=True, use_venv=False)
@@ -207,14 +199,11 @@ def test_disable_install_dependencies_to_target_directory_with_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_disable_install_dependencies_to_target_directory_without_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=False, use_venv=False)
@@ -226,14 +215,11 @@ def test_disable_install_dependencies_to_target_directory_without_requirements()
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_install_dependencies_to_venv_with_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=True, use_venv=True)
@@ -252,14 +238,11 @@ def test_install_dependencies_to_venv_with_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=False)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_install_dependencies_to_venv_without_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=False, use_venv=True)
@@ -278,14 +261,11 @@ def test_install_dependencies_to_venv_without_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_disable_install_dependencies_to_venv_with_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=True, use_venv=True)
@@ -297,14 +277,11 @@ def test_disable_install_dependencies_to_venv_with_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()
 
 
 def test_disable_install_dependencies_to_venv_without_requirements():
-    # Torchserve cleanup
-    test_utils.stop_torchserve()
-    test_utils.delete_all_snapshots()
+    test_utils.torchserve_cleanup()
 
     try:
         generate_model_archive(use_requirements=False, use_venv=True)
@@ -316,5 +293,4 @@ def test_disable_install_dependencies_to_venv_without_requirements():
         )
         register_model_and_make_inference_request(expect_model_load_failure=True)
     finally:
-        test_utils.stop_torchserve()
-        test_utils.delete_all_snapshots()
+        test_utils.torchserve_cleanup()

ts/utils/inherit_site_packages.py

@@ -0,0 +1,45 @@
+# This script enables the virtual environment that is passed in as an argument to inherit
+# site-packages directories of the environment from which this script is run
+
+import glob
+import os
+import site
+import sys
+
+
+def inherit_site_packages(venv_path):
+    # Identify target venv site-packages directory
+    target_venv_glob_matches = glob.glob(
+        os.path.join(
+            venv_path,
+            "lib",
+            f"python{sys.version_info[0]}.{sys.version_info[1]}",
+            "site-packages",
+        )
+    )
+    assert (
+        len(target_venv_glob_matches) == 1
+    ), f"{__file__} expected to find one supported python version in venv {venv_path} but found: {target_venv_glob_matches}"
+
+    # Create a .pth file with site-packages directories to inherit, in the target venv site-packages directory
+    # Ref: https://docs.python.org/3/library/site.html#module-site
+    with open(
+        os.path.join(target_venv_glob_matches[0], "inherited-site-packages.pth"), "w"
+    ) as f:
+        if site.ENABLE_USER_SITE:
+            user_site_packages_dir = site.getusersitepackages()
+            if os.path.exists(user_site_packages_dir):
+                f.write(f"{user_site_packages_dir}\n")
+                print(user_site_packages_dir)
+
+        for global_site_packages_dir in site.getsitepackages():
+            if os.path.exists(global_site_packages_dir):
+                f.write(f"{global_site_packages_dir}\n")
+                print(global_site_packages_dir)
+
+
+if __name__ == "__main__":
+    assert (
+        len(sys.argv) == 2
+    ), f"{__file__} expects one argument: path to venv that should inherit site-packages of the current environment but got {sys.argv}"
+    inherit_site_packages(sys.argv[1])