Enable per model useVenv configuration option

Author: namannandan

frontend/archive/src/main/java/org/pytorch/serve/archive/model/ModelConfig.java

@@ -65,9 +65,14 @@ public class ModelConfig {
     private int maxSequenceJobQueueSize = 1;
     /** the max number of sequences can be accepted. The default value is 1. */
     private int maxNumSequence = 1;
-
     /** continuousBatching is a flag to enable continuous batching. */
     private boolean continuousBatching;
+    /**
+     * Create python virtual environment when using python backend to: 1) Install model dependencies
+     * (if enabled globally using install_py_dep_per_model=true) 2) Run workers for model loading
+     * and inference
+     */
+    private boolean useVenv;
 
     public static ModelConfig build(Map<String, Object> yamlMap) {
         ModelConfig modelConfig = new ModelConfig();
@@ -207,6 +212,13 @@ public static ModelConfig build(Map<String, Object> yamlMap) {
                                         v);
                             }
                             break;
+                        case "useVenv":
+                            if (v instanceof Boolean) {
+                                modelConfig.setUseVenv((boolean) v);
+                            } else {
+                                logger.warn("Invalid useVenv: {}, should be true or false", v);
+                            }
+                            break;
                         default:
                             break;
                     }
@@ -379,6 +391,14 @@ public void setMaxNumSequence(int maxNumSequence) {
         this.maxNumSequence = Math.max(1, maxNumSequence);
     }
 
+    public boolean getUseVenv() {
+        return useVenv;
+    }
+
+    public void setUseVenv(boolean useVenv) {
+        this.useVenv = useVenv;
+    }
+
     public enum ParallelType {
         NONE(""),
         PP("pp"),

frontend/server/src/main/java/org/pytorch/serve/util/messages/EnvironmentUtils.java

@@ -10,6 +10,7 @@
 import java.util.Map;
 import java.util.regex.Pattern;
 import org.pytorch.serve.archive.model.Manifest;
+import org.pytorch.serve.archive.model.ModelConfig;
 import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.wlm.Model;
 import org.slf4j.Logger;
@@ -79,8 +80,11 @@ public static String getPythonRunTime(Model model) {
         Manifest.RuntimeType runtime = model.getModelArchive().getManifest().getRuntime();
         if (runtime == Manifest.RuntimeType.PYTHON) {
             pythonRuntime = configManager.getPythonExecutable();
+            ModelConfig modelConfig = model.getModelArchive().getModelConfig();
             Path pythonVenvRuntime = Paths.get(getPythonVenvPath(model), "bin", "python");
-            if (Files.exists(pythonVenvRuntime)) {
+            if (modelConfig != null
+                    && modelConfig.getUseVenv() == true
+                    && Files.exists(pythonVenvRuntime)) {
                 pythonRuntime = pythonVenvRuntime.toString();
             }
         } else {

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

@@ -6,6 +6,7 @@
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.net.HttpURLConnection;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
@@ -99,6 +100,8 @@ public void registerAndUpdateModel(String modelName, JsonObject modelInfo)
 
         createVersionedModel(tempModel, versionId);
 
+        setupModelVenv(tempModel);
+
         setupModelDependencies(tempModel);
         if (defaultVersion) {
             modelManager.setDefaultVersion(modelName, versionId);
@@ -152,6 +155,8 @@ public ModelArchive registerModel(
             }
         }
 
+        setupModelVenv(tempModel);
+
         setupModelDependencies(tempModel);
 
         logger.info("Model {} loaded.", tempModel.getModelName());
@@ -206,9 +211,16 @@ private ModelArchive createModelArchive(
 
     private void setupModelVenv(Model model)
             throws IOException, InterruptedException, ModelException {
+        ModelConfig modelConfig = model.getModelArchive().getModelConfig();
+        if (model.getModelArchive().getManifest().getRuntime() != Manifest.RuntimeType.PYTHON
+                || modelConfig == null
+                || modelConfig.getUseVenv() != true) {
+            return;
+        }
+
         String venvPath = EnvironmentUtils.getPythonVenvPath(model);
         List<String> commandParts = new ArrayList<>();
-        commandParts.add(EnvironmentUtils.getPythonRunTime(model));
+        commandParts.add(configManager.getPythonExecutable());
         commandParts.add("-m");
         commandParts.add("venv");
         commandParts.add("--clear");
@@ -217,7 +229,7 @@ private void setupModelVenv(Model model)
 
         ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
 
-        if (isValidVenvPath(venvPath)) {
+        if (isValidDependencyPath(venvPath)) {
             processBuilder.directory(Paths.get(venvPath).toFile().getParentFile());
         } else {
             throw new ModelException(
@@ -266,20 +278,29 @@ private void setupModelDependencies(Model model)
         String requirementsFile =
                 model.getModelArchive().getManifest().getModel().getRequirementsFile();
 
-        if (configManager.getInstallPyDepPerModel() && requirementsFile != null) {
-            setupModelVenv(model);
-            String pythonRuntime = EnvironmentUtils.getPythonRunTime(model);
-            if (!isValidVenvPath(pythonRuntime)) {
+        if (!configManager.getInstallPyDepPerModel() || requirementsFile == null) {
+            return;
+        }
+
+        ModelConfig modelConfig = model.getModelArchive().getModelConfig();
+        String pythonRuntime = EnvironmentUtils.getPythonRunTime(model);
+        Path requirementsFilePath =
+                Paths.get(model.getModelDir().getAbsolutePath(), requirementsFile);
+        List<String> commandParts = new ArrayList<>();
+        ProcessBuilder processBuilder = new ProcessBuilder();
+
+        if (modelConfig != null && modelConfig.getUseVenv() == true) {
+            if (!isValidDependencyPath(pythonRuntime)) {
                 throw new ModelException(
                         "Invalid python venv runtime path for model "
                                 + model.getModelName()
                                 + ": "
                                 + pythonRuntime);
             }
 
-            Path requirementsFilePath =
-                    Paths.get(model.getModelDir().getAbsolutePath(), requirementsFile);
-            List<String> commandParts = new ArrayList<>();
+            processBuilder.directory(
+                    Paths.get(EnvironmentUtils.getPythonVenvPath(model)).toFile().getParentFile());
+
             commandParts.add(pythonRuntime);
             commandParts.add("-m");
             commandParts.add("pip");
@@ -289,54 +310,72 @@ private void setupModelDependencies(Model model)
             commandParts.add("only-if-needed");
             commandParts.add("-r");
             commandParts.add(requirementsFilePath.toString());
-
-            ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
-
-            String[] envp =
-                    EnvironmentUtils.getEnvString(
-                            configManager.getModelServerHome(),
-                            model.getModelDir().getAbsolutePath(),
-                            null);
-            Map<String, String> environment = processBuilder.environment();
-            for (String envVar : envp) {
-                String[] parts = envVar.split("=", 2);
-                if (parts.length == 2) {
-                    environment.put(parts[0], parts[1]);
-                }
+        } else {
+            File dependencyPath = model.getModelDir();
+            if (Files.isSymbolicLink(dependencyPath.toPath())) {
+                dependencyPath = dependencyPath.getParentFile();
+            }
+            if (!isValidDependencyPath(dependencyPath.getPath())) {
+                throw new ModelException(
+                        "Invalid 3rd party package installation path "
+                                + dependencyPath.getCanonicalPath());
             }
-            processBuilder.directory(
-                    Paths.get(EnvironmentUtils.getPythonVenvPath(model)).toFile().getParentFile());
-            processBuilder.redirectErrorStream(true);
 
-            Process process = processBuilder.start();
+            processBuilder.directory(dependencyPath);
 
-            int exitCode = process.waitFor();
-            String line;
-            StringBuilder outputString = new StringBuilder();
-            BufferedReader brdr =
-                    new BufferedReader(new InputStreamReader(process.getInputStream()));
-            while ((line = brdr.readLine()) != null) {
-                outputString.append(line + "\n");
-            }
+            commandParts.add(pythonRuntime);
+            commandParts.add("-m");
+            commandParts.add("pip");
+            commandParts.add("install");
+            commandParts.add("-U");
+            commandParts.add("-t");
+            commandParts.add(dependencyPath.getAbsolutePath());
+            commandParts.add("-r");
+            commandParts.add(requirementsFilePath.toString());
+        }
 
-            if (exitCode == 0) {
-                logger.info(
-                        "Installed custom pip packages for model {}:\n{}",
-                        model.getModelName(),
-                        outputString.toString());
-            } else {
-                logger.error(
-                        "Custom pip package installation failed for model {}:\n{}",
-                        model.getModelName(),
-                        outputString.toString());
-                throw new ModelException(
-                        "Custom pip package installation failed for model " + model.getModelName());
+        processBuilder.command(commandParts);
+        String[] envp =
+                EnvironmentUtils.getEnvString(
+                        configManager.getModelServerHome(),
+                        model.getModelDir().getAbsolutePath(),
+                        null);
+        Map<String, String> environment = processBuilder.environment();
+        for (String envVar : envp) {
+            String[] parts = envVar.split("=", 2);
+            if (parts.length == 2) {
+                environment.put(parts[0], parts[1]);
             }
         }
+        processBuilder.redirectErrorStream(true);
+
+        Process process = processBuilder.start();
+
+        int exitCode = process.waitFor();
+        String line;
+        StringBuilder outputString = new StringBuilder();
+        BufferedReader brdr = new BufferedReader(new InputStreamReader(process.getInputStream()));
+        while ((line = brdr.readLine()) != null) {
+            outputString.append(line + "\n");
+        }
+
+        if (exitCode == 0) {
+            logger.info(
+                    "Installed custom pip packages for model {}:\n{}",
+                    model.getModelName(),
+                    outputString.toString());
+        } else {
+            logger.error(
+                    "Custom pip package installation failed for model {}:\n{}",
+                    model.getModelName(),
+                    outputString.toString());
+            throw new ModelException(
+                    "Custom pip package installation failed for model " + model.getModelName());
+        }
     }
 
-    private boolean isValidVenvPath(String venvPath) {
-        if (Paths.get(venvPath)
+    private boolean isValidDependencyPath(String dependencyPath) {
+        if (Paths.get(dependencyPath)
                 .normalize()
                 .startsWith(FileUtils.getTempDirectory().toPath().normalize())) {
             return true;