@@ -46,7 +46,9 @@ struct TORCH_API InterpreterSession {
46
46
// NOLINTNEXTLINE(bugprone-exception-escape)
47
47
~InterpreterSession ();
48
48
49
- // global imports a python object from the specified module.
49
+ // `global` imports a python object from the specified module.
50
+ // Specifically `global` is analogous to "from `module` import `name`"
51
+ // in python.
50
52
Obj global (const char * module, const char * name) {
51
53
return impl_->global (module, name);
52
54
}
@@ -57,6 +59,8 @@ struct TORCH_API InterpreterSession {
57
59
// InterpreterSession* I)' instead. We will have no backwards compatibility
58
60
// guarentees for this function.
59
61
ReplicatedObj createMovable (Obj obj);
62
+
63
+ // Converts a `ReplicatedObj` to an `Obj` on this InterpreterSession.
60
64
Obj fromMovable (const ReplicatedObj& obj);
61
65
62
66
protected:
@@ -73,6 +77,9 @@ struct TORCH_API InterpreterSession {
73
77
std::function<void ()> deconstruction_callback_ = nullptr ;
74
78
};
75
79
80
+ // An `Interpreter` represents an invidual subinterpreter created by
81
+ // `torch::deploy`. It allows for the creation of `InterpreterSession` objects
82
+ // which allow users to interact with python objects.
76
83
class TORCH_API Interpreter {
77
84
private:
78
85
void * handle_;
@@ -84,17 +91,24 @@ class TORCH_API Interpreter {
84
91
multipy::optional<EmbeddedFile> torchPluginFile_;
85
92
86
93
public:
94
+ // Creates an Interpreter which is managed by `manager` and using the
95
+ // environment `env`
87
96
Interpreter (InterpreterManager* manager, std::shared_ptr<Environment> env);
97
+
98
+ // Creates an Interpreter manager using environment `env` which is not tied to
99
+ // an Interpreter Manager.
88
100
explicit Interpreter (std::shared_ptr<Environment> env)
89
101
: Interpreter(nullptr , env) {}
90
102
103
+ // Gets a new `InterpreterSession` from this Interpreter.
91
104
InterpreterSession acquireSession () const {
92
105
if (manager_) {
93
106
return InterpreterSession (pImpl_->acquireSession (), manager_);
94
107
} else {
95
108
return InterpreterSession (pImpl_->acquireSession ());
96
109
}
97
110
}
111
+
98
112
~Interpreter ();
99
113
Interpreter (Interpreter&& rhs) noexcept
100
114
: handle_(rhs.handle_),
@@ -113,17 +127,28 @@ class TORCH_API Interpreter {
113
127
114
128
struct Package ;
115
129
130
+ // The default LoadBalancer for torch::deploy which handles allocating and
131
+ // freeing subinterpreters.
116
132
struct TORCH_API LoadBalancer {
133
+ // Creates a Loadbalancer which handles `n` interpreters.
117
134
explicit LoadBalancer (size_t n)
118
135
: uses_(new uint64_t [8 * n]), allocated_(n), n_(n) {
119
136
// 8*... to avoid false sharing of atomics on the same cache line
120
137
memset (uses_.get (), 0 , 8 * n_ * sizeof (uint64_t ));
121
138
}
139
+
140
+ // Changes the amount of subinterpreters which is handled by the load
141
+ // balancer.
122
142
void setResourceLimit (size_t n) {
123
143
MULTIPY_INTERNAL_ASSERT (n <= allocated_);
124
144
n_ = n;
125
145
}
146
+
147
+ // Allocates an subinterpreter, and return its ID which is used to free it.
126
148
int acquire ();
149
+
150
+ // Frees the subinterpreter with ID `where`. This ID is returned by
151
+ // `LoadBalancer::acquire()`
127
152
void free (int where);
128
153
129
154
private:
@@ -134,13 +159,20 @@ struct TORCH_API LoadBalancer {
134
159
size_t n_;
135
160
};
136
161
162
+ // An `InterpreterManager` handles the interaction of multiple subinterpreters
163
+ // such as allocating subinterpreters, or load balancing the subinterpreters.
137
164
struct TORCH_API InterpreterManager {
165
+ // constructor for `InterpreterManager` which takes the number of interpreters
166
+ // (usually correlates to number of cores on your cpu), and a pointer to an
167
+ // `Environment`. The default uses the local python env.
138
168
explicit InterpreterManager (
139
169
size_t nInterp = 2 ,
140
170
std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141
171
142
- // get a free model, guarenteed that no other user of acquireOne has the same
143
- // model. It _is_ possible that other users will be using the interpreter.
172
+ // Returns a free interpreter or an arbitrary interpreter if there are
173
+ // none free. To ensure data safety it's best to match the number of
174
+ // calling threads to the size of the interpreter pool to avoid
175
+ // sharing an interpreter.
144
176
InterpreterSession acquireOne () {
145
177
int where = resources_.acquire ();
146
178
InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +186,19 @@ struct TORCH_API InterpreterManager {
154
186
at::ArrayRef<Interpreter> allInstances () {
155
187
return instances_;
156
188
}
189
+
190
+ // debugging tool to control the size of the loadBalancer
191
+ // and change the number of interpreters on the fly
157
192
void debugLimitInterpreters (size_t N) {
158
193
AT_ASSERT (N <= instances_.size ());
159
194
resources_.setResourceLimit (N);
160
195
}
196
+
197
+ // loads a package from a file with name `uri`
161
198
Package loadPackage (const std::string& uri);
199
+
200
+ // loads a package from a `PyTorchStreamReader` or any class other which uses
201
+ // `ReadAdapterInterface`
162
202
Package loadPackage (
163
203
std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164
204
@@ -171,10 +211,12 @@ struct TORCH_API InterpreterManager {
171
211
registeredModuleSource_[std::move (name)] = std::move (src);
172
212
}
173
213
174
- // Util function for debugging.
214
+ // Util function for debugging which outputs the number of registered modules .
175
215
size_t countRegisteredModuleSources () {
176
216
return registeredModuleSource_.size ();
177
217
}
218
+
219
+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178
220
ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179
221
InterpreterManager (const InterpreterManager&) = delete ;
180
222
InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -204,33 +246,57 @@ struct TORCH_API ReplicatedObjImpl {
204
246
InterpreterManager* manager_;
205
247
};
206
248
249
+ // ReplicatedObj represents a python object that can be used on multiple
250
+ // interpreters. Calling methods on this will pick an arbitrary interpreter
251
+ // to run on, transfer it there if not already and run the method. A
252
+ // replicated object can be converted to an interpreter specific `Obj`
253
+ // using `InterpreterSession::fromMovable(ReplicatedObj)`
207
254
struct TORCH_API ReplicatedObj {
255
+ // Default constructor for `ReplicatedObj`
208
256
ReplicatedObj () : pImpl_(nullptr ) {}
257
+
258
+ // Creates a new InterpreterSession on onThisInterpreter if specified else
259
+ // uses an arbitrary one from InteprreterManager.
209
260
InterpreterSession acquireSession (
210
261
const Interpreter* onThisInterpreter = nullptr ) const ;
211
262
at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212
263
auto I = acquireSession ();
213
264
return I.self (args).toIValue ();
214
265
}
215
266
267
+ // Invokes the Python function or class on an arbitrary interpreter with
268
+ // arguments given by the tuple args and named arguments given by the
269
+ // dictionary kwargs (equivalent to python's `__call__`).
216
270
[[nodiscard]] at::IValue callKwargs (
217
271
std::vector<at::IValue> args,
218
272
std::unordered_map<std::string, c10::IValue> kwargs) const {
219
273
auto I = acquireSession ();
220
274
return I.self .callKwargs (std::move (args), std::move (kwargs)).toIValue ();
221
275
}
222
276
277
+ // Invokes the Python function or class on an arbitrary interpreter.with
278
+ // named arguments given by the dictionary kwargs (equivalent to python's
279
+ // `__call__`).
223
280
[[nodiscard]] at::IValue callKwargs (
224
281
std::unordered_map<std::string, c10::IValue> kwargs) const {
225
282
auto I = acquireSession ();
226
283
return I.self .callKwargs (std::move (kwargs)).toIValue ();
227
284
}
228
285
229
- [[nodiscard]] bool hasattr (const char * name) const {
286
+ // Returns true if `ReplicatedObj` has attribute with name `attr` and false
287
+ // otherwise. This is done on an arbitrary `InterpreterSession` which belongs
288
+ // to the `ReplicatedObj`'s manager.
289
+ [[nodiscard]] bool hasattr (const char * attr) const {
230
290
auto I = acquireSession ();
231
- return I.self .hasattr (name );
291
+ return I.self .hasattr (attr );
232
292
}
293
+
294
+ // Deletes `ReplicatedObj` from onThisInterpreter, if onThisInterpreter is
295
+ // `nullptr`, unload is called on all interpreters belonging to the
296
+ // ReplicatedObject's InterpreterManager
233
297
void unload (const Interpreter* onThisInterpreter = nullptr );
298
+
299
+ // Converts `ReplicatedObj` to `Obj` on `InterpreterSession` `I`
234
300
Obj toObj (InterpreterSession* I);
235
301
236
302
private:
@@ -242,21 +308,24 @@ struct TORCH_API ReplicatedObj {
242
308
friend struct InterpreterManager ;
243
309
};
244
310
311
+ // PythonMethodWrapper is a more specific instance of a
312
+ // ReplicatedObj which represents a python method, and
313
+ // is therefore callable and has argument names accessible.
245
314
class PythonMethodWrapper : public torch ::IMethod {
246
- // PythonMethodWrapper is a more specific instance of a
247
- // ReplicatedObj which represents a python method, and
248
- // is therefore callable and has argument names accessible.
249
315
public:
250
316
// TODO(whc) make bound method pickleable, then directly construct from that
317
+
251
318
PythonMethodWrapper (
252
319
torch::deploy::ReplicatedObj model,
253
320
std::string methodName)
254
321
: model_(std::move(model)), methodName_(std::move(methodName)) {}
255
322
323
+ // return the name of the python method.
256
324
const std::string& name () const override {
257
325
return methodName_;
258
326
}
259
327
328
+ // overrides the `()` operater to call the underlying python method.
260
329
c10::IValue operator ()(
261
330
std::vector<c10::IValue> args,
262
331
const IValueMap& kwargs = IValueMap()) const override {
@@ -274,6 +343,8 @@ class PythonMethodWrapper : public torch::IMethod {
274
343
std::string methodName_;
275
344
};
276
345
346
+ // Package is a wrapper around `torch.package` which allows loading a
347
+ // PyTorch model and its dependencies from a package.
277
348
struct TORCH_API Package {
278
349
// shorthand for getting the object as a pickle resource in the package
279
350
ReplicatedObj loadPickle (const std::string& module, const std::string& file) {
@@ -308,12 +379,16 @@ struct TORCH_API Package {
308
379
}
309
380
#endif
310
381
382
+ // Allocates an `InterpreterSession` and load the appropriate torch.package
383
+ // with it.
311
384
InterpreterSession acquireSession () {
312
385
auto I = manager_->acquireOne ();
313
386
I.self =
314
387
I.impl_ ->createOrGetPackageImporterFromContainerFile (containerFile_);
315
388
return I;
316
389
}
390
+
391
+ // Converts an `Obj` from `InterpreterSession` `I` into a `ReplicatedObj`.
317
392
ReplicatedObj createMovable (Obj obj, InterpreterSession* I) {
318
393
return manager_->createMovable (obj, I);
319
394
}
0 commit comments