yunfengzhou-hub commented on code in PR #97: URL: https://github.com/apache/flink-ml/pull/97#discussion_r891140907
########## flink-ml-iteration/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCache.java: ########## @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.iteration.datacache.nonkeyed; + +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.core.fs.Path; +import org.apache.flink.core.memory.ManagedMemoryUseCase; +import org.apache.flink.iteration.operator.OperatorUtils; +import org.apache.flink.runtime.jobgraph.OperatorID; +import org.apache.flink.runtime.memory.MemoryManager; +import org.apache.flink.runtime.state.StateInitializationContext; +import org.apache.flink.runtime.state.StatePartitionStreamProvider; +import org.apache.flink.runtime.state.StateSnapshotContext; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.apache.flink.streaming.runtime.tasks.StreamTask; +import org.apache.flink.table.runtime.util.LazyMemorySegmentPool; +import org.apache.flink.table.runtime.util.MemorySegmentPool; +import org.apache.flink.util.Preconditions; + +import org.apache.commons.collections.IteratorUtils; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** + * A {@link ListState} child class that records data and replays them on required. + * + * <p>This class basically stores data in file system, and provides the option to cache them in + * memory. In order to use the memory caching option, users need to allocate certain managed memory + * for the wrapper operator through {@link + * org.apache.flink.api.dag.Transformation#declareManagedMemoryUseCaseAtOperatorScope}. + * + * <p>NOTE: Users need to explicitly invoke this class's {@link + * #snapshotState(StateSnapshotContext)} method in order to store the recorded data in snapshot. + */ +public class ListStateWithCache<T> implements ListState<T> { + + /** The tool to serialize/deserialize records. */ + private final TypeSerializer<T> serializer; + + /** The path of the directory that holds the files containing recorded data. */ + private final Path basePath; + + /** The data cache writer for the received records. */ + private final DataCacheWriter<T> dataCacheWriter; + + @SuppressWarnings("unchecked") + public ListStateWithCache( + TypeSerializer<T> serializer, + StreamTask<?, ?> containingTask, + StreamingRuntimeContext runtimeContext, + StateInitializationContext stateInitializationContext, + OperatorID operatorID) + throws IOException { + this.serializer = serializer; + + MemorySegmentPool segmentPool = null; + double fraction = + containingTask + .getConfiguration() + .getManagedMemoryFractionOperatorUseCaseOfSlot( + ManagedMemoryUseCase.OPERATOR, + runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(), + runtimeContext.getUserCodeClassLoader()); + if (fraction > 0) { + MemoryManager memoryManager = containingTask.getEnvironment().getMemoryManager(); + segmentPool = + new LazyMemorySegmentPool( + containingTask, + memoryManager, + memoryManager.computeNumberOfPages(fraction)); + } + + basePath = + OperatorUtils.getDataCachePath( + containingTask.getEnvironment().getTaskManagerInfo().getConfiguration(), + containingTask + .getEnvironment() + .getIOManager() + .getSpillingDirectoriesPaths()); + + List<StatePartitionStreamProvider> inputs = + IteratorUtils.toList( + stateInitializationContext.getRawOperatorStateInputs().iterator()); + Preconditions.checkState( + inputs.size() < 2, "The input from raw operator state should be one or zero."); + + List<Segment> priorFinishedSegments = new ArrayList<>(); + if (inputs.size() > 0) { + DataCacheSnapshot dataCacheSnapshot = + DataCacheSnapshot.recover( + inputs.get(0).getStream(), + basePath.getFileSystem(), + OperatorUtils.createDataCacheFileGenerator( + basePath, "cache", operatorID)); + + if (segmentPool != null) { + dataCacheSnapshot.tryReadSegmentsToMemory(serializer, segmentPool); + } + + priorFinishedSegments = dataCacheSnapshot.getSegments(); + } + + this.dataCacheWriter = + new DataCacheWriter<>( + serializer, + basePath.getFileSystem(), + OperatorUtils.createDataCacheFileGenerator(basePath, "cache", operatorID), + segmentPool, + priorFinishedSegments); + } + + public void snapshotState(StateSnapshotContext context) throws Exception { Review Comment: I have also confirmed with @gaoyunhaii that the `processElement` and `snapshotState` methods of an operator would not be invoked at the same time, so there should not be a parallel computing issue. More descriptions of Flink's resource usage can be found at https://nightlies.apache.org/flink/flink-docs-master/docs/concepts/flink-architecture/#task-slots-and-resources. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@flink.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
