tair-kvcache/kv_cache_manager/client/src/internal/stub/stub.h at 494581b6f44a453b8e8b34958a8a8b935ef8985c · alibaba/tair-kvcache · GitHub

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
#pragma once

#include "kv_cache_manager/client/include/common.h"
#include "kv_cache_manager/config/instance_info.h"

namespace kv_cache_manager {

class Stub {
public:
    using KeyType = int64_t;
    using KeyVector = std::vector<KeyType>;
    using TokenIds = int64_t;
    using TokenIdsVector = std::vector<KeyType>;
    using LocationSpecInfoMap = std::map<std::string, int64_t>;
    using LocationSpecGroups = std::map<std::string, std::vector<std::string>>;

    virtual ~Stub() = default;

    virtual ClientErrorCode AddConnection(const std::string &address, uint32_t connection_timeout) = 0;

    virtual void RemoveAllConnections() = 0;

    virtual 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) = 0;
    virtual std::pair<ClientErrorCode, InstanceInfo> GetInstanceInfo(const std::string &trace_id,
                                                                     const std::string &instance_id) = 0;

    // TODO : remove this
    virtual 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) = 0;

    virtual 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) = 0;

    virtual 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) = 0;

    virtual 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) = 0;

    virtual 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) = 0;
    virtual 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) = 0;

    virtual ClientErrorCode RemoveCache(const std::string &trace_id,
                                        const std::string &instance_id,
                                        const KeyVector &keys,
                                        const TokenIdsVector &tokens,
                                        const BlockMask &block_mask) = 0;
    virtual bool TrimCache() = 0;

    // leader 节点信息查询失败时 ClusterInfo::leader_endpoint 可能为空，调用方应据此决定是否重试。
    virtual std::pair<ClientErrorCode, ClusterInfo> GetClusterInfo(const std::string &trace_id,
                                                                   const std::string &instance_id) = 0;
};

} // namespace kv_cache_manager