From 7761416f00aed73e93bf919036625fc3f7941cbf Mon Sep 17 00:00:00 2001
From: eqy <eqy@cs.washington.edu>
Date: Mon, 19 Nov 2018 12:35:15 -0800
Subject: [PATCH] [WIP] [RPC] clean up uploaded modules (#2121)
[RPC] clean up uploaded modules
---
python/tvm/autotvm/measure/measure_methods.py | 6 ++++++
python/tvm/rpc/client.py | 13 +++++++++++++
src/runtime/file_util.cc | 4 ++++
src/runtime/file_util.h | 6 ++++++
src/runtime/rpc/rpc_server_env.cc | 7 +++++++
5 files changed, 36 insertions(+)
diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py
index 802abe019..ff93704ed 100644
--- a/python/tvm/autotvm/measure/measure_methods.py
+++ b/python/tvm/autotvm/measure/measure_methods.py
@@ -467,6 +467,12 @@ def run_through_rpc(measure_input, build_result,
ctx.sync()
costs = time_f(*args).results
+
+ # clean up remote files
+ remote.remove(build_result.filename)
+ remote.remove(os.path.splitext(build_result.filename)[0] + '.so')
+ remote.remove('')
+
if len(costs) > 2: # remove largest and smallest value to reduce variance
costs = list(costs)
costs.sort()
diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py
index ae44e5a79..c975ec64a 100644
--- a/python/tvm/rpc/client.py
+++ b/python/tvm/rpc/client.py
@@ -103,6 +103,19 @@ class RPCSession(object):
"tvm.rpc.server.download")
return self._remote_funcs["download"](path)
+ def remove(self, path):
+ """Remove file from remote temp folder.
+
+ Parameters
+ ----------
+ path: str
+ The relative location to remote temp folder.
+ """
+ if "remove" not in self._remote_funcs:
+ self._remote_funcs["remove"] = self.get_function(
+ "tvm.rpc.server.remove")
+ self._remote_funcs["remove"](path)
+
def load_module(self, path):
"""Load a remote module, the file need to be uploaded first.
diff --git a/src/runtime/file_util.cc b/src/runtime/file_util.cc
index 4df335a54..ff579d121 100644
--- a/src/runtime/file_util.cc
+++ b/src/runtime/file_util.cc
@@ -142,5 +142,9 @@ void LoadMetaDataFromFile(
fs.close();
}
+void RemoveFile(const std::string& file_name) {
+ std::remove(file_name.c_str());
+}
+
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/file_util.h b/src/runtime/file_util.h
index de520fa31..2b7976142 100644
--- a/src/runtime/file_util.h
+++ b/src/runtime/file_util.h
@@ -71,6 +71,12 @@ void SaveMetaDataToFile(
void LoadMetaDataFromFile(
const std::string& file_name,
std::unordered_map<std::string, FunctionInfo>* fmap);
+
+/*!
+ * \brief Remove (unlink) a file.
+ * \param file_name The file name.
+ */
+void RemoveFile(const std::string& file_name);
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_FILE_UTIL_H_
diff --git a/src/runtime/rpc/rpc_server_env.cc b/src/runtime/rpc/rpc_server_env.cc
index ca91b8824..fb8d95d60 100644
--- a/src/runtime/rpc/rpc_server_env.cc
+++ b/src/runtime/rpc/rpc_server_env.cc
@@ -35,5 +35,12 @@ TVM_REGISTER_GLOBAL("tvm.rpc.server.download")
*rv = arr;
});
+TVM_REGISTER_GLOBAL("tvm.rpc.server.remove")
+.set_body([](TVMArgs args, TVMRetValue *rv) {
+ std::string file_name = RPCGetPath(args[0]);
+ LOG(INFO) << "Remove " << file_name;
+ RemoveFile(file_name);
+ });
+
} // namespace runtime
} // namespace tvm
--
GitLab