54 lines
2.1 KiB
Scala
54 lines
2.1 KiB
Scala
package com.twitter.timelines.prediction.common.aggregates.real_time
|
|
|
|
import com.twitter.finagle.mtls.authentication.ServiceIdentifier
|
|
import com.twitter.finagle.stats.StatsReceiver
|
|
import com.twitter.ml.featurestore.catalog.datasets.magicrecs.UserFeaturesDataset
|
|
import com.twitter.ml.featurestore.catalog.datasets.geo.GeoUserLocationDataset
|
|
import com.twitter.ml.featurestore.lib.dataset.DatasetParams
|
|
import com.twitter.ml.featurestore.lib.export.strato.FeatureStoreAppNames
|
|
import com.twitter.ml.featurestore.lib.online.FeatureStoreClient
|
|
import com.twitter.ml.featurestore.lib.params.FeatureStoreParams
|
|
import com.twitter.strato.client.{Client, Strato}
|
|
import com.twitter.strato.opcontext.Attribution.ManhattanAppId
|
|
import com.twitter.util.Duration
|
|
|
|
private[real_time] object FeatureStoreUtils {
|
|
private def mkStratoClient(serviceIdentifier: ServiceIdentifier): Client =
|
|
Strato.client
|
|
.withMutualTls(serviceIdentifier)
|
|
.withRequestTimeout(Duration.fromMilliseconds(50))
|
|
.build()
|
|
|
|
private val featureStoreParams: FeatureStoreParams =
|
|
FeatureStoreParams(
|
|
perDataset = Map(
|
|
UserFeaturesDataset.id ->
|
|
DatasetParams(
|
|
stratoSuffix = Some(FeatureStoreAppNames.Timelines),
|
|
attributions = Seq(ManhattanAppId("athena", "timelines_aggregates_v2_features_by_user"))
|
|
),
|
|
GeoUserLocationDataset.id ->
|
|
DatasetParams(
|
|
attributions = Seq(ManhattanAppId("starbuck", "timelines_geo_features_by_user"))
|
|
)
|
|
)
|
|
)
|
|
|
|
def mkFeatureStoreClient(
|
|
serviceIdentifier: ServiceIdentifier,
|
|
statsReceiver: StatsReceiver
|
|
): FeatureStoreClient = {
|
|
com.twitter.server.Init() // necessary in order to use WilyNS path
|
|
|
|
val stratoClient: Client = mkStratoClient(serviceIdentifier)
|
|
val featureStoreClient: FeatureStoreClient = FeatureStoreClient(
|
|
featureSet =
|
|
UserFeaturesAdapter.UserFeaturesSet ++ AuthorFeaturesAdapter.UserFeaturesSet ++ TweetFeaturesAdapter.TweetFeaturesSet,
|
|
client = stratoClient,
|
|
statsReceiver = statsReceiver,
|
|
featureStoreParams = featureStoreParams
|
|
)
|
|
featureStoreClient
|
|
}
|
|
}
|