[tpu_raiden] make PacketHeader safe for breaking changes (2/N)

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",

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"],

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)
+}

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

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 <cstdint>
+#include <cstring>
+#include <string>
+
+#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<uint16_t>(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<const char*>(&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<const char*>(&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<ChunkHeader> 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<ChunkMetadata> 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

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 <cstddef>
+#include <cstdint>
+#include <string>
+
+#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<ChunkHeader> 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<ChunkMetadata> DeserializeChunkMetadata(absl::string_view s,
+                                                       uint16_t ver);
+
+}  // namespace tpu_raiden::transport::lib
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TRANSPORT_LIB_CHUNK_SERIALIZER_H_