You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by ja...@apache.org on 2026/03/27 09:36:38 UTC
(systemds) branch main updated: [MINOR] OOC Bugfix Cache Reference Management + Return Right BlockKey on Externally Managed Grouped Callbacks
This is an automated email from the ASF dual-hosted git repository.
janniklinde pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new 1f2e2eab00 [MINOR] OOC Bugfix Cache Reference Management + Return Right BlockKey on Externally Managed Grouped Callbacks
1f2e2eab00 is described below
commit 1f2e2eab00a835f83a71e07b18f6b2b369ba1b97
Author: Jannik Lindemann <ja...@gmail.com>
AuthorDate: Fri Mar 27 10:29:30 2026 +0100
[MINOR] OOC Bugfix Cache Reference Management + Return Right BlockKey on Externally Managed Grouped Callbacks
Closes #2454.
---
.../runtime/instructions/ooc/CachingStream.java | 49 +++++++++++++++++++++-
.../apache/sysds/runtime/ooc/cache/BlockEntry.java | 10 +++++
.../sysds/runtime/ooc/cache/OOCCacheScheduler.java | 8 ++++
.../runtime/ooc/cache/OOCLRUCacheScheduler.java | 39 +++++++++++++++--
4 files changed, 101 insertions(+), 5 deletions(-)
diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java
index b3f5e57aaf..b7c4e2aa64 100644
--- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java
+++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java
@@ -130,6 +130,7 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
boolean ownsEntry = true;
if(tmp instanceof OOCCacheManager.CachedGroupCallback<?> cachedGroup) {
baseKey = cachedGroup.getBlockKey();
+ ensureReferencedOrRematerialize(baseKey, cachedGroup);
ownsEntry = false;
if(mSubscribers != null && mSubscribers.length > 0)
mCallback = tmp.keepOpen();
@@ -183,12 +184,14 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
if(tmp instanceof OOCCacheManager.CachedQueueCallback<?> cachedQueue) {
blockKey = cachedQueue.getBlockKey();
+ ensureReferencedOrRematerialize(blockKey, task);
ownsEntry = false;
if(mSubscribers != null && mSubscribers.length > 0)
mCallback = tmp.keepOpen();
}
else if(tmp instanceof OOCCacheManager.CachedSubCallback<?> cachedSub) {
BlockKey parent = cachedSub.getParent().getBlockKey();
+ ensureReferencedOrRematerialize(parent, cachedSub.getParent());
blockKey = new GroupedBlockKey(parent.getStreamId(), (int) parent.getSequenceNumber(),
cachedSub.getGroupIndex());
ownsEntry = false;
@@ -297,6 +300,49 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
});
}
+
+ private void ensureReferencedOrRematerialize(BlockKey key, IndexedMatrixValue value) {
+ try {
+ OOCCacheManager.getCache().addReference(key);
+ }
+ catch(IllegalArgumentException ex) {
+ try {
+ OOCCacheManager.putRaw(key, value, ((MatrixBlock) value.getValue()).getExactSerializedSize());
+ }
+ catch(IllegalStateException putEx) {
+ // Another downstream stream may have re-materialized the same entry first.
+ OOCCacheManager.getCache().addReference(key);
+ }
+ }
+ }
+
+ private void ensureReferencedOrRematerialize(BlockKey key, OOCCacheManager.CachedGroupCallback<?> group) {
+ try {
+ OOCCacheManager.getCache().addReference(key);
+ }
+ catch(IllegalArgumentException ex) {
+ try {
+ List<IndexedMatrixValue> values = new ArrayList<>(group.size());
+ long totalSize = 0;
+ for(int gi = 0; gi < group.size(); gi++) {
+ @SuppressWarnings("unchecked")
+ OOCStream.QueueCallback<IndexedMatrixValue> sub =
+ (OOCStream.QueueCallback<IndexedMatrixValue>) group.getCallback(gi);
+ try(sub) {
+ IndexedMatrixValue imv = sub.get();
+ values.add(imv);
+ totalSize += ((MatrixBlock) imv.getValue()).getExactSerializedSize();
+ }
+ }
+ OOCCacheManager.putRaw(key, values, totalSize);
+ }
+ catch(IllegalStateException putEx) {
+ // Another downstream stream may have re-materialized the same entry first.
+ OOCCacheManager.getCache().addReference(key);
+ }
+ }
+ }
+
private String getCtxMsg() {
StackTraceElement[] st = new Exception().getStackTrace();
// Skip the first few frames (constructor, createWritableStream, etc.)
@@ -687,7 +733,7 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
if(groupIdx > 0)
continue; // only replay grouped blocks once at the base index
- BlockKey replayKey = (groupSize > 1 && groupIdx == 0) ? new BlockKey(_streamId, idx) : getBlockKey(i);
+ BlockKey replayKey = (groupSize > 1 && groupIdx == 0) ? getEntryBlockKey(idx) : getBlockKey(i);
OOCCacheManager.requestBlock(replayKey).whenComplete((cb, r) -> {
if(r != null) {
subscriber.accept(OOCStream.eos(DMLRuntimeException.of(r)));
@@ -697,7 +743,6 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
synchronized(CachingStream.this) {
if(_index != null) {
if(cb instanceof OOCStream.GroupQueueCallback<?> && groupSize > 1) {
- @SuppressWarnings("unchecked")
OOCStream.GroupQueueCallback<IndexedMatrixValue> group =
(OOCStream.GroupQueueCallback<IndexedMatrixValue>) cb;
for(int gi = 0; gi < groupSize; gi++) {
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java
index fbc7d64223..c0604d017d 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java
@@ -30,6 +30,7 @@ public final class BlockEntry {
private volatile BlockState _state;
private Object _data;
private int _retainHintCount;
+ private int _referenceCount; // The number of references from different managing instances (e.g. CachingStream)
BlockEntry(BlockKey key, long size, Object data) {
this._key = key;
@@ -38,6 +39,7 @@ public final class BlockEntry {
this._state = BlockState.HOT;
this._data = data;
this._retainHintCount = 0;
+ this._referenceCount = 1;
}
public BlockKey getKey() {
@@ -84,6 +86,14 @@ public final class BlockEntry {
return _pinCount > 0;
}
+ synchronized int addReference() {
+ return ++_referenceCount;
+ }
+
+ synchronized int forget() {
+ return --_referenceCount;
+ }
+
synchronized void setState(BlockState state) {
_state = state;
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheScheduler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheScheduler.java
index 9cc108db5e..dbbd73d53a 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheScheduler.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheScheduler.java
@@ -103,6 +103,14 @@ public interface OOCCacheScheduler {
BlockEntry putAndPinSourceBacked(BlockKey key, Object data, long size,
OOCIOHandler.SourceBlockDescriptor descriptor);
+ /**
+ * Notifies the cache that there is another reference to the same block key.
+ * This will prevent forget(key) from removing the block from cache.
+ * A block will only be forgotten after all referencing instances called forget(key).
+ * @param key
+ */
+ void addReference(BlockKey key);
+
/**
* Forgets a block from the cache.
* @param key the associated key of the block
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCLRUCacheScheduler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCLRUCacheScheduler.java
index a204cd16db..cc7aa7bcd1 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCLRUCacheScheduler.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCLRUCacheScheduler.java
@@ -19,6 +19,7 @@
package org.apache.sysds.runtime.ooc.cache;
+import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.sysds.api.DMLScript;
@@ -291,6 +292,18 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
return put(key, data, size, true, descriptor);
}
+ @Override
+ public void addReference(BlockKey key) {
+ synchronized(this) {
+ BlockEntry entry = _cache.get(key);
+ if(entry == null)
+ entry = _evictionCache.get(key);
+ if(entry == null)
+ throw new IllegalArgumentException("Could not find requested block with key " + key);
+ entry.addReference();
+ }
+ }
+
private BlockEntry put(BlockKey key, Object data, long size, boolean pin, OOCIOHandler.SourceBlockDescriptor descriptor) {
if (!this._running)
throw new IllegalStateException();
@@ -324,14 +337,34 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
public void forget(BlockKey key) {
if (!this._running)
return;
+ final MutableObject<BlockEntry> mEntry = new MutableObject<>();
BlockEntry entry;
boolean shouldScheduleDeletion = false;
long cacheSizeDelta = 0;
synchronized(this) {
- entry = _cache.remove(key);
+ _cache.compute(key, (k, e) -> {
+ if(e == null)
+ return null;
+ if(e.forget() == 0) {
+ mEntry.setValue(e);
+ return null;
+ }
+ return e;
+ });
- if (entry == null)
- entry = _evictionCache.remove(key);
+ if (mEntry.getValue() == null) {
+ _evictionCache.compute(key, (k, e) -> {
+ if(e == null)
+ return null;
+ if(e.forget() == 0) {
+ mEntry.setValue(e);
+ return null;
+ }
+ return e;
+ });
+ }
+
+ entry = mEntry.getValue();
if (entry != null) {
synchronized(entry) {