wehub-resource-sync
249 lines
8.5 KiB
C++
249 lines
8.5 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing,
|
|
* software distributed under the License is distributed on an
|
|
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
* KIND, either express or implied. See the License for the
|
|
* specific language governing permissions and limitations
|
|
* under the License.
|
|
*/
|
|
|
|
/*!
|
|
* \file rpc_tracker_client.h
|
|
* \brief RPC Tracker client to report resources.
|
|
*/
|
|
#ifndef TVM_APPS_CPP_RPC_TRACKER_CLIENT_H_
|
|
#define TVM_APPS_CPP_RPC_TRACKER_CLIENT_H_
|
|
|
|
#include <chrono>
|
|
#include <iostream>
|
|
#include <random>
|
|
#include <set>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "../../src/runtime/rpc/rpc_endpoint.h"
|
|
#include "../../src/support/socket.h"
|
|
|
|
namespace tvm {
|
|
namespace runtime {
|
|
|
|
/*!
|
|
* \brief TrackerClient Tracker client class.
|
|
* \param tracker The address of RPC tracker in host:port format e.g. 10.77.1.234:9190 Default=""
|
|
* \param key The key used to identify the device type in tracker. Default=""
|
|
* \param custom_addr Custom IP Address to Report to RPC Tracker. Default=""
|
|
*/
|
|
class TrackerClient {
|
|
public:
|
|
/*!
|
|
* \brief Constructor.
|
|
*/
|
|
TrackerClient(const std::string& tracker_addr, const std::string& key,
|
|
const std::string& custom_addr, int port)
|
|
: tracker_addr_(tracker_addr),
|
|
key_(key),
|
|
custom_addr_(custom_addr),
|
|
port_(port),
|
|
gen_(std::random_device{}()),
|
|
dis_(0.0, 1.0) {
|
|
if (custom_addr_.empty()) {
|
|
custom_addr_ = "null";
|
|
} else {
|
|
// Since custom_addr_ can be either the json value null which is not surrounded by quotes
|
|
// or a string containing the custom value which json does required to be quoted then we
|
|
// need to set custom_addr_ to be a string containing the quotes here.
|
|
custom_addr_ = "\"" + custom_addr_ + "\"";
|
|
}
|
|
}
|
|
/*!
|
|
* \brief Destructor.
|
|
*/
|
|
~TrackerClient() {
|
|
// Free the resources
|
|
Close();
|
|
}
|
|
/*!
|
|
* \brief IsValid Check tracker is valid.
|
|
*/
|
|
bool IsValid() { return (!tracker_addr_.empty() && !tracker_sock_.IsClosed()); }
|
|
/*!
|
|
* \brief TryConnect Connect to tracker if the tracker address is valid.
|
|
*/
|
|
void TryConnect() {
|
|
if (!tracker_addr_.empty() && (tracker_sock_.IsClosed())) {
|
|
tracker_sock_ = ConnectWithRetry();
|
|
|
|
int code = kRPCTrackerMagic;
|
|
TVM_FFI_ICHECK_EQ(tracker_sock_.SendAll(&code, sizeof(code)), sizeof(code));
|
|
TVM_FFI_ICHECK_EQ(tracker_sock_.RecvAll(&code, sizeof(code)), sizeof(code));
|
|
TVM_FFI_ICHECK_EQ(code, kRPCTrackerMagic) << tracker_addr_.c_str() << " is not RPC Tracker";
|
|
|
|
std::ostringstream ss;
|
|
ss << "[" << static_cast<int>(TrackerCode::kUpdateInfo) << ", {\"key\": \"server:" << key_
|
|
<< "\", \"addr\": [" << custom_addr_ << ", \"" << port_ << "\"]}]";
|
|
tracker_sock_.SendBytes(ss.str());
|
|
|
|
// Receive status and validate
|
|
std::string remote_status = tracker_sock_.RecvBytes();
|
|
TVM_FFI_ICHECK_EQ(std::stoi(remote_status), static_cast<int>(TrackerCode::kSuccess));
|
|
}
|
|
}
|
|
/*!
|
|
* \brief Close Clean up tracker resources.
|
|
*/
|
|
void Close() {
|
|
// close tracker resource
|
|
if (!tracker_sock_.IsClosed()) {
|
|
tracker_sock_.Close();
|
|
}
|
|
}
|
|
/*!
|
|
* \brief ReportResourceAndGetKey Report resource to tracker.
|
|
* \param port listening port.
|
|
* \param matchkey Random match key output.
|
|
*/
|
|
void ReportResourceAndGetKey(int port, std::string* matchkey) {
|
|
if (!tracker_sock_.IsClosed()) {
|
|
*matchkey = RandomKey(key_ + ":", old_keyset_);
|
|
|
|
std::ostringstream ss;
|
|
ss << "[" << static_cast<int>(TrackerCode::kPut) << ", \"" << key_ << "\", [" << port
|
|
<< ", \"" << *matchkey << "\"], " << custom_addr_ << "]";
|
|
|
|
tracker_sock_.SendBytes(ss.str());
|
|
|
|
// Receive status and validate
|
|
std::string remote_status = tracker_sock_.RecvBytes();
|
|
TVM_FFI_ICHECK_EQ(std::stoi(remote_status), static_cast<int>(TrackerCode::kSuccess));
|
|
} else {
|
|
*matchkey = key_;
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* \brief ReportResourceAndGetKey Report resource to tracker.
|
|
* \param listen_sock Listen socket details for select.
|
|
* \param port listening port.
|
|
* \param ping_period Select wait time.
|
|
* \param matchkey Random match key output.
|
|
*/
|
|
void WaitConnectionAndUpdateKey(support::TCPSocket listen_sock, int port, int ping_period,
|
|
std::string* matchkey) {
|
|
int unmatch_period_count = 0;
|
|
int unmatch_timeout = 4;
|
|
while (true) {
|
|
if (!tracker_sock_.IsClosed()) {
|
|
support::PollHelper poller;
|
|
poller.WatchRead(listen_sock.sockfd);
|
|
poller.Poll(ping_period * 1000);
|
|
if (!poller.CheckRead(listen_sock.sockfd)) {
|
|
std::ostringstream ss;
|
|
ss << "[" << int(TrackerCode::kGetPendingMatchKeys) << "]";
|
|
tracker_sock_.SendBytes(ss.str());
|
|
|
|
// Receive status and validate
|
|
std::string pending_keys = tracker_sock_.RecvBytes();
|
|
old_keyset_.insert(*matchkey);
|
|
|
|
// if match key not in pending key set
|
|
// it means the key is acquired by a client but not used.
|
|
if (pending_keys.find(*matchkey) == std::string::npos) {
|
|
unmatch_period_count += 1;
|
|
} else {
|
|
unmatch_period_count = 0;
|
|
}
|
|
// regenerate match key if key is acquired but not used for a while
|
|
if (unmatch_period_count * ping_period > unmatch_timeout + ping_period) {
|
|
LOG(INFO) << "no incoming connections, regenerate key ...";
|
|
|
|
*matchkey = RandomKey(key_ + ":", old_keyset_);
|
|
|
|
std::ostringstream ss;
|
|
ss << "[" << static_cast<int>(TrackerCode::kPut) << ", \"" << key_ << "\", [" << port
|
|
<< ", \"" << *matchkey << "\"], " << custom_addr_ << "]";
|
|
tracker_sock_.SendBytes(ss.str());
|
|
|
|
std::string remote_status = tracker_sock_.RecvBytes();
|
|
TVM_FFI_ICHECK_EQ(std::stoi(remote_status), static_cast<int>(TrackerCode::kSuccess));
|
|
unmatch_period_count = 0;
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
|
|
private:
|
|
/*!
|
|
* \brief Connect to a RPC address with retry.
|
|
This function is only reliable to short period of server restart.
|
|
* \param timeout Timeout during retry
|
|
* \param retry_period Number of seconds before we retry again.
|
|
* \return TCPSocket The socket information if connect is success.
|
|
*/
|
|
support::TCPSocket ConnectWithRetry(int timeout = 60, int retry_period = 5) {
|
|
auto tbegin = std::chrono::system_clock::now();
|
|
while (true) {
|
|
support::SockAddr addr(tracker_addr_);
|
|
support::TCPSocket sock;
|
|
sock.Create();
|
|
if (sock.Connect(addr)) {
|
|
LOG(INFO) << "Connected to tracker " << addr.AsString();
|
|
return sock;
|
|
}
|
|
|
|
auto period = (std::chrono::duration_cast<std::chrono::seconds>(
|
|
std::chrono::system_clock::now() - tbegin))
|
|
.count();
|
|
TVM_FFI_ICHECK(period < timeout) << "Failed to connect to tracker " << addr.AsString();
|
|
LOG(WARNING) << "Cannot connect to tracker " << addr.AsString() << " retry in "
|
|
<< retry_period << " seconds.";
|
|
std::this_thread::sleep_for(std::chrono::seconds(retry_period));
|
|
}
|
|
}
|
|
/*!
|
|
* \brief Random Generate a random number between 0 and 1.
|
|
* \return random float value.
|
|
*/
|
|
float Random() { return dis_(gen_); }
|
|
/*!
|
|
* \brief Generate a random key.
|
|
* \param prefix The string prefix.
|
|
* \return cmap The conflict map set.
|
|
*/
|
|
std::string RandomKey(const std::string& prefix, const std::set<std::string>& cmap) {
|
|
if (!cmap.empty()) {
|
|
while (true) {
|
|
std::string key = prefix + std::to_string(Random());
|
|
if (cmap.find(key) == cmap.end()) {
|
|
return key;
|
|
}
|
|
}
|
|
}
|
|
return prefix + std::to_string(Random());
|
|
}
|
|
|
|
std::string tracker_addr_;
|
|
std::string key_;
|
|
std::string custom_addr_;
|
|
int port_;
|
|
support::TCPSocket tracker_sock_;
|
|
std::set<std::string> old_keyset_;
|
|
std::mt19937 gen_;
|
|
std::uniform_real_distribution<float> dis_;
|
|
};
|
|
} // namespace runtime
|
|
} // namespace tvm
|
|
#endif // TVM_APPS_CPP_RPC_TRACKER_CLIENT_H_
|