From a0cbd5303e50743e29f672a69b7e6a14665c530e Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Fri, 3 Jul 2026 14:38:15 +0200 Subject: [PATCH] [OOC] IO Handler Performance Improvements and Generalization --- .../spark/data/IndexedMatrixValue.java | 51 ++- .../sysds/runtime/ooc/cache/OOCFuture.java | 275 ++++++++++++++++ .../cache/io/OOCBufferedDataInputStream.java | 307 ++++++++++++++++++ .../cache/io/OOCBufferedDataOutputStream.java | 277 ++++++++++++++++ .../runtime/ooc/cache/io/OOCIOHandler.java | 3 +- .../ooc/cache/io/OOCMatrixIOHandler.java | 118 ++++--- .../runtime/ooc/cache/io/SpillableObject.java | 32 ++ .../ooc/cache/io/SpillableObjectRegistry.java | 55 ++++ .../cache/legacy/OOCLRUCacheScheduler.java | 3 +- .../ooc/cache/OOCLRUCacheSchedulerTest.java | 11 +- 10 files changed, 1060 insertions(+), 72 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataInputStream.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataOutputStream.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObject.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java diff --git a/src/main/java/org/apache/sysds/runtime/instructions/spark/data/IndexedMatrixValue.java b/src/main/java/org/apache/sysds/runtime/instructions/spark/data/IndexedMatrixValue.java index 7b20fe2f9e5..bd96bbb614f 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/spark/data/IndexedMatrixValue.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/spark/data/IndexedMatrixValue.java @@ -20,34 +20,39 @@ package org.apache.sysds.runtime.instructions.spark.data; +import java.io.DataInput; +import java.io.DataOutput; +import java.io.IOException; import java.io.Serializable; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.matrix.data.MatrixValue; +import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; -public class IndexedMatrixValue implements Serializable +public class IndexedMatrixValue implements SpillableObject, Serializable { private static final long serialVersionUID = 6723389820806752110L; private MatrixIndexes _indexes = null; private MatrixValue _value = null; - + public IndexedMatrixValue() { _indexes = new MatrixIndexes(); } - + public IndexedMatrixValue(Class cls) { this(); - + //create new value object for given class try { _value=cls.getDeclaredConstructor().newInstance(); - } + } catch (Exception e) { throw new RuntimeException(e); } } - + public IndexedMatrixValue(MatrixIndexes ind, MatrixValue b) { this(); _indexes.setIndexes(ind); @@ -55,14 +60,14 @@ public IndexedMatrixValue(MatrixIndexes ind, MatrixValue b) { } public IndexedMatrixValue(IndexedMatrixValue that) { - this(that._indexes, that._value); + this(that._indexes, that._value); } - + public MatrixIndexes getIndexes() { return _indexes; } - + public MatrixValue getValue() { return _value; } @@ -70,14 +75,38 @@ public MatrixValue getValue() { public void setValue(MatrixValue value) { _value = value; } - + public void set(MatrixIndexes indexes2, MatrixValue block2) { _indexes.setIndexes(indexes2); _value = block2; } - + @Override public String toString() { return "("+_indexes.getRowIndex()+", "+_indexes.getColumnIndex()+"): \n"+_value; } + + @Override + public boolean tryWrite(DataOutput dataOutput) throws IOException { + MatrixIndexes ix = _indexes; + MatrixValue value = _value; + if(ix == null || value == null) + return false; + ix.write(dataOutput); + value.write(dataOutput); + return true; + } + + @Override + public void discard() { + _value = null; + } + + @Override + public void read(DataInput dataInput) throws IOException { + _indexes = new MatrixIndexes(); + _value = new MatrixBlock(); + _indexes.readFields(dataInput); + _value.readFields(dataInput); + } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java new file mode 100644 index 00000000000..491fefaad87 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java @@ -0,0 +1,275 @@ +/* + * 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.cache; + +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Function; + +/** + * Small future implementation for OOC hot paths. It supports multiple synchronous subscribers without + * the completion-stage support of {@link java.util.concurrent.CompletableFuture}. + */ +public class OOCFuture { + private Subscriber _subscribers; + private T _value; + private Throwable _error; + private boolean _done; + + public static OOCFuture completed(T value) { + OOCFuture future = new OOCFuture<>(); + future._value = value; + future._done = true; + return future; + } + + public static OOCFuture failed(Throwable error) { + OOCFuture future = new OOCFuture<>(); + future._error = error; + future._done = true; + return future; + } + + public boolean complete(T value) { + return finish(value, null); + } + + public boolean completeExceptionally(Throwable error) { + if(error == null) + throw new NullPointerException("error"); + return finish(null, error); + } + + public void thenAccept(Consumer action) { + subscribe(null, action, null); + } + + public void whenComplete(BiConsumer action) { + subscribe(null, null, action); + } + + public OOCFuture map(Function mapper) { + return new MappedFuture<>(this, mapper); + } + + public synchronized boolean isDone() { + return _done; + } + + public synchronized T getNow(T fallback) { + if(!_done) + return fallback; + if(_error != null) + throw new CompletionException(_error); + return _value; + } + + public T get() throws InterruptedException, ExecutionException { + synchronized(this) { + while(!_done) + wait(); + if(_error != null) + throw new ExecutionException(_error); + return _value; + } + } + + public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { + long remaining = unit.toNanos(timeout); + long deadline = System.nanoTime() + remaining; + synchronized(this) { + while(!_done) { + if(remaining <= 0) + throw new TimeoutException(); + TimeUnit.NANOSECONDS.timedWait(this, remaining); + remaining = deadline - System.nanoTime(); + } + if(_error != null) + throw new ExecutionException(_error); + return _value; + } + } + + private void subscribe(Function mapper, Consumer action, + BiConsumer completion) { + T value; + Throwable error; + synchronized(this) { + if(!_done) { + _subscribers = new Subscriber<>(mapper, action, completion, _subscribers); + return; + } + value = _value; + error = _error; + } + accept(mapper, action, completion, value, error); + } + + private boolean finish(T value, Throwable error) { + Subscriber subscribers; + synchronized(this) { + if(_done) + return false; + _value = value; + _error = error; + _done = true; + subscribers = _subscribers; + _subscribers = null; + notifyAll(); + } + while(subscribers != null) { + Subscriber next = subscribers.next; + subscribers.accept(value, error); + subscribers = next; + } + return true; + } + + private static void accept(Function mapper, Consumer action, + BiConsumer completion, T value, Throwable error) { + R result = null; + Throwable resultError = error; + if(resultError == null) { + try { + @SuppressWarnings("unchecked") + R mapped = mapper == null ? (R)value : mapper.apply(value); + result = mapped; + } + catch(Throwable t) { + resultError = t; + } + } + try { + if(completion != null) + completion.accept(result, resultError); + else if(resultError == null) + action.accept(result); + } + catch(Throwable ignored) { + // Subscribers are independent; one failed callback must not prevent the remaining notifications. + } + } + + private static final class Subscriber { + private final Function mapper; + private final Consumer action; + private final BiConsumer completion; + private final Subscriber next; + + @SuppressWarnings("unchecked") + private Subscriber(Function mapper, Consumer action, + BiConsumer completion, Subscriber next) { + this.mapper = mapper; + this.action = (Consumer)action; + this.completion = (BiConsumer)(BiConsumer)completion; + this.next = next; + } + + private void accept(T value, Throwable error) { + OOCFuture.accept(mapper, action, completion, value, error); + } + } + + private static final class MappedFuture extends OOCFuture { + private final OOCFuture source; + private final Function mapper; + + private MappedFuture(OOCFuture source, Function mapper) { + this.source = source; + this.mapper = mapper; + } + + @Override + public boolean complete(T value) { + throw new UnsupportedOperationException("Cannot complete a mapped OOCFuture"); + } + + @Override + public boolean completeExceptionally(Throwable error) { + throw new UnsupportedOperationException("Cannot complete a mapped OOCFuture"); + } + + @Override + public void thenAccept(Consumer action) { + source.subscribe(mapper, action, null); + } + + @Override + public void whenComplete(BiConsumer action) { + source.subscribe(mapper, null, action); + } + + @Override + public OOCFuture map(Function nextMapper) { + return new MappedFuture<>(source, value -> nextMapper.apply(mapper.apply(value))); + } + + @Override + public boolean isDone() { + return source.isDone(); + } + + @Override + public T getNow(T fallback) { + if(!source.isDone()) + return fallback; + try { + return mapper.apply(source.getNow(null)); + } + catch(CompletionException ex) { + throw ex; + } + catch(Throwable t) { + throw new CompletionException(t); + } + } + + @Override + public T get() throws InterruptedException, ExecutionException { + try { + return mapper.apply(source.get()); + } + catch(InterruptedException | ExecutionException ex) { + throw ex; + } + catch(Throwable t) { + throw new ExecutionException(t); + } + } + + @Override + public T get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + try { + return mapper.apply(source.get(timeout, unit)); + } + catch(InterruptedException | ExecutionException | TimeoutException ex) { + throw ex; + } + catch(Throwable t) { + throw new ExecutionException(t); + } + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataInputStream.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataInputStream.java new file mode 100644 index 00000000000..ee02b28404c --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataInputStream.java @@ -0,0 +1,307 @@ +/* + * 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.cache.io; + +import org.apache.sysds.runtime.data.SparseBlock; +import org.apache.sysds.runtime.data.SparseBlockCSR; +import org.apache.sysds.runtime.io.IOUtilFunctions; +import org.apache.sysds.runtime.matrix.data.MatrixBlockDataInput; + +import java.io.DataInput; +import java.io.DataInputStream; +import java.io.EOFException; +import java.io.IOException; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.DoubleBuffer; + +import jdk.incubator.vector.DoubleVector; +import jdk.incubator.vector.VectorOperators; +import jdk.incubator.vector.VectorSpecies; + +class OOCBufferedDataInputStream implements DataInput, MatrixBlockDataInput { + private static final int PAGE_SIZE = 4096; + private static final int PAGE_MASK = PAGE_SIZE - 1; + private static final int DEFAULT_BUFFER_SIZE = 64 * 1024; + private static final VectorSpecies DOUBLE_SPECIES = DoubleVector.SPECIES_PREFERRED; + private static final int DOUBLE_VECTOR_LENGTH = DOUBLE_SPECIES.length(); + + private final RandomAccessFile _in; + private final byte[] _buff; + private final byte[] _tmp; + private final DoubleBuffer[] _doubleDecodeBuffers; + private final int _bufflen; + private long _filePos; + private int _pos; + private int _count; + + OOCBufferedDataInputStream(RandomAccessFile in) throws IOException { + this(in, DEFAULT_BUFFER_SIZE); + } + + OOCBufferedDataInputStream(RandomAccessFile in, int size) throws IOException { + if(size <= 0) + throw new IllegalArgumentException("Buffer size <= 0."); + if(size % 8 != 0) + throw new IllegalArgumentException("Buffer size not a multiple of 8."); + _in = in; + _buff = new byte[size]; + _tmp = new byte[8]; + _doubleDecodeBuffers = createDoubleDecodeBuffers(_buff); + _bufflen = size; + _filePos = in.getFilePointer(); + _pos = 0; + _count = 0; + } + + @Override + public void readFully(byte[] b) throws IOException { + readFully(b, 0, b.length); + } + + @Override + public void readFully(byte [] b, int off, int len) throws IOException { + if(len < 0) + throw new IndexOutOfBoundsException(); + + while(len > 0) { + int avail = _count - _pos; + if(avail > 0) { + int n = Math.min(avail, len); + System.arraycopy(_buff, _pos, b, off, n); + _pos += n; + off += n; + len -= n; + } + else if(len >= _bufflen) { + _in.readFully(b, off, len); + _filePos += len; + return; + } + else { + refill(); + } + } + } + + @Override + public int skipBytes(int n) throws IOException { + throw new IOException("Not supported."); + } + + @Override + public boolean readBoolean() throws IOException { + return readByte() != 0; + } + + @Override + public byte readByte() throws IOException { + if(_pos >= _count) + refill(); + return _buff[_pos++]; + } + + @Override + public int readUnsignedByte() throws IOException { + return readByte() & 0xFF; + } + + @Override + public short readShort() throws IOException { + if(_count - _pos >= 2) { + short ret = (short)baToShort(_buff, _pos); + _pos += 2; + return ret; + } + readFully(_tmp, 0, 2); + return (short)baToShort(_tmp, 0); + } + + @Override + public int readUnsignedShort() throws IOException { + return readShort() & 0xFFFF; + } + + @Override + public char readChar() throws IOException { + return (char)readUnsignedShort(); + } + + @Override + public int readInt() throws IOException { + if(_count - _pos >= 4) { + int ret = baToInt(_buff, _pos); + _pos += 4; + return ret; + } + readFully(_tmp, 0, 4); + return baToInt(_tmp, 0); + } + + @Override + public long readLong() throws IOException { + if(_count - _pos >= 8) { + long ret = baToLong(_buff, _pos); + _pos += 8; + return ret; + } + readFully(_tmp, 0, 8); + return baToLong(_tmp, 0); + } + + @Override + public float readFloat() throws IOException { + return Float.intBitsToFloat(readInt()); + } + + @Override + public double readDouble() throws IOException { + return Double.longBitsToDouble(readLong()); + } + + @Override + public String readLine() throws IOException { + throw new IOException("Not supported."); + } + + @Override + public String readUTF() throws IOException { + return DataInputStream.readUTF(this); + } + + @Override + public long readDoubleArray(int len, double[] varr) throws IOException { + if(len <= 0 || len > varr.length) + throw new IndexOutOfBoundsException("len=" + len + ", varr.length=" + varr.length); + + long nnz = 0; + int ix = 0; + while(ix < len) { + int avail = _count - _pos; + if(avail <= 0) { + refill(); + continue; + } + if(avail < 8) { + readFully(_tmp, 0, 8); + double v = Double.longBitsToDouble(baToLong(_tmp, 0)); + varr[ix] = v; + nnz += (v != 0) ? 1 : 0; + ix++; + continue; + } + + int ndbl = Math.min(len - ix, avail / 8); + int end = _pos + ndbl * 8; + readDoubles(_pos, ndbl, varr, ix); + nnz += countNonZeros(varr, ix, ndbl); + ix += ndbl; + _pos = end; + } + return nnz; + } + + @Override + public long readSparseRows(int rlen, long nnz, SparseBlock rows) throws IOException { + if(rows instanceof SparseBlockCSR) { + ((SparseBlockCSR)rows).initSparse(rlen, (int)nnz, this); + return nnz; + } + + long gnnz = 0; + for(int i = 0; i < rlen; i++) { + int lnnz = readInt(); + if(lnnz > 0) { + rows.allocate(i, lnnz); + + for(int j = 0; j < lnnz; j++) { + int aix = readInt(); + double aval = readDouble(); + rows.append(i, aix, aval); + } + gnnz += lnnz; + } + } + + if(gnnz != nnz) + throw new IOException("Invalid number of read nnz: " + gnnz + " vs " + nnz); + return nnz; + } + + private void refill() throws IOException { + int len = getRefillLength(); + _count = _in.read(_buff, 0, len); + _pos = 0; + if(_count < 0) + throw new EOFException(); + _filePos += _count; + } + + private int getRefillLength() { + int pageOffset = (int)(_filePos & PAGE_MASK); + if(pageOffset == 0) + return _bufflen; + return Math.min(_bufflen, PAGE_SIZE - pageOffset); + } + + private static int baToShort(byte[] ba, final int off) { + return IOUtilFunctions.baToShort(ba, off); + } + + private static int baToInt(byte[] ba, final int off) { + return IOUtilFunctions.baToInt(ba, off); + } + + private static long baToLong(byte[] ba, final int off) { + return IOUtilFunctions.baToLong(ba, off); + } + + private void readDoubles(int srcPos, int len, double[] dest, int destPos) { + int alignment = srcPos & 7; + DoubleBuffer dbuff = _doubleDecodeBuffers[alignment]; + dbuff.position((srcPos - alignment) >>> 3); + dbuff.get(dest, destPos, len); + } + + private static DoubleBuffer[] createDoubleDecodeBuffers(byte[] buff) { + DoubleBuffer[] ret = new DoubleBuffer[8]; + for(int i = 0; i < ret.length; i++) { + ByteBuffer bbuff = ByteBuffer.wrap(buff); + bbuff.position(i); + ret[i] = bbuff.slice().order(ByteOrder.BIG_ENDIAN).asDoubleBuffer(); + } + return ret; + } + + private static long countNonZeros(double[] values, int off, int len) { + long nnz = 0; + int i = 0; + int upper = DOUBLE_SPECIES.loopBound(len); + DoubleVector vzero = DoubleVector.zero(DOUBLE_SPECIES); + for(; i < upper; i += DOUBLE_VECTOR_LENGTH) { + DoubleVector v = DoubleVector.fromArray(DOUBLE_SPECIES, values, off + i); + nnz += v.compare(VectorOperators.NE, vzero).trueCount(); + } + for(; i < len; i++) + nnz += (values[off + i] != 0) ? 1 : 0; + return nnz; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataOutputStream.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataOutputStream.java new file mode 100644 index 00000000000..9854a94586b --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCBufferedDataOutputStream.java @@ -0,0 +1,277 @@ +/* + * 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.cache.io; + +import org.apache.sysds.runtime.data.SparseBlock; +import org.apache.sysds.runtime.io.IOUtilFunctions; +import org.apache.sysds.runtime.matrix.data.MatrixBlockDataOutput; + +import java.io.DataOutput; +import java.io.FilterOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.io.UTFDataFormatException; + +class OOCBufferedDataOutputStream extends FilterOutputStream implements DataOutput, MatrixBlockDataOutput { + private final byte[] _buff; + private final int _bufflen; + private int _count; + private long _position; + private long _flushedPosition; + + OOCBufferedDataOutputStream(OutputStream out) { + this(out, 8192); + } + + OOCBufferedDataOutputStream(OutputStream out, int size) { + super(out); + if(size <= 0) + throw new IllegalArgumentException("Buffer size <= 0."); + if(size % 8 != 0) + throw new IllegalArgumentException("Buffer size not a multiple of 8."); + _buff = new byte[size]; + _bufflen = size; + _count = 0; + _position = 0; + _flushedPosition = 0; + } + + long getPosition() { + return _position; + } + + long getFlushedPosition() { + return _flushedPosition; + } + + @Override + public void write(int b) throws IOException { + if(_count >= _bufflen) + flushBuffer(); + _buff[_count++] = (byte)b; + _position++; + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + if(len > _bufflen) { + flushBuffer(); + out.write(b, off, len); + _position += len; + _flushedPosition += len; + } + else { + if(len > _bufflen - _count) + flushBuffer(); + System.arraycopy(b, off, _buff, _count, len); + _count += len; + _position += len; + } + } + + @Override + public void flush() throws IOException { + flushBuffer(); + out.flush(); + } + + private void flushBuffer() throws IOException { + if(_count > 0) { + out.write(_buff, 0, _count); + _flushedPosition += _count; + _count = 0; + } + } + + @Override + public void close() throws IOException { + super.close(); + } + + @Override + public void writeBoolean(boolean v) throws IOException { + if(_count >= _bufflen) + flushBuffer(); + _buff[_count++] = (byte)(v ? 1 : 0); + _position++; + } + + @Override + public void writeInt(int v) throws IOException { + if(_count + 4 > _bufflen) + flushBuffer(); + intToBa(v, _buff, _count); + _count += 4; + _position += 4; + } + + @Override + public void writeLong(long v) throws IOException { + if(_count + 8 > _bufflen) + flushBuffer(); + longToBa(v, _buff, _count); + _count += 8; + _position += 8; + } + + @Override + public void writeDouble(double v) throws IOException { + if(_count + 8 > _bufflen) + flushBuffer(); + longToBa(Double.doubleToRawLongBits(v), _buff, _count); + _count += 8; + _position += 8; + } + + @Override + public void writeFloat(float v) throws IOException { + if(_count + 4 > _bufflen) + flushBuffer(); + intToBa(Float.floatToIntBits(v), _buff, _count); + _count += 4; + _position += 4; + } + + @Override + public void writeByte(int v) throws IOException { + if(_count + 1 > _bufflen) + flushBuffer(); + _buff[_count++] = (byte)v; + _position++; + } + + @Override + public void writeShort(int v) throws IOException { + if(_count + 2 > _bufflen) + flushBuffer(); + shortToBa(v, _buff, _count); + _count += 2; + _position += 2; + } + + @Override + public void writeBytes(String s) throws IOException { + throw new IOException("Not supported."); + } + + @Override + public void writeChar(int v) throws IOException { + writeShort(v); + } + + @Override + public void writeChars(String s) throws IOException { + throw new IOException("Not supported."); + } + + @Override + public void writeUTF(String s) throws IOException { + int slen = s.length(); + int utflen = IOUtilFunctions.getUTFSize(s) - 2; + if(utflen - 2 > 65535) + throw new UTFDataFormatException("encoded string too long: " + utflen); + + writeShort(utflen); + for(int i = 0; i < slen; i++) { + if(_count + 3 > _bufflen) + flushBuffer(); + final char c = s.charAt(i); + if(c >= 0x0001 && c <= 0x007F) { + _buff[_count++] = (byte)c; + _position++; + } + else if(c >= 0x0800) { + _buff[_count++] = (byte)(0xE0 | ((c >> 12) & 0x0F)); + _buff[_count++] = (byte)(0x80 | ((c >> 6) & 0x3F)); + _buff[_count++] = (byte)(0x80 | (c & 0x3F)); + _position += 3; + } + else { + _buff[_count++] = (byte)(0xC0 | ((c >> 6) & 0x1F)); + _buff[_count++] = (byte)(0x80 | (c & 0x3F)); + _position += 2; + } + } + } + + @Override + public void writeDoubleArray(int len, double[] varr) throws IOException { + for(int i = 0; i < len; ) { + if(_count >= _bufflen) + flushBuffer(); + int lblen = Math.min(len - i, (_bufflen - _count) / 8); + if(lblen == 0) { + flushBuffer(); + continue; + } + for(int j = 0; j < lblen; j++) { + longToBa(Double.doubleToRawLongBits(varr[i + j]), _buff, _count); + _count += 8; + } + _position += 8L * lblen; + i += lblen; + if(_count >= _bufflen) + flushBuffer(); + } + } + + @Override + public void writeSparseRows(int rlen, SparseBlock rows) throws IOException { + int lrlen = Math.min(rows.numRows(), rlen); + for(int i = 0; i < lrlen; i++) { + if(!rows.isEmpty(i)) { + int apos = rows.pos(i); + int alen = rows.size(i); + int[] aix = rows.indexes(i); + double[] avals = rows.values(i); + + writeInt(alen); + + for(int j = apos; j < apos + alen; j++) { + if(_count + 12 > _bufflen) + flushBuffer(); + long tmp = Double.doubleToRawLongBits(avals[j]); + intToBa(aix[j], _buff, _count); + longToBa(tmp, _buff, _count + 4); + _count += 12; + _position += 12; + } + } + else { + writeInt(0); + } + } + + for(int i = lrlen; i < rlen; i++) + writeInt(0); + } + + private static void shortToBa(final int val, byte[] ba, final int off) { + IOUtilFunctions.shortToBa(val, ba, off); + } + + private static void intToBa(final int val, byte[] ba, final int off) { + IOUtilFunctions.intToBa(val, ba, off); + } + + private static void longToBa(final long val, byte[] ba, final int off) { + IOUtilFunctions.longToBa(val, ba, off); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCIOHandler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCIOHandler.java index 21085626a71..ab28df0ef0f 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCIOHandler.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCIOHandler.java @@ -25,6 +25,7 @@ import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.BlockKey; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import java.util.concurrent.CompletableFuture; import java.util.List; @@ -34,7 +35,7 @@ public interface OOCIOHandler { CompletableFuture scheduleEviction(BlockEntry block); - CompletableFuture scheduleRead(BlockEntry block); + OOCFuture scheduleRead(BlockEntry block); /** * Increase priority for a pending scheduled read if it has not started yet. diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java index 029c9e8060f..bfe1565ab3a 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java @@ -34,10 +34,9 @@ import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.BlockKey; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; import org.apache.sysds.runtime.ooc.stream.SourceOOCStream; -import org.apache.sysds.runtime.util.FastBufferedDataInputStream; -import org.apache.sysds.runtime.util.FastBufferedDataOutputStream; import org.apache.sysds.runtime.util.LocalFileUtils; import org.apache.sysds.utils.Statistics; import scala.Tuple2; @@ -46,9 +45,7 @@ import java.io.DataInput; import java.io.FileOutputStream; import java.io.IOException; -import java.io.OutputStream; import java.io.RandomAccessFile; -import java.nio.channels.Channels; import java.nio.channels.ClosedByInterruptException; import java.util.ArrayList; import java.util.Arrays; @@ -69,13 +66,14 @@ import java.util.concurrent.atomic.AtomicReference; public class OOCMatrixIOHandler implements OOCIOHandler { - private static final int WRITER_SIZE = 4; - private static final int READER_SIZE = 10; + private static final int WRITER_SIZE = 8; + private static final int READER_SIZE = 16; private static final long OVERFLOW = 8192 * 1024; private static final long MAX_PARTITION_SIZE = 8192 * 8192; private static final long GROUP_TARGET_BYTES = 8L * 1024 * 1024; private static final long GROUP_MAX_BYTES = 16L * 1024 * 1024; private static final int GROUP_MAX_COUNT = 64; + private static final long IDLE_FLUSH_MS = 1; private final String _spillDir; private final ThreadPoolExecutor _writeExec; @@ -147,11 +145,11 @@ public void shutdown() { if (started) { try { for(int i = 0; i < WRITER_SIZE; i++) { - _q[i].close(); + if(_q[i] != null) + _q[i].close(); } } - catch(InterruptedException e) { - Thread.currentThread().interrupt(); + catch(InterruptedException ignored) { } } _writeExec.getQueue().clear(); @@ -175,18 +173,20 @@ public CompletableFuture scheduleEviction(BlockEntry block) { try { long q = _wCtr.getAndAdd(block.getSize()) / OVERFLOW; int i = (int)(q % WRITER_SIZE); - _q[i].enqueueIfOpen(new Tuple2<>(block, future)); + if(!_q[i].enqueueIfOpen(new Tuple2<>(block, future))) + future.completeExceptionally(new DMLRuntimeException("OOC writer queue is closed")); } - catch(InterruptedException e) { + catch(InterruptedException ignored) { Thread.currentThread().interrupt(); + future.completeExceptionally(new DMLRuntimeException("Interrupted while scheduling OOC eviction")); } return future; } @Override - public CompletableFuture scheduleRead(final BlockEntry block) { - final CompletableFuture future = new CompletableFuture<>(); + public OOCFuture scheduleRead(final BlockEntry block) { + final OOCFuture future = new OOCFuture<>(); int pinnedPartitionId = pinPartitionForRead(block.getKey()); try { ReadTask task = new ReadTask(block, future, _readSeq.getAndIncrement(), pinnedPartitionId); @@ -532,25 +532,23 @@ private void loadFromDisk(BlockEntry block) { String filename = partFile.filePath; - // Create an empty object to read data into. - MatrixIndexes ix = new MatrixIndexes(); - MatrixBlock mb = new MatrixBlock(); + SpillableObject obj; try (RandomAccessFile raf = new RandomAccessFile(filename, "r")) { raf.seek(sloc.offset); - DataInput dis = new FastBufferedDataInputStream(Channels.newInputStream(raf.getChannel())); + DataInput dis = new OOCBufferedDataInputStream(raf); long ioStart = DMLScript.OOC_STATISTICS ? System.nanoTime() : 0; - ix.readFields(dis); // 1. Read Indexes - mb.readFields(dis); // 2. Read Block + obj = SpillableObjectRegistry.read(dis); if (DMLScript.OOC_STATISTICS) ioDuration = System.nanoTime() - ioStart; } catch (ClosedByInterruptException ignored) { + return; } catch (IOException e) { throw new RuntimeException(e); } - block.setDataUnsafe(new IndexedMatrixValue(ix, mb)); + block.setDataUnsafe(obj); if (DMLScript.OOC_STATISTICS) { Statistics.incrementOOCLoadFromDisk(); @@ -617,22 +615,27 @@ private void evictTask(CloseableQueue partFile.incrementRefCount(); // Writer pin; released when partition closes FileOutputStream fos = null; - CountableFastBufferedDataOutputStream dos = null; + OOCBufferedDataOutputStream dos = null; ConcurrentLinkedDeque>> waitingForFlush = null; try { fos = new FileOutputStream(filename); - dos = new CountableFastBufferedDataOutputStream(fos); + dos = new OOCBufferedDataOutputStream(fos); Tuple2> tpl; waitingForFlush = new ConcurrentLinkedDeque<>(); boolean closePartition = false; - while((tpl = q.take()) != null) { + while(!q.isFinished()) { + tpl = q.poll(IDLE_FLUSH_MS, TimeUnit.MILLISECONDS); + if(tpl == null) { + flushReadable(dos, waitingForFlush); + continue; + } long ioStart = DMLScript.OOC_STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; - BlockEntry entry = tpl._1; - CompletableFuture future = tpl._2; - long wrote = writeOut(partitionId, entry, future, fos, dos, waitingForFlush); + BlockEntry entry = tpl._1(); + CompletableFuture future = tpl._2(); + long wrote = writeOut(partitionId, entry, future, dos, waitingForFlush); if(DMLScript.OOC_STATISTICS && wrote > 0) { Statistics.incrementOOCEvictionWrite(); @@ -654,9 +657,9 @@ private void evictTask(CloseableQueue if (!closePartition && q.close()) { while((tpl = q.take()) != null) { long ioStart = DMLScript.OOC_STATISTICS ? System.nanoTime() : 0; - BlockEntry entry = tpl._1; - CompletableFuture future = tpl._2; - long wrote = writeOut(partitionId, entry, future, fos, dos, waitingForFlush); + BlockEntry entry = tpl._1(); + CompletableFuture future = tpl._2(); + long wrote = writeOut(partitionId, entry, future, dos, waitingForFlush); byteCtr += wrote; if(DMLScript.OOC_STATISTICS && wrote > 0) { @@ -685,45 +688,62 @@ private void evictTask(CloseableQueue } } - private long writeOut(int partitionId, BlockEntry entry, CompletableFuture future, FileOutputStream fos, - CountableFastBufferedDataOutputStream dos, ConcurrentLinkedDeque>> flushQueue) throws IOException { + private long writeOut(int partitionId, BlockEntry entry, CompletableFuture future, + OOCBufferedDataOutputStream dos, + ConcurrentLinkedDeque>> flushQueue) throws IOException { + String key = entry.getKey().toFileKey(); boolean alreadySpilled = _spillLocations.containsKey(key); if (!alreadySpilled) { - // 1. get the current file position. this is the offset. - // flush any buffered data to the file - //dos.flush(); - long offsetBefore = fos.getChannel().position() + dos.getCount(); + long offsetBefore = dos.getPosition(); + + if(future.isCancelled()) + return 0; // 2. write indexes and block - IndexedMatrixValue imv = (IndexedMatrixValue) entry.getDataUnsafe(); // Get data without requiring pin - if(imv == null) + SpillableObject so = (SpillableObject) entry.getDataUnsafe(); // Get data without requiring pin + if(so == null) + return 0; + if(!SpillableObjectRegistry.tryWrite(dos, so)) return 0; - imv.getIndexes().write(dos); // write Indexes - imv.getValue().write(dos); - long offsetAfter = fos.getChannel().position() + dos.getCount(); + long offsetAfter = dos.getPosition(); + if(future.isCancelled()) + return offsetAfter - offsetBefore; flushQueue.offer(new Tuple3<>(offsetBefore, offsetAfter, future)); // 3. create the spillLocation SpillLocation sloc = new SpillLocation(partitionId, offsetBefore); addSpillLocation(key, sloc); - flushQueue(fos.getChannel().position(), flushQueue); + if(future.isCancelled()) { + removeSpillLocation(key); + return offsetAfter - offsetBefore; + } + flushQueue(dos.getFlushedPosition(), flushQueue); return offsetAfter - offsetBefore; } + future.completeExceptionally(new DMLRuntimeException("Duplicate OOC spill location for: " + key)); return 0; } private void flushQueue(long offset, ConcurrentLinkedDeque>> flushQueue) { Tuple3> tmp; - while ((tmp = flushQueue.peek()) != null && tmp._2() < offset) { + while ((tmp = flushQueue.peek()) != null && tmp._2() <= offset) { flushQueue.poll(); tmp._3().complete(null); } } + private void flushReadable(OOCBufferedDataOutputStream dos, + ConcurrentLinkedDeque>> flushQueue) throws IOException { + if(dos == null || flushQueue.isEmpty()) + return; + dos.flush(); + flushQueue(dos.getFlushedPosition(), flushQueue); + } + private void addSpillLocation(String key, SpillLocation sloc) { synchronized(_spillLock) { SpillLocation existing = _spillLocations.putIfAbsent(key, sloc); @@ -806,12 +826,12 @@ private void unpinPartitionForRead(int partitionId) { private class ReadTask implements Runnable, Comparable { private final BlockEntry _block; - private final CompletableFuture _future; + private final OOCFuture _future; private final long _sequence; private final int _pinnedPartitionId; private double _priority; - private ReadTask(BlockEntry block, CompletableFuture future, long sequence, int pinnedPartitionId) { + private ReadTask(BlockEntry block, OOCFuture future, long sequence, int pinnedPartitionId) { this._block = block; this._future = future; this._sequence = sequence; @@ -880,16 +900,6 @@ int decrementRefCount() { } } - private static class CountableFastBufferedDataOutputStream extends FastBufferedDataOutputStream { - public CountableFastBufferedDataOutputStream(OutputStream out) { - super(out); - } - - public int getCount() { - return _count; - } - } - private static class SourceReadState implements SourceReadContinuation { final SourceReadRequest request; final Path[] paths; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObject.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObject.java new file mode 100644 index 00000000000..434f70601d6 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObject.java @@ -0,0 +1,32 @@ +/* + * 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.cache.io; + +import java.io.DataInput; +import java.io.DataOutput; +import java.io.IOException; + +public interface SpillableObject { + boolean tryWrite(DataOutput out) throws IOException; + void read(DataInput in) throws IOException; + + default void discard() { + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java new file mode 100644 index 00000000000..6cd8fded9e7 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java @@ -0,0 +1,55 @@ +/* + * 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.cache.io; + +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; + +import java.io.DataInput; +import java.io.DataOutput; +import java.io.IOException; + +public final class SpillableObjectRegistry { + private static final byte INDEXED_MATRIX_VALUE = 1; + + private SpillableObjectRegistry() { + } + + public static boolean tryWrite(DataOutput out, SpillableObject obj) throws IOException { + byte type = typeOf(obj); + out.writeByte(type); + return obj.tryWrite(out); + } + + public static SpillableObject read(DataInput in) throws IOException { + byte type = in.readByte(); + SpillableObject obj = switch(type) { + case INDEXED_MATRIX_VALUE -> new IndexedMatrixValue(); + default -> throw new IOException("Unknown spillable object type: " + type); + }; + obj.read(in); + return obj; + } + + private static byte typeOf(SpillableObject obj) throws IOException { + if(obj instanceof IndexedMatrixValue) + return INDEXED_MATRIX_VALUE; + throw new IOException("Unsupported spillable object type: " + obj.getClass().getName()); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java index c1f7058b5dc..3f5601adbae 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java @@ -28,6 +28,7 @@ import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.BlockKey; import org.apache.sysds.runtime.ooc.cache.BlockState; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; @@ -860,7 +861,7 @@ else if(allReserved && reading && req.isComplete()) { for(Tuple2 tpl : toRead) { final BlockEntry entry = tpl._2; - CompletableFuture future = _ioHandler.scheduleRead(entry); + OOCFuture future = _ioHandler.scheduleRead(entry); future.whenComplete((r, t) -> { if(t != null) { BlockReadState state; diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCLRUCacheSchedulerTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCLRUCacheSchedulerTest.java index 002b19e57be..cd05d97b317 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCLRUCacheSchedulerTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCLRUCacheSchedulerTest.java @@ -24,6 +24,7 @@ import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.BlockKey; import org.apache.sysds.runtime.ooc.cache.BlockState; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; import org.apache.sysds.runtime.ooc.cache.legacy.OOCLRUCacheScheduler; import org.junit.After; @@ -302,7 +303,7 @@ private static List snapshotDeferredOrder(OOCLRUCacheScheduler schedul } private static class FakeIOHandler implements OOCIOHandler { - private final Map> _readFutures = new HashMap<>(); + private final Map> _readFutures = new HashMap<>(); private final Map _readEntries = new HashMap<>(); private final Map _readCounts = new HashMap<>(); @@ -319,8 +320,8 @@ public CompletableFuture scheduleEviction(BlockEntry block) { } @Override - public CompletableFuture scheduleRead(BlockEntry block) { - CompletableFuture future = new CompletableFuture<>(); + public OOCFuture scheduleRead(BlockEntry block) { + OOCFuture future = new OOCFuture<>(); _readFutures.put(block.getKey(), future); _readEntries.put(block.getKey(), block); _readCounts.computeIfAbsent(block.getKey(), k -> new AtomicInteger(0)).incrementAndGet(); @@ -355,7 +356,7 @@ public int getReadCount(BlockKey key) { } public void completeRead(BlockKey key) { - CompletableFuture future = _readFutures.get(key); + OOCFuture future = _readFutures.get(key); if (future == null) throw new IllegalStateException("No scheduled read for " + key); BlockEntry entry = _readEntries.get(key); @@ -364,5 +365,5 @@ public void completeRead(BlockKey key) { BlockEntryTestAccess.setDataUnsafe(entry, new Object()); future.complete(entry); } - } } +}