You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by zh...@apache.org on 2023/09/05 03:23:24 UTC
[flink-ml] branch master updated: [FLINK-32810] Improve memory allocation in ListStateWithCache
This is an automated email from the ASF dual-hosted git repository.
zhangzp pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink-ml.git
The following commit(s) were added to refs/heads/master by this push:
new ebdf3626 [FLINK-32810] Improve memory allocation in ListStateWithCache
ebdf3626 is described below
commit ebdf362664283bacad476126ee18c4059593950a
Author: Fan Hong <ho...@gmail.com>
AuthorDate: Tue Sep 5 11:23:20 2023 +0800
[FLINK-32810] Improve memory allocation in ListStateWithCache
This closes #250.
---
.../ml/common/datastream/DataStreamUtils.java | 14 +-
.../datacache/nonkeyed/ListStateWithCacheTest.java | 151 +++++++++++++++++++++
.../datacache/nonkeyed/ListStateWithCache.java | 61 ++++++---
.../OperatorScopeManagedMemoryManager.java | 105 ++++++++++++++
.../OperatorScopeManagedMemoryManagerTest.java | 86 ++++++++++++
.../apache/flink/ml/clustering/kmeans/KMeans.java | 15 +-
6 files changed, 402 insertions(+), 30 deletions(-)
diff --git a/flink-ml-core/src/main/java/org/apache/flink/ml/common/datastream/DataStreamUtils.java b/flink-ml-core/src/main/java/org/apache/flink/ml/common/datastream/DataStreamUtils.java
index e4cbcd52..88ff0d31 100644
--- a/flink-ml-core/src/main/java/org/apache/flink/ml/common/datastream/DataStreamUtils.java
+++ b/flink-ml-core/src/main/java/org/apache/flink/ml/common/datastream/DataStreamUtils.java
@@ -38,6 +38,7 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.core.memory.ManagedMemoryUseCase;
import org.apache.flink.iteration.datacache.nonkeyed.ListStateWithCache;
+import org.apache.flink.iteration.datacache.nonkeyed.OperatorScopeManagedMemoryManager;
import org.apache.flink.iteration.operator.OperatorStateUtils;
import org.apache.flink.ml.common.datastream.sort.CoGroupOperator;
import org.apache.flink.ml.common.window.CountTumblingWindows;
@@ -47,6 +48,7 @@ import org.apache.flink.ml.common.window.GlobalWindows;
import org.apache.flink.ml.common.window.ProcessingTimeSessionWindows;
import org.apache.flink.ml.common.window.ProcessingTimeTumblingWindows;
import org.apache.flink.ml.common.window.Windows;
+import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.runtime.state.VoidNamespace;
@@ -72,6 +74,7 @@ import org.apache.flink.streaming.api.windowing.windows.GlobalWindow;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.table.api.TableException;
import org.apache.flink.util.Collector;
@@ -483,13 +486,18 @@ public class DataStreamUtils {
public void initializeState(StateInitializationContext context) throws Exception {
super.initializeState(context);
+ final StreamTask<?, ?> containingTask = getContainingTask();
+ final OperatorID operatorID = config.getOperatorID();
+ final OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(containingTask, operatorID);
+ final String stateKey = "values-state";
+ manager.register(stateKey, 1.);
valuesState =
new ListStateWithCache<>(
getOperatorConfig().getTypeSerializerIn(0, getClass().getClassLoader()),
- getContainingTask(),
- getRuntimeContext(),
+ stateKey,
context,
- config.getOperatorID());
+ this);
}
@Override
diff --git a/flink-ml-core/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCacheTest.java b/flink-ml-core/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCacheTest.java
new file mode 100644
index 00000000..372b4a1a
--- /dev/null
+++ b/flink-ml-core/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCacheTest.java
@@ -0,0 +1,151 @@
+/*
+ * 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.typeinfo.Types;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.common.typeutils.base.StringSerializer;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.ml.common.datastream.DataStreamUtils;
+import org.apache.flink.ml.util.TestUtils;
+import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.runtime.minicluster.MiniCluster;
+import org.apache.flink.runtime.minicluster.MiniClusterConfiguration;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.ExecutionCheckpointingOptions;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.sink.SinkFunction;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.BoundedOneInput;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.StreamOperatorStateHandler;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+
+import org.apache.commons.lang3.RandomStringUtils;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+/** Tests {@link ListStateWithCache}. */
+public class ListStateWithCacheTest {
+
+ @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
+
+ private MiniClusterConfiguration createMiniClusterConfiguration() {
+ Configuration configuration = new Configuration();
+ // Set managed memory size to a small value, so when the instance of ListStateWithCache
+ // tries to allocate memory, it will use up all managed memory assigned to itself.
+ configuration.set(TaskManagerOptions.MANAGED_MEMORY_SIZE, MemorySize.ofMebiBytes(16));
+ configuration.set(
+ ExecutionCheckpointingOptions.ENABLE_CHECKPOINTS_AFTER_TASKS_FINISH, true);
+ return new MiniClusterConfiguration.Builder()
+ .setConfiguration(configuration)
+ .setNumTaskManagers(2)
+ .setNumSlotsPerTaskManager(2)
+ .build();
+ }
+
+ @Test
+ public void testWithMemoryWeights() throws Exception {
+ try (MiniCluster miniCluster = new MiniCluster(createMiniClusterConfiguration())) {
+ miniCluster.start();
+ double[] weights = new double[] {10, 20, 20, 40, 0};
+ JobGraph jobGraph = getJobGraph(weights);
+ miniCluster.executeJobBlocking(jobGraph);
+ }
+ }
+
+ private JobGraph getJobGraph(double[] weights) {
+ Configuration configuration = new Configuration();
+ StreamExecutionEnvironment env = TestUtils.getExecutionEnvironment(configuration);
+ env.setParallelism(4);
+
+ final int n = 10;
+ DataStream<String> data =
+ env.fromSequence(1, n).map(d -> RandomStringUtils.randomAlphabetic(1024 * 1024));
+ DataStream<Integer> counter =
+ data.transform("cache", Types.INT, new CacheDataOperator(weights));
+ DataStreamUtils.setManagedMemoryWeight(counter, 100);
+ DataStream<Integer> sum = DataStreamUtils.reduce(counter, Integer::sum);
+ sum.addSink(
+ new SinkFunction<Integer>() {
+ @Override
+ public void invoke(Integer value, Context context) {
+ Assert.assertEquals((Integer) (n * weights.length), value);
+ }
+ });
+ return env.getStreamGraph().getJobGraph();
+ }
+
+ private static class CacheDataOperator extends AbstractStreamOperator<Integer>
+ implements OneInputStreamOperator<String, Integer>,
+ BoundedOneInput,
+ StreamOperatorStateHandler.CheckpointedStreamOperator {
+ private final double[] weights;
+ private transient ListStateWithCache<String>[] cached;
+
+ public CacheDataOperator(double[] weights) {
+ this.weights = weights;
+ }
+
+ @Override
+ public void processElement(StreamRecord<String> element) throws Exception {
+ for (int i = 0; i < weights.length; i += 1) {
+ cached[i].add(element.getValue());
+ }
+ }
+
+ @Override
+ public void endInput() throws Exception {
+ int counter = 0;
+ for (int i = 0; i < weights.length; i += 1) {
+ for (String ignored : cached[i].get()) {
+ counter += 1;
+ }
+ cached[i].clear();
+ }
+ output.collect(new StreamRecord<>(counter));
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public void initializeState(StateInitializationContext context) throws Exception {
+ StreamTask<?, ?> containingTask = getContainingTask();
+ final OperatorID operatorID = getOperatorID();
+ TypeSerializer<String>[] serializers =
+ (TypeSerializer<String>[]) new TypeSerializer[weights.length];
+ final OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(containingTask, operatorID);
+ for (int i = 0; i < weights.length; i += 1) {
+ serializers[i] = StringSerializer.INSTANCE;
+ manager.register("state-" + i, weights[i]);
+ }
+ cached = new ListStateWithCache[weights.length];
+ for (int i = 0; i < weights.length; i += 1) {
+ cached[i] = new ListStateWithCache<>(serializers[i], "state-" + i, context, this);
+ }
+ }
+ }
+}
diff --git a/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCache.java b/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCache.java
index 3853f622..788e91ec 100644
--- a/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCache.java
+++ b/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/ListStateWithCache.java
@@ -23,11 +23,14 @@ 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.execution.Environment;
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.graph.StreamConfig;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.table.runtime.util.LazyMemorySegmentPool;
@@ -46,7 +49,8 @@ import java.util.List;
* <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}.
+ * org.apache.flink.api.dag.Transformation#declareManagedMemoryUseCaseAtOperatorScope} and register
+ * usage through {@link OperatorScopeManagedMemoryManager#register}.
*
* <p>NOTE: Users need to explicitly invoke this class's {@link
* #snapshotState(StateSnapshotContext)} method in order to store the recorded data in snapshot.
@@ -62,40 +66,53 @@ public class ListStateWithCache<T> implements ListState<T> {
/** The data cache writer for the received records. */
private final DataCacheWriter<T> dataCacheWriter;
+ /**
+ * Creates an instance of {@link ListStateWithCache}.
+ *
+ * @param serializer The type serializer of data.
+ * @param key The key registered in the manager.
+ * @param stateInitializationContext The state initialization state.
+ * @param operator The operator to new the instance.
+ */
@SuppressWarnings("unchecked")
public ListStateWithCache(
TypeSerializer<T> serializer,
- StreamTask<?, ?> containingTask,
- StreamingRuntimeContext runtimeContext,
+ String key,
StateInitializationContext stateInitializationContext,
- OperatorID operatorID)
+ AbstractStreamOperator<?> operator)
throws IOException {
this.serializer = serializer;
+ final StreamConfig config = operator.getOperatorConfig();
+ final OperatorID operatorID = config.getOperatorID();
+ StreamTask<?, ?> containingTask = operator.getContainingTask();
+ final Environment environment = containingTask.getEnvironment();
+ final StreamingRuntimeContext runtimeContext = operator.getRuntimeContext();
+
MemorySegmentPool segmentPool = null;
double fraction =
- containingTask
- .getConfiguration()
- .getManagedMemoryFractionOperatorUseCaseOfSlot(
- ManagedMemoryUseCase.OPERATOR,
- runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(),
- runtimeContext.getUserCodeClassLoader());
+ config.getManagedMemoryFractionOperatorUseCaseOfSlot(
+ ManagedMemoryUseCase.OPERATOR,
+ runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(),
+ runtimeContext.getUserCodeClassLoader());
if (fraction > 0) {
- MemoryManager memoryManager = containingTask.getEnvironment().getMemoryManager();
- segmentPool =
- new LazyMemorySegmentPool(
- containingTask,
- memoryManager,
- memoryManager.computeNumberOfPages(fraction));
+ OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(containingTask, operatorID);
+ double memorySubFraction = manager.getFraction(key);
+ if (memorySubFraction > 0) {
+ MemoryManager memoryManager = environment.getMemoryManager();
+ segmentPool =
+ new LazyMemorySegmentPool(
+ operator,
+ memoryManager,
+ memoryManager.computeNumberOfPages(fraction * memorySubFraction));
+ }
}
basePath =
OperatorUtils.getDataCachePath(
- containingTask.getEnvironment().getTaskManagerInfo().getConfiguration(),
- containingTask
- .getEnvironment()
- .getIOManager()
- .getSpillingDirectoriesPaths());
+ environment.getTaskManagerInfo().getConfiguration(),
+ environment.getIOManager().getSpillingDirectoriesPaths());
List<StatePartitionStreamProvider> inputs =
IteratorUtils.toList(
@@ -104,7 +121,7 @@ public class ListStateWithCache<T> implements ListState<T> {
inputs.size() < 2, "The input from raw operator state should be one or zero.");
List<Segment> priorFinishedSegments = new ArrayList<>();
- if (inputs.size() > 0) {
+ if (!inputs.isEmpty()) {
DataCacheSnapshot dataCacheSnapshot =
DataCacheSnapshot.recover(
inputs.get(0).getStream(),
diff --git a/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManager.java b/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManager.java
new file mode 100644
index 00000000..908c328a
--- /dev/null
+++ b/flink-ml-iteration/flink-ml-iteration-common/src/main/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManager.java
@@ -0,0 +1,105 @@
+/*
+ * 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.annotation.PublicEvolving;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.util.Preconditions;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.WeakHashMap;
+
+/**
+ * A manager for operator-scope managed memory.
+ *
+ * <p>Every operator must call {@link #getOrCreate} to get an instance to manage usages of
+ * operator-scope managed memory.
+ *
+ * <p>In the operator, for every usage of operator-scope managed memory, {@link #register} is called
+ * to declare its weight of memory usage with a key. Then, the fraction of memory can be obtained by
+ * calling {@link #getFraction}. Note that all calls of {@link #register} must be before those of
+ * {@link #getFraction}.
+ */
+@PublicEvolving
+public class OperatorScopeManagedMemoryManager {
+
+ /**
+ * Stores instances corresponding to operators. The instance is expected to be released after
+ * some point after the corresponding task or operator ID is garbage-collected.
+ */
+ private static final Map<StreamTask<?, ?>, Map<OperatorID, OperatorScopeManagedMemoryManager>>
+ managers = Collections.synchronizedMap(new WeakHashMap<>());
+
+ /** Stores keys and weights of memory usages. */
+ private final Map<String, Double> weights = new HashMap<>();
+ /** Indicates whether the `weights` is frozen. */
+ private boolean frozen = false;
+ /** Stores sum of weights of all usages. */
+ private double sum;
+
+ private OperatorScopeManagedMemoryManager() {}
+
+ /**
+ * Gets or creates an instance identified by the containing task and operator ID.
+ *
+ * @param containingTask The containing task.
+ * @param operatorID The operator ID.
+ * @return An instance of {@link OperatorScopeManagedMemoryManager}.
+ */
+ public static OperatorScopeManagedMemoryManager getOrCreate(
+ StreamTask<?, ?> containingTask, OperatorID operatorID) {
+ Map<OperatorID, OperatorScopeManagedMemoryManager> m =
+ managers.computeIfAbsent(containingTask, (key) -> new WeakHashMap<>());
+ return m.computeIfAbsent(operatorID, (key) -> new OperatorScopeManagedMemoryManager());
+ }
+
+ /**
+ * Registers a usage of operator-scope managed memory with memory weight declared.
+ *
+ * @param key The key to identify the usage.
+ * @param weight The weight of memory usage.
+ */
+ public void register(String key, double weight) {
+ Preconditions.checkState(!frozen, "Cannot call register after getFraction is called.");
+ Preconditions.checkState(
+ !weights.containsKey(key),
+ String.format("Cannot set a same key \"%s\" twice.", key));
+ Preconditions.checkArgument(weight >= 0, "The weight must be non-negative.");
+ weights.put(key, weight);
+ }
+
+ /**
+ * Gets the fraction of operator-scope managed memory for the usage.
+ *
+ * @param key The key to identify the usage.
+ * @return The fraction.
+ */
+ public double getFraction(String key) {
+ Preconditions.checkArgument(
+ weights.containsKey(key), String.format("%s is not registered yet.", key));
+ if (!frozen) {
+ frozen = true;
+ sum = weights.values().stream().mapToDouble(d -> d).sum();
+ }
+ return sum > 0 ? weights.get(key) / sum : 0;
+ }
+}
diff --git a/flink-ml-iteration/flink-ml-iteration-common/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManagerTest.java b/flink-ml-iteration/flink-ml-iteration-common/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManagerTest.java
new file mode 100644
index 00000000..c5234dcb
--- /dev/null
+++ b/flink-ml-iteration/flink-ml-iteration-common/src/test/java/org/apache/flink/iteration/datacache/nonkeyed/OperatorScopeManagedMemoryManagerTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.runtime.execution.Environment;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+
+import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+/** Tests the behavior of the {@link OperatorScopeManagedMemoryManager}. */
+public class OperatorScopeManagedMemoryManagerTest {
+
+ private static final double EPS = 1e-9;
+
+ private static class MockStreamTask
+ extends StreamTask<Integer, OneInputStreamOperator<Integer, Integer>> {
+ protected MockStreamTask(Environment env) throws Exception {
+ super(env);
+ }
+
+ @Override
+ protected void init() {}
+ }
+
+ @Test
+ public void testUsage() throws Exception {
+ OperatorID operatorID = new OperatorID();
+ OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(
+ new MockStreamTask(new DummyEnvironment()), operatorID);
+ manager.register("state-1", 100);
+ manager.register("state-2", 400);
+ Assert.assertEquals(manager.getFraction("state-1"), 0.2, EPS);
+ Assert.assertEquals(manager.getFraction("state-2"), 0.8, EPS);
+ }
+
+ @Test
+ public void testZeroUsage() throws Exception {
+ OperatorID operatorID = new OperatorID();
+ OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(
+ new MockStreamTask(new DummyEnvironment()), operatorID);
+ manager.register("state-1", 0);
+ manager.register("state-2", 0);
+ Assert.assertEquals(manager.getFraction("state-1"), 0, EPS);
+ Assert.assertEquals(manager.getFraction("state-2"), 0, EPS);
+ }
+
+ @Test
+ public void testInvalidUsage() throws Exception {
+ OperatorID operatorID = new OperatorID();
+ OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(
+ new MockStreamTask(new DummyEnvironment()), operatorID);
+ try {
+ manager.register("state-1", 100);
+ Assert.assertEquals(manager.getFraction("state-1"), 1., EPS);
+ manager.register("state-2", 400);
+ Assert.fail();
+ } catch (Exception e) {
+ Assert.assertEquals(
+ IllegalStateException.class, ExceptionUtils.getRootCause(e).getClass());
+ }
+ }
+}
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/clustering/kmeans/KMeans.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/clustering/kmeans/KMeans.java
index 6fbf39d5..8aeb6b60 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/clustering/kmeans/KMeans.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/clustering/kmeans/KMeans.java
@@ -36,6 +36,7 @@ import org.apache.flink.iteration.IterationListener;
import org.apache.flink.iteration.Iterations;
import org.apache.flink.iteration.ReplayableDataStreamList;
import org.apache.flink.iteration.datacache.nonkeyed.ListStateWithCache;
+import org.apache.flink.iteration.datacache.nonkeyed.OperatorScopeManagedMemoryManager;
import org.apache.flink.iteration.operator.OperatorStateUtils;
import org.apache.flink.ml.api.Estimator;
import org.apache.flink.ml.common.datastream.DataStreamUtils;
@@ -51,12 +52,14 @@ import org.apache.flink.ml.linalg.typeinfo.VectorWithNormSerializer;
import org.apache.flink.ml.param.Param;
import org.apache.flink.ml.util.ParamUtils;
import org.apache.flink.ml.util.ReadWriteUtils;
+import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.internal.TableImpl;
@@ -239,13 +242,15 @@ public class KMeans implements Estimator<KMeans, KMeansModel>, KMeansParams<KMea
context.getOperatorStateStore()
.getListState(new ListStateDescriptor<>("centroids", type));
+ final StreamTask<?, ?> containingTask = getContainingTask();
+ final OperatorID operatorID = config.getOperatorID();
+ final OperatorScopeManagedMemoryManager manager =
+ OperatorScopeManagedMemoryManager.getOrCreate(containingTask, operatorID);
+ final String stateKey = "points-state";
+ manager.register(stateKey, 1.);
points =
new ListStateWithCache<>(
- new VectorWithNormSerializer(),
- getContainingTask(),
- getRuntimeContext(),
- context,
- config.getOperatorID());
+ new VectorWithNormSerializer(), stateKey, context, this);
}
@Override