1
1
import logging
2
+ import threading
2
3
from concurrent import futures
4
+ from typing import Optional
3
5
4
6
import grpc
5
7
import pandas as pd
6
8
from grpc_health .v1 import health , health_pb2_grpc
7
9
8
10
from feast .data_source import PushMode
9
- from feast .errors import PushSourceNotFoundException
11
+ from feast .errors import FeatureServiceNotFoundException , PushSourceNotFoundException
10
12
from feast .feature_store import FeatureStore
11
13
from feast .protos .feast .serving .GrpcServer_pb2 import (
12
14
PushResponse ,
16
18
GrpcFeatureServerServicer ,
17
19
add_GrpcFeatureServerServicer_to_server ,
18
20
)
21
+ from feast .protos .feast .serving .ServingService_pb2 import (
22
+ GetOnlineFeaturesRequest ,
23
+ GetOnlineFeaturesResponse ,
24
+ )
25
+
26
+ logger = logging .getLogger (__name__ )
19
27
20
28
21
29
def parse (features ):
@@ -28,10 +36,16 @@ def parse(features):
28
36
class GrpcFeatureServer (GrpcFeatureServerServicer ):
29
37
fs : FeatureStore
30
38
31
- def __init__ (self , fs : FeatureStore ):
39
+ _shuting_down : bool = False
40
+ _active_timer : Optional [threading .Timer ] = None
41
+
42
+ def __init__ (self , fs : FeatureStore , registry_ttl_sec : int = 5 ):
32
43
self .fs = fs
44
+ self .registry_ttl_sec = registry_ttl_sec
33
45
super ().__init__ ()
34
46
47
+ self ._async_refresh ()
48
+
35
49
def Push (self , request , context ):
36
50
try :
37
51
df = parse (request .features )
@@ -53,19 +67,19 @@ def Push(self, request, context):
53
67
to = to ,
54
68
)
55
69
except PushSourceNotFoundException as e :
56
- logging .exception (str (e ))
70
+ logger .exception (str (e ))
57
71
context .set_code (grpc .StatusCode .INVALID_ARGUMENT )
58
72
context .set_details (str (e ))
59
73
return PushResponse (status = False )
60
74
except Exception as e :
61
- logging .exception (str (e ))
75
+ logger .exception (str (e ))
62
76
context .set_code (grpc .StatusCode .INTERNAL )
63
77
context .set_details (str (e ))
64
78
return PushResponse (status = False )
65
79
return PushResponse (status = True )
66
80
67
81
def WriteToOnlineStore (self , request , context ):
68
- logging .warning (
82
+ logger .warning (
69
83
"write_to_online_store is deprecated. Please consider using Push instead"
70
84
)
71
85
try :
@@ -76,16 +90,55 @@ def WriteToOnlineStore(self, request, context):
76
90
allow_registry_cache = request .allow_registry_cache ,
77
91
)
78
92
except Exception as e :
79
- logging .exception (str (e ))
93
+ logger .exception (str (e ))
80
94
context .set_code (grpc .StatusCode .INTERNAL )
81
95
context .set_details (str (e ))
82
96
return PushResponse (status = False )
83
97
return WriteToOnlineStoreResponse (status = True )
84
98
99
+ def GetOnlineFeatures (self , request : GetOnlineFeaturesRequest , context ):
100
+ if request .HasField ("feature_service" ):
101
+ logger .info (f"Requesting feature service: { request .feature_service } " )
102
+ try :
103
+ features = self .fs .get_feature_service (
104
+ request .feature_service , allow_cache = True
105
+ )
106
+ except FeatureServiceNotFoundException as e :
107
+ logger .error (f"Feature service { request .feature_service } not found" )
108
+ context .set_code (grpc .StatusCode .INTERNAL )
109
+ context .set_details (str (e ))
110
+ return GetOnlineFeaturesResponse ()
111
+ else :
112
+ features = list (request .features .val )
113
+
114
+ result = self .fs ._get_online_features (
115
+ features ,
116
+ request .entities ,
117
+ request .full_feature_names ,
118
+ ).proto
119
+
120
+ return result
121
+
122
+ def _async_refresh (self ):
123
+ self .fs .refresh_registry ()
124
+ if self ._shuting_down :
125
+ return
126
+ self ._active_timer = threading .Timer (self .registry_ttl_sec , self ._async_refresh )
127
+ self ._active_timer .start ()
85
128
86
- def get_grpc_server (address : str , fs : FeatureStore , max_workers : int ):
129
+
130
+ def get_grpc_server (
131
+ address : str ,
132
+ fs : FeatureStore ,
133
+ max_workers : int ,
134
+ registry_ttl_sec : int ,
135
+ ):
136
+ logger .info (f"Initializing gRPC server on { address } " )
87
137
server = grpc .server (futures .ThreadPoolExecutor (max_workers = max_workers ))
88
- add_GrpcFeatureServerServicer_to_server (GrpcFeatureServer (fs ), server )
138
+ add_GrpcFeatureServerServicer_to_server (
139
+ GrpcFeatureServer (fs , registry_ttl_sec = registry_ttl_sec ),
140
+ server ,
141
+ )
89
142
health_servicer = health .HealthServicer (
90
143
experimental_non_blocking = True ,
91
144
experimental_thread_pool = futures .ThreadPoolExecutor (max_workers = max_workers ),
0 commit comments