Clean up the reshard API from the IFRT Proxy server

It has been six months since we switched from `Reshard` to `CopyArrays`. Per compatibility contract, it is now safe to remove the Reshard emulation code on the proxy server.

PiperOrigin-RevId: 707182393

Author: junwhanahn

xla/python/ifrt_proxy/client/rpc_helper.cc

@@ -330,7 +330,6 @@ RPC(CopyToHostBuffer, copy_to_host_buffer);
 RPC(IsArrayDeleted, is_array_deleted);
 RPC(DestructArray, destruct_array)
 RPC(CopyArrays, copy_arrays);
-RPC(Reshard, reshard);
 RPC(FullyReplicatedShard, fully_replicated_shard);
 RPC(DeleteArray, delete_array);
 RPC(Compile, compile);

xla/python/ifrt_proxy/client/rpc_helper.h

@@ -112,7 +112,6 @@ class RpcHelper {
       std::unique_ptr<CopyToHostBufferRequest> req);
   ResponseFuture<CopyArraysResponse> CopyArrays(
       std::unique_ptr<CopyArraysRequest> req);
-  ResponseFuture<ReshardResponse> Reshard(std::unique_ptr<ReshardRequest> req);
   ResponseFuture<FullyReplicatedShardResponse> FullyReplicatedShard(
       std::unique_ptr<FullyReplicatedShardRequest> req);
   ResponseFuture<IsArrayDeletedResponse> IsArrayDeleted(

xla/python/ifrt_proxy/common/ifrt_service.proto

@@ -56,7 +56,6 @@ message IfrtRequest {
         disassemble_into_single_device_arrays_request = 7;
     DeleteArrayRequest delete_array_request = 9;
     CopyArraysRequest copy_arrays_request = 24;
-    ReshardRequest reshard_request = 10 [deprecated = true];
     FullyReplicatedShardRequest fully_replicated_shard_request = 20;
     IsArrayDeletedRequest is_array_deleted_request = 11;
     DestructArrayRequest destruct_array_request = 12;
@@ -79,6 +78,8 @@ message IfrtRequest {
     GetDefaultDeviceAssignmentRequest get_default_device_assignment_request =
         19;
   }
+
+  reserved 10;
 }
 
 message IfrtResponse {
@@ -103,7 +104,6 @@ message IfrtResponse {
         disassemble_into_single_device_arrays_response = 7;
     DeleteArrayResponse delete_array_response = 9;
     CopyArraysResponse copy_arrays_response = 24;
-    ReshardResponse reshard_response = 10 [deprecated = true];
     FullyReplicatedShardResponse fully_replicated_shard_response = 20;
     IsArrayDeletedResponse is_array_deleted_response = 11;
     DestructArrayResponse destruct_array_response = 12;
@@ -127,6 +127,8 @@ message IfrtResponse {
     GetDefaultDeviceAssignmentResponse get_default_device_assignment_response =
         19;
   }
+
+  reserved 10;
 }
 
 // Metadata of an IFRT Request.
@@ -323,15 +325,6 @@ message CopyArraysResponse {
   repeated fixed64 array_handles = 1;
 }
 
-message ReshardRequest {
-  fixed64 array_handle = 1;
-  ShardingProto sharding = 2;
-  proto.ArrayCopySemantics copy_semantics = 3;
-}
-message ReshardResponse {
-  fixed64 array_handle = 1;
-}
-
 message FullyReplicatedShardRequest {
   fixed64 array_handle = 1;
   proto.ArrayCopySemantics copy_semantics = 2;

xla/python/ifrt_proxy/common/versions.h

@@ -26,7 +26,7 @@ namespace protocol_version {
 inline constexpr int kClientMin = 3;
 
 // The minimum protocol_version that the current server code understands.
-inline constexpr int kServerMin = 1;
+inline constexpr int kServerMin = 3;
 
 enum {
   // Versions kAncient are named and are only referred to by their numbers. See

xla/python/ifrt_proxy/server/ifrt_backend.cc

@@ -342,8 +342,6 @@ Future<BackendInterface::Response> IfrtBackend::ProcessInternal(
       return Future<Response>(HandleCheckValueReadyRequest(std::move(request)));
     case IfrtRequest::RequestCase::kCopyArraysRequest:
       return Future<Response>(HandleCopyArraysRequest(std::move(request)));
-    case IfrtRequest::RequestCase::kReshardRequest:
-      return Future<Response>(HandleReshardRequest(std::move(request)));
     case IfrtRequest::RequestCase::kFullyReplicatedShardRequest:
       return Future<Response>(
           HandleFullyReplicatedShardRequest(std::move(request)));
@@ -1029,44 +1027,6 @@ absl::StatusOr<BackendInterface::Response> IfrtBackend::HandleCopyArraysRequest(
   return ifrt_resp;
 }
 
-absl::StatusOr<BackendInterface::Response> IfrtBackend::HandleReshardRequest(
-    std::unique_ptr<IfrtRequest> request) {
-  const auto& reshard_request = request->reshard_request();
-  TF_ASSIGN_OR_RETURN(auto array, GetArray(reshard_request.array_handle()));
-  TF_ASSIGN_OR_RETURN(
-      std::shared_ptr<const Sharding> sharding,
-      Sharding::FromProto(
-          absl::bind_front(&Client::LookupDevice, client_.get()),
-          reshard_request.sharding()));
-  TF_ASSIGN_OR_RETURN(auto semantics, FromArrayCopySemanticsProto(
-                                          reshard_request.copy_semantics()));
-
-  // Emulate the old `Array::Reshard` behavior using `Client::CopyArrays`. No
-  // existing IFRT implementations before `Array::Reshard` was deleted actually
-  // supported resharding, so this should be safe.
-  if (!array->sharding().HasSamePartitioning(*sharding)) {
-    return absl::InvalidArgumentError(absl::StrCat(
-        "IFRT Proxy does not support resharding, but got ",
-        array->sharding().DebugString(), " as the original sharding and ",
-        sharding->DebugString(), " as the target sharding"));
-  }
-  TF_ASSIGN_OR_RETURN(
-      auto copied_arrays,
-      client_->CopyArrays(absl::MakeSpan(&array, 1), sharding->devices(),
-                          sharding->memory_kind(), semantics));
-
-  uint64_t resharded_array_handle = handle_generator_.GenerateAtServer();
-  {
-    absl::MutexLock lock(&arrays_mutex_);
-    arrays_.insert({resharded_array_handle, std::move(copied_arrays[0])});
-  }
-
-  auto ifrt_resp = NewIfrtResponse(request->request_metadata().op_id());
-  ifrt_resp->mutable_reshard_response()->set_array_handle(
-      resharded_array_handle);
-  return ifrt_resp;
-}
-
 absl::StatusOr<BackendInterface::Response>
 IfrtBackend::HandleFullyReplicatedShardRequest(
     std::unique_ptr<IfrtRequest> request) {

xla/python/ifrt_proxy/server/ifrt_backend_test.cc

@@ -970,74 +970,6 @@ TEST_P(IfrtBackendHandlerTest, CopyArrays) {
               SizeIs(copied_arrays.size()));
 }
 
-TEST_P(IfrtBackendHandlerTest, ReshardSuccess) {
-  auto src_mock_array = tsl::MakeRef<xla::ifrt::MockArray>();
-  TF_ASSERT_OK_AND_ASSIGN(auto* device,
-                          mock_client_->LookupDevice(DeviceId(0)));
-  auto src_sharding = SingleDeviceSharding::Create(device, MemoryKind());
-  ON_CALL(*src_mock_array, sharding()).WillByDefault(ReturnRef(*src_sharding));
-  TF_ASSERT_OK_AND_ASSIGN(auto src_array_handle,
-                          MakeTestArray(std::move(src_mock_array)));
-
-  auto copied_mock_array = tsl::MakeRef<xla::ifrt::MockArray>();
-  EXPECT_CALL(*mock_client_, CopyArrays(_, _, _, _))
-      .WillOnce(Return(std::vector<tsl::RCReference<xla::ifrt::Array>>(
-          {copied_mock_array})));
-
-  auto ifrt_request = NewIfrtRequest(NewOpId());
-  auto* reshard_request = ifrt_request->mutable_reshard_request();
-  reshard_request->set_array_handle(src_array_handle);
-  reshard_request->set_copy_semantics(proto::ARRAY_COPY_SEMANTICS_ALWAYS_COPY);
-  TF_ASSERT_OK_AND_ASSIGN(auto* new_device,
-                          mock_client_->LookupDevice(DeviceId(1)));
-  TF_ASSERT_OK_AND_ASSIGN(
-      *ifrt_request->mutable_reshard_request()->mutable_sharding(),
-      SingleDeviceSharding::Create(new_device, MemoryKind())->ToProto());
-
-  TF_ASSERT_OK_AND_ASSIGN(auto response, CallBackend(std::move(ifrt_request)));
-
-  EXPECT_THAT(tsl::StatusFromProto(response->response_metadata().status()),
-              IsOk());
-  EXPECT_NE(response->reshard_response().array_handle(), 0);
-}
-
-TEST_P(IfrtBackendHandlerTest, ReshardFailsWhenTheBackendFails) {
-  auto mock_array = tsl::MakeRef<xla::ifrt::MockArray>();
-  TF_ASSERT_OK_AND_ASSIGN(auto* device,
-                          mock_client_->LookupDevice(DeviceId(1)));
-  auto sharding = SingleDeviceSharding::Create(device, MemoryKind());
-  ON_CALL(*mock_array, sharding()).WillByDefault(ReturnRef(*sharding));
-  TF_ASSERT_OK_AND_ASSIGN(auto array_handle,
-                          MakeTestArray(std::move(mock_array)));
-
-  EXPECT_CALL(*mock_client_, CopyArrays(_, _, _, _))
-      .WillOnce(Return(absl::UnknownError("injected error")));
-
-  auto ifrt_request = NewIfrtRequest(NewOpId());
-  auto* reshard_request = ifrt_request->mutable_reshard_request();
-  reshard_request->set_array_handle(array_handle);
-  reshard_request->set_copy_semantics(proto::ARRAY_COPY_SEMANTICS_ALWAYS_COPY);
-  TF_ASSERT_OK_AND_ASSIGN(auto* new_device,
-                          mock_client_->LookupDevice(DeviceId(1)));
-  TF_ASSERT_OK_AND_ASSIGN(
-      *ifrt_request->mutable_reshard_request()->mutable_sharding(),
-      SingleDeviceSharding::Create(new_device, MemoryKind())->ToProto());
-
-  EXPECT_THAT(CallBackend(std::move(ifrt_request)),
-              StatusIs(absl::StatusCode::kUnknown, StrEq("injected error")));
-}
-
-TEST_P(IfrtBackendHandlerTest, ReshardFailsWithNonExistentArrayHandle) {
-  auto ifrt_request = NewIfrtRequest(NewOpId());
-  auto* reshard_request = ifrt_request->mutable_reshard_request();
-  reshard_request->set_array_handle(0);
-  reshard_request->set_copy_semantics(proto::ARRAY_COPY_SEMANTICS_ALWAYS_COPY);
-  reshard_request->mutable_sharding();
-
-  EXPECT_THAT(CallBackend(std::move(ifrt_request)),
-              StatusIs(absl::StatusCode::kNotFound));
-}
-
 TEST_P(IfrtBackendHandlerTest, FullyReplicatedShardSuccess) {
   auto fully_replicated_mock_array = tsl::MakeRef<xla::ifrt::MockArray>();
   auto resultant_array = tsl::MakeRef<xla::ifrt::MockArray>();