[IFRT] Use SerDes versioning in most IFRT and PjRt-IFRT types

The current serialization format is anchored as `SerDesVersionNumber(0)`.
Existing serialization protos created by the previous version of IFRT will be
compatible because the version field will be read as 0 (proto default value),
which will be recognized as `SerDesVersionNumber(0)` and use the identical
deserialization ligic as before.

All `ToProto()`s get a `version_number` parameter (defaults to
`SerDesVersion::current().version_number().value()`), and use it for internal
format selection and forward it to nested serialization.

All `Serialize()`s recognize `options` and get a version if provided (if
omitted, `SerDesVersion::current()` is assumed), and behave similarly to
`ToProto()`.

`Deserialize()`/`FromProto()` will explicitly reject the serialized data if its
version is no longer supported, if the updated serialization format tracks the
serailzation version.

Some places that have limited versioning support:

`HloProgram` and `IfrtIRProgram` currently pin its internal serialization
including VHLO and IFRT IR to a at-least-4-week-old version. This may change to
accept a broader range of versions in case it is beneficial to be able to use
the latest SerDes version instead of 4-week-old version within a
single-versioned system.

`PluginProgram` ignores versioning. Any change to its format currently is a
breaking change because its serialization format is opaque/not extensible. Its
use will also be replaced with `CustomCallProgram`, so it seems not very
beneficial to add versioning to it.

IFRT Proxy does not yet use SerDes versioning. It will be done in a subsequent
change.

PiperOrigin-RevId: 771280809

Author: hyeontaek

xla/python/ifrt/BUILD

@@ -85,6 +85,7 @@ cc_library(
         ":ref_wrapper",
         ":remap_plan_proto_cc",
         ":serdes",
+        ":serdes_version",
         ":shape_proto_cc",
         ":sharding_proto_cc",
         ":user_context",
@@ -152,7 +153,9 @@ cc_library(
     ]),
     deps = [
         ":attribute_map_proto_cc",
+        ":serdes_version",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -165,6 +168,8 @@ xla_cc_test(
     srcs = ["attribute_map_test.cc"],
     deps = [
         ":attribute_map",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/platform:statusor",
         "@com_google_googletest//:gtest_main",
     ],
@@ -211,6 +216,8 @@ xla_cc_test(
         ":attribute_map",
         ":execute_options_proto_cc",
         ":ifrt",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
         "@com_google_googletest//:gtest_main",
@@ -296,8 +303,10 @@ cc_library(
         ":ifrt",
         ":layout_serdes_proto_cc",
         ":serdes",
+        ":serdes_version",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
     ],
@@ -312,6 +321,8 @@ xla_cc_test(
         ":ifrt",
         ":layout_serdes",
         ":serdes",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/platform:statusor",
         "@com_google_googletest//:gtest_main",
         "@llvm-project//llvm:Support",
@@ -363,6 +374,8 @@ xla_cc_test(
     srcs = ["shape_test.cc"],
     deps = [
         ":ifrt",
+        ":serdes_test_util",
+        ":serdes_version",
         ":shape_proto_cc",
         "//xla/tsl/platform:status_matchers",
         "//xla/tsl/platform:statusor",
@@ -611,6 +624,7 @@ cc_library(
     deps = [
         ":serdes_any_version_accessor",
         ":serdes_version",
+        ":serdes_week_4_old_version_accessor",
     ],
 )
 
@@ -697,11 +711,13 @@ cc_library(
     deps = [
         ":ifrt",
         ":serdes",
+        ":serdes_version",
         ":sharding_serdes_proto_cc",
         "//xla/python/ifrt/ir:sharding_param",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
     ],
@@ -716,6 +732,8 @@ xla_cc_test(
         ":ifrt",
         ":serdes",
         ":serdes_proto_cc",
+        ":serdes_test_util",
+        ":serdes_version",
         ":sharding_serdes",
         "//xla/python/ifrt/ir:sharding_param",
         "//xla/tsl/platform:statusor",
@@ -747,6 +765,8 @@ xla_cc_test(
         ":array_spec_proto_cc",
         ":device_test_util",
         ":ifrt",
+        ":serdes_test_util",
+        ":serdes_version",
         ":sharding_serdes",
         "//xla:shape_util",
         "//xla/pjrt:pjrt_layout",
@@ -776,6 +796,8 @@ xla_cc_test(
         ":device_proto_cc",
         ":device_test_util",
         ":ifrt",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/types:span",
@@ -824,6 +846,8 @@ xla_cc_test(
     deps = [
         ":dtype_proto_cc",
         ":ifrt",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
         "@com_google_googletest//:gtest_main",
@@ -849,6 +873,8 @@ xla_cc_test(
         ":device_test_util",
         ":ifrt",
         ":remap_plan_proto_cc",
+        ":serdes_test_util",
+        ":serdes_version",
         ":sharding_serdes",
         "//xla:shape_util",
         "//xla/pjrt:pjrt_layout",
@@ -970,6 +996,7 @@ xla_cc_test(
         ":plugin_program_serdes",
         ":serdes",
         ":serdes_proto_cc",
+        ":serdes_version",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/protobuf:error_codes_proto_impl_cc",
@@ -1026,12 +1053,13 @@ cc_library(
         ":ifrt",
         ":program_serdes",
         ":serdes",
+        ":serdes_version",
         ":sharding_proto_cc",
         ":sharding_serdes",
-        "//xla/tsl/concurrency:ref_count",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
@@ -1065,6 +1093,8 @@ xla_cc_test(
         ":program_serdes",
         ":serdes",
         ":serdes_proto_cc",
+        ":serdes_test_util",
+        ":serdes_version",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:status_matchers",
         "//xla/tsl/platform:statusor",

xla/python/ifrt/array_spec.cc

@@ -19,13 +19,14 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array_spec.pb.h"
 #include "xla/python/ifrt/client.h"
-#include "xla/python/ifrt/device_list.h"
 #include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/serdes_version.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
 #include "xla/tsl/platform/statusor.h"
@@ -35,6 +36,12 @@ namespace ifrt {
 
 absl::StatusOr<ArraySpec> ArraySpec::FromProto(Client* client,
                                                const ArraySpecProto& proto) {
+  const SerDesVersionNumber version_number(proto.version_number());
+  if (version_number != SerDesVersionNumber(0)) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "Unsupported ", version_number, " for ArraySpec deserialization"));
+  }
+
   TF_ASSIGN_OR_RETURN(auto dtype, DType::FromProto(proto.dtype()));
   TF_ASSIGN_OR_RETURN(auto shape, Shape::FromProto(proto.shape()));
   TF_ASSIGN_OR_RETURN(auto sharding,
@@ -51,11 +58,18 @@ absl::StatusOr<ArraySpec> ArraySpec::FromProto(Client* client,
   };
 }
 
-absl::StatusOr<ArraySpecProto> ArraySpec::ToProto() const {
+absl::StatusOr<ArraySpecProto> ArraySpec::ToProto(SerDesVersion version) const {
+  if (version.version_number() < SerDesVersionNumber(0)) {
+    return absl::FailedPreconditionError(
+        absl::StrCat("Unsupported ", version.version_number(),
+                     " for ArraySpec serialization"));
+  }
+
   ArraySpecProto proto;
-  *proto.mutable_dtype() = dtype.ToProto();
-  *proto.mutable_shape() = shape.ToProto();
-  TF_ASSIGN_OR_RETURN(*proto.mutable_sharding(), sharding->ToProto());
+  proto.set_version_number(SerDesVersionNumber(0).value());
+  *proto.mutable_dtype() = dtype.ToProto(version);
+  *proto.mutable_shape() = shape.ToProto(version);
+  TF_ASSIGN_OR_RETURN(*proto.mutable_sharding(), sharding->ToProto(version));
   if (layout != nullptr) {
     proto.set_layout(layout->Serialize());
   }

xla/python/ifrt/array_spec.h

@@ -25,6 +25,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array_spec.pb.h"
 #include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/serdes_version.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
 
@@ -77,7 +78,8 @@ struct ArraySpec {
                                              const ArraySpecProto& proto);
 
   // Returns a `ArraySpecProto` representation.
-  absl::StatusOr<ArraySpecProto> ToProto() const;
+  absl::StatusOr<ArraySpecProto> ToProto(
+      SerDesVersion version = SerDesVersion::current()) const;
 
   // TODO(hyeontaek): Remove this method in favor of AbslStringify.
   std::string DebugString() const;

xla/python/ifrt/array_spec.proto

@@ -23,6 +23,8 @@ import "xla/python/ifrt/sharding.proto";
 
 // Proto equivalent of C++ `ArraySpec`.
 message ArraySpecProto {
+  int32 version_number = 5;
+
   DTypeProto dtype = 1;
   ShapeProto shape = 2;
   ShardingProto sharding = 3;

xla/python/ifrt/array_spec_test.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/python/ifrt/array_spec.h"
 
 #include <memory>
+#include <tuple>
 
 #include <gtest/gtest.h>
 #include "absl/hash/hash_testing.h"
@@ -28,6 +29,8 @@ limitations under the License.
 #include "xla/python/ifrt/device_test_util.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/memory.h"
+#include "xla/python/ifrt/serdes_test_util.h"
+#include "xla/python/ifrt/serdes_version.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
 #include "xla/tsl/platform/statusor.h"
@@ -36,18 +39,23 @@ namespace xla {
 namespace ifrt {
 namespace {
 
-using ArraySpecTestParam = test_util::DeviceTestParam;
+using ArraySpecTestParam =
+    std::tuple<SerDesVersion, test_util::DeviceTestParam>;
 
 class ArraySpecTest : public testing::TestWithParam<ArraySpecTestParam> {
  public:
-  ArraySpecTest() : fixture_(GetParam()) {}
+  ArraySpecTest()
+      : version_(std::get<0>(GetParam())), fixture_(std::get<1>(GetParam())) {}
+
+  SerDesVersion version() const { return version_; }
 
   Client* client() { return fixture_.client(); }
   DeviceListRef GetDevices(absl::Span<const int> device_indices) {
     return fixture_.GetDevices(device_indices);
   }
 
  private:
+  SerDesVersion version_;
   test_util::DeviceTestFixture fixture_;
 };
 
@@ -77,7 +85,7 @@ TEST_P(ArraySpecTest, ToFromProto) {
                                               /*shape=*/shape,
                                               /*shard_shape=*/shard_shape)};
 
-  TF_ASSERT_OK_AND_ASSIGN(const ArraySpecProto proto, spec.ToProto());
+  TF_ASSERT_OK_AND_ASSIGN(const ArraySpecProto proto, spec.ToProto(version()));
   TF_ASSERT_OK_AND_ASSIGN(const ArraySpec array_spec_copy,
                           ArraySpec::FromProto(client(), proto));
 
@@ -93,10 +101,12 @@ TEST_P(ArraySpecTest, ToFromProto) {
   EXPECT_EQ(sharding->shard_shape(), shard_shape);
 }
 
-INSTANTIATE_TEST_SUITE_P(NumDevices, ArraySpecTest,
-                         testing::Values(test_util::DeviceTestParam{
-                             /*num_devices=*/2,
-                             /*num_addressable_devices=*/2}));
+INSTANTIATE_TEST_SUITE_P(
+    SerDesVersion_NumDevices, ArraySpecTest,
+    testing::Combine(testing::ValuesIn(test_util::AllSupportedSerDesVersions()),
+                     testing::Values(test_util::DeviceTestParam{
+                         /*num_devices=*/2,
+                         /*num_addressable_devices=*/2})));
 
 }  // namespace
 }  // namespace ifrt

xla/python/ifrt/attribute_map.cc

@@ -23,17 +23,25 @@ limitations under the License.
 #include <variant>
 #include <vector>
 
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "xla/python/ifrt/attribute_map.pb.h"
+#include "xla/python/ifrt/serdes_version.h"
 
 namespace xla {
 namespace ifrt {
 
 absl::StatusOr<AttributeMap> AttributeMap::FromProto(
     const AttributeMapProto& proto) {
+  const SerDesVersionNumber version_number(proto.version_number());
+  if (version_number != SerDesVersionNumber(0)) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "Unsupported ", version_number, " for AttributeMap deserialization"));
+  }
+
   AttributeMap::Map map;
   map.reserve(proto.attributes_size());
   for (const auto& [key, value] : proto.attributes()) {
@@ -63,8 +71,16 @@ absl::StatusOr<AttributeMap> AttributeMap::FromProto(
   return AttributeMap(std::move(map));
 }
 
-AttributeMapProto AttributeMap::ToProto() const {
+AttributeMapProto AttributeMap::ToProto(SerDesVersion version) const {
+  // TODO(b/423702568): Change the return type to `absl::StatusOr<...>` for
+  // graceful error handling.
+  CHECK_GE(version.version_number(), SerDesVersionNumber(0))
+      << "Unsupported " << version.version_number()
+      << " for AttributeMap serialization";
+
   AttributeMapProto proto;
+  proto.set_version_number(SerDesVersionNumber(0).value());
+
   for (const auto& [key, value] : map_) {
     AttributeMapProto::Value value_proto;
     std::visit(

xla/python/ifrt/attribute_map.h

@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
 #include "xla/python/ifrt/attribute_map.pb.h"
+#include "xla/python/ifrt/serdes_version.h"
 
 namespace xla {
 namespace ifrt {
@@ -88,7 +89,8 @@ class AttributeMap {
   static absl::StatusOr<AttributeMap> FromProto(const AttributeMapProto& proto);
 
   // Serializes `AttributeMap` into `AttributeMapProto`.
-  AttributeMapProto ToProto() const;
+  AttributeMapProto ToProto(
+      SerDesVersion version = SerDesVersion::current()) const;
 
   std::string DebugString(size_t max_string_length = 64,
                           size_t max_int64_list_size = 16) const;

xla/python/ifrt/attribute_map.proto

@@ -19,6 +19,8 @@ package xla.ifrt;
 
 // Proto equivalent of C++ `AttributeMap`.
 message AttributeMapProto {
+  int32 version_number = 2;
+
   message Value {
     message Int64List {
       repeated sfixed64 elements = 1;