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.

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

[api] Fixes memory leak in hybrid engine #1518

Merged
merged 1 commit into from
Mar 2, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
13 changes: 7 additions & 6 deletions api/src/main/java/ai/djl/ndarray/BaseNDManager.java
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ public abstract class BaseNDManager implements NDManager {
private static final Logger logger = LoggerFactory.getLogger(BaseNDManager.class);

protected NDManager parent;
protected NDManager alternativeManager;
protected String uid;
protected String name;
protected Device device;
Expand All @@ -53,6 +54,10 @@ protected BaseNDManager(NDManager parent, Device device) {
resources = new ConcurrentHashMap<>();
tempResources = new ConcurrentHashMap<>();
uid = UUID.randomUUID().toString();
Engine engine = getEngine().getAlternativeEngine();
if (engine != null) {
alternativeManager = engine.newBaseManager(Device.cpu());
}
}

/** {@inheritDoc} */
Expand Down Expand Up @@ -368,12 +373,8 @@ public void debugDump(int level) {
}
}

protected NDManager getAlternativeManager() {
Engine engine = getEngine().getAlternativeEngine();
if (engine != null) {
return engine.newBaseManager(Device.cpu());
}
return null;
NDManager getAlternativeManager() {
return alternativeManager;
}

/**
Expand Down
4 changes: 4 additions & 0 deletions api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,10 @@ public void attach(NDManager manager) {
detach();
this.manager = manager;
manager.attachInternal(getUid(), this);
alternativeManager = ((BaseNDManager) manager).getAlternativeManager();
if (alternativeManager == null) {
alternativeManager = manager;
}
}

/** {@inheritDoc} */
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,8 @@ public class DlrNDManager extends BaseNDManager {

private static final DlrNDManager SYSTEM_MANAGER = new SystemManager();

private NDManager alternativeManager;

private DlrNDManager(NDManager parent, Device device) {
super(parent, device);
alternativeManager = getAlternativeManager();
}

static DlrNDManager getSystemManager() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,8 @@ public class XgbNDManager extends BaseNDManager {

private static final XgbNDManager SYSTEM_MANAGER = new SystemManager();

private NDManager alternativeManager;

private XgbNDManager(NDManager parent, Device device) {
super(parent, device);
alternativeManager = getAlternativeManager();
}

static XgbNDManager getSystemManager() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,10 @@ public class OrtNDManager extends BaseNDManager {
private static final OrtNDManager SYSTEM_MANAGER = new SystemManager();

private OrtEnvironment env;
private NDManager alternativeManager;

private OrtNDManager(NDManager parent, Device device, OrtEnvironment env) {
super(parent, device);
this.env = env;
alternativeManager = getAlternativeManager();
}

static OrtNDManager getSystemManager() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,8 @@ public class PpNDManager extends BaseNDManager {

private static final PpNDManager SYSTEM_MANAGER = new SystemManager();

private NDManager alternativeManager;

private PpNDManager(NDManager parent, Device device) {
super(parent, device);
alternativeManager = getAlternativeManager();
}

static PpNDManager getSystemManager() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,8 @@ public class TrtNDManager extends BaseNDManager {

private static final TrtNDManager SYSTEM_MANAGER = new SystemManager();

private NDManager alternativeManager;

private TrtNDManager(NDManager parent, Device device) {
super(parent, device);
alternativeManager = getAlternativeManager();
}

static TrtNDManager getSystemManager() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,8 @@ public class TfLiteNDManager extends BaseNDManager {

private static final TfLiteNDManager SYSTEM_MANAGER = new SystemManager();

private NDManager alternativeManager;

private TfLiteNDManager(NDManager parent, Device device) {
super(parent, device);
alternativeManager = getAlternativeManager();
}

static TfLiteNDManager getSystemManager() {
Expand Down