remove context object and use class members instead

Author: zhipeng93

flink-ml-lib/src/main/java/org/apache/flink/ml/classification/logisticregression/LogisticRegressionWithFtrl.java

@@ -34,6 +34,7 @@
 import org.apache.flink.ml.common.ps.training.ProcessStage;
 import org.apache.flink.ml.common.ps.training.PullStage;
 import org.apache.flink.ml.common.ps.training.PushStage;
+import org.apache.flink.ml.common.ps.training.SerializableConsumer;
 import org.apache.flink.ml.common.ps.training.TrainingContext;
 import org.apache.flink.ml.common.ps.training.TrainingUtils;
 import org.apache.flink.ml.common.updater.FTRL;
@@ -51,6 +52,7 @@
 import org.apache.flink.types.Row;
 import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.function.SerializableFunction;
+import org.apache.flink.util.function.SerializableSupplier;
 
 import it.unimi.dsi.fastutil.longs.Long2DoubleOpenHashMap;
 import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
@@ -128,15 +130,25 @@ public LogisticRegressionModel fit(Table... inputs) {
                             .map((MapFunction<Long, Long>) value -> value + 1);
         }
 
-        IterationStageList<LogisticRegressionTrainingContext> iterationStages =
-                new IterationStageList<>(new LogisticRegressionTrainingContext(getParamMap()));
+        LogisticRegressionWithFtrlTrainingContext trainingContext =
+                new LogisticRegressionWithFtrlTrainingContext(getParamMap());
+
+        IterationStageList<LogisticRegressionWithFtrlTrainingContext> iterationStages =
+                new IterationStageList<>(trainingContext);
         iterationStages
                 .addTrainingStage(new ComputeIndices())
-                .addTrainingStage(new PullStage("pullIndices"))
+                .addTrainingStage(
+                        new PullStage(
+                                (SerializableSupplier<long[]>) () -> trainingContext.pullIndices,
+                                (SerializableConsumer<double[]>)
+                                        x -> trainingContext.pulledValues = x))
                 .addTrainingStage(new ComputeGradients(BinaryLogisticLoss.INSTANCE))
-                .addTrainingStage(new PushStage("pushGradient"))
+                .addTrainingStage(
+                        new PushStage(
+                                (SerializableSupplier<long[]>) () -> trainingContext.pushIndices,
+                                (SerializableSupplier<double[]>) () -> trainingContext.pushValues))
                 .setTerminationCriteria(
-                        (SerializableFunction<LogisticRegressionTrainingContext, Boolean>)
+                        (SerializableFunction<LogisticRegressionWithFtrlTrainingContext, Boolean>)
                                 o -> o.iterationId >= getMaxIter());
         FTRL ftrl = new FTRL(getAlpha(), getBeta(), getReg(), getElasticNet());
 
@@ -183,16 +195,15 @@ public Map<Param<?>, Object> getParamMap() {
 }
 
 /**
- * A stage that samples a batch of training data and computes the indices needed to compute
- * gradients.
+ * An iteration stage that samples a batch of training data and computes the indices needed to
+ * compute gradients.
  */
-class ComputeIndices extends ProcessStage<LogisticRegressionTrainingContext> {
+class ComputeIndices extends ProcessStage<LogisticRegressionWithFtrlTrainingContext> {
 
     @Override
-    public void process(LogisticRegressionTrainingContext context) throws Exception {
+    public void process(LogisticRegressionWithFtrlTrainingContext context) throws Exception {
         context.readInNextBatchData();
-        long[] indices = computeIndices(context.batchData);
-        context.put("pullIndices", indices);
+        context.pullIndices = computeIndices(context.batchData);
     }
 
     public static long[] computeIndices(List<LabeledLargePointWithWeight> dataPoints) {
@@ -216,24 +227,24 @@ public static long[] computeIndices(List<LabeledLargePointWithWeight> dataPoints
 }
 
 /**
- * A Stage that uses the pulled model parameters and batch data to compute the gradients. The
- * gradients are stored in context for later push.
+ * An iteration stage that uses the pulled model values and sampled batch data to compute the
+ * gradients.
  */
-class ComputeGradients extends ProcessStage<LogisticRegressionTrainingContext> {
+class ComputeGradients extends ProcessStage<LogisticRegressionWithFtrlTrainingContext> {
     private final LossFunc lossFunc;
 
     public ComputeGradients(LossFunc lossFunc) {
         this.lossFunc = lossFunc;
     }
 
     @Override
-    public void process(LogisticRegressionTrainingContext context) throws IOException {
+    public void process(LogisticRegressionWithFtrlTrainingContext context) throws IOException {
         long[] indices = ComputeIndices.computeIndices(context.batchData);
-        double[] pulledModelValues = (double[]) context.get("pullIndices");
+        double[] pulledModelValues = context.pulledValues;
         double[] gradients = computeGradient(context.batchData, indices, pulledModelValues);
 
-        // Saves the gradient for push.
-        context.put("pushGradient", Tuple2.of(indices, gradients));
+        context.pushIndices = indices;
+        context.pushValues = gradients;
     }
 
     private double[] computeGradient(
@@ -267,37 +278,34 @@ private static double dot(
     }
 }
 
-class LogisticRegressionTrainingContext
+/** The context information of local computing process. */
+class LogisticRegressionWithFtrlTrainingContext
         implements TrainingContext,
-                LogisticRegressionWithFtrlParams<LogisticRegressionTrainingContext> {
+                LogisticRegressionWithFtrlParams<LogisticRegressionWithFtrlTrainingContext> {
+    /** Parameters of LogisticRegressionWithFtrl. */
+    private final Map<Param<?>, Object> paramMap;
+    /** Current iteration id. */
     int iterationId;
-    int workerId;
-    private int numWorkers = -1;
+    /** The local batch size. */
     private int localBatchSize = -1;
-
-    ResettableIterator<LabeledLargePointWithWeight> trainData;
-    private final Map<String, Object> contextObjs = new HashMap<>();
-
-    ListState<LabeledLargePointWithWeight> batchDataState;
+    /** The training data. */
+    private ResettableIterator<LabeledLargePointWithWeight> trainData;
+    /** The batch of training data for computing gradients. */
     List<LabeledLargePointWithWeight> batchData;
 
-    private final Map<Param<?>, Object> paramMap;
+    private ListState<LabeledLargePointWithWeight> batchDataState;
 
-    public LogisticRegressionTrainingContext(Map<Param<?>, Object> paramMap) {
-        this.paramMap = paramMap;
-    }
+    /** The indices to pull for each iteration. */
+    long[] pullIndices;
+    /** The placeholder for the pulled values for each iteration. */
+    double[] pulledValues;
+    /** The indices to push for each iteration. */
+    long[] pushIndices;
+    /** The values to push for each iteration. */
+    double[] pushValues;
 
-    /** Reads in next batch of training data. */
-    public void readInNextBatchData() throws IOException {
-        batchData.clear();
-        int i = 0;
-        while (i < localBatchSize && trainData.hasNext()) {
-            batchData.add(trainData.next());
-            i++;
-        }
-        if (!trainData.hasNext()) {
-            trainData.reset();
-        }
+    public LogisticRegressionWithFtrlTrainingContext(Map<Param<?>, Object> paramMap) {
+        this.paramMap = paramMap;
     }
 
     @Override
@@ -307,8 +315,6 @@ public void setIterationId(int iterationId) {
 
     @Override
     public void setWorldInfo(int workerId, int numWorkers) {
-        this.workerId = workerId;
-        this.numWorkers = numWorkers;
         int globalBatchSize = getGlobalBatchSize();
         this.localBatchSize = globalBatchSize / numWorkers;
         if (globalBatchSize % numWorkers > workerId) {
@@ -322,16 +328,6 @@ public void setTrainData(ResettableIterator<?> trainData) {
         this.trainData = (ResettableIterator<LabeledLargePointWithWeight>) trainData;
     }
 
-    @Override
-    public void put(String key, Object value) {
-        contextObjs.put(key, value);
-    }
-
-    @Override
-    public Object get(String key) {
-        return contextObjs.get(key);
-    }
-
     @Override
     public void initializeState(StateInitializationContext context) throws Exception {
         batchDataState =
@@ -343,7 +339,7 @@ public void initializeState(StateInitializationContext context) throws Exception
 
         Iterator<LabeledLargePointWithWeight> batchDataIterator = batchDataState.get().iterator();
         if (batchDataIterator.hasNext()) {
-            contextObjs.put("batchData", IteratorUtils.toList(batchDataIterator));
+            batchData = IteratorUtils.toList(batchDataIterator);
         }
     }
 
@@ -359,4 +355,17 @@ public void snapshotState(StateSnapshotContext context) throws Exception {
     public Map<Param<?>, Object> getParamMap() {
         return paramMap;
     }
+
+    /** Reads in next batch of training data. */
+    public void readInNextBatchData() throws IOException {
+        batchData.clear();
+        int i = 0;
+        while (i < localBatchSize && trainData.hasNext()) {
+            batchData.add(trainData.next());
+            i++;
+        }
+        if (!trainData.hasNext()) {
+            trainData.reset();
+        }
+    }
 }

flink-ml-lib/src/main/java/org/apache/flink/ml/common/ps/RangeModelPartitioner.java renamed to flink-ml-lib/src/main/java/org/apache/flink/ml/common/ps/RangePartitioner.java

@@ -18,24 +18,27 @@
 
 package org.apache.flink.ml.common.ps;
 
-import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.api.java.tuple.Tuple3;
-import org.apache.flink.ml.util.Bits;
 import org.apache.flink.util.Preconditions;
 
 import javax.annotation.Nullable;
 
 import java.util.Arrays;
 import java.util.Iterator;
 
-/** Range partitioner for vector. */
-public class RangeModelPartitioner {
-    public long dim;
-    public int numServers;
-    private long[] ranges;
+/** Range partitioner for model data. */
+public class RangePartitioner {
+    public final long dim;
+    public final int numServers;
+    public final long[] ranges;
+
+    public RangePartitioner(long dim, int numServers) {
+        Preconditions.checkArgument(
+                dim > 0,
+                String.format(
+                        "Illegal dimension when using %s: %d",
+                        RangePartitioner.class.getSimpleName(), dim));
 
-    public RangeModelPartitioner(long dim, int numServers) {
-        Preconditions.checkArgument(dim > 0 && numServers > 0);
         this.dim = dim;
         this.numServers = numServers;
         this.ranges = new long[numServers + 1];
@@ -47,38 +50,6 @@ public RangeModelPartitioner(long dim, int numServers) {
         ranges[numServers] = dim;
     }
 
-    public RangeModelPartitioner() {}
-
-    public Tuple2<Long, Long> getStartAndEnd(int serverId) {
-        return Tuple2.of(ranges[serverId], ranges[serverId + 1]);
-    }
-
-    public static int getNumBytes() {
-        return Long.BYTES + Integer.BYTES + Integer.BYTES;
-    }
-
-    public int writeToBytes(byte[] bytesData, int offset) {
-        Bits.putLong(bytesData, offset, dim);
-        offset += Long.BYTES;
-        Bits.putInt(bytesData, offset, numServers);
-        offset += Integer.BYTES;
-        return offset;
-    }
-
-    public byte[] toBytes() {
-        byte[] buffer = new byte[getNumBytes()];
-        writeToBytes(buffer, 0);
-        return buffer;
-    }
-
-    public static RangeModelPartitioner readFromBytes(byte[] bytesData, int offset) {
-        long dim = Bits.getLong(bytesData, offset);
-        offset += Long.BYTES;
-        int numPss = Bits.getInt(bytesData, offset);
-        offset += Integer.BYTES;
-        return new RangeModelPartitioner(dim, numPss);
-    }
-
     /**
      * Splits the push/pull request according to the given sorted indices and the corresponding
      * values.
@@ -93,45 +64,44 @@ public Iterator<Tuple3<Integer, long[], double[]>> splitRequest(
     }
 
     private static class RequestsIterator implements Iterator<Tuple3<Integer, long[], double[]>> {
-        private final int numPss;
+        private final int numServers;
         private final long[] indices;
         private final double[] values;
         private final long[] ranges;
 
-        private int psId = 0;
+        private int serverId = 0;
 
         private int s = 0;
 
         public RequestsIterator(
                 int numPss, long[] indices, @Nullable double[] values, long[] ranges) {
-            this.numPss = numPss;
+            this.numServers = numPss;
             this.indices = indices;
             this.values = values;
             this.ranges = ranges;
         }
 
         @Override
         public boolean hasNext() {
-            return psId < numPss;
+            return serverId < numServers;
         }
 
         @Override
         public Tuple3<Integer, long[], double[]> next() {
             int e = s;
-            while (e < indices.length && indices[e] < ranges[psId + 1]) {
+            while (e < indices.length && indices[e] < ranges[serverId + 1]) {
                 e++;
             }
 
-            // Also pushes the empty message for atomic of push/pull in async setting.
             long[] splitIndices = new long[0];
             double[] splitValues = values == null ? null : new double[0];
             if (s < e) {
                 splitIndices = Arrays.copyOfRange(indices, s, e);
                 splitValues = values == null ? null : Arrays.copyOfRange(values, s, e);
             }
             s = e;
-            psId++;
-            return Tuple3.of(psId - 1, splitIndices, splitValues);
+            serverId++;
+            return Tuple3.of(serverId - 1, splitIndices, splitValues);
         }
     }
 }