Skip to content

Commit 79f908c

Browse files
justinlucopybara-github
justinlu
authored and
copybara-github
committed
Deprecate legacy raw-offset TransferBuffers overloads in RaidenController.
PiperOrigin-RevId: 951669979
1 parent bc34c1b commit 79f908c

8 files changed

Lines changed: 108 additions & 216 deletions

File tree

‎tpu_raiden/core/controller/raiden_controller.cc‎

Lines changed: 53 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -151,8 +151,51 @@ void RaidenController::Init(absl::Span<const std::string> worker_addresses,
151151
absl::Span<const int64_t> dst_offsets,
152152
absl::Span<const int64_t> copy_sizes,
153153
absl::Span<const std::string> peers) {
154-
return this->TransferBuffers(src_mem_type, dst_mem_type, src_offsets,
155-
dst_offsets, copy_sizes, peers);
154+
if (src_offsets.empty() || src_offsets.size() != dst_offsets.size()) {
155+
return tsl::Future<>(
156+
absl::InvalidArgumentError("Source and destination offsets must "
157+
"have the same non-zero length"));
158+
}
159+
if (!copy_sizes.empty() && copy_sizes.size() != src_offsets.size()) {
160+
return tsl::Future<>(absl::InvalidArgumentError(
161+
"copy_sizes, if provided, must match the length of src_offsets"));
162+
}
163+
auto workers = worker_registry_->GetRegisteredWorkers();
164+
if (workers.empty()) {
165+
return tsl::Future<>(absl::FailedPreconditionError(
166+
"No registered workers available for TransferBuffers"));
167+
}
168+
std::sort(workers.begin(), workers.end(),
169+
[](const core::controller::WorkerRegistration& a,
170+
const core::controller::WorkerRegistration& b) {
171+
return CompareWorkerIds(a.worker_id, b.worker_id);
172+
});
173+
if (!peers.empty() && workers.size() != peers.size()) {
174+
return tsl::Future<>(absl::InvalidArgumentError(
175+
absl::StrCat("Peers count mismatch: workers has ", workers.size(),
176+
", peers has ", peers.size())));
177+
}
178+
std::vector<tsl::Future<>> worker_futures;
179+
worker_futures.reserve(workers.size());
180+
for (size_t i = 0; i < workers.size(); ++i) {
181+
std::optional<std::string> peer =
182+
peers.empty() ? std::nullopt : std::make_optional(peers[i]);
183+
std::vector<Buffer> src_buffers;
184+
src_buffers.reserve(src_offsets.size());
185+
for (int64_t offset : src_offsets) {
186+
src_buffers.emplace_back(offset, std::vector<BufferShard>{},
187+
std::nullopt, src_mem_type);
188+
}
189+
std::vector<Buffer> dst_buffers;
190+
dst_buffers.reserve(dst_offsets.size());
191+
for (int64_t offset : dst_offsets) {
192+
dst_buffers.emplace_back(offset, std::vector<BufferShard>{}, peer,
193+
dst_mem_type);
194+
}
195+
worker_futures.push_back(this->TransferBuffers(
196+
workers[i].worker_id, src_buffers, dst_buffers, copy_sizes));
197+
}
198+
return tsl::JoinFutures(absl::MakeSpan(worker_futures));
156199
});
157200

158201
// 3. Register static workers
@@ -172,8 +215,10 @@ void RaidenController::Init(absl::Span<const std::string> worker_addresses,
172215
orchestrator_client_ = std::make_unique<OrchestratorServiceClient>(
173216
grpc::CreateChannel(std::string(raiden_orchestrator_address),
174217
grpc::InsecureChannelCredentials()));
175-
std::string my_endpoint = absl::StrCat("localhost:", raiden_controller_port_);
176-
absl::Status status = orchestrator_client_->RegisterController(unit_, my_endpoint);
218+
std::string my_endpoint =
219+
absl::StrCat("localhost:", raiden_controller_port_);
220+
absl::Status status =
221+
orchestrator_client_->RegisterController(unit_, my_endpoint);
177222
if (!status.ok()) {
178223
throw std::runtime_error(absl::StrCat(
179224
"Failed to register with orchestrator: ", status.message()));
@@ -190,7 +235,8 @@ RaidenController::RaidenController(
190235
shard_size_bytes_(shard_size_bytes),
191236
num_total_blocks_(num_blocks),
192237
worker_registry_(std::make_shared<core::controller::WorkerRegistry>()),
193-
block_manager_(std::make_unique<kv_cache::LogicalBlockManager>(num_blocks)),
238+
block_manager_(
239+
std::make_unique<kv_cache::LogicalBlockManager>(num_blocks)),
194240
raiden_controller_port_(raiden_controller_port) {
195241
Init(/*worker_addresses=*/{}, raiden_orchestrator_address);
196242
}
@@ -205,7 +251,8 @@ RaidenController::RaidenController(
205251
shard_size_bytes_(shard_size_bytes),
206252
num_total_blocks_(num_blocks),
207253
worker_registry_(std::make_shared<core::controller::WorkerRegistry>()),
208-
block_manager_(std::make_unique<kv_cache::LogicalBlockManager>(num_blocks)),
254+
block_manager_(
255+
std::make_unique<kv_cache::LogicalBlockManager>(num_blocks)),
209256
raiden_controller_port_(raiden_controller_port) {
210257
Init(worker_addresses, raiden_orchestrator_address);
211258
}
@@ -372,22 +419,11 @@ RaidenController::BuildTransferBuffersRequest(
372419

373420
rpc::MemoryType src_mem_type = src_buffers[0].memory_type();
374421
rpc::MemoryType dst_mem_type = dst_buffers[0].memory_type();
375-
std::string peer;
376-
if (dst_buffers[0].remote_address().has_value() &&
377-
!dst_buffers[0].remote_address()->empty()) {
378-
peer = *dst_buffers[0].remote_address();
379-
} else if (src_buffers[0].remote_address().has_value() &&
380-
!src_buffers[0].remote_address()->empty()) {
381-
peer = *src_buffers[0].remote_address();
382-
}
383422

384423
proto::TransferBuffersRequest request;
385424
auto* transfer = request.mutable_transfer();
386425
transfer->set_src_mem_type(src_mem_type);
387426
transfer->set_dst_mem_type(dst_mem_type);
388-
if (!peer.empty()) {
389-
transfer->set_peer(peer);
390-
}
391427

392428
for (const auto& buf : src_buffers) {
393429
if (buf.index() < 0) {
@@ -414,41 +450,6 @@ RaidenController::BuildTransferBuffersRequest(
414450
return request;
415451
}
416452

417-
absl::StatusOr<proto::TransferBuffersRequest>
418-
RaidenController::BuildRawTransferBuffersRequest(
419-
rpc::MemoryType src_mem_type, rpc::MemoryType dst_mem_type,
420-
absl::Span<const int64_t> src_offsets,
421-
absl::Span<const int64_t> dst_offsets, absl::Span<const int64_t> copy_sizes,
422-
absl::string_view peer) {
423-
if (src_offsets.empty() || src_offsets.size() != dst_offsets.size()) {
424-
return absl::InvalidArgumentError(
425-
"Source and destination offsets must have the same non-zero length");
426-
}
427-
if (!copy_sizes.empty() && copy_sizes.size() != src_offsets.size()) {
428-
return absl::InvalidArgumentError(
429-
"copy_sizes, if provided, must match the length of src_offsets");
430-
}
431-
432-
proto::TransferBuffersRequest request;
433-
auto* transfer = request.mutable_transfer();
434-
transfer->set_src_mem_type(src_mem_type);
435-
transfer->set_dst_mem_type(dst_mem_type);
436-
if (!peer.empty()) {
437-
transfer->set_peer(std::string(peer));
438-
}
439-
440-
for (int64_t offset : src_offsets) {
441-
transfer->add_src_offsets(offset);
442-
}
443-
for (int64_t offset : dst_offsets) {
444-
transfer->add_dst_offsets(offset);
445-
}
446-
for (int64_t size : copy_sizes) {
447-
transfer->add_copy_sizes(size);
448-
}
449-
return request;
450-
}
451-
452453
tsl::Future<> RaidenController::TransferBuffers(
453454
absl::string_view worker_id, absl::Span<const Buffer> src_buffers,
454455
absl::Span<const Buffer> dst_buffers,
@@ -504,74 +505,7 @@ tsl::Future<> RaidenController::TransferBuffers(
504505
return tsl::JoinFutures(absl::MakeSpan(worker_futures));
505506
}
506507

507-
tsl::Future<> RaidenController::TransferBuffers(
508-
absl::string_view worker_id, rpc::MemoryType src_mem_type,
509-
rpc::MemoryType dst_mem_type, absl::Span<const int64_t> src_offsets,
510-
absl::Span<const int64_t> dst_offsets, absl::Span<const int64_t> copy_sizes,
511-
absl::string_view peer) {
512-
auto request_or = BuildRawTransferBuffersRequest(
513-
src_mem_type, dst_mem_type, src_offsets, dst_offsets, copy_sizes, peer);
514-
if (!request_or.ok()) {
515-
return tsl::Future<>(request_or.status());
516-
}
517508

518-
auto worker_or = worker_registry_->GetWorker(worker_id);
519-
if (!worker_or.ok()) {
520-
return tsl::Future<>(absl::FailedPreconditionError(
521-
absl::StrCat("Worker ", worker_id, " is not registered. ",
522-
"Did you wait for the worker to register?")));
523-
}
524-
auto worker_client = worker_or->worker_service_client;
525-
if (!worker_client) {
526-
return tsl::Future<>(absl::FailedPreconditionError(absl::StrCat(
527-
"WorkerServiceClient for ", worker_id, " is not initialized.")));
528-
}
529-
530-
return worker_client->TransferBuffers(*request_or);
531-
}
532-
533-
tsl::Future<> RaidenController::TransferBuffers(
534-
rpc::MemoryType src_mem_type, rpc::MemoryType dst_mem_type,
535-
absl::Span<const int64_t> src_offsets,
536-
absl::Span<const int64_t> dst_offsets, absl::Span<const int64_t> copy_sizes,
537-
absl::Span<const std::string> peers) {
538-
if (src_offsets.empty() || src_offsets.size() != dst_offsets.size()) {
539-
return tsl::Future<>(absl::InvalidArgumentError(
540-
"Source and destination offsets must have the same non-zero length"));
541-
}
542-
if (!copy_sizes.empty() && copy_sizes.size() != src_offsets.size()) {
543-
return tsl::Future<>(absl::InvalidArgumentError(
544-
"copy_sizes, if provided, must match the length of src_offsets"));
545-
}
546-
auto workers = worker_registry_->GetRegisteredWorkers();
547-
if (workers.empty()) {
548-
return tsl::Future<>(absl::FailedPreconditionError(
549-
"No registered workers available for TransferBuffers"));
550-
}
551-
552-
std::sort(workers.begin(), workers.end(),
553-
[](const core::controller::WorkerRegistration& a,
554-
const core::controller::WorkerRegistration& b) {
555-
return CompareWorkerIds(a.worker_id, b.worker_id);
556-
});
557-
558-
if (!peers.empty() && workers.size() != peers.size()) {
559-
return tsl::Future<>(absl::InvalidArgumentError(
560-
absl::StrCat("Peers count mismatch: workers has ", workers.size(),
561-
", peers has ", peers.size())));
562-
}
563-
564-
std::vector<tsl::Future<>> worker_futures;
565-
worker_futures.reserve(workers.size());
566-
for (size_t i = 0; i < workers.size(); ++i) {
567-
std::string peer = peers.empty() ? "" : peers[i];
568-
worker_futures.push_back(TransferBuffers(workers[i].worker_id, src_mem_type,
569-
dst_mem_type, src_offsets,
570-
dst_offsets, copy_sizes, peer));
571-
}
572-
573-
return tsl::JoinFutures(absl::MakeSpan(worker_futures));
574-
}
575509

576510
absl::StatusOr<std::string> RaidenController::ResolvePeerController(
577511
const rpc::RaidenIdProto& peer_id) {

‎tpu_raiden/core/controller/raiden_controller.h‎

Lines changed: 3 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -127,22 +127,7 @@ class RaidenController {
127127
absl::Span<const Buffer> dst_buffers,
128128
absl::Span<const int64_t> copy_sizes = {});
129129

130-
// Legacy targeted worker transfer using raw block offsets
131-
tsl::Future<> TransferBuffers(absl::string_view worker_id,
132-
rpc::MemoryType src_mem_type,
133-
rpc::MemoryType dst_mem_type,
134-
absl::Span<const int64_t> src_offsets,
135-
absl::Span<const int64_t> dst_offsets,
136-
absl::Span<const int64_t> copy_sizes = {},
137-
absl::string_view peer = "");
138-
139-
// Legacy broadcast transfer using raw block offsets
140-
tsl::Future<> TransferBuffers(rpc::MemoryType src_mem_type,
141-
rpc::MemoryType dst_mem_type,
142-
absl::Span<const int64_t> src_offsets,
143-
absl::Span<const int64_t> dst_offsets,
144-
absl::Span<const int64_t> copy_sizes = {},
145-
absl::Span<const std::string> peers = {});
130+
146131

147132
// Initiates remote read from source controller.
148133
tsl::Future<> ReadRemote(const kv_cache::RaidenId& src_raiden_id,
@@ -176,16 +161,12 @@ class RaidenController {
176161
absl::Span<const Buffer> dst_buffers,
177162
absl::Span<const int64_t> copy_sizes);
178163

179-
absl::StatusOr<proto::TransferBuffersRequest> BuildRawTransferBuffersRequest(
180-
rpc::MemoryType src_mem_type, rpc::MemoryType dst_mem_type,
181-
absl::Span<const int64_t> src_offsets,
182-
absl::Span<const int64_t> dst_offsets,
183-
absl::Span<const int64_t> copy_sizes, absl::string_view peer);
184164

185165
void Init(absl::Span<const std::string> worker_addresses,
186166
absl::string_view raiden_orchestrator_address);
187167

188-
absl::Status InitializeWorkerBuffers(core::controller::WorkerRegistration& reg);
168+
absl::Status InitializeWorkerBuffers(
169+
core::controller::WorkerRegistration& reg);
189170
rpc::RaidenIdProto unit_;
190171

191172
int num_shards_;

‎tpu_raiden/core/controller/raiden_controller_test.cc‎

Lines changed: 8 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -681,71 +681,37 @@ TEST_F(RaidenControllerTest, TransferBuffersBufferProtoSuccess) {
681681
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(0, 1));
682682
}
683683

684-
TEST_F(RaidenControllerTest, LegacyTransferBuffersTargetedSuccess) {
684+
TEST_F(RaidenControllerTest, TransferBuffersBroadcastSuccess) {
685685
MockTransferManager mock_mgr;
686686
test_server_->service->SetTransferManager(KVManagerHolder(&mock_mgr));
687687

688688
RaidenController controller(unit_, /*num_blocks=*/5, /*num_shards=*/1,
689689
/*shard_size_bytes=*/512);
690690
RegisterAndInitWorker(controller, "worker_0", test_server_->server_address);
691691

692-
std::vector<int64_t> src_offsets = {10, 30};
693-
std::vector<int64_t> dst_offsets = {20, 40};
694-
std::vector<int64_t> copy_sizes = {1, 2};
695-
696-
auto status = controller
697-
.TransferBuffers("worker_0", rpc::MEMORY_TYPE_HBM,
698-
rpc::MEMORY_TYPE_DRAM, src_offsets,
699-
dst_offsets, copy_sizes, "")
700-
.Await();
701-
ASSERT_TRUE(status.ok());
702-
EXPECT_EQ(mock_mgr.d2h_calls, 1);
703-
EXPECT_EQ(mock_mgr.h2d_calls, 0);
704-
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(10, 30));
705-
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(20, 40));
706-
EXPECT_THAT(mock_mgr.last_copy_sizes, ElementsAre(1, 2));
707-
}
708-
709-
TEST_F(RaidenControllerTest, LegacyTransferBuffersBroadcastSuccess) {
710-
MockTransferManager mock_mgr;
711-
test_server_->service->SetTransferManager(KVManagerHolder(&mock_mgr));
712-
713-
RaidenController controller(unit_, /*num_blocks=*/5, /*num_shards=*/1,
714-
/*shard_size_bytes=*/512);
715-
RegisterAndInitWorker(controller, "worker_0", test_server_->server_address);
716-
717-
std::vector<int64_t> src_offsets = {100};
718-
std::vector<int64_t> dst_offsets = {200};
692+
Buffer src_buf(100, {}, std::nullopt, rpc::MEMORY_TYPE_DRAM);
693+
Buffer dst_buf(200, {}, std::nullopt, rpc::MEMORY_TYPE_HBM);
719694

720-
auto status =
721-
controller
722-
.TransferBuffers(rpc::MEMORY_TYPE_DRAM, rpc::MEMORY_TYPE_HBM,
723-
src_offsets, dst_offsets, {}, {})
724-
.Await();
695+
auto status = controller.TransferBuffers({src_buf}, {dst_buf}).Await();
725696
ASSERT_TRUE(status.ok());
726697
EXPECT_EQ(mock_mgr.d2h_calls, 0);
727698
EXPECT_EQ(mock_mgr.h2d_calls, 1);
728699
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
729700
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(200));
730701
}
731702

732-
TEST_F(RaidenControllerTest, LegacyTransferBuffersH2HSuccess) {
703+
TEST_F(RaidenControllerTest, TransferBuffersBroadcastH2HSuccess) {
733704
MockTransferManager mock_mgr;
734705
test_server_->service->SetTransferManager(KVManagerHolder(&mock_mgr));
735706

736707
RaidenController controller(unit_, /*num_blocks=*/5, /*num_shards=*/1,
737708
/*shard_size_bytes=*/512);
738709
RegisterAndInitWorker(controller, "worker_0", test_server_->server_address);
739710

740-
std::vector<int64_t> src_offsets = {5};
741-
std::vector<int64_t> dst_offsets = {6};
711+
Buffer src_buf(5, {}, std::nullopt, rpc::MEMORY_TYPE_DRAM);
712+
Buffer dst_buf(6, {}, "localhost:8080", rpc::MEMORY_TYPE_DRAM);
742713

743-
auto status =
744-
controller
745-
.TransferBuffers(rpc::MEMORY_TYPE_DRAM, rpc::MEMORY_TYPE_DRAM,
746-
src_offsets, dst_offsets, {},
747-
std::vector<std::string>{"localhost:8080"})
748-
.Await();
714+
auto status = controller.TransferBuffers({src_buf}, {dst_buf}).Await();
749715
ASSERT_TRUE(status.ok());
750716
EXPECT_EQ(mock_mgr.h2h_write_calls, 1);
751717
EXPECT_EQ(mock_mgr.last_peer, "localhost:8080");

‎tpu_raiden/core/controller/worker_service_impl.cc‎

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -226,12 +226,9 @@ grpc::Status WorkerServiceImpl::TransferBuffers(
226226
std::string src_peer = transfer.src_buffers(0).remote_address();
227227
future_or = transfer_manager_.D2hRead(src_peer, src_offsets, dst_offsets,
228228
copy_sizes);
229-
} else if (!transfer.peer().empty() ||
230-
(transfer.dst_buffers_size() > 0 &&
231-
!transfer.dst_buffers(0).remote_address().empty())) {
232-
std::string dst_peer = !transfer.peer().empty()
233-
? transfer.peer()
234-
: transfer.dst_buffers(0).remote_address();
229+
} else if (transfer.dst_buffers_size() > 0 &&
230+
!transfer.dst_buffers(0).remote_address().empty()) {
231+
std::string dst_peer = transfer.dst_buffers(0).remote_address();
235232
future_or = transfer_manager_.D2hWrite(dst_peer, src_offsets, dst_offsets,
236233
copy_sizes);
237234
} else {
@@ -243,12 +240,9 @@ grpc::Status WorkerServiceImpl::TransferBuffers(
243240
std::string src_peer = transfer.src_buffers(0).remote_address();
244241
future_or = transfer_manager_.H2dRead(src_peer, src_offsets, dst_offsets,
245242
copy_sizes);
246-
} else if (!transfer.peer().empty() ||
247-
(transfer.dst_buffers_size() > 0 &&
248-
!transfer.dst_buffers(0).remote_address().empty())) {
249-
std::string dst_peer = !transfer.peer().empty()
250-
? transfer.peer()
251-
: transfer.dst_buffers(0).remote_address();
243+
} else if (transfer.dst_buffers_size() > 0 &&
244+
!transfer.dst_buffers(0).remote_address().empty()) {
245+
std::string dst_peer = transfer.dst_buffers(0).remote_address();
252246
future_or = transfer_manager_.H2dWrite(dst_peer, src_offsets, dst_offsets,
253247
copy_sizes);
254248
} else {
@@ -262,9 +256,6 @@ grpc::Status WorkerServiceImpl::TransferBuffers(
262256
!transfer.dst_buffers(0).remote_address().empty()) {
263257
future_or = transfer_manager_.H2hWrite(
264258
transfer.dst_buffers(0).remote_address(), src_offsets, dst_offsets);
265-
} else if (!transfer.peer().empty()) {
266-
future_or =
267-
transfer_manager_.H2hWrite(transfer.peer(), src_offsets, dst_offsets);
268259
} else {
269260
response->set_success(false);
270261
response->set_message("Peer address must be provided for H2H transfers");

0 commit comments

Comments
 (0)