@@ -57,6 +57,8 @@ struct TORCH_API InterpreterSession {
57
57
// InterpreterSession* I)' instead. We will have no backwards compatibility
58
58
// guarentees for this function.
59
59
ReplicatedObj createMovable (Obj obj);
60
+
61
+ // Converts a `ReplicatedObj` to an `Obj` on this InterpreterSession.
60
62
Obj fromMovable (const ReplicatedObj& obj);
61
63
62
64
protected:
@@ -73,6 +75,9 @@ struct TORCH_API InterpreterSession {
73
75
std::function<void ()> deconstruction_callback_ = nullptr ;
74
76
};
75
77
78
+ // An `Interpreter` represents an invidual subinterpreter created by
79
+ // `torch::deploy`. It allows for the creation of `InterpreterSession` objects
80
+ // which allow users to interact with python objects.
76
81
class TORCH_API Interpreter {
77
82
private:
78
83
void * handle_;
@@ -84,17 +89,24 @@ class TORCH_API Interpreter {
84
89
multipy::optional<EmbeddedFile> torchPluginFile_;
85
90
86
91
public:
92
+ // Creates an Interpreter which is managed by `manager` and using the
93
+ // environment `env`
87
94
Interpreter (InterpreterManager* manager, std::shared_ptr<Environment> env);
95
+
96
+ // Creates an Interpreter manager using environment `env` which is not tied to
97
+ // an Interpreter Manager.
88
98
explicit Interpreter (std::shared_ptr<Environment> env)
89
99
: Interpreter(nullptr , env) {}
90
100
101
+ // Gets a new `InterpreterSession` from this Interpreter.
91
102
InterpreterSession acquireSession () const {
92
103
if (manager_) {
93
104
return InterpreterSession (pImpl_->acquireSession (), manager_);
94
105
} else {
95
106
return InterpreterSession (pImpl_->acquireSession ());
96
107
}
97
108
}
109
+
98
110
~Interpreter ();
99
111
Interpreter (Interpreter&& rhs) noexcept
100
112
: handle_(rhs.handle_),
@@ -113,17 +125,28 @@ class TORCH_API Interpreter {
113
125
114
126
struct Package ;
115
127
128
+ // The default LoadBalancer for torch::deploy which handles allocating and
129
+ // freeing subinterpreters.
116
130
struct TORCH_API LoadBalancer {
131
+ // Creates a Loadbalancer which handles `n` interpreters.
117
132
explicit LoadBalancer (size_t n)
118
133
: uses_(new uint64_t [8 * n]), allocated_(n), n_(n) {
119
134
// 8*... to avoid false sharing of atomics on the same cache line
120
135
memset (uses_.get (), 0 , 8 * n_ * sizeof (uint64_t ));
121
136
}
137
+
138
+ // Changes the amount of subinterpreters which is handled by the load
139
+ // balancer.
122
140
void setResourceLimit (size_t n) {
123
141
MULTIPY_INTERNAL_ASSERT (n <= allocated_);
124
142
n_ = n;
125
143
}
144
+
145
+ // Allocates an subinterpreter, and return its ID which is used to free it.
126
146
int acquire ();
147
+
148
+ // Frees the subinterpreter with ID `where`. This ID is returned by
149
+ // `LoadBalancer::acquire()`
127
150
void free (int where);
128
151
129
152
private:
@@ -134,13 +157,21 @@ struct TORCH_API LoadBalancer {
134
157
size_t n_;
135
158
};
136
159
160
+ // An `InterpreterManager` handles the interaction of multiple subinterpreters
161
+ // such as allocating subinterpreters, or load balancing the subinterpreters.
137
162
struct TORCH_API InterpreterManager {
163
+ // constructor for `InterpreterManager` which takes the number of interpreters
164
+ // (usually correlates to number of cores on your cpu), and a pointer to an
165
+ // `Environment`. The default uses the local python env.
138
166
explicit InterpreterManager (
139
167
size_t nInterp = 2 ,
140
168
std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141
169
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.
170
+ // get a free InterpreterSession, guarenteed that no other user of acquireOne
171
+ // has the same InterpreterSession. It _is_ possible that other users will be
172
+ // using the interpreter if there are no free InterpreterSessions. Unless you
173
+ // are very careful to only use free interpreters, then do not assume that the
174
+ // `Obj`s are isolated from each other.
144
175
InterpreterSession acquireOne () {
145
176
int where = resources_.acquire ();
146
177
InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +185,19 @@ struct TORCH_API InterpreterManager {
154
185
at::ArrayRef<Interpreter> allInstances () {
155
186
return instances_;
156
187
}
188
+
189
+ // debugging tool to control the size of the loadBalancer
190
+ // and change the number of interpreters on the fly
157
191
void debugLimitInterpreters (size_t N) {
158
192
AT_ASSERT (N <= instances_.size ());
159
193
resources_.setResourceLimit (N);
160
194
}
195
+
196
+ // loads a package from a file with name `uri`
161
197
Package loadPackage (const std::string& uri);
198
+
199
+ // loads a package from a `PyTorchStreamReader` or any class other which uses
200
+ // `ReadAdapterInterface`
162
201
Package loadPackage (
163
202
std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164
203
@@ -171,10 +210,12 @@ struct TORCH_API InterpreterManager {
171
210
registeredModuleSource_[std::move (name)] = std::move (src);
172
211
}
173
212
174
- // Util function for debugging.
213
+ // Util function for debugging which outputs the number of registered modules .
175
214
size_t countRegisteredModuleSources () {
176
215
return registeredModuleSource_.size ();
177
216
}
217
+
218
+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178
219
ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179
220
InterpreterManager (const InterpreterManager&) = delete ;
180
221
InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -204,33 +245,58 @@ struct TORCH_API ReplicatedObjImpl {
204
245
InterpreterManager* manager_;
205
246
};
206
247
248
+ // ReplicatedObj represents a python object that can be used on multiple interpreters. Calling
249
+ // methods on this will pick an arbitrary interpreter to run on, transfer it there if not already
250
+ // and run the method. A replicated object can be converted to an interpreter specific `Obj` using
251
+ // `InterpreterSession::fromMovable(ReplicatedObj)`
207
252
struct TORCH_API ReplicatedObj {
253
+ // Default constructor for `ReplicatedObj`
208
254
ReplicatedObj () : pImpl_(nullptr ) {}
255
+
256
+ // Creates an `InterpreterSession` using `onThisInterpreter`. If
257
+ // `onThisInterpreter` is a `nullptr', then the associated
258
+ // `InterpreterManager` allocates it.
209
259
InterpreterSession acquireSession (
210
260
const Interpreter* onThisInterpreter = nullptr ) const ;
211
261
at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212
262
auto I = acquireSession ();
213
263
return I.self (args).toIValue ();
214
264
}
215
265
266
+ // Calls an `ReplicatedObj` callable, with arguments given by the tuple args
267
+ // and named arguments given by the dictionary kwargs (equivalent to python's `__call__`). This is done on an
268
+ // arbitrary `InterpreterSession` which belongs to the `ReplicatedObj`'s
269
+ // manager.
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
+ // Calls an `ReplicatedObj` callable, with named arguments given by the
278
+ // dictionary kwargs (equivalent to python's `__call__`). This is done on an arbitrary `InterpreterSession` which
279
+ // belongs to the `ReplicatedObj`'s manager.
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
+ // An object to encapsulate a `torch::package` which can act as part (or entire)
347
+ // environment for subinterpreters.
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