MME - deviceId while creating workers (deepjavalibrary#1257)

Co-authored-by: sindhuso <sindhuso@3c0630156b9b.ant.amazon.com> Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
maaquib · Nov 4, 2023 · 61c2a23 · 61c2a23
1 parent 415d017
commit 61c2a23
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/PyModel.java b/engines/python/src/main/java/ai/djl/python/engine/PyModel.java
@@ -316,9 +316,10 @@ private void createAllPyProcesses(int mpiWorkers, int tp) {
             pool = Executors.newFixedThreadPool(mpiWorkers);
         }
         logger.info("Start {} mpiWorkers ...", mpiWorkers);
+        int deviceId = manager.getDevice().getDeviceId();
         for (int i = 0; i < mpiWorkers; ++i) {
             logger.debug("Pre-creating python worker: {} ", i);
-            PyProcess worker = new PyProcess(this, pyEnv, i * tp);
+            PyProcess worker = new PyProcess(this, pyEnv, deviceId + i * tp);
             workerQueue.offer(worker);
             if (pool != null) {
                 logger.debug("Submitting to pool: {}", i);

diff --git a/serving/src/main/java/ai/djl/serving/ModelServer.java b/serving/src/main/java/ai/djl/serving/ModelServer.java
@@ -324,9 +324,17 @@ private void initModelStore() throws IOException, BadWorkflowException {
             }
 
             if (Files.isDirectory(modelStore)) {
-                // Check if root model store folder contains a model
-                String url = mapModelUrl(modelStore);
-                if (url == null) {
+                // contains only directory or archive files
+                boolean isMultiModelsDirectory =
+                        Files.list(modelStore)
+                                .filter(p -> !p.getFileName().toString().startsWith("."))
+                                .allMatch(
+                                        p ->
+                                                Files.isDirectory(p)
+                                                        || FilenameUtils.isArchiveFile(
+                                                                p.toString()));
+
+                if (isMultiModelsDirectory) {
                     // Check folders to see if they can be models as well
                     try (Stream<Path> stream = Files.list(modelStore)) {
                         urls.addAll(
@@ -335,7 +343,11 @@ private void initModelStore() throws IOException, BadWorkflowException {
                                         .collect(Collectors.toList()));
                     }
                 } else {
-                    urls.add(url);
+                    // Check if root model store folder contains a model
+                    String url = mapModelUrl(modelStore);
+                    if (url != null) {
+                        urls.add(url);
+                    }
                 }
             } else {
                 logger.warn("Model store path is not found: {}", modelStore);