Skip to content

Commit 16fb7d3

Browse files
Ly Caoachals
authored andcommitted
working version of python connector
Signed-off-by: Felix Wang <wangfelix98@gmail.com> Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 1b4c457 commit 16fb7d3

16 files changed

Lines changed: 1950 additions & 52 deletions

go/feast/connector.go

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@ import (
88
"os"
99
"os/exec"
1010
"github.com/hashicorp/go-plugin"
11+
"github.com/hashicorp/go-hclog"
12+
"path/filepath"
13+
"runtime"
14+
// "time"
1115
// "github.com/feast-dev/feast/go/protos/feast/third_party/grpc/connector"
1216
)
1317

@@ -22,7 +26,6 @@ func getOnlineStore(config *RepoConfig) (OnlineStore, error) {
2226
} else {
2327
// TODO(willem): Python connectors here
2428
KV_PLUGIN := config.OnlineStore["KV_PLUGIN"].(string)
25-
fmt.Println("Hello world")
2629
return connectorClient(KV_PLUGIN)
2730
}
2831
}
@@ -32,37 +35,54 @@ func connectorClient(KV_PLUGIN string) (OnlineStore, error) {
3235
// log.SetOutput(ioutil.Discard)
3336

3437
// We're a host. Start by launching the plugin process.
38+
_, filename, _, ok := runtime.Caller(0)
39+
if !ok {
40+
panic("couldn't find file path of the connector file")
41+
}
3542
cmd := exec.Command("sh", "-c", KV_PLUGIN )
3643
cmd.Env = os.Environ()
37-
cmd.Env = append(cmd.Env, "PYTHONPATH=test_repo/connector_python:$PYTHONPATH")
44+
connectorPythonPath := filepath.Join(filename, "..", "..", "test_repo/connector_python")
45+
cmd.Env = append(cmd.Env, fmt.Sprintf("PYTHONPATH=%s:$PYTHONPATH", connectorPythonPath))
46+
47+
logger := hclog.New(&hclog.LoggerOptions{
48+
Name: "plugin",
49+
Output: os.Stdout,
50+
Level: hclog.Debug,
51+
})
3852

3953
client := plugin.NewClient(&plugin.ClientConfig{
4054
HandshakeConfig: Handshake,
4155
Plugins: PluginMap,
4256
Cmd: cmd,
4357
AllowedProtocols: []plugin.Protocol{
4458
plugin.ProtocolGRPC},
59+
Logger: logger,
4560
})
46-
// defer client.Kill()
61+
4762
// Connect via RPC
4863
rpcClient, err := client.Client()
4964
if err != nil {
5065
return nil, err
5166
}
52-
fmt.Println("here")
53-
5467
// Request the plugin
5568
raw, err := rpcClient.Dispense("onlinestore_grpc")
5669
if err != nil {
5770
return nil, err
5871
}
59-
fmt.Println("here 2")
6072

6173
// We should have a OnlineStore now! This feels like a normal interface
6274
// implementation but is in fact over an RPC connection.
6375
if onlineStore, ok := raw.(OnlineStore); !ok {
6476
return nil, errors.New("Error creating a Connector OnlineStore")
6577
} else {
78+
grpcClient, ok := onlineStore.(*GRPCClient)
79+
if !ok {
80+
return nil, errors.New("Connector is not a *connector.GrpcClient")
81+
}
82+
grpcClient.destructor = func() {
83+
client.Kill()
84+
}
6685
return onlineStore, nil
6786
}
87+
// return raw, nil
6888
}

go/feast/featurestore.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010
timestamppb "google.golang.org/protobuf/types/known/timestamppb"
1111
"io/ioutil"
1212
"strings"
13-
"sort"
13+
// "sort"
1414
"fmt"
1515
)
1616

@@ -104,7 +104,6 @@ func (fs *FeatureStore) GetOnlineFeatures(request *serving.GetOnlineFeaturesRequ
104104
// and for each join key, create a EntityKey
105105
// and add to entity_keys
106106
entitiesInFeatureView := featureViewSpec.GetEntities()
107-
sort.Strings(entitiesInFeatureView)
108107
featuresInFeatureView := featureViewSpec.GetFeatures()
109108
// Validate that all features asked for are inside this feature view
110109
featuresInFeatureViewMap := make(map[string]bool)
@@ -197,5 +196,8 @@ func populateResultRowsFromColumnar(response *serving.GetOnlineFeaturesResponse,
197196
featureVector.EventTimestamps = append(featureVector.EventTimestamps, &featureTimeStamp)
198197
}
199198
}
200-
// fmt.Println(response.Metadata.FeatureNames.Val)
199+
}
200+
201+
func (fs *FeatureStore) DestructOnlineStore() {
202+
fs.onlineStore.Destruct()
201203
}

go/feast/grpcplugin.go

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package feast
2+
3+
import (
4+
"context"
5+
"github.com/feast-dev/feast/go/protos/feast/types"
6+
"github.com/feast-dev/feast/go/protos/feast/third_party/grpc/connector"
7+
)
8+
9+
// GRPCClient is an implementation of KV that talks over RPC.
10+
type GRPCClient struct{ client connector.OnlineStoreClient
11+
destructor func() }
12+
13+
func (m *GRPCClient) OnlineRead(entityKeys []types.EntityKey, view string, features []string) ([][]Feature, error) {
14+
entityKeysRef := make([]*types.EntityKey, len(entityKeys))
15+
for i := 0; i < len(entityKeys); i++ {
16+
entityKeysRef[i] = &entityKeys[i]
17+
}
18+
results, err := m.client.OnlineRead(context.Background(), &connector.OnlineReadRequest{
19+
EntityKeys: entityKeysRef,
20+
View: view,
21+
Features: features,
22+
})
23+
if err != nil {
24+
return nil, err
25+
}
26+
feature2D := results.GetResults()
27+
featureResults := make([][]Feature, len(feature2D))
28+
for entityIndex, featureList := range feature2D {
29+
connectorList := featureList.GetFeatureList()
30+
featureResults[entityIndex] = make([]Feature, len(connectorList))
31+
for featureIndex, feature := range connectorList {
32+
featureResults[entityIndex][featureIndex] = Feature{ reference: *feature.GetReference(),
33+
timestamp: *feature.GetTimestamp(),
34+
value: *feature.GetValue() }
35+
}
36+
}
37+
return featureResults, nil
38+
}
39+
40+
func (m *GRPCClient) Destruct() {
41+
m.destructor()
42+
}
43+
44+
// Here is the gRPC server that GRPCClient talks to.
45+
type GRPCServer struct {
46+
// This is the real implementation
47+
Impl OnlineStore
48+
connector.UnimplementedOnlineStoreServer
49+
}
50+
51+
func (m *GRPCServer) OnlineRead(
52+
ctx context.Context,
53+
req *connector.OnlineReadRequest) (*connector.OnlineReadResponse, error) {
54+
numEntityKeys := len(req.EntityKeys)
55+
entityKeys := make([]types.EntityKey, numEntityKeys)
56+
for i := 0; i < numEntityKeys; i++ {
57+
entityKeys[i] = *req.EntityKeys[i]
58+
}
59+
features, err := m.Impl.OnlineRead(entityKeys, req.View, req.Features)
60+
if err != nil {
61+
return nil, err
62+
}
63+
response := connector.OnlineReadResponse{Results: make([]*connector.ConnectorFeatureList, len(features))}
64+
65+
for entityIndex, featureList := range features {
66+
response.Results[entityIndex] = &connector.ConnectorFeatureList{FeatureList: make([]*connector.ConnectorFeature, len(featureList))}
67+
for featureIndex, feature := range featureList {
68+
reference := feature.reference
69+
value := feature.value
70+
timestamp := feature.timestamp
71+
response.Results[entityIndex].FeatureList[featureIndex] = &connector.ConnectorFeature{ Reference: &reference,
72+
Value: &value,
73+
Timestamp: &timestamp}
74+
}
75+
}
76+
return &response, nil
77+
}

go/feast/onlinestore.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ type OnlineStore interface {
2222
// The inner array will have the same size as featureReferences,
2323
// while the outer array will have the same size as entityKeys.
2424
OnlineRead(entityKeys []types.EntityKey, view string, features []string) ([][]Feature, error)
25+
// Destruct must be call once user is done using OnlineStore
26+
Destruct()
2527
}
2628

2729
func getOnlineStoreType(onlineStoreConfig map[string]interface{}) (string, bool) {

go/feast/plugin.go

Lines changed: 6 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,8 @@ package feast
33
import (
44
"context"
55
"github.com/hashicorp/go-plugin"
6-
"github.com/feast-dev/feast/go/protos/feast/types"
76
"github.com/feast-dev/feast/go/protos/feast/third_party/grpc/connector"
87
"google.golang.org/grpc"
9-
// "net/rpc"
108
)
119

1210
// Handshake is a common handshake that is shared by plugin and host.
@@ -22,6 +20,8 @@ var PluginMap = map[string]plugin.Plugin{
2220
"onlinestore_grpc": &OnlineStoreGRPCPlugin{},
2321
}
2422

23+
24+
2525
// // // This is the implementation of plugin.GRPCPlugin so we can serve/consume this.
2626
type OnlineStoreGRPCPlugin struct {
2727
// GRPCPlugin must still implement the Plugin interface
@@ -31,36 +31,11 @@ type OnlineStoreGRPCPlugin struct {
3131
Impl OnlineStore
3232
}
3333

34-
// GRPCClient is an implementation of KV that talks over RPC.
35-
type GRPCClient struct{ client connector.OnlineStoreClient }
36-
37-
func (m *GRPCClient) OnlineRead(entityKeys []types.EntityKey, view string, features []string) ([][]Feature, error) {
38-
entityKeysRef := make([]*types.EntityKey, len(entityKeys))
39-
for i := 0; i < len(entityKeys); i++ {
40-
entityKeysRef[i] = &entityKeys[i]
41-
}
42-
results, err := m.client.OnlineRead(context.Background(), &connector.OnlineReadRequest{
43-
EntityKeys: entityKeysRef,
44-
View: view,
45-
Features: features,
46-
})
47-
if err != nil {
48-
return nil, err
49-
}
50-
feature2D := results.GetResults()
51-
featureResults := make([][]Feature, len(feature2D))
52-
for entityIndex, featureList := range feature2D {
53-
connectorList := featureList.GetFeatureList()
54-
featureResults[entityIndex] = make([]Feature, len(connectorList))
55-
for featureIndex, feature := range connectorList {
56-
featureResults[entityIndex][featureIndex] = Feature{ reference: *feature.GetReference(),
57-
timestamp: *feature.GetTimestamp(),
58-
value: *feature.GetValue() }
59-
}
60-
}
61-
return featureResults, nil
34+
func (p *OnlineStoreGRPCPlugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error {
35+
connector.RegisterOnlineStoreServer(s, &GRPCServer{Impl: p.Impl})
36+
return nil
6237
}
6338

6439
func (p *OnlineStoreGRPCPlugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
65-
return &GRPCClient{client: connector.NewOnlineStoreClient(c)}, nil
40+
return &GRPCClient{ client: connector.NewOnlineStoreClient(c)}, nil
6641
}

go/feast/redisonlinestore.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,11 @@ func (r *RedisOnlineStore) OnlineRead(entityKeys []types.EntityKey, view string,
189189
return results, nil
190190
}
191191

192+
// Dummy destruct function to conform with plugin OnlineStore interface
193+
func (r *RedisOnlineStore) Destruct() {
194+
195+
}
196+
192197
func BuildRedisKey(project string, entityKey types.EntityKey) (*[]byte, error) {
193198
serKey, err := SerializeEntityKey(entityKey)
194199
if err != nil {

go/feast/rpcplugin.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package feast
2+
3+
// import (
4+
// "net/rpc"
5+
// )
6+
7+
// // RPCClient is an implementation of KV that talks over RPC.
8+
// type RPCClient struct{ client *rpc.Client }
9+
10+
// func (m *RPCClient) Put(key string, value []byte) error {
11+
// // We don't expect a response, so we can just use interface{}
12+
// var resp interface{}
13+
14+
// // The args are just going to be a map. A struct could be better.
15+
// return m.client.Call("Plugin.Put", map[string]interface{}{
16+
// "key": key,
17+
// "value": value,
18+
// }, &resp)
19+
// }
20+
21+
// func (m *RPCClient) Get(key string) ([]byte, error) {
22+
// var resp []byte
23+
// err := m.client.Call("Plugin.Get", key, &resp)
24+
// return resp, err
25+
// }
26+
27+
// // Here is the RPC server that RPCClient talks to, conforming to
28+
// // the requirements of net/rpc
29+
// type RPCServer struct {
30+
// // This is the real implementation
31+
// Impl KV
32+
// }
33+
34+
// func (m *RPCServer) Put(args map[string]interface{}, resp *interface{}) error {
35+
// return m.Impl.Put(args["key"].(string), args["value"].([]byte))
36+
// }
37+
38+
// func (m *RPCServer) Get(key string, resp *[]byte) error {
39+
// v, err := m.Impl.Get(key)
40+
// *resp = v
41+
// return err
42+
// }

go/server/main.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ func main() {
4343
if err != nil {
4444
log.Fatalln(err)
4545
}
46+
defer fs.DestructOnlineStore()
4647

4748
grpcPort, ok := os.LookupEnv(flagFeastGrpcPort)
4849
if !ok {

go/test_repo/connector_python/Connector_pb2.py

Lines changed: 3 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

go/test_repo/connector_python/Connector_pb2_grpc.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
"""Client and server classes corresponding to protobuf-defined services."""
33
import grpc
44

5-
import Connector_pb2 as feast_dot_third__party_dot_grpc_dot_connector_dot_Connector__pb2
5+
from connector_python import Connector_pb2 as feast_dot_third__party_dot_grpc_dot_connector_dot_Connector__pb2
66

77

88
class OnlineStoreStub(object):

0 commit comments

Comments
 (0)