From c2c4128659744b7caaa4b1b6f04c60e28d596ecc Mon Sep 17 00:00:00 2001 From: EunTaik Lee Date: Tue, 16 Mar 2021 18:53:26 +0900 Subject: [PATCH 1/2] [CPP_RPC] allow user supplied work dir --- apps/cpp_rpc/main.cc | 10 +++++++++- apps/cpp_rpc/rpc_env.cc | 35 ++++++++++++++++++++--------------- apps/cpp_rpc/rpc_env.h | 4 ++-- apps/cpp_rpc/rpc_server.cc | 20 +++++++++++--------- apps/cpp_rpc/rpc_server.h | 4 +++- 5 files changed, 45 insertions(+), 28 deletions(-) diff --git a/apps/cpp_rpc/main.cc b/apps/cpp_rpc/main.cc index e381dd2b261b..0663c378819e 100644 --- a/apps/cpp_rpc/main.cc +++ b/apps/cpp_rpc/main.cc @@ -55,6 +55,7 @@ static const string kUsage = "--tracker - The RPC tracker address in host:port format e.g. 10.1.1.2:9190 Default=\"\"\n" "--key - The key used to identify the device type in tracker. Default=\"\"\n" "--custom-addr - Custom IP Address to Report to RPC Tracker. Default=\"\"\n" + "--work-dir - Custom work directory. Default=\"\"\n" "--silent - Whether to run in silent mode. Default=False\n" "\n" " Example\n" @@ -70,6 +71,7 @@ static const string kUsage = * \arg tracker The address of RPC tracker in host:port format e.g. 10.77.1.234:9190 Default="" * \arg key The key used to identify the device type in tracker. Default="" * \arg custom_addr Custom IP Address to Report to RPC Tracker. Default="" + * \arg work_dir Custom work directory. Default="" * \arg silent Whether run in silent mode. Default=False */ struct RpcServerArgs { @@ -79,6 +81,7 @@ struct RpcServerArgs { string tracker; string key; string custom_addr; + string work_dir; bool silent = false; #if defined(WIN32) std::string mmap_path; @@ -96,6 +99,7 @@ void PrintArgs(const RpcServerArgs& args) { LOG(INFO) << "tracker = " << args.tracker; LOG(INFO) << "key = " << args.key; LOG(INFO) << "custom_addr = " << args.custom_addr; + LOG(INFO) << "work_dir = " << args.work_dir; LOG(INFO) << "silent = " << ((args.silent) ? ("True") : ("False")); } @@ -238,6 +242,10 @@ void ParseCmdArgs(int argc, char* argv[], struct RpcServerArgs& args) { dmlc::InitLogging("--minloglevel=0"); } #endif + const string work_dir = GetCmdOption(argc, argv, "--work-dir="); + if (!work_dir.empty()) { + args.work_dir = work_dir; + } } /*! @@ -274,7 +282,7 @@ int RpcServer(int argc, char* argv[]) { #endif RPCServerCreate(args.host, args.port, args.port_end, args.tracker, args.key, args.custom_addr, - args.silent); + args.work_dir, args.silent); return 0; } diff --git a/apps/cpp_rpc/rpc_env.cc b/apps/cpp_rpc/rpc_env.cc index ea19cfa3979d..5f703e1dc2b0 100644 --- a/apps/cpp_rpc/rpc_env.cc +++ b/apps/cpp_rpc/rpc_env.cc @@ -39,7 +39,6 @@ int mkdir(const char* path, int /* ignored */) { return _mkdir(path); } #include #include #include - #include "../../src/support/utils.h" #include "rpc_env.h" @@ -85,25 +84,31 @@ void CleanDir(const std::string& dirname); */ std::string BuildSharedLibrary(std::string file_in); -RPCEnv::RPCEnv() { +RPCEnv::RPCEnv(const std::string& wd) { + if (wd != "") { + base_ = wd + "/.cache"; + mkdir(wd.c_str(), 0777); + mkdir(base_.c_str(), 0777); + } else { #if defined(ANDROID) || defined(__ANDROID__) - char cwd[PATH_MAX]; - auto cmdline = fopen("/proc/self/cmdline", "r"); - fread(cwd, 1, sizeof(cwd), cmdline); - fclose(cmdline); - base_ = "/data/data/" + std::string(cwd) + "/cache/rpc"; + char cwd[PATH_MAX]; + auto cmdline = fopen("/proc/self/cmdline", "r"); + fread(cwd, 1, sizeof(cwd), cmdline); + fclose(cmdline); + base_ = "/data/data/" + std::string(cwd) + "/cache/rpc"; #elif !defined(_WIN32) - char cwd[PATH_MAX]; - if (getcwd(cwd, sizeof(cwd))) { - base_ = std::string(cwd) + "/rpc"; - } else { - base_ = "./rpc"; - } + char cwd[PATH_MAX]; + if (getcwd(cwd, sizeof(cwd))) { + base_ = std::string(cwd) + "/rpc"; + } else { + base_ = "./rpc"; + } #else - base_ = "./rpc"; + base_ = "./rpc"; #endif + mkdir(base_.c_str(), 0777); + } - mkdir(base_.c_str(), 0777); TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body([this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetPath(args[0]); }); diff --git a/apps/cpp_rpc/rpc_env.h b/apps/cpp_rpc/rpc_env.h index 50ef3835e015..1ca6ddaf79b3 100644 --- a/apps/cpp_rpc/rpc_env.h +++ b/apps/cpp_rpc/rpc_env.h @@ -36,10 +36,10 @@ namespace runtime { */ struct RPCEnv { public: - /*! + /*! * \brief Constructor Init The RPC Environment initialize function */ - RPCEnv(); + RPCEnv(const std::string& word_dir = ""); /*! * \brief GetPath To get the workpath from packed function * \param name The file name diff --git a/apps/cpp_rpc/rpc_server.cc b/apps/cpp_rpc/rpc_server.cc index a4028ff61eca..1a5474a71100 100644 --- a/apps/cpp_rpc/rpc_server.cc +++ b/apps/cpp_rpc/rpc_server.cc @@ -98,14 +98,15 @@ class RPCServer { * \brief Constructor. */ RPCServer(std::string host, int port, int port_end, std::string tracker_addr, std::string key, - std::string custom_addr) + std::string custom_addr, std::string work_dir) : host_(std::move(host)), port_(port), my_port_(0), port_end_(port_end), tracker_addr_(std::move(tracker_addr)), key_(std::move(key)), - custom_addr_(std::move(custom_addr)) {} + custom_addr_(std::move(custom_addr)), + work_dir_(std::move(work_dir)) {} /*! * \brief Destructor. @@ -174,7 +175,7 @@ class RPCServer { const pid_t worker_pid = fork(); if (worker_pid == 0) { // Worker process - ServerLoopProc(conn, addr); + ServerLoopProc(conn, addr, work_dir_); _exit(0); } @@ -201,7 +202,7 @@ class RPCServer { } else { auto pid = fork(); if (pid == 0) { - ServerLoopProc(conn, addr); + ServerLoopProc(conn, addr, work_dir_); exit(0); } // Wait for the result @@ -308,9 +309,9 @@ class RPCServer { * \param sock The socket information * \param addr The socket address information */ - static void ServerLoopProc(support::TCPSocket sock, support::SockAddr addr) { + static void ServerLoopProc(support::TCPSocket sock, support::SockAddr addr, std::string work_dir) { // Server loop - const auto env = RPCEnv(); + const auto env = RPCEnv(work_dir); RPCServerLoop(int(sock.sockfd)); LOG(INFO) << "Finish serving " << addr.AsString(); env.CleanUp(); @@ -339,6 +340,7 @@ class RPCServer { std::string tracker_addr_; std::string key_; std::string custom_addr_; + std::string work_dir_; support::TCPSocket listen_sock_; support::TCPSocket tracker_sock_; }; @@ -370,19 +372,19 @@ void ServerLoopFromChild(SOCKET socket) { * silent mode. Default=True */ void RPCServerCreate(std::string host, int port, int port_end, std::string tracker_addr, - std::string key, std::string custom_addr, bool silent) { + std::string key, std::string custom_addr, std::string work_dir, bool silent) { if (silent) { // Only errors and fatal is logged dmlc::InitLogging("--minloglevel=2"); } // Start the rpc server RPCServer rpc(std::move(host), port, port_end, std::move(tracker_addr), std::move(key), - std::move(custom_addr)); + std::move(custom_addr), std::move(work_dir)); rpc.Start(); } TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body([](TVMArgs args, TVMRetValue* rv) { - RPCServerCreate(args[0], args[1], args[2], args[3], args[4], args[5], args[6]); + RPCServerCreate(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7]); }); } // namespace runtime } // namespace tvm diff --git a/apps/cpp_rpc/rpc_server.h b/apps/cpp_rpc/rpc_server.h index 7a4bda5d65c4..3937ff252f48 100644 --- a/apps/cpp_rpc/rpc_server.h +++ b/apps/cpp_rpc/rpc_server.h @@ -48,11 +48,13 @@ void ServerLoopFromChild(SOCKET socket); * \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="" + * \param work_dir Custom work directory. Default="" * \param silent Whether run in silent mode. Default=True */ void RPCServerCreate(std::string host = "", int port = 9090, int port_end = 9099, std::string tracker_addr = "", std::string key = "", - std::string custom_addr = "", bool silent = true); + std::string custom_addr = "", + std::string work_dir = "", bool silent = true); } // namespace runtime } // namespace tvm #endif // TVM_APPS_CPP_RPC_SERVER_H_ From 606f8fe5f5ff19f42bb77b8bf8b14dfdd6820fd7 Mon Sep 17 00:00:00 2001 From: EunTaik Lee Date: Tue, 16 Mar 2021 19:47:53 +0900 Subject: [PATCH 2/2] clang format --- apps/cpp_rpc/rpc_env.h | 2 +- apps/cpp_rpc/rpc_server.cc | 3 ++- apps/cpp_rpc/rpc_server.h | 3 +-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/apps/cpp_rpc/rpc_env.h b/apps/cpp_rpc/rpc_env.h index 1ca6ddaf79b3..dbb0a62d2c5d 100644 --- a/apps/cpp_rpc/rpc_env.h +++ b/apps/cpp_rpc/rpc_env.h @@ -36,7 +36,7 @@ namespace runtime { */ struct RPCEnv { public: - /*! + /*! * \brief Constructor Init The RPC Environment initialize function */ RPCEnv(const std::string& word_dir = ""); diff --git a/apps/cpp_rpc/rpc_server.cc b/apps/cpp_rpc/rpc_server.cc index 1a5474a71100..52b5da965b4c 100644 --- a/apps/cpp_rpc/rpc_server.cc +++ b/apps/cpp_rpc/rpc_server.cc @@ -309,7 +309,8 @@ class RPCServer { * \param sock The socket information * \param addr The socket address information */ - static void ServerLoopProc(support::TCPSocket sock, support::SockAddr addr, std::string work_dir) { + static void ServerLoopProc(support::TCPSocket sock, support::SockAddr addr, + std::string work_dir) { // Server loop const auto env = RPCEnv(work_dir); RPCServerLoop(int(sock.sockfd)); diff --git a/apps/cpp_rpc/rpc_server.h b/apps/cpp_rpc/rpc_server.h index 3937ff252f48..e4565d095b2e 100644 --- a/apps/cpp_rpc/rpc_server.h +++ b/apps/cpp_rpc/rpc_server.h @@ -53,8 +53,7 @@ void ServerLoopFromChild(SOCKET socket); */ void RPCServerCreate(std::string host = "", int port = 9090, int port_end = 9099, std::string tracker_addr = "", std::string key = "", - std::string custom_addr = "", - std::string work_dir = "", bool silent = true); + std::string custom_addr = "", std::string work_dir = "", bool silent = true); } // namespace runtime } // namespace tvm #endif // TVM_APPS_CPP_RPC_SERVER_H_