forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfeast_serving_server.py
More file actions
182 lines (162 loc) · 6.62 KB
/
Copy pathfeast_serving_server.py
File metadata and controls
182 lines (162 loc) · 6.62 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
from concurrent import futures
import time
import logging
import grpc
import threading
import feast.serving.ServingService_pb2_grpc as Serving
from feast.serving.ServingService_pb2 import (
GetBatchFeaturesResponse,
GetOnlineFeaturesRequest,
GetOnlineFeaturesResponse,
GetFeastServingInfoResponse,
)
import fake_kafka
from typing import Dict
import sqlite3
from feast.core.CoreService_pb2_grpc import CoreServiceStub
from feast.core.CoreService_pb2 import (
ListFeatureSetsResponse,
ListStoresRequest,
ListStoresResponse,
)
from feast.core import FeatureSet_pb2 as FeatureSetProto
import stores
from feast.types import (
FeatureRow_pb2 as FeatureRowProto,
Field_pb2 as FieldProto,
Value_pb2 as ValueProto,
)
from google.protobuf.timestamp_pb2 import Timestamp
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
class ServingServicer(Serving.ServingServiceServicer):
def __init__(self, kafka: fake_kafka = None, core_url: str = None):
if kafka and core_url:
self.__core_channel = None
self.__connect_core(core_url)
self._feature_sets = (
dict()
) # type: Dict[str, FeatureSetProto.FeatureSetSpec]
self._kafka = kafka
self._store = stores.SQLiteDatabase()
thread = threading.Thread(target=self.__consume, args=())
thread.daemon = True
thread.start()
def __connect_core(self, core_url: str):
if not core_url:
raise ValueError("Please set Feast Core URL.")
if self.__core_channel is None:
self.__core_channel = grpc.insecure_channel(core_url)
try:
grpc.channel_ready_future(self.__core_channel).result(timeout=5)
except grpc.FutureTimeoutError:
raise ConnectionError(
"connection timed out while attempting to connect to Feast Core gRPC server "
+ core_url
)
else:
self._core_service_stub = CoreServiceStub(self.__core_channel)
def __get_feature_sets_from_core(self):
# Get updated list of feature sets
feature_sets = (
self._core_service_stub.ListFeatureSets
) # type: ListFeatureSetsResponse
# Store each feature set locally
for feature_set in list(feature_sets.feature_sets):
self._feature_sets[feature_set.name] = feature_set
def __consume(self):
"""
Consume message in the background from Fake Kafka
"""
while True:
self.__get_feature_sets_from_core()
self.__register_feature_sets_with_store()
for feature_set in list(self._feature_sets.values()):
message = self._kafka.get(feature_set.source.kafka_source_config.topic)
if message is None:
break
self._store.upsert_feature_row(feature_set, message)
time.sleep(1)
def __register_feature_sets_with_store(self):
for feature_set in list(self._feature_sets.values()):
self._store.register_feature_set(feature_set)
def GetFeastServingVersion(self, request, context):
return GetFeastServingInfoResponse(version="0.3.2")
def GetOnlineFeatures(self, request: GetOnlineFeaturesRequest, context):
response = GetOnlineFeaturesResponse(
feature_data_sets=[
GetOnlineFeaturesResponse.FeatureDataSet(
name="feature_set_1",
version="1",
feature_rows=[
FeatureRowProto.FeatureRow(
feature_set="feature_set_1",
event_timestamp=Timestamp(),
fields=[
FieldProto.Field(
name="feature_1",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_2",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_3",
value=ValueProto.Value(float_val=1.2),
),
],
),
FeatureRowProto.FeatureRow(
feature_set="feature_set_1",
event_timestamp=Timestamp(),
fields=[
FieldProto.Field(
name="feature_1",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_2",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_3",
value=ValueProto.Value(float_val=1.2),
),
],
),
FeatureRowProto.FeatureRow(
feature_set="feature_set_1",
event_timestamp=Timestamp(),
fields=[
FieldProto.Field(
name="feature_1",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_2",
value=ValueProto.Value(float_val=1.2),
),
FieldProto.Field(
name="feature_3",
value=ValueProto.Value(float_val=1.2),
),
],
),
],
)
]
)
return response
def serve():
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Serving.add_ServingServiceServicer_to_server(ServingServicer(), server)
server.add_insecure_port("[::]:50052")
server.start()
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
logging.basicConfig()
serve()