Col_lin
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

tcp support ssl

This commit is contained in:
chendongsheng 2021-04-22 21:21:51 +08:00
parent ebb41de91f
commit 2aae0b01ec
15 changed files with 114 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mindspore/ccsrc/pipeline/jit/init.cc
View File

@ -351,7 +351,8 @@ PYBIND11_MODULE(_c_expression, m) {
.def("set_client_epoch_num", &PSContext::set_client_epoch_num, "Set federated learning client epoch number.")
.def("set_client_batch_size", &PSContext::set_client_batch_size, "Set federated learning client batch size.")
.def("set_secure_aggregation", &PSContext::set_secure_aggregation,
"Set federated learning client using secure aggregation.");
"Set federated learning client using secure aggregation.")
.def("set_enable_ssl", &PSContext::enable_ssl, "Set PS SSL mode enabled or disabled.");
(void)py::class_<OpInfoLoaderPy, std::shared_ptr<OpInfoLoaderPy>>(m, "OpInfoLoaderPy")
.def(py::init())

5
mindspore/ccsrc/ps/constants.h
View File

@ -54,7 +54,8 @@ constexpr char kSparseLazyAdamOp[] = "LazyAdam";
constexpr char kSparseFtrlOp[] = "FTRL";
constexpr char kCertificateChain[] = "server.crt";
constexpr char kPrivateKey[] = "server.key.unsecure";
constexpr char kPrivateKey[] = "server.key";
constexpr char kCAcrt[] = "ca.crt";
constexpr int64_t kInitWeightsCmd = 10;
constexpr int64_t kInitWeightToOptimIdCmd = 11;
@ -82,6 +83,8 @@ constexpr int64_t kSubmitTimeOutInMs = 30000;
constexpr int64_t kRetryCount = 60;
constexpr int64_t kRetryIntervalInMs = 10;
constexpr int64_t kThreadNum = 32;
using DataPtr = std::shared_ptr<unsigned char[]>;
using VectorPtr = std::shared_ptr<std::vector<unsigned char>>;
using Key = uint64_t;

1
mindspore/ccsrc/ps/core/communicator/http_communicator.cc
View File

@ -23,7 +23,6 @@ namespace ps {
namespace core {
bool HttpCommunicator::Start() {
MS_LOG(INFO) << "Initialize http server IP:" << ip_ << ", PORT:" << port_;
http_server_ = std::make_shared<HttpServer>(ip_, port_, 32);
http_server_->InitServer();
MS_EXCEPTION_IF_NULL(http_server_);
if (!http_server_->Start()) {

5
mindspore/ccsrc/ps/core/communicator/http_communicator.h
View File

@ -34,7 +34,10 @@ class HttpCommunicator : public CommunicatorBase {
public:
explicit HttpCommunicator(const std::string &ip, std::int16_t port,
const std::shared_ptr<TaskExecutor> &task_executor)
: task_executor_(task_executor), http_server_(nullptr), ip_(ip), port_(port) {}
: task_executor_(task_executor), http_server_(nullptr), ip_(ip), port_(port) {
http_server_ = std::make_shared<HttpServer>(ip_, port_, kThreadNum);
}
~HttpCommunicator() = default;
bool Start() override;

40
mindspore/ccsrc/ps/core/communicator/http_request_handler.cc
View File

@ -25,28 +25,30 @@ bool HttpRequestHandler::Initialize(int fd, const std::unordered_map<std::string
struct evhttp *http = evhttp_new(evbase_);
MS_EXCEPTION_IF_NULL(http);
SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(),
SSL_OP_SINGLE_DH_USE | SSL_OP_SINGLE_ECDH_USE | SSL_OP_NO_SSLv2);
EC_KEY *ecdh = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
MS_EXCEPTION_IF_NULL(ecdh);
if (PSContext::instance()->enable_ssl()) {
SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(),
SSL_OP_SINGLE_DH_USE | SSL_OP_SINGLE_ECDH_USE | SSL_OP_NO_SSLv2);
EC_KEY *ecdh = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
MS_EXCEPTION_IF_NULL(ecdh);
if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) {
MS_LOG(ERROR) << "SSL use certificate chain file failed!";
return false;
if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) {
MS_LOG(ERROR) << "SSL use certificate chain file failed!";
return false;
}
if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) {
MS_LOG(ERROR) << "SSL use private key file failed!";
return false;
}
if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) {
MS_LOG(ERROR) << "SSL check private key file failed!";
return false;
}
evhttp_set_bevcb(http, BuffereventCallback, SSLWrapper::GetInstance().GetSSLCtx());
}
if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) {
MS_LOG(ERROR) << "SSL use private key file failed!";
return false;
}
if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) {
MS_LOG(ERROR) << "SSL check private key file failed!";
return false;
}
evhttp_set_bevcb(http, BuffereventCallback, SSLWrapper::GetInstance().GetSSLCtx());
int result = evhttp_accept_socket(http, fd);
if (result < 0) {
MS_LOG(ERROR) << "Evhttp accept socket failed!";

1
mindspore/ccsrc/ps/core/communicator/http_request_handler.h
View File

@ -31,6 +31,7 @@
#include "ps/core/communicator/http_message_handler.h"
#include "ps/core/communicator/ssl_wrapper.h"
#include "ps/constants.h"
#include "ps/ps_context.h"
namespace mindspore {
namespace ps {

14
mindspore/ccsrc/ps/core/communicator/ssl_wrapper.cc
View File

@ -20,7 +20,7 @@
namespace mindspore {
namespace ps {
namespace core {
SSLWrapper::SSLWrapper() : ssl_ctx_(nullptr) { InitSSL(); }
SSLWrapper::SSLWrapper() : ssl_ctx_(nullptr), client_ssl_ctx_(nullptr) { InitSSL(); }
SSLWrapper::~SSLWrapper() { CleanSSL(); }
@ -41,6 +41,10 @@ void SSLWrapper::InitSSL() {
if (X509_STORE_set_default_paths(store) != 1) {
MS_LOG(ERROR) << "X509_STORE_set_default_paths failed";
}
client_ssl_ctx_ = SSL_CTX_new(SSLv23_client_method());
if (!ssl_ctx_) {
MS_LOG(ERROR) << "SSL_CTX_new failed";
}
}
void SSLWrapper::CleanSSL() {
@ -53,7 +57,13 @@ void SSLWrapper::CleanSSL() {
CRYPTO_cleanup_all_ex_data();
}
SSL_CTX *SSLWrapper::GetSSLCtx() { return ssl_ctx_; }
SSL_CTX *SSLWrapper::GetSSLCtx(bool is_server) {
if (is_server) {
return ssl_ctx_;
} else {
return client_ssl_ctx_;
}
}
} // namespace core
} // namespace ps
} // namespace mindspore

3
mindspore/ccsrc/ps/core/communicator/ssl_wrapper.h
View File

@ -32,7 +32,7 @@ class SSLWrapper {
static SSLWrapper instance;
return instance;
}
SSL_CTX *GetSSLCtx();
SSL_CTX *GetSSLCtx(bool is_server = true);
private:
SSLWrapper();
@ -44,6 +44,7 @@ class SSLWrapper {
void CleanSSL();
SSL_CTX *ssl_ctx_;
SSL_CTX *client_ssl_ctx_;
};
} // namespace core
} // namespace ps

27
mindspore/ccsrc/ps/core/communicator/tcp_client.cc
View File

@ -106,7 +106,32 @@ void TcpClient::Init() {
sin.sin_addr.s_addr = inet_addr(server_address_.c_str());
sin.sin_port = htons(server_port_);
buffer_event_ = bufferevent_socket_new(event_base_, -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
if (!PSContext::instance()->enable_ssl()) {
buffer_event_ = bufferevent_socket_new(event_base_, -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
} else {
SSL *ssl = SSL_new(SSLWrapper::GetInstance().GetSSLCtx(false));
if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(false), kCertificateChain)) {
MS_LOG(EXCEPTION) << "SSL use certificate chain file failed!";
}
if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(false), kPrivateKey, SSL_FILETYPE_PEM)) {
MS_LOG(EXCEPTION) << "SSL use private key file failed!";
}
if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx(false))) {
MS_LOG(EXCEPTION) << "SSL check private key file failed!";
}
if (!SSL_CTX_load_verify_locations(SSLWrapper::GetInstance().GetSSLCtx(false), kCAcrt, nullptr)) {
MS_LOG(EXCEPTION) << "SSL load ca location failed!";
}
SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(false), SSL_OP_NO_SSLv2);
buffer_event_ = bufferevent_openssl_socket_new(event_base_, -1, ssl, BUFFEREVENT_SSL_CONNECTING,
BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
}
MS_EXCEPTION_IF_NULL(buffer_event_);
bufferevent_setcb(buffer_event_, ReadCallback, nullptr, EventCallback, this);

4
mindspore/ccsrc/ps/core/communicator/tcp_client.h
View File

@ -22,6 +22,7 @@
#include <event2/event.h>
#include <event2/bufferevent.h>
#include <event2/thread.h>
#include <event2/bufferevent_ssl.h>
#include <functional>
#include <string>
@ -35,6 +36,9 @@
#include "ps/core/cluster_metadata.h"
#include "utils/convert_utils_base.h"
#include "ps/core/comm_util.h"
#include "ps/core/communicator/ssl_wrapper.h"
#include "ps/constants.h"
#include "ps/ps_context.h"
namespace mindspore {
namespace ps {

25
mindspore/ccsrc/ps/core/communicator/tcp_server.cc
View File

@ -282,7 +282,30 @@ void TcpServer::ListenerCallback(struct evconnlistener *, evutil_socket_t fd, st
MS_EXCEPTION_IF_NULL(base);
MS_EXCEPTION_IF_NULL(sockaddr);
struct bufferevent *bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
struct bufferevent *bev = nullptr;
if (!PSContext::instance()->enable_ssl()) {
bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
} else {
SSL *ssl = SSL_new(SSLWrapper::GetInstance().GetSSLCtx());
if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) {
MS_LOG(EXCEPTION) << "SSL use certificate chain file failed!";
}
if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) {
MS_LOG(EXCEPTION) << "SSL use private key file failed!";
}
if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) {
MS_LOG(EXCEPTION) << "SSL check private key file failed!";
}
SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(), SSL_OP_NO_SSLv2);
bev = bufferevent_openssl_socket_new(base, fd, ssl, BUFFEREVENT_SSL_ACCEPTING,
BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
}
if (bev == nullptr) {
MS_LOG(ERROR) << "Error constructing buffer event!";
int ret = event_base_loopbreak(base);

4
mindspore/ccsrc/ps/core/communicator/tcp_server.h
View File

@ -22,6 +22,7 @@
#include <event2/event.h>
#include <event2/listener.h>
#include <event2/thread.h>
#include <event2/bufferevent_ssl.h>
#include <exception>
#include <functional>
@ -35,9 +36,12 @@
#include <atomic>
#include "ps/core/communicator/tcp_message_handler.h"
#include "ps/core/communicator/ssl_wrapper.h"
#include "ps/core/cluster_metadata.h"
#include "utils/convert_utils_base.h"
#include "ps/core/comm_util.h"
#include "ps/constants.h"
#include "ps/ps_context.h"
namespace mindspore {
namespace ps {

4
mindspore/ccsrc/ps/ps_context.cc
View File

@ -286,5 +286,9 @@ uint64_t PSContext::worker_upload_weights() const { return worker_upload_weights
void PSContext::set_secure_aggregation(bool secure_aggregation) { secure_aggregation_ = secure_aggregation; }
bool PSContext::secure_aggregation() const { return secure_aggregation_; }
bool PSContext::enable_ssl() const { return enable_ssl_; }
void PSContext::set_enable_ssl(bool enabled) { enable_ssl_ = enabled; }
} // namespace ps
} // namespace mindspore

4
mindspore/ccsrc/ps/ps_context.h
View File

@ -80,6 +80,8 @@ class PSContext {
void CloneHashTable(const std::string &dest_param_name, const std::string &src_param_name) const;
void set_cache_enable(bool cache_enable) const;
void set_rank_id(int rank_id) const;
bool enable_ssl() const;
void set_enable_ssl(bool enabled);
// In new server framework, process role, worker number, server number, scheduler ip and scheduler port should be set
// by ps_context.
@ -145,6 +147,7 @@ class PSContext {
is_worker_(false),
is_pserver_(false),
is_sched_(false),
enable_ssl_(false),
rank_id_(-1),
worker_num_(0),
server_num_(0),
@ -166,6 +169,7 @@ class PSContext {
bool is_worker_;
bool is_pserver_;
bool is_sched_;
bool enable_ssl_;
int rank_id_;
uint32_t worker_num_;
uint32_t server_num_;

3
mindspore/parallel/_ps_context.py
View File

@ -47,7 +47,8 @@ _set_ps_context_func_map = {
"fl_iteration_num": ps_context().set_fl_iteration_num,
"client_epoch_num": ps_context().set_client_epoch_num,
"client_batch_size": ps_context().set_client_batch_size,
"secure_aggregation": ps_context().set_secure_aggregation
"secure_aggregation": ps_context().set_secure_aggregation,
"enable_ps_ssl": ps_context().set_enable_ssl
}
_get_ps_context_func_map = {