...
- The DataSet::iterate(...) only supports iteration on the bounded data streams. And as described in FLIP-131, we will deprecate the DataSet API in favor of the Table API/SQL in the future.
- The DataStream::iterate(...) has a few design issues that prevents it from being reliably used in production jobs. Many of these issues, such as possibility of deadlock, are described in FLIP-15.
And there is another performance issue with the DataSet::iterate(...) API: in order to replay the user-provided data streams multiple times, it requires the runtime to always dump the user-provided data streams to disk. This introduces storage and disk I/O overhead even if user's algorithm may prefer to cache those values in-memory and in possibly a more compact format.
In order to address these all the issues and improve the relevant Flink design as much as possibledescribed above, and make Flink ML available for more iteration use-case in the long run, this FLIP proposes to add a couple APIs in the flink-ml repository to achieve the following goals:
- Provide solution for all the iteration use-cases (see the use-case section below for more detail) supported by the existing APIs, without having the issues described above.
- Provide solution for a few use-cases (e.g. bounded streams + async mode + per-round variable update) not supported by the existing APIs.
- Decouple the iteration-related APIs from core Flink core runtime (by moving them to the flink-ml repo) so that we can keep the Flink core runtime as simple and maintainable as possible.
Terminology
- Provide iteration API that does not enforce the disk I/O overhead described above, so that users can optimize an iterative algorithm for best possible performance.
Terminology
We explain a few terminologies We explain a few terminologies in the following to facilitate the understanding of this doc.
...
Public Interfaces
We propose to add make the following APIs based on API changes to support the iteration paradigm described above.
...
Code Block | ||||||
---|---|---|---|---|---|---|
| ||||||
package org.apache.flink.ml.iteration; /** * A helper class to apply {@link IterationBody} to data streams. */ @PublicEvolving public class IterationUtils { /** * This method can use an iteration body to process records in unbounded data streams. * * This method invokes the iteration body with the following parameters: * 1) The 1st parameter is a list of input variable streams, which are created as the union of the initial variable * streams and the corresponding feedback variable streams (returned by the iteration body). * 2) The 2nd parameter is the data streams given to this method. * * The epoch values are determined as described below. See IterationListener for how the epoch values are used. * 1) All records in the initial variable streams has epoch=0. * 2) All records in the data streams has epoch=MAX_LONG. In this case, records in the data stream won't affect * any operator's low watermark. * 3) For any record emitted by this operator into a non-feedback stream, the epoch of this emitted record = the * epoch of the input record that triggers this emission. If this record is emitted by onEpochLWIncremented(), then * the epoch of this record = incrementedEpochLW - 1. * 4) For any record emitted by this operator into a feedback variable stream, the epoch of the emitted record = * min(the epoch of the input record that triggers this emission, MAX_LONG - 1) + 1. If this record is emitted by * onEpochLWIncremented(), then the epoch of this record = incrementedEpochLW. * * The execution of the graph created by the iteration body will not terminate by itself. This is because at least * one of its data streams is unbounded. * * Required: * 1) All the init variable streams must be bounded. * 2) There is at least one unbounded stream in the data streams list. * 3) The parallelism of any stream in the initial variable streams must equal the parallelism of the stream at the * same index of the feedback variable streams returned by the IterationBody. * * @param initVariableStreams The initial variable streams. These streams will be merged with the feedback variable * streams before being used as the 1st parameter to invoke the iteration body. * @param dataStreams The data streams. These streams will be used as the 2nd parameter to invoke the iteration * body. * @param body The computation logic which takes variable/data streams and returns variable/output streams. * @return The list of output streams returned by the iteration boy. */ static DataStreamList iterateUnboundedStreams(DataStreamList initVariableStreams, DataStreamList dataStreams, IterationBody body) {...} /** * This method can use an iteration body to process records in some bounded data streams iteratively until a * termination criteria is reached (e.g. the given number of rounds is completed or no further variable update is * needed). Because this method does not replay records in the data streams, the iteration body needs to cache those * records in order to visit those records repeatedly. * * This method invokes the iteration body with the following parameters: * 1) The 1st parameter is a list of input variable streams, which are created as the union of the initial variable * streams and the corresponding feedback variable streams (returned by the iteration body). * 2) The 2nd parameter is the data streams given to this method. * * The epoch values are determined as described below. See IterationListener for how the epoch values are used. * 1) All records in the initial variable streams has epoch=0. * 2) All records in the data streams has epoch=0. * 3) For any record emitted by this operator into a non-feedback stream, the epoch of this emitted record = the * epoch of the input record that triggers this emission. If this record is emitted by onEpochLWIncremented(), then * the epoch of this record = incrementedEpochLW - 1. * 4) For any record emitted by this operator into a feedback variable stream, the epoch of the emitted record = the * epoch of the input record that triggers this emission + 1. If this record is emitted by onEpochLWIncremented(), * then the epoch of this record = incrementedEpochLW. * * Suppose there is a coordinator operator which takes all feedback variable streams (emitted by the iteration body) * and the termination criteria stream (if not null) as inputs. The execution of the graph created by the * iteration body will terminate when all input streams have been fully AND any of the following conditions is met: * consumed: * 1) The termination criteria stream is not null. And the coordinator operator has not observed any new value from * the termination criteria stream between two consecutive onEpochLWIncremented invocations. * 2) The coordinator operator has not observed any new value from any feedback variable stream between two * consecutive onEpochLWIncremented invocations. * * Required: * 1) All the init variable streams and the data streams must be bounded. * 2) The parallelism of any stream in the initial variable streams must equal the parallelism of the stream at the * same index of the feedback variable streams returned by the IterationBody. * * @param initVariableStreams The initial variable streams. These streams will be merged with the feedback variable * streams before being used as the 1st parameter to invoke the iteration body. * @param dataStreams The data streams. These streams will be used as the 2nd parameter to invoke the iteration * body. * @param body The computation logic which takes variable/data streams and returns variable/output streams. * @return The list of output streams returned by the iteration boy. */ static DataStreamList iterateBoundedStreamsUntilTermination(DataStreamList initVariableStreams, DataStreamList dataStreams, IterationBody body) {...} /** * This method can use an iteration body to process records in some bounded data streams iteratively until a * termination criteria is reached (e.g. the given number of rounds is completed or no further variable update is * needed). Because this method replays records in the data streams, the iteration body does not need to cache those * records to visit those records repeatedly. * * This method invokes the iteration body with the following parameters: * 1) The 1st parameter is a list of input variable streams, which are created as the union of the initial variable * streams and the corresponding feedback variable streams (returned by the iteration body). * 2) The 2nd parameter is a list of replayed data streams, which are created by replaying the initial data streams * round by round until the iteration terminates. The records in the Nth round will be emitted into the iteration * body only if the low watermark of the first operator in the iteration body >= N - 1. * * The epoch values are determined as described below. See IterationListener for how the epoch values are used. * 1) All records in the initial variable streams has epoch=0. * 2) The records from the initial data streams will be replayed round by round into the iteration body. The records * in the first round have epoch=0. And records in the Nth round have epoch = N - 1. * 3) For any record emitted by this operator into a non-feedback stream, the epoch of this emitted record = the * epoch of the input record that triggers this emission. If this record is emitted by onEpochLWIncremented(), then * the epoch of this record = incrementedEpochLW - 1. * 4) For any record emitted by this operator into a feedback stream, the epoch of the emitted record = the epoch * of the input record that triggers this emission + 1. If this record is emitted by onEpochLWIncremented(), then * the epoch of this record = incrementedEpochLW. * * Suppose there is a coordinator operator which takes all feedback variable streams (emitted by the iteration body) * and the termination criteria stream (if not null) as inputs. The execution of the graph created by the * iteration body will terminate when all input streams have been fully consumed AND any of the following conditions * is met: * 1) The termination criteria stream is not null. And the coordinator operator has not observed any new value from * the termination criteria stream between two consecutive onEpochLWIncremented invocations. * 2) The coordinator operator has not observed any new value from any feedback variable stream between two * consecutive onEpochLWIncremented invocations. * * Required: * 1) All the init variable streams and the data streams must be bounded. * 2) The parallelism of any stream in the initial variable streams must equal the parallelism of the stream at the * same index of the feedback variable streams returned by the IterationBody. * * @param initVariableStreams The initial variable streams. These streams will be merged with the feedback variable * streams before being used as the 1st parameter to invoke the iteration body. * @param initDataStreams The initial data streams. Records from these streams will be repeatedly replayed and used * as the 2nd parameter to invoke the iteration body. * @param body The computation logic which takes variable/data streams and returns variable/output streams. * @return The list of output streams returned by the iteration boy. */ static DataStreamList iterateAndReplayBoundedStreamsUntilTermination(DataStreamList initVariableStreams, DataStreamList initDataStreams, IterationBody body) {...} } |
...
Code Block | ||||||
---|---|---|---|---|---|---|
| ||||||
package org.apache.flink.ml.iteration; public class DataStreamList { // Returns the number of data streams in this list. public int size() {...} // Returns the data stream at the given index in this list. public <T> DataStream<T> get(int index) {...} } |
Overall Design
To reduce the development and maintenance overhead, it would be preferred to have a unified implementation for different types of iterations. In fact, the different iteration types shares the same requirements in runtime implementation:
- All the iteration types should support multiple inputs and multiple outputs.
- All the iteration types require some kind of back edges that transfer the data back to the iteration head. Since Flink does not support cycles in scheduler and network stack, the back edges should not be visible in the StreamGraph and JobGraph.
- All the iteration should support checkpoints mechanism in Stream execution mode.
Different types of iterations differ in their requirements for Progress tracking. Progress tracking is analogous to the watermark outside the iteration and it tracks the “progress” inside the iteration:
- For bounded iteration, we could track if we have processed all the records for a specific round. This is necessary for operators like aggregation inside the iteration: if it is notified all the records of the current round is processed, it could output the result of this round. We could also track if the whole iteration is end, namely all the inputs are finished and no pending records inside the iteration.
- For unbounded iteration, there is no concept of global rounds, and the only progress tracking is at the end of iteration.
The difference of the progress tracking would also affect the API. For example, for bounded iteration, we could allow users to specify the termination condition based on number of rounds, but it is meaningless for the unbounded iteration.
To make the API easy to use, we propose to have dedicated API for different types of iteration, and underlying we will translate them onto the same framework. would implements the basic functionality like iteration StreamGraph building, runtime structure and checkpoint, and it allows to implement different iterations to implement different types of progress tracking support.
Public Interfaces
As shown in Figure 1, an iteration is composed of
- The inputs from outside of the iteration.
- An iteration body specify the structure inside the iteration.
- The subgraph inside the iteration.
- Some input have corresponding feedbacks to update the underlying data stream. The feedbacks are union with the corresponding inputs: the original inputs are emitted into the iteration body for only once, and the feedbacks are also emitted to the same set of operators.
- The outputs going out of the iteration. The outputs could be emitted from arbitrary data stream.
Unbounded Iteration
Similar to FLIP-15, we would more tend to provide a structural iteration API to make it easier to be understand. With this method, users are required to specify an IterationBody that generates the part of JobGraph inside the iteration. The iteration body should specify the DAG inside the iteration, and also the list of feedback streams and the output streams. The feedback streams would be union with the corresponding inputs and the output streams would be provided to the caller routine.
However, since we do not know the accurate number and type of input streams, it is not easy to define a unified interface for the iteration body without type casting. Thus we would propose to use the annotation to allows for arbitrary number of inputs:
The interface for the unbounded iteration is straightforward:
To avoid more data is read from the inputs while too much data accumulate inside the iteration, the iteration would first process the feedback data if both side of data is available.
For termination detection, the iteration would continue until
- All the inputs are terminated.
- And there is no records inside the iteration subgraph.
Then the iteration terminates.
Bounded Iteration
As mentioned in the motivation, the existing dataset iteration API uses the "per-round" semantics: it views the iteration as a repeat execution of the same DAG, thus underlying it would automatically merge the inputs and feedbacks and replay the inputs without feedbacks, and the operators inside the iteration live only for one-round. This might cause bad performance for some algorithms who could cache these data in a more efficient way.
To avoid this issue, similar to the unbounded iteration, by default we use the "per-iteration" semantics:
- Operators inside the iteration would live till the whole iteration is finished.
- We do not automatically merge the inputs and feedbacks. Instead, we union the original inputs and the feedbacks so that users could decide how the merge them.
- We do not replay the inputs without feedbacks. Users could decide to how to cache them more efficiently.
Besides, to cooperate with the "per-round" semantics, previously the iteration is by default synchronous: before the current round fully finished, the feedback data is cached and would not be emitted. Thus it could not support some algorithms like asynchronous regression. To cope with this issue, we view synchronous iteration as a special case of asynchronous iteration with additional synchronization. Thus by default the iteration is asynchronous.
Based on the above assumption, the API to add iteration to a job is nearly the same compared to the unbounded iteration. The only difference is that bounded iteration supports more sophisticated termination conditions: a function is evaluated when each round ends based on the round or the records of a specified data stream. If it returns true, the iteration would deserts all the following feedback records, ends all the ongoing rounds and finish.
Since now the operators would live across multiple rounds and multiple rounds might be concurrent, the operators inside the iteration needs to know the rounds of the current record and when one round is fully finished, namely the progress tracking. For example, an operator computes the sum of the records in each rounds would like to add the record to the corresponding partial sum, and when one round is finished, it would emit the sum for this round. To support the progress track, UDFs / operators inside the iteration could implementation `BoundedIterationProgressListener` to acquire the additional information about the progress.
Based on the progress tracking interface, if users want to implement a synchronous method, some operators inside the subgraph needs to be synchronous: they only emits the records in `onRoundEnd`, namely after all the data of the current round is received. If for the subgraph of iteration body, every path from input to the feedbacks has at least such an operator, then the iteration would be synchronous.
For users still want to use the iteration with the "per-round" semantics, a utility `forEachRound()` is provided. With the utility users could add a subgraph inside the iteration body that
- The operators inside the subgraph would live only for one round.
- If an input stream without feedback is referenced, the input stream would be replayed for each round.
For input stream with feedbacks, we also provide two utility processFunction that automatically merge the original inputs and feedbacks. Both the existing bulk and delta method is supported. Then users would be able to implement a per-round iteration with input.process(bulkCache()).forEachRound(() → {...}).
The API for the bounded iteration is as follows:
Key Implementation
The runtime physical structure of the iteration is shown in Figure 1, which is similar to the current implementation. The head & tail is added by the framework. They would be colocated so that we could implement the feedback edge with the local queue. The head could coordinator with an operator coordinator bind to a virtual operator ID for synchronization, including progress tracking and termination condition calculating.
Figure 1. The physical runtime structure for the iteration.
To support the progress tracking, we would introduce new events inside the iteration body, like how watermark is implemented. However, since the normal operators could not identify these event, we would wrap the operators inside the iteration to parse these events.
To wrap the operators for the part of DAG inside the iteration, when building the stream graph we would introduce a mock execution environment and build the iteration DAG inside this environment first, then when apply() method is called, we would translate the DAG into the real execution environment with the suitable wrapper. Besides, all the edges inside the iteration should be PIPELINE, we would also set the edge property when translating.
The operator wrapper needs to simulates the context that an operator executes. Specially, for operators with single-round lifecycle in bounded iteration, we would need to isolate the states used for each round and cleanup the corresponding state after the round end.
Examples
This sections shows how general used ML algorithms could be implemented with the iteration API.
Offline Training with Bounded Iteration
We would like to first show the usage of the bounded iteration with the linear regression case: the model is Y = XA, and we would like to acquire the best estimation of A with the SGD algorithm. To simplify we assume the parameters could be held in the memory of one task.
The job graph of the algorithm could be shown in the Figure 3: in each round, we use the latest parameters to calculate the update to the parameters: ΔA = ∑(Y - XA)X. To achieve this, the Parameters vertex would broadcast the latest parameters to the Train vertex. Each subtask of the Train vertex holds a part of dataset. Follow the sprite of SGD, it would sample a small batch of training records, and calculate the update with the above equation. Then the Train vertex emit ΔA to the Parameters node to update the parameters.
Figure 3. The JobGraph for the offline training of the linear regression case.
We will start with the synchronous training. The synchronous training requires the updates from all the Train vertex subtask is merged before the next round of training. It could be done by only emit the next round of parameters on the end of round. The code is shown as follows:
6) Deprecate the existing DataStream::iterate() and the DataStream::iterate(long maxWaitTimeMillis) methods.
We plan to remove both methods after the APIs added in this doc is ready for production use. This change is needed to decouple the iteration-related APIs from core Flink core runtime so that we can keep the Flink core runtime as simple and maintainable as possible.
Proposed Changes
1) Termination of the iteration execution.
See the Java doc of those APIs in the IterationUtils for how each API determine the iteration termination.
2) Execution mode.
If all inputs are bounded streams, then the iteration body can be executed in either the stream mode or the batch mode.
If some inputs are unbounded streams, then the iteration body must be executed in the stream mode.
3) Type of edges inside the iteration body.
All edges inside the iteration body are required to have the PIPELINE type.
If the user-defined iteration body contains an edge that does not have the PIPELINE type, methods that create the subgraph from the iteration body, such as iterateBoundedStreamsUntilTermination, will throw exception upon invocation.
4) Implementation of the feedback edge.
The Flink core runtime supports only DAG of operators. Thus it does not provide native support for feedback edges since feedback edges introduce circle in the operator graph.
Same as the implementation of the DataSet::iterate() API, the proposed APIs are implemented with the following trick:
- Automatically insert HEAD and TAIL operators as the first and last operators in the iteration body.
- Co-locate HEAD and TAIL operators on the same task manager.
- Have HEAD and TAIL operators transmit the records of the feedback edges using an in-memory queue.
5) Lifetime of the operators inside the iteration body.
The operator inside the iteration body are only created once and destroyed after the iteration terminates. In contrast, the existing DataSet::iterate(..) would re-create the iteration body (together with all states inside it) once for every round of execution, which in general could introduce more runtime overhead then the approach adopted in this FLIP.
6) Failover
To be described.
7) Support for synchronous iteration.
To be descried
Example Usages
Offline Training with Bounded Iteration
We would like to first show the usage of the bounded iteration with the linear regression case: the model is Y = XA, and we would like to acquire the best estimation of A with the SGD algorithm. To simplify we assume the parameters could be held in the memory of one task.
The job graph of the algorithm could be shown in the Figure 3: in each round, we use the latest parameters to calculate the update to the parameters: ΔA = ∑(Y - XA)X. To achieve this, the Parameters vertex would broadcast the latest parameters to the Train vertex. Each subtask of the Train vertex holds a part of dataset. Follow the sprite of SGD, it would sample a small batch of training records, and calculate the update with the above equation. Then the Train vertex emit ΔA to the Parameters node to update the parameters.
We will start with the synchronous training. The synchronous training requires the updates from all the Train vertex subtask is merged before the next round of training. It could be done by only emit the next round of parameters on the end of round. The code is shown as follows:
Code Block | ||||
---|---|---|---|---|
| ||||
public class SynchronousBoundedLinearRegression {
private static final N_DIM | ||||
Code Block | ||||
| ||||
public class SynchronousBoundedLinearRegression {
private static final N_DIM = 50;
private static final OutputTag<double[]> FINAL_MODEL_OUTPUT_TAG = new OutputTag<double[]>{};
public static void main(String[] args) {
DataStream<double[]> initParameters = loadParameters().setParallelism(1);
DataStream<Tuple2<double[], Double>> dataset = loadDataSet().setParallelism(1);
int batch = 5;
int epochEachBatch = 10;
ResultStreams resultStreams = new BoundedIteration()
.withBody(new IterationBody(
@IterationInput("model") DataStream<double[]> model,
@IterationInput("dataset") DataStream<Tuple2<double[], Double>> dataset
) {
SingleOutputStreamOperator<double[]> parameters = model.process(new ParametersCacheFunction());
DataStream<double[]> modelUpdate = parameters.setParallelism(1)
.broadcast()
.connect(dataset)
.coProcess(new TrainFunction())
.setParallelism(10)
return new BoundedIterationDeclarationBuilder()
.withFeedback("model", modelUpdate)
.withOutput("final_model", parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG))
.until(new TerminationCondition(null, context -> context.getRound() >= batch * epochEachBatch))
.build();
})
.build();
DataStream<double[]> finalModel = resultStreams.get("final_model");
finalModel.print();
}
public static class ParametersCacheFunction extends ProcessFunction<double[], double[]>
implements BoundedIterationProgressListener<double[]> {
private final double[] parameters = new double[N_DIM];
public void processElement(double[] update, Context ctx, Collector<O> output) {
// Suppose we have a util to add the second array to the first.
ArrayUtils.addWith(parameters, update);
}
public void onRoundEnd(int[] round, Context context, Collector<T> collector) {
collector.collect(parameters);
}
public void onIterationEnd(int[] round, Context context) {
context.output(FINAL_MODEL_OUTPUT_TAG, parameters);
}
}
public static class TrainFunction extends CoProcessFunction<double[], Tuple2<double[], Double>, double[]> implements BoundedIterationProgressListener<double[]> {
private final List<Tuple2<double[], Double>> dataset = new ArrayList<>();
private double[] firstRoundCachedParameter;
private Supplier<int[]> recordRoundQuerier;
public void setCurrentRecordRoundsQuerier(Supplier<int[]> querier) {
this.recordRoundQuerier = querier;
}
public void processElement1(double[] parameter, Context context, Collector<O> output) {
int[] round = recordRoundQuerier.get();
if (round[0] == 0) {
firstRoundCachedParameter = parameter;
return;
}
calculateModelUpdate(parameter, output);
}
public void processElement2(Tuple2<double[], Double> trainSample, Context context, Collector<O> output) {
dataset.add(trainSample)
}
public void onRoundEnd(int[] round, Context context, Collector<T> output) {
if (round[0] == 0) {
calculateModelUpdate(firstRoundCachedParameter, output);
firstRoundCachedParameter = null;
}
}
private void calculateModelUpdate(double[] parameters, Collector<O> output) {
List<Tuple2<double[], Double>> samples = sample(dataset);
double[] modelUpdate = new double[N_DIM];
for (Tuple2<double[], Double> record : samples) {
double diff = (ArrayUtils.muladd(record.f0, parameters) - record.f1);
ArrayUtils.addWith(modelUpdate, ArrayUtils.multiply(record.f0, diff));
}
output.collect(modelUpdate);
}
}
} |
...
Code Block | ||
---|---|---|
| ||
public static class ParametersCacheFunction extends ProcessFunction<Tuple2<Integer, double[]>, Tuple2<Integer, double[]>> {
private final double[] parameters = new double[N_DIM];
public void processElement(Tuple2<Integer, double[]> update, Context ctx, Collector<Tuple2<Integer, double[]>> output) {
ArrayUtils.addWith(parameters, update);
if (update.f0 < 0) {
// Received the initialized parameter values, broadcast to all the downstream tasks
for (int i = 0; i < 10; ++i) {
output.collect(new Tuple2<>(i, parameters))
}
} else {
output.collect(new Tuple2<>(update.f0, parameters))
}
}
} |
Implementation Plan
Logically all the iteration types would support both BATCH and STREAM execution mode. However, according to the algorithms' requirements, we would implement
- Unbounded iteration + STREAM mode.
- Bounded iteration + BATCH mode.
...
(new Tuple2<>(update.f0, parameters))
}
}
} |
Compatibility, Deprecation, and Migration Plan
The API is added as a library inside flink-ml repository, thus it does not have compatibility problem. However, it has some difference with the existing iteration API and the algorithms would need some re-implementation.
For the long run, the new iteration implementation might provide an alternative for the iteration functionality, and we may consider deprecating and removing the existing API to reduce the complexity of core flink code. TBD
Rejected Alternatives
Naiad has proposed a unified model for watermark mechanism (namely progress tracking outside of the iteration) and the progress tracking inside the iteration. It extends the event time and watermark to be a vector (long timestamp, int[] rounds) and implements a vectorized alignment algorithm. Although Naiad provides an elegant model, the direct implementation on Flink would requires a large amount of modification to the flink runtime, which would cause a lot of complexity and maintenance overhead. Thus we would choose to implement a simplified version on top of FLINK, as a part of the flink-ml library.
...