|
25 | 25 | import org.apache.log4j.Logger; |
26 | 26 | import org.apache.sysds.api.jmlc.Connection; |
27 | 27 |
|
| 28 | +import org.apache.sysds.common.Types; |
| 29 | +import org.apache.sysds.runtime.DMLRuntimeException; |
| 30 | +import org.apache.sysds.runtime.matrix.data.MatrixBlock; |
| 31 | +import org.apache.sysds.runtime.util.CommonThreadPool; |
| 32 | +import org.apache.sysds.runtime.util.UnixPipeUtils; |
28 | 33 | import py4j.DefaultGatewayServerListener; |
29 | 34 | import py4j.GatewayServer; |
30 | 35 | import py4j.Py4JNetworkException; |
31 | 36 |
|
| 37 | +import java.io.BufferedInputStream; |
| 38 | +import java.io.BufferedOutputStream; |
| 39 | +import java.io.IOException; |
| 40 | +import java.util.ArrayList; |
| 41 | +import java.util.HashMap; |
| 42 | +import java.util.List; |
| 43 | +import java.util.concurrent.Callable; |
| 44 | +import java.util.concurrent.ExecutionException; |
| 45 | +import java.util.concurrent.ExecutorService; |
| 46 | +import java.util.concurrent.Future; |
| 47 | + |
32 | 48 |
|
33 | 49 | public class PythonDMLScript { |
34 | 50 |
|
35 | 51 | private static final Log LOG = LogFactory.getLog(PythonDMLScript.class.getName()); |
36 | 52 | final private Connection _connection; |
37 | 53 | public static GatewayServer GwS; |
38 | 54 |
|
| 55 | + private static String fromPythonBase = "py2java"; |
| 56 | + private static String toPythonBase = "java2py"; |
| 57 | + public HashMap<Integer, BufferedInputStream> fromPython = null; |
| 58 | + public HashMap<Integer, BufferedOutputStream> toPython = null; |
| 59 | + public String baseDir; |
| 60 | + private static int BATCH_SIZE = 32*1024; |
| 61 | + |
39 | 62 | /** |
40 | 63 | * Entry point for Python API. |
41 | 64 | * |
@@ -78,6 +101,103 @@ public Connection getConnection() { |
78 | 101 | return _connection; |
79 | 102 | } |
80 | 103 |
|
| 104 | + |
| 105 | + public void openPipes(String path, int num) throws IOException { |
| 106 | + fromPython = new HashMap<>(num * 2); |
| 107 | + toPython = new HashMap<>(num * 2); |
| 108 | + baseDir = path; |
| 109 | + for (int i = 0; i < num; i++) { |
| 110 | + BufferedInputStream pipe_in = UnixPipeUtils.openInput(path + "/" + fromPythonBase + "-" + i, i); |
| 111 | + LOG.debug("PY2JAVA pipe "+i+" is ready!"); |
| 112 | + fromPython.put(i, pipe_in); |
| 113 | + |
| 114 | + BufferedOutputStream pipe_out = UnixPipeUtils.openOutput(path + "/" + toPythonBase + "-" + i, i); |
| 115 | + toPython.put(i, pipe_out); |
| 116 | + } |
| 117 | + } |
| 118 | + |
| 119 | + public MatrixBlock startReadingMbFromPipe(int id, int rlen, int clen, Types.ValueType type) throws IOException { |
| 120 | + long limit = (long) rlen * clen; |
| 121 | + LOG.debug("trying to read matrix from "+id+" with "+rlen+" rows and "+clen+" columns. Total size: "+limit); |
| 122 | + if(limit > Integer.MAX_VALUE) |
| 123 | + throw new DMLRuntimeException("Dense NumPy array of size " + limit + |
| 124 | + " cannot be converted to MatrixBlock"); |
| 125 | + MatrixBlock mb = new MatrixBlock(rlen, clen, false, -1); |
| 126 | + if(fromPython != null){ |
| 127 | + BufferedInputStream pipe = fromPython.get(id); |
| 128 | + double[] denseBlock = new double[(int) limit]; |
| 129 | + UnixPipeUtils.readNumpyArrayInBatches(pipe, id, BATCH_SIZE, (int) limit, type, denseBlock, 0); |
| 130 | + mb.init(denseBlock, rlen, clen); |
| 131 | + } else { |
| 132 | + throw new DMLRuntimeException("FIFO Pipes are not initialized."); |
| 133 | + } |
| 134 | + mb.recomputeNonZeros(); |
| 135 | + mb.examSparsity(); |
| 136 | + LOG.debug("Reading from Python finished"); |
| 137 | + return mb; |
| 138 | + } |
| 139 | + |
| 140 | + public MatrixBlock startReadingMbFromPipes(int[] blockSizes, int rlen, int clen, Types.ValueType type) throws ExecutionException, InterruptedException { |
| 141 | + long limit = (long) rlen * clen; |
| 142 | + if(limit > Integer.MAX_VALUE) |
| 143 | + throw new DMLRuntimeException("Dense NumPy array of size " + limit + |
| 144 | + " cannot be converted to MatrixBlock"); |
| 145 | + MatrixBlock mb = new MatrixBlock(rlen, clen, false, -1); |
| 146 | + if(fromPython != null){ |
| 147 | + ExecutorService pool = CommonThreadPool.get(); |
| 148 | + double[] denseBlock = new double[(int) limit]; |
| 149 | + int offsetOut = 0; |
| 150 | + List<Future<Void>> futures = new ArrayList<>(); |
| 151 | + for (int i = 0; i < blockSizes.length; i++) { |
| 152 | + BufferedInputStream pipe = fromPython.get(i); |
| 153 | + int id = i, blockSize = blockSizes[i], _offsetOut = offsetOut; |
| 154 | + Callable<Void> task = () -> { |
| 155 | + UnixPipeUtils.readNumpyArrayInBatches(pipe, id, BATCH_SIZE, blockSize, type, denseBlock, _offsetOut); |
| 156 | + return null; |
| 157 | + }; |
| 158 | + |
| 159 | + futures.add(pool.submit(task)); |
| 160 | + offsetOut += blockSize; |
| 161 | + } |
| 162 | + // Wait for all tasks and propagate exceptions |
| 163 | + for (Future<Void> f : futures) { |
| 164 | + f.get(); |
| 165 | + } |
| 166 | + |
| 167 | + mb.init(denseBlock, rlen, clen); |
| 168 | + } else { |
| 169 | + throw new DMLRuntimeException("FIFO Pipes are not initialized."); |
| 170 | + } |
| 171 | + mb.recomputeNonZeros(); |
| 172 | + mb.examSparsity(); |
| 173 | + return mb; |
| 174 | + } |
| 175 | + |
| 176 | + public void startWritingMbToPipe(int id, MatrixBlock mb) throws IOException { |
| 177 | + if (toPython != null) { |
| 178 | + int rlen = mb.getNumRows(); |
| 179 | + int clen = mb.getNumColumns(); |
| 180 | + int numElem = rlen * clen; |
| 181 | + LOG.debug("Trying to write matrix ["+baseDir + "-"+ id+"] with "+rlen+" rows and "+clen+" columns. Total size: "+numElem*8); |
| 182 | + |
| 183 | + BufferedOutputStream out = toPython.get(id); |
| 184 | + long bytes = UnixPipeUtils.writeNumpyArrayInBatches(out, id, BATCH_SIZE, numElem, Types.ValueType.FP64, mb); |
| 185 | + |
| 186 | + LOG.debug("Writing of " + bytes +" Bytes to Python ["+baseDir + "-"+ id+"] finished"); |
| 187 | + } else { |
| 188 | + throw new DMLRuntimeException("FIFO Pipes are not initialized."); |
| 189 | + } |
| 190 | + } |
| 191 | + |
| 192 | + public void closePipes() throws IOException { |
| 193 | + LOG.debug("Closing all pipes in Java"); |
| 194 | + for (BufferedInputStream pipe : fromPython.values()) |
| 195 | + pipe.close(); |
| 196 | + for (BufferedOutputStream pipe : toPython.values()) |
| 197 | + pipe.close(); |
| 198 | + LOG.debug("Closed all pipes in Java"); |
| 199 | + } |
| 200 | + |
81 | 201 | protected static class DMLGateWayListener extends DefaultGatewayServerListener { |
82 | 202 | private static final Log LOG = LogFactory.getLog(DMLGateWayListener.class.getName()); |
83 | 203 |
|
|
0 commit comments