Skip to content

Redirect ReadRemote to Load in KVCacheStore. #855

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Redirect ReadRemote to Load in KVCacheStore. #855

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .bazelrc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,5 @@ build:ci --remote_local_fallback
# targets set per-target and the flags kokoro presubmit passes globally.
build --cxxopt=-std=c++20
build --host_cxxopt=-std=c++20
build --copt=-DABSL_DEFINE_UNQUALIFIED_STATUS_MACROS
build --host_copt=-DABSL_DEFINE_UNQUALIFIED_STATUS_MACROS
9 changes: 2 additions & 7 deletions tpu_sync/api/jax/kv_cache_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -521,12 +521,7 @@ def read_remote(
Returns:
True if successfully launched.
"""
raw_slices = []
for s in slices:
if isinstance(s, RaidenId):
s = RaidenBlockId(raiden_id=s)
raw_slices.append(s._impl) # pylint: disable=protected-access
return self._impl.read_remote(block_hashes, raw_slices, device_block_ids)
return self.load(block_hashes, device_block_ids, slices=slices)

def poll_remote_read_status(
self,
Expand All @@ -539,4 +534,4 @@ def poll_remote_read_status(
failed: List of block hashes whose remote read failed.
pending: List of block hashes whose remote read is still in progress.
"""
return self._impl.poll_remote_read_status()
return self.poll_load_status()
9 changes: 2 additions & 7 deletions tpu_sync/api/torch/kv_cache_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,12 +532,7 @@ def read_remote(
Returns:
True if successfully launched.
"""
raw_slices = []
for s in slices:
if isinstance(s, RaidenId):
s = RaidenBlockId(raiden_id=s)
raw_slices.append(s._impl) # pylint: disable=protected-access
return self._impl.read_remote(block_hashes, raw_slices, device_block_ids)
return self.load(block_hashes, device_block_ids, slices=slices)

def poll_remote_read_status(
self,
Expand All @@ -550,4 +545,4 @@ def poll_remote_read_status(
failed: List of block hashes whose remote read failed.
pending: List of block hashes whose remote read is still in progress.
"""
return self._impl.poll_remote_read_status()
return self.poll_load_status()
23 changes: 3 additions & 20 deletions tpu_sync/core/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,23 +38,6 @@ cc_library(
visibility = ["//visibility:public"],
)

cc_library(
name = "status_macros",
hdrs = [
"status_macros.h",
],
copts = [
"-fno-strict-aliasing",
"-fexceptions",
],
features = ["-use_header_modules"],
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
],
)

cc_library(
name = "buffer",
srcs = [
Expand Down Expand Up @@ -301,9 +284,9 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
":raw_transfer_core",
":status_macros",
":xla_raw_transfer_headers",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
],
Expand All @@ -317,9 +300,9 @@ cc_library(
deps = [
":raiden_transfer_endpoint",
":raw_transfer_core",
":status_macros",
"//tpu_sync/rpc:raiden_service_cc_proto",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
Expand Down Expand Up @@ -429,7 +412,6 @@ cc_library(
":raiden_manager_base",
":raiden_transfer_endpoint",
":raw_transfer_core",
":status_macros",
":tpu_utils",
"//tpu_sync/kv_cache:kv_cache_manager_base",
"//tpu_sync/kv_cache:pool_layout",
Expand All @@ -441,6 +423,7 @@ cc_library(
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
Expand Down
2 changes: 1 addition & 1 deletion tpu_sync/core/controller/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,6 @@ cc_library(
"//tpu_sync/common:raiden_id",
"//tpu_sync/core:buffer",
"//tpu_sync/core:raiden_transfer_endpoint",
"//tpu_sync/core:status_macros",
"//tpu_sync/kv_cache:logical_block_manager",
"//tpu_sync/proto:controller_service_cc_grpc",
"//tpu_sync/proto:controller_service_cc_proto",
Expand All @@ -102,6 +101,7 @@ cc_library(
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
Expand Down
10 changes: 5 additions & 5 deletions tpu_sync/core/controller/raiden_controller.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
Expand All @@ -54,7 +55,6 @@
#include "tpu_sync/core/controller/worker_registry.h"
#include "tpu_sync/core/controller/worker_service_client.h"
#include "tpu_sync/core/raiden_transfer_endpoint.h"
#include "tpu_sync/core/status_macros.h"
#include "tpu_sync/kv_cache/logical_block_manager.h"
#include "tpu_sync/proto/controller_service.grpc.pb.h"
#include "tpu_sync/proto/controller_service.pb.h"
Expand Down Expand Up @@ -416,8 +416,8 @@ RaidenController::Allocate(int num_blocks) {
"that Physical/BufferProto mode is unavailable when the controller "
"was built with preprovision_worker_buffers=false");
}
ASSIGN_OR_RETURN(std::vector<int> block_ids,
block_manager_->Allocate(num_blocks, /*lock=*/true));
ABSL_ASSIGN_OR_RETURN(std::vector<int> block_ids,
block_manager_->Allocate(num_blocks, /*lock=*/true));
std::vector<::tpu_sync::proto::BufferProto> result;
result.reserve(num_blocks);
for (int block_id : block_ids) {
Expand All @@ -428,8 +428,8 @@ RaidenController::Allocate(int num_blocks) {

absl::StatusOr<std::vector<Buffer>> RaidenController::AllocateBuffers(
int num_blocks) {
ASSIGN_OR_RETURN(std::vector<::tpu_sync::proto::BufferProto> protos,
Allocate(num_blocks));
ABSL_ASSIGN_OR_RETURN(std::vector<::tpu_sync::proto::BufferProto> protos,
Allocate(num_blocks));
std::vector<Buffer> buffers;
buffers.reserve(protos.size());
for (const auto& proto : protos) {
Expand Down
8 changes: 4 additions & 4 deletions tpu_sync/core/kv_cache_manager_with_transfer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
Expand All @@ -71,7 +72,6 @@
#include "tpu_sync/core/raiden_manager_base.h"
#include "tpu_sync/core/raiden_transfer_endpoint.h"
#include "tpu_sync/core/raw_transfer_core.h"
#include "tpu_sync/core/status_macros.h"
#include "tpu_sync/core/tpu_utils.h"
#include "tpu_sync/kv_cache/kv_cache_manager_base.h"
#include "tpu_sync/kv_cache/pool_layout.h"
Expand Down Expand Up @@ -2104,9 +2104,9 @@ absl::Status KVCacheManagerWithTransfer::InitializeSlotPool(int64_t num_slots) {
all_slots_.clear();
all_slots_.reserve(num_slots);
for (int64_t i = 0; i < num_slots; ++i) {
ASSIGN_OR_RETURN(std::vector<int> allocated_ids,
host_block_manager_->Allocate(max_blocks_,
/*lock=*/true));
ABSL_ASSIGN_OR_RETURN(std::vector<int> allocated_ids,
host_block_manager_->Allocate(max_blocks_,
/*lock=*/true));
if (allocated_ids.size() != max_blocks_) {
return absl::InternalError(absl::StrCat(
"Slot pool allocation returned incorrect number of blocks: ",
Expand Down
51 changes: 30 additions & 21 deletions tpu_sync/core/kv_manager_holder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,13 @@
#include <vector>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tpu_sync/core/raiden_transfer_endpoint.h"
#include "tpu_sync/core/raw_transfer_core.h"
#include "tpu_sync/core/status_macros.h"
#include "tpu_sync/rpc/raiden_service.pb.h"

namespace tpu_raiden {
Expand Down Expand Up @@ -280,83 +280,92 @@ class KVManagerHolder {
absl::StatusOr<raiden::PjRtCopyFuture> H2hRead(
absl::string_view peer, const std::vector<int64_t>& src_offsets,
const std::vector<int64_t>& dst_offsets) override {
ASSIGN_OR_RETURN(std::vector<int> src_ids, SafeCastOffsets(src_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> src_ids,
SafeCastOffsets(src_offsets));
// When the caller named its destination blocks, land the data THERE:
// plain H2hRead auto-allocates destination blocks from the manager's
// own accounting, which neither matches the ids the caller reserved
// and committed to its directory nor respects blocks the controller
// already handed out.
if constexpr (internal::has_peer_h2h_read_explicit_v<T>) {
if (!dst_offsets.empty()) {
ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
return impl_->H2hReadExplicit(std::string(peer), src_ids, dst_ids,
/*explicit_dst_ptrs=*/{});
}
}
ASSIGN_OR_RETURN(auto res, impl_->H2hRead(std::string(peer), src_ids));
ABSL_ASSIGN_OR_RETURN(auto res,
impl_->H2hRead(std::string(peer), src_ids));
return res.second;
}
absl::StatusOr<raiden::PjRtCopyFuture> H2hWrite(
absl::string_view peer, const std::vector<int64_t>& src_offsets,
const std::vector<int64_t>& dst_offsets) override {
ASSIGN_OR_RETURN(std::vector<int> src_ids, SafeCastOffsets(src_offsets));
ASSIGN_OR_RETURN(std::vector<int> dst_ids, SafeCastOffsets(dst_offsets));
ASSIGN_OR_RETURN(auto res,
impl_->H2hWrite(std::string(peer), src_ids, dst_ids));
ABSL_ASSIGN_OR_RETURN(std::vector<int> src_ids,
SafeCastOffsets(src_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
ABSL_ASSIGN_OR_RETURN(
auto res, impl_->H2hWrite(std::string(peer), src_ids, dst_ids));
return res.second;
}
absl::StatusOr<raiden::PjRtCopyFuture> H2hRead(
const std::vector<RaidenTransferEndpoint>& remote_descriptors,
const std::vector<int64_t>& src_offsets,
const std::vector<int64_t>& dst_offsets) override {
ASSIGN_OR_RETURN(std::vector<int> src_ids, SafeCastOffsets(src_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> src_ids,
SafeCastOffsets(src_offsets));
// When the caller named its destination blocks, land the data THERE.
// Plain H2hRead auto-allocates, which is fine for a fire-and-forget pull
// but wrong for a store-level read: the store already reserved landing
// blocks and commits those ids into its directory, so auto-allocated
// blocks would leave the directory pointing at the wrong memory.
if constexpr (internal::has_vector_h2h_read_explicit_v<T>) {
if (!dst_offsets.empty()) {
ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
return impl_->H2hReadExplicit(remote_descriptors, src_ids, dst_ids);
}
} else if constexpr (internal::has_peer_h2h_read_explicit_v<T>) {
// No descriptor-shaped explicit read; the peer-string one lands the
// blocks just as precisely.
if (!dst_offsets.empty() && !remote_descriptors.empty()) {
ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
return impl_->H2hReadExplicit(remote_descriptors[0].endpoint,
src_ids, dst_ids,
/*explicit_dst_ptrs=*/{});
}
}
if constexpr (internal::has_vector_h2h_read_v<T>) {
ASSIGN_OR_RETURN(auto res, impl_->H2hRead(remote_descriptors, src_ids));
ABSL_ASSIGN_OR_RETURN(auto res,
impl_->H2hRead(remote_descriptors, src_ids));
return res.second;
} else {
std::string peer =
remote_descriptors.empty() ? "" : remote_descriptors[0].endpoint;
ASSIGN_OR_RETURN(auto res, impl_->H2hRead(peer, src_ids));
ABSL_ASSIGN_OR_RETURN(auto res, impl_->H2hRead(peer, src_ids));
return res.second;
}
}
absl::StatusOr<raiden::PjRtCopyFuture> H2hWrite(
const std::vector<RaidenTransferEndpoint>& remote_descriptors,
const std::vector<int64_t>& src_offsets,
const std::vector<int64_t>& dst_offsets) override {
ASSIGN_OR_RETURN(std::vector<int> src_ids, SafeCastOffsets(src_offsets));
ASSIGN_OR_RETURN(std::vector<int> dst_ids, SafeCastOffsets(dst_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> src_ids,
SafeCastOffsets(src_offsets));
ABSL_ASSIGN_OR_RETURN(std::vector<int> dst_ids,
SafeCastOffsets(dst_offsets));
if constexpr (internal::has_vector_h2h_write_v<T>) {
ASSIGN_OR_RETURN(auto res,
impl_->H2hWrite(remote_descriptors, src_ids, dst_ids));
ABSL_ASSIGN_OR_RETURN(
auto res, impl_->H2hWrite(remote_descriptors, src_ids, dst_ids));
return res.second;
} else {
std::string peer =
remote_descriptors.empty() ? "" : remote_descriptors[0].endpoint;
ASSIGN_OR_RETURN(auto res, impl_->H2hWrite(peer, src_ids, dst_ids));
ABSL_ASSIGN_OR_RETURN(auto res,
impl_->H2hWrite(peer, src_ids, dst_ids));
return res.second;
}
}
Expand Down
1 change: 0 additions & 1 deletion tpu_sync/core/raw_transfer_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@
#include "xla/tsl/platform/logging.h"
#include "xla/tsl/platform/statusor.h"
#include "tpu_sync/core/raw_transfer_core.h"
#include "tpu_sync/core/status_macros.h"

namespace raiden {

Expand Down
Loading
Loading