wehub-resource-sync
289 lines
9.6 KiB
C++
289 lines
9.6 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_env.cc
|
|
* \brief Server environment of the RPC.
|
|
*/
|
|
#include "rpc_env.h"
|
|
|
|
#include <tvm/ffi/extra/module.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/runtime/logging.h>
|
|
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <system_error>
|
|
#include <vector>
|
|
|
|
#include "../../src/support/utils.h"
|
|
|
|
namespace {
|
|
std::string GenerateUntarCommand(const std::string& tar_file, const std::string& output_dir) {
|
|
std::string untar_cmd;
|
|
untar_cmd.reserve(512);
|
|
#if defined(__linux__) || defined(__ANDROID__) || defined(__APPLE__)
|
|
untar_cmd += "tar -C ";
|
|
untar_cmd += output_dir;
|
|
untar_cmd += " -zxf ";
|
|
untar_cmd += tar_file;
|
|
#elif defined(_WIN32)
|
|
untar_cmd += "python -m tarfile -e ";
|
|
untar_cmd += tar_file;
|
|
untar_cmd += " ";
|
|
untar_cmd += output_dir;
|
|
#endif
|
|
return untar_cmd;
|
|
}
|
|
|
|
} // Anonymous namespace
|
|
|
|
namespace tvm {
|
|
namespace runtime {
|
|
|
|
RPCEnv::RPCEnv(const std::string& wd) {
|
|
std::error_code ec;
|
|
if (!wd.empty()) {
|
|
base_ = wd + "/.cache";
|
|
std::filesystem::create_directories(base_, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Failed to create directory " << base_ << " : " << ec.message();
|
|
}
|
|
} else {
|
|
#if defined(ANDROID) || defined(__ANDROID__)
|
|
std::string pkg_name;
|
|
if (std::ifstream cmdline("/proc/self/cmdline"); cmdline) {
|
|
std::getline(cmdline, pkg_name, '\0');
|
|
}
|
|
std::string android_base_ = "/data/data/" + pkg_name + "/cache";
|
|
// Check if application data directory exist. If not exist, usually means we run tvm_rpc from
|
|
// adb shell terminal.
|
|
if (!std::filesystem::is_directory(android_base_)) {
|
|
// Tmp directory is always writable for 'shell' user.
|
|
android_base_ = "/data/local/tmp";
|
|
}
|
|
base_ = android_base_ + "/rpc";
|
|
#elif !defined(_WIN32)
|
|
base_ = std::filesystem::current_path(ec).string() + "/rpc";
|
|
if (ec) {
|
|
base_ = "./rpc";
|
|
}
|
|
#else
|
|
base_ = "./rpc";
|
|
#endif
|
|
std::filesystem::create_directories(base_, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Failed to create directory " << base_ << " : " << ec.message();
|
|
}
|
|
}
|
|
std::filesystem::permissions(base_, std::filesystem::perms::all,
|
|
std::filesystem::perm_options::replace, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Failed to grant permissions to " << base_ << " : " << ec.message();
|
|
}
|
|
|
|
ffi::Function::SetGlobal(
|
|
"tvm.rpc.server.workpath",
|
|
ffi::Function::FromTyped([this](const std::string& path) { return this->GetPath(path); }));
|
|
|
|
ffi::Function::SetGlobal("tvm.rpc.server.listdir",
|
|
ffi::Function::FromTyped([this](const std::string& path) {
|
|
std::string dir = this->GetPath(path);
|
|
std::ostringstream os;
|
|
for (auto d : ListDir(dir)) {
|
|
os << d << ",";
|
|
}
|
|
return os.str();
|
|
}));
|
|
|
|
ffi::Function::SetGlobal("tvm.rpc.server.load_module",
|
|
ffi::Function::FromTyped([this](const std::string& path) {
|
|
std::string file_name = this->GetPath(path);
|
|
file_name = BuildSharedLibrary(file_name);
|
|
LOG(INFO) << "Load module from " << file_name << " ...";
|
|
return ffi::Module::LoadFromFile(file_name);
|
|
}));
|
|
|
|
ffi::Function::SetGlobal("tvm.rpc.server.download_linked_module",
|
|
ffi::Function::FromTyped([this](const std::string& path) {
|
|
std::string file_name = this->GetPath(path);
|
|
file_name = BuildSharedLibrary(file_name);
|
|
std::string bin;
|
|
|
|
std::ifstream fs(file_name, std::ios::in | std::ios::binary);
|
|
TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << file_name;
|
|
fs.seekg(0, std::ios::end);
|
|
size_t size = static_cast<size_t>(fs.tellg());
|
|
fs.seekg(0, std::ios::beg);
|
|
bin.resize(size);
|
|
fs.read(bin.data(), size);
|
|
LOG(INFO) << "Send linked module " << file_name << " to client";
|
|
return ffi::Bytes(bin);
|
|
}));
|
|
}
|
|
|
|
/*!
|
|
* \brief GetPath To get the work path from packed function
|
|
* \param file_name The file name
|
|
* \return The full path of file.
|
|
*/
|
|
std::string RPCEnv::GetPath(const std::string& file_name) const {
|
|
// we assume file_name starts with "/" means file_name is the exact path
|
|
// and does not create /.rpc/
|
|
return !file_name.empty() && file_name[0] == '/' ? file_name : base_ + "/" + file_name;
|
|
}
|
|
|
|
/*!
|
|
* \brief Remove The RPC Environment cleanup function
|
|
*/
|
|
void RPCEnv::CleanUp() const {
|
|
std::error_code ec;
|
|
std::filesystem::remove_all(base_, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Cleanup " << base_ << " failed: " << ec.message();
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* \brief ListDir get the list of files in a directory
|
|
* \param dirname The root directory name
|
|
* \return vector Files in directory.
|
|
*/
|
|
std::vector<std::string> ListDir(const std::string& dirname) {
|
|
std::error_code ec;
|
|
std::vector<std::string> vec;
|
|
auto iter = std::filesystem::directory_iterator(dirname, ec);
|
|
if (ec) {
|
|
if (ec == std::errc::no_such_file_or_directory) return vec;
|
|
TVM_FFI_THROW(InternalError) << "ListDir " << dirname << " error: " << ec.message();
|
|
}
|
|
for (const auto& entry : iter) {
|
|
vec.push_back(entry.path().generic_string());
|
|
}
|
|
return vec;
|
|
}
|
|
|
|
#if defined(__linux__) || defined(__ANDROID__) || defined(__APPLE__)
|
|
/*!
|
|
* \brief LinuxShared Creates a linux shared library
|
|
* \param output The output file name
|
|
* \param files The files for building
|
|
* \param options The compiler options
|
|
* \param cc The compiler
|
|
*/
|
|
void LinuxShared(const std::string output, const std::vector<std::string>& files,
|
|
std::string options = "", std::string cc = "g++") {
|
|
std::string cmd = cc;
|
|
cmd += " -shared -fPIC ";
|
|
cmd += " -o " + output;
|
|
for (auto f = files.begin(); f != files.end(); ++f) {
|
|
cmd += " " + *f;
|
|
}
|
|
cmd += " " + options;
|
|
std::string err_msg;
|
|
auto executed_status = support::Execute(cmd, &err_msg);
|
|
if (executed_status) {
|
|
TVM_FFI_THROW(InternalError) << err_msg;
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef _WIN32
|
|
/*!
|
|
* \brief WindowsShared Creates a Windows shared library
|
|
* \param output The output file name
|
|
* \param files The files for building
|
|
* \param options The compiler options
|
|
* \param cc The compiler
|
|
*/
|
|
void WindowsShared(const std::string& output, const std::vector<std::string>& files,
|
|
const std::string& options = "", const std::string& cc = "clang") {
|
|
std::string cmd = cc;
|
|
cmd += " -O2 -flto=full -fuse-ld=lld-link -shared ";
|
|
cmd += " -o " + output;
|
|
for (const auto& file : files) {
|
|
cmd += " " + file;
|
|
}
|
|
cmd += " " + options;
|
|
std::string err_msg;
|
|
const auto executed_status = support::Execute(cmd, &err_msg);
|
|
if (executed_status) {
|
|
TVM_FFI_THROW(InternalError) << err_msg;
|
|
}
|
|
}
|
|
#endif
|
|
|
|
/*!
|
|
* \brief CreateShared Creates a shared library
|
|
* \param output The output file name
|
|
* \param files The files for building
|
|
*/
|
|
void CreateShared(const std::string& output, const std::vector<std::string>& files) {
|
|
#if defined(__linux__) || defined(__ANDROID__) || defined(__APPLE__)
|
|
LinuxShared(output, files);
|
|
#elif defined(_WIN32)
|
|
WindowsShared(output, files);
|
|
#else
|
|
TVM_FFI_THROW(InternalError) << "Operating system not supported";
|
|
#endif
|
|
}
|
|
|
|
std::string BuildSharedLibrary(std::string file) {
|
|
if (support::EndsWith(file, ".so") || support::EndsWith(file, ".dll") ||
|
|
support::EndsWith(file, ".dylib")) {
|
|
return file;
|
|
}
|
|
|
|
std::string file_name = file + ".so";
|
|
if (support::EndsWith(file, ".o")) {
|
|
CreateShared(file_name, {file});
|
|
} else if (support::EndsWith(file, ".tar")) {
|
|
const std::string tmp_dir = "./rpc/tmp/";
|
|
std::error_code ec;
|
|
std::filesystem::create_directories(tmp_dir, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Failed to create directory " << tmp_dir << " : " << ec.message();
|
|
}
|
|
std::filesystem::permissions(tmp_dir, std::filesystem::perms::all,
|
|
std::filesystem::perm_options::replace, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Failed to grant permissions to " << tmp_dir << " : " << ec.message();
|
|
}
|
|
|
|
const std::string cmd = GenerateUntarCommand(file, tmp_dir);
|
|
|
|
std::string err_msg;
|
|
const int executed_status = support::Execute(cmd, &err_msg);
|
|
if (executed_status) {
|
|
TVM_FFI_THROW(InternalError) << err_msg;
|
|
}
|
|
CreateShared(file_name, ListDir(tmp_dir));
|
|
std::filesystem::remove_all(tmp_dir, ec);
|
|
if (ec) {
|
|
LOG(WARNING) << "Remove " << tmp_dir << " failed: " << ec.message();
|
|
}
|
|
} else {
|
|
file_name = file;
|
|
}
|
|
return file_name;
|
|
}
|
|
|
|
} // namespace runtime
|
|
} // namespace tvm
|