[SYSTEMDS-3946] Enable sending of large (>2GiB) federated requests and responses

Closes #2591.
Closes #2496.

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkDecoder.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkDecoder.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkDecoder.java
new file mode 100644
index 0000000..df00e2a
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkDecoder.java
@@ -0,0 +1,241 @@
+/*
+ * 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.controlprogram.federated;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectStreamClass;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.apache.sysds.runtime.util.CommonThreadPool;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.MessageToMessageDecoder;
+
+public class FederatedChunkDecoder extends MessageToMessageDecoder<ByteBuf> {
+	private static final Object END_OF_STREAM = new Object();
+	// stop reading at QUEUE_DEPTH, resume at half: gap avoids autoRead thrash
+	private static final int LOW_WATERMARK = FederatedChunkProtocol.QUEUE_DEPTH / 2;
+
+	private final BlockingQueue<Object> _payloads = new LinkedBlockingQueue<>();
+	private boolean _started;
+	private volatile boolean _throttled;
+
+	@Override
+	protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out) {
+		startReader(ctx);
+		byte type = buf.readByte();
+		int len = buf.readInt();
+		switch(type) {
+			case FederatedChunkProtocol.TYPE_DATA:
+				_payloads.add(readBytes(buf, len));
+				break;
+			case FederatedChunkProtocol.TYPE_END:
+				_payloads.add(END_OF_STREAM);
+				break;
+			case FederatedChunkProtocol.TYPE_ERROR:
+				_payloads.add(new IOException(buf.toString(buf.readerIndex(), len, StandardCharsets.UTF_8)));
+				break;
+			default:
+				_payloads.add(new IOException("Unknown federated chunk frame type: " + type));
+				break;
+		}
+		if(_payloads.size() >= FederatedChunkProtocol.QUEUE_DEPTH) {
+			_throttled = true;
+			ctx.channel().config().setAutoRead(false);
+		}
+	}
+
+	@Override
+	public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+		_payloads.add(new IOException("Channel closed before the federated chunk stream ended."));
+		super.channelInactive(ctx);
+	}
+
+	/**
+	 * Start the deserializer on a pool thread, at most once per channel.
+	 *
+	 * @param ctx handler context
+	 */
+	private void startReader(ChannelHandlerContext ctx) {
+		if(_started)
+			return;
+		_started = true;
+		CommonThreadPool.getDynamicPool().execute(() -> runDeserializer(ctx));
+	}
+
+	/**
+	 * Read one object from the queued payloads and fire it up the pipeline. A failure is fired as an exception on the
+	 * event loop instead.
+	 *
+	 * @param ctx handler context
+	 */
+	private void runDeserializer(ChannelHandlerContext ctx) {
+		try(PayloadInputStream in = new PayloadInputStream(this, ctx);
+			ObjectInputStream ois = getObjectInputStream(in)) {
+			Object msg = ois.readObject();
+			in.skipToEndOfStream();
+			ctx.channel().eventLoop().execute(() -> ctx.fireChannelRead(msg));
+		}
+		catch(Throwable t) {
+			ctx.channel().eventLoop().execute(() -> ctx.fireExceptionCaught(t));
+		}
+	}
+
+	/**
+	 * Take the next queued payload, blocking until one arrives.
+	 *
+	 * @return payload bytes, the end of stream marker or a queued failure
+	 * @throws InterruptedException on interrupt while the queue is empty
+	 */
+	private Object nextPayload() throws InterruptedException {
+		return _payloads.take();
+	}
+
+	/**
+	 * Re-enable channel reads once the payload queue has drained to the low watermark.
+	 *
+	 * @param ctx handler context
+	 */
+	private void resumeReadingIfDrained(ChannelHandlerContext ctx) {
+		if(_throttled && _payloads.size() <= LOW_WATERMARK) {
+			_throttled = false;
+			ctx.channel().eventLoop().execute(() -> ctx.channel().config().setAutoRead(true));
+		}
+	}
+
+	/**
+	 * Create an object input stream using the system class loader.
+	 *
+	 * @param in stream of payload bytes
+	 * @return object input stream
+	 * @throws IOException on stream header failure
+	 */
+	private static ObjectInputStream getObjectInputStream(InputStream in) throws IOException {
+		return new ObjectInputStream(in) {
+			@Override
+			protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
+				try {
+					return Class.forName(desc.getName(), false, ClassLoader.getSystemClassLoader());
+				}
+				catch(ClassNotFoundException e) {
+					return super.resolveClass(desc);
+				}
+			}
+		};
+	}
+
+	private static byte[] readBytes(ByteBuf buf, int len) {
+		byte[] bytes = new byte[len];
+		buf.readBytes(bytes);
+		return bytes;
+	}
+
+	private static final class PayloadInputStream extends InputStream {
+		private static final byte[] EMPTY = new byte[0];
+
+		private final FederatedChunkDecoder _decoder;
+		private final ChannelHandlerContext _ctx;
+		private byte[] _current = EMPTY;
+		private int _pos;
+		private boolean _eof;
+
+		PayloadInputStream(FederatedChunkDecoder decoder, ChannelHandlerContext ctx) {
+			_decoder = decoder;
+			_ctx = ctx;
+		}
+
+		@Override
+		public int read() throws IOException {
+			if(!ensureCurrent())
+				return -1;
+			return _current[_pos++] & 0xff;
+		}
+
+		@Override
+		public int read(byte[] b, int off, int len) throws IOException {
+			if(!ensureCurrent())
+				return -1;
+			int n = Math.min(len, _current.length - _pos);
+			System.arraycopy(_current, _pos, b, off, n);
+			_pos += n;
+			return n;
+		}
+
+		/**
+		 * Take the next payload and resume reading if the queue has drained.
+		 *
+		 * @return payload bytes, the end of stream marker or a queued failure
+		 * @throws IOException on interrupt
+		 */
+		private Object take() throws IOException {
+			try {
+				Object next = _decoder.nextPayload();
+				_decoder.resumeReadingIfDrained(_ctx);
+				return next;
+			}
+			catch(InterruptedException e) {
+				Thread.currentThread().interrupt();
+				throw new IOException(e);
+			}
+		}
+
+		/**
+		 * Advance to the next payload when the current one is exhausted. A queued failure is rethrown as an
+		 * IOException.
+		 *
+		 * @return true if bytes are available, false at end of stream
+		 * @throws IOException on a queued failure or on interrupt
+		 */
+		private boolean ensureCurrent() throws IOException {
+			while(_pos == _current.length) {
+				if(_eof)
+					return false;
+				Object next = take();
+				if(next == END_OF_STREAM) {
+					_eof = true;
+					return false;
+				}
+				if(next instanceof Throwable)
+					throw new IOException((Throwable) next);
+				_current = (byte[]) next;
+				_pos = 0;
+			}
+			return true;
+		}
+
+		/**
+		 * Consume the remaining frames up to and including the end of stream marker.
+		 *
+		 * @throws IOException on a queued failure or on interrupt
+		 */
+		void skipToEndOfStream() throws IOException {
+			while(!_eof) {
+				_pos = _current.length;
+				ensureCurrent();
+			}
+		}
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkEncoder.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkEncoder.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkEncoder.java
new file mode 100644
index 0000000..71782cb
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkEncoder.java
@@ -0,0 +1,250 @@
+/*
+ * 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.controlprogram.federated;
+
+import java.io.IOException;
+import java.io.ObjectOutputStream;
+import java.io.OutputStream;
+import java.io.Serializable;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+
+import org.apache.sysds.runtime.util.CommonThreadPool;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.stream.ChunkedInput;
+import io.netty.handler.stream.ChunkedWriteHandler;
+
+public final class FederatedChunkEncoder {
+	private FederatedChunkEncoder() {
+	}
+
+	/**
+	 * Create a chunked input over the serialized form of the given message. Serialization starts immediately on a pool
+	 * thread.
+	 *
+	 * @param msg       message to serialize
+	 * @param chunkSize payload bytes per data frame
+	 * @param alloc     allocator for the frame buffers
+	 * @param writer    write handler resumed when a frame becomes available
+	 * @return chunked input over the serialized message
+	 */
+	public static ChunkedInput<ByteBuf> chunkedInput(Serializable msg, int chunkSize, ByteBufAllocator alloc,
+		ChunkedWriteHandler writer) {
+		return new SerializedChunks(msg, chunkSize, alloc, writer);
+	}
+
+	private static final class SerializedChunks implements ChunkedInput<ByteBuf> {
+		private final BlockingQueue<ByteBuf> _frames = new ArrayBlockingQueue<>(FederatedChunkProtocol.QUEUE_DEPTH);
+		private final ByteBufAllocator _alloc;
+		private final ChunkedWriteHandler _writer;
+		private volatile boolean _closed;
+		private boolean _done;
+
+		/**
+		 * Create the frame source and start serialization on a pool thread.
+		 *
+		 * @param msg       message to serialize
+		 * @param chunkSize payload bytes per data frame
+		 * @param alloc     allocator for the frame buffers
+		 * @param writer    write handler resumed when a frame becomes available
+		 */
+		SerializedChunks(Serializable msg, int chunkSize, ByteBufAllocator alloc, ChunkedWriteHandler writer) {
+			_alloc = alloc;
+			_writer = writer;
+			CommonThreadPool.getDynamicPool().execute(() -> produceFrames(msg, chunkSize));
+		}
+
+		/**
+		 * Serialize the message into data frames followed by an end frame. A failure is turned into an error frame
+		 * instead.
+		 *
+		 * @param msg       message to serialize
+		 * @param chunkSize payload bytes per data frame
+		 */
+		private void produceFrames(Serializable msg, int chunkSize) {
+			try(FrameOutputStream out = new FrameOutputStream(this, _alloc, chunkSize);
+				ObjectOutputStream oos = new ObjectOutputStream(out)) {
+				oos.writeObject(msg);
+				oos.flush();
+				out.flushFrame();
+				enqueueControlFrame(controlFrame(FederatedChunkProtocol.TYPE_END));
+			}
+			catch(Throwable t) {
+				enqueueControlFrame(errorFrame(t));
+			}
+		}
+
+		/**
+		 * Create a header only frame of the given type.
+		 *
+		 * @param type frame type
+		 * @return control frame with an empty payload
+		 */
+		private ByteBuf controlFrame(byte type) {
+			return _alloc.buffer(FederatedChunkProtocol.HEADER_LEN).writeByte(type).writeInt(0);
+		}
+
+		/**
+		 * Create an error frame whose payload is the UTF-8 text of the failure.
+		 *
+		 * @param t serialization failure
+		 * @return error frame
+		 */
+		private ByteBuf errorFrame(Throwable t) {
+			byte[] cause = String.valueOf(t).getBytes(StandardCharsets.UTF_8);
+			return _alloc.buffer(FederatedChunkProtocol.HEADER_LEN + cause.length)
+				.writeByte(FederatedChunkProtocol.TYPE_ERROR).writeInt(cause.length).writeBytes(cause);
+		}
+
+		/**
+		 * Append a frame to the queue and resume the write handler. A frame arriving after close is released instead.
+		 *
+		 * @param frame frame to append
+		 * @throws InterruptedException on interrupt while the queue is full
+		 */
+		void enqueueFrame(ByteBuf frame) throws InterruptedException {
+			if(_closed) {
+				frame.release();
+				return;
+			}
+			_frames.put(frame);
+			_writer.resumeTransfer();
+		}
+
+		/**
+		 * Append a control frame, releasing it on interrupt.
+		 *
+		 * @param frame frame to append
+		 */
+		private void enqueueControlFrame(ByteBuf frame) {
+			try {
+				enqueueFrame(frame);
+			}
+			catch(InterruptedException e) {
+				frame.release();
+				Thread.currentThread().interrupt();
+			}
+		}
+
+		@Override
+		public ByteBuf readChunk(ByteBufAllocator allocator) {
+			if(_done)
+				return null;
+			ByteBuf frame = _frames.poll();
+			if(frame == null)
+				return null;
+			_done = frame.getByte(frame.readerIndex()) != FederatedChunkProtocol.TYPE_DATA;
+			return frame;
+		}
+
+		@Override
+		public ByteBuf readChunk(ChannelHandlerContext ctx) {
+			return readChunk(ctx.alloc());
+		}
+
+		@Override
+		public boolean isEndOfInput() {
+			return _done;
+		}
+
+		@Override
+		public long length() {
+			return -1;
+		}
+
+		@Override
+		public long progress() {
+			return 0;
+		}
+
+		@Override
+		public void close() {
+			_closed = true;
+			ByteBuf frame;
+			while((frame = _frames.poll()) != null)
+				frame.release();
+		}
+	}
+
+	private static final class FrameOutputStream extends OutputStream {
+		private final SerializedChunks _sink;
+		private final ByteBufAllocator _alloc;
+		private final byte[] _buffer;
+		private int _len;
+
+		/**
+		 * Create a stream that buffers written bytes and emits one data frame per full chunk.
+		 *
+		 * @param sink      chunked input receiving the frames
+		 * @param alloc     allocator for the frame buffers
+		 * @param chunkSize payload bytes per data frame
+		 */
+		FrameOutputStream(SerializedChunks sink, ByteBufAllocator alloc, int chunkSize) {
+			_sink = sink;
+			_alloc = alloc;
+			_buffer = new byte[chunkSize];
+		}
+
+		@Override
+		public void write(int b) throws IOException {
+			_buffer[_len++] = (byte) b;
+			if(_len == _buffer.length)
+				flushFrame();
+		}
+
+		@Override
+		public void write(byte[] b, int off, int len) throws IOException {
+			while(len > 0) {
+				int n = Math.min(len, _buffer.length - _len);
+				System.arraycopy(b, off, _buffer, _len, n);
+				_len += n;
+				off += n;
+				len -= n;
+				if(_len == _buffer.length)
+					flushFrame();
+			}
+		}
+
+		/**
+		 * Emit the buffered bytes as one data frame or nothing when the buffer is empty.
+		 *
+		 * @throws IOException on interrupt while the queue is full
+		 */
+		void flushFrame() throws IOException {
+			if(_len == 0)
+				return;
+			ByteBuf frame = _alloc.buffer(FederatedChunkProtocol.HEADER_LEN + _len)
+				.writeByte(FederatedChunkProtocol.TYPE_DATA).writeInt(_len).writeBytes(_buffer, 0, _len);
+			_len = 0;
+			try {
+				_sink.enqueueFrame(frame);
+			}
+			catch(InterruptedException e) {
+				frame.release();
+				Thread.currentThread().interrupt();
+				throw new IOException(e);
+			}
+		}
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkProtocol.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkProtocol.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkProtocol.java
new file mode 100644
index 0000000..9a31bcc
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedChunkProtocol.java
@@ -0,0 +1,65 @@
+/*
+ * 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.controlprogram.federated;
+
+import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
+
+public final class FederatedChunkProtocol {
+	public static final byte TYPE_DATA = 0;
+	public static final byte TYPE_END = 1;
+	public static final byte TYPE_ERROR = 2;
+
+	public static final byte MARKER_OBJECT_ENCODER = 0;
+	public static final byte MARKER_CHUNKED = 1;
+
+	public static final long STREAM_THRESHOLD = 2000L << 20;
+
+	public static final int HEADER_LEN = 5;
+	public static final int DEFAULT_CHUNK_SIZE = 1 << 22;
+	public static final int QUEUE_DEPTH = 16;
+
+	public static final int LENGTH_FIELD_OFFSET = 1;
+	public static final int LENGTH_FIELD_LENGTH = 4;
+	public static final int LENGTH_ADJUSTMENT = 0;
+	public static final int INITIAL_BYTES_TO_STRIP = 0;
+
+	/**
+	 * Get the largest frame a chunk of the given size can produce.
+	 *
+	 * @param chunkSize payload bytes per data frame
+	 * @return frame size including the header
+	 */
+	static int maxChunkFrameLength(int chunkSize) {
+		return chunkSize + HEADER_LEN;
+	}
+
+	/**
+	 * Create a length based frame decoder for the chunked wire format.
+	 *
+	 * @return frame decoder sized for the default chunk size
+	 */
+	static LengthFieldBasedFrameDecoder newChunkFrameDecoder() {
+		return new LengthFieldBasedFrameDecoder(maxChunkFrameLength(DEFAULT_CHUNK_SIZE), LENGTH_FIELD_OFFSET,
+			LENGTH_FIELD_LENGTH, LENGTH_ADJUSTMENT, INITIAL_BYTES_TO_STRIP);
+	}
+
+	private FederatedChunkProtocol() {
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedData.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedData.java
index 3c6e64a..02f4bf9 100644
--- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedData.java
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedData.java
@@ -33,13 +33,16 @@
 import java.util.concurrent.Future;
 
 import io.netty.bootstrap.Bootstrap;
+import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
+import io.netty.channel.ChannelOption;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.EventLoopGroup;
+import io.netty.channel.ChannelFutureListener;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.commons.logging.Log;
@@ -60,11 +63,11 @@
 import org.apache.sysds.runtime.controlprogram.paramserv.NetworkTrafficCounter;
 import org.apache.sysds.runtime.controlprogram.caching.CacheBlock;
 import org.apache.sysds.conf.DMLConfig;
-import io.netty.buffer.ByteBuf;
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.channel.socket.nio.NioSocketChannel;
 import io.netty.handler.codec.serialization.ObjectEncoder;
+import io.netty.handler.stream.ChunkedWriteHandler;
 import io.netty.handler.timeout.ReadTimeoutHandler;
 import io.netty.util.concurrent.DefaultThreadFactory;
 import io.netty.util.concurrent.Promise;
@@ -205,6 +208,7 @@
 				createWorkGroup();
 			b.group(workerGroup);
 			b.channel(NioSocketChannel.class);
+			b.option(ChannelOption.ALLOW_HALF_CLOSURE, true);
 			final DataRequestHandler handler = new DataRequestHandler();
 			// Client Netty
 
@@ -213,7 +217,12 @@
 			ChannelFuture f = b.connect(address).sync();
 			Promise<FederatedResponse> promise = f.channel().eventLoop().newPromise();
 			handler.setPromise(promise);
-			f.channel().writeAndFlush(request);
+			f.channel().writeAndFlush(request).addListener((ChannelFutureListener) future -> {
+				if(!future.isSuccess()) {
+					LOG.error("Federated network write failed: " + future.cause().getMessage());
+					promise.setFailure(future.cause());
+				}
+			});
 
 			return handler.getProm();
 		}
@@ -258,9 +267,11 @@
 					cp.addLast(new ReadTimeoutHandler(timeout));
 
 				compressionStrategy.ifPresent(strategy -> cp.addLast(strategy.left));
-				cp.addLast(FederationUtils.decoder());
+				cp.addLast(new FederatedFormatDecoder());
 				compressionStrategy.ifPresent(strategy -> cp.addLast(strategy.right));
+				cp.addLast(new ChunkedWriteHandler());
 				cp.addLast(new FederatedRequestEncoder());
+				cp.addLast(new FederatedFormatEncoder());
 				cp.addLast(handler);
 			}
 		};
@@ -350,6 +361,15 @@
 			ctx.close();
 		}
 
+		@Override
+		public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+			// Fail (rather than leave hanging) any request whose connection closed before its response
+			// was delivered, so a waiting caller gets an exception instead of blocking until timeout.
+			if(_prom != null && !_prom.isDone())
+				_prom.tryFailure(new IOException("Channel closed before federated response was received"));
+			super.channelInactive(ctx);
+		}
+
 		public Promise<FederatedResponse> getProm() {
 			return _prom;
 		}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatDecoder.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatDecoder.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatDecoder.java
new file mode 100644
index 0000000..7990f92
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatDecoder.java
@@ -0,0 +1,45 @@
+/*
+ * 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.controlprogram.federated;
+
+import java.util.List;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelPipeline;
+import io.netty.handler.codec.ByteToMessageDecoder;
+
+public final class FederatedFormatDecoder extends ByteToMessageDecoder {
+	@Override
+	protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) {
+		if(in.readableBytes() < 1)
+			return;
+		byte marker = in.readByte();
+		ChannelPipeline cp = ctx.pipeline();
+		if(marker == FederatedChunkProtocol.MARKER_CHUNKED) {
+			cp.addAfter(ctx.name(), "FederatedChunkFrameDecoder", FederatedChunkProtocol.newChunkFrameDecoder());
+			cp.addAfter("FederatedChunkFrameDecoder", "FederatedChunkDecoder", new FederatedChunkDecoder());
+		}
+		else {
+			cp.addAfter(ctx.name(), "FederatedObjectDecoder", FederationUtils.decoder());
+		}
+		cp.remove(this);
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatEncoder.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatEncoder.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatEncoder.java
new file mode 100644
index 0000000..06e49be
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedFormatEncoder.java
@@ -0,0 +1,87 @@
+/*
+ * 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.controlprogram.federated;
+
+import java.io.IOException;
+import java.io.Serializable;
+
+import org.apache.sysds.runtime.lineage.LineageCacheConfig.ReuseCacheType;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
+import io.netty.handler.stream.ChunkedWriteHandler;
+
+public class FederatedFormatEncoder extends ChannelOutboundHandlerAdapter {
+	private final int _chunkSize;
+	private final long _streamThreshold;
+
+	public FederatedFormatEncoder() {
+		this(FederatedChunkProtocol.DEFAULT_CHUNK_SIZE, FederatedChunkProtocol.STREAM_THRESHOLD);
+	}
+
+	public FederatedFormatEncoder(int chunkSize, long streamThreshold) {
+		_chunkSize = chunkSize;
+		_streamThreshold = streamThreshold;
+	}
+
+	@Override
+	public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws IOException {
+		if(!(msg instanceof Serializable)) {
+			throw new IOException("Network message must be serializable.");
+		}
+		if(useObjectEncoder(msg)) {
+			ctx.write(markerBuffer(ctx, FederatedChunkProtocol.MARKER_OBJECT_ENCODER), ctx.voidPromise());
+			ctx.write(msg, promise);
+		}
+		else {
+			ctx.write(markerBuffer(ctx, FederatedChunkProtocol.MARKER_CHUNKED), ctx.voidPromise());
+			ctx.write(FederatedChunkEncoder.chunkedInput((Serializable) msg, _chunkSize, ctx.alloc(),
+				ctx.pipeline().get(ChunkedWriteHandler.class)), promise);
+		}
+	}
+
+	/**
+	 * Check whether a message must take the object encoder path.
+	 *
+	 * @param msg outbound message
+	 * @return true if lineage reuse is enabled and the message is a reusable response below the stream threshold
+	 */
+	private boolean useObjectEncoder(Object msg) {
+		if(ReuseCacheType.isNone() || !(msg instanceof FederatedResponse))
+			return false;
+		FederatedResponse resp = (FederatedResponse) msg;
+		// Chunks are not supported by the lineage cache yet; use regular object encoder in this case if
+		// the response size does not exceed the upper limits of the object encoder
+		return resp.isLineageReusable() && resp.estimateSerializationBufferSize() < _streamThreshold;
+	}
+
+	/**
+	 * Create a one byte buffer holding the given format marker.
+	 *
+	 * @param ctx  handler context supplying the allocator
+	 * @param type marker byte
+	 * @return buffer holding the marker
+	 */
+	private static ByteBuf markerBuffer(ChannelHandlerContext ctx, byte type) {
+		return ctx.alloc().buffer(1).writeByte(type);
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedResponse.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedResponse.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedResponse.java
index 9a4b59a..d208d75 100644
--- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedResponse.java
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedResponse.java
@@ -121,6 +121,10 @@
 		return _linItem;
 	}
 
+	public boolean isLineageReusable() {
+		return _linItem != null;
+	}
+
 	@Override
 	public String toString() {
 		StringBuilder sb = new StringBuilder();

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorker.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorker.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorker.java
index 682cc8e..8ed74c4 100644
--- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorker.java
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorker.java
@@ -55,9 +55,8 @@
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.channel.socket.nio.NioServerSocketChannel;
-import io.netty.handler.codec.serialization.ClassResolvers;
-import io.netty.handler.codec.serialization.ObjectDecoder;
 import io.netty.handler.codec.serialization.ObjectEncoder;
+import io.netty.handler.stream.ChunkedWriteHandler;
 import io.netty.handler.ssl.SslContext;
 import io.netty.util.concurrent.DefaultThreadFactory;
 
@@ -205,13 +204,13 @@
 				cp.addLast("CompressionDecodingStartStatistics", new CompressionDecoderStartStatisticsHandler());
 				compressionStrategy.ifPresent(strategy -> cp.addLast("CompressionDecoder", strategy.left));
 				cp.addLast("CompressionDecoderEndStatistics", new CompressionDecoderEndStatisticsHandler());
-				cp.addLast("ObjectDecoder", new ObjectDecoder(Integer.MAX_VALUE,
-					ClassResolvers.weakCachingResolver(ClassLoader.getSystemClassLoader())));
+				cp.addLast("FederatedFormatDecoder", new FederatedFormatDecoder());
 				cp.addLast("CompressionEncodingEndStatistics", new CompressionEncoderEndStatisticsHandler());
 				compressionStrategy.ifPresent(strategy -> cp.addLast("CompressionEncoder", strategy.right));
 				cp.addLast("CompressionEncodingStartStatistics", new CompressionEncoderStartStatisticsHandler());
-				cp.addLast("ObjectEncoder", new ObjectEncoder());
-				cp.addLast(FederationUtils.decoder(), new FederatedResponseEncoder());
+				cp.addLast("ChunkedWriteHandler", new ChunkedWriteHandler());
+				cp.addLast("FederatedResponseEncoder", new FederatedResponseEncoder());
+				cp.addLast("FederatedFormatEncoder", new FederatedFormatEncoder());
 				cp.addLast(new FederatedWorkerHandler(_flt, _frc, _fan, networkTimer));
 			}
 		};

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java

diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java
index 2cd9e8a..1233b70 100644
--- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java
+++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java
@@ -23,6 +23,7 @@
 import java.io.InputStreamReader;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
+import java.nio.channels.ClosedChannelException;
 import java.time.LocalDateTime;
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
@@ -734,15 +735,21 @@
 
 	private static class CloseListener implements ChannelFutureListener {
 		@Override
-		public void operationComplete(ChannelFuture channelFuture) throws InterruptedException {
-			if(!channelFuture.isSuccess()) {
-				LOG.error("Federated Worker Write failed");
-				channelFuture.channel().writeAndFlush(new FederatedResponse(ResponseType.ERROR,
-					new FederatedWorkerHandlerException("Error while sending response."))).channel().close().sync();
+		public void operationComplete(ChannelFuture channelFuture) {
+			if(channelFuture.isSuccess()) {
+				channelFuture.channel().close();
+				return;
 			}
-			else {
-				channelFuture.channel().close().sync();
+			Throwable cause = channelFuture.cause();
+			if(cause instanceof ClosedChannelException || !channelFuture.channel().isActive()) {
+				channelFuture.channel().close();
+				return;
 			}
+			LOG.error("Federated Worker Write failed", cause);
+			channelFuture.channel()
+				.writeAndFlush(new FederatedResponse(ResponseType.ERROR,
+					new FederatedWorkerHandlerException("Error while sending response.")))
+				.addListener(ChannelFutureListener.CLOSE);
 		}
 	}
 }

src/test/java/org/apache/sysds/test/component/federated/FederatedChunkCodecTest.java

diff --git a/src/test/java/org/apache/sysds/test/component/federated/FederatedChunkCodecTest.java b/src/test/java/org/apache/sysds/test/component/federated/FederatedChunkCodecTest.java
new file mode 100644
index 0000000..ad62a0a
--- /dev/null
+++ b/src/test/java/org/apache/sysds/test/component/federated/FederatedChunkCodecTest.java
@@ -0,0 +1,279 @@
+/*
+ * 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.test.component.federated;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import org.apache.sysds.runtime.controlprogram.federated.FederatedChunkDecoder;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedChunkEncoder;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedChunkProtocol;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse.ResponseType;
+import org.junit.Assert;
+import org.junit.Test;
+
+import io.netty.buffer.AbstractByteBufAllocator;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.buffer.UnpooledByteBufAllocator;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
+import io.netty.channel.embedded.EmbeddedChannel;
+import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
+import io.netty.handler.codec.compression.JdkZlibDecoder;
+import io.netty.handler.codec.compression.JdkZlibEncoder;
+import io.netty.handler.codec.compression.ZlibWrapper;
+import io.netty.handler.stream.ChunkedInput;
+import io.netty.handler.stream.ChunkedWriteHandler;
+
+public class FederatedChunkCodecTest {
+	private static final int CHUNK_SIZE = 4096; // tiny on purpose: forces a multi-frame stream
+	private static final int MAX_FRAME = 1 << 20; // 1 MB: frame-decoder ceiling, must exceed CHUNK_SIZE + header
+	private static final int PAYLOAD_DOUBLES = 20000; // ~160 KB serialized, many CHUNK_SIZE frames
+	private static final int QUEUED_DOUBLES = 2000; // ~16 KB serialized, few enough frames to never block the queue
+	// mirrors the package-private FederatedChunkProtocol, which this test cannot import
+
+	@Test
+	public void roundTripPlainSplitsIntoManyFrames() throws Exception {
+		FederatedResponse original = sampleResponse();
+		List<ByteBuf> frames = encode(original, false);
+		Assert.assertTrue("expected multiple frames", frames.size() > 2);
+		assertSamePayload(original, decode(frames, false));
+	}
+
+	@Test
+	public void roundTripThroughCompression() throws Exception {
+		FederatedResponse original = sampleResponse();
+		assertSamePayload(original, decode(encode(original, true), true));
+	}
+
+	@Test
+	public void producerFailureEmitsErrorFrame() throws Exception {
+		List<ByteBuf> frames = encode(new Unserializable(), false);
+		Assert.assertFalse("expected an error frame", frames.isEmpty());
+		ByteBuf last = frames.get(frames.size() - 1);
+		Assert.assertEquals(FederatedChunkProtocol.TYPE_ERROR, last.getByte(0));
+		String message = last.toString(FederatedChunkProtocol.HEADER_LEN, last.getInt(1), StandardCharsets.UTF_8);
+		Assert.assertTrue(message.contains("NotSerializableException"));
+		for(ByteBuf frame : frames)
+			frame.release();
+	}
+
+	@Test
+	public void errorFrameThrowsException() throws Exception {
+		final String ERROR_MSG = "remote failure";
+		byte[] cause = ERROR_MSG.getBytes(StandardCharsets.UTF_8);
+		ByteBuf errorFrame = Unpooled.buffer(FederatedChunkProtocol.HEADER_LEN + cause.length)
+			.writeByte(FederatedChunkProtocol.TYPE_ERROR).writeInt(cause.length).writeBytes(cause);
+		Throwable caught = writeAndAwaitException(errorFrame);
+		Assert.assertTrue("expected an IOException, got " + caught, caught instanceof IOException);
+		Assert.assertTrue(String.valueOf(caught).contains("remote failure"));
+	}
+
+	@Test
+	public void unknownFrameTypeThrowsException() throws Exception {
+		byte unknownType = 7; // not part of the protocol, must fail fast instead of stalling
+		ByteBuf unknownTypeFrame = Unpooled.buffer(FederatedChunkProtocol.HEADER_LEN).writeByte(unknownType)
+			.writeInt(0);
+		Throwable caught = writeAndAwaitException(unknownTypeFrame);
+		Assert.assertTrue("expected an IOException, got " + caught, caught instanceof IOException);
+		Assert.assertTrue(String.valueOf(caught).contains("Unknown federated chunk frame type: " + unknownType));
+	}
+
+	@Test
+	public void closeReleasesQueuedFrames() throws Exception {
+		RecordingAllocator alloc = new RecordingAllocator();
+		EmbeddedChannel channel = new EmbeddedChannel(new ChunkedWriteHandler());
+		ChunkedInput<ByteBuf> input = FederatedChunkEncoder.chunkedInput(sampleResponse(QUEUED_DOUBLES), CHUNK_SIZE,
+			alloc, channel.pipeline().get(ChunkedWriteHandler.class));
+		awaitProducerFinished(alloc);
+		input.close();
+		for(ByteBuf frame : alloc.frames())
+			Assert.assertEquals("frame left unreleased by close()", 0, frame.refCnt());
+	}
+
+	private static FederatedResponse sampleResponse() {
+		return sampleResponse(PAYLOAD_DOUBLES);
+	}
+
+	private static FederatedResponse sampleResponse(int doubles) {
+		double[] data = new double[doubles];
+		for(int i = 0; i < data.length; i++)
+			data[i] = i;
+		return new FederatedResponse(ResponseType.SUCCESS, data);
+	}
+
+	private static Throwable writeAndAwaitException(ByteBuf frame) throws InterruptedException {
+		EmbeddedChannel channel = new EmbeddedChannel(frameDecoder(), new FederatedChunkDecoder());
+		try {
+			// the deserializer thread reports the failure asynchronously: writeInbound throws it if it already arrived,
+			// otherwise awaitException below waits for it
+			channel.writeInbound(frame);
+		}
+		catch(Throwable t) {
+			return t;
+		}
+		return awaitException(channel);
+	}
+
+	private static Throwable awaitException(EmbeddedChannel channel) throws InterruptedException {
+		for(int i = 0; i < 20; i++) {
+			channel.runPendingTasks();
+			try {
+				channel.checkException();
+			}
+			catch(Throwable t) {
+				return t;
+			}
+			Thread.sleep(50);
+		}
+		throw new AssertionError("no exception propagated");
+	}
+
+	private static void awaitProducerFinished(RecordingAllocator alloc) throws InterruptedException {
+		for(int i = 0; i < 20; i++) {
+			if(alloc.sawFinalFrame())
+				return;
+			Thread.sleep(50);
+		}
+		throw new AssertionError("producer did not finish");
+	}
+
+	private static List<ByteBuf> encode(Serializable response, boolean compress) throws Exception {
+		EmbeddedChannel channel = compress ? new EmbeddedChannel(new JdkZlibEncoder(ZlibWrapper.ZLIB),
+			new ChunkedWriteHandler(), chunkEncoder()) : new EmbeddedChannel(new ChunkedWriteHandler(), chunkEncoder());
+		channel.config().setWriteBufferHighWaterMark(MAX_FRAME * 64);
+		List<ByteBuf> frames = new ArrayList<>();
+		ChannelFuture done = channel.write(response);
+		channel.flush();
+		pumpOutbound(channel, done, frames);
+		return frames;
+	}
+
+	private static ChannelOutboundHandlerAdapter chunkEncoder() {
+		return new ChannelOutboundHandlerAdapter() {
+			@Override
+			public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) {
+				ctx.write(FederatedChunkEncoder.chunkedInput((Serializable) msg, CHUNK_SIZE, ctx.alloc(),
+					ctx.pipeline().get(ChunkedWriteHandler.class)), promise);
+			}
+		};
+	}
+
+	private static void pumpOutbound(EmbeddedChannel channel, ChannelFuture done, List<ByteBuf> out) throws Exception {
+		for(int i = 0; i < 200; i++) {
+			channel.runPendingTasks();
+			drainOutbound(channel, out);
+			if(done.isDone())
+				break;
+			Thread.sleep(8);
+		}
+		drainOutbound(channel, out);
+	}
+
+	private static void drainOutbound(EmbeddedChannel channel, List<ByteBuf> out) {
+		ByteBuf buf;
+		while((buf = channel.readOutbound()) != null)
+			out.add(buf);
+	}
+
+	private static FederatedResponse decode(List<ByteBuf> frames, boolean compress) throws Exception {
+		EmbeddedChannel channel = compress ? new EmbeddedChannel(new JdkZlibDecoder(ZlibWrapper.ZLIB), frameDecoder(),
+			new FederatedChunkDecoder()) : new EmbeddedChannel(frameDecoder(), new FederatedChunkDecoder());
+		for(ByteBuf frame : frames)
+			channel.writeInbound(frame);
+		return awaitResponse(channel);
+	}
+
+	private static LengthFieldBasedFrameDecoder frameDecoder() {
+		return new LengthFieldBasedFrameDecoder(MAX_FRAME, 1, 4, 0, 0);
+	}
+
+	private static FederatedResponse awaitResponse(EmbeddedChannel channel) throws InterruptedException {
+		for(int i = 0; i < 20; i++) {
+			channel.runPendingTasks();
+			FederatedResponse response = channel.readInbound();
+			if(response != null)
+				return response;
+			Thread.sleep(50);
+		}
+		throw new AssertionError("no decoded response");
+	}
+
+	private static void assertSamePayload(FederatedResponse expected, FederatedResponse actual) throws Exception {
+		Assert.assertNotNull(actual);
+		Assert.assertTrue(actual.isSuccessful());
+		Assert.assertArrayEquals((double[]) expected.getData()[0], (double[]) actual.getData()[0], 0.0);
+	}
+
+	private static class Unserializable implements Serializable {
+		private static final long serialVersionUID = 1L;
+		private final Object _payload = new Object();
+
+		@Override
+		public String toString() {
+			return String.valueOf(_payload);
+		}
+	}
+
+	private static final class RecordingAllocator extends AbstractByteBufAllocator {
+		private final List<ByteBuf> _frames = Collections.synchronizedList(new ArrayList<>());
+
+		@Override
+		protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) {
+			return record(UnpooledByteBufAllocator.DEFAULT.heapBuffer(initialCapacity, maxCapacity));
+		}
+
+		@Override
+		protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) {
+			return record(UnpooledByteBufAllocator.DEFAULT.directBuffer(initialCapacity, maxCapacity));
+		}
+
+		@Override
+		public boolean isDirectBufferPooled() {
+			return false;
+		}
+
+		private ByteBuf record(ByteBuf buf) {
+			_frames.add(buf);
+			return buf;
+		}
+
+		List<ByteBuf> frames() {
+			return _frames;
+		}
+
+		boolean sawFinalFrame() {
+			synchronized(_frames) {
+				for(ByteBuf frame : _frames)
+					if(frame.isReadable() && frame.getByte(0) != FederatedChunkProtocol.TYPE_DATA)
+						return true;
+			}
+			return false;
+		}
+	}
+}

src/test/java/org/apache/sysds/test/component/federated/FederatedFormatRoutingTest.java

diff --git a/src/test/java/org/apache/sysds/test/component/federated/FederatedFormatRoutingTest.java b/src/test/java/org/apache/sysds/test/component/federated/FederatedFormatRoutingTest.java
new file mode 100644
index 0000000..153b3a6
--- /dev/null
+++ b/src/test/java/org/apache/sysds/test/component/federated/FederatedFormatRoutingTest.java
@@ -0,0 +1,174 @@
+/*
+ * 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.test.component.federated;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.sysds.api.DMLScript;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedFormatDecoder;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedFormatEncoder;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedChunkProtocol;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse.ResponseType;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedWorker.FederatedResponseEncoder;
+import org.apache.sysds.runtime.lineage.LineageCacheConfig.ReuseCacheType;
+import org.apache.sysds.runtime.lineage.LineageItem;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.junit.Assert;
+import org.junit.Test;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.embedded.EmbeddedChannel;
+import io.netty.handler.stream.ChunkedWriteHandler;
+
+public class FederatedFormatRoutingTest {
+	private static final int CHUNK_SIZE = 4096; // tiny on purpose: forces a multi-frame chunk stream
+	private static final long THRESHOLD_NEVER = Long.MAX_VALUE; // size guard never trips
+	private static final long THRESHOLD_ALWAYS = 1; // size guard always trips -> stream
+	private static final int PAYLOAD_DOUBLES = 20000; // ~160 KB serialized
+
+	@Test
+	public void nonCacheableRoutesChunked() throws Exception {
+		// plain double[] payload is not lineage-cacheable -> streams regardless of threshold
+		FederatedResponse original = sampleResponse();
+		List<ByteBuf> wire = encode(original, THRESHOLD_NEVER);
+		Assert.assertEquals(FederatedChunkProtocol.MARKER_CHUNKED, marker(wire));
+
+		EmbeddedChannel in = new EmbeddedChannel(new FederatedFormatDecoder());
+		FederatedResponse decoded = decode(in, wire);
+		Assert.assertNotNull(in.pipeline().get("FederatedChunkDecoder"));
+		assertDetectorRemoved(in);
+		assertSamePayload(original, decoded);
+	}
+
+	@Test
+	public void lineageCacheableRoutesObjectEncoder() throws Exception {
+		ReuseCacheType prev = DMLScript.LINEAGE_REUSE;
+		DMLScript.LINEAGE_REUSE = ReuseCacheType.REUSE_FULL;
+		try {
+			List<ByteBuf> wire = encode(cacheableResponse(), THRESHOLD_NEVER);
+			Assert.assertEquals(FederatedChunkProtocol.MARKER_OBJECT_ENCODER, marker(wire));
+
+			EmbeddedChannel in = new EmbeddedChannel(new FederatedFormatDecoder());
+			FederatedResponse decoded = decode(in, wire);
+			Assert.assertNotNull(in.pipeline().get("FederatedObjectDecoder"));
+			assertDetectorRemoved(in);
+			Assert.assertTrue(decoded.isSuccessful());
+		}
+		finally {
+			DMLScript.LINEAGE_REUSE = prev;
+		}
+	}
+
+	@Test
+	public void lineageCacheableOverThresholdRoutesChunked() throws Exception {
+		ReuseCacheType prev = DMLScript.LINEAGE_REUSE;
+		DMLScript.LINEAGE_REUSE = ReuseCacheType.REUSE_FULL;
+		try {
+			List<ByteBuf> wire = encode(cacheableResponse(), THRESHOLD_ALWAYS);
+			Assert.assertEquals(FederatedChunkProtocol.MARKER_CHUNKED, marker(wire));
+		}
+		finally {
+			DMLScript.LINEAGE_REUSE = prev;
+		}
+	}
+
+	@Test
+	public void reuseDisabledCacheableRoutesChunked() throws Exception {
+		// cacheable payload but lineage reuse off -> no cache to feed -> streams
+		ReuseCacheType prev = DMLScript.LINEAGE_REUSE;
+		DMLScript.LINEAGE_REUSE = ReuseCacheType.NONE;
+		try {
+			List<ByteBuf> wire = encode(cacheableResponse(), THRESHOLD_NEVER);
+			Assert.assertEquals(FederatedChunkProtocol.MARKER_CHUNKED, marker(wire));
+		}
+		finally {
+			DMLScript.LINEAGE_REUSE = prev;
+		}
+	}
+
+	private static FederatedResponse sampleResponse() {
+		double[] data = new double[PAYLOAD_DOUBLES];
+		for(int i = 0; i < data.length; i++)
+			data[i] = i;
+		return new FederatedResponse(ResponseType.SUCCESS, data);
+	}
+
+	private static FederatedResponse cacheableResponse() {
+		MatrixBlock mb = new MatrixBlock(16, 16, 1.0);
+		return new FederatedResponse(ResponseType.SUCCESS, new Object[] {mb}, new LineageItem("routing-test"));
+	}
+
+	private static List<ByteBuf> encode(FederatedResponse response, long threshold) throws Exception {
+		EmbeddedChannel out = new EmbeddedChannel(new ChunkedWriteHandler(), new FederatedResponseEncoder(),
+			new FederatedFormatEncoder(CHUNK_SIZE, threshold));
+		out.config().setWriteBufferHighWaterMark((CHUNK_SIZE + 64) * 64);
+		List<ByteBuf> wire = new ArrayList<>();
+		ChannelFuture done = out.write(response);
+		out.flush();
+		for(int i = 0; i < 800; i++) {
+			out.runPendingTasks();
+			drainInto(out, wire);
+			if(done.isDone())
+				break;
+			Thread.sleep(2);
+		}
+		drainInto(out, wire);
+		return wire;
+	}
+
+	private static void drainInto(EmbeddedChannel out, List<ByteBuf> wire) {
+		ByteBuf buf;
+		while((buf = out.readOutbound()) != null)
+			wire.add(buf);
+	}
+
+	private static byte marker(List<ByteBuf> wire) {
+		ByteBuf first = wire.get(0);
+		Assert.assertEquals("marker is a standalone 1-byte frame", 1, first.readableBytes());
+		return first.getByte(first.readerIndex());
+	}
+
+	private static FederatedResponse decode(EmbeddedChannel in, List<ByteBuf> wire) throws Exception {
+		for(ByteBuf buf : wire)
+			in.writeInbound(buf);
+		for(int i = 0; i < 200; i++) {
+			in.runPendingTasks();
+			FederatedResponse response = in.readInbound();
+			if(response != null)
+				return response;
+			Thread.sleep(5);
+		}
+		throw new AssertionError("no decoded response");
+	}
+
+	private static void assertDetectorRemoved(EmbeddedChannel in) {
+		Assert.assertNull("detector must remove itself after the first message",
+			in.pipeline().get(FederatedFormatDecoder.class));
+	}
+
+	private static void assertSamePayload(FederatedResponse expected, FederatedResponse actual) throws Exception {
+		Assert.assertNotNull(actual);
+		Assert.assertTrue(actual.isSuccessful());
+		Assert.assertArrayEquals((double[]) expected.getData()[0], (double[]) actual.getData()[0], 0.0);
+	}
+}

src/test/java/org/apache/sysds/test/functions/federated/io/FederatedMaxPayloadTest.java

diff --git a/src/test/java/org/apache/sysds/test/functions/federated/io/FederatedMaxPayloadTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/FederatedMaxPayloadTest.java
new file mode 100644
index 0000000..6ee00c6
--- /dev/null
+++ b/src/test/java/org/apache/sysds/test/functions/federated/io/FederatedMaxPayloadTest.java
@@ -0,0 +1,77 @@
+/*
+ * 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.test.functions.federated.io;
+
+import java.net.InetSocketAddress;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.runtime.controlprogram.federated.FederatedData;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest;
+import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.test.AutomatedTestBase;
+import org.apache.sysds.test.TestConfiguration;
+import org.junit.Assert;
+import org.junit.Ignore;
+import org.junit.Test;
+
+@Ignore("Heavy test: Transfer a matrix of size greater than 2GB over network. "
+	+ "This test needs ~9GB heap, since client and worker share one JVM; to run "
+	+ "this test, set the 'argLine' property inside the pom to '-Xmx9g'.")
+public class FederatedMaxPayloadTest extends AutomatedTestBase {
+
+	private final static String TEST_NAME = "FederatedMaxPayloadTest";
+	private final static String TEST_DIR = "functions/federated/network/";
+	private final static String TEST_CLASS_DIR = TEST_DIR + FederatedMaxPayloadTest.class.getSimpleName() + "/";
+
+	@Override
+	public void setUp() {
+		addTestConfiguration(TEST_NAME, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME, new String[] {""}));
+	}
+
+	@Test
+	public void transferOverTwoGigabytePayload() {
+		int port = getRandomAvailablePort();
+		startLocalFedWorkerThread(port);
+		InetSocketAddress address = new InetSocketAddress("localhost", port);
+		MatrixBlock mb = denseMatrixExceedingTwoGigabytes();
+		try {
+			FederatedRequest request = new FederatedRequest(FederatedRequest.RequestType.PUT_VAR, 1, mb);
+
+			Future<FederatedResponse> response = FederatedData.executeFederatedOperation(address, request);
+			Assert.assertTrue("Network send was not successful.", response.get().isSuccessful());
+		}
+		catch(Exception e) {
+			Assert.fail("Federated transfer failed: " + e.getMessage());
+		}
+		finally {
+			FederatedData.clearFederatedWorkers();
+		}
+	}
+
+	private static MatrixBlock denseMatrixExceedingTwoGigabytes() {
+		int rows = 30000;
+		int cols = 8950;
+		MatrixBlock mb = new MatrixBlock(rows, cols, false);
+		mb.allocateDenseBlock();
+		mb.setNonZeros((long) rows * cols);
+		return mb;
+	}
+}