From 7784d8b5a946e702c3607560025ceae7662a529c Mon Sep 17 00:00:00 2001 From: Googler Date: Fri, 31 Jul 2026 19:33:37 -0700 Subject: [PATCH] [tpu_raiden] make PacketHeader safe for breaking changes (2/N) PiperOrigin-RevId: 957448015 --- MODULE.bazel | 1 + tpu_raiden/transport/lib/BUILD | 35 ++++ tpu_raiden/transport/lib/chunk.fbs | 43 +++++ tpu_raiden/transport/lib/chunk.h | 8 + tpu_raiden/transport/lib/chunk_serializer.cc | 136 ++++++++++++++++ tpu_raiden/transport/lib/chunk_serializer.h | 59 +++++++ .../transport/lib/chunk_serializer_test.cc | 151 ++++++++++++++++++ 7 files changed, 433 insertions(+) create mode 100644 tpu_raiden/transport/lib/chunk.fbs create mode 100644 tpu_raiden/transport/lib/chunk_serializer.cc create mode 100644 tpu_raiden/transport/lib/chunk_serializer.h create mode 100644 tpu_raiden/transport/lib/chunk_serializer_test.cc diff --git a/MODULE.bazel b/MODULE.bazel index dca8611c..f1810982 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -96,6 +96,7 @@ bazel_dep(name = "nanobind_bazel", version = "2.12.0") bazel_dep(name = "nanobind_abseil") bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "com_google_absl_py") bazel_dep(name = "googletest", version = "1.17.0.bcr.2", repo_name = "com_google_googletest") +bazel_dep(name = "flatbuffers", version = "25.12.19", repo_name = "com_github_google_flatbuffers") pybind11_internal_configure = use_extension( "@pybind11_bazel//:internal_configure.bzl", diff --git a/tpu_raiden/transport/lib/BUILD b/tpu_raiden/transport/lib/BUILD index 6126f5aa..4f62ce14 100644 --- a/tpu_raiden/transport/lib/BUILD +++ b/tpu_raiden/transport/lib/BUILD @@ -15,15 +15,50 @@ # BUILD file for tpu_raiden transport library package. # It is only visible to the BlockTransport package. +load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library") load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") package(default_visibility = ["//visibility:public"]) +flatbuffer_cc_library( + name = "chunk_cc_fbs", + srcs = ["chunk.fbs"], +) + cc_library( name = "chunk", hdrs = ["chunk.h"], ) +cc_library( + name = "chunk_serializer", + srcs = ["chunk_serializer.cc"], + hdrs = ["chunk_serializer.h"], + deps = [ + ":chunk", + ":chunk_cc_fbs", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + ], +) + +cc_test( + name = "chunk_serializer_test", + srcs = ["chunk_serializer_test.cc"], + deps = [ + ":chunk", + ":chunk_serializer", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "raw_buffer_transport", srcs = ["raw_buffer_transport.cc"], diff --git a/tpu_raiden/transport/lib/chunk.fbs b/tpu_raiden/transport/lib/chunk.fbs new file mode 100644 index 00000000..b3319c3b --- /dev/null +++ b/tpu_raiden/transport/lib/chunk.fbs @@ -0,0 +1,43 @@ +namespace tpu_raiden.transport.lib.flatbuf; + +// Do not change. +enum Constant:uint16 { + MAGIC = 0x4452, // 'RD' in little endian. +} + +struct ChunkHeader { + // Two fixed fields. Do not change. + magic:uint16; + ver:uint16; + + // The following fields can change. + // When they do, the `ver` field must be incremented in the new serializer. + // LINT.IfChange + op:uint8; + flags:uint8; + buffer_id:uint16; + reserved:uint16; + padding:uint16; + remote_id:uint32; + local_id:uint32; + count_or_size:uint32; + uuid:uint64; + + // Paddings to keep the header size always 64 bytes. + padding0:uint64; + padding1:uint64; + padding2:uint64; + padding3:uint64; +// LINT.ThenChange(//depot/google3/third_party/tpu_raiden/tpu_raiden/transport/lib/chunk.h) +} + +struct ChunkMetadata { + // The following fields can change. + // When they do, the `ver` field in ChunkHeader must be incremented. + // LINT.IfChange + layer_idx:uint32; + dst_shard_idx:uint32; + dst_offset_bytes:uint32; + size_bytes:uint32; +// LINT.ThenChange(//depot/google3/third_party/tpu_raiden/tpu_raiden/transport/lib/chunk.h) +} diff --git a/tpu_raiden/transport/lib/chunk.h b/tpu_raiden/transport/lib/chunk.h index 4f02e5e7..8e2c3115 100644 --- a/tpu_raiden/transport/lib/chunk.h +++ b/tpu_raiden/transport/lib/chunk.h @@ -27,6 +27,7 @@ inline constexpr uint8_t kOpBufferPushBatched = 7; // TODO(swasthi): serialization using flatbuffer. // TODO(swasthi): add version field to prevent breaking changes. struct alignas(8) ChunkHeader { + // LINT.IfChange uint8_t op; // OP code. See kOp* constants. uint8_t flags; // Holds major_order or protocol flags uint16_t buffer_id; // Multidimensional Buffer / Layer ID coordinate @@ -36,13 +37,20 @@ struct alignas(8) ChunkHeader { uint32_t local_id; // Local block ID or target shard index uint32_t count_or_size; // Number of blocks or continuous payload bytes uint64_t uuid; // Globally unique transaction routing ID + // LINT.ThenChange(//depot/google3/tpu_raiden/transport/lib/chunk.fbs) + + bool operator==(const ChunkHeader&) const = default; }; struct ChunkMetadata { + // LINT.IfChange uint32_t layer_idx; uint32_t dst_shard_idx; uint32_t dst_offset_bytes; uint32_t size_bytes; + // LINT.ThenChange(//depot/google3/tpu_raiden/transport/lib/chunk.fbs) + + bool operator==(const ChunkMetadata&) const = default; }; } // namespace tpu_raiden::transport::lib diff --git a/tpu_raiden/transport/lib/chunk_serializer.cc b/tpu_raiden/transport/lib/chunk_serializer.cc new file mode 100644 index 00000000..ad9d790e --- /dev/null +++ b/tpu_raiden/transport/lib/chunk_serializer.cc @@ -0,0 +1,136 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/transport/lib/chunk_serializer.h" + +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/transport/lib/chunk.h" +#include "tpu_raiden/transport/lib/chunk_generated.h" + +namespace tpu_raiden::transport::lib { + +namespace { + +constexpr uint16_t kMagicRaiden = + static_cast(flatbuf::Constant_MAGIC); +static_assert(kMagicRaiden == 0x4452); + +std::string SerializeHeaderV1(const ChunkHeader& header) { + constexpr uint16_t kVer = 1; + const flatbuf::ChunkHeader h( + kMagicRaiden, kVer, header.op, header.flags, header.buffer_id, + header.reserved, /*padding=*/0, header.remote_id, header.local_id, + header.count_or_size, header.uuid, /*padding0=*/0, /*padding1=*/0, + /*padding2=*/0, /*padding3=*/0); + return std::string(reinterpret_cast(&h), sizeof(h)); +} + +void DeserializeHeaderV1(const flatbuf::ChunkHeader& h, ChunkHeader& header) { + DCHECK_EQ(h.ver(), 1); + header.op = h.op(); + header.flags = h.flags(); + header.buffer_id = h.buffer_id(); + header.reserved = h.reserved(); + header.remote_id = h.remote_id(); + header.local_id = h.local_id(); + header.count_or_size = h.count_or_size(); + header.uuid = h.uuid(); +} + +std::string SerializeMetadataV1(const ChunkMetadata& meta) { + const flatbuf::ChunkMetadata m(meta.layer_idx, meta.dst_shard_idx, + meta.dst_offset_bytes, meta.size_bytes); + return std::string(reinterpret_cast(&m), sizeof(m)); +} + +void DeserializeMetadataV1(const flatbuf::ChunkMetadata& m, + ChunkMetadata& meta) { + meta.layer_idx = m.layer_idx(); + meta.dst_shard_idx = m.dst_shard_idx(); + meta.dst_offset_bytes = m.dst_offset_bytes(); + meta.size_bytes = m.size_bytes(); +} + +} // namespace + +std::string SerializeChunkHeader(const ChunkHeader& header) { + const std::string s = SerializeHeaderV1(header); + DCHECK_EQ(s.size(), kChunkHeaderSize); + return s; +} + +absl::StatusOr DeserializeChunkHeader(absl::string_view s) { + flatbuf::ChunkHeader h; + DCHECK_EQ(sizeof(h), kChunkHeaderSize); + if (s.size() != kChunkHeaderSize) { + return absl::InvalidArgumentError("Invalid chunk header size"); + } + + std::memcpy(&h, s.data(), sizeof(h)); + + if (h.magic() != kMagicRaiden) { + return absl::InvalidArgumentError( + absl::StrCat("Chunk header magic mismatch: expected ", kMagicRaiden, + ", got ", h.magic())); + } + + const uint16_t ver = h.ver(); + switch (ver) { + case 1: { + ChunkHeader header = {}; + DeserializeHeaderV1(h, header); + return header; + } + default: + return absl::FailedPreconditionError( + absl::StrCat("Unsupported chunk header flatbuf version: ", ver)); + } +} + +std::string SerializeChunkMetadata(const ChunkMetadata& meta) { + const std::string s = SerializeMetadataV1(meta); + return s; +} + +absl::StatusOr DeserializeChunkMetadata(absl::string_view s, + uint16_t ver) { + const size_t meta_size = GetChunkMetadataSize(ver); + if (s.size() != meta_size) { + return absl::InvalidArgumentError("Invalid chunk metadata size"); + } + + flatbuf::ChunkMetadata m; + std::memcpy(&m, s.data(), meta_size); + + ChunkMetadata metadata = {}; + switch (ver) { + case 1: + DeserializeMetadataV1(m, metadata); + break; + default: + return absl::FailedPreconditionError( + absl::StrCat("Unsupported chunk metadata flatbuf version: ", ver)); + } + return metadata; +} + +} // namespace tpu_raiden::transport::lib diff --git a/tpu_raiden/transport/lib/chunk_serializer.h b/tpu_raiden/transport/lib/chunk_serializer.h new file mode 100644 index 00000000..d4d056ff --- /dev/null +++ b/tpu_raiden/transport/lib/chunk_serializer.h @@ -0,0 +1,59 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TRANSPORT_LIB_CHUNK_SERIALIZER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TRANSPORT_LIB_CHUNK_SERIALIZER_H_ + +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/transport/lib/chunk.h" +#include "tpu_raiden/transport/lib/chunk_generated.h" + +namespace tpu_raiden::transport::lib { + +inline constexpr size_t kChunkHeaderSize = 64; + +static_assert(sizeof(flatbuf::ChunkHeader) == kChunkHeaderSize); + +// Returns the size of a chunk metadata for the given version. +constexpr size_t GetChunkMetadataSize(uint16_t ver) { + switch (ver) { + case 1: + return 16; + default: + return 0; + } +} + +// Serializes the chunk header to a binary string. +std::string SerializeChunkHeader(const ChunkHeader& header); + +// Parses the chunk header from its serialized binary string. +absl::StatusOr DeserializeChunkHeader(absl::string_view s); + +// Serializes the chunk metadata to a binary string. +std::string SerializeChunkMetadata(const ChunkMetadata& meta); + +// Parses the chunk metadata from its serialized binary string. +absl::StatusOr DeserializeChunkMetadata(absl::string_view s, + uint16_t ver); + +} // namespace tpu_raiden::transport::lib + +#endif // THIRD_PARTY_TPU_RAIDEN_TRANSPORT_LIB_CHUNK_SERIALIZER_H_ diff --git a/tpu_raiden/transport/lib/chunk_serializer_test.cc b/tpu_raiden/transport/lib/chunk_serializer_test.cc new file mode 100644 index 00000000..195e7b7d --- /dev/null +++ b/tpu_raiden/transport/lib/chunk_serializer_test.cc @@ -0,0 +1,151 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/transport/lib/chunk_serializer.h" + +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_matchers.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/transport/lib/chunk.h" + +namespace tpu_raiden::transport::lib { +namespace { + +using ::absl_testing::IsOkAndHolds; +using ::absl_testing::StatusIs; + +ChunkHeader MakeSampleHeaderV1() { + return ChunkHeader{ + .op = 0xAB, + .flags = 0xCD, + .buffer_id = 0x1234, + .reserved = 0x5678, + .remote_id = 0x12345678, + .local_id = 0x9ABCDEF0, + .count_or_size = 0x11223344, + .uuid = 0x0123456789ABCDEFULL, + }; +} + +ChunkMetadata MakeSampleMetadataV1() { + return ChunkMetadata{ + .layer_idx = 0x12345678, + .dst_shard_idx = 0x9ABCDEF0, + .dst_offset_bytes = 0x11223344, + .size_bytes = 0x55667788, + }; +} + +TEST(ChunkHeaderSerializerTest, SerializeAndDeserialize) { + const ChunkHeader original = MakeSampleHeaderV1(); + const std::string bytes = SerializeChunkHeader(original); + + EXPECT_THAT(DeserializeChunkHeader(bytes), IsOkAndHolds(original)); +} + +TEST(ChunkHeaderSerializerTest, SerializeToLittleEndian) { + const std::string wire = SerializeChunkHeader(MakeSampleHeaderV1()); + ASSERT_EQ(wire.size(), kChunkHeaderSize); + + alignas(8) const uint8_t expected_wire[64] = { + 0x52, 0x44, 0x01, 0x00, 0xAB, 0xCD, 0x34, 0x12, 0x78, 0x56, 0x00, + 0x00, 0x78, 0x56, 0x34, 0x12, 0xF0, 0xDE, 0xBC, 0x9A, 0x44, 0x33, + 0x22, 0x11, 0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01, + }; + + EXPECT_EQ(wire, + absl::string_view(reinterpret_cast(expected_wire), + sizeof(expected_wire))); +} + +TEST(ChunkHeaderSerializerTest, VerifyMagicBytes) { + const std::string s = SerializeChunkHeader(MakeSampleHeaderV1()); + ASSERT_GE(s.size(), 2); + ASSERT_EQ(s[0], 'R'); + ASSERT_EQ(s[1], 'D'); +} + +TEST(ChunkHeaderSerializerTest, DeserializeLittleEndian) { + alignas(8) const uint8_t raw_wire[64] = { + 0x52, 0x44, 0x01, 0x00, 0xAB, 0xCD, 0x34, 0x12, 0x78, 0x56, 0x00, + 0x00, 0x78, 0x56, 0x34, 0x12, 0xF0, 0xDE, 0xBC, 0x9A, 0x44, 0x33, + 0x22, 0x11, 0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01, + }; + + const std::string wire(reinterpret_cast(raw_wire), + sizeof(raw_wire)); + + EXPECT_THAT(DeserializeChunkHeader(wire), IsOkAndHolds(MakeSampleHeaderV1())); +} + +TEST(ChunkHeaderSerializerTest, DeserializeRejectsInvalidMagic) { + std::string bytes = SerializeChunkHeader(MakeSampleHeaderV1()); + bytes[0] ^= 0xFF; // Corrupt the magic field. + + EXPECT_THAT(DeserializeChunkHeader(bytes), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST(ChunkHeaderSerializerTest, DeserializeRejectsInvalidVersion) { + std::string bytes = SerializeChunkHeader(MakeSampleHeaderV1()); + // The `ver` field is a little-endian uint16. + bytes[2] = 0x02; + bytes[3] = 0x00; + + EXPECT_THAT(DeserializeChunkHeader(bytes), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + +TEST(ChunkMetadataSerializerTest, SerializeAndDeserialize) { + const ChunkMetadata original = MakeSampleMetadataV1(); + const std::string bytes = SerializeChunkMetadata(original); + + EXPECT_THAT(DeserializeChunkMetadata(bytes, /*ver=*/1), + IsOkAndHolds(original)); +} + +TEST(ChunkMetadataSerializerTest, SerializeToLittleEndian) { + const std::string wire = SerializeChunkMetadata(MakeSampleMetadataV1()); + ASSERT_EQ(wire.size(), GetChunkMetadataSize(1)); + + alignas(4) const uint8_t expected_wire[16] = { + 0x78, 0x56, 0x34, 0x12, 0xF0, 0xDE, 0xBC, 0x9A, + 0x44, 0x33, 0x22, 0x11, 0x88, 0x77, 0x66, 0x55, + }; + + EXPECT_EQ(wire, + absl::string_view(reinterpret_cast(expected_wire), + sizeof(expected_wire))); +} + +TEST(ChunkMetadataSerializerTest, DeserializeLittleEndian) { + alignas(4) const uint8_t raw_wire[16] = { + 0x78, 0x56, 0x34, 0x12, 0xF0, 0xDE, 0xBC, 0x9A, + 0x44, 0x33, 0x22, 0x11, 0x88, 0x77, 0x66, 0x55, + }; + + const std::string wire(reinterpret_cast(raw_wire), + sizeof(raw_wire)); + + EXPECT_THAT(DeserializeChunkMetadata(wire, /*ver=*/1), + IsOkAndHolds(MakeSampleMetadataV1())); +} + +} // namespace +} // namespace tpu_raiden::transport::lib