Implement Serialize and Deserialize methods for CheckpointAggregator.

PiperOrigin-RevId: 631908037

tensorflow_federated/cc/core/impl/aggregation/protocol/BUILD

@@ -41,6 +41,16 @@ proto_library(
     deps = ["//tensorflow_federated/cc/core/impl/aggregation/core:tensor_proto"],
 )
 
+proto_library(
+    name = "checkpoint_aggregator_proto",
+    srcs = ["checkpoint_aggregator.proto"],
+)
+
+cc_proto_library(
+    name = "checkpoint_aggregator_cc_proto",
+    deps = [":checkpoint_aggregator_proto"],
+)
+
 cc_proto_library(
     name = "configuration_cc_proto",
     visibility = ["//visibility:public"],
@@ -99,6 +109,7 @@ cc_library(
     hdrs = ["checkpoint_aggregator.h"],
     visibility = ["//visibility:public"],
     deps = [
+        ":checkpoint_aggregator_cc_proto",
         ":checkpoint_builder",
         ":checkpoint_parser",
         ":config_converter",

tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_aggregator.cc

@@ -36,6 +36,7 @@
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_aggregator_factory.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_aggregator_registry.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_spec.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_aggregator.pb.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_builder.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_parser.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/config_converter.h"
@@ -67,13 +68,50 @@ absl::Status CheckpointAggregator::ValidateConfig(
 
 absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
 CheckpointAggregator::Create(const Configuration& configuration) {
+  return CreateInternal(configuration, nullptr);
+}
+
+absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
+CheckpointAggregator::Create(const std::vector<Intrinsic>* intrinsics) {
+  return CreateInternal(intrinsics, nullptr);
+}
+
+absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
+CheckpointAggregator::Deserialize(const Configuration& configuration,
+                                  std::string serialized_state) {
+  CheckpointAggregatorState aggregator_state;
+  if (!aggregator_state.ParseFromString(serialized_state)) {
+    return absl::InvalidArgumentError("Failed to parse serialized state.");
+  }
+  return CreateInternal(configuration, &aggregator_state);
+}
+
+absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
+CheckpointAggregator::Deserialize(const std::vector<Intrinsic>* intrinsics,
+                                  std::string serialized_state) {
+  CheckpointAggregatorState aggregator_state;
+  if (!aggregator_state.ParseFromString(serialized_state)) {
+    return absl::InvalidArgumentError("Failed to parse serialized state.");
+  }
+  return CreateInternal(intrinsics, &aggregator_state);
+}
+
+absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
+CheckpointAggregator::CreateInternal(
+    const Configuration& configuration,
+    const CheckpointAggregatorState* aggregator_state) {
   TFF_ASSIGN_OR_RETURN(std::vector<Intrinsic> intrinsics,
                        ParseFromConfig(configuration));
 
   std::vector<std::unique_ptr<TensorAggregator>> aggregators;
-  for (const Intrinsic& intrinsic : intrinsics) {
+  for (int i = 0; i < intrinsics.size(); ++i) {
+    const Intrinsic& intrinsic = intrinsics[i];
+    const std::string* serialized_aggregator = nullptr;
+    if (aggregator_state != nullptr) {
+      serialized_aggregator = &aggregator_state->aggregators(i);
+    }
     TFF_ASSIGN_OR_RETURN(std::unique_ptr<TensorAggregator> aggregator,
-                         CreateAggregator(intrinsic));
+                         CreateAggregator(intrinsic, serialized_aggregator));
     aggregators.push_back(std::move(aggregator));
   }
 
@@ -82,11 +120,18 @@ CheckpointAggregator::Create(const Configuration& configuration) {
 }
 
 absl::StatusOr<std::unique_ptr<CheckpointAggregator>>
-CheckpointAggregator::Create(const std::vector<Intrinsic>* intrinsics) {
+CheckpointAggregator::CreateInternal(
+    const std::vector<Intrinsic>* intrinsics,
+    const CheckpointAggregatorState* aggregator_state) {
   std::vector<std::unique_ptr<TensorAggregator>> aggregators;
-  for (const Intrinsic& intrinsic : *intrinsics) {
+  for (int i = 0; i < intrinsics->size(); ++i) {
+    const Intrinsic& intrinsic = (*intrinsics)[i];
+    const std::string* serialized_aggregator = nullptr;
+    if (aggregator_state != nullptr) {
+      serialized_aggregator = &aggregator_state->aggregators(i);
+    }
     TFF_ASSIGN_OR_RETURN(std::unique_ptr<TensorAggregator> aggregator,
-                         CreateAggregator(intrinsic));
+                         CreateAggregator(intrinsic, serialized_aggregator));
     aggregators.push_back(std::move(aggregator));
   }
 
@@ -201,14 +246,33 @@ absl::Status CheckpointAggregator::Report(
 
 void CheckpointAggregator::Abort() { aggregation_finished_ = true; }
 
+absl::StatusOr<std::string> CheckpointAggregator::Serialize() && {
+  absl::MutexLock lock(&aggregation_mu_);
+  if (aggregation_finished_) {
+    return absl::AbortedError("Aggregation has already been finished.");
+  }
+  CheckpointAggregatorState state;
+  google::protobuf::RepeatedPtrField<std::string>* aggregators_proto =
+      state.mutable_aggregators();
+  aggregators_proto->Reserve(aggregators_.size());
+  for (const auto& aggregator : aggregators_) {
+    aggregators_proto->Add(std::move(*aggregator).Serialize().value());
+  }
+  return state.SerializeAsString();
+}
+
 absl::StatusOr<std::unique_ptr<TensorAggregator>>
-CheckpointAggregator::CreateAggregator(const Intrinsic& intrinsic) {
+CheckpointAggregator::CreateAggregator(
+    const Intrinsic& intrinsic, const std::string* serialized_aggregator) {
   // Resolve the intrinsic_uri to the registered TensorAggregatorFactory.
   TFF_ASSIGN_OR_RETURN(const TensorAggregatorFactory* factory,
                        GetAggregatorFactory(intrinsic.uri));
 
   // Use the factory to create the TensorAggregator instance.
-  return factory->Create(intrinsic);
+  if (serialized_aggregator == nullptr) {
+    return factory->Create(intrinsic);
+  }
+  return factory->Deserialize(intrinsic, *serialized_aggregator);
 }
 
 std::vector<std::unique_ptr<TensorAggregator>>

tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_aggregator.h

@@ -22,6 +22,7 @@
 #include <atomic>
 #include <memory>
 #include <optional>
+#include <string>
 #include <vector>
 
 #include "absl/base/attributes.h"
@@ -31,6 +32,7 @@
 #include "absl/synchronization/mutex.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/intrinsic.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_aggregator.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_aggregator.pb.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_builder.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_parser.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/configuration.pb.h"
@@ -62,6 +64,20 @@ class CheckpointAggregator {
   static absl::StatusOr<std::unique_ptr<CheckpointAggregator>> Create(
       const std::vector<Intrinsic>* intrinsics ABSL_ATTRIBUTE_LIFETIME_BOUND);
 
+  // Creates an instance of CheckpointAggregator based on the given
+  // configuration and serialized state.
+  static absl::StatusOr<std::unique_ptr<CheckpointAggregator>> Deserialize(
+      const Configuration& configuration, std::string serialized_state);
+
+  // Creates an instance of CheckpointAggregator based on the given intrinsics
+  // and serialized state.
+  // The `intrinsics` are expected to be created using `ParseFromConfig` which
+  // validates the configuration. CheckpointAggregator does not take any
+  // ownership, and `intrinsics` must outlive it.
+  static absl::StatusOr<std::unique_ptr<CheckpointAggregator>> Deserialize(
+      const std::vector<Intrinsic>* intrinsics ABSL_ATTRIBUTE_LIFETIME_BOUND,
+      std::string serialized_state);
+
   // Accumulates a checkpoint via nested tensor aggregators. The tensors are
   // provided by the CheckpointParser instance.
   absl::Status Accumulate(CheckpointParser& checkpoint_parser);
@@ -75,6 +91,8 @@ class CheckpointAggregator {
   // Signal that the aggregation must be aborted and the report can't be
   // produced.
   void Abort();
+  // Serialize the internal state of the checkpoint aggregator as a string.
+  absl::StatusOr<std::string> Serialize() &&;
 
  private:
   CheckpointAggregator(
@@ -85,9 +103,18 @@ class CheckpointAggregator {
       std::vector<Intrinsic> intrinsics,
       std::vector<std::unique_ptr<TensorAggregator>> aggregators);
 
-  // Creates an aggregation intrinsic based on the intrinsic configuration.
+  // Creates an aggregation intrinsic based on the intrinsic configuration and
+  // optional serialized state.
   static absl::StatusOr<std::unique_ptr<TensorAggregator>> CreateAggregator(
-      const Intrinsic& intrinsic);
+      const Intrinsic& intrinsic, const std::string* serialized_aggregator);
+
+  static absl::StatusOr<std::unique_ptr<CheckpointAggregator>> CreateInternal(
+      const Configuration& configuration,
+      const CheckpointAggregatorState* aggregator_state);
+
+  static absl::StatusOr<std::unique_ptr<CheckpointAggregator>> CreateInternal(
+      const std::vector<Intrinsic>* intrinsics,
+      const CheckpointAggregatorState* aggregator_state);
 
   // Used by the implementation of Merge.
   std::vector<std::unique_ptr<TensorAggregator>> TakeAggregators() &&;

tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_aggregator.proto

@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+package tensorflow_federated.aggregation;
+
+// Internal state representation of a CheckpointAggregator.
+message CheckpointAggregatorState {
+  repeated bytes aggregators = 1;
+}