add placeholder logic for iotrace

Author: j143

src/main/java/org/apache/sysds/runtime/controlprogram/caching/UnifiedMemoryManager.java

@@ -112,6 +112,7 @@ public class UnifiedMemoryManager
 	// Prescient policy
 	private static PrescientPolicy _prescientPolicy;
 	private static IOTrace _ioTrace;
+	private static long _currentTime = 0;
 
 	// Pinned size of physical memory. Starts from 0 for each operation. Max is 70% of heap
 	// This increases only if the input is not present in the cache and read from FS/rdd/fed/gpu
@@ -207,6 +208,66 @@ public static void cleanup() {
 		_pinnedVirtualMemSize = 0;
 	}
 
+	/**
+	 * Sets the I/O trace for the prescient policy.
+	 * This should be called once by the ExecutionContext after trace generation.
+	 * @param trace The generated IOTrace
+	 */
+	public static void setTrace(IOTrace trace) {
+		_ioTrace = trace;
+		if (_evictionPolicy instanceof PrescientPolicy) {
+			_prescientPolicy = (PrescientPolicy) _evictionPolicy;
+			_prescientPolicy.setTrace(_ioTrace);
+		}
+		else {
+			// Optional: Log a warning if the trace is set but the policy isn't Prescient
+			// LOG.warn("IOTrace was provided, but eviction policy is not prescient!");
+		}
+	}
+
+	/**
+	 * Updates the UMM's logical time.
+	 * This should be called by the ExecutionContext *before* each instruction.
+	 * @param logicalTime The new logical time
+	 */
+	public static void updateTime(long logicalTime) {
+		_currentTime = logicalTime;
+		prefetch();
+	}
+
+	/**
+	 * Prefetches blocks that will be needed soon, based on the I/O trace.
+	 */
+	private static void prefetch() {
+		if (_ioTrace == null || _prescientPolicy == null) {
+			return; // No trace or policy, cannot prefetch
+		}
+
+		// Get the list of blocks to prefetch from our policy
+		List<String> blocksToPrefetch = _prescientPolicy.getBlocksToPrefetch(_currentTime);
+
+		// A real implementation MUST use an asynchronous thread pool
+		// (e.g., from _fClean) to load these blocks without blocking the main thread.
+
+		for (String blockID : blocksToPrefetch) {
+			synchronized (_mQueue) {
+				// Check again inside lock if block was already loaded or pinned
+				if (_mQueue.containsKey(blockID) || _pinnedEntries.contains(blockID)) {
+					continue; // Already in memory
+				}
+			}
+
+			// --- This is a simplified version for now ---
+			// TODO: Submit an async prefetch task to _fClean's thread pool
+			// The task should: 1. Get block size (from metadata)
+			//                 2. Call makeSpace(blockSize)
+			//                 3. Load block from disk
+			//                 4. Add block to _mQueue (synchronized)
+
+			System.out.println("UMM PREFETCH [T="+_currentTime+"]: Planning to prefetch " + blockID);
+		}
+	}
+
 	/**
 	 * Print current status of UMM, including all entries.
 	 * NOTE: use only for debugging or testing.
@@ -312,10 +373,35 @@ public static int makeSpace(long reqSpace) {
 			synchronized(_mQueue) {
 				// Evict blobs to make room (by default FIFO)
 				while (getUMMFree() < reqSpace && !_mQueue.isEmpty()) {
-					//remove first unpinned entry from eviction queue
-					var entry = _mQueue.removeFirstUnpinned(_pinnedEntries);
-					String ftmp = entry.getKey();
-					ByteBuffer bb = entry.getValue();
+					// --- NEW PRESCIENT LOGIC ---
+					String ftmp; // Block ID / filename to evict
+
+					if (_prescientPolicy != null && _ioTrace != null) {
+						// Use prescient policy to find the best block to evict
+						ftmp = _prescientPolicy.evict(_mQueue.keySet(), _pinnedEntries, _currentTime);
+					} else {
+						// Fallback to default LRU if prescient policy isn't set or has no trace
+						var entry = _mQueue.removeFirstUnpinned(_pinnedEntries);
+						ftmp = (entry != null) ? entry.getKey() : null;
+					}
+
+					if (ftmp == null) {
+						// Policy couldn't find a block to evict (e.g., all are pinned)
+						if(!_pinnedEntries.containsAll(_mQueue.keySet())) {
+							// This case should ideally not be reached if unpinned blocks exist
+							throw new DMLRuntimeException("UMM: Eviction policy failed to find a candidate.");
+						}
+						// If we are here, all blocks are pinned, and we cannot make space.
+						// The original exception will be thrown later.
+						break; // Exit the while loop
+					}
+
+					// Remove the chosen block from the queue
+					ByteBuffer bb = _mQueue.remove(ftmp);
+//					//remove first unpinned entry from eviction queue
+//					var entry = _mQueue.removeFirstUnpinned(_pinnedEntries);
+//					String ftmp = entry.getKey();
+//					ByteBuffer bb = entry.getValue();
 
 					if(bb != null) {
 						// Wait for pending serialization

src/main/java/org/apache/sysds/runtime/controlprogram/caching/prescientbuffer/IOTrace.java

@@ -42,9 +42,8 @@ public IOTrace() {
 	/**
 	 * Access to the block at a current time
 	 */
-	public void recordAccess(String blockID) {
-		_trace.computeIfAbsent(blockID, k -> new ArrayList<>()).add(_currentTime);
-		_currentTime++;
+	public void recordAccess(String blockID, long logicalTime) {
+		_trace.computeIfAbsent(blockID, k -> new ArrayList<>()).add(logicalTime);
 	}
 
 	/**

src/main/java/org/apache/sysds/runtime/controlprogram/caching/prescientbuffer/IOTraceGenerator.java

@@ -145,7 +145,7 @@ private static void processInstruction(Instruction inst, ExecutionContext ec, IO
 				for (long j = 1; j <= numColBlocks; j++) {
 
 					String blockID = createBlockID(fname, i, j);
-					trace.recordAccess(blockID);
+					trace.recordAccess(blockID, logicalTime);
 				}
 			}
 		}

src/main/java/org/apache/sysds/runtime/controlprogram/caching/prescientbuffer/PrescientPolicy.java

@@ -22,6 +22,7 @@
 import org.apache.sysds.runtime.controlprogram.caching.EvictionPolicy;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -32,6 +33,7 @@ public class PrescientPolicy implements EvictionPolicy {
 
 	// Map of block ID, access times
 	private final Map<String, Long> accessTimeMap = new HashMap<>();
+	private IOTrace _trace;
 
 	// register blocks with access time
 	public void setAccessTime(String blockId, long accessTime) {
@@ -66,4 +68,41 @@ public String selectBlockForEviction(Set<String> candidates) {
 		return selected;
 	}
 
+	/**
+	 * Called by UMM's makeSpace() to decide which block to evict.
+	 * * @param cache The set of all block IDs currently in the buffer
+	 * @param pinned The list of all block IDs that are pinned
+	 * @param currentTime The current logical time
+	 * @return The block ID to evict
+	 */
+	public String evict(Set<String> cache, List<String> pinned, long currentTime) {
+		// TODO: Implement "evict-furthest-in-future" logic here
+		// 1. Iterate through every 'blockID' in 'cache'
+		// 2. If 'blockID' is in 'pinned', ignore it.
+		// 3. Use '_trace.getAccessTime(blockID)' to find its next access time > currentTime
+		// 4. The block with the (largest next access time) or (no future access) is the winner.
+		// 5. Return the winner's blockID.
+
+		return null; // Placeholder
+	}
+
+	/**
+	 * Called by UMM's prefetch() to decide which blocks to load.
+	 * * @param currentTime The current logical time
+	 * @return A list of block IDs to prefetch
+	 */
+	public List<String> getBlocksToPrefetch(long currentTime) {
+		// TODO: Implement prefetch logic here
+		// 1. Define a "prefetch window" (e.g., time T+1 to T+5)
+		// 2. Iterate through all blocks in '_trace.getTrace()'
+		// 3. Check if a block has an access time within that window
+		// 4. If yes, add it to a list.
+		// 5. Return the list of blocks.
+
+		return java.util.Collections.emptyList(); // Placeholder
+	}
+
+	public void setTrace(IOTrace ioTrace) {
+		_trace = ioTrace;
+	}
 }

src/main/java/org/apache/sysds/runtime/controlprogram/context/ExecutionContext.java

@@ -21,6 +21,7 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.yarn.webapp.hamlet2.HamletSpec;
 import org.apache.sysds.api.DMLScript;
 import org.apache.sysds.common.Types;
 import org.apache.sysds.common.Types.FileFormat;
@@ -37,6 +38,7 @@
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject.UpdateType;
 import org.apache.sysds.runtime.controlprogram.caching.TensorObject;
+import org.apache.sysds.runtime.controlprogram.caching.prescientbuffer.IOTrace;
 import org.apache.sysds.runtime.controlprogram.federated.MatrixLineagePair;
 import org.apache.sysds.runtime.controlprogram.paramserv.homomorphicEncryption.SEALClient;
 import org.apache.sysds.runtime.data.TensorBlock;
@@ -93,6 +95,19 @@ public class ExecutionContext {
 	//parfor temporary functions (created by eval)
 	protected Set<String> _fnNames;
 
+	private IOTrace _ioTrace;
+
+	public IOTrace getIOTrace() {
+		if (_ioTrace == null) {
+			_ioTrace = new IOTrace();
+		}
+		return _ioTrace;
+	}
+
+	public void setIOTrace(IOTrace ioTrace) {
+		_ioTrace = ioTrace;
+	}
+
 	/**
 	 * List of {@link GPUContext}s owned by this {@link ExecutionContext}
 	 */