Skip to content

Commit 60db2ff

Browse files
davidheryantofeast-ci-bot
davidheryanto
authored and
feast-ci-bot
committed
Ensure ImportJobTest is not flaky by checking WriteToStore metric and requesting adequate resources for testing (#332)
* Print Redis INFO when key not found in ImportJobTest Also print random Redis element to debug that some FeatureRow has been ingested properly * Request CPU and memory for the pod running tests Some tests (like ingestion test) expect the operation to complete in certain amount of time. This can only be guaranteed if the process have adequate CPU and memory. Without it, when the test cluster is overloaded, the test process may get little CPU time allocated and the expected completion time is no longer valid * Add metric WriteToStore:elements_written So we can obtain information about no of elements have been written in the pipeline without resorting to external metrics collector This method makes use built in metrics util in Apache Beam * Add a check for all elements to be written to store before checking the ingestion result in ImportJobTest
1 parent e8d1b01 commit 60db2ff

4 files changed

Lines changed: 147 additions & 37 deletions

File tree

‎.prow/config.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ presubmits:
5757
containers:
5858
- image: maven:3.6-jdk-8
5959
command: [".prow/scripts/test-core-ingestion.sh"]
60+
resources:
61+
requests:
62+
cpu: "1000m"
63+
memory: "512Mi"
64+
limit:
65+
memory: "4096Mi"
6066

6167
- name: test-serving
6268
decorate: true
@@ -97,6 +103,12 @@ presubmits:
97103
containers:
98104
- image: maven:3.6-jdk-8
99105
command: [".prow/scripts/test-end-to-end.sh"]
106+
resources:
107+
requests:
108+
cpu: "1000m"
109+
memory: "1024Mi"
110+
limit:
111+
memory: "4096Mi"
100112

101113
# TODO: do a release when a git tag is pushed
102114
#

‎ingestion/src/main/java/feast/ingestion/transform/WriteToStore.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,18 +41,29 @@
4141
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryInsertError;
4242
import org.apache.beam.sdk.io.gcp.bigquery.InsertRetryPolicy;
4343
import org.apache.beam.sdk.io.gcp.bigquery.WriteResult;
44+
import org.apache.beam.sdk.metrics.Counter;
45+
import org.apache.beam.sdk.metrics.Metrics;
4446
import org.apache.beam.sdk.transforms.DoFn;
47+
import org.apache.beam.sdk.transforms.MapElements;
4548
import org.apache.beam.sdk.transforms.PTransform;
4649
import org.apache.beam.sdk.transforms.ParDo;
4750
import org.apache.beam.sdk.values.PCollection;
4851
import org.apache.beam.sdk.values.PDone;
52+
import org.apache.beam.sdk.values.TypeDescriptors;
53+
import org.apache.beam.sdk.values.ValueInSingleWindow;
4954
import org.slf4j.Logger;
5055

5156
@AutoValue
5257
public abstract class WriteToStore extends PTransform<PCollection<FeatureRow>, PDone> {
5358

5459
private static final Logger log = org.slf4j.LoggerFactory.getLogger(WriteToStore.class);
5560

61+
public static final String METRIC_NAMESPACE = "WriteToStore";
62+
public static final String ELEMENTS_WRITTEN_METRIC = "elements_written";
63+
64+
private static final Counter elementsWritten = Metrics
65+
.counter(METRIC_NAMESPACE, ELEMENTS_WRITTEN_METRIC);
66+
5667
public abstract Store getStore();
5768

5869
public abstract Map<String, FeatureSetSpec> getFeatureSetSpecs();
@@ -140,6 +151,12 @@ public void processElement(ProcessContext context) {
140151
break;
141152
}
142153

154+
input.apply("IncrementWriteToStoreElementsWrittenCounter",
155+
MapElements.into(TypeDescriptors.booleans()).via((FeatureRow row) -> {
156+
elementsWritten.inc();
157+
return true;
158+
}));
159+
143160
return PDone.in(input.getPipeline());
144161
}
145162
}

‎ingestion/src/test/java/feast/ingestion/ImportJobTest.java‎

Lines changed: 52 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
import feast.types.ValueProto.ValueType.Enum;
3939
import java.io.IOException;
4040
import java.nio.charset.StandardCharsets;
41-
import java.time.Duration;
4241
import java.util.ArrayList;
4342
import java.util.Collections;
4443
import java.util.HashMap;
@@ -49,6 +48,7 @@
4948
import org.apache.beam.sdk.PipelineResult.State;
5049
import org.apache.beam.sdk.options.PipelineOptionsFactory;
5150
import org.apache.kafka.common.serialization.ByteArraySerializer;
51+
import org.joda.time.Duration;
5252
import org.junit.AfterClass;
5353
import org.junit.Assert;
5454
import org.junit.BeforeClass;
@@ -77,10 +77,15 @@ public class ImportJobTest {
7777
private static final String REDIS_HOST = "localhost";
7878
private static final int REDIS_PORT = 6380;
7979

80-
// Expected time taken for the import job to be ready to receive Feature Row input
81-
private static final int IMPORT_JOB_READY_DURATION_SEC = 5;
82-
// Expected time taken for the import job to finish writing to Store
83-
private static final int IMPORT_JOB_RUN_DURATION_SEC = 30;
80+
// No of samples of feature row that will be generated and used for testing.
81+
// Note that larger no of samples will increase completion time for ingestion.
82+
private static final int IMPORT_JOB_SAMPLE_FEATURE_ROW_SIZE = 128;
83+
// Expected time taken for the import job to be ready to receive Feature Row input.
84+
private static final int IMPORT_JOB_READY_DURATION_SEC = 10;
85+
// The interval between checks for import job to finish writing elements to store.
86+
private static final int IMPORT_JOB_CHECK_INTERVAL_DURATION_SEC = 5;
87+
// Max duration to wait until the import job finishes writing to Store.
88+
private static final int IMPORT_JOB_MAX_RUN_DURATION_SEC = 300;
8489

8590
@BeforeClass
8691
public static void setup() throws IOException, InterruptedException {
@@ -161,50 +166,60 @@ public void runPipeline_ShouldWriteToRedisCorrectlyGivenValidSpecAndFeatureRow()
161166
options.setProject("");
162167
options.setBlockOnRun(false);
163168

164-
int inputSize = 128;
165169
List<FeatureRow> input = new ArrayList<>();
166170
Map<RedisKey, FeatureRow> expected = new HashMap<>();
167171

168172
LOGGER.info("Generating test data ...");
169-
IntStream.range(0, inputSize)
170-
.forEach(
171-
i -> {
172-
FeatureRow randomRow = TestUtil.createRandomFeatureRow(spec);
173-
RedisKey redisKey = TestUtil.createRedisKey(spec, randomRow);
174-
input.add(randomRow);
175-
expected.put(redisKey, randomRow);
176-
});
173+
IntStream.range(0, IMPORT_JOB_SAMPLE_FEATURE_ROW_SIZE).forEach(i -> {
174+
FeatureRow randomRow = TestUtil.createRandomFeatureRow(spec);
175+
RedisKey redisKey = TestUtil.createRedisKey(spec, randomRow);
176+
input.add(randomRow);
177+
expected.put(redisKey, randomRow);
178+
});
177179

178180
LOGGER.info("Starting Import Job with the following options: {}", options.toString());
179181
PipelineResult pipelineResult = ImportJob.runPipeline(options);
180-
Thread.sleep(Duration.ofSeconds(IMPORT_JOB_READY_DURATION_SEC).toMillis());
182+
Thread.sleep(Duration.standardSeconds(IMPORT_JOB_READY_DURATION_SEC).getMillis());
181183
Assert.assertEquals(pipelineResult.getState(), State.RUNNING);
182184

183185
LOGGER.info("Publishing {} Feature Row messages to Kafka ...", input.size());
184-
TestUtil.publishFeatureRowsToKafka(
185-
KAFKA_BOOTSTRAP_SERVERS,
186-
KAFKA_TOPIC,
187-
input,
188-
ByteArraySerializer.class,
189-
KAFKA_PUBLISH_TIMEOUT_SEC);
190-
Thread.sleep(Duration.ofSeconds(IMPORT_JOB_RUN_DURATION_SEC).toMillis());
186+
TestUtil.publishFeatureRowsToKafka(KAFKA_BOOTSTRAP_SERVERS, KAFKA_TOPIC, input,
187+
ByteArraySerializer.class, KAFKA_PUBLISH_TIMEOUT_SEC);
188+
TestUtil.waitUntilAllElementsAreWrittenToStore(pipelineResult,
189+
Duration.standardSeconds(IMPORT_JOB_MAX_RUN_DURATION_SEC),
190+
Duration.standardSeconds(IMPORT_JOB_CHECK_INTERVAL_DURATION_SEC));
191191

192192
LOGGER.info("Validating the actual values written to Redis ...");
193193
Jedis jedis = new Jedis(REDIS_HOST, REDIS_PORT);
194-
expected.forEach(
195-
(key, expectedValue) -> {
196-
byte[] actualByteValue = jedis.get(key.toByteArray());
197-
Assert.assertNotNull("Key not found in Redis: " + key, actualByteValue);
198-
FeatureRow actualValue = null;
199-
try {
200-
actualValue = FeatureRow.parseFrom(actualByteValue);
201-
} catch (InvalidProtocolBufferException e) {
202-
Assert.fail(
203-
String.format(
204-
"Actual Redis value cannot be parsed as FeatureRow, key: %s, value :%s",
205-
key, new String(actualByteValue, StandardCharsets.UTF_8)));
206-
}
207-
Assert.assertEquals(expectedValue, actualValue);
208-
});
194+
expected.forEach((key, expectedValue) -> {
195+
196+
// Ensure ingested key exists.
197+
byte[] actualByteValue = jedis.get(key.toByteArray());
198+
if (actualByteValue == null) {
199+
LOGGER.error("Key not found in Redis: " + key);
200+
LOGGER.info("Redis INFO:");
201+
LOGGER.info(jedis.info());
202+
String randomKey = jedis.randomKey();
203+
if (randomKey != null) {
204+
LOGGER.info("Sample random key, value (for debugging purpose):");
205+
LOGGER.info("Key: " + randomKey);
206+
LOGGER.info("Value: " + jedis.get(randomKey));
207+
}
208+
Assert.fail("Missing key in Redis.");
209+
}
210+
211+
// Ensure value is a valid serialized FeatureRow object.
212+
FeatureRow actualValue = null;
213+
try {
214+
actualValue = FeatureRow.parseFrom(actualByteValue);
215+
} catch (InvalidProtocolBufferException e) {
216+
Assert.fail(String
217+
.format("Actual Redis value cannot be parsed as FeatureRow, key: %s, value :%s",
218+
key, new String(actualByteValue, StandardCharsets.UTF_8)));
219+
}
220+
221+
// Ensure the retrieved FeatureRow is equal to the ingested FeatureRow.
222+
Assert.assertEquals(expectedValue, actualValue);
223+
});
209224
}
210225
}

‎ingestion/src/test/java/feast/test/TestUtil.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.protobuf.ByteString;
2020
import com.google.protobuf.util.Timestamps;
2121
import feast.core.FeatureSetProto.FeatureSetSpec;
22+
import feast.ingestion.transform.WriteToStore;
2223
import feast.storage.RedisProto.RedisKey;
2324
import feast.types.FeatureRowProto.FeatureRow;
2425
import feast.types.FeatureRowProto.FeatureRow.Builder;
@@ -41,13 +42,18 @@
4142
import java.util.concurrent.TimeoutException;
4243
import kafka.server.KafkaConfig;
4344
import kafka.server.KafkaServerStartable;
45+
import org.apache.beam.sdk.PipelineResult;
46+
import org.apache.beam.sdk.PipelineResult.State;
47+
import org.apache.beam.sdk.metrics.MetricResult;
48+
import org.apache.beam.sdk.metrics.MetricResults;
4449
import org.apache.commons.lang3.RandomStringUtils;
4550
import org.apache.kafka.clients.producer.KafkaProducer;
4651
import org.apache.kafka.clients.producer.Producer;
4752
import org.apache.kafka.clients.producer.ProducerRecord;
4853
import org.apache.kafka.common.serialization.LongSerializer;
4954
import org.apache.zookeeper.server.ServerConfig;
5055
import org.apache.zookeeper.server.ZooKeeperServerMain;
56+
import org.joda.time.Duration;
5157
import redis.embedded.RedisServer;
5258

5359
@SuppressWarnings("WeakerAccess")
@@ -343,4 +349,64 @@ public static Field field(String name, Object value, ValueType.Enum valueType) {
343349
throw new IllegalStateException("Unexpected valueType: " + value.getClass());
344350
}
345351
}
352+
353+
/**
354+
* This blocking method waits until an ImportJob pipeline has written all elements to the store.
355+
* <p>
356+
* The pipeline must be in the RUNNING state before calling this method.
357+
*
358+
* @param pipelineResult result of running the Pipeline
359+
* @param maxWaitDuration wait until this max amount of duration
360+
* @throws InterruptedException if the thread is interruped while waiting
361+
*/
362+
public static void waitUntilAllElementsAreWrittenToStore(PipelineResult pipelineResult,
363+
Duration maxWaitDuration, Duration checkInterval) throws InterruptedException {
364+
if (pipelineResult.getState().isTerminal()) {
365+
return;
366+
}
367+
368+
if (!pipelineResult.getState().equals(State.RUNNING)) {
369+
throw new IllegalArgumentException(
370+
"Pipeline must be in RUNNING state before calling this method.");
371+
}
372+
373+
MetricResults metricResults;
374+
try {
375+
metricResults = pipelineResult.metrics();
376+
} catch (UnsupportedOperationException e) {
377+
// Runner does not support metrics so we just wait as long as we are allowed to.
378+
Thread.sleep(maxWaitDuration.getMillis());
379+
return;
380+
}
381+
382+
String writeToStoreMetric =
383+
WriteToStore.METRIC_NAMESPACE + ":" + WriteToStore.ELEMENTS_WRITTEN_METRIC;
384+
long committed = 0;
385+
long maxSystemTimeMillis = System.currentTimeMillis() + maxWaitDuration.getMillis();
386+
387+
while (System.currentTimeMillis() <= maxSystemTimeMillis) {
388+
Thread.sleep(checkInterval.getMillis());
389+
390+
for (MetricResult<Long> metricResult : metricResults.allMetrics().getCounters()) {
391+
// We are only concerned with the metric: count of elements that have been
392+
// written to the store.
393+
if (!metricResult.getName().toString().contains(writeToStoreMetric)) {
394+
continue;
395+
}
396+
try {
397+
// If between check interval, no more changes in the no of committed elements
398+
// we can assume the pipeline has finished writing all the elements to store.
399+
if (metricResult.getCommitted() == committed) {
400+
return;
401+
}
402+
committed = metricResult.getCommitted();
403+
break;
404+
} catch (UnsupportedOperationException e) {
405+
// Runner does not support committed metrics so we just wait as long as we are allowed to.
406+
Thread.sleep(maxWaitDuration.getMillis());
407+
return;
408+
}
409+
}
410+
}
411+
}
346412
}

0 commit comments

Comments
 (0)