Skip to content

Commit aa7c667

Browse files
justinlu13copybara-github
authored andcommitted
Fix compilation errors and test failures across tpu_sync targets
- Qualify PollSaveStatusResult and include kv_cache_store.h in kv_cache_store_mock.h/cc, adding missing kv_cache dependency. - Rename mock registration macro in kv_cache_store_backend_mock.h to REGISTER_MOCK_KV_CACHE_STORE_BACKEND to prevent macro collision with the production kv_cache_store_backend_factory.h macro. - Add torch_tpu_device_buffer_headers dependency to weight_synchronizer_torch_mock to fix layering check failure. - Update PollSaveStatus() structured bindings in storage_offload_test.cc to unpack all 5 fields of PollSaveStatusResult. - Add default metric instance registration in RegisterKnownFamilies() in prometheus_exporter.cc to ensure text snapshot serializes registered metric families on initialization. - Configure pathways_qwen_weight_synchronization_perf_test with reduced layers/iterations and appropriate perf/hardware test tags (notap, cpu:8, requires-accel-hbm, requires-net:external). Tested: - blaze test --nocheck_visibility //third_party/tpu_raiden/tpu_sync/experimental/feat_storage_extension:kv_cache_store_mock - blaze test --nocheck_visibility //third_party/tpu_raiden/tpu_sync/experimental/feat_storage_extension:storage_offload_test - blaze test --nocheck_visibility //third_party/tpu_raiden/tpu_sync/frameworks/torch:weight_synchronizer_torch_mock - blaze test --nocheck_visibility //third_party/tpu_raiden/tpu_sync/telemetry/python:telemetry_binding_test - blaze test --nocheck_visibility //third_party/tpu_raiden/tpu_sync/telemetry:prometheus_exporter_test - Full test pass across //third_party/tpu_raiden/tpu_sync/... PiperOrigin-RevId: 969954875
1 parent d07d955 commit aa7c667

2 files changed

Lines changed: 4 additions & 0 deletions

File tree

tpu_sync/frameworks/torch/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -429,6 +429,7 @@ cc_library(
429429
tags = ["nobuilder"],
430430
visibility = ["//visibility:public"],
431431
deps = [
432+
":torch_tpu_device_buffer_headers",
432433
":torch_tpu_tensor_to_buffer_headers",
433434
":torch_utils_mock",
434435
"//tpu_sync/weight_sync:weight_synchronizer_base",

tpu_sync/telemetry/prometheus_exporter.cc

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ void PrometheusExporter::RegisterKnownFamilies() {
7575
.Name(prometheus_name)
7676
.Help(std::string(meta.description))
7777
.Register(*registry_);
78+
family->Add({});
7879
counter_families_.emplace(meta.name, family);
7980
break;
8081
}
@@ -83,6 +84,7 @@ void PrometheusExporter::RegisterKnownFamilies() {
8384
.Name(prometheus_name)
8485
.Help(std::string(meta.description))
8586
.Register(*registry_);
87+
family->Add({});
8688
gauge_families_.emplace(meta.name, family);
8789
break;
8890
}
@@ -91,6 +93,7 @@ void PrometheusExporter::RegisterKnownFamilies() {
9193
.Name(prometheus_name)
9294
.Help(std::string(meta.description))
9395
.Register(*registry_);
96+
family->Add({}, default_buckets_);
9497
histogram_families_.emplace(meta.name, family);
9598
break;
9699
}

0 commit comments

Comments
 (0)