1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
/**
* \file src/core/include/megbrain/system.h
* MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
*
* Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
#pragma once
#include "megbrain/utils/metahelper.h"
#include "megbrain/utils/thin/function.h"
#include <memory>
#include <string>
#include <thread>
#include <vector>
#include <cstring>
namespace mgb {
namespace sys {
//! set name of caller thread
MGE_WIN_DECLSPEC_FUC void set_thread_name(const std::string& name);
#if !__DEPLOY_ON_XP_SP2__
/*!
* \brief get name of of given thread
* \param tid thread id, or None to for the caller thread
*/
MGE_WIN_DECLSPEC_FUC std::string get_thread_name(Maybe<std::thread::id> tid = None);
#endif
//! get number of CPU cores on this system
MGE_WIN_DECLSPEC_FUC int get_cpu_count();
//! set cpu affinity for caller thread
MGE_WIN_DECLSPEC_FUC void set_cpu_affinity(const std::vector<int>& cpuset);
//! whether stderr supports ansi color code
MGE_WIN_DECLSPEC_FUC bool stderr_ansi_color();
//! get total ram and free ram in bytes
MGE_WIN_DECLSPEC_FUC std::pair<size_t, size_t> get_ram_status_bytes();
/*!
* \brief invoke a function with time limit
*
* This class should be accessed via the singleton ins().
*
* It is currently used to implement algorithm profiling because:
*
* 1. Some algos may be much slower (sometimes even more than 1000x)
* than others. Therefore we want to set a time limit so current
* algo can take no longer best known time.
* 2. There is no portable and elegant way to interrupt an asynchronous
* function. So here we proceed by invoking the function in a child
* process and kill the whole process on timeout.
* 3. We use fork-exec to launch the child process rather than using a
* simple fork because some device drivers (e.g. CUDA) would be
* broken if we fork without exec.
*
* For SDK developers (i.e. MegBrain users):
*
* 1. TimedFuncInvoker is currently only implemented for linux
* platforms when MGB_BUILD_SLIM_SERVING is disabled.
* 2. You need to implement a fork-exec entry point and call
* TimedFuncInvoker::ins().set_fork_exec_impl() to setup the system.
* 3. An example implementation is available in the python module.
*
* For algorithm profiling implementations (i.e. MegBrain developers):
*
* 1. Register the function to be profiled via register_func(). Use
* invoke() to call it with a timeout.
* 2. You may need AlgoChooserProfileCache to save the profiling
* result.
*
*/
class TimedFuncInvoker : public NonCopyableObj {
friend class TimedFuncInvokerTest;
struct Del {
void operator()(TimedFuncInvoker* p) { delete p; }
};
//! make an instance for test purpose
static std::unique_ptr<TimedFuncInvoker, Del> make_test_ins();
protected:
virtual ~TimedFuncInvoker() = default;
public:
struct Result {
size_t size = 0;
std::unique_ptr<uint8_t[]> data;
template <typename T>
static Result from_pod(const T& val) {
Result ret{sizeof(T), std::make_unique<uint8_t[]>(sizeof(T))};
memcpy(ret.data.get(), &val, sizeof(T));
return ret;
}
template <typename T>
const T& as_single_pod() const {
static_assert(is_location_invariant<T>::value, "bad type");
mgb_assert(sizeof(T) == size);
return *reinterpret_cast<const T*>(data.get());
}
};
struct Param {
size_t size = 0;
const uint8_t* data = nullptr;
// param is non-const ref to ensure caller has ownership; it
// would not be modified
template <typename T>
static Param from_pod(T& val) {
return {sizeof(T), reinterpret_cast<uint8_t*>(&val)};
}
template <typename T>
const T& as_single_pod() const {
static_assert(is_location_invariant<T>::value, "bad type");
mgb_assert(sizeof(T) == size);
return *reinterpret_cast<const T*>(data);
}
};
//! exception thrown by invoke()
class RemoteError final : public MegBrainError {
public:
using MegBrainError::MegBrainError;
};
using Func = thin_function<Result(const Param& param)>;
using FuncInit = thin_function<void(const Param& param)>;
using FuncId = size_t;
/*!
* \brief call fork() and exec(), and pass *arg* to the child
* process, and the child process should pass *arg* back to
* fork_exec_impl_mainloop()
*
* \param arg a null-terminated string argument
* \return child process PID
*/
using ForkExecImpl = thin_function<int(const std::string& arg)>;
/*!
* \brief set the function to implement fork-exec
*
* ForkExecImpl can not be implemented by TimedFuncInvoker because
* it does not know the entry point of the compiled ELF.
*
* This method must be called from this server process, before any
* call to invoke()
*/
virtual void set_fork_exec_impl(const ForkExecImpl& impl) = 0;
/*!
* \brief to be called in the child process by ForkExecImpl
* registered by set_fork_exec_impl()
*
* \param arg the argument passed to ForkExecImpl
*/
[[noreturn]] virtual void fork_exec_impl_mainloop(const char* arg) = 0;
/*!
* \brief register a function that can be invoked
*
* This method must be called both from the server process and from
* the client process. It is usually called during global setup.
*
* \param func the function associated with \p id; its execution
* time can not exceed given timeout
* \param func_init an initializer whose time would not be counted
* in the timeout setting
*/
virtual void register_func(
FuncId id, const Func& func, const FuncInit& func_init = {}) = 0;
/*!
* \brief invoke a function with given timeout
*
* This method must be called from the server process (i.e. main
* mebrain process). This method is thread-safe.
*
* \param timeout timeout in seconds; if it is 0, timeout is
* disabled, and if worker not started yet, the function is
* invoked inplace
* \return the return value if function finishes within given
* timeout; if the function could not finish in time, None
* would be returned.
* \exception RemoteError thrown when function fails on the worker(
* function throws exception, not due to timeout); the error
* message would be forwarded to RemoteError::what()
*/
virtual Maybe<Result> invoke(FuncId id, const Param& param, double timeout) = 0;
/*!
* \brief kill the worker process
*/
virtual void kill_worker() = 0;
//! global unique instance
MGE_WIN_DECLSPEC_FUC static TimedFuncInvoker& ins();
};
} // namespace sys
} // namespace mgb
// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}