Repository: deepjavalibrary/djl
Pull request #3225 (2024-05-29 20:04:28)

add close method to all nd manager
## Description ##

Close all alternative ndmanager 



Patch:
diff --git a/.github/workflows/publish-job-success.yml b/.github/workflows/publish-job-success.yml
index ec75dd635..91b481f8f 100644
--- a/.github/workflows/publish-job-success.yml
+++ b/.github/workflows/publish-job-success.yml
@@ -14,7 +14,7 @@ permissions:
 
 jobs:
   publish-job-success-to-cloudwatch:
-    if: ${{ github.event.workflow_run.event == 'schedule' }}
+    if: ${{ github.repository == 'deepjavalibrary/djl' && github.event.workflow_run.event == 'schedule' }}
     runs-on: ubuntu-latest
     steps:
       - name: Configure AWS Credentials
diff --git a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
index 1de074ea6..859596a6b 100644
--- a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
+++ b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
@@ -61,7 +61,11 @@ public final class CudaUtils {
             if (gpuInfo == null) {
                 gpuInfo = execute(-1); // NOPMD
             }
-            return Integer.parseInt(gpuInfo[0]);
+            try {
+                return Integer.parseInt(gpuInfo[0]);
+            } catch (NumberFormatException e) {
+                return 0;
+            }
         }
 
         if (LIB == null) {
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java
index b9f79a124..80351ee79 100644
--- a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java
+++ b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java
@@ -104,6 +104,16 @@ public class LgbmNDManager extends BaseNDManager {
         return new NDList(new LgbmDataset(this, null, path));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link LgbmNDManager} of which all others are children. */
     private static final class SystemManager extends LgbmNDManager implements SystemNDManager {
 
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
index 81f9708e7..d1406b140 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
@@ -180,6 +180,16 @@ public class XgbNDManager extends BaseNDManager {
         return new XgbNDArray(this, alternativeManager, handle, shape, SparseFormat.CSR);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link XgbNDManager} of which all others are children. */
     private static final class SystemManager extends XgbNDManager implements SystemNDManager {
 
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java
index f69276b61..ebb4de96c 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java
@@ -126,6 +126,16 @@ public class TrtNDManager extends BaseNDManager {
         return create(bb, shape, dataType);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link TrtNDManager} of which all others are children. */
     private static final class SystemManager extends TrtNDManager implements SystemNDManager {
 
diff --git a/extensions/tokenizers/rust/src/models/mod.rs b/extensions/tokenizers/rust/src/models/mod.rs
index a4db801bb..b915f9229 100644
--- a/extensions/tokenizers/rust/src/models/mod.rs
+++ b/extensions/tokenizers/rust/src/models/mod.rs
@@ -124,7 +124,8 @@ pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_loadModel<'local>(
 
 #[no_mangle]
 pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_deleteModel<'local>(
-    _env: JObject,
+    _: JNIEnv,
+    _: JObject,
     handle: jlong,
 ) {
     drop_handle::<Box<dyn Model>>(handle);


Test patch:
diff --git a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
index e09417fa5..df135938a 100644
--- a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
+++ b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
@@ -13,6 +13,7 @@
 package ai.djl.util.cuda;
 
 import ai.djl.Device;
+import ai.djl.engine.EngineException;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -30,17 +31,13 @@ public class CudaUtilsTest {
         if (!CudaUtils.hasCuda()) {
             try {
                 CudaUtils.getCudaVersionString();
-            } catch (IllegalStateException | IllegalArgumentException ignore) {
+            } catch (IllegalStateException | IllegalArgumentException | EngineException ignore) {
                 // ignore exception
             }
             Assert.assertThrows(() -> CudaUtils.getComputeCapability(0));
             Assert.assertThrows(() -> CudaUtils.getGpuMemory(Device.gpu()));
             return;
         }
-        // Possible to have CUDA and not have a GPU.
-        if (CudaUtils.getGpuCount() == 0) {
-            return;
-        }
 
         String cudaVersion = CudaUtils.getCudaVersionString();
         String smVersion = CudaUtils.getComputeCapability(0);