Refactor createVirtualEnv()

engines/python/src/main/java/ai/djl/python/engine/PyEnv.java

@@ -172,37 +172,40 @@ public Map<String, String> getInitParameters() {
      *
      * @param name the virtual environment name
      */
-    public synchronized void createVirtualEnv(String name) {
-        if (venvCreated) {
-            return;
-        }
-        Path path = getVenvDir().resolve(name).toAbsolutePath();
-        if (Files.exists(path)) {
-            logger.info("Virtual environment already exists at {}.", path);
-            setPythonExecutable(path.resolve("bin").resolve("python").toString());
-            venvCreated = true;
-            return;
-        }
-        String[] cmd = {pythonExecutable, "-m", "venv", path.toString(), "--system-site-packages"};
-
-        try {
-            Process process = new ProcessBuilder(cmd).redirectErrorStream(true).start();
-            String logOutput;
-            try (InputStream is = process.getInputStream()) {
-                logOutput = Utils.toString(is);
+    public void createVirtualEnv(String name) {
+        synchronized (PyEnv.class) {
+            if (venvCreated) {
+                return;
             }
-            int ret = process.waitFor();
-            logger.debug("{}", logOutput);
-            if (ret != 0) {
-                throw new EngineException(
-                        "Failed to create virtual environment with error code: " + ret);
+            Path path = getVenvDir().resolve(name).toAbsolutePath();
+            if (Files.exists(path)) {
+                logger.info("Virtual environment already exists at {}.", path);
+                setPythonExecutable(path.resolve("bin").resolve("python").toString());
+                venvCreated = true;
+                return;
             }
+            String[] cmd = {
+                pythonExecutable, "-m", "venv", path.toString(), "--system-site-packages"
+            };
+
+            try {
+                Process process = new ProcessBuilder(cmd).redirectErrorStream(true).start();
+                String logOutput;
+                try (InputStream is = process.getInputStream()) {
+                    logOutput = Utils.toString(is);
+                }
+                int ret = process.waitFor();
+                logger.debug("{}", logOutput);
+                if (ret != 0) {
+                    throw new EngineException("Failed to create venv with error code: " + ret);
+                }
 
-            logger.info("Python virtual environment created successfully at {}!", path);
-            setPythonExecutable(path.resolve("bin").resolve("python").toString());
-            venvCreated = true;
-        } catch (IOException | InterruptedException e) {
-            throw new EngineException("Python virtual failed", e);
+                logger.info("Python virtual environment created successfully at {}!", path);
+                setPythonExecutable(path.resolve("bin").resolve("python").toString());
+                venvCreated = true;
+            } catch (IOException | InterruptedException e) {
+                throw new EngineException("Create venv failed", e);
+            }
         }
     }
 
@@ -212,9 +215,6 @@ public synchronized void createVirtualEnv(String name) {
      * @param name the virtual environment name
      */
     public synchronized void deleteVirtualEnv(String name) {
-        if (!venvCreated) {
-            return;
-        }
         Path path = getVenvDir().resolve(name);
         Utils.deleteQuietly(path);
     }

engines/python/src/main/java/ai/djl/python/engine/PyModel.java

@@ -185,7 +185,7 @@ public void load(Path modelPath, String prefix, Map<String, ?> options) throws I
         }
         pyEnv.setEntryPoint(entryPoint);
         if (pyEnv.isEnableVenv()) {
-            pyEnv.createVirtualEnv(getName());
+            pyEnv.createVirtualEnv(Utils.hash(modelDir.toString()));
         }
 
         if (pyEnv.isMpiMode()) {
@@ -325,12 +325,12 @@ private void createAllPyProcesses(int mpiWorkers) {
     }
 
     private void shutdown() {
-        if (pyEnv.isEnableVenv()) {
-            pyEnv.deleteVirtualEnv(getName());
-        }
         for (PyProcess process : workerQueue) {
             process.stopPythonProcess();
         }
         workerQueue.clear();
+        if (pyEnv.isEnableVenv()) {
+            pyEnv.deleteVirtualEnv(Utils.hash(modelDir.toString()));
+        }
     }
 }

engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java

@@ -400,30 +400,6 @@ public void testRollingBatch() throws TranslateException, IOException, ModelExce
         }
     }
 
-    @Test
-    public void testPythonVenv() throws IOException, ModelException {
-        Criteria<Input, Output> criteria =
-                Criteria.builder()
-                        .setTypes(Input.class, Output.class)
-                        .optEngine("Python")
-                        .optModelPath(Paths.get("src/test/resources/venv"))
-                        .build();
-
-        String venvDir = "build/venv";
-        System.setProperty("DJL_VENV_DIR", venvDir);
-        Path path;
-        try (ZooModel<Input, Output> model = criteria.loadModel()) {
-            path = Paths.get(venvDir).resolve(model.getName());
-            Assert.assertTrue(Files.isDirectory(path));
-        }
-        Assert.assertFalse(Files.isDirectory(path));
-
-        // Test exception
-        venvDir = "/COM1"; // Invalid directory
-        System.setProperty("DJL_VENV_DIR", venvDir);
-        Assert.assertThrows(EngineException.class, criteria::loadModel);
-    }
-
     @Test
     public void testModelException() throws TranslateException, IOException, ModelException {
         Criteria<Input, Output> criteria =

engines/python/src/test/java/ai/djl/python/engine/PyEnvTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.python.engine;
+
+import ai.djl.MalformedModelException;
+import ai.djl.engine.EngineException;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ModelNotFoundException;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.util.Utils;
+
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+public class PyEnvTest {
+
+    @BeforeClass
+    public void setUp() {
+        System.setProperty("DJL_VENV_DIR", "build/venv");
+    }
+
+    @AfterClass
+    public void tierDown() {
+        System.clearProperty("DJL_VENV_DIR");
+        Utils.deleteQuietly(Paths.get("build/venv"));
+    }
+
+    @Test
+    public void testPythonVenv()
+            throws ModelNotFoundException, MalformedModelException, IOException {
+        Criteria<Input, Output> criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optEngine("Python")
+                        .optModelPath(Paths.get("src/test/resources/echo"))
+                        .optOption("enable_venv", "true")
+                        .build();
+
+        Path venvDir;
+        try (ZooModel<Input, Output> model = criteria.loadModel();
+                ZooModel<Input, Output> model2 = criteria.loadModel()) {
+            String venvName = Utils.hash(model.getModelPath().toString());
+            venvDir = Paths.get("build/venv").resolve(venvName);
+            Assert.assertTrue(Files.exists(venvDir));
+            Assert.assertNotNull(model2.getModelPath());
+        }
+        Assert.assertFalse(Files.exists(venvDir));
+
+        // Test exception
+        if (!System.getProperty("os.name").startsWith("Win")) {
+            System.setProperty("DJL_VENV_DIR", "/non_exist/venv");
+            Assert.assertThrows(EngineException.class, criteria::loadModel);
+        }
+    }
+}