Page properties | ||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Status
Current state: <TODO>
Discussion thread: <TODO>
JIRA: <TODO>
This method is easier to use, but it also limit some possible optimizations
...
Please keep the discussion on the mailing list rather than commenting on the wiki (wiki discussions get unwieldy fast).
Table of Contents |
---|
[This FLIP proposal is a joint work between Yun Gao, Dong Lin and Zhipeng Zhang]
Motivation
Iteration is a basic building block for a ML library. It is required for training ML models for both offline and online Training. In general, two types of iterations is required:
...
Previously Flink supported bounded iteration with DataSet API and supported the unbounded iteration with DataStream API. However, since Flink aims to deprecate the DataSet API and the iteration in the DataStream API is rather incomplete, thus we would require to re-implement a new iteration library in the Flink-ml repository to support the algorithms.
The Goals
The Types of the Algorithms
In general a ML algorithm would update the model according to the data in iteration until the model is converged. According to the granularity of the dataset used to update the model, in general ML algorithms could be classified into two types:
- Epoch-based: each epoch means the algorithm goes through all the training dataset. The epoch-based algorithm must work with the bounded dataset.
- Batch-based: each batch means the algorithm samples a subset from all the records and used the sampled records to update the model. The batch-based algorithm could be work with the bounded or unbounded dataset.
In a distributed settings, the dataset might be partitioned onto multiple subtasks, then for each epoch we refer to each subtask goes through all the assigned records, and for batch we refer to each subtask sample from its assigned records. When update the model with the referred data in a distributed settings, there are further two styles:
- Synchronous: In synchronous pattern the model must wait for the updates from all the subtasks before it could be used in the computation of the next update in each subtask.
- Asynchronous: In asynchronous pattern the model could directly apply the updates from some subtasks, and uses the updated value in the following computation immediately.
In general the synchronous pattern would have higher accuracy and the asynchronous pattern would convergent faster.
Based on the above dimensions, the algorithms could be classified into the following types:
...
SGD-Based Synchronous Offline algorithm
...
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 all the issues described 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.
- 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.
Target Use-cases
In general a ML algorithm would update the model according to the data in iteration until the model is converged. The target algorithms could be classified w.r.t. three properties: boundedness of input datasets, amount of data required for each variable update and the synchronization policy.
1) Algorithms have different needs for whether the input data streams should be bounded or unbounded. We classify those algorithms into online algorithm and offline algorithms as below.
- For online training algorithms, the training samples will be unbounded streams of data. The corresponding iteration body should ingest these unbounded streams of data, read each value in each stream once, and update machine learning model repeatedly in near real-time. The iteration will never terminate in this case. The algorithm should be executed as a streaming job.
- For offline training algorithms, the training samples will be bounded streams of data. The corresponding iteration body should read these bounded data streams for arbitrary number of rounds and update machine learning model repeatedly until a termination criteria is met (e.g. a given number of rounds is reached or the model has converged). The algorithm should be executed as a batch job.
2) An algorithm may have additional requirements in how much data should be consumed each time before a subtask can update variables. There are two categories of choices here:
- Per-batch variable update: The algorithm wants to update variables every time an arbitrary subset of the user-provided data streams (either bounded or unbounded) is processed.
- Per-round variable update: The algorithm wants to update variables every time all data of the user-provided bounded data streams is processed.
In the machine learning domain, some algorithms allow users to configure a batch size and the model will be updated every time each subtask processes a batch of data. Those algorithms fits into the first category. And such an algorithm can be either online or offline.
Other algorithms only update variables every time the entire data is consumed for one round. Those algorithms fit into the second category. And such an algorithm must be offline because, by this definition, the user-provided dataset must be bounded.
3) Algorithms (either online or offline) have different needs of how their parallel subtasks.
- In the sync mode, parallel subtasks, which execute the iteration body, update the model variables in a coordinated manner. There exists global epochs, such that all subtasks fetch the shared model variable values at the beginning of an epoch, calculate model variable updates based on the fetched variable values, and updates the model variable values at the end of this epoch. In other words, all subtasks read and update model variables in global lock steps.
- In the async mode, each parallel subtask, which execute the iteration body, could read/update the shared model variables without waiting for variable updates from other subtask. For example, a subtask could have updated model variables 10 times when another subtask has updated model variables only 3 times.
The sync mode is useful when an algorithm should be executed in a deterministic way to achieve best possible accuracy, and the straggler issue (i.e. there is subtask which is considerably slower than others) does not cause slow down the algorithm execution too much. In comparison, the async mode is useful for algorithms which want to be parallelized and executed across many subtasks as fast as possible, without worrying about performance issue caused by stragglers, at the possible cost of reduced accuracy.
Based on the above dimensions, the algorithms could be classified into the following types:
Type | Bounded / Unbounded | Data Granularity | Synchronization Pattern | Support in the existing APIs | Support in the proposed API | Examples |
---|---|---|---|---|---|---|
Non-SGD-based | Bounded | Epoch | Mostly Synchronous | Yes | Yes | Offline K-Means, Apriori, Decision Tree |
SGD-Based Synchronous Offline algorithm | Bounded | Batch → Epoch* | Synchronous | Yes | Yes | Linear Regression, Logistic Regression, Deep Learning algorithms |
SGD-Based Asynchronous Offline algorithm | Bounded | Batch → Epoch* | Asynchronous | No | Yes | Same to the above |
SGD-Based Synchronous Online algorithm | Unbounded | Batch | Synchronous | Yes | Yes | Online version of the above algorithm |
SGD-Based Asynchronous Online algorithm | Unbounded | Batch | Asynchronous | No | Yes | Online version of the above algorithm |
*Although SGD-based algorithms are also batch-based, they could be implemented with an Epoch-based method if intermediate state is allowed: the subtasks could sample a batch from all the records from the position of the last batch.
Based on the above classification and the replacement implementation for SGD-based algorithms with bounded dataset, we mainly need to support
- The synchronous / asynchronous epoch-based algorithms on the bounded dataset.
- The synchronous / asynchronous batch-based algorithms on the unbounded dataset.
Overview of the Iteration Paradigm
Based on the types of algorithms, we explain the iteration paradigm that has motivated our choices of the proposed APIs.
An iterative algorithm has the following behavior pattern:
- The iterative algorithm has an iteration body that is repeatedly invoked until some termination criteria is reached (e.g. after a user-specified number of epochs has been reached). An iteration body is a subgraph of operators that implements the computation logic of e.g. an iterative machine learning algorithm, whose outputs might be be fed back as the inputs of this subgraph.
- In each invocation, the iteration body updates the model parameters based on the user-provided data as well as the most recent model parameters.
- The iterative algorithm takes as inputs the user-provided data and the initial model parameters.
- The iterative algorithm could output arbitrary user-defined information, such as the loss after each epoch, or the final model parameters.
Therefore, the behavior of an iterative algorithm could be characterized with the following iteration paradigm (w.r.t. Flink concepts):
- An iteration-body is a Flink subgraph with the following inputs and outputs:
- Inputs: model-variables (as a list of data streams) and user-provided-data (as another list of data streams)
- Outputs: feedback-model-variables (as a list of data streams) and user-observed-outputs (as a list of data streams)
- A termination-condition that specifies when the iterative execution of the iteration body should terminate.
- In order to execute an iteration body, a user needs to execute an iteration body the following inputs, and gets the following outputs.
- Inputs: initial-model-variables (as a list of bounded data streams) and user-provided-data (as a list of data streams)
- Outputs: the user-observed-output emitted by the iteration body.
It is important to note that the users typically should not invoke the IterationBody::process directly because the model-variables expected by the iteration body is not the same as the initial-model-variables provided by the user. Instead, model-variables are computed as the union of the feedback-model-variables (emitted by the iteration body) and the initial-model-variables (provided by the caller of the iteration body). To relieve user from creating this union operator, we have added utility class (see Iterations) to run an iteration-body with the user-provided inputs.
The figure below summarizes the iteration paradigm described above. The streams in the red color are inputs provided by the user to the iteration body, as well as outputs emitted by the iteration body to the user.
draw.io Diagram | ||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Public Interfaces
We propose to make the following API changes to support the iteration paradigm described above.
1) Add the IterationBody interface.
This interface corresponds to the iteration-body with the inputs and outputs described in the iteration paradigm. This interface should be implemented by the developer of the algorithm.
Note that the IterationBody also outputs the termination criteria which corresponds to the termination-condition described in the iteration paradigm. This allows the algorithm developer to use a stream created inside the IterationBody as the terminationCriteria.
Code Block | ||||
---|---|---|---|---|
| ||||
package org.apache.flink.iteration;
@PublicEvolving
public interface IterationBody {
/**
* This method creates the graph for the iteration body.
*
* See Utils::iterate, Utils::iterateBoundedStreams and Utils::iterateAndReplayBoundedStreams for how the iteration
* body can be executed and when execution of the corresponding graph should terminate.
*
* Required: the number of feedback variable streams returned by this method must equal the number of variable
* streams given to this method.
*
* @param variableStreams the variable streams.
* @param dataStreams the data streams.
* @return a IterationBodyResult.
*/
IterationBodyResult process(DataStreamList variableStreams, DataStreamList dataStreams);
} |
2) Add IterationBodyResult class.
This is a helper class which contains the objects returned by the IterationBody::process(...).
Code Block | ||||
---|---|---|---|---|
| ||||
package org.apache.flink.iteration;
/**
* A helper class that contains the streams returned by the iteration body.
*/
class IterationBodyResult {
/**
* A list of feedback variable streams. These streams will only be used during the iteration execution and will
* not be returned to the caller of the iteration body. It is assumed that the method which executes the
* iteration body will feed the records of the feedback variable streams back to the corresponding input variable
* streams.
*/
DataStreamList feedbackVariableStreams;
/**
* A list of output streams. These streams will be returned to the caller of the methods that execute the
* iteration body.
*/
DataStreamList outputStreams;
/**
* An optional termination criteria stream. If this stream is not null, it will be used together with the
* feedback variable streams to determine when the iteration should terminate.
*/
@Nullable DataStream<?> terminationCriteria;
}
|
3) Add the IterationListener interface.
If an UDF (a.k.a user-defined function) used inside the IterationBody implements this interface, the callbacks on this interface will be invoked when corresponding events happen.
This interface allows users to achieve the following goals:
- Run an algorithm in sync mode, i.e. each subtask will wait for model parameters updates from all other subtasks before reading the aggregated model parameters and starting the next epoch of execution. As long as we could find a cut in the subgraph of the iteration body that all the operators in the cut only emit records in onEpochWatermarkIncremented, the algorithm would be synchronous. The detailed proof could be found in the appendix.
- Emit final output after the iteration terminates.
Code Block | ||||
---|---|---|---|---|
| ||||
package org.apache.flink.iteration
/**
* The callbacks defined below will be invoked only if the operator instance which implements this interface is used
* within an iteration body.
*/
@PublicEvolving
public interface IterationListener<T> {
/**
* This callback is invoked every time the epoch watermark of this operator increments. The initial epoch watermark
* of an operator is 0.
*
* The epochWatermark is the maximum integer that meets this requirement: every record that arrives at the operator
* going forward should have an epoch larger than the epochWatermark. See Java docs in Iterations for how epoch
* is determined for records ingested into the iteration body and for records emitted by operators within the
* iteration body.
*
* If all inputs are bounded, the maximum epoch of all records ingested into this operator is used as the
* epochWatermark parameter for the last invocation of this callback.
*
* @param epochWatermark The incremented epoch watermark.
* @param context A context that allows emitting side output. The context is only valid during the invocation of
* this method.
* @param collector The collector for returning result values.
*/
void onEpochWatermarkIncremented(int epochWatermark, Context context, Collector<T> collector);
/**
* This callback is invoked after the execution of the iteration body has terminated.
*
* See Java doc of methods in Iterations for the termination conditions.
*
* @param context A context that allows emitting side output. The context is only valid during the invocation of
* this method.
* @param collector The collector for returning result values.
*/
void onIterationTermination(Context context, Collector<T> collector);
/**
* Information available in an invocation of the callbacks defined in the IterationProgressListener.
*/
interface Context {
/**
* Emits a record to the side output identified by the {@link OutputTag}.
*
* @param outputTag the {@code OutputTag} that identifies the side output to emit to.
* @param value The record to emit.
*/
<X> void output(OutputTag<X> outputTag, X value);
}
} |
4) Add the Iterations class.
This class provides APIs to execute an iteration body with the user-provided inputs. This class provides three APIs to run an iteration body, each with different input types (e.g. bounded data streams vs. unbounded data streams) and data replay semantics (i.e. whether to replay the user-provided data streams).
Each of these three APIs provide the functionality as described in the iteration paradigm: Union the feedback variables streams (returned by the iteration body) with the initial variable streams (provided by the user) and use the merged streams as inputs to invoke IterationBody::process(...).
Code Block | ||||
---|---|---|---|---|
| ||||
package org.apache.flink.iteration;
/**
* A helper class to create iterations. To construct an iteration, Users are required to provide
*
* <ul>
* <li>initVariableStreams: the initial values of the variable data streams which would be updated
* in each round.
* <li>dataStreams: the other data streams used inside the iteration, but would not be updated.
* <li>iterationBody: specifies the subgraph to update the variable streams and the outputs.
* </ul>
*
* <p>The iteration body will be invoked with two parameters: The first 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); The second parameter is
* the data streams given to this method.
*
* <p>During the execution of iteration body, each of the records involved in the iteration has an
* epoch attached, which is mark the progress of the iteration. The epoch is computed as:
*
* <ul>
* <li>All records in the initial variable streams and initial data streams has epoch = 0.
* <li>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 onEpochWatermarkIncremented(), then the epoch of this record =
* epochWatermark.
* <li>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.
* </ul>
*
* <p>The framework would given the notification at the end of each epoch for operators and UDFs
* that implements {@link IterationListener}.
*
* <p>The limitation of constructing the subgraph inside the iteration body could be refer in {@link
* IterationBody}.
*
* <p>An example of the iteration is like:
*
* <pre>{@code
* DataStreamList result = Iterations.iterateUnboundedStreams(
* DataStreamList.of(first, second),
* DataStreamList.of(third),
* (variableStreams, dataStreams) -> {
* ...
* return new IterationBodyResult(
* DataStreamList.of(firstFeedback, secondFeedback),
* DataStreamList.of(output));
* }
* result.<Integer>get(0).addSink(...);
* }</pre>
*/
public class Iterations {
/**
* This method uses an iteration body to process records in possibly unbounded data streams. The
* iteration would not terminate if at least one of its inputs is unbounded. Otherwise it will
* terminated after all the inputs are terminated and no more records are iterating.
*
* @param initVariableStreams The initial variable streams, which is merged with the feedback
* variable streams before being used as the 1st parameter to invoke the iteration body.
* @param dataStreams The non-variable streams also refered in the {@code body}.
* @param body The computation logic which takes variable/data streams and returns
* feedback/output streams.
* @return The list of output streams returned by the iteration boy.
*/
public static DataStreamList iterateUnboundedStreams(
DataStreamList initVariableStreams, DataStreamList dataStreams, IterationBody body) {
return null;
}
/**
* This method uses 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.
*
* @param initVariableStreams The initial variable streams, which is merged with the feedback
* variable streams before being used as the 1st parameter to invoke the iteration body.
* @param dataStreams The non-variable streams also refered in the {@code body}.
* @param config The config for the iteration, like whether to re-create the operator on each
* round.
* @param body The computation logic which takes variable/data streams and returns
* feedback/output streams.
* @return The list of output streams returned by the iteration boy.
*/
public static DataStreamList iterateBoundedStreamsUntilTermination(
DataStreamList initVariableStreams, |
*Although SGD-based algorithms are also batch-based, they could be implemented with an Epoch-based method if intermediate state is allowed: the subtasks could sample a batch from all the records from the position of the last batch.
Based on the above classification and the replacement implementation for SGD-based algorithms with bounded dataset, we mainly need to support
- The synchronous / asynchronous epoch-based algorithms on the bounded dataset.
- The synchronous / asynchronous batch-based algorithms on the unbounded dataset.
The Goals of the Iteration Library
If we directly copy the current implementation of iteration on the DataSet And DataStream API, we would still meet with some problem, thus we would like to have some optimization to the existing iteration functionality.
The Iteration Body and Round Semantics
At the iteration API level, we would need the corresponding concept corresponding to Epoch and Batch. We would call processing one epoch as a round: users would specify a subgraph as the body of the iteration to specify how to calculate the update, after the iteration body process the whole dataset for one time (namely one Epoch). Apparently the round is meaningful only for the bounded cases.
...
Figure 1. The structure of an iteration body.
To process the inputs for multiple rounds, we would need feedback edges to emit the outputs of the last round to the inputs of the iteration body, and we union the initial inputs and the feedbacks. For bounded dataset when processing the first epoch, the data would be from the initial input edges and when processing the remaining epochs, the data would be from the feedback edges.
There are also inputs to the iteration body that do not have feedbacks. For example, an ML algorithm might have two inputs, one is the initialized model and the other is the training data. The input corresponding to the initial model will have a feedback after each epoch about the update to the model, but the training data would not need to be updated.
The iteration body also have one or multiple output streams. The iteration body might output records at each round, and the records emitted in all the rounds composed the final outputs.
Therefore, an iteration body composed of
- Variable inputs with feedbacks.
- Constant inputs without feedbacks.
- Outputs.
For the unbounded case, logically we should have
Per-Round v.s. All-Rounds Semantics
How users could specify the iteration body ? If we first consider the bounded cases, there are two options
- Per-round: Users specify a subgraph, and for each round, the framework would recreate the operators and do the same computation.
- All-rounds: Users specify a subgraph, and the operators inside the subgraph would process the epochs of all the rounds.
The DataSet iteration choose the per-round semantics. to support this semantics, in addition to re-create operators for each round, the framework also needs:
- For the inputs outside the iteration,
The benefits of this method is that writing an iteration body is no difference from constructing a DAG outside of the iteration.
Synchronization
Since for the bounded dataset, all the algorithms, to the best of out extend, are all able to be converted into epoch-based algorithms, thus we could only support the synchronization between epoch, namely between rounds.
How to
Besides, the previous DataStream and DataSet iteration APIs also have some caveats to support algorithm implementation:
- Lack of the support for multiple inputs, arbitrary outputs and nested iteration for both iteration APIs, which is required by scenarios like Metapath (multiple-inputs), boost algorithms (nested iteration) or when we want to output both loss and model (multiple-outputs). In the new iteration we would support these functionalities.
- Lack of asynchronous iteration support for the DataSet iteration, which is required by algorithms like asynchronous linear regression, in the new iterations we would support both synchronous and asynchronous modes for the bounded iteration.
- The current DataSet iteration by default provides a "for each round" semantics, namely users only need to specify the computation logic in each round, and the framework would executes the subgraph multiple times until convergence. To cooperate with the semantics, the DataSet iteration framework would merge the initial input and the feedback (bulk style and delta style), and replay the datasets comes from outside of the iteration. This method is easier to use, but it also limit some possible optimizations.
We also would like to improve these caveats in the new iteration library.
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:
Code Block | ||||||
---|---|---|---|---|---|---|
| ||||||
/** The iteration body specify the sub-graph inside the iteration. */
public interface IterationBody {
/** This annotation marks the function as it would builds the subgraph. */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface IterationFunction {}
/**
* This annotation marks a parameter of the iteration function as an input to the subgraph.
* The input is the union of the initial inputs bound to the iteration and the corresponding
* feedback.
*/
@Target(ElementType.PARAMETER)
@Retention(RetentionPolicy.RUNTIME)
public @interface IterationInput {
String value();
}
}
/** An example usage for the iteration body with two inputs. */
new IterationBody() {
@IterationFunction
UnboundedIterationDeclarative iterate(
@IterationInput("first") DataStream<Integer> first
@IterationInput("second") DataStream<> second,
) {
DataStream<Integer> feedBack1 = ...;
DataStream<String> output1 = ...;
return new UnboundedIterationDeclarationBuilder()
.withFeedback("first", feedBack1)
.withOutput("output1", output1)
.build();
}
} |
The interface for the unbounded iteration is straightforward:
Code Block | ||||||
---|---|---|---|---|---|---|
| ||||||
/** Builder for the unbounded iteration. */
public class UnboundedIteration {
/** Set the body of the iteration. */
UnboundedIteration withBody(IterationBody body) {...}
/** Bind the initial input with the specific name. */
UnboundedIteration bindInput(String name, DataStream<?> input) {...}
/** Generates and adds the subgraph corresponding to the iteration. */
ResultStreams build() {...}
}
/** The expected return type of the iteration function, which specifies the feedbacks and outputs. */
public class UnboundedIterationDeclaration {
public static class Builder {
/**
* Specify the feedback corresponding to the specific name. The feedback would
* be union with the initial input with the same name to provide to the iteration
* body.
*/
public Builder withFeedback(String name, DataStream<?> feedback) {...}
/** Specify one output with the specific name. */
public Builder withOutput(String name, DataStream<?> output) {...}
/** Generate the Declaration. */
UnboundedIterationDeclaration build() {...}
}
}
/** The map of the output streams of an iteration. */
public class ResultStreams {
/** Gets the DataStream with the specific name. */
public <T> DataStream<T> getStream(String name) {...}
} |
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:
Code Block | ||||||
---|---|---|---|---|---|---|
| ||||||
/** Builder for the bounded iteration. */
public class BoundedIteration {
/** Set the body of the iteration. */
BoundedIteration withBody(IterationBody body) {...}
/** Bind the initial input with the specific name. */
BoundedIteration bindInput(String name, DataStream<?> input) {...}
/** Generates and adds the subgraph corresponding to the iteration. */
ResultStreams build() {...}
}
/** The expected return type of the iteration function, which specifies the feedbacks, outputs and termination conditions. */
public class BoundedIterationDeclaration {
public static class Builder {
/**
* Specify the feedback corresponding to the specific name. The feedback would
* be union with the initial input with the same name to provide to the iteration
* body.
*/
public Builder withFeedback(String name, DataStream<?> feedback) {...}
/** Specify one output with the specific name. */
public Builder withOutput(String name, DataStream<?> output) {...}
/** Specify the termination condition of the iteration. */
<U> Builder until(TerminationCondition terminationCondition) { ... }
/** Generate the Declaration. */
BoundedIterationDeclaration build() {...}
}
}
/**
* The termination condition judges if iteration should stop based on the round number
* or the records of a given data stream in one round.
*
* Since there might be asynchronous iteration that multiple rounds are executed in parallel,
* the condition is evaluated at the end of each round. If it is evaluated to true for one
* round, then the iteration would desert the feedback records for the following rounds, but
* the already emitted records would not be withdrawal.
*/
public class TerminationCondition {
/**
* The records of the given DataStream will be collected in each
* round to be used in judging whether the loop should terminate.
*/
@Nullable DataStream<?> refStream;
/**
* A user-defined function that is evaluated at the end of each round.
*/
Function<Context, Boolean> isConverged;
interface Context {
/** The round number. */
int[] getRound();
/** The list of records of the referred stream. */
<T> List<T> getStreamRecords();
}
}
/** The progress tracking interface for UDF / Operator */
public interface BoundedIterationProgressListener<T> {
/** Sets a tool to be used to query the round number of the current record. */
default void setCurrentRecordRoundsQuerier(Supplier<int[]> querier) {}
/** Notified at the end of each round. */
void onRoundEnd(int[] round, Context context, Collector<T> collector);
/** Notified at the end of the whole iteration. */
default void onIterationEnd(int[] rounds, Context context) {}
public interface Context {
<X> void output(OutputTag<X> outputTag, X value);
Long timestamp();
TimerService timerService();
}
}
/** The utility methods to support per-round semantics. */
public class BoundedIterationPerRoundUtils {
/** The builder that creates the subgraph that executed with the per-round semantics **/
public interface EachRound {
Map<String, DataStream<?>> executeInEachRound();
}
/** Create a subgraph inside the iteration that execute with per-round semantics */
static ResultStreams forEachRound(EachRound eachRoundBuilder);
/** A cache to a Datastream. For each round it replaces the Datastream with the records received in this round and outputs all these records. */
static <T> ProcessFunction bulkCache(DataStream<T> inputStream);
/** A cache to a KeyedStream. For each round it updates the KeyedStream with the records received in this round and outputs the updated records. */
static <K, T> ProcessFunction deltaCache(KeyedStream<K, T> inputStream);
} |
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:
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()); ReplayableDataStreamList dataStreams, DataStream<double[]> modelUpdate = parameters.setParallelism(1) IterationConfig config, IterationBody .broadcast(body) { return null; } } |
5) Introduce the forEachRound utility method.
forEachRound allows the users to specify a sub-graph that executes as per-round mode, namely all the operators would be re-created for each round.
Code Block | ||
---|---|---|
| ||
public interface IterationBody { .connect(dataset) .... /** * @param inputs The inputs of the .coProcess(new TrainFunction())subgraph. * @param perRoundSubBody The computational logic that want to be executed as per-round. .setParallelism(10) * @return The output of the subgraph. */ returnstatic newDataStreamList BoundedIterationDeclarationBuilderforEachRound() DataStreamList inputs, PerRoundSubBody perRoundSubBody) { .withFeedback("model", modelUpdate)return null; } /** The sub-graph inside the iteration body that should be executed as .withOutput("final_model", parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG))per-round. */ interface PerRoundSubBody { DataStreamList process(DataStreamList input); } } |
6) Add the DataStreamList and ReplayableDataStreamList class.
DataStreamList is a helper class that contains a list of data streams with possibly different elements types and ReplayableDataStreamList is a helper class that contains a list of data streams and whether they need replay for each round.
Code Block | ||||
---|---|---|---|---|
| ||||
package org.apache.flink.iteration; public class DataStreamList { public static DataStreamList of(DataStream<?>... streams); // Returns the number of data streams in this list..until(new TerminationCondition(null, context -> context.getRound() >= batch * epochEachBatch)) .build(); public int size() {...} // }) Returns the data stream at the given index in this .build();list. public <T> DataStream<T> get(int index) {...} } public class ReplayableDataStreamList { DataStream<double[]> finalModel = resultStreams.get("final_model"); private final List<DataStream<?>> replayedDataStreams; private final List<DataStream<?>> nonReplayedStreams; finalModel.print(); private }ReplayableDataStreamList( public static class ParametersCacheFunction extends ProcessFunction<double[], double[]> List<DataStream<?>> replayedDataStreams, List<DataStream<?>> implements BoundedIterationProgressListener<double[]> {nonReplayedStreams) { this.replayedDataStreams = replayedDataStreams; private final double[] parametersthis.nonReplayedStreams = new double[N_DIM]; nonReplayedStreams; } public static voidReplayedDataStreamList processElement(double[] update, Context ctx, Collector<O> outputreplay(DataStream<?>... dataStreams) { return // Suppose we have a util to add the second array to the first.new ReplayedDataStreamList(Arrays.asList(dataStreams)); } public static NonReplayedDataStreamList notReplay(DataStream<?>... dataStreams) { return ArrayUtils.addWith(parameters, updatenew NonReplayedDataStreamList(Arrays.asList(dataStreams)); } List<DataStream<?>> getReplayedDataStreams() }{ public void onRoundEnd(int[] round, Context context, Collector<T> collectorreturn Collections.unmodifiableList(replayedDataStreams); } List<DataStream<?>> getNonReplayedStreams() { return collector.collect(parametersCollections.unmodifiableList(nonReplayedStreams); } private static class ReplayedDataStreamList extends public void onIterationEnd(int[] round, Context contextReplayableDataStreamList { public ReplayedDataStreamList(List<DataStream<?>> replayedDataStreams) { context.output(FINAL_MODEL_OUTPUT_TAG, parameterssuper(replayedDataStreams, Collections.emptyList()); } } public static class TrainFunction extends CoProcessFunction<double[], Tuple2<double[], Double>, double[]> implements BoundedIterationProgressListener<double[]> { ReplayableDataStreamList andNotReplay(DataStream<?>... nonReplayedStreams) { return new ReplayableDataStreamList( private final List<Tuple2<double[], Double>> dataset = new ArrayList<>(); private double[] firstRoundCachedParameter; getReplayedDataStreams(), Arrays.asList(nonReplayedStreams)); } } private Supplier<int[]> recordRoundQuerier; static class NonReplayedDataStreamList extends ReplayableDataStreamList { public void setCurrentRecordRoundsQuerier(Supplier<int[]> querierNonReplayedDataStreamList(List<DataStream<?>> nonReplayedDataStreams) { this.recordRoundQuerier = queriersuper(Collections.emptyList(), nonReplayedDataStreams); } public voidReplayableDataStreamList processElement1(double[] parameter, Context context, Collector<O> outputandReplay(DataStream<?>... replayedStreams) { int[]return round = recordRoundQuerier.get();new ReplayableDataStreamList( if (round[0] == 0) { Arrays.asList(replayedStreams), getNonReplayedStreams()); } firstRoundCachedParameter = parameter; } } |
7) Introduce the IterationConfig
IterationConfig allows users to specify the config for each iteration. For now users could only specify the default operator lifecycle inside the iteration, but it ensures the forward compatibility if we have more options in the future.
Code Block | ||
---|---|---|
| ||
package org.apache.flink.iteration; /** The config for an iteration. */ public class IterationConfig { private final OperatorLifeCycle returnoperatorRoundMode; private IterationConfig(OperatorLifeCycle operatorRoundMode) { } this.operatorRoundMode = operatorRoundMode; } calculateModelUpdate(parameter, output); public static IterationConfigBuilder newBuilder() { } return new IterationConfigBuilder(); public void processElement2(Tuple2<double[], Double> trainSample, Context context, Collector<O> output)} public static class IterationConfigBuilder { private OperatorLifeCycle operatorRoundMode = dataset.add(trainSample) OperatorLifeCycle.ALL_ROUND; private IterationConfigBuilder() {} public voidIterationConfigBuilder onRoundEnd(int[] round, Context context, Collector<T> output) { setOperatorRoundMode(OperatorLifeCycle operatorRoundMode) { if (round[0] == 0) {this.operatorRoundMode = operatorRoundMode; calculateModelUpdate(firstRoundCachedParameter, output)return this; } firstRoundCachedParameterpublic = null; IterationConfig build() { return new IterationConfig(operatorRoundMode); } } public private void calculateModelUpdate(double[] parameters, Collector<O> output)enum OperatorLifeCycle { ALL_ROUND, List<Tuple2<double[], Double>> samples = sample(dataset); PER_ROUND 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); } } } |
If instead we want to do asynchronous training, we would need to do the following change:
- The Parameters vertex would not wait till round end to ensure received all the updates from the iteration. Instead, it would immediately output the current parameters values once it received the model update from one train subtask.
- To label the source of the update, we would like to change the input type to be Tuple2<Integer, double[]>. The Parameters would only output the new parameters values to the Train task that send the update.
We omit the change to the graph building code since the change is trivial (change the output type and the partitioner to be customized one). The change to the Parameters vertex is the follows:
}
} |
8) 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.
Example Usage
This sections shows how general used ML algorithms could be implemented with the iteration API.
Offline Algorithms
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 Reducer would merge ΔA from all the train tasks and emit the reduced value to the Parameters node to update the parameters. The Reducer is required since we required the feedback streams have the same parallelism with the corresponding initialized streams.
draw.io Diagram | ||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Figure 2. 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:
Code Block | ||||
---|---|---|---|---|
| ||||
public class SynchronousBoundedLinearRegression {
private static final int N_DIM = 50;
private static final int N_EPOCH = 5;
private static final int N_BATCH_PER_EPOCH = 10;
private static final OutputTag<double[]> FINAL_MODEL_OUTPUT_TAG = new OutputTag<double[]>{};
public static void main(String[] args) {
DataStream<double[]> initParameters = loadParameters().setParallelism(1); | ||||
Code Block | ||||
| ||||
public static class ParametersCacheFunction extends ProcessFunction<Tuple2<Integer, double[]>, Tuple2<Integer, double[]>> implements BoundedIterationProgressListener<double[]> { private final double[] parameters = new double[N_DIM]; public void processElement(Tuple2<Integer, double[]> update, Context ctx, Collector<Tuple2<Integer, double[]>> output) { // Suppose we have a util to add the second array to the first. ArrayUtils.addWith(parametersDataStream<Tuple2<double[], update); Double>> dataset output.collect(new Tuple2<>(update.f0, parameters))= loadDataSet().setParallelism(1); } DataStreamList publicresultStreams void onIterationEnd(int[] round, Context context) { = context.output(FINAL_MODEL_OUTPUT_TAG, parameters); } } |
Online Training with Unbounded Iteration
Suppose now we would change the algorithm to unbounded iteration, compared to the offline, the differences is that
- The dataset is unbounded. The Train operator could not cache all the data in the first round.
- The training algorithm might be changed to others like FTRL. But we keep using SGD in this example since it does not affect showing the usage of the iteration.
We also start with the synchronous case. for online training, the Train vertex usually do one update after accumulating one mini-batch. This is to ensure the distribution of the samples is similar to the global statistics. In this example we omit the complex data re-sample process and just fetch the next several records as one mini-batch.
The JobGraph for online training is still shown in Figure 1, with the training dataset become unbounded. Similar to the bounded cases, for the synchronous training, the process would be expected like
- The Parameters broadcast the initialized values on received the input values.
- All the Train task read the next mini-batch of records, Calculating an update and emit to the Parameters vertex. Then it would wait till received update parameters from the Parameters Vertex before it head to process the next mini-batch.
- The Parameter vertex would wait received the updates from all the Train tasks before it broadcast the updated parameters.
Since in the unbounded case there is not the concept of round, and we do update per-mini-batch, thus we could instead use the InputSelectable functionality to implement the algorithm:
Code Block | ||
---|---|---|
| ||
public class SynchronousUnboundedLinearRegression { private static final N_DIM = 50; private static final OutputTag<double[]> MODEL_UPDATE_OUTPUT_TAG = new OutputTag<double[]>{}; public static void main(String[] args) { DataStream<double[]> initParameters = loadParameters().setParallelism(1Iterations.iterateBoundedStreamsUntilTermination( DataStreamList.of(initParameters), ReplayableDataStreamList.notReplay(dataset), IterationConfig.newBuilder().setOperatorRoundMode(ALL_ROUND).build(); (variableStreams, dataStreams) -> { DataStream<double[]> parameterUpdates = variableStreams.get(0); DataStream<Tuple2<double[], Double>> dataset = dataStreams.get(0); SingleOutputStreamOperator<double[]> parameters = parameterUpdates.process(new ParametersCacheFunction()); DataStream<double[]> modelUpdate = parameters.setParallelism(1) .broadcast() .connect(dataset) .coProcess(new TrainFunction()) .setParallelism(10) .process(new ReduceFunction()) .setParallelism(1) return new IterationBodyResult(DataStreamList.of(modelUpdate), DataStreamList.of(parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG))); }); DataStream<double[]> finalModel = resultStreams.get("final_model"); DataStream<Tuple2<double[], Double>> dataset = loadDataSet().setParallelism(1finalModel.print(); } public static ResultStreamsclass resultStreamsParametersCacheFunction = new UnboundedIteration()extends ProcessFunction<double[], double[]> implements IterationListener<double[]> { .withBody(new IterationBody( @IterationInput("model") DataStream<doubleprivate final double[]> model, parameters = new double[N_DIM]; public void @IterationInput("dataset") DataStream<Tuple2<doubleprocessElement(double[] update, Double>>Context dataset ctx, Collector<O> output) { // Suppose we have SingleOutputStreamOperator<double[]> parameters = model.process(new ParametersCacheFunction(10)); a util to add the second array to the first. DataStream<double[]> modelUpdate = parameters.setParallelism(1) ArrayUtils.addWith(parameters, update); } void onEpochWatermarkIncremented(int epochWatermark, .broadcast() Context context, Collector<T> collector) { if (epochWatermark < N_EPOCH .connect(dataset)* N_BATCH_PER_EPOCH) { .transform(collector.collect(parameters); } } public void "operator", onIterationEnd(int[] round, Context context) { context.output(FINAL_MODEL_OUTPUT_TAG, parameters); } BasicTypeInfo.INT_TYPE_INFO,} public static class TrainFunction extends CoProcessFunction<double[], Tuple2<double[], Double>, double[]> implements IterationListener<double[]> { private final List<Tuple2<double[], Double>> dataset = new TrainOperatorsArrayList<>(50)); private double[] firstRoundCachedParameter; private .setParallelism(10)Supplier<int[]> recordRoundQuerier; public void setCurrentRecordRoundsQuerier(Supplier<int[]> querier) { return new UnBoundedIterationDeclarationBuilder() this.recordRoundQuerier = querier; .withFeedback("model", modelUpdate)} public void processElement1(double[] parameter, Context context, Collector<O> output) { .withOutput("model_update", parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG)) int[] round = recordRoundQuerier.get(); .build(); if (round[0] == 0) { }) firstRoundCachedParameter = .build()parameter; DataStream<double[]> finalModel = resultStreams.get("model_update")return; finalModel.addSink(...) } public static class ParametersCacheFunction extends ProcessFunction<double[], double[]> { calculateModelUpdate(parameter, output); } privatepublic final int numOfTrainTasks; private final int numOfUpdatesReceived = 0; void processElement2(Tuple2<double[], Double> trainSample, Context context, Collector<O> output) { dataset.add(trainSample) private final double[] parameters = new double[N_DIM]; } publicvoid ParametersCacheFunctiononEpochWatermarkIncremented(int numOfTrainTasks epochWatermark, Context context, Collector<T> collector) { if this.numOfTrainTasks(epochWatermark == numOfTrainTasks;0) { } public void processElement(double[] update, Context ctx, Collector<O> output) { calculateModelUpdate(firstRoundCachedParameter, output); firstRoundCachedParameter = null; // Suppose we have a util to add the second array to} the first. } ArrayUtils.addWith( private void calculateModelUpdate(double[] parameters, update); Collector<O> output) { List<Tuple2<double[], Double>> samples numOfUpdatesReceived++= sample(dataset); if (numOfUpdatesReceived =double[] modelUpdate = numOfTrainTasks) {new double[N_DIM]; output.collect(parameters); for (Tuple2<double[], Double> record : samples) { numOfUpdatesReceived = 0; double diff = (ArrayUtils.muladd(record.f0, parameters) }- record.f1); } } public static class TrainOperators extends AbstractStreamOperator<double[]> implements TwoInputStreamOperator<double[], Tuple2<double[], Double>, double[]>, InputSelectable { ArrayUtils.addWith(modelUpdate, ArrayUtils.multiply(record.f0, diff)); private final int miniBatchSize; } private final List<Tuple2<double[], Double>> miniBatch = new ArrayList<>( output.collect(modelUpdate); private double[] firstRoundCachedParameter; } } public static class public TrainOperators(int miniBatchSize) ReduceFunction { private double[] mergedValue this.miniBatchSize = miniBatchSize; } = ArrayUtils.newArray(N_DIM); public void processElement1processElement(double[] parameter, Context context, Collector<O> output) { mergedValue calculateModelUpdate(parameter, output); miniBatchSize.clear(= ArrayUtils.add(mergedValue, parameter); } public void processElement2(Tuple2<double[], Double> trainSampleonEpochWatermarkIncremented(int epochWatermark, Context context, Collector<O>Collector<T> outputcollector) { dataset.add(trainSamplecollector.collect(mergedValue); mergedValue = ArrayUtils.newArray(N_DIM); } public InputSelection nextSelection() { } } |
Here we use a specialized reduce function that handles the boundary of the round. However, outside of the iteration we might already be able to do the reduce with windowAllI().reduce(), to reuse the operators outside of the iteration, we could use the forEachRound utility method.
Code Block | ||
---|---|---|
| ||
DataStreamList resultStreams = if (miniBatch.size() < miniBatchSize) { Iterations.iterateBoundedStreamsUntilTermination( DataStreamList.of(initParameters), ReplayableDataStreamList.notReplay(dataset), IterationConfig.newBuilder().setOperatorRoundMode(ALL_ROUND).build(); (variableStreams, dataStreams) -> { return InputSelection.SECOND; DataStream<double[]> parameterUpdates = variableStreams.get(0); } else { DataStream<Tuple2<double[], Double>> dataset return InputSelection.FIRST; = dataStreams.get(0); } SingleOutputStreamOperator<double[]> parameters = } parameterUpdates.process(new ParametersCacheFunction()); private void calculateModelUpdate(double[] parameters, Collector<O> output) { DataStream<double[]> modelUpdate = parameters.setParallelism(1) double[] modelUpdate = new double[N_DIM]; .broadcast() for (Tuple2<double[], Double> record : miniBatchSize) { .connect(dataset) .coProcess(new TrainFunction()) double diff = (ArrayUtils.muladd(record.f0, parameters) - record.f1); ArrayUtils.addWith(modelUpdate, ArrayUtils.multiply(record.f0, diff)); } output.collect(modelUpdate); } } } |
Also similar to the bounded case, for the asynchronous training the Parameters vertex would not wait for received updates from all the Train tasks. Instead, it would directly response to the task sending update:
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.
Currently we do not see requirements on Bounded iteration + STREAM mode, if there are additional requirement in the future we would implement this mode, and it could also be supported with the current framework.
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.
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.
...
setParallelism(10)
DataStream<double[]> reduced = forEachRound(DataStreamList.of(modelUpdate), streams -> {
return streams.<double[]>get(0).windowAll().reduce((x, y) -> ArrayUtils.add(x, y));
}).<double[]>get(0);
return new IterationBodyResult(DataStreamList.of(modelUpdate), DataStreamList.of(parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG)));
}); |
This would be very helpful for complex scenarios that we could reuse the ability of the the current datastream and table operators.
If instead we want to do asynchronous training, we would need to do the following change:
- The Parameters vertex would not wait till round end to ensure received all the updates from the iteration. Instead, it would immediately output the current parameters values once it received the model update from one train subtask.
- To label the source of the update, we would like to change the input type to be Tuple2<Integer, double[]>. The Parameters would only output the new parameters values to the Train task that send the update.
- The Reducer would not merge the received updates any more, instead, it would directly redirect the received updates.
We omit the change to the graph building code since the change is trivial (change the output type and the partitioner to be customized one). The change to the Parameters vertex is the follows:
Code Block | ||
---|---|---|
| ||
public static class ParametersCacheFunction extends ProcessFunction<Tuple2<Integer, double[]>, Tuple2<Integer, double[]>>
implements BoundedIterationProgressListener<double[]> {
private final double[] parameters = new double[N_DIM];
public void processElement(Tuple2<Integer, double[]> update, Context ctx, Collector<Tuple2<Integer, double[]>> output) {
// Suppose we have a util to add the second array to the first.
ArrayUtils.addWith(parameters, update);
output.collect(new Tuple2<>(update.f0, parameters))
}
public void onIterationEnd(int[] round, Context context) {
context.output(FINAL_MODEL_OUTPUT_TAG, parameters);
}
}
public class AsynchronousBoundedLinearRegression {
...
DataStreamList resultStreams =
Iterations.iterateBoundedStreamsUntilTermination(DataStreamList.of(initParameters), DataStreamList.of(dataset), (variableStreams, dataStreams) -> {
DataStream<double> parameterUpdates = variableStreams.get(0);
DataStream<Tuple2<double[], Double>> dataset = dataStreams.get(0);
SingleOutputStreamOperator<double[]> parameters = parameterUpdates.process(new ParametersCacheFunction());
DataStream<double[]> modelUpdate = parameters.setParallelism(1)
.partitionCustom((key, numPartitions) -> key % numPartitions, update -> update.f0)
.connect(dataset)
.coProcess(new TrainFunction())
.setParallelism(10)
return new IterationBodyResult(DataStreamList.of(modelUpdate), DataStreamList.of(parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG)));
});
...
} |
Online Algorithms
Suppose now we would change the algorithm to unbounded iteration, compared to the offline, the differences is that
- The dataset is unbounded. The Train operator could not cache all the data in the first round.
- The training algorithm might be changed to others like FTRL. But we keep using SGD in this example since it does not affect showing the usage of the iteration.
We also start with the synchronous case. for online training, the Train vertex usually do one update after accumulating one mini-batch. This is to ensure the distribution of the samples is similar to the global statistics. In this example we omit the complex data re-sample process and just fetch the next several records as one mini-batch.
The JobGraph for online training is still shown in Figure 1, with the training dataset become unbounded. Similar to the bounded cases, for the synchronous training, the process would be expected like
- The Parameters broadcast the initialized values on received the input values.
- All the Train task read the next mini-batch of records, Calculating an update and emit to the Parameters vertex. Then it would wait till received update parameters from the Parameters Vertex before it head to process the next mini-batch.
- The Parameter vertex would wait received the updates from all the Train tasks before it broadcast the updated parameters.
Since in the unbounded case there is not the concept of round, and we do update per-mini-batch, thus we could instead use the InputSelectable functionality to implement the algorithm:
Code Block | ||
---|---|---|
| ||
public class SynchronousUnboundedLinearRegression {
private static final N_DIM = 50;
private static final OutputTag<double[]> MODEL_UPDATE_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);
DataStreamList resultStreams = Iterations.iterateUnboundedStreams(DataStreamList.of(initParameters), DataStreamList.of(dataset), (variableStreams, dataStreams) -> {
DataStream<double> parameterUpdates = variableStreams.get(0);
DataStream<Tuple2<double[], Double>> dataset = dataStreams.get(0);
SingleOutputStreamOperator<double[]> parameters = model.process(new ParametersCacheFunction(10));
DataStream<double[]> modelUpdate = parameters.setParallelism(1)
.broadcast()
.connect(dataset)
.transform(
"operator",
BasicTypeInfo.INT_TYPE_INFO,
new TrainOperators(50));
.setParallelism(10);
return new IterationBodyResult(DataStreamList.of(modelUpdate), DataStreamList.of(parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG)));
})
DataStream<double[]> finalModel = resultStreams.get("model_update");
finalModel.addSink(...)
}
public static class ParametersCacheFunction extends ProcessFunction<double[], double[]> {
private final int numOfTrainTasks;
private final int numOfUpdatesReceived = 0;
private final double[] parameters = new double[N_DIM];
public ParametersCacheFunction(int numOfTrainTasks) {
this.numOfTrainTasks = numOfTrainTasks;
}
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);
numOfUpdatesReceived++;
if (numOfUpdatesReceived == numOfTrainTasks) {
output.collect(parameters);
numOfUpdatesReceived = 0;
}
}
}
public static class TrainOperators extends AbstractStreamOperator<double[]> implements TwoInputStreamOperator<double[], Tuple2<double[], Double>, double[]>, InputSelectable {
private final int miniBatchSize;
private final List<Tuple2<double[], Double>> miniBatch = new ArrayList<>();
private double[] firstRoundCachedParameter;
public TrainOperators(int miniBatchSize) {
this.miniBatchSize = miniBatchSize;
}
public void processElement1(double[] parameter, Context context, Collector<O> output) {
calculateModelUpdate(parameter, output);
miniBatchSize.clear();
}
public void processElement2(Tuple2<double[], Double> trainSample, Context context, Collector<O> output) {
dataset.add(trainSample);
}
public InputSelection nextSelection() {
if (miniBatch.size() < miniBatchSize) {
return InputSelection.SECOND;
} else {
return InputSelection.FIRST;
}
}
private void calculateModelUpdate(double[] parameters, Collector<O> output) {
double[] modelUpdate = new double[N_DIM];
for (Tuple2<double[], Double> record : miniBatchSize) {
double diff = (ArrayUtils.muladd(record.f0, parameters) - record.f1);
ArrayUtils.addWith(modelUpdate, ArrayUtils.multiply(record.f0, diff));
}
output.collect(modelUpdate);
}
}
} |
Also similar to the bounded case, for the asynchronous training the Parameters vertex would not wait for received updates from all the Train tasks. Instead, it would directly response to the task sending update:
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))
}
}
}
|
Proposed Changes
In this section, we discuss a few design choices related to the implementation and usage of the proposed APIs.
Support the Feedback Edges
1) How the feedback edge is supported.
The Flink core runtime can only execute a DAG of operators that does not involve cycles. Thus extra work needs to be done to support feedback edges (which effectively introduces cycles in the data flow).
Similar to the existing iterative API, this FLIP plans to implement the feedback edge using the following approach:
- Automatically insert the HEAD and the TAIL operators as the first and the last operators in the iteration body.
- Co-locate the HEAD and the TAIL operators on the same task manager.
- Have the HEAD and the TAIL operators transmit the records of the feedback edges using an in-memory queue data structure.
Since all the forward edges have limited buffers, to avoid the deadlocks the feedback queues must have unlimited size. To avoid unlimited memory footprint, when the queued size exceeds a threshold the records would be spilled to the disk. To avoid such spilling the HEAD operators would read the data in a "feedback-first" manner, namely it would always process the feedback records first if there are records from both initial input and feedback edges.
2) How the termination of the iteration execution is determined.
Having the feedback edges also complicate the termination detection of the job. Since the feedback edges is not visible in the JobGraph, the HEAD operators, as the first operators in the DAG of the iteration body, would decide when the whole iteration body could be terminated and initiate the termination process. The termination happens when
- All the inputs to the iteration body have been finished.
- AND
- If users have specify a reference stream, the number of records in each epoch would be counted after the epoch is done, and if the count is 0, HEAD would start to terminate.
- OR If users have not specify a reference stream, then when no records are still being processed inside the iteration body. This is detected by have a special event travel through the whole iteration body.
The Execution Mode
Currently Flink provides two execution mode for the bounded job, name blocking mode and stream mode. The two modes differ in the edge types, scheduler policy and failover policy. For the jobs with iteration, to reduce the implementation complexity of coordination and failover, we expect the whole iteration body would be always run together, thus at least initially we would not support blocking mode.
Users may have jobs that requires preprocessing and post-process. For such jobs, currently we expected users to split the whole pipeline into multiple jobs so that users could use blocking modes for the other jobs.
Dealing With Failure
Since currently we only support stream mode, we could rely on the checkpoint mechanism of Flink directly. The existing checkpoint mechanism does not support feedback edges, we would enhance them by also snapshotting the feedback records with the checkpoint.
Compatibility, Deprecation, and Migration Plan
Deprecation plan
The following APIs will be deprecated and removed in the future Flink release:
- The entire DataSet class. See FLIP-131 for its motivation and the migration plan. The deprecation of DataSet::iterate(...) proposed by this FLIP is covered by FLIP-131.
- The DataStream::iterate(...) and DataStream::iterate(long).
Compatibility
- This FLIP introduces backward in-compatible changes by proposing to remove DataStream::iterate(...) in the future.
- We expect the APIs proposed in this FLIP can address most of the use-cases supported by the DataStream::iterate(...) and DataSet::iterate(...). The only use-cases that we have dropped support are those that require nested iteration on bounded data streams in sync mode. We have made this choice because we are not aware of any reasonable use-cases that require nested iteration. This support can be added if any user provide a good use-case for nested iteration.
Migration plan
- Users will need to re-write their application code in order to migrate from the existing iterative APIs to the proposed APIs.
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.
For the iteration DAG build graph, it would be more simpler if we could directly refer to the data stream variables outside of the closure of iteration body. However, since we need to make the iteration DAG creation first happen in the mock execution environment, we could not use these variables directly, otherwise we would directly modify the real environment and won't have chance to add wrappers to the operators.
The existing FLIP does not explicitly support nested iterations. This is because we have not seen clear use-case that require nested iteration. We would prefer to only introduce additional complexity that are required by some reasonable use-cases, In the future if we decide to support nested iterations, we will need extra APIs to run the iteration body, with the semantics that operators inside the iteration body will be re-created for every round of iteration.
Appendix
1) In the following, we prove that the proposed solution can be used to implement an iterative algorithm in the sync mode.
Refer to the "Proposed Changes" section for the definition of sync mode and the description of the solution. In the following, we prove that the solution does work as expected.
Proof
In the following, we will prove that the solution described above could enforce the sync-mode execution. Note that the calculation of the record's epoch and the semantics of onEpochWatermarkIncremented(...) are described in the Java doc of the corresponding APIs.
Lemma-1: For any operator OpB defined in the IterationBody, at the time its Nth invocation of onEpochWatermarkIncremented(...) starts, it is guaranteed that:
- If an input edge is a non-feedback edge from OpA, then OpA's Nth invocation of onEpochWatermarkIncremented(...) has been completed.
- If an input edge is a feedback edge from OpA, then OpA's (N-1)th invocation of onEpochWatermarkIncremented(...) has been completed.
Let's prove the lemma-1 by contradiction:
- At the time the OpB's Nth invocation starts, its epoch watermark has incremented to N, which means OpB will no longer receive any record with epoch <= N.
- Suppose there is a non-feedback edge from OpA AND OpA's Nth invocation has not been completed. Then when OpA's Nth invocation completes, OpA can generate a record with epoch=N and send it to OpB via this non-feedback edge, which contradicts the guarantee described above.
- Suppose there is a feedback edge from OpA AND OpA's (N-1)th invocation has not been completed. Then when OpA's (N-1)th invocation completes, OpA can generate a record with epoch=N and send it to OpB via this feedback edge, which contradicts the guarantee described above.
Lemma-2: For any operator OpB defined in the IterationBody, at the time its Nth invocation of onEpochWatermarkIncremented(...) starts, it is guaranteed that:
- If an edge is a non-feedback input edge from OpA and this edge is part of a feedback loop, then OpA's (N+1)th invocation of onEpochWatermarkIncremented(...) has not started.
- If an edge is a feedback input edge from OpA and this edge is part of a feedback loop, then OpA's Nth invocation of onEpochWatermarkIncremented(...) has not started.
Let's prove this lemma by contradiction:
- Suppose there is a non-feedback edge from OpA, this edge is part of a feedback loop, and OpA's (N+1)th invocation has started. Since this non-feedback edge is part of a feedback loop, there is a backward path from OpA to OpB with exactly 1 feedback edge on this path. By applying the lemma-1 recursively for operators on this path, we can tell that OpB's Nth invocation has been completed. This contradicts the assumption that OpB's Nth invocation just started.
- Suppose there is a feedback edge from OpA, this edge is part of a feedback loop, and OpA's Nth invocation has started. Since this feedback edge is part of feedback loop, there is a backward path from OpA to OpB with no feedback edge on this path. By applying lemma-1 recursively for operators on this path, we can tell that OpB's Nth invocation has been completed. This contradicts the assumption that OpB's Nth invocation just started.
Let's now prove that the sync-mode is achieved:
- For any operator in the IterationBody, we define its output for the Nth epoch as the output emitted by the Nth invocation of onEpochWatermarkIncremented(). This definition is well-defined because operators only emit records in onEpochWatermarkIncremented().
- At the time an operator OpB computes its output for the Nth epoch, this operator must have received exactly the following records from its input edges:
- Suppose an edge is a non-feedback input edge from OpA and this edge is part of a feedback loop. It follows that OpA has emitted records for its Nth epoch (by lemma-1) and has not started to emit records for its (N+1)th epoch (by lemma-2).
- Suppose an edge is a feedback input edge from OpA and this edge is part of a feedback loop. It follows that OpA has emitted records for its (N-1)th epoch (by lemma-1) and has not started to emit records for its Nth epoch (by lemma-2).