[fix] prevent requests being sent to python model until model is full… (

#2400)
deepjavalibrary · Sep 20, 2024 · 1459f46 · 1459f46
1 parent 7181117
commit 1459f46
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/PyPredictor.java b/engines/python/src/main/java/ai/djl/python/engine/PyPredictor.java
@@ -63,7 +63,7 @@ public PyPredictor(
     @Override
     @SuppressWarnings("unchecked")
     public List<O> batchPredict(List<I> inputs) throws TranslateException {
-        if (process.isStopped()) {
+        if (!process.isReady()) {
             // TODO: wait for restart
             throw new TranslateException("Backend Python process is stopped.");
         }

diff --git a/engines/python/src/main/java/ai/djl/python/engine/PyProcess.java b/engines/python/src/main/java/ai/djl/python/engine/PyProcess.java
@@ -50,6 +50,7 @@ class PyProcess {
     private List<Connection> connections;
     private CountDownLatch latch;
     private volatile boolean started; // NOPMD
+    private volatile boolean modelLoaded; // NOPMD
     private AtomicInteger restartCount;
     private CompletableFuture<Void> restartFuture;
     private boolean trtLlmMode;
@@ -156,6 +157,7 @@ Output predict(Input inputs, int timeout, boolean initialLoad) {
 
     synchronized void startPythonProcess() {
         try {
+            modelLoaded = false;
             int id = restartCount.get();
             int port = connections.get(0).getPort();
             logger.info("Start process: {} - retry: {}", port, id);
@@ -191,6 +193,7 @@ synchronized void startPythonProcess() {
             Input init = new Input();
             init.setProperties(pyEnv.getInitParameters());
             predict(init, pyEnv.getModelLoadingTimeout(), true);
+            modelLoaded = true;
         } catch (EngineException e) {
             started = false;
             throw e;
@@ -256,8 +259,8 @@ void setStarted(boolean started, int id) {
         }
     }
 
-    boolean isStopped() {
-        return !started;
+    boolean isReady() {
+        return started && modelLoaded;
     }
 
     private static String[] getHosts(int clusterSize) {