1 /*
2 * Copyright 2022 Google LLC
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <pybind11/pybind11.h>
18
19 #include <memory>
20 #include <string>
21
22 #include "absl/status/statusor.h"
23 #include "fcp/aggregation/protocol/aggregation_protocol.h"
24 #include "fcp/aggregation/protocol/configuration.pb.h"
25 #include "fcp/aggregation/protocol/simple_aggregation/simple_aggregation_protocol.h"
26 #include "fcp/aggregation/tensorflow/tensorflow_checkpoint_builder_factory.h"
27 #include "fcp/aggregation/tensorflow/tensorflow_checkpoint_parser_factory.h"
28 #include "pybind11_abseil/status_casters.h"
29 #include "pybind11_protobuf/native_proto_caster.h"
30
31 namespace py = ::pybind11;
32
33 using ::fcp::aggregation::AggregationProtocol;
34 using ::fcp::aggregation::Configuration;
35 using ::fcp::aggregation::ResourceResolver;
36 using ::fcp::aggregation::tensorflow::TensorflowCheckpointBuilderFactory;
37 using ::fcp::aggregation::tensorflow::TensorflowCheckpointParserFactory;
38
PYBIND11_MODULE(aggregation_protocols,m)39 PYBIND11_MODULE(aggregation_protocols, m) {
40 class DefaultResourceResolver : public ResourceResolver {
41 absl::StatusOr<absl::Cord> RetrieveResource(
42 int64_t client_id, const std::string& uri) override {
43 return absl::UnimplementedError("RetrieveResource() is not supported.");
44 }
45 };
46
47 pybind11::google::ImportStatusModule();
48 pybind11_protobuf::ImportNativeProtoCasters();
49
50 static const TensorflowCheckpointBuilderFactory* const
51 kCheckpointBuilderFactory = new TensorflowCheckpointBuilderFactory();
52 static const TensorflowCheckpointParserFactory* const
53 kCheckpointParserFactory = new TensorflowCheckpointParserFactory();
54 static ResourceResolver* kResourceResolver = new DefaultResourceResolver();
55
56 m.def(
57 "create_simple_aggregation_protocol",
58 [](const Configuration& configuration,
59 AggregationProtocol::Callback* callback)
60 -> absl::StatusOr<std::unique_ptr<AggregationProtocol>> {
61 return fcp::aggregation::SimpleAggregationProtocol::Create(
62 configuration, callback, kCheckpointParserFactory,
63 kCheckpointBuilderFactory, kResourceResolver);
64 },
65 // Ensure the Callback object outlives the AggregationProtocol.
66 py::keep_alive<0, 2>());
67 }
68