Add Tasks validation to C++ SDK

examples/quickstart-cpp/include/simple_client.h

@@ -1,6 +1,6 @@
 /***********************************************************************************************************
  * 
- * @file libtorch_client.h
+ * @file simple_client.h
  *
  * @brief Define an example flower client, train and test method
  *

src/cc/flwr/include/grpc_rere.h

@@ -1,22 +1,22 @@
 /*************************************************************************************************
  *
- * @file start.h
+ * @file grpc-rere.h
  *
- * @brief Create a gRPC channel to connect to the server and enable message
- *communication
+ * @brief Provide functions for establishing gRPC request-response communication
  *
- * @author Lekang Jiang
+ * @author The Flower Authors
  *
  * @version 1.0
  *
- * @date 06/09/2021
+ * @date 06/11/2023
  *
  *************************************************************************************************/
 
 #ifndef GRPC_RERE_H
 #define GRPC_RERE_H
 #pragma once
 #include "message_handler.h"
+#include "task_handler.h"
 #include <grpcpp/grpcpp.h>
 
 void create_node(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub);

src/cc/flwr/include/task_handler.h

@@ -0,0 +1,24 @@
+/*************************************************************************************************
+ *
+ * @file task_handler.h
+ *
+ * @brief Handle incoming or outgoing tasks
+ *
+ * @author The Flower Authors
+ *
+ * @version 1.0
+ *
+ * @date 06/11/2023
+ *
+ *************************************************************************************************/
+
+#pragma once
+#include "client.h"
+#include "serde.h"
+
+bool validate_task_ins(const flwr::proto::TaskIns &task_ins,
+ const bool discard_reconnect_ins);
+bool validate_task_res(const flwr::proto::TaskRes &task_res);
+flwr::proto::TaskRes configure_task_res(const flwr::proto::TaskRes &task_res,
+ const flwr::proto::TaskIns &task_ins,
+ const flwr::proto::Node &node);

src/cc/flwr/src/grpc_rere.cc

@@ -110,18 +110,14 @@ receive(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub) {
 
  if (response.task_ins_list_size() > 0) {
  flwr::proto::TaskIns task_ins = response.task_ins_list().at(0);
- // TODO: Validate TaskIns
-
- {
+ if (validate_task_ins(task_ins, true)) {
  std::lock_guard<std::mutex> state_lock(state_mutex);
  state[KEY_TASK_INS] = task_ins;
+ return task_ins;
  }
-
- return task_ins;
- } else {
- std::cerr << "TaskIns list is empty." << std::endl;
- return std::nullopt;
  }
+ std::cerr << "TaskIns list is empty." << std::endl;
+ return std::nullopt;
 }
 
 void send(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub,
@@ -136,7 +132,12 @@ void send(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub,
  return;
  }
 
- // TODO: Validate TaskIns
+ if (!validate_task_res(task_res)) {
+ std::cerr << "TaskRes is invalid" << std::endl;
+ std::lock_guard<std::mutex> state_lock(state_mutex);
+ state[KEY_TASK_INS].reset();
+ return;
+ }
 
  flwr::proto::TaskRes new_task_res =
  configure_task_res(task_res, *task_ins, *node);
@@ -151,8 +152,8 @@ void send(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub,
  if (!status.ok()) {
  std::cerr << "PushTaskRes RPC failed with status: "
  << status.error_message() << std::endl;
-  return;
- } else {
+ }
+ {
  std::lock_guard<std::mutex> state_lock(state_mutex);
  state[KEY_TASK_INS].reset();
  }

src/cc/flwr/src/message_handler.cc

@@ -102,34 +102,3 @@ handle_task(flwr_local::Client *client, const flwr::proto::TaskIns &task_ins) {
  return std::make_tuple(task_res, std::get<1>(legacy_res),
  std::get<2>(legacy_res));
 }
-
-flwr::proto::TaskRes
-configure_task_res(const flwr::proto::TaskRes &task_res,
- const flwr::proto::TaskIns &ref_task_ins,
- const flwr::proto::Node &producer) {
- flwr::proto::TaskRes result_task_res;
-
- // Setting scalar fields
- result_task_res.set_task_id(""); // This will be generated by the server
- result_task_res.set_group_id(ref_task_ins.group_id());
- result_task_res.set_workload_id(ref_task_ins.workload_id());
-
- // Merge the task from the input task_res
- *result_task_res.mutable_task() = task_res.task();
-
- // Construct and set the producer and consumer for the task
- std::unique_ptr<flwr::proto::Node> new_producer =
- std::make_unique<flwr::proto::Node>(producer);
- result_task_res.mutable_task()->set_allocated_producer(
- new_producer.release());
-
- std::unique_ptr<flwr::proto::Node> new_consumer =
- std::make_unique<flwr::proto::Node>(ref_task_ins.task().producer());
- result_task_res.mutable_task()->set_allocated_consumer(
- new_consumer.release());
-
- // Set ancestry in the task
- result_task_res.mutable_task()->add_ancestry(ref_task_ins.task_id());
-
- return result_task_res;
-}

src/cc/flwr/src/task_handler.cc

@@ -0,0 +1,52 @@
+#include "task_handler.h"
+
+bool validate_task_ins(const flwr::proto::TaskIns &task_ins,
+ const bool discard_reconnect_ins) {
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+ return !(!task_ins.has_task() ||
+ (!task_ins.task().has_legacy_server_message() &&
+ !task_ins.task().has_sa()) ||
+ (discard_reconnect_ins &&
+ task_ins.task().legacy_server_message().has_reconnect_ins()));
+#pragma GCC diagnostic pop
+}
+
+bool validate_task_res(const flwr::proto::TaskRes &task_res) {
+ // Retrieve initialized fields in TaskRes
+ return (task_res.task_id().empty() && task_res.group_id().empty() &&
+ task_res.workload_id() == 0 && !task_res.task().has_producer() &&
+ !task_res.task().has_producer() && !task_res.task().has_consumer() &&
+ task_res.task().ancestry_size() == 0);
+}
+
+flwr::proto::TaskRes
+configure_task_res(const flwr::proto::TaskRes &task_res,
+ const flwr::proto::TaskIns &ref_task_ins,
+ const flwr::proto::Node &producer) {
+ flwr::proto::TaskRes result_task_res;
+
+ // Setting scalar fields
+ result_task_res.set_task_id(""); // This will be generated by the server
+ result_task_res.set_group_id(ref_task_ins.group_id());
+ result_task_res.set_workload_id(ref_task_ins.workload_id());
+
+ // Merge the task from the input task_res
+ *result_task_res.mutable_task() = task_res.task();
+
+ // Construct and set the producer and consumer for the task
+ std::unique_ptr<flwr::proto::Node> new_producer =
+ std::make_unique<flwr::proto::Node>(producer);
+ result_task_res.mutable_task()->set_allocated_producer(
+ new_producer.release());
+
+ std::unique_ptr<flwr::proto::Node> new_consumer =
+ std::make_unique<flwr::proto::Node>(ref_task_ins.task().producer());
+ result_task_res.mutable_task()->set_allocated_consumer(
+ new_consumer.release());
+
+ // Set ancestry in the task
+ result_task_res.mutable_task()->add_ancestry(ref_task_ins.task_id());
+
+ return result_task_res;
+}