Skip to content

Commit 629a9a7

Browse files
sumingZeroflesher0813
sumingZero
authored and
flesher0813
committed
[feat] change the log position of UCM metrics
1 parent adfe5ba commit 629a9a7

File tree

12 files changed

+196
-101
lines changed

12 files changed

+196
-101
lines changed

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
88

99
option(BUILD_UCM_STORE "build ucm store module." ON)
1010
option(BUILD_UCM_SPARSE "build ucm sparse module." ON)
11+
option(BUILD_UCM_METRICS "build ucm metrics module." ON)
1112
option(BUILD_UNIT_TESTS "build all unit test suits." OFF)
1213
option(BUILD_NUMA "build numactl library." OFF)
1314
option(DOWNLOAD_DEPENDENCE "download dependence by cmake." ON)

‎ucm/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,6 @@ endif()
55
if(BUILD_UCM_SPARSE)
66
add_subdirectory(sparse)
77
endif()
8+
if(BUILD_UCM_METRICS)
9+
add_subdirectory(metrics)
10+
endif()

‎ucm/integration/vllm/ucm_connector.py‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919
from vllm.v1.request import Request
2020

2121
from ucm.logger import init_logger
22+
from ucm.metrics.ucm_obser import UCMStatsLogger
2223
from ucm.metrics.ucmmonitor import UCMStatsMonitor
2324
from ucm.store.factory import UcmConnectorFactory
2425
from ucm.store.ucmstore import Task, UcmKVStoreBase
25-
from ucm.metrics.ucm_obser import UCMStatsLogger
2626
from ucm.utils import Config
2727

2828
if TYPE_CHECKING:
@@ -131,6 +131,8 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
131131

132132
connector_configs = self.launch_config.get("ucm_connectors", [])
133133
assert len(connector_configs) > 0, "no storage connector name in config."
134+
self.io_size = config["io_size"]
135+
self.num_layers = num_layers
134136

135137
name = connector_configs[0].get("ucm_connector_name")
136138
config = connector_configs[0].get("ucm_connector_config") or {}
@@ -445,17 +447,21 @@ def _broadcast(self, dst_tensor_addr: list[torch.Tensor]):
445447
tensor.copy_(rec_tensor[i])
446448

447449
def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
448-
449450
metadata = self._get_connector_metadata()
450451
assert isinstance(metadata, UCMConnectorMetadata)
451452

452453
self._init_kv_caches_from_forward_context(forward_context)
453454

454455
request_to_task: dict[str, Optional[Task]] = {}
455456
req_broadcast_addr = {}
457+
is_load = False
458+
num_loaded_block = 0
459+
load_start_time = time.perf_counter() * 1000
456460
for request_id, request in metadata.request_meta.items():
457461
if len(request.load_block_ids[0]) == 0:
458462
continue
463+
is_load = True
464+
num_loaded_block += len(request.load_block_ids[0])
459465

460466
ucm_block_ids, vllm_block_ids = request.load_block_ids
461467
if self.rank != 0 and not self.is_mla:
@@ -479,6 +485,20 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
479485
logger.error(f"request {request_id} load kv cache failed.")
480486
if self.load_only_first_rank:
481487
self._broadcast(req_broadcast_addr[request_id])
488+
load_end_time = time.perf_counter() * 1000
489+
if is_load:
490+
UCMStatsMonitor.get_instance().update_stats(
491+
"UCMStats",
492+
{
493+
"load_duration": load_end_time - load_start_time,
494+
"load_speed": num_loaded_block
495+
* self.io_size
496+
* self.num_layers
497+
/ (load_end_time - load_start_time)
498+
/ 1024
499+
/ 1024, # GB/s
500+
},
501+
)
482502

483503
def wait_for_layer_load(self, layer_name: str) -> None:
484504
pass
@@ -493,7 +513,6 @@ def save_kv_layer(
493513
pass
494514

495515
def wait_for_save(self) -> None:
496-
497516
if self.is_mla and self.rank != 0:
498517
return
499518

@@ -502,9 +521,14 @@ def wait_for_save(self) -> None:
502521

503522
request_to_task: dict[str, Task] = {}
504523
request_to_blocks: dict[str, list[str]] = {}
524+
is_dump = False
525+
dump_start_time = time.perf_counter() * 1000
526+
num_dumped_block = 0
505527
for request_id, request in metadata.request_meta.items():
506528
if len(request.dump_block_ids[0]) == 0:
507529
continue
530+
is_dump = True
531+
num_dumped_block += len(request.dump_block_ids[0])
508532

509533
ucm_block_ids, vllm_block_ids = request.dump_block_ids
510534
if self.rank != 0:
@@ -539,6 +563,20 @@ def wait_for_save(self) -> None:
539563
else:
540564
logger.error(f"request {request_id} dump kv cache failed.")
541565
self.store.commit(ucm_block_ids, False)
566+
dump_end_time = time.perf_counter() * 1000
567+
if is_dump:
568+
UCMStatsMonitor.get_instance().update_stats(
569+
"UCMStats",
570+
{
571+
"save_duration": dump_end_time - dump_start_time,
572+
"save_speed": num_dumped_block
573+
* self.io_size
574+
* self.num_layers
575+
/ (dump_end_time - dump_start_time)
576+
/ 1024
577+
/ 1024, # GB/s
578+
},
579+
)
542580

543581
def clear_connector_metadata(self) -> None:
544582
super().clear_connector_metadata()

‎ucm/metrics/CMakeLists.txt‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,19 @@
1-
cmake_minimum_required(VERSION 3.16)
1+
cmake_minimum_required(VERSION 3.18)
22
project(ucm_metrics LANGUAGES CXX)
33

44
set(CMAKE_CXX_STANDARD 17)
55
set(CMAKE_CXX_STANDARD_REQUIRED ON)
66

7-
find_package(pybind11 REQUIRED)
7+
# find_package(pybind11 REQUIRED)
88

9-
add_library(metrics STATIC
9+
pybind11_add_module(ucmmonitor
10+
cpy/metrics.py.cc
11+
cc/stats/ucm_stats.cc
12+
cc/stats_registry.cc
1013
cc/stats_monitor.cc
1114
)
1215

13-
set_property(TARGET metrics PROPERTY POSITION_INDEPENDENT_CODE ON)
14-
15-
target_include_directories(metrics PUBLIC
16-
${CMAKE_CURRENT_SOURCE_DIR}/cc
17-
)
18-
19-
pybind11_add_module(ucmmonitor cpy/metrics.py.cc)
20-
target_link_libraries(ucmmonitor PRIVATE metrics)
16+
target_include_directories(ucmmonitor PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cc)
2117

2218
set_target_properties(ucmmonitor PROPERTIES
2319
LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}

‎ucm/metrics/cc/stats/istats.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
#pragma once
1+
#ifndef UCM_METRICS_ISTATS_H
2+
#define UCM_METRICS_ISTATS_H
3+
24
#include <string>
35
#include <memory>
46
#include <unordered_map>
@@ -12,3 +14,4 @@ class IStats {
1214
virtual void reset() = 0;
1315
virtual std::unordered_map<std::string, std::vector<double>> data() = 0;
1416
};
17+
#endif

‎ucm/metrics/cc/stats/ucm_stats.cc‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
#include "ucm_stats.h"
2+
#include "../stats_registry.h"
3+
#include <iostream>
4+
5+
UCMStats::UCMStats() = default;
6+
7+
std::string UCMStats::name() const {
8+
return "UCMStats";
9+
}
10+
11+
void UCMStats::reset() {
12+
for (auto& v : data_) v.clear();
13+
}
14+
15+
void UCMStats::update(const std::unordered_map<std::string, double>& params) {
16+
for (const auto& [k, v] : params) {
17+
Key id = key_from_string(k);
18+
if (id == Key::COUNT) continue;
19+
emplace_back(id, v);
20+
}
21+
}
22+
23+
std::unordered_map<std::string, std::vector<double>> UCMStats::data() {
24+
std::unordered_map<std::string, std::vector<double>> result;
25+
result["save_duration"] = data_[static_cast<std::size_t>(Key::save_duration)];
26+
result["save_speed"] = data_[static_cast<std::size_t>(Key::save_speed)];
27+
result["load_duration"] = data_[static_cast<std::size_t>(Key::load_duration)];
28+
result["load_speed"] = data_[static_cast<std::size_t>(Key::load_speed)];
29+
result["interval_lookup_hit_rates"] = data_[static_cast<std::size_t>(Key::interval_lookup_hit_rates)];
30+
return result;
31+
}
32+
33+
Key UCMStats::key_from_string(const std::string& k) {
34+
if (k == "save_duration") return Key::save_duration;
35+
if (k == "save_speed") return Key::save_speed;
36+
if (k == "load_duration") return Key::load_duration;
37+
if (k == "load_speed") return Key::load_speed;
38+
if (k == "interval_lookup_hit_rates")return Key::interval_lookup_hit_rates;
39+
return Key::COUNT;
40+
}
41+
42+
void UCMStats::emplace_back(Key id, double value) {
43+
data_[static_cast<std::size_t>(id)].push_back(value);
44+
}
45+
46+
struct Registrar {
47+
Registrar() {
48+
StatsRegistry::registerStats("UCMStats", []()->std::unique_ptr<IStats> {
49+
return std::make_unique<UCMStats>();
50+
});
51+
}
52+
} registrar;

‎ucm/metrics/cc/stats/ucm_stats.h‎

Lines changed: 17 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,38 @@
1-
#pragma once
1+
#ifndef UCM_METRICS_UCMSTATS_H
2+
#define UCM_METRICS_UCMSTATS_H
3+
24
#include "istats.h"
35
#include <array>
46
#include <vector>
57
#include <unordered_map>
68
#include <string>
9+
#include <cstdint>
710

8-
/* key → id */
911
enum class Key : uint8_t {
1012
save_duration = 0,
1113
save_speed,
1214
load_duration,
1315
load_speed,
1416
interval_lookup_hit_rates,
15-
COUNT // Total keys num
17+
COUNT
1618
};
1719

1820
class UCMStats : public IStats {
19-
static constexpr std::size_t N = static_cast<std::size_t>(Key::COUNT);
20-
std::array<std::vector<double>, N> data_;
21-
22-
static Key key_from_string(const std::string& k) {
23-
if (k == "save_duration") return Key::save_duration;
24-
if (k == "save_speed") return Key::save_speed;
25-
if (k == "load_duration") return Key::load_duration;
26-
if (k == "load_speed") return Key::load_speed;
27-
if (k == "interval_lookup_hit_rates")return Key::interval_lookup_hit_rates;
28-
return Key::COUNT; // Invalid key
29-
}
30-
3121
public:
32-
UCMStats() = default;
33-
34-
std::string name() const override {
35-
return "UCMStats";
36-
}
22+
UCMStats();
23+
~UCMStats() = default;
3724

38-
void reset() override {
39-
for (auto& v : data_) v.clear();
40-
}
41-
42-
void update(const std::unordered_map<std::string, double>& params) override {
43-
for (const auto& [k, v] : params) {
44-
Key id = key_from_string(k);
45-
if (id == Key::COUNT) continue;
46-
emplace_back(id, v);
47-
}
48-
}
49-
50-
std::unordered_map<std::string, std::vector<double>> data() override {
51-
std::unordered_map<std::string, std::vector<double>> result;
52-
result["save_duration"] = data_[static_cast<std::size_t>(Key::save_duration)];
53-
result["save_speed"] = data_[static_cast<std::size_t>(Key::save_speed)];
54-
result["load_duration"] = data_[static_cast<std::size_t>(Key::load_duration)];
55-
result["load_speed"] = data_[static_cast<std::size_t>(Key::load_speed)];
56-
result["interval_lookup_hit_rates"] = data_[static_cast<std::size_t>(Key::interval_lookup_hit_rates)];
57-
return result;
58-
}
25+
std::string name() const override;
26+
void reset() override;
27+
void update(const std::unordered_map<std::string, double>& params) override;
28+
std::unordered_map<std::string, std::vector<double>> data() override;
5929

6030
private:
61-
void emplace_back(Key id, double value) {
62-
data_[static_cast<std::size_t>(id)].push_back(value);
63-
}
31+
static constexpr std::size_t N = static_cast<std::size_t>(Key::COUNT);
32+
std::array<std::vector<double>, N> data_;
33+
34+
static Key key_from_string(const std::string& k);
35+
void emplace_back(Key id, double value);
6436
};
6537

38+
#endif // UCM_METRICS_UCMSTATS_H

‎ucm/metrics/cc/stats_monitor.cc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,22 @@
33
#include "stats_monitor.h"
44
#include <mutex>
55
#include <vector>
6+
#include <iostream>
67

78
UCMStatsMonitor::UCMStatsMonitor() {
89
auto& registry = StatsRegistry::getInstance();
910
for (const auto& name : registry.getRegisteredStatsNames()) {
11+
std::cout<<"Creating stats instance for: " << name << std::endl;
1012
stats_map_[name] = registry.createStats(name);
1113
}
1214
}
1315

16+
void UCMStatsMonitor::createStats(const std::string& name) {
17+
std::lock_guard<std::mutex> lock(mutex_);
18+
auto& registry = StatsRegistry::getInstance();
19+
stats_map_[name] = registry.createStats(name);
20+
}
21+
1422
std::unordered_map<std::string, std::vector<double>> UCMStatsMonitor::getStats(const std::string& name) {
1523
std::lock_guard<std::mutex> lock(mutex_);
1624
return stats_map_[name]->data();

‎ucm/metrics/cc/stats_monitor.h‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
#pragma once
1+
#ifndef UCM_METRICS_MONITOR_H
2+
#define UCM_METRICS_MONITOR_H
3+
24
#include "stats/istats.h"
35
#include <unordered_map>
46
#include <memory>
@@ -15,6 +17,8 @@ class UCMStatsMonitor {
1517

1618
~UCMStatsMonitor() = default;
1719

20+
void createStats(const std::string& name);
21+
1822
std::unordered_map<std::string, std::vector<double>>
1923
getStats(const std::string& name);
2024

@@ -35,4 +39,6 @@ class UCMStatsMonitor {
3539
UCMStatsMonitor();
3640
UCMStatsMonitor(const UCMStatsMonitor&) = delete;
3741
UCMStatsMonitor& operator=(const UCMStatsMonitor&) = delete;
38-
};
42+
};
43+
44+
#endif // UCM_METRICS_MONITOR_H

‎ucm/metrics/cc/stats_registry.cc‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
#include "stats_registry.h"
2+
3+
std::mutex StatsRegistry::mutex_;
4+
std::unordered_map<std::string, Creator> StatsRegistry::registry_;
5+
6+
StatsRegistry& StatsRegistry::getInstance() {
7+
static StatsRegistry inst;
8+
return inst;
9+
}
10+
11+
void StatsRegistry::registerStats(std::string name, Creator creator) {
12+
std::lock_guard lk(mutex_);
13+
registry_[name] = creator;
14+
}
15+
16+
std::unique_ptr<IStats> StatsRegistry::createStats(const std::string& name) {
17+
std::lock_guard lk(mutex_);
18+
if (auto it = registry_.find(name); it != registry_.end())
19+
return it->second();
20+
return nullptr;
21+
}
22+
23+
std::vector<std::string> StatsRegistry::getRegisteredStatsNames() {
24+
std::lock_guard lk(mutex_);
25+
std::vector<std::string> names;
26+
names.reserve(registry_.size());
27+
for (auto& [n, _] : registry_) names.push_back(n);
28+
return names;
29+
}

0 commit comments

Comments
 (0)