Enable model custom dependency installation using virtual environment #2910
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,7 +6,6 @@ | |
| import java.io.IOException; | ||
| import java.io.InputStreamReader; | ||
| import java.net.HttpURLConnection; | ||
| import java.nio.file.Path; | ||
| import java.nio.file.Paths; | ||
| import java.util.ArrayList; | ||
|
|
@@ -205,87 +204,139 @@ private ModelArchive createModelArchive( | |
| return archive; | ||
| } | ||
|
|
||
| private void setupModelVenv(Model model) | ||
| throws IOException, InterruptedException, ModelException { | ||
| String venvPath = EnvironmentUtils.getPythonVenvPath(model); | ||
| List<String> commandParts = new ArrayList<>(); | ||
| commandParts.add(EnvironmentUtils.getPythonRunTime(model)); | ||
| commandParts.add("-m"); | ||
| commandParts.add("venv"); | ||
| commandParts.add("--clear"); | ||
|
There was a problem hiding this comment. why clear? Seems beneficial to allow users to install their dependencies beforehand It was always quite weird how we pip installed a bunch of stuff on launching a model There was a problem hiding this comment. Sounds good, will check if From the official docs: https://docs.python.org/3/library/venv.html |
||
| commandParts.add("--system-site-packages"); | ||
|
There was a problem hiding this comment. This is to make sure that the |
||
| commandParts.add(venvPath); | ||
|
|
||
| ProcessBuilder processBuilder = new ProcessBuilder(commandParts); | ||
|
|
||
| if (isValidVenvPath(venvPath)) { | ||
| processBuilder.directory(Paths.get(venvPath).toFile().getParentFile()); | ||
| } else { | ||
| throw new ModelException( | ||
| "Invalid python venv path for model " + model.getModelName() + ": " + venvPath); | ||
| } | ||
| Map<String, String> environment = processBuilder.environment(); | ||
| String[] envp = | ||
| EnvironmentUtils.getEnvString( | ||
| configManager.getModelServerHome(), | ||
| model.getModelDir().getAbsolutePath(), | ||
| null); | ||
| for (String envVar : envp) { | ||
| String[] parts = envVar.split("=", 2); | ||
agunapal marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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( | ||
| "Created virtual environment for model {}: {}", model.getModelName(), venvPath); | ||
| } else { | ||
| logger.error( | ||
| "Virtual environment creation for model {} at {} failed:\n{}", | ||
| model.getModelName(), | ||
| venvPath, | ||
| outputString.toString()); | ||
| throw new ModelException( | ||
| "Virtual environment creation failed for model " + model.getModelName()); | ||
| } | ||
| } | ||
|
|
||
| private void setupModelDependencies(Model model) | ||
| throws IOException, InterruptedException, ModelException { | ||
| String requirementsFile = | ||
| model.getModelArchive().getManifest().getModel().getRequirementsFile(); | ||
|
|
||
| if (configManager.getInstallPyDepPerModel() && requirementsFile != null) { | ||
| setupModelVenv(model); | ||
| String pythonRuntime = EnvironmentUtils.getPythonRunTime(model); | ||
| if (!isValidVenvPath(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<>(); | ||
| commandParts.add(pythonRuntime); | ||
| commandParts.add("-m"); | ||
| commandParts.add("pip"); | ||
| commandParts.add("install"); | ||
| commandParts.add("-U"); | ||
| commandParts.add("--upgrade-strategy"); | ||
| commandParts.add("only-if-needed"); | ||
|
There was a problem hiding this comment. I think this is a good change but do you mind just telling me why you needed to make it? There was a problem hiding this comment. In prior versions of pip (i.e pip<10.0) when From the pip docs: https://pip.pypa.io/en/stable/user_guide/#only-if-needed-recursive-upgrade |
||
| 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]); | ||
| } | ||
| } | ||
| processBuilder.directory( | ||
| Paths.get(EnvironmentUtils.getPythonVenvPath(model)).toFile().getParentFile()); | ||
| 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) | ||
| .normalize() | ||
| .startsWith(FileUtils.getTempDirectory().toPath().normalize())) { | ||
| return true; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
should you make this name customizable? Part of the appeal of this PR is different workers should be able to have different virtual environments
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Currently this PR creates a virtual environment on a per model basis at model load time. All workers for a given model use the same virtual environment. This replaces installing dependencies on a per model basis in a target directory and is backwards compatible with the existing behavior with no change to customer experience. Although the same name
venvis used, they are located within the individual model directories, for ex:/tmp/models/test-model/venv.Would it be useful to extend this implementation to support separate
venvper worker?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No the isolation is fine as is I think