State handling on KeyedCoProcessFunction serving ML models - apache-flink

Question

I am working on a KeyedCoProcessFunction that looks like this:
class MyOperator extends KeyedCoProcessFunction[String, ModelDef, Data, Prediction]
with CheckpointedFunction {
// To hold loaded models
#transient private var models: HashMap[(String, String), Model] = _
// For serialization purposes
#transient private var modelsBytes: MapState[(String, String), Array[Bytes]] = _
...
override def snapshotState(context: FunctionSnapshotContext): Unit = {
modelsBytes.clear() // This raises an exception when there is no active key set
for ((k, model) <- models) {
modelsBytes.put(k, model.toBytes(v))
}
}
override def initializeState(context: FunctionInitializationContext): Unit = {
modelsBytes = context.getKeyedStateStore.getMapState[String, String](
new MapStateDescriptor("modelsBytes", classOf[String], classOf[String])
)
if (context.isRestored) {
// restore models from modelsBytes
}
}
}
The state consists of a collection of ML models built using a third party library. Before checkpoints, I need to dump the loaded models into byte arrays in snapshotState.
My question is, within snapshotState, modelsBytes.clear() raises an exception when there is no active key. This happens when I start the application from scratch without any data on the input streams. So, when the time for a checkpoint comes, I get this error:
java.lang.NullPointerException: No key set. This method should not be called outside of a keyed context.
However, when the input stream contains data, checkpoints work just fine. I am a bit confused about this because snapshotState does not provide a keyed context (contrary to processElement1 and processElement2, where the current key is accessible by doing ctx.getCurrentKey) so it seems to me that the calls to clear and put within snapshotState should fail always since they're supposed to work only within a keyed context. Can anyone clarify if this is the expected behaviour actually?

Answer 1

A keyed state can only be used on a keyed stream as written in the documentation.
* <p>The state is only accessible by functions applied on a {#code KeyedStream}. The key is
* automatically supplied by the system, so the function always sees the value mapped to the
* key of the current element. That way, the system can handle stream and state partitioning
* consistently together.
If you call clear(), you will not clear the whole map, but just reset the state of the current key. The key is always known in processElementX.
/**
* Removes the value mapped under the current key.
*/
void clear();
You should actually receive a better exception when you try to call clear in a function other than processElementX. In the end, you are using the keyed state incorrectly.
Now for your actual problem. I'm assuming you are using a KeyedCoProcessFunction because the models are updated in a separate input. If they are static, you could just load them open from a static source (for example, included in the jar). Furthermore, often there is only one model that is applied for all values with different keys, then you could use BroadCast state. So I'm assuming you have different models for different types of data separated by keys.
If they are coming in from input2, then you already serialize them upon invocation of processElement2.
override def processElement2(model: Model, ctx: Context, collector): Unit = {
models.put(ctx.getCurrentKey, model)
modelsBytes.put(ctx.getCurrentKey, model.toBytes(v))
}
Then you would not override snapshotState, as the state is already up-to-date. initializeState would deserialize models eagerly or you could also materialize them lazily in processElement1.