Refactor model venv support (#2954)

* Refactor model venv creation * Update assertion error message --------- Co-authored-by: Ankith Gunapal <agunapal@ischool.Berkeley.edu>
pytorch · Feb 23, 2024 · 048a123 · 048a123
1 parent ccc3f95
commit 048a123
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 57 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java b/frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java
@@ -218,9 +218,10 @@ private void setupModelVenv(Model model)
         File venvPath = EnvironmentUtils.getPythonVenvPath(model);
         List<String> commandParts = new ArrayList<>();
         commandParts.add(configManager.getPythonExecutable());
-        commandParts.add("-m");
-        commandParts.add("venv");
-        commandParts.add("--clear");
+        commandParts.add(
+                Paths.get(configManager.getModelServerHome(), "ts", "utils", "setup_model_venv.py")
+                        .toAbsolutePath()
+                        .toString());
         commandParts.add(venvPath.toString());
 
         ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
@@ -272,57 +273,6 @@ 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)

diff --git a/ts/utils/inherit_site_packages.py → ts/utils/setup_model_venv.py b/ts/utils/inherit_site_packages.py → ts/utils/setup_model_venv.py
@@ -1,10 +1,15 @@
-# 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
+# This script enables creation of the model virtual environment whose path is passed in as an argument.
+# The virtual environment inherits the site-packages directories of the environment from which this script is run.
 
 import glob
 import os
 import site
 import sys
+import venv
+
+
+def create_venv(venv_path):
+    venv.create(venv_path, clear=True, with_pip=True)
 
 
 def inherit_site_packages(venv_path):
@@ -41,5 +46,6 @@ def inherit_site_packages(venv_path):
 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}"
+    ), f"{__file__} expects one argument: path to venv that should be created, but got {sys.argv}"
+    create_venv(sys.argv[1])
     inherit_site_packages(sys.argv[1])