Connect drift metrics computation in Spark with Hub for serving to frontend (#95)

## Summary
Builds on a couple of the summary computation PRs and data generation to
wire things up so that Hub can serve them.
* Yanked out mock data based endpoints (model perf / drift, join &
feature skew) - decided it would be confusing to have a mix of mock and
generated data so we just have the generated data served
* Dropped a few of the scripts introduced in
#87. We bring up our
containers the way and we have a script `load_summaries.sh` that we can
trigger that leverages the existing app container to load data.
* DDB ingestion was taking too long and we were dropping a lot of data
due to rejected execution exceptions. To unblock for now, we've gone
with an approach of making a bulk put HTTP call from the
ObservabilityDemo app -> Hub and Hub utilizing a InMemoryKV store to
persist and serve up features.
* Added an endpoint to serve the join that are configured as we've
switched from the model based world.

There's still an issue to resolve around fetching individual feature
series data. Once I resolve that, we can switch this PR out of wip mode.

To test / run:
start up our docker containers:
```
$ docker-compose -f docker-init/compose.yaml up --build
...
```
In a different term load data:
```
$ ./docker-init/demo/load_summaries.sh 
Done uploading summaries! 🥳
```

You can now curl join & feature time series data.
Join drift (null ratios)
```
curl -X GET   'http://localhost:9000/api/v1/join/risk.user_transactions.txn_join/timeseries?startTs=1673308800000&endTs=1674172800000&metricType=drift&metrics=null&offset=10h&algorithm=psi'
```

Join drift (value drift)
```
curl -X GET   'http://localhost:9000/api/v1/join/risk.user_transactions.txn_join/timeseries?startTs=1673308800000&endTs=1674172800000&metricType=drift&metrics=value&offset=10h&algorithm=psi'
```

Feature drift:
```
curl -X GET   'http://localhost:9000/api/v1/join/risk.user_transactions.txn_join/feature/dim_user_account_type/timeseries?startTs=1673308800000&endTs=1674172800000&metricType=drift&metrics=value&offset=1D&algorithm=psi&granularity=aggregates'
```

Feature summaries:
```
curl -X GET   'http://localhost:9000/api/v1/join/risk.user_transactions.txn_join/feature/dim_user_account_type/timeseries?startTs=1673308800000&endTs=1674172800000&metricType=drift&metrics=value&offset=1D&algorithm=psi&granularity=percentile'
```

Join metadata
```
curl -X GET 'http://localhost:9000/api/v1/joins'
curl -X GET 'http://localhost:9000/api/v1/join/risk.user_transactions.txn_join'
```

## Checklist
- [X] Added Unit Tests
- [ ] Covered by existing CI
- [X] Integration tested
- [ ] Documentation update



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
- Introduced a new `JoinController` for managing joins with pagination
support.
- Added functionality for an in-memory key-value store with bulk data
upload capabilities.
- Implemented observability demo data loading within a Spark
application.
- Added a new `HTTPKVStore` class for remote key-value store
interactions over HTTP.

- **Improvements**
- Enhanced the `ModelController` and `SearchController` to align with
the new join data structure.
- Updated the `TimeSeriesController` to support asynchronous operations
and improved error handling.
- Refined dependency management in the build configuration for better
clarity and maintainability.
- Updated API routes to include new endpoints for listing and retrieving
joins.
- Updated configuration to replace the `DynamoDBModule` with
`ModelStoreModule`, adding `InMemoryKVStoreModule` and
`DriftStoreModule`.

- **Documentation**
- Revised README instructions for Docker container setup and demo data
loading.
- Updated API routes documentation to reflect new endpoints for joins
and in-memory data operations.

- **Bug Fixes**
- Resolved issues related to error handling in various controllers and
improved logging for better traceability.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: nikhil-zlai <[email protected]>

Author: piyush-zlai

build.sbt

@@ -80,6 +80,13 @@ val jackson = Seq(
   "com.fasterxml.jackson.module" %% "jackson-module-scala"
 ).map(_ % jackson_2_15)
 
+// Circe is used to ser / deser case class payloads for the Hub Play webservice
+val circe = Seq(
+  "io.circe" %% "circe-core",
+  "io.circe" %% "circe-generic",
+  "io.circe" %% "circe-parser",
+).map(_ % circeVersion)
+
 val flink_all = Seq(
   "org.apache.flink" %% "flink-streaming-scala",
   "org.apache.flink" % "flink-metrics-dropwizard",
@@ -129,6 +136,8 @@ lazy val online = project
       "com.github.ben-manes.caffeine" % "caffeine" % "3.1.8"
     ),
     libraryDependencies ++= jackson,
+    // dep needed for HTTPKvStore - yank when we rip this out
+    libraryDependencies += "com.softwaremill.sttp.client3" %% "core" % "3.9.7",
     libraryDependencies ++= spark_all.map(_ % "provided"),
     libraryDependencies ++= flink_all.map(_ % "provided")
   )
@@ -236,20 +245,18 @@ lazy val frontend = (project in file("frontend"))
 // build interop between one module solely on 2.13 and others on 2.12 is painful
 lazy val hub = (project in file("hub"))
   .enablePlugins(PlayScala)
-  .dependsOn(cloud_aws)
+  .dependsOn(cloud_aws, spark)
   .settings(
     name := "hub",
     libraryDependencies ++= Seq(
       guice,
       "org.scalatestplus.play" %% "scalatestplus-play" % "5.1.0" % Test,
       "org.scalatestplus" %% "mockito-3-4" % "3.2.10.0" % "test",
-      "io.circe" %% "circe-core" % circeVersion,
-      "io.circe" %% "circe-generic" % circeVersion,
-      "io.circe" %% "circe-parser" % circeVersion,
       "org.scala-lang.modules" %% "scala-xml" % "2.1.0",
       "org.scala-lang.modules" %% "scala-parser-combinators" % "2.3.0",
       "org.scala-lang.modules" %% "scala-java8-compat" % "1.0.2"
     ),
+    libraryDependencies ++= circe,
     libraryDependencySchemes ++= Seq(
       "org.scala-lang.modules" %% "scala-xml" % VersionScheme.Always,
       "org.scala-lang.modules" %% "scala-parser-combinators" % VersionScheme.Always,
@@ -258,7 +265,10 @@ lazy val hub = (project in file("hub"))
     excludeDependencies ++= Seq(
       ExclusionRule(organization = "org.slf4j", name = "slf4j-log4j12"),
       ExclusionRule(organization = "log4j", name = "log4j"),
-      ExclusionRule(organization = "org.apache.logging.log4j", name = "log4j-to-slf4j")
+      ExclusionRule(organization = "org.apache.logging.log4j", name = "log4j-to-slf4j"),
+      ExclusionRule("org.apache.logging.log4j", "log4j-slf4j-impl"),
+      ExclusionRule("org.apache.logging.log4j", "log4j-core"),
+      ExclusionRule("org.apache.logging.log4j", "log4j-api")
     ),
     // Ensure consistent versions of logging libraries
     dependencyOverrides ++= Seq(

online/src/main/scala/ai/chronon/online/HTTPKVStore.scala

@@ -0,0 +1,57 @@
+package ai.chronon.online
+
+import ai.chronon.online.KVStore.PutRequest
+import sttp.client3._
+import sttp.model.StatusCode
+
+import java.util.Base64
+import scala.concurrent.Future
+
+// Hacky test kv store that we use to send objects to the in-memory KV store that lives in a different JVM (e.g spark -> hub)
+class HTTPKVStore(host: String = "localhost", port: Int = 9000) extends KVStore with Serializable {
+
+  val backend: SttpBackend[Identity, Any] = HttpClientSyncBackend()
+  val baseUrl: String = s"http://$host:$port/api/v1/dataset"
+
+  override def multiGet(requests: collection.Seq[KVStore.GetRequest]): Future[collection.Seq[KVStore.GetResponse]] = ???
+
+  override def multiPut(putRequests: collection.Seq[KVStore.PutRequest]): Future[collection.Seq[Boolean]] = {
+    if (putRequests.isEmpty) {
+      Future.successful(Seq.empty)
+    } else {
+      Future {
+        basicRequest
+          .post(uri"$baseUrl/data")
+          .header("Content-Type", "application/json")
+          .body(jsonList(putRequests))
+          .send(backend)
+      }.map { response =>
+        response.code match {
+          case StatusCode.Ok => Seq(true)
+          case _ =>
+            logger.error(s"HTTP multiPut failed with status ${response.code}: ${response.body}")
+            Seq(false)
+        }
+      }
+    }
+  }
+
+  override def bulkPut(sourceOfflineTable: String, destinationOnlineDataSet: String, partition: String): Unit = ???
+
+  override def create(dataset: String): Unit = {
+    logger.warn(s"Skipping creation of $dataset in HTTP kv store implementation")
+  }
+
+  // wire up json conversion manually to side step serialization issues in spark executors
+  def jsonString(request: PutRequest): String = {
+    val keyBase64 = Base64.getEncoder.encodeToString(request.keyBytes)
+    val valueBase64 = Base64.getEncoder.encodeToString(request.valueBytes)
+    s"""{ "keyBytes": "${keyBase64}", "valueBytes": "${valueBase64}", "dataset": "${request.dataset}", "tsMillis": ${request.tsMillis.orNull}}""".stripMargin
+  }
+
+  def jsonList(requests: Seq[PutRequest]): String = {
+    val requestsJson = requests.map(jsonString(_)).mkString(", ")
+
+    s"[ $requestsJson ]"
+  }
+}

online/src/main/scala/ai/chronon/online/MetadataEndPoint.scala

@@ -22,7 +22,7 @@ case class MetadataEndPoint[Conf <: TBase[_, _]: Manifest: ClassTag](
 object MetadataEndPoint {
   @transient implicit lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
-  val ConfByKeyEndPointName = "ZIPLINE_METADATA"
+  val ConfByKeyEndPointName = "CHRONON_METADATA"
   val NameByTeamEndPointName = "CHRONON_ENTITY_BY_TEAM"
 
   private def getTeamFromMetadata(metaData: MetaData): String = {

online/src/main/scala/ai/chronon/online/stats/DriftStore.scala

@@ -12,8 +12,8 @@ import ai.chronon.api.thrift.protocol.TProtocolFactory
 import ai.chronon.online.KVStore
 import ai.chronon.online.KVStore.GetRequest
 import ai.chronon.online.MetadataStore
-import ai.chronon.online.stats.DriftStore.compactDeserializer
-import ai.chronon.online.stats.DriftStore.compactSerializer
+import ai.chronon.online.stats.DriftStore.binaryDeserializer
+import ai.chronon.online.stats.DriftStore.binarySerializer
 
 import java.io.Serializable
 import scala.concurrent.Future
@@ -52,8 +52,6 @@ class DriftStore(kvStore: KVStore,
     }
   }
 
-  private val deserializer: TDeserializer = compactDeserializer
-
   private case class SummaryRequestContext(request: GetRequest, tileKey: TileKey, groupName: String)
   private case class SummaryResponseContext(summaries: Array[(TileSummary, Long)], tileKey: TileKey, groupName: String)
 
@@ -76,8 +74,8 @@ class DriftStore(kvStore: KVStore,
                    endMs: Option[Long],
                    columnPrefix: Option[String]): Future[Seq[TileSummaryInfo]] = {
 
-    val serializer: TSerializer = compactSerializer
-    val tileKeyMap = tileKeysForJoin(joinConf, columnPrefix)
+    val serializer: TSerializer = binarySerializer.get()
+    val tileKeyMap = tileKeysForJoin(joinConf, None, columnPrefix)
     val requestContextMap: Map[GetRequest, SummaryRequestContext] = tileKeyMap.flatMap {
       case (group, keys) =>
         keys.map { key =>
@@ -90,6 +88,7 @@ class DriftStore(kvStore: KVStore,
     val responseFuture = kvStore.multiGet(requestContextMap.keys.toSeq)
 
     responseFuture.map { responses =>
+      val deserializer = binaryDeserializer.get()
       // deserialize the responses and surround with context
       val responseContextTries: Seq[Try[SummaryResponseContext]] = responses.map { response =>
         val valuesTry = response.values
@@ -200,7 +199,17 @@ object DriftStore {
   class SerializableSerializer(factory: TProtocolFactory) extends TSerializer(factory) with Serializable
 
   // crazy bug in compact protocol - do not change to compact
-  def compactSerializer: SerializableSerializer = new SerializableSerializer(new TBinaryProtocol.Factory())
 
-  def compactDeserializer: TDeserializer = new TDeserializer(new TBinaryProtocol.Factory())
+  @transient
+  lazy val binarySerializer: ThreadLocal[TSerializer] = new ThreadLocal[TSerializer] {
+    override def initialValue(): TSerializer = new TSerializer(new TBinaryProtocol.Factory())
+  }
+
+  @transient
+  lazy val binaryDeserializer: ThreadLocal[TDeserializer] = new ThreadLocal[TDeserializer] {
+    override def initialValue(): TDeserializer = new TDeserializer(new TBinaryProtocol.Factory())
+  }
+
+  // todo - drop this hard-coded list in favor of a well known list or exposing as part of summaries
+  def breaks(count: Int): Seq[String] = (0 to count).map(_ * (100 / count)).map("p" + _.toString)
 }

online/src/main/scala/ai/chronon/online/stats/TileDriftCalculator.scala

@@ -81,7 +81,7 @@ object TileDriftCalculator {
     result
   }
 
-  // for each summary with ts >= startMs, use spec.lookBack to find the previous summary and calculate dirft
+  // for each summary with ts >= startMs, use spec.lookBack to find the previous summary and calculate drift
   // we do this by first creating a map of summaries by timestamp
   def toTileDrifts(summariesWithTimestamps: Array[(TileSummary, Long)],
                    metric: DriftMetric,

spark/src/main/scala/ai/chronon/spark/Driver.scala

@@ -18,6 +18,7 @@ package ai.chronon.spark
 
 import ai.chronon.api
 import ai.chronon.api.Constants
+import ai.chronon.api.Constants.MetadataDataset
 import ai.chronon.api.Extensions.GroupByOps
 import ai.chronon.api.Extensions.MetadataOps
 import ai.chronon.api.Extensions.SourceOps
@@ -565,7 +566,7 @@ object Driver {
     lazy val api: Api = impl(serializableProps)
 
     def metaDataStore =
-      new MetadataStore(impl(serializableProps).genKvStore, "ZIPLINE_METADATA", timeoutMillis = 10000)
+      new MetadataStore(impl(serializableProps).genKvStore, MetadataDataset, timeoutMillis = 10000)
 
     def impl(props: Map[String, String]): Api = {
       val urls = Array(new File(onlineJar()).toURI.toURL)

spark/src/main/scala/ai/chronon/spark/scripts/DataServer.scala

@@ -5,7 +5,6 @@ import ai.chronon.api.TileSeriesKey
 import ai.chronon.api.TileSummarySeries
 import ai.chronon.api.thrift.TBase
 import ai.chronon.online.stats.DriftStore
-import ai.chronon.online.stats.DriftStore.SerializableSerializer
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.databind.SerializationFeature
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
@@ -19,7 +18,6 @@ import io.netty.handler.codec.http._
 import io.netty.util.CharsetUtil
 
 import java.util.Base64
-import java.util.function.Supplier
 import scala.reflect.ClassTag
 
 class DataServer(driftSeries: Seq[TileDriftSeries], summarySeries: Seq[TileSummarySeries], port: Int = 8181) {
@@ -35,15 +33,10 @@ class DataServer(driftSeries: Seq[TileDriftSeries], summarySeries: Seq[TileSumma
       ctx.flush()
     }
 
-    private val serializer: ThreadLocal[SerializableSerializer] =
-      ThreadLocal.withInitial(new Supplier[SerializableSerializer] {
-        override def get(): SerializableSerializer = DriftStore.compactSerializer
-      })
-
     private def convertToBytesMap[T <: TBase[_, _]: Manifest: ClassTag](
         series: T,
         keyF: T => TileSeriesKey): Map[String, String] = {
-      val serializerInstance = serializer.get()
+      val serializerInstance = DriftStore.binarySerializer.get()
       val encoder = Base64.getEncoder
       val keyBytes = serializerInstance.serialize(keyF(series))
       val valueBytes = serializerInstance.serialize(series)

spark/src/main/scala/ai/chronon/spark/scripts/ObservabilityDemoDataLoader.scala

@@ -0,0 +1,120 @@
+package ai.chronon.spark.scripts
+
+import ai.chronon.api.ColorPrinter.ColorString
+import ai.chronon.api.Constants
+import ai.chronon.api.Extensions.MetadataOps
+import ai.chronon.online.HTTPKVStore
+import ai.chronon.online.KVStore
+import ai.chronon.spark.SparkSessionBuilder
+import ai.chronon.spark.TableUtils
+import ai.chronon.spark.stats.drift.Summarizer
+import ai.chronon.spark.stats.drift.SummaryUploader
+import ai.chronon.spark.stats.drift.scripts.PrepareData
+import ai.chronon.spark.utils.InMemoryKvStore
+import ai.chronon.spark.utils.MockApi
+import org.rogach.scallop.ScallopConf
+import org.rogach.scallop.ScallopOption
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+
+object ObservabilityDemoDataLoader {
+  @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
+
+  def time(message: String)(block: => Unit): Unit = {
+    logger.info(s"$message..".yellow)
+    val start = System.currentTimeMillis()
+    block
+    val end = System.currentTimeMillis()
+    logger.info(s"$message took ${end - start} ms".green)
+  }
+
+  class Conf(arguments: Seq[String]) extends ScallopConf(arguments) {
+    val startDs: ScallopOption[String] = opt[String](
+      name = "start-ds",
+      default = Some("2023-01-01"),
+      descr = "Start date in YYYY-MM-DD format"
+    )
+
+    val endDs: ScallopOption[String] = opt[String](
+      name = "end-ds",
+      default = Some("2023-02-30"),
+      descr = "End date in YYYY-MM-DD format"
+    )
+
+    val rowCount: ScallopOption[Int] = opt[Int](
+      name = "row-count",
+      default = Some(700000),
+      descr = "Number of rows to generate"
+    )
+
+    val namespace: ScallopOption[String] = opt[String](
+      name = "namespace",
+      default = Some("observability_demo"),
+      descr = "Namespace for the demo"
+    )
+
+    verify()
+  }
+
+  def main(args: Array[String]): Unit = {
+
+    val config = new Conf(args)
+    val startDs = config.startDs()
+    val endDs = config.endDs()
+    val rowCount = config.rowCount()
+    val namespace = config.namespace()
+
+    val spark = SparkSessionBuilder.build(namespace, local = true)
+    implicit val tableUtils: TableUtils = TableUtils(spark)
+    tableUtils.createDatabase(namespace)
+
+    // generate anomalous data (join output)
+    val prepareData = PrepareData(namespace)
+    val join = prepareData.generateAnomalousFraudJoin
+
+    time("Preparing data") {
+      val df = prepareData.generateFraudSampleData(rowCount, startDs, endDs, join.metaData.loggedTable)
+      df.show(10, truncate = false)
+    }
+
+    // mock api impl for online fetching and uploading
+    val inMemKvStoreFunc: () => KVStore = () => {
+      // cannot reuse the variable - or serialization error
+      val result = InMemoryKvStore.build(namespace, () => null)
+      result
+    }
+    val inMemoryApi = new MockApi(inMemKvStoreFunc, namespace)
+
+    time("Summarizing data") {
+      // compute summary table and packed table (for uploading)
+      Summarizer.compute(inMemoryApi, join.metaData, ds = endDs, useLogs = true)
+    }
+
+    val packedTable = join.metaData.packedSummaryTable
+
+    // create necessary tables in kvstore - we now publish to the HTTP KV store as we need this available to the Hub
+    val httpKvStoreFunc: () => KVStore = () => {
+      // cannot reuse the variable - or serialization error
+      val result = new HTTPKVStore()
+      result
+    }
+    val hubApi = new MockApi(httpKvStoreFunc, namespace)
+
+    val kvStore = hubApi.genKvStore
+    kvStore.create(Constants.MetadataDataset)
+    kvStore.create(Constants.TiledSummaryDataset)
+
+    // upload join conf
+    hubApi.buildFetcher().putJoinConf(join)
+
+    time("Uploading summaries") {
+      val uploader = new SummaryUploader(tableUtils.loadTable(packedTable), hubApi)
+      uploader.run()
+    }
+
+    println("Done uploading summaries! \uD83E\uDD73".green)
+    // clean up spark session and force jvm exit
+    spark.stop()
+    System.exit(0)
+  }
+}

spark/src/main/scala/ai/chronon/spark/stats/drift/Summarizer.scala

@@ -6,7 +6,7 @@ import ai.chronon.api._
 import ai.chronon.online.Api
 import ai.chronon.online.KVStore.GetRequest
 import ai.chronon.online.KVStore.PutRequest
-import ai.chronon.online.stats.DriftStore.compactSerializer
+import ai.chronon.online.stats.DriftStore.binarySerializer
 import ai.chronon.spark.TableUtils
 import ai.chronon.spark.stats.drift.Expressions.CardinalityExpression
 import ai.chronon.spark.stats.drift.Expressions.SummaryExpression
@@ -322,9 +322,10 @@ class SummaryPacker(confPath: String,
     val func: sql.Row => Seq[TileRow] =
       Expressions.summaryPopulatorFunc(summaryExpressions, df.schema, keyBuilder, tu.partitionColumn)
 
-    val serializer = compactSerializer
     val packedRdd: RDD[sql.Row] = df.rdd.flatMap(func).map { tileRow =>
       // pack into bytes
+      val serializer = binarySerializer.get()
+
       val partition = tileRow.partition
       val timestamp = tileRow.tileTs
       val summaries = tileRow.summaries