Skip to content

Commit 158a887

Browse files
pyalexpyalex
committed
generalize ingestion pipeline metrics
Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent a498dae commit 158a887

10 files changed

Lines changed: 109 additions & 92 deletions

File tree

‎spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ trait BasePipeline {
5555
.set("spark.metrics.conf.*.sink.statsd.port", c.port.toString)
5656
.set("spark.metrics.conf.*.sink.statsd.period", "30")
5757
.set("spark.metrics.conf.*.sink.statsd.unit", "seconds")
58-
.set("spark.metrics.namespace", jobConfig.mode.toString.toLowerCase)
58+
.set("spark.metrics.namespace", s"feast_${jobConfig.mode.toString.toLowerCase}")
5959
// until proto parser udf will be fixed, we have to use this
6060
.set("spark.sql.legacy.allowUntypedScalaUDF", "true")
6161
case None => ()

‎spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
*/
1717
package feast.ingestion
1818

19+
import feast.ingestion.metrics.IngestionPipelineMetrics
1920
import feast.ingestion.sources.bq.BigQueryReader
2021
import feast.ingestion.sources.file.FileReader
21-
import feast.ingestion.stores.deadletters.DeadLetterMetrics
2222
import feast.ingestion.validation.{RowValidator, TypeCheck}
2323
import org.apache.commons.lang.StringUtils
2424
import org.apache.spark.SparkEnv
@@ -59,13 +59,16 @@ object BatchPipeline extends BasePipeline {
5959

6060
val projected = input.select(projection: _*).cache()
6161

62+
implicit def rowEncoder: Encoder[Row] = RowEncoder(projected.schema)
63+
6264
TypeCheck.allTypesMatch(projected.schema, featureTable) match {
6365
case Some(error) =>
6466
throw new RuntimeException(s"Dataframe columns don't match expected feature types: $error")
6567
case _ => ()
6668
}
6769

6870
val validRows = projected
71+
.mapPartitions(IngestionPipelineMetrics.incrementRead)
6972
.filter(rowValidator.allChecks)
7073

7174
validRows.write
@@ -77,12 +80,11 @@ object BatchPipeline extends BasePipeline {
7780
.option("max_age", config.featureTable.maxAge.getOrElse(0L))
7881
.save()
7982

80-
implicit def rowEncoder: Encoder[Row] = RowEncoder(projected.schema)
8183
config.deadLetterPath match {
8284
case Some(path) =>
8385
projected
8486
.filter(!rowValidator.allChecks)
85-
.mapPartitions(iter => DeadLetterMetrics.incrementCount(iter))
87+
.mapPartitions(IngestionPipelineMetrics.incrementDeadletters)
8688
.write
8789
.format("parquet")
8890
.mode(SaveMode.Append)

‎spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@ package feast.ingestion
1919
import java.io.File
2020
import java.util.concurrent.TimeUnit
2121

22+
import feast.ingestion.metrics.IngestionPipelineMetrics
2223
import feast.ingestion.registry.proto.ProtoRegistryFactory
2324
import org.apache.spark.sql.{DataFrame, Encoder, Row, SaveMode, SparkSession}
2425
import org.apache.spark.sql.functions.{expr, struct, udf}
25-
import feast.ingestion.stores.deadletters.DeadLetterMetrics
2626
import feast.ingestion.utils.ProtoReflection
2727
import feast.ingestion.utils.testing.MemoryStreamingSource
2828
import feast.ingestion.validation.{RowValidator, TypeCheck}
@@ -104,8 +104,10 @@ object StreamingPipeline extends BasePipeline with Serializable {
104104
batchDF.withColumn("_isValid", rowValidator.allChecks)
105105
}
106106
rowsAfterValidation.persist()
107+
implicit def rowEncoder: Encoder[Row] = RowEncoder(rowsAfterValidation.schema)
107108

108109
rowsAfterValidation
110+
.mapPartitions(IngestionPipelineMetrics.incrementRead)
109111
.filter(if (config.doNotIngestInvalidRows) expr("_isValid") else rowValidator.allChecks)
110112
.write
111113
.format("feast.ingestion.stores.redis")
@@ -116,13 +118,12 @@ object StreamingPipeline extends BasePipeline with Serializable {
116118
.option("max_age", config.featureTable.maxAge.getOrElse(0L))
117119
.save()
118120

119-
implicit def rowEncoder: Encoder[Row] = RowEncoder(projected.schema)
120121
config.deadLetterPath match {
121122
case Some(path) =>
122123

123124
rowsAfterValidation
124125
.filter("!_isValid")
125-
.mapPartitions(iter => DeadLetterMetrics.incrementCount(iter))
126+
.mapPartitions(IngestionPipelineMetrics.incrementDeadletters)
126127
.write
127128
.format("parquet")
128129
.mode(SaveMode.Append)

‎spark/ingestion/src/main/scala/feast/ingestion/metrics/IngestionPipelineMetrics.scala‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion.metrics
18+
19+
import org.apache.spark.SparkEnv
20+
import org.apache.spark.metrics.source.IngestionPipelineMetricSource
21+
22+
object IngestionPipelineMetrics {
23+
def incrementDeadletters[A](rowIterator: Iterator[A]): Iterator[A] = {
24+
if (metricSource.nonEmpty)
25+
metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc(rowIterator.length)
26+
27+
rowIterator
28+
}
29+
30+
def incrementRead[A](rowIterator: Iterator[A]): Iterator[A] = {
31+
if (metricSource.nonEmpty)
32+
metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc(rowIterator.length)
33+
34+
rowIterator
35+
}
36+
37+
private lazy val metricSource: Option[IngestionPipelineMetricSource] = {
38+
this.synchronized {
39+
if (
40+
SparkEnv.get.metricsSystem
41+
.getSourcesByName(IngestionPipelineMetricSource.sourceName)
42+
.isEmpty
43+
) {
44+
SparkEnv.get.metricsSystem.registerSource(new IngestionPipelineMetricSource)
45+
}
46+
}
47+
48+
SparkEnv.get.metricsSystem.getSourcesByName(IngestionPipelineMetricSource.sourceName) match {
49+
case Seq(head) => Some(head.asInstanceOf[IngestionPipelineMetricSource])
50+
case _ => None
51+
}
52+
}
53+
}

‎spark/ingestion/src/main/scala/feast/ingestion/stores/deadletters/DeadLetterMetrics.scala‎

Lines changed: 0 additions & 47 deletions
This file was deleted.

spark/ingestion/src/main/scala/org/apache/spark/metrics/source/DeadLetterSinkMetricSource.scala renamed to spark/ingestion/src/main/scala/org/apache/spark/metrics/source/IngestionPipelineMetricSource.scala

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,16 @@
1616
*/
1717
package org.apache.spark.metrics.source
1818

19-
class DeadLetterSinkMetricSource extends BaseMetricSource {
20-
override val sourceName: String = DeadLetterSinkMetricSource.sourceName
19+
class IngestionPipelineMetricSource extends BaseMetricSource {
20+
override val sourceName: String = IngestionPipelineMetricSource.sourceName
2121

2222
val METRIC_DEADLETTER_ROWS_INSERTED =
23-
metricRegistry.counter(counterWithLabels("feast_ingestion_deadletter_count"))
23+
metricRegistry.counter(counterWithLabels("deadletter_count"))
24+
25+
val METRIC_ROWS_READ_FROM_SOURCE =
26+
metricRegistry.counter(counterWithLabels("read_from_source_count"))
2427
}
2528

26-
object DeadLetterSinkMetricSource {
27-
val sourceName = "deadletter_sink"
29+
object IngestionPipelineMetricSource {
30+
val sourceName = "ingestion_pipeline"
2831
}

‎spark/ingestion/src/main/scala/org/apache/spark/metrics/source/RedisSinkMetricSource.scala‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ class RedisSinkMetricSource extends BaseMetricSource {
2020
override val sourceName: String = RedisSinkMetricSource.sourceName
2121

2222
val METRIC_TOTAL_ROWS_INSERTED =
23-
metricRegistry.counter(counterWithLabels("feast_ingestion_feature_row_ingested_count"))
23+
metricRegistry.counter(counterWithLabels("feature_row_ingested_count"))
2424

2525
val METRIC_ROWS_LAG =
26-
metricRegistry.histogram(metricWithLabels("feast_ingestion_feature_row_lag_ms"))
26+
metricRegistry.histogram(metricWithLabels("feature_row_lag_ms"))
2727
}
2828

2929
object RedisSinkMetricSource {

‎spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import org.scalatest._
3030
import redis.clients.jedis.Jedis
3131
import feast.ingestion.helpers.RedisStorageHelper._
3232
import feast.ingestion.helpers.DataHelper._
33+
import feast.ingestion.metrics.StatsDStub
3334
import feast.proto.storage.RedisProto.RedisKeyV2
3435
import feast.proto.types.ValueProto
3536
import org.apache.spark.sql.Encoder
@@ -55,6 +56,7 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
5556
jedis.flushAll()
5657

5758
implicit def testRowEncoder: Encoder[TestRow] = ExpressionEncoder()
59+
val statsDStub = new StatsDStub
5860

5961
def rowGenerator(start: DateTime, end: DateTime, customerGen: Option[Gen[String]] = None) =
6062
for {
@@ -95,7 +97,8 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
9597
)
9698
),
9799
startTime = DateTime.parse("2020-08-01"),
98-
endTime = DateTime.parse("2020-09-01")
100+
endTime = DateTime.parse("2020-09-01"),
101+
metrics = Some(StatsDConfig(host="localhost", port=statsDStub.port))
99102
)
100103
}
101104

‎spark/ingestion/src/test/scala/feast/ingestion/metrics/StatsDStub.scala‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package feast.ingestion.metrics
2+
3+
import java.net.{DatagramPacket, DatagramSocket, SocketTimeoutException}
4+
5+
import scala.collection.mutable.ArrayBuffer
6+
7+
class StatsDStub {
8+
val socket = new DatagramSocket()
9+
socket.setSoTimeout(100)
10+
11+
def port: Int = socket.getLocalPort
12+
13+
def receive: Array[String] = {
14+
val messages: ArrayBuffer[String] = ArrayBuffer()
15+
var finished = false
16+
17+
do {
18+
val buf = new Array[Byte](65535)
19+
val p = new DatagramPacket(buf, buf.length)
20+
try {
21+
socket.receive(p)
22+
} catch {
23+
case _: SocketTimeoutException =>
24+
finished = true
25+
}
26+
messages += new String(p.getData, 0, p.getLength)
27+
} while (!finished)
28+
29+
messages.toArray
30+
}
31+
}

‎spark/ingestion/src/test/scala/feast/ingestion/metrics/StatsReporterSpec.scala‎

Lines changed: 1 addition & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -16,46 +16,17 @@
1616
*/
1717
package feast.ingestion.metrics
1818

19-
import java.net.{DatagramPacket, DatagramSocket, SocketTimeoutException}
2019
import java.util
2120
import java.util.Collections
2221

2322
import com.codahale.metrics.{Gauge, Histogram, MetricRegistry, UniformReservoir}
2423
import feast.ingestion.UnitSpec
2524

26-
import scala.collection.mutable.ArrayBuffer
2725
import scala.jdk.CollectionConverters._
2826

2927
class StatsReporterSpec extends UnitSpec {
30-
class SimpleServer {
31-
val socket = new DatagramSocket()
32-
socket.setSoTimeout(100)
33-
34-
def port: Int = socket.getLocalPort
35-
36-
def receive: Array[String] = {
37-
val messages: ArrayBuffer[String] = ArrayBuffer()
38-
var finished = false
39-
40-
do {
41-
val buf = new Array[Byte](65535)
42-
val p = new DatagramPacket(buf, buf.length)
43-
try {
44-
socket.receive(p)
45-
} catch {
46-
case _: SocketTimeoutException => {
47-
finished = true
48-
}
49-
}
50-
messages += new String(p.getData, 0, p.getLength)
51-
} while (!finished)
52-
53-
messages.toArray
54-
}
55-
}
56-
5728
trait Scope {
58-
val server = new SimpleServer
29+
val server = new StatsDStub
5930
val reporter = new StatsdReporterWithTags(
6031
new MetricRegistry,
6132
"127.0.0.1",

0 commit comments

Comments
 (0)