diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java
new file mode 100644
index 00000000000..47e08202dfb
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.ooc.memory;
+
+import org.apache.sysds.runtime.ooc.cache.OOCFuture;
+
+public final class ReservationBudget implements MemoryAllowance, AutoCloseable {
+	private final MemoryAllowance _parent;
+	private long _outstanding;
+	private long _available;
+	private boolean _closed;
+
+	public ReservationBudget(MemoryAllowance parent, long bytes) {
+		if(parent == null)
+			throw new NullPointerException("parent");
+		if(bytes < 0)
+			throw new IllegalArgumentException("Budget must not be negative: " + bytes);
+		_parent = parent;
+		_outstanding = bytes;
+		_available = bytes;
+	}
+
+	@Override
+	public synchronized boolean tryReserve(long bytes) {
+		checkNonNegative(bytes);
+		if(bytes == 0)
+			return true;
+		if(_closed || _available < bytes)
+			return false;
+		_available -= bytes;
+		return true;
+	}
+
+	@Override
+	public void reserveBlocking(long bytes) {
+		if(!tryReserve(bytes))
+			throw insufficientBudget(bytes);
+	}
+
+	@Override
+	public OOCFuture<Void> reserveAsync(long bytes) {
+		return tryReserve(bytes) ? OOCFuture.completed(null) : OOCFuture.failed(insufficientBudget(bytes));
+	}
+
+	@Override
+	public void release(long bytes) {
+		checkNonNegative(bytes);
+		if(bytes == 0)
+			return;
+		synchronized(this) {
+			long used = _outstanding - _available;
+			if(bytes > used)
+				throw new IllegalStateException("Cannot release " + bytes + " bytes from a budget using " + used);
+			_outstanding -= bytes;
+		}
+		_parent.release(bytes);
+	}
+
+	@Override
+	public synchronized long getUsedMemory() {
+		return _outstanding - _available;
+	}
+
+	@Override
+	public synchronized long getGrantedMemory() {
+		return _outstanding;
+	}
+
+	@Override
+	public synchronized long getTargetMemory() {
+		return _outstanding;
+	}
+
+	@Override
+	public void setTargetMemory(long targetMemory) {
+		throw new UnsupportedOperationException("Reservation budgets have a fixed target");
+	}
+
+	@Override
+	public void shutdown() {
+		close();
+	}
+
+	@Override
+	public synchronized boolean isShutdown() {
+		return _closed || _parent.isShutdown();
+	}
+
+	@Override
+	public void close() {
+		long released;
+		synchronized(this) {
+			if(_closed)
+				return;
+			_closed = true;
+			released = _available;
+			_available = 0;
+			_outstanding -= released;
+		}
+		if(released > 0)
+			_parent.release(released);
+	}
+
+	private synchronized IllegalStateException insufficientBudget(long bytes) {
+		return new IllegalStateException(
+			"Cannot reserve " + bytes + " bytes from a budget with " + _available + " bytes available");
+	}
+
+	private static void checkNonNegative(long bytes) {
+		if(bytes < 0)
+			throw new IllegalArgumentException("Bytes must not be negative: " + bytes);
+	}
+}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/stream/AllocatedOOCStream.java b/src/main/java/org/apache/sysds/runtime/ooc/stream/AllocatedOOCStream.java
new file mode 100644
index 00000000000..6d3a7c5936f
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/ooc/stream/AllocatedOOCStream.java
@@ -0,0 +1,278 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.ooc.stream;
+
+import java.util.function.ToLongFunction;
+
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue;
+import org.apache.sysds.runtime.ooc.cache.OOCFuture;
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+import org.apache.sysds.runtime.ooc.memory.ReservationBudget;
+
+public final class AllocatedOOCStream<T> extends SubscribableTaskQueue<T> {
+	private final OOCStream<T> _source;
+	private final MemoryAllowance _allowance;
+	private final ToLongFunction<T> _reservationSize;
+	private volatile DMLRuntimeException _failure;
+	private int _pendingReservations;
+	private boolean _sourceComplete;
+	private boolean _outputClosed;
+
+	public AllocatedOOCStream(OOCStream<T> source, MemoryAllowance allowance, ToLongFunction<T> reservationSize) {
+		_source = source;
+		_allowance = allowance;
+		_reservationSize = reservationSize;
+		setData(source.getData());
+		source.setSubscriber(this::admit);
+	}
+
+	public static ReservationBudget detachBudget(OOCStream.QueueCallback<?> callback) {
+		return callback instanceof BudgetedQueueCallback<?> budgeted ? budgeted.detachBudget() : null;
+	}
+
+	private void admit(OOCStream.QueueCallback<T> callback) {
+		if(callback.isFailure()) {
+			try(callback) {
+				callback.get();
+			}
+			catch(DMLRuntimeException failure) {
+				fail(failure);
+			}
+			finishSource();
+			return;
+		}
+		if(callback.isEos()) {
+			callback.close();
+			finishSource();
+			return;
+		}
+		if(_failure != null) {
+			callback.close();
+			return;
+		}
+		try(callback) {
+			long bytes = _reservationSize.applyAsLong(callback.get());
+			if(bytes < 0)
+				throw new IllegalArgumentException("Cannot reserve negative bytes: " + bytes);
+			if(bytes == 0) {
+				enqueueOwned(callback.keepOpen(), null);
+				return;
+			}
+			if(_allowance.tryReserve(bytes)) {
+				enqueueOwned(callback.keepOpen(), new ReservationBudget(_allowance, bytes));
+				return;
+			}
+			retainUntilAllocated(callback, bytes);
+		}
+		catch(RuntimeException error) {
+			fail(DMLRuntimeException.of(error));
+		}
+	}
+
+	private void retainUntilAllocated(OOCStream.QueueCallback<T> callback, long bytes) {
+		OOCStream.QueueCallback<T> retained = callback.keepOpen();
+		OOCFuture<Void> reservation;
+		synchronized(this) {
+			_pendingReservations++;
+		}
+		try {
+			reservation = _allowance.reserveAsync(bytes);
+		}
+		catch(RuntimeException error) {
+			try {
+				retained.close();
+			}
+			finally {
+				releasePendingReservation();
+			}
+			throw error;
+		}
+		reservation.whenComplete((ignored, error) -> {
+			try {
+				if(error != null) {
+					fail(DMLRuntimeException.of(error));
+					retained.close();
+				}
+				else if(_failure != null) {
+					_allowance.release(bytes);
+					retained.close();
+				}
+				else
+					enqueueOwned(retained, new ReservationBudget(_allowance, bytes));
+			}
+			catch(RuntimeException completionError) {
+				fail(DMLRuntimeException.of(completionError));
+			}
+			finally {
+				releasePendingReservation();
+			}
+		});
+	}
+
+	private void enqueueOwned(OOCStream.QueueCallback<T> callback, ReservationBudget budget) {
+		OOCStream.QueueCallback<T> output = budget == null ? callback : new BudgetedQueueCallback<>(callback, budget);
+		try {
+			enqueue(output);
+		}
+		catch(RuntimeException error) {
+			output.close();
+			throw error;
+		}
+	}
+
+	private boolean fail(DMLRuntimeException failure) {
+		synchronized(this) {
+			if(_failure != null)
+				return false;
+			_failure = failure;
+		}
+		super.propagateFailure(failure);
+		return true;
+	}
+
+	private void releasePendingReservation() {
+		boolean close;
+		synchronized(this) {
+			if(_pendingReservations <= 0)
+				throw new IllegalStateException("Pending reservation count underflow");
+			_pendingReservations--;
+			close = _sourceComplete && _pendingReservations == 0 && !_outputClosed;
+			if(close)
+				_outputClosed = true;
+		}
+		if(close)
+			closeInput();
+	}
+
+	private void finishSource() {
+		boolean close;
+		synchronized(this) {
+			if(_sourceComplete)
+				return;
+			_sourceComplete = true;
+			close = _pendingReservations == 0 && !_outputClosed;
+			if(close)
+				_outputClosed = true;
+		}
+		if(close)
+			closeInput();
+	}
+
+	@Override
+	public void propagateFailure(DMLRuntimeException failure) {
+		if(fail(failure))
+			_source.propagateFailure(failure);
+	}
+
+	private static final class BudgetedQueueCallback<T> implements OOCStream.QueueCallback<T> {
+		private final OOCStream.QueueCallback<T> _callback;
+		private final BudgetedQueueCallback<T> _budgetOwner;
+		private ReservationBudget _budget;
+		private int _budgetReferences;
+		private boolean _closed;
+
+		private BudgetedQueueCallback(OOCStream.QueueCallback<T> callback, ReservationBudget budget) {
+			_callback = callback;
+			_budgetOwner = this;
+			_budget = budget;
+			_budgetReferences = 1;
+		}
+
+		private BudgetedQueueCallback(OOCStream.QueueCallback<T> callback, BudgetedQueueCallback<T> budgetOwner) {
+			_callback = callback;
+			_budgetOwner = budgetOwner;
+		}
+
+		private synchronized ReservationBudget detachBudget() {
+			if(_closed)
+				throw new IllegalStateException("Cannot detach from a closed callback");
+			return _budgetOwner.takeBudget();
+		}
+
+		private synchronized ReservationBudget takeBudget() {
+			ReservationBudget budget = _budget;
+			_budget = null;
+			return budget;
+		}
+
+		@Override
+		public T get() {
+			return _callback.get();
+		}
+
+		@Override
+		public synchronized OOCStream.QueueCallback<T> keepOpen() {
+			if(_closed)
+				throw new IllegalStateException("Cannot keep open a closed callback");
+			OOCStream.QueueCallback<T> retained = _callback.keepOpen();
+			_budgetOwner.retainBudget();
+			return new BudgetedQueueCallback<>(retained, _budgetOwner);
+		}
+
+		private synchronized void retainBudget() {
+			_budgetReferences++;
+		}
+
+		@Override
+		public void close() {
+			synchronized(this) {
+				if(_closed)
+					return;
+				_closed = true;
+			}
+			try {
+				_callback.close();
+			}
+			finally {
+				_budgetOwner.releaseBudget();
+			}
+		}
+
+		private void releaseBudget() {
+			ReservationBudget budget = null;
+			synchronized(this) {
+				_budgetReferences--;
+				if(_budgetReferences == 0) {
+					budget = _budget;
+					_budget = null;
+				}
+			}
+			if(budget != null)
+				budget.close();
+		}
+
+		@Override
+		public void fail(DMLRuntimeException failure) {
+			_callback.fail(failure);
+		}
+
+		@Override
+		public boolean isEos() {
+			return _callback.isEos();
+		}
+
+		@Override
+		public boolean isFailure() {
+			return _callback.isFailure();
+		}
+	}
+}
diff --git a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
index 979498502d1..c75458eef38 100644
--- a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
+++ b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
@@ -19,6 +19,7 @@
 
 package org.apache.sysds.test.component.ooc.memory;
 
+import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
 import org.apache.sysds.runtime.functionobjects.Plus;
 import org.apache.sysds.runtime.instructions.ooc.OOCInstruction;
@@ -36,7 +37,9 @@
 import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback;
 import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
 import org.apache.sysds.runtime.ooc.memory.MemoryBroker;
+import org.apache.sysds.runtime.ooc.memory.ReservationBudget;
 import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance;
+import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream;
 import org.junit.Assert;
 import org.junit.Test;
 import scala.Tuple3;
@@ -112,6 +115,66 @@ public void testReservationWaiters() throws Exception {
 		}
 	}
 
+	@Test
+	public void testAllocatedStreamReservations() {
+		GlobalMemoryBroker broker = new GlobalMemoryBroker(100);
+		SyncMemoryAllowance allowance = new SyncMemoryAllowance(broker);
+		SubscribableTaskQueue<Integer> source = new SubscribableTaskQueue<>();
+		AllocatedOOCStream<Integer> allocated = new AllocatedOOCStream<>(source, allowance, value -> 60);
+		try {
+			allowance.reserveBlocking(100);
+			source.enqueue(1);
+			Assert.assertEquals(100, allowance.getUsedMemory());
+
+			allowance.release(100);
+			OOCStream.QueueCallback<Integer> first = allocated.dequeueCB();
+			ReservationBudget budget = AllocatedOOCStream.detachBudget(first);
+			Assert.assertNotNull(budget);
+			first.close();
+			Assert.assertEquals(60, allowance.getUsedMemory());
+			budget.reserveBlocking(20);
+			budget.release(20);
+			Assert.assertEquals(40, allowance.getUsedMemory());
+			budget.close();
+			Assert.assertEquals(0, allowance.getUsedMemory());
+
+			source.enqueue(2);
+			OOCStream.QueueCallback<Integer> second = allocated.dequeueCB();
+			OOCStream.QueueCallback<Integer> retained = second.keepOpen();
+			second.close();
+			Assert.assertEquals(60, allowance.getUsedMemory());
+			retained.close();
+			Assert.assertEquals(0, allowance.getUsedMemory());
+			source.closeInput();
+			Assert.assertNull(allocated.dequeueCB());
+		}
+		finally {
+			if(allowance.getUsedMemory() > 0)
+				allowance.release(allowance.getUsedMemory());
+			allowance.destroy();
+		}
+	}
+
+	@Test
+	public void testAllocatedStreamFailure() {
+		GlobalMemoryBroker broker = new GlobalMemoryBroker(100);
+		SyncMemoryAllowance allowance = new SyncMemoryAllowance(broker);
+		SubscribableTaskQueue<Integer> source = new SubscribableTaskQueue<>();
+		new AllocatedOOCStream<>(source, allowance, value -> 60);
+		try {
+			allowance.reserveBlocking(100);
+			source.enqueue(1);
+			source.propagateFailure(new DMLRuntimeException("injected failure"));
+			allowance.release(100);
+			Assert.assertEquals(0, allowance.getUsedMemory());
+		}
+		finally {
+			if(allowance.getUsedMemory() > 0)
+				allowance.release(allowance.getUsedMemory());
+			allowance.destroy();
+		}
+	}
+
 	public void test(boolean optimal, int nWarmup, int nMeasure) {
 		//DMLScript.OOC_STATISTICS = true;
 		long millis;