Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[xgb] Fix xgb intern to close replaced array #3558

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[xgb] Fix xgb intern to close replaced array #3558

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ddfae0e
[xgb] Fix xgb intern to close replaced array
ewan0x79 Dec 10, 2024
4b9d4ae
[xgb] Fix xgb intern to close replaced array
ewan0x79 Dec 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 28 additions & 6 deletions engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,14 +88,36 @@ public ByteBuffer toByteBuffer(boolean tryDirect) {
/** {@inheritDoc} */
@Override
public void intern(NDArray replaced) {
if (handle != null && handle.get() != 0L) {
long pointer = handle.getAndSet(0L);
JniUtils.deleteDMatrix(pointer);
if (replaced == null) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not necessary, the instanceof will check null:

if (!(replaced instanceof XgbNDArray))

throw new IllegalArgumentException("The replaced NDArray cannot be null.");
}
if (!(replaced instanceof XgbNDArray)) {
throw new IllegalArgumentException(
"The replaced NDArray must be an instance of XgbNDArray.");
}
XgbNDArray array = (XgbNDArray) replaced;
data = array.data;
handle = array.handle;
format = array.format;

synchronized (this) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NDArray operation is not thread-safe. Adding synchronized here won't improve thread-safty

if (handle != null && handle.get() != 0L) {
long pointer = handle.getAndSet(0L);
JniUtils.deleteDMatrix(pointer);
}

data = array.data;
format = array.format;

if (array.handle != null && array.handle.get() != 0L) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure about this null check, if we want to prevent a closed NDArray, we should just do:

if (this.closed()

if (handle == null) {
handle = new AtomicLong();
}
handle.set(array.handle.getAndSet(0L));
}
}

array.data = null;
array.handle = null;
array.format = null;
array.close();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The memory leak issue may caused by alternativeArray leak, I think we should close current array's alternative array as will. The code can be re-implemented as the following:

        if (!(replaced instanceof XgbNDArray)) {
            throw new IllegalArgumentException(
                    "The replaced NDArray must be an instance of XgbNDArray.");
        }
        XgbNDArray array = (XgbNDArray) replaced;
        if (isReleased()) {
            throw new IllegalArgumentException("This array is already closed");
        }
        if (replaced.isReleased()) {
            throw new IllegalArgumentException("This target array is already closed");
        }

        long pointer = handle.getAndSet(0L);
        JniUtils.deleteDMatrix(pointer);
        if (alternativeArray != null) {
            alternativeArray.close();
        }

        data = array.data;
        handle = array.handle;
        format = array.format;
        alternativeArray = array.alternativeArray;
        array.handle = null;
        array.alternativeArray = null;
        array.close();

}

/** {@inheritDoc} */
Expand Down
Loading