-
Notifications
You must be signed in to change notification settings - Fork 54
Expand file tree
/
Copy pathgrpc_stub.h
More file actions
97 lines (79 loc) · 5.32 KB
/
Copy pathgrpc_stub.h
File metadata and controls
97 lines (79 loc) · 5.32 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
#pragma once
#include <deque>
#include <shared_mutex>
#include "kv_cache_manager/client/src/internal/stub/stub.h"
#include "kv_cache_manager/protocol/protobuf/meta_service.grpc.pb.h"
namespace kv_cache_manager {
class GrpcStub : public Stub {
public:
GrpcStub() = default;
GrpcStub(uint32_t retry_time, uint32_t call_timeout);
ClientErrorCode AddConnection(const std::string &address, uint32_t connection_timeout) override;
void RemoveAllConnections() override;
std::pair<ClientErrorCode, std::string> RegisterInstance(const std::string &trace_id,
const std::string instance_group,
const std::string &instance_id,
int32_t block_size,
const LocationSpecInfoMap &location_spec_infos,
const ModelDeployment &model_deployment,
const LocationSpecGroups &location_spec_groups,
QueryType default_query_type =
QueryType::QT_UNSPECIFIED) override;
std::pair<ClientErrorCode, InstanceInfo> GetInstanceInfo(const std::string &trace_id,
const std::string &instance_id) override;
std::pair<ClientErrorCode, Metas> GetCacheMeta(const std::string &trace_id,
const std::string &instance_id,
const KeyVector &keys,
const TokenIdsVector &tokens,
const BlockMask &block_mask,
int32_t detail_level) override;
std::pair<ClientErrorCode, CacheMetaDetails> GetCacheMetaDetail(const std::string &trace_id,
const std::string &instance_id,
const KeyVector &keys,
const TokenIdsVector &tokens,
const BlockMask &block_mask,
int32_t detail_level) override;
std::pair<ClientErrorCode, Locations>
GetCacheLocation(const std::string &trace_id,
const std::string &instance_id,
QueryType query_type,
const KeyVector &keys,
const TokenIdsVector &tokens,
const BlockMask &block_mask,
int32_t sw_size,
const std::vector<std::string> &location_spec_names) override;
std::pair<ClientErrorCode, int64_t> GetCacheLocationLen(const std::string &trace_id,
const std::string &instance_id,
QueryType query_type,
const KeyVector &keys,
const TokenIdsVector &tokens,
int32_t sw_size) override;
std::pair<ClientErrorCode, WriteLocation> StartWriteCache(const std::string &trace_id,
const std::string &instance_id,
const KeyVector &keys,
const TokenIdsVector &tokens,
const std::vector<std::string> &location_spec_group_names,
int64_t write_timeout_seconds) override;
ClientErrorCode FinishWriteCache(const std::string &trace_id,
const std::string &instance_id,
const std::string write_session_id,
const BlockMask &success_block,
const Locations &locations) override;
ClientErrorCode RemoveCache(const std::string &trace_id,
const std::string &instance_id,
const KeyVector &keys,
const TokenIdsVector &tokens,
const BlockMask &block_mask) override;
bool TrimCache() override;
std::pair<ClientErrorCode, ClusterInfo> GetClusterInfo(const std::string &trace_id,
const std::string &instance_id) override;
private:
std::shared_ptr<proto::meta::MetaService::Stub> GetStub() const;
uint32_t retry_time_ = 3;
uint32_t call_timeout_ = 100;
mutable std::shared_mutex stubs_mutex_;
// TODO : add stub proxy to manage stub
std::deque<std::shared_ptr<proto::meta::MetaService::Stub>> stubs_;
mutable size_t next_stub_ = 0;
};
} // namespace kv_cache_manager