@@ -46,7 +46,9 @@ struct TORCH_API InterpreterSession {
4646 // NOLINTNEXTLINE(bugprone-exception-escape)
4747 ~InterpreterSession ();
4848
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.
5052 Obj global (const char * module , const char * name) {
5153 return impl_->global (module , name);
5254 }
@@ -57,6 +59,8 @@ struct TORCH_API InterpreterSession {
5759 // InterpreterSession* I)' instead. We will have no backwards compatibility
5860 // guarentees for this function.
5961 ReplicatedObj createMovable (Obj obj);
62+
63+ // Converts a `ReplicatedObj` to an `Obj` on this InterpreterSession.
6064 Obj fromMovable (const ReplicatedObj& obj);
6165
6266 protected:
@@ -73,6 +77,9 @@ struct TORCH_API InterpreterSession {
7377 std::function<void ()> deconstruction_callback_ = nullptr ;
7478};
7579
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.
7683class TORCH_API Interpreter {
7784 private:
7885 void * handle_;
@@ -84,17 +91,24 @@ class TORCH_API Interpreter {
8491 multipy::optional<EmbeddedFile> torchPluginFile_;
8592
8693 public:
94+ // Creates an Interpreter which is managed by `manager` and using the
95+ // environment `env`
8796 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.
88100 explicit Interpreter (std::shared_ptr<Environment> env)
89101 : Interpreter(nullptr , env) {}
90102
103+ // Gets a new `InterpreterSession` from this Interpreter.
91104 InterpreterSession acquireSession () const {
92105 if (manager_) {
93106 return InterpreterSession (pImpl_->acquireSession (), manager_);
94107 } else {
95108 return InterpreterSession (pImpl_->acquireSession ());
96109 }
97110 }
111+
98112 ~Interpreter ();
99113 Interpreter (Interpreter&& rhs) noexcept
100114 : handle_(rhs.handle_),
@@ -113,17 +127,28 @@ class TORCH_API Interpreter {
113127
114128struct Package ;
115129
130+ // The default LoadBalancer for torch::deploy which handles allocating and
131+ // freeing subinterpreters.
116132struct TORCH_API LoadBalancer {
133+ // Creates a Loadbalancer which handles `n` interpreters.
117134 explicit LoadBalancer (size_t n)
118135 : uses_(new uint64_t [8 * n]), allocated_(n), n_(n) {
119136 // 8*... to avoid false sharing of atomics on the same cache line
120137 memset (uses_.get (), 0 , 8 * n_ * sizeof (uint64_t ));
121138 }
139+
140+ // Changes the amount of subinterpreters which is handled by the load
141+ // balancer.
122142 void setResourceLimit (size_t n) {
123143 MULTIPY_INTERNAL_ASSERT (n <= allocated_);
124144 n_ = n;
125145 }
146+
147+ // Allocates an subinterpreter, and return its ID which is used to free it.
126148 int acquire ();
149+
150+ // Frees the subinterpreter with ID `where`. This ID is returned by
151+ // `LoadBalancer::acquire()`
127152 void free (int where);
128153
129154 private:
@@ -134,13 +159,20 @@ struct TORCH_API LoadBalancer {
134159 size_t n_;
135160};
136161
162+ // An `InterpreterManager` handles the interaction of multiple subinterpreters
163+ // such as allocating subinterpreters, or load balancing the subinterpreters.
137164struct 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.
138168 explicit InterpreterManager (
139169 size_t nInterp = 2 ,
140170 std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141171
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.
144176 InterpreterSession acquireOne () {
145177 int where = resources_.acquire ();
146178 InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +186,19 @@ struct TORCH_API InterpreterManager {
154186 at::ArrayRef<Interpreter> allInstances () {
155187 return instances_;
156188 }
189+
190+ // debugging tool to control the size of the loadBalancer
191+ // and change the number of interpreters on the fly
157192 void debugLimitInterpreters (size_t N) {
158193 AT_ASSERT (N <= instances_.size ());
159194 resources_.setResourceLimit (N);
160195 }
196+
197+ // loads a package from a file with name `uri`
161198 Package loadPackage (const std::string& uri);
199+
200+ // loads a package from a `PyTorchStreamReader` or any class other which uses
201+ // `ReadAdapterInterface`
162202 Package loadPackage (
163203 std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164204
@@ -171,10 +211,12 @@ struct TORCH_API InterpreterManager {
171211 registeredModuleSource_[std::move (name)] = std::move (src);
172212 }
173213
174- // Util function for debugging.
214+ // Util function for debugging which outputs the number of registered modules .
175215 size_t countRegisteredModuleSources () {
176216 return registeredModuleSource_.size ();
177217 }
218+
219+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178220 ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179221 InterpreterManager (const InterpreterManager&) = delete ;
180222 InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -204,33 +246,57 @@ struct TORCH_API ReplicatedObjImpl {
204246 InterpreterManager* manager_;
205247};
206248
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)`
207254struct TORCH_API ReplicatedObj {
255+ // Default constructor for `ReplicatedObj`
208256 ReplicatedObj () : pImpl_(nullptr ) {}
257+
258+ // Creates a new InterpreterSession on onThisInterpreter if specified else
259+ // uses an arbitrary one from InteprreterManager.
209260 InterpreterSession acquireSession (
210261 const Interpreter* onThisInterpreter = nullptr ) const ;
211262 at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212263 auto I = acquireSession ();
213264 return I.self (args).toIValue ();
214265 }
215266
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__`).
216270 [[nodiscard]] at::IValue callKwargs (
217271 std::vector<at::IValue> args,
218272 std::unordered_map<std::string, c10::IValue> kwargs) const {
219273 auto I = acquireSession ();
220274 return I.self .callKwargs (std::move (args), std::move (kwargs)).toIValue ();
221275 }
222276
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__`).
223280 [[nodiscard]] at::IValue callKwargs (
224281 std::unordered_map<std::string, c10::IValue> kwargs) const {
225282 auto I = acquireSession ();
226283 return I.self .callKwargs (std::move (kwargs)).toIValue ();
227284 }
228285
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 {
230290 auto I = acquireSession ();
231- return I.self .hasattr (name );
291+ return I.self .hasattr (attr );
232292 }
293+
294+ // Deletes `ReplicatedObj` from onThisInterpreter, if onThisInterpreter is
295+ // `nullptr`, unload is called on all interpreters belonging to the
296+ // ReplicatedObject's InterpreterManager
233297 void unload (const Interpreter* onThisInterpreter = nullptr );
298+
299+ // Converts `ReplicatedObj` to `Obj` on `InterpreterSession` `I`
234300 Obj toObj (InterpreterSession* I);
235301
236302 private:
@@ -242,21 +308,24 @@ struct TORCH_API ReplicatedObj {
242308 friend struct InterpreterManager ;
243309};
244310
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.
245314class 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.
249315 public:
250316 // TODO(whc) make bound method pickleable, then directly construct from that
317+
251318 PythonMethodWrapper (
252319 torch::deploy::ReplicatedObj model,
253320 std::string methodName)
254321 : model_(std::move(model)), methodName_(std::move(methodName)) {}
255322
323+ // return the name of the python method.
256324 const std::string& name () const override {
257325 return methodName_;
258326 }
259327
328+ // overrides the `()` operater to call the underlying python method.
260329 c10::IValue operator ()(
261330 std::vector<c10::IValue> args,
262331 const IValueMap& kwargs = IValueMap()) const override {
@@ -274,6 +343,8 @@ class PythonMethodWrapper : public torch::IMethod {
274343 std::string methodName_;
275344};
276345
346+ // Package is a wrapper around `torch.package` which allows loading a
347+ // PyTorch model and its dependencies from a package.
277348struct TORCH_API Package {
278349 // shorthand for getting the object as a pickle resource in the package
279350 ReplicatedObj loadPickle (const std::string& module , const std::string& file) {
@@ -308,12 +379,16 @@ struct TORCH_API Package {
308379 }
309380#endif
310381
382+ // Allocates an `InterpreterSession` and load the appropriate torch.package
383+ // with it.
311384 InterpreterSession acquireSession () {
312385 auto I = manager_->acquireOne ();
313386 I.self =
314387 I.impl_ ->createOrGetPackageImporterFromContainerFile (containerFile_);
315388 return I;
316389 }
390+
391+ // Converts an `Obj` from `InterpreterSession` `I` into a `ReplicatedObj`.
317392 ReplicatedObj createMovable (Obj obj, InterpreterSession* I) {
318393 return manager_->createMovable (obj, I);
319394 }
0 commit comments